Fix gateway run ID polling and pytest plugin loading

Author: anth-volk

libs/policyengine-fastapi/tests/conftest.py

@@ -1,6 +1,8 @@
 import sys
 from pathlib import Path
 
+pytest_plugins = ("fixtures.ping.shared",)
+
 library_root = Path(__file__).parent.parent
 if str(library_root) not in sys.path:
     sys.path.insert(0, str(library_root))

libs/policyengine-fastapi/tests/ping/conftest.py

@@ -60,10 +60,15 @@ def _build_policyengine_bundle(
     )
 
 
-def _serialize_job_metadata(resolved_app_name: str, bundle: PolicyEngineBundle) -> dict:
+def _serialize_job_metadata(
+    resolved_app_name: str,
+    bundle: PolicyEngineBundle,
+    run_id: str | None = None,
+) -> dict:
     return {
         "resolved_app_name": resolved_app_name,
         "policyengine_bundle": bundle.model_dump(),
+        "run_id": run_id,
     }
 
 
@@ -136,7 +141,7 @@ async def submit_simulation(request: SimulationRequest):
     call = sim_func.spawn(payload)
 
     bundle = _build_policyengine_bundle(request.country, resolved_version, payload)
-    job_metadata = _serialize_job_metadata(app_name, bundle)
+    job_metadata = _serialize_job_metadata(app_name, bundle, run_id)
     _job_metadata_store()[call.object_id] = job_metadata
 
     return JobSubmitResponse(
@@ -151,7 +156,11 @@ async def submit_simulation(request: SimulationRequest):
     )
 
 
-@router.get("/jobs/{job_id}")
+@router.get(
+    "/jobs/{job_id}",
+    response_model=JobStatusResponse,
+    response_model_exclude_none=True,
+)
 async def get_job_status(job_id: str):
     """
     Poll for job status.

projects/policyengine-api-simulation/src/modal/gateway/endpoints.py

@@ -63,6 +63,7 @@ class JobStatusResponse(BaseModel):
     error: Optional[str] = None
     resolved_app_name: Optional[str] = None
     policyengine_bundle: Optional[PolicyEngineBundle] = None
+    run_id: Optional[str] = None
 
 
 class PingRequest(BaseModel):

projects/policyengine-api-simulation/src/modal/gateway/models.py

@@ -363,10 +363,36 @@ def test__given_submitted_job__then_job_status_includes_bundle_metadata(
         assert response.status_code == 200
         data = response.json()
         assert data["status"] == "complete"
+        assert "run_id" not in data
         assert data["resolved_app_name"] == "policyengine-simulation-us1-500-0-uk2-66-0"
         assert data["policyengine_bundle"] == {
             "model_version": "1.500.0",
-            "policyengine_version": None,
-            "data_version": None,
             "dataset": "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5@1.77.0",
         }
+
+    def test__given_submitted_job_with_telemetry__then_polling_echoes_run_id(
+        self, mock_modal, client: TestClient
+    ):
+        mock_modal["dicts"]["simulation-api-us-versions"] = {
+            "latest": "1.500.0",
+            "1.500.0": "policyengine-simulation-us1-500-0-uk2-66-0",
+        }
+
+        submit_response = client.post(
+            "/simulate/economy/comparison",
+            json={
+                "country": "us",
+                "scope": "macro",
+                "reform": {},
+                "_telemetry": {
+                    "run_id": "run-123",
+                    "process_id": "proc-123",
+                    "capture_mode": "disabled",
+                },
+            },
+        )
+
+        response = client.get(f"/jobs/{submit_response.json()['job_id']}")
+
+        assert response.status_code == 200
+        assert response.json()["run_id"] == "run-123"