Lil' cleanup

PrefectHQ · Feb 25, 2025 · 8b3d8d3 · 8b3d8d3
1 parent 29934e9
commit 8b3d8d3
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/...ernetes/prefect_kubernetes/_decorators.py → ...ect_kubernetes/experimental/decorators.py b/...ernetes/prefect_kubernetes/_decorators.py → ...ect_kubernetes/experimental/decorators.py
@@ -13,13 +13,13 @@
     overload,
 )
 
+from prefect_kubernetes.worker import KubernetesWorker
 from typing_extensions import Literal, ParamSpec
 
 from prefect import Flow, State
 from prefect.futures import PrefectFuture
 from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.callables import get_call_parameters
-from prefect_kubernetes.worker import KubernetesWorker
 
 P = ParamSpec("P")
 R = TypeVar("R")
@@ -106,6 +106,7 @@ def __call__(
         async def modified_call(
             *args: P.args,
             return_state: bool = False,
+            # TODO: Handle wait_for once we have an asynchronous way to wait for futures
             wait_for: Optional[Iterable[PrefectFuture[Any]]] = None,
             **kwargs: P.kwargs,
         ) -> R | State[R]:

diff --git a/src/integrations/prefect-kubernetes/tests/experimental/test_decorator.py b/src/integrations/prefect-kubernetes/tests/experimental/test_decorator.py
@@ -2,7 +2,7 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from prefect_kubernetes._decorators import kubernetes
+from prefect_kubernetes.experimental.decorators import kubernetes
 from prefect_kubernetes.worker import KubernetesWorker
 
 from prefect import State, flow