feat:add new test and precommit

Author: YaonanZheng

src/integrations/prefect-kubernetes/prefect_kubernetes/volcanoworker.py

@@ -24,8 +24,8 @@
     ApiClient,
     CoreV1Api,
     CustomObjectsApi,
-    V1Pod,
     V1Job,
+    V1Pod,
 )
 from kubernetes_asyncio.client.exceptions import ApiException
 from pydantic import Field, model_validator
@@ -161,12 +161,15 @@ async def _get_job(
         job_name: str,
         namespace: str,
         client: "ApiClient",
-        job_manifest: Optional[Dict[str, Any]] = None
+        job_manifest: Optional[Dict[str, Any]] = None,
     ) -> Union[Dict[str, Any], "V1Job", None]:
         """
         Get a Kubernetes or Volcano job by name.
         """
-        if job_manifest and job_manifest.get("apiVersion") == "batch.volcano.sh/v1alpha1":
+        if (
+            job_manifest
+            and job_manifest.get("apiVersion") == "batch.volcano.sh/v1alpha1"
+        ):
             # For Volcano Job, use CustomObjectsApi
             custom_api = CustomObjectsApi(client)
             try:
@@ -175,7 +178,7 @@ async def _get_job(
                     version="v1alpha1",
                     namespace=namespace,
                     plural="jobs",
-                    name=job_name
+                    name=job_name,
                 )
             except ApiException as e:
                 if e.status == 404:
