Add integration tests for pod events

Author: desertaxle

src/integrations/prefect-kubernetes/integration_tests/src/prefect_kubernetes_integration_tests/test_events.py

@@ -0,0 +1,54 @@
+import asyncio
+
+import anyio
+import pytest
+
+from prefect import get_client
+from prefect.states import StateType
+from prefect_kubernetes_integration_tests.utils import display, prefect_core
+
+DEFAULT_JOB_VARIABLES = {
+    "image": "prefecthq/prefect:3.2.11-python3.12",
+}
+DEFAULT_PARAMETERS = {"n": 5}
+# Default source is a simple flow that sleeps
+DEFAULT_FLOW_SOURCE = "https://gist.github.com/772d095672484b76da40a4e6158187f0.git"
+DEFAULT_FLOW_ENTRYPOINT = "sleeping.py:sleepy"
+DEFAULT_FLOW_NAME = "pod-eviction-test"
+
+
+@pytest.mark.usefixtures("kind_cluster")
+async def test_happy_path_events(
+    work_pool_name: str,
+):
+    """Test that flow runs properly handle pod evictions."""
+    flow_run = await prefect_core.create_flow_run(
+        source=DEFAULT_FLOW_SOURCE,
+        entrypoint=DEFAULT_FLOW_ENTRYPOINT,
+        name=DEFAULT_FLOW_NAME,
+        work_pool_name=work_pool_name,
+        job_variables=DEFAULT_JOB_VARIABLES,
+        parameters=DEFAULT_PARAMETERS,
+    )
+
+    display.print_flow_run_created(flow_run)
+
+    prefect_core.start_worker(work_pool_name, run_once=True)
+
+    async with get_client() as client:
+        updated_flow_run = await client.read_flow_run(flow_run.id)
+
+        assert updated_flow_run.state is not None, "Flow run state should not be None"
+        assert updated_flow_run.state.type == StateType.COMPLETED, (
+            "Expected flow run to be COMPLETED. Got "
+            f"{updated_flow_run.state.type} instead."
+        )
+
+        display.print_flow_run_result(updated_flow_run)
+
+    events = []
+    with anyio.move_on_after(10):
+        while len(events) < 3:
+            events = await prefect_core.read_pod_events_for_flow_run(flow_run.id)
+            await asyncio.sleep(1)
+    assert len(events) == 3, "Expected 3 events"

src/integrations/prefect-kubernetes/integration_tests/src/prefect_kubernetes_integration_tests/utils/prefect_core.py

@@ -9,6 +9,7 @@
 
 from prefect import flow, get_client
 from prefect.client.schemas.objects import FlowRun
+from prefect.events.schemas.events import Event
 from prefect.states import StateType
 
 console = Console()
@@ -80,3 +81,21 @@ def wait_for_flow_run_state(
             raise TimeoutError(
                 f"Flow run {flow_run_id} did not reach state {target_state} within {timeout} seconds"
             )
+
+
+async def read_pod_events_for_flow_run(flow_run_id: UUID) -> list[Event]:
+    """Read events for a flow run."""
+    async with get_client() as client:
+        response = await client.request(
+            "POST",
+            "/events/filter",
+            json={
+                "filter": {
+                    "event": {"prefix": ["prefect.kubernetes.pod"]},
+                    "related": {
+                        "id": [f"prefect.flow-run.{flow_run_id}"],
+                    },
+                },
+            },
+        )
+        return [Event.model_validate(event) for event in response.json()["events"]]