From 8b3d8d3f1ea713779db89830adc535d1496f677a Mon Sep 17 00:00:00 2001
From: Alex Streed <alex.s@prefect.io>
Date: Tue, 25 Feb 2025 11:26:45 -0600
Subject: [PATCH] Lil' cleanup

---
 .../{_decorators.py => experimental/decorators.py}             | 3 ++-
 .../prefect-kubernetes/tests/experimental/test_decorator.py    | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)
 rename src/integrations/prefect-kubernetes/prefect_kubernetes/{_decorators.py => experimental/decorators.py} (98%)

diff --git a/src/integrations/prefect-kubernetes/prefect_kubernetes/_decorators.py b/src/integrations/prefect-kubernetes/prefect_kubernetes/experimental/decorators.py
similarity index 98%
rename from src/integrations/prefect-kubernetes/prefect_kubernetes/_decorators.py
rename to src/integrations/prefect-kubernetes/prefect_kubernetes/experimental/decorators.py
index 23ba83e26fc9..dcb061585c28 100644
--- a/src/integrations/prefect-kubernetes/prefect_kubernetes/_decorators.py
+++ b/src/integrations/prefect-kubernetes/prefect_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
index 2ef43e638609..fc7c3ac5f222 100644
--- 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