From bfb4fe29948ea91da98bc94964342a190368d01a Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Tue, 13 Oct 2026 14:37:28 -0400 Subject: [PATCH 1/2] Relay simulation run stage and progress on computing economy responses While an annual economy run is computing, the economy route now relays the optional stage fields that the simulation entrypoint reports for the job (PolicyEngine/policyengine-sim-api#781): stage, progress, progress_max, submitted_at, started_at, stage_started_at and segments. Malformed fields are dropped and never fail a poll. Completed, error and field-less computing bodies keep their existing shape. Co-Authored-By: Claude Opus 5.5 --- .../relay-simulation-run-progress.added.md | 1 + .../libs/simulation_entrypoint.py | 42 +++++++++ policyengine_api/openapi_spec.yaml | 37 +++++++- policyengine_api/routes/economy_routes.py | 18 ++-- policyengine_api/services/economy_service.py | 21 ++++- tests/unit/libs/test_simulation_entrypoint.py | 90 +++++++++++++++++++ .../routes/test_economy_dataset_validation.py | 50 +++++++++++ tests/unit/services/test_economy_service.py | 47 ++++++++++ 8 files changed, 294 insertions(+), 12 deletions(-) create mode 100644 changelog.d/relay-simulation-run-progress.added.md diff --git a/changelog.d/relay-simulation-run-progress.added.md b/changelog.d/relay-simulation-run-progress.added.md new file mode 100644 index 000000000..490594884 --- /dev/null +++ b/changelog.d/relay-simulation-run-progress.added.md @@ -0,0 +1 @@ +Annual economy responses relay the simulation entrypoint's optional stage fields (`stage`, `progress`, `progress_max`, timestamps and `segments`) while a run is computing. diff --git a/policyengine_api/libs/simulation_entrypoint.py b/policyengine_api/libs/simulation_entrypoint.py index aa2a2bd17..56d8c6091 100644 --- a/policyengine_api/libs/simulation_entrypoint.py +++ b/policyengine_api/libs/simulation_entrypoint.py @@ -8,6 +8,7 @@ import httpx from policyengine_observability import instrument_httpx +from pydantic import BaseModel, ConfigDict, Field, ValidationError from policyengine_api.gcp_logging import logger from policyengine_api.libs.gateway_auth import ( GatewayAuthError, @@ -69,6 +70,45 @@ def _attach_current_request_id(request: httpx.Request) -> None: request.headers[OBSERVABILITY_ID_HEADER] = observability_id +class SimulationRunSegments(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + completed: int = Field(ge=0) + total: int = Field(ge=1) + + +class SimulationRunProgress(BaseModel): + """Stage fields of a running Simulation Entrypoint job. + + The entrypoint adds these optional fields to ``GET /jobs/{job_id}`` for + jobs on workers that report stages (policyengine-sim-api#781). The API + relays them while an economy run is computing. ``stage`` stays a string: + later entrypoint versions may add stages, and clients must accept them. + """ + + model_config = ConfigDict(extra="ignore", frozen=True) + + stage: Optional[str] = None + progress: Optional[int] = Field(default=None, ge=0, le=100) + progress_max: Optional[int] = Field(default=None, ge=0, le=100) + submitted_at: Optional[str] = None + started_at: Optional[str] = None + stage_started_at: Optional[str] = None + segments: Optional[SimulationRunSegments] = None + + @classmethod + def from_status(cls, data: dict) -> Optional["SimulationRunProgress"]: + """Read the stage fields; never fail a poll because of them.""" + try: + run_progress = cls.model_validate(data) + except ValidationError: + return None + return run_progress if run_progress.to_dict() else None + + def to_dict(self) -> dict: + return self.model_dump(mode="json", exclude_none=True) + + @dataclass class ModalSimulationExecution: """ @@ -81,6 +121,7 @@ class ModalSimulationExecution: error: Optional[str] = None policyengine_bundle: Optional[dict] = None resolved_app_name: Optional[str] = None + run_progress: Optional[SimulationRunProgress] = None @property def name(self) -> str: @@ -441,6 +482,7 @@ def get_execution_by_id(self, job_id: str) -> ModalSimulationExecution: error=data.get("error"), policyengine_bundle=data.get("policyengine_bundle"), resolved_app_name=data.get("resolved_app_name"), + run_progress=SimulationRunProgress.from_status(data), ) except httpx.HTTPStatusError as e: diff --git a/policyengine_api/openapi_spec.yaml b/policyengine_api/openapi_spec.yaml index 3155fe341..7925f6935 100644 --- a/policyengine_api/openapi_spec.yaml +++ b/policyengine_api/openapi_spec.yaml @@ -879,7 +879,7 @@ paths: # Query parameters are generated from policyengine_api.query_parameters. responses: 200: - description: Annual economic impact, progress, or persisted error state. Canonical completed results include actual baseline/reform receipts containing the selected SPM configuration. + description: Annual economic impact, progress, or persisted error state. Canonical completed results include actual baseline/reform receipts containing the selected SPM configuration. While computing, the optional stage fields come from the simulation entrypoint and appear only when its worker reports stages. content: application/json: schema: @@ -897,6 +897,41 @@ paths: properties: spm_provenance: $ref: '#/components/schemas/SPMComparisonProvenance' + stage: + type: string + description: Current stage of a computing run, such as queued, loading_data, baseline, reform, simulating or outputs. Later versions may add stages; clients must accept a value they do not know. + progress: + type: integer + minimum: 0 + maximum: 100 + description: Coarse completion percent at the start of the current stage. It changes only at stage boundaries. + progress_max: + type: integer + minimum: 0 + maximum: 100 + description: Highest percent the run reaches before its next stage starts. A client can move a time-based estimate between progress and progress_max. + submitted_at: + type: string + format: date-time + description: When the simulation entrypoint accepted the run. + started_at: + type: string + format: date-time + description: When a worker started the run. + stage_started_at: + type: string + format: date-time + description: When the current stage started. + segments: + type: object + description: Completed region segments of a segmented US national run. + properties: + completed: + type: integer + minimum: 0 + total: + type: integer + minimum: 1 400: $ref: '#/components/responses/SPMValidationError' 404: diff --git a/policyengine_api/routes/economy_routes.py b/policyengine_api/routes/economy_routes.py index 858650511..dcec56d34 100644 --- a/policyengine_api/routes/economy_routes.py +++ b/policyengine_api/routes/economy_routes.py @@ -101,13 +101,17 @@ def get_economic_impact(country_id: str, policy_id: int, baseline_policy_id: int result=None, ) - return _json_response( - { - "status": result_dict["status"], - "message": None, - "result": result_dict["data"], - } - ) + body = { + "status": result_dict["status"], + "message": None, + "result": result_dict["data"], + } + # Stage fields from the Simulation Entrypoint, only while computing and + # only when the worker reports them, so other bodies keep their shape. + run_progress = getattr(economic_impact_result, "run_progress", None) + if result_dict["status"] == "computing" and isinstance(run_progress, dict): + body.update(run_progress) + return _json_response(body) @economy_bp.route( diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index ce7112851..1f8731f2c 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -23,7 +23,10 @@ from policyengine_api.data.v1_models import ReformImpact from policyengine_api.data.places import validate_place_code from policyengine_api.gcp_logging import logger -from policyengine_api.libs.simulation_entrypoint import simulation_entrypoint +from policyengine_api.libs.simulation_entrypoint import ( + SimulationRunProgress, + simulation_entrypoint, +) from policyengine_api.observability import ( runtime as observability_runtime, set_runtime_context, @@ -128,6 +131,9 @@ class EconomicImpactResult(BaseModel): status: ImpactStatus data: Optional[dict] = None message: Optional[str] = None + # Stage fields relayed from the Simulation Entrypoint while computing. + # Kept out of to_dict(); the route adds them only when present. + run_progress: Optional[dict] = None model_config = {"frozen": True} # Make model immutable @@ -149,11 +155,11 @@ def completed(cls, data: dict) -> "EconomicImpactResult": return cls(status=ImpactStatus.OK, data=data) @classmethod - def computing(cls) -> "EconomicImpactResult": + def computing(cls, run_progress: Optional[dict] = None) -> "EconomicImpactResult": """ Create an EconomicImpactResult for a computing impact calculation. """ - return cls(status=ImpactStatus.COMPUTING, data=None) + return cls(status=ImpactStatus.COMPUTING, data=None, run_progress=run_progress) @classmethod def error(cls, message: str) -> "EconomicImpactResult": @@ -1278,7 +1284,14 @@ def _handle_execution_state( {"message": "Sim API execution is still running"}, severity="INFO", ) - return EconomicImpactResult.computing() + run_progress = getattr(execution, "run_progress", None) + return EconomicImpactResult.computing( + run_progress=( + run_progress.to_dict() + if isinstance(run_progress, SimulationRunProgress) + else None + ) + ) else: raise ValueError(f"Unexpected sim API execution state: {execution_state}") diff --git a/tests/unit/libs/test_simulation_entrypoint.py b/tests/unit/libs/test_simulation_entrypoint.py index 833c85f68..8deaa05aa 100644 --- a/tests/unit/libs/test_simulation_entrypoint.py +++ b/tests/unit/libs/test_simulation_entrypoint.py @@ -26,6 +26,7 @@ ModalSimulationExecution, SimulationAPIModal, SimulationEntrypointClient, + SimulationRunProgress, resolve_simulation_entrypoint_url, ) from policyengine_api.request_context import ( # noqa: E402 @@ -961,6 +962,84 @@ def test__given_failed_job__then_returns_error( assert execution.status == MODAL_EXECUTION_STATUS_FAILED assert execution.error == "Simulation timed out" + def test__given_running_job_with_stage_fields__then_reads_run_progress( + self, + mock_httpx_client, + mock_modal_logger, + ): + # Given + mock_httpx_client.get.return_value = create_mock_httpx_response( + status_code=202, + json_data={ + **MOCK_POLL_RESPONSE_RUNNING, + "stage": "simulating", + "progress": 5, + "progress_max": 96, + "submitted_at": "2026-10-13T12:00:00Z", + "started_at": "2026-10-13T12:00:10Z", + "stage_started_at": "2026-10-13T12:00:17Z", + "segments": {"completed": 19, "total": 20}, + "stage_durations": {"loading_data": 7.0}, + }, + ) + api = SimulationAPIModal() + + # When + execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID) + + # Then + assert execution.status == MODAL_EXECUTION_STATUS_RUNNING + assert execution.run_progress.to_dict() == { + "stage": "simulating", + "progress": 5, + "progress_max": 96, + "submitted_at": "2026-10-13T12:00:00Z", + "started_at": "2026-10-13T12:00:10Z", + "stage_started_at": "2026-10-13T12:00:17Z", + "segments": {"completed": 19, "total": 20}, + } + + def test__given_running_job_without_stage_fields__then_has_no_run_progress( + self, + mock_httpx_client, + mock_modal_logger, + ): + # Given + mock_httpx_client.get.return_value = create_mock_httpx_response( + status_code=202, + json_data=MOCK_POLL_RESPONSE_RUNNING, + ) + api = SimulationAPIModal() + + # When + execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID) + + # Then + assert execution.run_progress is None + + def test__given_malformed_stage_fields__then_still_returns_status( + self, + mock_httpx_client, + mock_modal_logger, + ): + # Given + mock_httpx_client.get.return_value = create_mock_httpx_response( + status_code=202, + json_data={ + **MOCK_POLL_RESPONSE_RUNNING, + "stage": "reform", + "progress": 150, + }, + ) + api = SimulationAPIModal() + + # When + execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID) + + # Then + assert execution.status == MODAL_EXECUTION_STATUS_RUNNING + assert execution.run_progress is None + def test__given_job_id__then_polls_correct_endpoint( self, mock_httpx_client, @@ -1008,6 +1087,17 @@ def test__given_network_error__then_raises_exception( log_payload = mock_modal_logger.log_struct.call_args.args[0] assert MOCK_MODAL_JOB_ID in log_payload["message"] + class TestSimulationRunProgress: + def test__given_unknown_stage__then_keeps_it(self): + run_progress = SimulationRunProgress.from_status( + {"status": "running", "stage": "a_future_stage", "progress": 40} + ) + + assert run_progress.to_dict() == { + "stage": "a_future_stage", + "progress": 40, + } + class TestGetBudgetWindowBatchById: def test__given_running_batch__then_returns_running_status( self, diff --git a/tests/unit/routes/test_economy_dataset_validation.py b/tests/unit/routes/test_economy_dataset_validation.py index 4054448dc..a03362d6c 100644 --- a/tests/unit/routes/test_economy_dataset_validation.py +++ b/tests/unit/routes/test_economy_dataset_validation.py @@ -98,6 +98,56 @@ def test_economy_route_keeps_computing_response_as_http_200(mock_get_economic_im response = client.get("/us/economy/123/over/456?region=us&time_period=2026") assert response.status_code == 200 + assert json.loads(response.data) == { + "status": "computing", + "message": None, + "result": None, + } + + +@patch("policyengine_api.routes.economy_routes.economy_service.get_economic_impact") +def test_economy_route_relays_stage_fields_while_computing(mock_get_economic_impact): + mock_result = Mock() + mock_result.to_dict.return_value = { + "status": "computing", + "data": None, + "message": None, + } + mock_result.run_progress = { + "stage": "reform", + "progress": 67, + "progress_max": 96, + "stage_started_at": "2026-10-13T12:01:00Z", + } + mock_get_economic_impact.return_value = mock_result + client = _client_with_economy_blueprint() + + response = client.get("/us/economy/123/over/456?region=us&time_period=2026") + + assert response.status_code == 200 + assert json.loads(response.data) == { + "status": "computing", + "message": None, + "result": None, + "stage": "reform", + "progress": 67, + "progress_max": 96, + "stage_started_at": "2026-10-13T12:01:00Z", + } + + +@patch("policyengine_api.routes.economy_routes.economy_service.get_economic_impact") +def test_economy_route_never_adds_stage_fields_to_completed_results( + mock_get_economic_impact, +): + mock_result = _mock_economic_result() + mock_result.run_progress = {"stage": "outputs", "progress": 97} + mock_get_economic_impact.return_value = mock_result + client = _client_with_economy_blueprint() + + response = client.get("/us/economy/123/over/456?region=us&time_period=2026") + + assert set(json.loads(response.data)) == {"status", "message", "result"} @patch( diff --git a/tests/unit/services/test_economy_service.py b/tests/unit/services/test_economy_service.py index e9c30264f..0334145f3 100644 --- a/tests/unit/services/test_economy_service.py +++ b/tests/unit/services/test_economy_service.py @@ -20,6 +20,10 @@ ImpactStatus, REFORM_IMPACT_START_CLAIM_ATTEMPTS, ) +from policyengine_api.libs.simulation_entrypoint import ( + ModalSimulationExecution, + SimulationRunProgress, +) from policyengine_api.services.policy_service import PolicyService from policyengine_api.spm import SPMValidationError from tests.fixtures.services.economy_service import ( @@ -2269,6 +2273,49 @@ def test__given_modal_running_state__then_returns_computing_result( assert result.status == ImpactStatus.COMPUTING assert result.data is None + def test__given_modal_running_state_with_run_progress__then_relays_it( + self, economy_service, setup_options, mock_logger + ): + # Given + reform_impact = create_mock_reform_impact(status="computing") + execution = ModalSimulationExecution( + job_id="fc-job-1", + status="running", + run_progress=SimulationRunProgress( + stage="reform", progress=67, progress_max=96 + ), + ) + + # When + result = economy_service._handle_execution_state( + setup_options, "running", reform_impact, execution + ) + + # Then + assert result.status == ImpactStatus.COMPUTING + assert result.run_progress == { + "stage": "reform", + "progress": 67, + "progress_max": 96, + } + assert "run_progress" not in result.to_dict() + + def test__given_modal_running_state_without_run_progress__then_has_none( + self, economy_service, setup_options, mock_logger + ): + # Given + reform_impact = create_mock_reform_impact(status="computing") + execution = ModalSimulationExecution(job_id="fc-job-1", status="running") + + # When + result = economy_service._handle_execution_state( + setup_options, "running", reform_impact, execution + ) + + # Then + assert result.status == ImpactStatus.COMPUTING + assert result.run_progress is None + def test__given_modal_submitted_state__then_returns_computing_result( self, economy_service, setup_options, mock_logger ): From 1517a7116dd738af2873c5b3ad1e2548310eff7b Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Thu, 15 Oct 2026 09:40:11 -0400 Subject: [PATCH 2/2] Type relayed run progress and validate its timestamps Review follow-up: - EconomicImpactResult.run_progress is now SimulationRunProgress | None, and the route calls to_dict() once. The service passes execution.run_progress through without getattr or isinstance, so a wrong type fails loudly instead of dropping the fields. Test doubles for computing results now carry run_progress. - submitted_at, started_at and stage_started_at are AwareDatetime, so a malformed timestamp drops the stage fields; they still serialise as ISO 8601. - A test pins the body for a job whose worker reports no stages: once policyengine-sim-api#781 deploys, it still gains submitted_at. Co-Authored-By: Claude Opus 5.5 --- .../libs/simulation_entrypoint.py | 9 ++-- policyengine_api/routes/economy_routes.py | 10 ++-- policyengine_api/services/economy_service.py | 13 ++--- tests/fixtures/services/economy_service.py | 2 + tests/unit/libs/test_simulation_entrypoint.py | 47 +++++++++++++++++++ .../routes/test_economy_dataset_validation.py | 44 ++++++++++++++--- tests/unit/routes/test_economy_spm_query.py | 1 + tests/unit/services/test_economy_service.py | 10 ++-- 8 files changed, 108 insertions(+), 28 deletions(-) diff --git a/policyengine_api/libs/simulation_entrypoint.py b/policyengine_api/libs/simulation_entrypoint.py index 7f721ec4e..c76e4a58e 100644 --- a/policyengine_api/libs/simulation_entrypoint.py +++ b/policyengine_api/libs/simulation_entrypoint.py @@ -7,7 +7,8 @@ import httpx from policyengine_observability import instrument_httpx -from pydantic import BaseModel, ConfigDict, Field, ValidationError +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, ValidationError + from policyengine_api.gcp_logging import logger from policyengine_api.libs.gateway_auth import ( GatewayAuthError, @@ -90,9 +91,9 @@ class SimulationRunProgress(BaseModel): stage: str | None = None progress: int | None = Field(default=None, ge=0, le=100) progress_max: int | None = Field(default=None, ge=0, le=100) - submitted_at: str | None = None - started_at: str | None = None - stage_started_at: str | None = None + submitted_at: AwareDatetime | None = None + started_at: AwareDatetime | None = None + stage_started_at: AwareDatetime | None = None segments: SimulationRunSegments | None = None @classmethod diff --git a/policyengine_api/routes/economy_routes.py b/policyengine_api/routes/economy_routes.py index 5748d202a..000cb35e9 100644 --- a/policyengine_api/routes/economy_routes.py +++ b/policyengine_api/routes/economy_routes.py @@ -109,10 +109,12 @@ def get_economic_impact(country_id: str, policy_id: int, baseline_policy_id: int "result": result_dict["data"], } # Stage fields from the Simulation Entrypoint, only while computing and - # only when the worker reports them, so other bodies keep their shape. - run_progress = getattr(economic_impact_result, "run_progress", None) - if result_dict["status"] == "computing" and isinstance(run_progress, dict): - body.update(run_progress) + # only when the entrypoint reports them, so other bodies keep their shape. + if ( + result_dict["status"] == "computing" + and economic_impact_result.run_progress is not None + ): + body.update(economic_impact_result.run_progress.to_dict()) return _json_response(body) diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index 6b33cf79f..173a3a4bc 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -136,7 +136,7 @@ class EconomicImpactResult(BaseModel): message: str | None = None # Stage fields relayed from the Simulation Entrypoint while computing. # Kept out of to_dict(); the route adds them only when present. - run_progress: dict | None = None + run_progress: SimulationRunProgress | None = None model_config = {"frozen": True} # Make model immutable @@ -158,7 +158,9 @@ def completed(cls, data: dict) -> "EconomicImpactResult": return cls(status=ImpactStatus.OK, data=data) @classmethod - def computing(cls, run_progress: dict | None = None) -> "EconomicImpactResult": + def computing( + cls, run_progress: SimulationRunProgress | None = None + ) -> "EconomicImpactResult": """ Create an EconomicImpactResult for a computing impact calculation. """ @@ -1282,13 +1284,8 @@ def _handle_execution_state( {"message": "Sim API execution is still running"}, severity="INFO", ) - run_progress = getattr(execution, "run_progress", None) return EconomicImpactResult.computing( - run_progress=( - run_progress.to_dict() - if isinstance(run_progress, SimulationRunProgress) - else None - ) + run_progress=execution.run_progress if execution is not None else None ) raise ValueError(f"Unexpected sim API execution state: {execution_state}") diff --git a/tests/fixtures/services/economy_service.py b/tests/fixtures/services/economy_service.py index 4f3738d96..0ec3481ee 100644 --- a/tests/fixtures/services/economy_service.py +++ b/tests/fixtures/services/economy_service.py @@ -264,6 +264,7 @@ def create_mock_modal_execution( result=None, error=None, policyengine_bundle=None, + run_progress=None, ): """ Helper function to create mock Modal execution objects. @@ -292,6 +293,7 @@ def create_mock_modal_execution( mock_execution.error = error mock_execution.policyengine_bundle = policyengine_bundle or MOCK_RUNTIME_BUNDLE mock_execution.resolved_app_name = MOCK_RESOLVED_APP_NAME + mock_execution.run_progress = run_progress return mock_execution diff --git a/tests/unit/libs/test_simulation_entrypoint.py b/tests/unit/libs/test_simulation_entrypoint.py index 1879129df..bb0c452b2 100644 --- a/tests/unit/libs/test_simulation_entrypoint.py +++ b/tests/unit/libs/test_simulation_entrypoint.py @@ -1016,6 +1016,53 @@ def test__given_running_job_without_stage_fields__then_has_no_run_progress( # Then assert execution.run_progress is None + def test__given_only_a_submission_time__then_relays_it( + self, + mock_httpx_client, + mock_modal_logger, + ): + # Given - a job on a worker app without stages, after + # policyengine-sim-api#781 deploys its gateway + mock_httpx_client.get.return_value = create_mock_httpx_response( + status_code=202, + json_data={ + **MOCK_POLL_RESPONSE_RUNNING, + "submitted_at": "2026-10-13T12:00:00Z", + }, + ) + api = SimulationAPIModal() + + # When + execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID) + + # Then + assert execution.run_progress.to_dict() == { + "submitted_at": "2026-10-13T12:00:00Z" + } + + def test__given_a_malformed_timestamp__then_drops_the_stage_fields( + self, + mock_httpx_client, + mock_modal_logger, + ): + # Given + mock_httpx_client.get.return_value = create_mock_httpx_response( + status_code=202, + json_data={ + **MOCK_POLL_RESPONSE_RUNNING, + "stage": "reform", + "stage_started_at": "yesterday", + }, + ) + api = SimulationAPIModal() + + # When + execution = api.get_execution_by_id(MOCK_MODAL_JOB_ID) + + # Then + assert execution.status == MODAL_EXECUTION_STATUS_RUNNING + assert execution.run_progress is None + def test__given_malformed_stage_fields__then_still_returns_status( self, mock_httpx_client, diff --git a/tests/unit/routes/test_economy_dataset_validation.py b/tests/unit/routes/test_economy_dataset_validation.py index a03362d6c..237b79377 100644 --- a/tests/unit/routes/test_economy_dataset_validation.py +++ b/tests/unit/routes/test_economy_dataset_validation.py @@ -3,6 +3,7 @@ from flask import Flask +from policyengine_api.libs.simulation_entrypoint import SimulationRunProgress from policyengine_api.routes.economy_routes import economy_bp @@ -92,6 +93,7 @@ def test_economy_route_keeps_computing_response_as_http_200(mock_get_economic_im "data": None, "message": None, } + mock_result.run_progress = None mock_get_economic_impact.return_value = mock_result client = _client_with_economy_blueprint() @@ -113,12 +115,12 @@ def test_economy_route_relays_stage_fields_while_computing(mock_get_economic_imp "data": None, "message": None, } - mock_result.run_progress = { - "stage": "reform", - "progress": 67, - "progress_max": 96, - "stage_started_at": "2026-10-13T12:01:00Z", - } + mock_result.run_progress = SimulationRunProgress( + stage="reform", + progress=67, + progress_max=96, + stage_started_at="2026-10-13T12:01:00Z", + ) mock_get_economic_impact.return_value = mock_result client = _client_with_economy_blueprint() @@ -136,12 +138,40 @@ def test_economy_route_relays_stage_fields_while_computing(mock_get_economic_imp } +@patch("policyengine_api.routes.economy_routes.economy_service.get_economic_impact") +def test_economy_route_relays_submission_time_for_workers_without_stages( + mock_get_economic_impact, +): + # After policyengine-sim-api#781 deploys, the entrypoint reports + # submitted_at for every new job, also on worker apps without stages. + mock_result = Mock() + mock_result.to_dict.return_value = { + "status": "computing", + "data": None, + "message": None, + } + mock_result.run_progress = SimulationRunProgress( + submitted_at="2026-10-13T12:00:00Z" + ) + mock_get_economic_impact.return_value = mock_result + client = _client_with_economy_blueprint() + + response = client.get("/us/economy/123/over/456?region=us&time_period=2026") + + assert json.loads(response.data) == { + "status": "computing", + "message": None, + "result": None, + "submitted_at": "2026-10-13T12:00:00Z", + } + + @patch("policyengine_api.routes.economy_routes.economy_service.get_economic_impact") def test_economy_route_never_adds_stage_fields_to_completed_results( mock_get_economic_impact, ): mock_result = _mock_economic_result() - mock_result.run_progress = {"stage": "outputs", "progress": 97} + mock_result.run_progress = SimulationRunProgress(stage="outputs", progress=97) mock_get_economic_impact.return_value = mock_result client = _client_with_economy_blueprint() diff --git a/tests/unit/routes/test_economy_spm_query.py b/tests/unit/routes/test_economy_spm_query.py index 90a14ced2..a65f39e06 100644 --- a/tests/unit/routes/test_economy_spm_query.py +++ b/tests/unit/routes/test_economy_spm_query.py @@ -46,6 +46,7 @@ def calculate(**kwargs): setups.append(setup) return SimpleNamespace( cache_status=None, + run_progress=None, to_dict=lambda: { "status": "computing", "message": None, diff --git a/tests/unit/services/test_economy_service.py b/tests/unit/services/test_economy_service.py index 6bd1afd24..3110b3512 100644 --- a/tests/unit/services/test_economy_service.py +++ b/tests/unit/services/test_economy_service.py @@ -5,6 +5,10 @@ import httpx import pytest +from policyengine_api.libs.simulation_entrypoint import ( + ModalSimulationExecution, + SimulationRunProgress, +) from policyengine_api.runtime_cache.core import CacheCoordinationError from policyengine_api.runtime_cache.reform_impacts import ReformImpactStartClaim from policyengine_api.services.budget_window_cache import BudgetWindowCacheState @@ -18,10 +22,6 @@ ImpactAction, ImpactStatus, ) -from policyengine_api.libs.simulation_entrypoint import ( - ModalSimulationExecution, - SimulationRunProgress, -) from policyengine_api.services.policy_service import PolicyService from policyengine_api.services.reform_impacts_service import ( ReformImpactHandoffError, @@ -2294,7 +2294,7 @@ def test__given_modal_running_state_with_run_progress__then_relays_it( # Then assert result.status == ImpactStatus.COMPUTING - assert result.run_progress == { + assert result.run_progress.to_dict() == { "stage": "reform", "progress": 67, "progress_max": 96,