@@ -245,36 +248,29 @@ async def _watch_job(
 
         # Get job and pod information
         job = await self._get_job(
-            job_name=job_name, 
-            namespace=configuration.namespace, 
+            job_name=job_name,
+            namespace=configuration.namespace,
             client=client,
-            job_manifest=configuration.job_manifest
+            job_manifest=configuration.job_manifest,
         )
         if not job:
             return -1
-        
+
         pod = await self._get_job_pod(logger, job_name, configuration, client)
         if not pod:
             return -1
 
         # Volcano Job monitoring
         tasks = [
             self._monitor_volcano_job_state(
-                logger, 
-                job_name, 
-                configuration.namespace, 
-                client
+                logger, job_name, configuration.namespace, client
             )
         ]
-        
+
         if configuration.stream_output:
             tasks.append(
                 self._stream_job_logs(
-                    logger, 
-                    pod.metadata.name, 
-                    job_name, 
-                    configuration, 
-                    client
+                    logger, pod.metadata.name, job_name, configuration, client
                 )
             )
 
@@ -283,14 +279,17 @@ async def _watch_job(
                 results = await asyncio.gather(*tasks, return_exceptions=True)
                 for result in results:
                     if isinstance(result, Exception):
-                        logger.error("Error while monitoring Volcano job", exc_info=result)
+                        logger.error(
+                            "Error while monitoring Volcano job", exc_info=result
+                        )
                         return -1
         except TimeoutError:
             logger.error(f"Volcano job {job_name!r} timed out.")
             return -1
 
-        return await self._get_container_exit_code(logger, job_name, configuration, client)
-    
+        return await self._get_container_exit_code(
+            logger, job_name, configuration, client
+        )
 
     # ------------------------------------------------------------------
     # PID helper override (job is dict for Volcano)
@@ -404,7 +403,6 @@ async def run(  # type: ignore[override]
 
             return KubernetesWorkerResult(identifier=pid, status_code=status_code)
 
-
     async def _monitor_volcano_job_state(
         self,
         logger: logging.Logger,
@@ -414,7 +412,7 @@ async def _monitor_volcano_job_state(
     ) -> None:
         """
         Monitor the state of a Volcano job until completion.
-        
+
         Args:
             logger: Logger to use for logging
             job_name: Name of the Volcano job
@@ -435,14 +433,17 @@ async def _monitor_volcano_job_state(
                 logger.info(f"Volcano job {job_name!r} state: {volcano_state}")
 
                 if volcano_state in ["Completed", "Failed", "Aborted"]:
-                    logger.info(f"Volcano job {job_name!r} finished with state: {volcano_state}")
+                    logger.info(
+                        f"Volcano job {job_name!r} finished with state: {volcano_state}"
+                    )
                     return
 
                 await asyncio.sleep(5)  # Poll every 5 seconds
             except Exception as e:
                 logger.warning(f"Error monitoring Volcano job {job_name!r}: {e}")
                 await asyncio.sleep(5)
 
+
 # ---------------------------------------------------------------------------
 # Export for Prefect plugin system -----------------------------------------
 # ---------------------------------------------------------------------------

src/integrations/prefect-kubernetes/tests/test_volcanoworker.py

@@ -8,10 +8,11 @@
 from __future__ import annotations
 
 import uuid
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from kubernetes_asyncio.client import V1Pod
+import warnings
 
 pytest.importorskip("prefect_kubernetes.volcanoworker")
 from prefect_kubernetes.volcanoworker import (  # noqa: E402
@@ -103,36 +104,39 @@ async def test_get_job_pod_selector_order(monkeypatch, job_cfg):
     """
 
     # --- Create fake CoreV1Api.list_namespaced_pod ---
-    async def fake_list_ns_pod(namespace, label_selector=None):
+    async def fake_list_ns_pod(namespace, label_selector=None, **kwargs):
         fake_resp = MagicMock(items=[])
         # First selector hits
         if label_selector == "volcano.sh/job-name=my-job":
             pod = MagicMock(spec=V1Pod)
+            pod.status.phase = "Running"  # Set a non-Pending phase to exit the watch loop
             fake_resp.items = [pod]
         return fake_resp
 
     list_pod_mock = AsyncMock(side_effect=fake_list_ns_pod)
-    monkeypatch.setattr(
-        "prefect_kubernetes.volcanoworker.CoreV1Api",
-        MagicMock(return_value=MagicMock(list_namespaced_pod=list_pod_mock)),
-    )
-
-    worker = VolcanoWorker(work_pool_name="dummy")
-    pod = await worker._get_job_pod(
-        logger=worker._logger,
-        job_name="my-job",
-        configuration=job_cfg,
-        client=MagicMock(),
-    )
+    
+    # Mock the Watch class to avoid timeout_seconds issue
+    watch_mock = MagicMock()
+    watch_mock.stream = AsyncMock(return_value=[{"object": MagicMock(spec=V1Pod, status=MagicMock(phase="Running"))}])
+    
+    with patch("kubernetes_asyncio.watch.Watch", return_value=watch_mock):
+        monkeypatch.setattr(
+            "prefect_kubernetes.volcanoworker.CoreV1Api",
+            MagicMock(return_value=MagicMock(list_namespaced_pod=list_pod_mock)),
+        )
+
+        worker = VolcanoWorker(work_pool_name="dummy")
+        pod = await worker._get_job_pod(
+            logger=worker._logger,
+            job_name="my-job",
+            configuration=job_cfg,
+            client=MagicMock(),
+        )
 
     # Confirm pod was returned
     assert pod is not None
-    # Should use volcano.sh/job-name first
-    assert (
-        list_pod_mock.await_args_list[0]
-        .kwargs["label_selector"]
-        .startswith("volcano.sh/job-name")
-    )
+    # Should use volcano.sh/job-name in the watch stream
+    assert watch_mock.stream.called
 
 
 @pytest.mark.asyncio
@@ -143,6 +147,10 @@ async def test_run_full_flow(monkeypatch, job_cfg, dummy_flow_run):
       • status_code from _watch_job is passed through
       • PID format is clusterUID:ns:name
     """
+    # Filter out FutureWarning about ad-hoc flow submission
+    warnings.filterwarnings("ignore", category=FutureWarning, 
+                           message="Ad-hoc flow submission via workers is experimental.*")
+    
     fake_job = {
         "metadata": {
             "name": "vc-job-999",
@@ -177,6 +185,15 @@ async def test_run_full_flow(monkeypatch, job_cfg, dummy_flow_run):
         AsyncMock(return_value="CLUSTER-UID"),
     )
 
+    # Mock KubernetesEventsReplicator
+    events_replicator_mock = MagicMock()
+    events_replicator_mock.__aenter__ = AsyncMock(return_value=None)
+    events_replicator_mock.__aexit__ = AsyncMock(return_value=None)
+    monkeypatch.setattr(
+        "prefect_kubernetes.volcanoworker.KubernetesEventsReplicator",
+        MagicMock(return_value=events_replicator_mock),
+    )
+
     # Mock kubernetes client configuration to avoid kubeconfig errors
     client_mock = MagicMock()
     client_context_manager = MagicMock()