From a2decade69138c6af8a15a957b82cfbc5cfbc294 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 1 Oct 2026 17:54:45 +0400 Subject: [PATCH 1/9] Define compact typed SPM provenance contracts --- .../policyengine_simulation_contract/spm.py | 224 ++++++++++++++++-- .../stage12_execution.py | 3 +- .../tests/test_spm_provenance.py | 171 +++++++++++++ 3 files changed, 372 insertions(+), 26 deletions(-) create mode 100644 libs/policyengine-simulation-contract/tests/test_spm_provenance.py diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py index bbde945da..d958aafe4 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -1,6 +1,7 @@ """Dependency-light public models for US SPM selection and calculation receipts.""" import re +from collections.abc import Sequence from datetime import date from typing import Any, Literal, Optional @@ -98,19 +99,139 @@ def validate_location(self): class SPMProvenance(BaseModel): - """Detached calculation receipt; data certification is a separate claim.""" + """Compact detached receipt for one SPM calculation configuration.""" + + model_config = ConfigDict(extra="forbid", frozen=True, serialize_by_alias=True) + + schema_version: Literal["canonical-spm-provenance-v2"] + forecast_id: str = Field(min_length=1, pattern=r"^\S+$") + forecast_sha256: str = Field(pattern=r"^[0-9a-f]{64}$") + scenario: str = Field(min_length=1, pattern=r"^\S+$") + geography_kind: Literal["county", "national", "metro"] + geography_id: Optional[str] = Field(min_length=1) + county_vintage: str = Field(pattern=r"^[0-9]{4}$") + as_of: Optional[date] + years: tuple[str, ...] + runtime_versions: "SPMRuntimeVersions" + + @field_validator("years") + @classmethod + def validate_years(cls, value: tuple[str, ...]) -> tuple[str, ...]: + if any(re.fullmatch(r"[0-9]{4}", year) is None for year in value): + raise ValueError("SPM provenance years must use four-digit years") + if value != tuple(sorted(set(value))): + raise ValueError("SPM provenance years must be sorted and unique") + return value - model_config = ConfigDict(extra="forbid") + @model_validator(mode="after") + def validate_provenance_location(self) -> "SPMProvenance": + if self.geography_kind == "metro": + if self.geography_id is None or not self.geography_id.strip(): + raise ValueError("An SPM area receipt requires geography_id") + elif self.geography_id is not None: + raise ValueError("Only an SPM area receipt accepts geography_id") + return self + + +class SPMRuntimeVersions(BaseModel): + """Package versions needed to reproduce an SPM calculation.""" + + model_config = ConfigDict( + extra="forbid", + frozen=True, + populate_by_name=True, + serialize_by_alias=True, + ) - forecast_id: str - forecast_sha256: str - scenario: str - geography_kind: str - runtime_versions: dict[str, Optional[str]] - years: dict[str, dict[str, Any]] - geographies: list[dict[str, Any]] - composition_method: str - storage_method: str + policyengine: Optional[str] + policyengine_core: Optional[str] = Field(alias="policyengine-core") + policyengine_us: Optional[str] = Field(alias="policyengine-us") + spm_calculator: Optional[str] = Field(alias="spm-calculator") + + +class SPMCalculationProvenance(BaseModel): + """Resolved SPM selection paired with its compact calculation receipt.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + spm_config: SPMSelection + spm_provenance: SPMProvenance + + @model_validator(mode="after") + def require_matching_receipt(self) -> "SPMCalculationProvenance": + _require_complete_selection(self.spm_config) + expected = { + "forecast_sha256": self.spm_config.forecast_content_sha256, + "scenario": self.spm_config.scenario, + "geography_kind": self.spm_config.geography_kind, + "geography_id": self.spm_config.geography_id, + "county_vintage": self.spm_config.county_vintage, + "as_of": ( + date.fromisoformat(self.spm_config.as_of) + if self.spm_config.as_of is not None + else None + ), + } + for field_name, expected_value in expected.items(): + if getattr(self.spm_provenance, field_name) != expected_value: + raise ValueError( + f"SPM receipt {field_name} differs from the resolved selection" + ) + return self + + +def _require_complete_selection(selection: SPMSelection) -> None: + required = set(SPMSelection.model_fields) + missing = required.difference(selection.model_fields_set) + if missing: + raise ValueError( + "Resolved SPM selection is missing " + ", ".join(sorted(missing)) + ) + if selection.forecast_content_sha256 is None: + raise ValueError("Resolved SPM selection has no artifact hash") + if selection.scenario is None: + raise ValueError("Resolved SPM selection has no scenario") + + +def build_spm_provenance( + *, + forecast_id: str, + forecast_sha256: str, + selection: SPMSelection, + years: Sequence[str], + runtime_versions: SPMRuntimeVersions, +) -> SPMProvenance: + """Build one compact receipt from resolved calculation inputs.""" + + _require_complete_selection(selection) + if forecast_sha256 != selection.forecast_content_sha256: + raise ValueError("SPM receipt artifact hash differs from the selection") + return SPMProvenance( + schema_version="canonical-spm-provenance-v2", + forecast_id=forecast_id, + forecast_sha256=forecast_sha256, + scenario=selection.scenario, + geography_kind=selection.geography_kind, + geography_id=selection.geography_id, + county_vintage=selection.county_vintage, + as_of=( + date.fromisoformat(selection.as_of) + if selection.as_of is not None + else None + ), + years=tuple(years), + runtime_versions=runtime_versions, + ) + + +def build_spm_calculation_provenance( + *, + config: SPMSelection, + receipt: SPMProvenance, +) -> SPMCalculationProvenance: + """Pair a resolved selection with a matching compact receipt.""" + + return SPMCalculationProvenance(spm_config=config, spm_provenance=receipt) SPM_CONTRACT_VERSION = "canonical-spm-v1" @@ -177,12 +298,53 @@ def pinned_defaults(self): return self +class SPMExecutionProvenance(BaseModel): + """One shared receipt and the number of executions that produced it.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + receipt: SPMProvenance + execution_count: int = Field(ge=1) + + class SPMComparisonProvenance(BaseModel): - """One receipt per executed regional segment, separately for each policy.""" + """Compact receipts for the baseline and reform execution sets.""" - model_config = ConfigDict(extra="forbid") - baseline: list[SPMProvenance] = Field(min_length=1) - reform: list[SPMProvenance] = Field(min_length=1) + model_config = ConfigDict(extra="forbid", frozen=True) + + schema_version: Literal["canonical-spm-comparison-v2"] + baseline: SPMExecutionProvenance + reform: SPMExecutionProvenance + + +def _collapse_spm_receipts( + side: str, + receipts: Sequence[SPMProvenance], +) -> SPMExecutionProvenance: + if not receipts: + raise ValueError(f"{side} receipts must not be empty") + first = receipts[0] + if any(receipt != first for receipt in receipts[1:]): + raise ValueError(f"{side} receipts differ") + return SPMExecutionProvenance(receipt=first, execution_count=len(receipts)) + + +def build_spm_comparison_provenance( + *, + baseline_receipts: Sequence[SPMProvenance], + reform_receipts: Sequence[SPMProvenance], +) -> SPMComparisonProvenance: + """Collapse identical child receipts into one receipt per policy side.""" + + baseline = _collapse_spm_receipts("baseline", baseline_receipts) + reform = _collapse_spm_receipts("reform", reform_receipts) + if baseline.receipt != reform.receipt: + raise ValueError("baseline and reform SPM receipts differ") + return SPMComparisonProvenance( + schema_version="canonical-spm-comparison-v2", + baseline=baseline, + reform=reform, + ) def resolve_spm_selection( @@ -295,14 +457,15 @@ def validate_spm_result( provenance = SPMComparisonProvenance.model_validate( result.get("spm_provenance") ) - for receipt in provenance.baseline + provenance.reform: + for execution in (provenance.baseline, provenance.reform): + build_spm_calculation_provenance( + config=chosen, + receipt=execution.receipt, + ) if ( - receipt.forecast_sha256 != selection["forecast_content_sha256"] - or receipt.scenario != selection["scenario"] - or receipt.geography_kind != selection["geography_kind"] + expected_year is not None + and str(expected_year) not in execution.receipt.years ): - raise ValueError("Result SPM provenance differs from the request") - if expected_year is not None and str(expected_year) not in receipt.years: raise ValueError( "Result SPM provenance does not cover the requested year" ) @@ -330,8 +493,19 @@ def combine_spm_results( "SPM_CONFIGURATION_UNAVAILABLE", "Missing canonical SPM receipt" ) validated.append(receipt) - combined = SPMComparisonProvenance( - baseline=[r for item in validated for r in item.baseline], - reform=[r for item in validated for r in item.reform], + combined = build_spm_comparison_provenance( + baseline_receipts=[ + item.baseline.receipt + for item in validated + for _ in range(item.baseline.execution_count) + ], + reform_receipts=[ + item.reform.receipt + for item in validated + for _ in range(item.reform.execution_count) + ], ) - return {"spm_config": selection, "spm_provenance": combined.model_dump(mode="json")} + return { + "spm_config": selection, + "spm_provenance": combined.model_dump(mode="json", by_alias=True), + } diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py index e5258864d..37069da46 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py @@ -18,6 +18,7 @@ field_validator, model_validator, ) +from policyengine_simulation_contract.spm import SPMCalculationProvenance from policyengine_simulation_contract.uk_geography import ( UKLocalAuthorityBoundaryVersion as UKLocalAuthorityBoundaryVersion, UKLocalAuthorityMetadata as UKLocalAuthorityMetadata, @@ -332,7 +333,7 @@ class SimulationArtifactDescriptor(StrictContractModel): output_plan_sha256: Sha256Digest row_identity: RowIdentity bundle: BundleProvenance - calculation_provenance: dict[str, JsonValue] | None = None + calculation_provenance: SPMCalculationProvenance | None = None class ReportExecutionInput(StrictContractModel): diff --git a/libs/policyengine-simulation-contract/tests/test_spm_provenance.py b/libs/policyengine-simulation-contract/tests/test_spm_provenance.py new file mode 100644 index 000000000..59edc8b3f --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_spm_provenance.py @@ -0,0 +1,171 @@ +"""Compact, versioned SPM calculation-receipt contracts.""" + +import json + +import pytest +from pydantic import ValidationError + +from policyengine_simulation_contract.spm import ( + SPMCalculationProvenance, + SPMComparisonProvenance, + SPMProvenance, + SPMRuntimeVersions, + SPMSelection, + build_spm_calculation_provenance, + build_spm_comparison_provenance, + build_spm_provenance, +) + + +SELECTION = SPMSelection( + forecast_content_sha256="a" * 64, + scenario="ce_trend", + geography_kind="national", + geography_id=None, + county_vintage="2020", + as_of=None, +) +VERSIONS = SPMRuntimeVersions.model_validate( + { + "policyengine": "6.2.1", + "policyengine-core": "3.32.10", + "policyengine-us": "2.2.1", + "spm-calculator": "1.0.0", + } +) + + +def receipt(*, year: str = "2026", forecast_id: str = "forecast-2026"): + return build_spm_provenance( + forecast_id=forecast_id, + forecast_sha256="a" * 64, + selection=SELECTION, + years=(year,), + runtime_versions=VERSIONS, + ) + + +def test_compact_receipt_has_exact_public_shape() -> None: + value = receipt().model_dump(mode="json", by_alias=True) + + assert value == { + "schema_version": "canonical-spm-provenance-v2", + "forecast_id": "forecast-2026", + "forecast_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + "years": ["2026"], + "runtime_versions": { + "policyengine": "6.2.1", + "policyengine-core": "3.32.10", + "policyengine-us": "2.2.1", + "spm-calculator": "1.0.0", + }, + } + assert len(json.dumps(value, separators=(",", ":")).encode()) < 1_024 + + +@pytest.mark.parametrize( + "updates", + [ + {"forecast_sha256": "A" * 64}, + {"forecast_sha256": "a" * 63}, + {"years": ["2027", "2026"]}, + {"years": ["2026", "2026"]}, + {"years": ["26"]}, + {"geography_kind": "state"}, + {"geography_kind": "metro", "geography_id": None}, + {"geography_kind": "national", "geography_id": "35620"}, + {"county_vintage": "20"}, + {"as_of": "09/01/2026"}, + ], +) +def test_compact_receipt_rejects_invalid_values(updates) -> None: + value = receipt().model_dump(mode="json", by_alias=True) + value.update(updates) + + with pytest.raises(ValidationError): + SPMProvenance.model_validate(value) + + +def test_compact_receipt_rejects_old_rich_shape() -> None: + old = { + "forecast_id": "forecast-2026", + "forecast_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "runtime_versions": {"policyengine-us": "2.2.1"}, + "years": {"2026": {"median_diagnostics": {}}}, + "geographies": [], + "composition_method": "classified-inputs", + "storage_method": "formula", + } + + with pytest.raises(ValidationError): + SPMProvenance.model_validate(old) + + +def test_builder_requires_receipt_to_match_resolved_selection() -> None: + with pytest.raises(ValueError, match="artifact hash"): + build_spm_provenance( + forecast_id="forecast-2026", + forecast_sha256="b" * 64, + selection=SELECTION, + years=("2026",), + runtime_versions=VERSIONS, + ) + + +def test_calculation_builder_requires_complete_matching_config() -> None: + calculated = build_spm_calculation_provenance( + config=SELECTION, + receipt=receipt(), + ) + assert isinstance(calculated, SPMCalculationProvenance) + + with pytest.raises(ValueError, match="scenario"): + build_spm_calculation_provenance( + config=SELECTION.model_copy(update={"scenario": "zero_real"}), + receipt=receipt(), + ) + + +def test_comparison_collapses_identical_children_and_counts_them() -> None: + comparison = build_spm_comparison_provenance( + baseline_receipts=[receipt()] * 20, + reform_receipts=[receipt()] * 20, + ) + value = comparison.model_dump(mode="json", by_alias=True) + + assert value["schema_version"] == "canonical-spm-comparison-v2" + assert value["baseline"]["execution_count"] == 20 + assert value["reform"]["execution_count"] == 20 + assert len(json.dumps(value, separators=(",", ":")).encode()) < 2_048 + + +def test_comparison_rejects_missing_or_mismatched_children() -> None: + with pytest.raises(ValueError, match="baseline receipts"): + build_spm_comparison_provenance( + baseline_receipts=[], + reform_receipts=[receipt()], + ) + with pytest.raises(ValueError, match="baseline receipts differ"): + build_spm_comparison_provenance( + baseline_receipts=[receipt(), receipt(forecast_id="other")], + reform_receipts=[receipt()], + ) + with pytest.raises(ValueError, match="baseline and reform"): + build_spm_comparison_provenance( + baseline_receipts=[receipt()], + reform_receipts=[receipt(year="2027")], + ) + + +def test_comparison_contract_rejects_old_receipt_lists() -> None: + with pytest.raises(ValidationError): + SPMComparisonProvenance.model_validate( + {"baseline": [receipt()], "reform": [receipt()]} + ) From 6b93fc370134157f41c644dad13fdc080d413e27 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 1 Oct 2026 18:29:35 +0400 Subject: [PATCH 2/9] Emit compact SPM provenance across simulation paths --- .../fixtures/spm_doubles.py | 7 +- .../policyengine_simulation_executor/spm.py | 140 ++++++++++++++---- .../stage12_artifacts.py | 33 ++--- .../stage12_runtime/aggregation.py | 35 ++--- .../stage12_runtime/simulation.py | 23 +-- .../tests/test_baseline_artifacts.py | 9 +- .../tests/test_budget_window_scheduler.py | 22 ++- .../tests/test_canonical_spm.py | 64 ++++++-- .../tests/test_canonical_spm_native.py | 2 +- .../tests/test_segmented_national.py | 47 +++--- .../tests/test_simulation_output_builder.py | 6 +- .../tests/test_stage12_artifacts.py | 39 ++++- .../tests/test_stage12_runtime.py | 83 ++++++----- 13 files changed, 343 insertions(+), 167 deletions(-) diff --git a/projects/policyengine-simulation-executor/fixtures/spm_doubles.py b/projects/policyengine-simulation-executor/fixtures/spm_doubles.py index b9c2771d1..790e2ac9b 100644 --- a/projects/policyengine-simulation-executor/fixtures/spm_doubles.py +++ b/projects/policyengine-simulation-executor/fixtures/spm_doubles.py @@ -18,7 +18,12 @@ def spm_receipt(selection, year="2026"): "forecast_sha256": selection["forecast_content_sha256"], "scenario": selection["scenario"], "geography_kind": selection["geography_kind"], - "runtime_versions": {"policyengine-us": "test-only"}, + "runtime_versions": { + "policyengine": "test-only", + "policyengine-core": "test-only", + "policyengine-us": "test-only", + "spm-calculator": "test-only", + }, "years": {str(year): {"status": "forecast"}}, "geographies": [], "composition_method": "classified-inputs", diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py index bfb1c311a..2859434fc 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py @@ -1,13 +1,19 @@ """Certified SPM runtime selection, independent of caller-supplied metadata.""" import json +from collections.abc import Mapping from functools import lru_cache from policyengine_simulation_contract.spm import ( SPMCapability, + SPMCalculationProvenance, SPMInputError, SPMProvenance, - combine_spm_results, + SPMRuntimeVersions, + SPMSelection, + build_spm_calculation_provenance, + build_spm_comparison_provenance, + build_spm_provenance, resolve_spm_selection, spm_error_detail, ) @@ -193,34 +199,112 @@ def normalize_runtime_spm(params): return selection +def _required_string(source: Mapping[str, object], field_name: str) -> str: + value = source.get(field_name) + if not isinstance(value, str) or not value: + raise ValueError(f"SPM calculation receipt has no valid {field_name}") + return value + + +def compact_spm_provenance( + country_receipt: object, + selection: object, +) -> SPMProvenance: + """Extract the stable public receipt from country-owned diagnostics. + + The country package retains rich in-process diagnostic data. This adapter + reads only the scalar identity fields, executed-year keys, and explicit + package versions that belong on the simulation API wire contract. + """ + + if not isinstance(country_receipt, Mapping): + raise TypeError("SPM calculation receipt must be an object") + resolved = SPMSelection.model_validate(selection) + forecast_sha256 = _required_string(country_receipt, "forecast_sha256") + for field_name, expected_value in ( + ("forecast_sha256", resolved.forecast_content_sha256), + ("scenario", resolved.scenario), + ("geography_kind", resolved.geography_kind), + ): + if country_receipt.get(field_name) != expected_value: + raise ValueError( + f"SPM calculation receipt {field_name} differs from the selection" + ) + raw_versions = country_receipt.get("runtime_versions") + if not isinstance(raw_versions, Mapping): + raise TypeError("SPM calculation receipt has no runtime versions") + runtime_versions = SPMRuntimeVersions.model_validate(raw_versions) + if any( + version is None + for version in ( + runtime_versions.policyengine, + runtime_versions.policyengine_core, + runtime_versions.policyengine_us, + runtime_versions.spm_calculator, + ) + ): + raise ValueError("Certified SPM calculation runtime versions must be populated") + raw_years = country_receipt.get("years") + if not isinstance(raw_years, Mapping): + raise TypeError("SPM calculation receipt has no executed years") + if any(not isinstance(year, str) for year in raw_years): + raise TypeError("SPM calculation receipt year keys must be strings") + return build_spm_provenance( + forecast_id=_required_string(country_receipt, "forecast_id"), + forecast_sha256=forecast_sha256, + selection=resolved, + years=tuple(raw_years), + runtime_versions=runtime_versions, + ) + + +def simulation_spm_calculation_provenance( + simulation: object, + selection: object, +) -> SPMCalculationProvenance: + """Read and compact one simulation's resolved SPM receipt.""" + + resolved = SPMSelection.model_validate(selection) + if getattr(simulation, "spm_config", None) != resolved.model_dump(mode="json"): + raise SPMInputError( + "SPM_CONFIGURATION_UNAVAILABLE", + "Simulation ignored the requested SPM selection", + ) + receipt = getattr(simulation, "spm_provenance", None) + if not callable(receipt): + raise SPMInputError( + "SPM_CONFIGURATION_UNAVAILABLE", + "Simulation has no valid SPM calculation receipt", + ) + try: + return build_spm_calculation_provenance( + config=resolved, + receipt=compact_spm_provenance(receipt(), resolved), + ) + except (TypeError, ValueError) as exc: + raise SPMInputError( + "SPM_CONFIGURATION_UNAVAILABLE", + "Simulation has no valid SPM calculation receipt", + ) from exc + + def simulation_spm_result(baseline, reform, selection, *, expected_year=None): if selection is None: return {} - receipts = [] - for simulation in (baseline, reform): - if simulation.spm_config != selection: - raise SPMInputError( - "SPM_CONFIGURATION_UNAVAILABLE", - "Simulation ignored the requested SPM selection", - ) - try: - receipts.append( - SPMProvenance.model_validate(simulation.spm_provenance()).model_dump( - mode="json" - ) - ) - except ValueError as exc: - raise SPMInputError( - "SPM_CONFIGURATION_UNAVAILABLE", - "Simulation has no valid SPM calculation receipt", - ) from exc - return combine_spm_results( - [ - { - "spm_config": selection, - "spm_provenance": {"baseline": [receipts[0]], "reform": [receipts[1]]}, - } - ], - selection, - expected_year=expected_year, + baseline_provenance = simulation_spm_calculation_provenance(baseline, selection) + reform_provenance = simulation_spm_calculation_provenance(reform, selection) + comparison = build_spm_comparison_provenance( + baseline_receipts=[baseline_provenance.spm_provenance], + reform_receipts=[reform_provenance.spm_provenance], ) + if expected_year is not None and str(expected_year) not in ( + comparison.baseline.receipt.years + ): + raise SPMInputError( + "SPM_CONFIGURATION_UNAVAILABLE", + "Result SPM provenance does not cover the requested year", + ) + return { + "spm_config": baseline_provenance.spm_config.model_dump(mode="json"), + "spm_provenance": comparison.model_dump(mode="json", by_alias=True), + } diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py index c3a57613c..7be267fd9 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py @@ -7,7 +7,7 @@ from datetime import datetime from hashlib import sha256 from io import BytesIO -from typing import Any, cast +from typing import Any from uuid import UUID import pandas as pd @@ -24,8 +24,8 @@ SimulationArtifactDescriptor, stage12_output_plan_sha256, ) +from policyengine_simulation_contract.spm import SPMCalculationProvenance from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata -from pydantic import JsonValue from policyengine_simulation_executor.artifact_store import ArtifactStore @@ -118,7 +118,7 @@ def _frame_payload( def serialize_simulation_frames( frames: Mapping[str, pd.DataFrame], *, - calculation_provenance: Mapping[str, Any] | None = None, + calculation_provenance: SPMCalculationProvenance | None = None, uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, ) -> tuple[bytes, RowIdentity]: combined, dtypes, identifier_values = _frame_payload(frames) @@ -140,7 +140,9 @@ def serialize_simulation_frames( ) if calculation_provenance is not None: metadata[PARQUET_CONTRACT.calculation_provenance_metadata_key.encode()] = ( - canonical_json_bytes(calculation_provenance) + canonical_json_bytes( + calculation_provenance.model_dump(mode="json", by_alias=True) + ) ) if uk_local_authority_metadata is not None: metadata[PARQUET_CONTRACT.uk_local_authority_metadata_key.encode()] = ( @@ -210,16 +212,15 @@ def _parquet_schema_metadata(payload: bytes) -> Mapping[bytes, bytes]: return parquet_file.schema_arrow.metadata or {} -def deserialize_calculation_provenance(payload: bytes) -> dict[str, Any] | None: +def deserialize_calculation_provenance( + payload: bytes, +) -> SPMCalculationProvenance | None: raw = _parquet_schema_metadata(payload).get( PARQUET_CONTRACT.calculation_provenance_metadata_key.encode() ) if raw is None: return None - value = json.loads(raw) - if not isinstance(value, dict): - raise TypeError("Stage 12 calculation provenance must be an object") - return value + return SPMCalculationProvenance.model_validate_json(raw) def deserialize_uk_local_authority_metadata( @@ -288,20 +289,12 @@ def write_simulation( prefix: str, simulation: PlannedSimulationExecutionInput, frames: Mapping[str, pd.DataFrame], - calculation_provenance: Mapping[str, Any] | None = None, + calculation_provenance: SPMCalculationProvenance | None = None, uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, ) -> SimulationArtifactDescriptor: - normalized_provenance = cast( - dict[str, JsonValue] | None, - ( - dict(calculation_provenance) - if calculation_provenance is not None - else None - ), - ) payload, row_identity = serialize_simulation_frames( frames, - calculation_provenance=normalized_provenance, + calculation_provenance=calculation_provenance, uk_local_authority_metadata=uk_local_authority_metadata, ) artifact = self._write_immutable( @@ -317,7 +310,7 @@ def write_simulation( output_plan_sha256=stage12_output_plan_sha256(simulation.output_plan), row_identity=row_identity, bundle=simulation.bundle, - calculation_provenance=normalized_provenance, + calculation_provenance=calculation_provenance, ) def write_aggregate( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py index cbc19e59a..7e6d5d87c 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py @@ -195,27 +195,28 @@ def build_spm_result( requested_spm = report.baseline.options.get("spm") if requested_spm is None: return {} - from policyengine_simulation_contract.spm import combine_spm_results + from policyengine_simulation_contract.spm import ( + build_spm_comparison_provenance, + ) baseline_provenance = baseline_descriptor.calculation_provenance reform_provenance = reform_descriptor.calculation_provenance if baseline_provenance is None or reform_provenance is None: raise ValueError("SPM calculation provenance is missing") - selection = baseline_provenance.get("spm_config") - if not isinstance(selection, dict): - raise TypeError("SPM calculation selection is invalid") - if selection != reform_provenance.get("spm_config"): + selection = baseline_provenance.spm_config + if selection != reform_provenance.spm_config: raise ValueError("SPM calculation selections do not match") - return combine_spm_results( - [ - { - "spm_config": selection, - "spm_provenance": { - "baseline": [baseline_provenance.get("spm_provenance")], - "reform": [reform_provenance.get("spm_provenance")], - }, - } - ], - selection, - expected_year=report.baseline.year, + for receipt in ( + baseline_provenance.spm_provenance, + reform_provenance.spm_provenance, + ): + if str(report.baseline.year) not in receipt.years: + raise ValueError("SPM calculation provenance does not cover report year") + comparison = build_spm_comparison_provenance( + baseline_receipts=[baseline_provenance.spm_provenance], + reform_receipts=[reform_provenance.spm_provenance], ) + return { + "spm_config": selection.model_dump(mode="json"), + "spm_provenance": comparison.model_dump(mode="json", by_alias=True), + } diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 20af9ab7d..7bf43c91b 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -21,6 +21,7 @@ Stage12InvocationContext, stage12_output_plan_sha256, ) +from policyengine_simulation_contract.spm import SPMCalculationProvenance from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata from policyengine_simulation_observability.stages import ( STAGE12_SIMULATION_STAGES, @@ -43,7 +44,7 @@ @dataclass(frozen=True) class SimulationCalculation: frames: Mapping[str, pd.DataFrame] - calculation_provenance: dict[str, Any] | None = None + calculation_provenance: SPMCalculationProvenance | None = None uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None @@ -221,22 +222,14 @@ def calculate_simulation_frames( selection = getattr(model, "spm_config", None) calculation_provenance = None if selection is not None: - from policyengine_simulation_contract.spm import ( - SPMProvenance, - SPMSelection, + from policyengine_simulation_executor.spm import ( + simulation_spm_calculation_provenance, ) - receipt = getattr(model, "spm_provenance", None) - if not callable(receipt): - raise RuntimeError("simulation produced no SPM receipt") - calculation_provenance = { - "spm_config": SPMSelection.model_validate(selection).model_dump( - mode="json" - ), - "spm_provenance": SPMProvenance.model_validate(receipt()).model_dump( - mode="json" - ), - } + calculation_provenance = simulation_spm_calculation_provenance( + model, + selection, + ) return SimulationCalculation( frames=frames, calculation_provenance=calculation_provenance, diff --git a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py index 2f0557a2a..f3db74ed3 100644 --- a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py @@ -638,7 +638,12 @@ def _spm_receipt(selection=None, *, year="2026"): "forecast_sha256": selection["forecast_content_sha256"], "scenario": selection["scenario"], "geography_kind": selection["geography_kind"], - "runtime_versions": {"policyengine-us": "test-only"}, + "runtime_versions": { + "policyengine": "test-only", + "policyengine-core": "test-only", + "policyengine-us": "test-only", + "spm-calculator": "test-only", + }, "years": {year: {"status": "forecast"}}, "geographies": [], "composition_method": "classified-inputs", @@ -819,7 +824,7 @@ def test_artifact_without_a_receipt_recomputes(self, fresh_cache): def test_malformed_receipt_recomputes(self, fresh_cache): malformed = _spm_receipt() - del malformed["composition_method"] + del malformed["runtime_versions"] model = SPMModelVersion( stored_config=deepcopy(SPM_SELECTION), stored_receipt=malformed ) diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py b/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py index 69df30493..abfe10e5a 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py @@ -41,19 +41,29 @@ def spm_child_result(runtime, simulation_year, *, receipt_year): """A child result carrying a canonical SPM receipt for ``receipt_year``.""" receipt = { + "schema_version": "canonical-spm-provenance-v2", "forecast_id": "test-only", "forecast_sha256": SPM_SELECTION["forecast_content_sha256"], "scenario": SPM_SELECTION["scenario"], "geography_kind": SPM_SELECTION["geography_kind"], - "runtime_versions": {"policyengine-us": "test-only"}, - "years": {receipt_year: {"status": "forecast"}}, - "geographies": [], - "composition_method": "classified-inputs", - "storage_method": "formula", + "geography_id": SPM_SELECTION["geography_id"], + "county_vintage": SPM_SELECTION["county_vintage"], + "as_of": SPM_SELECTION["as_of"], + "years": [receipt_year], + "runtime_versions": { + "policyengine": "test-only", + "policyengine-core": "test-only", + "policyengine-us": "test-only", + "spm-calculator": "test-only", + }, } return { "spm_config": dict(SPM_SELECTION), - "spm_provenance": {"baseline": [receipt], "reform": [dict(receipt)]}, + "spm_provenance": { + "schema_version": "canonical-spm-comparison-v2", + "baseline": {"receipt": receipt, "execution_count": 1}, + "reform": {"receipt": dict(receipt), "execution_count": 1}, + }, **runtime.child_result_for_year(simulation_year), } diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py index 5fd9503a3..85a26df60 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py @@ -24,6 +24,7 @@ ) from policyengine_simulation_executor import artifact_keys, simulation_runtime from policyengine_simulation_executor.simulation import create_router +from policyengine_simulation_executor.spm import compact_spm_provenance from src.modal.budget_window_results import ( extract_annual_impact, build_budget_window_result, @@ -41,21 +42,64 @@ CAPABILITY = SPMCapability(defaults=SELECTION).model_dump() +def test_country_receipt_is_compacted_at_the_executor_boundary(): + compact = compact_spm_provenance( + { + "forecast_id": "forecast-2026", + "forecast_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "runtime_versions": { + "policyengine": "6.2.1", + "policyengine-core": "3.32.10", + "policyengine-us": "2.2.1", + "spm-calculator": "1.0.0", + }, + "years": {"2026": {"median_diagnostics": {"large": "value"}}}, + "geographies": [{"large": "value"}], + "composition_method": "internal", + "storage_method": "internal", + }, + SELECTION, + ).model_dump(mode="json", by_alias=True) + + assert compact["years"] == ["2026"] + assert compact["geography_id"] is None + assert compact["county_vintage"] == "2020" + for excluded in ( + "geographies", + "median_diagnostics", + "composition_method", + "storage_method", + ): + assert excluded not in json.dumps(compact) + + def result(selection=SELECTION, year="2026"): receipt = dict( + schema_version="canonical-spm-provenance-v2", forecast_id="test-only", forecast_sha256=selection["forecast_content_sha256"], scenario=selection["scenario"], geography_kind=selection["geography_kind"], - runtime_versions={"policyengine-us": "test-only"}, - years={year: {"status": "forecast"}}, - geographies=[], - composition_method="classified-inputs", - storage_method="formula", + geography_id=selection["geography_id"], + county_vintage=selection["county_vintage"], + as_of=selection["as_of"], + years=[year], + runtime_versions={ + "policyengine": "test-only", + "policyengine-core": "test-only", + "policyengine-us": "test-only", + "spm-calculator": "test-only", + }, ) return { "spm_config": deepcopy(selection), - "spm_provenance": {"baseline": [receipt], "reform": [deepcopy(receipt)]}, + "spm_provenance": { + "schema_version": "canonical-spm-comparison-v2", + "baseline": {"receipt": receipt, "execution_count": 1}, + "reform": {"receipt": deepcopy(receipt), "execution_count": 1}, + }, "budget": { "tax_revenue_impact": 1, "benefit_spending_impact": 2, @@ -534,11 +578,11 @@ def test_fanout_and_budget_window_preserve_settings_and_receipts(): start_year="2026", window_size=2, annual_impacts=rows ) dumped = json.loads(window.model_dump_json()) - assert dumped["annualImpacts"][1]["spm_provenance"]["baseline"][0]["years"] == { - "2027": {"status": "forecast"} - } + assert dumped["annualImpacts"][1]["spm_provenance"]["baseline"]["receipt"][ + "years" + ] == ["2027"] merged = combine_spm_results([result(), result()], SELECTION) - assert len(merged["spm_provenance"]["baseline"]) == 2 + assert merged["spm_provenance"]["baseline"]["execution_count"] == 2 changed = result({**SELECTION, "scenario": "zero_real"}) for outputs in ([result(), changed], [result(), {"budget": {}}]): with pytest.raises(SPMInputError): diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py index c26e37086..b03de5121 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py @@ -64,7 +64,7 @@ def test_worker_baseline_reform_national_local_cache_and_receipts( result = simulation_spm_result(baseline, reform, baseline.spm_config) dumped = json.loads(json.dumps(result)) assert dumped["spm_config"] == baseline.spm_config == reform.spm_config - assert dumped["spm_provenance"]["reform"][0]["years"]["2024"] + assert dumped["spm_provenance"]["reform"]["receipt"]["years"] == ["2024"] original = baseline.spm_provenance() original["years"].clear() assert baseline.spm_provenance()["years"] diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national.py b/projects/policyengine-simulation-executor/tests/test_segmented_national.py index e76e953bc..9328b078e 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national.py @@ -399,17 +399,23 @@ def test__everything_else_takes_the_monolithic_path(self, monkeypatch, params): } -def _spm_receipt(label, *, selection=SPM_SELECTION, year="2026"): +def _spm_receipt(*, selection=SPM_SELECTION, year="2026"): return { - "forecast_id": label, + "schema_version": "canonical-spm-provenance-v2", + "forecast_id": "test-only", "forecast_sha256": selection["forecast_content_sha256"], "scenario": selection["scenario"], "geography_kind": selection["geography_kind"], - "runtime_versions": {"policyengine-us": "test-only"}, - "years": {year: {"status": "forecast"}}, - "geographies": [], - "composition_method": "classified-inputs", - "storage_method": "formula", + "geography_id": selection["geography_id"], + "county_vintage": selection["county_vintage"], + "as_of": selection["as_of"], + "years": [year], + "runtime_versions": { + "policyengine": "test-only", + "policyengine-core": "test-only", + "policyengine-us": "test-only", + "spm-calculator": "test-only", + }, } @@ -419,10 +425,15 @@ def _spm_child(index, *, selection=SPM_SELECTION, year="2026"): "child": index, "spm_config": dict(selection), "spm_provenance": { - "baseline": [ - _spm_receipt(f"baseline-{index}", selection=selection, year=year) - ], - "reform": [_spm_receipt(f"reform-{index}", selection=selection, year=year)], + "schema_version": "canonical-spm-comparison-v2", + "baseline": { + "receipt": _spm_receipt(selection=selection, year=year), + "execution_count": 1, + }, + "reform": { + "receipt": _spm_receipt(selection=selection, year=year), + "execution_count": 1, + }, }, } @@ -456,16 +467,12 @@ def test__child_receipts_combine_into_one_national_receipt( # back on the parent as the one national selection. assert all(p["spm"] == SPM_SELECTION for p in fake.spawned_payloads) assert output["spm_config"] == SPM_SELECTION - # Every executed segment's receipts survive, baseline and reform - # concatenated in group order: a national SPM result must account - # for all 20 segments, not just the first. + # One receipt describes all identical child executions. The counts + # prove that every segment contributed without duplicating metadata. provenance = output["spm_provenance"] - assert [r["forecast_id"] for r in provenance["baseline"]] == [ - f"baseline-{i}" for i in range(20) - ] - assert [r["forecast_id"] for r in provenance["reform"]] == [ - f"reform-{i}" for i in range(20) - ] + assert provenance["baseline"]["execution_count"] == 20 + assert provenance["reform"]["execution_count"] == 20 + assert provenance["baseline"]["receipt"]["forecast_id"] == "test-only" assert output["budget"] == {} @pytest.mark.parametrize( diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py index 784b8353c..490c24194 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -102,7 +102,11 @@ def _assert_macro_result(result, selection): assert all(result[key] is None for key in spm_keys) return assert result["spm_config"] == selection - assert set(result["spm_provenance"]) == {"baseline", "reform"} + assert set(result["spm_provenance"]) == { + "schema_version", + "baseline", + "reform", + } def _load_dataset(params, *, country_module=None, region_resolution=None): diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py b/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py index ca85dd067..07df8001a 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py @@ -13,6 +13,12 @@ UKLocalAuthorityBoundaryVersion, UKLocalAuthorityMetadata, ) +from policyengine_simulation_contract.spm import ( + SPMRuntimeVersions, + SPMSelection, + build_spm_calculation_provenance, + build_spm_provenance, +) from policyengine_simulation_executor.stage12_artifacts import ( Stage12ArtifactStore, @@ -42,6 +48,32 @@ def _frames(): } +def _calculation_provenance(): + selection = SPMSelection( + forecast_content_sha256="a" * 64, + scenario="official", + geography_kind="national", + geography_id=None, + county_vintage="2020", + as_of=None, + ) + receipt = build_spm_provenance( + forecast_id="forecast-2026", + forecast_sha256="a" * 64, + selection=selection, + years=("2026",), + runtime_versions=SPMRuntimeVersions.model_validate( + { + "policyengine": "6.2.1", + "policyengine-core": "3.32.10", + "policyengine-us": "2.2.1", + "spm-calculator": "1.0.0", + } + ), + ) + return build_spm_calculation_provenance(config=selection, receipt=receipt) + + def test_parquet_encoding_is_deterministic_and_preserves_rows_and_dtypes() -> None: first, first_identity = serialize_simulation_frames(_frames()) second, second_identity = serialize_simulation_frames(_frames()) @@ -71,10 +103,7 @@ def test_parquet_round_trip_preserves_an_all_null_entity_column() -> None: def test_parquet_retains_detached_calculation_provenance() -> None: - provenance = { - "spm_config": {"scenario": "official"}, - "spm_provenance": {"forecast_sha256": "a" * 64}, - } + provenance = _calculation_provenance() payload, _ = serialize_simulation_frames( _frames(), calculation_provenance=provenance, @@ -103,7 +132,7 @@ def test_parquet_omits_uk_metadata_for_non_uk_simulations() -> None: def test_metadata_reads_do_not_materialize_simulation_rows() -> None: - provenance = {"spm_config": {"scenario": "official"}} + provenance = _calculation_provenance() metadata = UKLocalAuthorityMetadata( boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 3963458b8..3059df0a3 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -40,6 +40,12 @@ Stage12OutputPlan, stage12_output_plan_sha256, ) +from policyengine_simulation_contract.spm import ( + SPMRuntimeVersions, + SPMSelection, + build_spm_calculation_provenance, + build_spm_provenance, +) from policyengine_simulation_contract.uk_geography import ( UKLocalAuthorityBoundaryVersion, UKLocalAuthorityMetadata, @@ -70,6 +76,32 @@ REFORM_ID = UUID("00000000-0000-0000-0000-000000000003") +def _spm_calculation_provenance(): + selection = SPMSelection( + forecast_content_sha256="f" * 64, + scenario="official", + geography_kind="national", + geography_id=None, + county_vintage="2020", + as_of=None, + ) + receipt = build_spm_provenance( + forecast_id="forecast-1", + forecast_sha256="f" * 64, + selection=selection, + years=("2026",), + runtime_versions=SPMRuntimeVersions.model_validate( + { + "policyengine": "6.2.1", + "policyengine-core": "3.32.10", + "policyengine-us": "2.2.1", + "spm-calculator": "1.0.0", + } + ), + ) + return build_spm_calculation_provenance(config=selection, receipt=receipt) + + def _bundle() -> BundleProvenance: return BundleProvenance( policyengine_version="5.2.0", @@ -561,7 +593,7 @@ def test_single_worker_retains_detached_calculation_provenance() -> None: simulation = _planned_simulation(SimulationRole.BASELINE) store.children[simulation.simulation_execution_id] = _child(simulation) artifacts = FakeArtifacts() - provenance = {"receipt": {"version": 1}} + provenance = _spm_calculation_provenance() result = run_single_simulation( simulation.model_dump(mode="json"), @@ -575,7 +607,9 @@ def test_single_worker_retains_detached_calculation_provenance() -> None: ), ) - assert result["calculation_provenance"] == provenance + assert result["calculation_provenance"] == provenance.model_dump( + mode="json", by_alias=True + ) def test_single_worker_exposes_only_a_bounded_failure() -> None: @@ -628,27 +662,8 @@ def test_single_worker_rejects_frames_that_do_not_satisfy_the_output_plan() -> N def test_aggregate_combines_detached_spm_receipts() -> None: - selection = { - "forecast_content_sha256": "f" * 64, - "scenario": "official", - "geography_kind": "national", - "geography_id": None, - "county_vintage": "2020", - "as_of": None, - } - - def receipt(): - return { - "forecast_id": "forecast-1", - "forecast_sha256": "f" * 64, - "scenario": "official", - "geography_kind": "national", - "runtime_versions": {"spm-calculator": "0.3.1"}, - "years": {"2026": {}}, - "geographies": [], - "composition_method": "direct", - "storage_method": "detached", - } + provenance = _spm_calculation_provenance() + selection = provenance.spm_config.model_dump(mode="json") report = _report() report = report.model_copy( @@ -664,26 +679,12 @@ def receipt(): baseline = ( FakeArtifacts() .add_simulation(report.baseline) - .model_copy( - update={ - "calculation_provenance": { - "spm_config": selection, - "spm_provenance": receipt(), - } - } - ) + .model_copy(update={"calculation_provenance": provenance}) ) reform = ( FakeArtifacts() .add_simulation(report.reform) - .model_copy( - update={ - "calculation_provenance": { - "spm_config": selection, - "spm_provenance": receipt(), - } - } - ) + .model_copy(update={"calculation_provenance": provenance}) ) result = _build_spm_result( @@ -693,8 +694,8 @@ def receipt(): ) assert result["spm_config"] == selection - assert len(result["spm_provenance"]["baseline"]) == 1 - assert len(result["spm_provenance"]["reform"]) == 1 + assert result["spm_provenance"]["baseline"]["execution_count"] == 1 + assert result["spm_provenance"]["reform"]["execution_count"] == 1 def test_aggregate_stand_ins_preserve_policy_and_cliff_options(monkeypatch) -> None: From f38f3c5bd1a3acfd7a9cfdcc30dc3c6cf41f3264 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 1 Oct 2026 18:31:44 +0400 Subject: [PATCH 3/9] Verify compact SPM clients and artifact boundaries --- .../policyengine_simulation_contract/spm.py | 6 ++-- .../tests/test_spm_provenance.py | 15 ++++++++++ .../tests/test_spm_generated_client.py | 29 ++++++++++++++----- .../tests/test_stage12_artifacts.py | 24 +++++++++++++++ 4 files changed, 64 insertions(+), 10 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py index d958aafe4..9a0812b95 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -204,6 +204,8 @@ def build_spm_provenance( """Build one compact receipt from resolved calculation inputs.""" _require_complete_selection(selection) + assert selection.forecast_content_sha256 is not None + assert selection.scenario is not None if forecast_sha256 != selection.forecast_content_sha256: raise ValueError("SPM receipt artifact hash differs from the selection") return SPMProvenance( @@ -215,9 +217,7 @@ def build_spm_provenance( geography_id=selection.geography_id, county_vintage=selection.county_vintage, as_of=( - date.fromisoformat(selection.as_of) - if selection.as_of is not None - else None + date.fromisoformat(selection.as_of) if selection.as_of is not None else None ), years=tuple(years), runtime_versions=runtime_versions, diff --git a/libs/policyengine-simulation-contract/tests/test_spm_provenance.py b/libs/policyengine-simulation-contract/tests/test_spm_provenance.py index 59edc8b3f..fa10cc32b 100644 --- a/libs/policyengine-simulation-contract/tests/test_spm_provenance.py +++ b/libs/policyengine-simulation-contract/tests/test_spm_provenance.py @@ -73,6 +73,7 @@ def test_compact_receipt_has_exact_public_shape() -> None: [ {"forecast_sha256": "A" * 64}, {"forecast_sha256": "a" * 63}, + {"forecast_id": "forecast with spaces"}, {"years": ["2027", "2026"]}, {"years": ["2026", "2026"]}, {"years": ["26"]}, @@ -108,6 +109,20 @@ def test_compact_receipt_rejects_old_rich_shape() -> None: SPMProvenance.model_validate(old) +def test_compact_receipt_serializes_as_of_as_an_iso_date() -> None: + dated = build_spm_provenance( + forecast_id="forecast-2026", + forecast_sha256="a" * 64, + selection=SELECTION.model_copy( + update={"as_of": "2026-09-09"}, + ), + years=("2026",), + runtime_versions=VERSIONS, + ) + + assert dated.model_dump(mode="json", by_alias=True)["as_of"] == "2026-09-09" + + def test_builder_requires_receipt_to_match_resolved_selection() -> None: with pytest.raises(ValueError, match="artifact hash"): build_spm_provenance( diff --git a/projects/policyengine-apis-integ/tests/test_spm_generated_client.py b/projects/policyengine-apis-integ/tests/test_spm_generated_client.py index 31c9e9d49..2cd685864 100644 --- a/projects/policyengine-apis-integ/tests/test_spm_generated_client.py +++ b/projects/policyengine-apis-integ/tests/test_spm_generated_client.py @@ -42,17 +42,32 @@ def test_generated_capability_keeps_and_enforces_contract_version(): SPMCapability.from_dict({**capability, "contract_version": "unknown-v2"}) -def test_generated_receipts_preserve_arbitrary_dated_metadata(): +def test_generated_receipts_use_only_compact_typed_metadata(): receipt = { + "schema_version": "canonical-spm-provenance-v2", "forecast_id": "test-only", "forecast_sha256": "a" * 64, "scenario": "ce_trend", "geography_kind": "national", - "runtime_versions": {"policyengine-us": "test-only", "optional": None}, - "years": {"2026": {"status": "forecast", "window": [2021, 2025]}}, - "geographies": [{"year": 2026, "county_assignment": None}], - "composition_method": "test-only", - "storage_method": "test-only", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + "years": ["2026"], + "runtime_versions": { + "policyengine": "6.2.1", + "policyengine-core": "3.32.10", + "policyengine-us": "2.2.1", + "spm-calculator": "1.0.0", + }, + } + provenance = { + "schema_version": "canonical-spm-comparison-v2", + "baseline": {"receipt": receipt, "execution_count": 20}, + "reform": {"receipt": receipt, "execution_count": 20}, } - provenance = {"baseline": [receipt], "reform": [receipt]} assert SPMComparisonProvenance.from_dict(provenance).to_dict() == provenance + + +def test_generated_comparison_rejects_old_receipt_lists(): + with pytest.raises((KeyError, TypeError, ValueError)): + SPMComparisonProvenance.from_dict({"baseline": [], "reform": []}) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py b/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py index 07df8001a..b54b1a938 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py @@ -2,12 +2,15 @@ from __future__ import annotations +import json from datetime import UTC, datetime +from io import BytesIO from unittest.mock import patch from uuid import UUID import numpy as np import pandas as pd +import pyarrow.parquet as pq import pytest from policyengine_simulation_contract.stage12_execution import ( UKLocalAuthorityBoundaryVersion, @@ -19,6 +22,7 @@ build_spm_calculation_provenance, build_spm_provenance, ) +from pydantic import ValidationError from policyengine_simulation_executor.stage12_artifacts import ( Stage12ArtifactStore, @@ -112,6 +116,26 @@ def test_parquet_retains_detached_calculation_provenance() -> None: assert deserialize_calculation_provenance(payload) == provenance +def test_parquet_rejects_old_rich_calculation_provenance() -> None: + payload, _ = serialize_simulation_frames(_frames()) + table = pq.read_table(BytesIO(payload)) + metadata = dict(table.schema.metadata or {}) + metadata[b"policyengine.stage12.calculation_provenance"] = json.dumps( + { + "spm_config": {"scenario": "official"}, + "spm_provenance": { + "forecast_id": "old-rich-receipt", + "geographies": [], + }, + } + ).encode() + output = BytesIO() + pq.write_table(table.replace_schema_metadata(metadata), output) + + with pytest.raises(ValidationError): + deserialize_calculation_provenance(output.getvalue()) + + def test_parquet_retains_typed_uk_local_authority_metadata() -> None: metadata = UKLocalAuthorityMetadata( boundary_version=UKLocalAuthorityBoundaryVersion.LAD22 From f7619fd6687fecbae1669de71e4f44b0ecf9f2c7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 1 Oct 2026 19:35:01 +0400 Subject: [PATCH 4/9] Fix compact SPM gateway contract fixtures --- .../policyengine_simulation_contract/spm.py | 8 +- .../tests/test_spm_provenance.py | 12 ++ .../policyengine_simulation_executor/spm.py | 10 -- .../tests/golden/openapi.json | 161 ++++++++++++------ .../tests/test_spm_routes.py | 22 ++- 5 files changed, 144 insertions(+), 69 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py index 9a0812b95..d9151a839 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -143,10 +143,10 @@ class SPMRuntimeVersions(BaseModel): serialize_by_alias=True, ) - policyengine: Optional[str] - policyengine_core: Optional[str] = Field(alias="policyengine-core") - policyengine_us: Optional[str] = Field(alias="policyengine-us") - spm_calculator: Optional[str] = Field(alias="spm-calculator") + policyengine: str + policyengine_core: str = Field(alias="policyengine-core") + policyengine_us: str = Field(alias="policyengine-us") + spm_calculator: str = Field(alias="spm-calculator") class SPMCalculationProvenance(BaseModel): diff --git a/libs/policyengine-simulation-contract/tests/test_spm_provenance.py b/libs/policyengine-simulation-contract/tests/test_spm_provenance.py index fa10cc32b..6dad439b2 100644 --- a/libs/policyengine-simulation-contract/tests/test_spm_provenance.py +++ b/libs/policyengine-simulation-contract/tests/test_spm_provenance.py @@ -109,6 +109,18 @@ def test_compact_receipt_rejects_old_rich_shape() -> None: SPMProvenance.model_validate(old) +@pytest.mark.parametrize( + "package", + ["policyengine", "policyengine-core", "policyengine-us", "spm-calculator"], +) +def test_runtime_versions_reject_null_package_versions(package: str) -> None: + versions = VERSIONS.model_dump(mode="json", by_alias=True) + invalid_versions: dict[str, str | None] = {**versions, package: None} + + with pytest.raises(ValidationError): + SPMRuntimeVersions.model_validate(invalid_versions) + + def test_compact_receipt_serializes_as_of_as_an_iso_date() -> None: dated = build_spm_provenance( forecast_id="forecast-2026", diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py index 2859434fc..90b735d49 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py @@ -234,16 +234,6 @@ def compact_spm_provenance( if not isinstance(raw_versions, Mapping): raise TypeError("SPM calculation receipt has no runtime versions") runtime_versions = SPMRuntimeVersions.model_validate(raw_versions) - if any( - version is None - for version in ( - runtime_versions.policyengine, - runtime_versions.policyengine_core, - runtime_versions.policyengine_us, - runtime_versions.spm_calculator, - ) - ): - raise ValueError("Certified SPM calculation runtime versions must be populated") raw_years = country_receipt.get("years") if not isinstance(raw_years, Mapping): raise TypeError("SPM calculation receipt has no executed years") diff --git a/projects/policyengine-simulation-gateway/tests/golden/openapi.json b/projects/policyengine-simulation-gateway/tests/golden/openapi.json index bddec43f3..00a61112a 100644 --- a/projects/policyengine-simulation-gateway/tests/golden/openapi.json +++ b/projects/policyengine-simulation-gateway/tests/golden/openapi.json @@ -1758,31 +1758,27 @@ }, "SPMComparisonProvenance": { "properties": { + "schema_version": { + "type": "string", + "const": "canonical-spm-comparison-v2", + "title": "Schema Version" + }, "baseline": { - "items": { - "$ref": "#/components/schemas/SPMProvenance" - }, - "type": "array", - "minItems": 1, - "title": "Baseline" + "$ref": "#/components/schemas/SPMExecutionProvenance" }, "reform": { - "items": { - "$ref": "#/components/schemas/SPMProvenance" - }, - "type": "array", - "minItems": 1, - "title": "Reform" + "$ref": "#/components/schemas/SPMExecutionProvenance" } }, "additionalProperties": false, "type": "object", "required": [ + "schema_version", "baseline", "reform" ], "title": "SPMComparisonProvenance", - "description": "One receipt per executed regional segment, separately for each policy." + "description": "Compact receipts for the baseline and reform execution sets." }, "SPMErrorDetail": { "properties": { @@ -1802,78 +1798,145 @@ ], "title": "SPMErrorDetail" }, + "SPMExecutionProvenance": { + "properties": { + "receipt": { + "$ref": "#/components/schemas/SPMProvenance" + }, + "execution_count": { + "type": "integer", + "minimum": 1.0, + "title": "Execution Count" + } + }, + "additionalProperties": false, + "type": "object", + "required": [ + "receipt", + "execution_count" + ], + "title": "SPMExecutionProvenance", + "description": "One shared receipt and the number of executions that produced it." + }, "SPMProvenance": { "properties": { + "schema_version": { + "type": "string", + "const": "canonical-spm-provenance-v2", + "title": "Schema Version" + }, "forecast_id": { "type": "string", + "minLength": 1, + "pattern": "^\\S+$", "title": "Forecast Id" }, "forecast_sha256": { "type": "string", + "pattern": "^[0-9a-f]{64}$", "title": "Forecast Sha256" }, "scenario": { "type": "string", + "minLength": 1, + "pattern": "^\\S+$", "title": "Scenario" }, "geography_kind": { "type": "string", + "enum": [ + "county", + "national", + "metro" + ], "title": "Geography Kind" }, - "runtime_versions": { - "additionalProperties": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ] - }, - "type": "object", - "title": "Runtime Versions" + "geography_id": { + "anyOf": [ + { + "type": "string", + "minLength": 1 + }, + { + "type": "null" + } + ], + "title": "Geography Id" }, - "years": { - "additionalProperties": { - "additionalProperties": true, - "type": "object" - }, - "type": "object", - "title": "Years" + "county_vintage": { + "type": "string", + "pattern": "^[0-9]{4}$", + "title": "County Vintage" }, - "geographies": { + "as_of": { + "anyOf": [ + { + "type": "string", + "format": "date" + }, + { + "type": "null" + } + ], + "title": "As Of" + }, + "years": { "items": { - "additionalProperties": true, - "type": "object" + "type": "string" }, "type": "array", - "title": "Geographies" - }, - "composition_method": { - "type": "string", - "title": "Composition Method" + "title": "Years" }, - "storage_method": { - "type": "string", - "title": "Storage Method" + "runtime_versions": { + "$ref": "#/components/schemas/SPMRuntimeVersions" } }, "additionalProperties": false, "type": "object", "required": [ + "schema_version", "forecast_id", "forecast_sha256", "scenario", "geography_kind", - "runtime_versions", + "geography_id", + "county_vintage", + "as_of", "years", - "geographies", - "composition_method", - "storage_method" + "runtime_versions" ], "title": "SPMProvenance", - "description": "Detached calculation receipt; data certification is a separate claim." + "description": "Compact detached receipt for one SPM calculation configuration." + }, + "SPMRuntimeVersions": { + "properties": { + "policyengine": { + "type": "string", + "title": "Policyengine" + }, + "policyengine-core": { + "type": "string", + "title": "Policyengine-Core" + }, + "policyengine-us": { + "type": "string", + "title": "Policyengine-Us" + }, + "spm-calculator": { + "type": "string", + "title": "Spm-Calculator" + } + }, + "additionalProperties": false, + "type": "object", + "required": [ + "policyengine", + "policyengine-core", + "policyengine-us", + "spm-calculator" + ], + "title": "SPMRuntimeVersions", + "description": "Package versions needed to reproduce an SPM calculation." }, "SPMSelection": { "properties": { diff --git a/projects/policyengine-simulation-gateway/tests/test_spm_routes.py b/projects/policyengine-simulation-gateway/tests/test_spm_routes.py index 6ceaffe94..a500b1310 100644 --- a/projects/policyengine-simulation-gateway/tests/test_spm_routes.py +++ b/projects/policyengine-simulation-gateway/tests/test_spm_routes.py @@ -238,20 +238,30 @@ def test_versions_omits_malformed_capability_and_submission_rejects_it( def spm_receipt(year): return { + "schema_version": "canonical-spm-provenance-v2", "forecast_id": "ce-forecast", "forecast_sha256": "a" * 64, "scenario": "ce_trend", "geography_kind": "national", - "runtime_versions": {"spm_calculator": "0.3.1"}, - "years": {str(year): {"entry": f"{year}-01-01"}}, - "geographies": [{"kind": "national"}], - "composition_method": "national", - "storage_method": "artifact", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + "years": [str(year)], + "runtime_versions": { + "policyengine": "6.2.1", + "policyengine-core": "3.32.10", + "policyengine-us": "2.2.1", + "spm-calculator": "1.0.0", + }, } def spm_provenance(year): - return {"baseline": [spm_receipt(year)], "reform": [spm_receipt(year)]} + return { + "schema_version": "canonical-spm-comparison-v2", + "baseline": {"receipt": spm_receipt(year), "execution_count": 1}, + "reform": {"receipt": spm_receipt(year), "execution_count": 1}, + } def test_completed_result_body_keeps_resolved_nulls(mock_modal, client): From 253e86fe8a0bcdf0844c5be5f9ec72c7a023edb9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 1 Oct 2026 19:59:58 +0400 Subject: [PATCH 5/9] Require resolved SPM result metadata --- .../gateway_models.py | 3 +- .../macro_output.py | 7 +- .../policyengine_simulation_contract/spm.py | 113 +++++++++++++----- .../tests/test_spm_provenance.py | 36 +++++- .../tests/test_spm_generated_client.py | 28 +++++ .../tests/test_openapi.py | 17 +++ .../src/modal/budget_window_results.py | 10 +- .../tests/test_canonical_spm.py | 27 ++--- .../tests/golden/openapi.json | 74 +++++++++++- 9 files changed, 263 insertions(+), 52 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py index 4969e14e6..a5c8bf202 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py @@ -15,6 +15,7 @@ ) from policyengine_simulation_contract.spm import ( + SPMResolvedConfiguration, SPMSelection, SPMCapability, SPMComparisonProvenance, @@ -237,7 +238,7 @@ def validate_end_year(self) -> "BudgetWindowBatchRequest": class BudgetWindowAnnualImpact(BaseModel): """Annual budget-window impact row.""" - spm_config: Optional[SPMSelection] = None + spm_config: Optional[SPMResolvedConfiguration] = None spm_provenance: Optional[SPMComparisonProvenance] = None year: str taxRevenueImpact: float diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py index 7b6eae156..8986c0dbc 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py @@ -6,7 +6,10 @@ from pydantic import BaseModel, ConfigDict, RootModel -from policyengine_simulation_contract.spm import SPMSelection, SPMComparisonProvenance +from policyengine_simulation_contract.spm import ( + SPMComparisonProvenance, + SPMResolvedConfiguration, +) T = TypeVar("T") @@ -189,7 +192,7 @@ class CongressionalDistrictImpactOutput(MacroOutputModel): class SingleYearMacroOutput(MacroOutputModel): """Completed response returned by a single-year macro simulation.""" - spm_config: SPMSelection | None = None + spm_config: SPMResolvedConfiguration | None = None spm_provenance: SPMComparisonProvenance | None = None model_version: str diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py index d9151a839..bd455aa10 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -49,13 +49,11 @@ class SPMSelection(BaseModel): def serialize_selection( self, handler: SerializerFunctionWrapHandler, info: SerializationInfo ): - """Preserve inherited options through ordinary and nested JSON. + """Preserve explicitly selected options through ordinary and nested JSON. Presence is the contract: an omitted option inherits the bundle default, so an option explicitly selected as null has to stay on the - wire. ``exclude_none`` would otherwise turn a completed result's - resolved selection back into a partial request, and the poll routes - apply it to every body via ``response_model_exclude_none``. + request wire even when a containing transport excludes null values. """ dumped = handler(self) if info.exclude_none: @@ -98,6 +96,52 @@ def validate_location(self): return self +class SPMResolvedConfiguration(BaseModel): + """The complete six-field SPM selection recorded beside a receipt.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + forecast_content_sha256: str = Field(pattern=r"^[0-9a-f]{64}$") + scenario: str = Field(min_length=1, pattern=r"^\S+$") + geography_kind: Literal["county", "national", "metro"] + geography_id: str | None = Field(min_length=1) + county_vintage: str = Field(pattern=r"^[0-9]{4}$") + as_of: str | None + + @model_serializer(mode="wrap") + def serialize_resolved_configuration( + self, handler: SerializerFunctionWrapHandler, info: SerializationInfo + ): + """Keep required nullable fields in completed response bodies.""" + + dumped = handler(self) + if info.exclude_none: + for name in ("geography_id", "as_of"): + if ( + (info.include is None or name in info.include) + and (info.exclude is None or name not in info.exclude) + and getattr(self, name) is None + ): + dumped[name] = None + return dumped + + @field_validator("as_of") + @classmethod + def validate_as_of(cls, value: str | None) -> str | None: + if value is not None and date.fromisoformat(value).isoformat() != value: + raise ValueError("as_of must be an ISO calendar date (YYYY-MM-DD)") + return value + + @model_validator(mode="after") + def validate_location(self) -> "SPMResolvedConfiguration": + if self.geography_kind == "metro": + if self.geography_id is None or not self.geography_id.strip(): + raise ValueError("An SPM area selection requires geography_id") + elif self.geography_id is not None: + raise ValueError("Only an SPM area selection accepts geography_id") + return self + + class SPMProvenance(BaseModel): """Compact detached receipt for one SPM calculation configuration.""" @@ -134,7 +178,7 @@ def validate_provenance_location(self) -> "SPMProvenance": class SPMRuntimeVersions(BaseModel): - """Package versions needed to reproduce an SPM calculation.""" + """Package versions required to reproduce one certified SPM execution.""" model_config = ConfigDict( extra="forbid", @@ -143,10 +187,10 @@ class SPMRuntimeVersions(BaseModel): serialize_by_alias=True, ) - policyengine: str - policyengine_core: str = Field(alias="policyengine-core") - policyengine_us: str = Field(alias="policyengine-us") - spm_calculator: str = Field(alias="spm-calculator") + policyengine: str = Field(min_length=1) + policyengine_core: str = Field(min_length=1, alias="policyengine-core") + policyengine_us: str = Field(min_length=1, alias="policyengine-us") + spm_calculator: str = Field(min_length=1, alias="spm-calculator") class SPMCalculationProvenance(BaseModel): @@ -154,12 +198,11 @@ class SPMCalculationProvenance(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) - spm_config: SPMSelection + spm_config: SPMResolvedConfiguration spm_provenance: SPMProvenance @model_validator(mode="after") def require_matching_receipt(self) -> "SPMCalculationProvenance": - _require_complete_selection(self.spm_config) expected = { "forecast_sha256": self.spm_config.forecast_content_sha256, "scenario": self.spm_config.scenario, @@ -193,6 +236,24 @@ def _require_complete_selection(selection: SPMSelection) -> None: raise ValueError("Resolved SPM selection has no scenario") +def _resolved_spm_configuration( + selection: SPMSelection, +) -> SPMResolvedConfiguration: + """Convert a resolved request selection to the completed-result contract.""" + + _require_complete_selection(selection) + assert selection.forecast_content_sha256 is not None + assert selection.scenario is not None + return SPMResolvedConfiguration( + forecast_content_sha256=selection.forecast_content_sha256, + scenario=selection.scenario, + geography_kind=selection.geography_kind, + geography_id=selection.geography_id, + county_vintage=selection.county_vintage, + as_of=selection.as_of, + ) + + def build_spm_provenance( *, forecast_id: str, @@ -231,7 +292,10 @@ def build_spm_calculation_provenance( ) -> SPMCalculationProvenance: """Pair a resolved selection with a matching compact receipt.""" - return SPMCalculationProvenance(spm_config=config, spm_provenance=receipt) + return SPMCalculationProvenance( + spm_config=_resolved_spm_configuration(config), + spm_provenance=receipt, + ) SPM_CONTRACT_VERSION = "canonical-spm-v1" @@ -439,27 +503,18 @@ def validate_spm_result( ) return None try: - chosen = SPMSelection.model_validate(selection) - result_selection = SPMSelection.model_validate(result.get("spm_config")) - required = { - "forecast_content_sha256", - "scenario", - "geography_kind", - "county_vintage", - } - if not required <= result_selection.model_fields_set: - raise ValueError("Result has no complete resolved SPM selection") - selection = {name: getattr(chosen, name) for name in SPMSelection.model_fields} - if { - name: getattr(result_selection, name) for name in SPMSelection.model_fields - } != selection: + chosen = _resolved_spm_configuration(SPMSelection.model_validate(selection)) + result_selection = SPMResolvedConfiguration.model_validate( + result.get("spm_config") + ) + if result_selection != chosen: raise ValueError("Result SPM selection differs from the request") provenance = SPMComparisonProvenance.model_validate( result.get("spm_provenance") ) for execution in (provenance.baseline, provenance.reform): build_spm_calculation_provenance( - config=chosen, + config=SPMSelection.model_validate(chosen.model_dump(mode="json")), receipt=execution.receipt, ) if ( @@ -506,6 +561,8 @@ def combine_spm_results( ], ) return { - "spm_config": selection, + "spm_config": _resolved_spm_configuration( + SPMSelection.model_validate(selection) + ).model_dump(mode="json"), "spm_provenance": combined.model_dump(mode="json", by_alias=True), } diff --git a/libs/policyengine-simulation-contract/tests/test_spm_provenance.py b/libs/policyengine-simulation-contract/tests/test_spm_provenance.py index 6dad439b2..121b27cea 100644 --- a/libs/policyengine-simulation-contract/tests/test_spm_provenance.py +++ b/libs/policyengine-simulation-contract/tests/test_spm_provenance.py @@ -9,6 +9,7 @@ SPMCalculationProvenance, SPMComparisonProvenance, SPMProvenance, + SPMResolvedConfiguration, SPMRuntimeVersions, SPMSelection, build_spm_calculation_provenance, @@ -113,9 +114,12 @@ def test_compact_receipt_rejects_old_rich_shape() -> None: "package", ["policyengine", "policyengine-core", "policyengine-us", "spm-calculator"], ) -def test_runtime_versions_reject_null_package_versions(package: str) -> None: +@pytest.mark.parametrize("invalid_value", [None, ""]) +def test_runtime_versions_reject_invalid_package_versions( + package: str, invalid_value: str | None +) -> None: versions = VERSIONS.model_dump(mode="json", by_alias=True) - invalid_versions: dict[str, str | None] = {**versions, package: None} + invalid_versions: dict[str, str | None] = {**versions, package: invalid_value} with pytest.raises(ValidationError): SPMRuntimeVersions.model_validate(invalid_versions) @@ -152,6 +156,12 @@ def test_calculation_builder_requires_complete_matching_config() -> None: receipt=receipt(), ) assert isinstance(calculated, SPMCalculationProvenance) + assert isinstance(calculated.spm_config, SPMResolvedConfiguration) + assert calculated.model_dump(mode="json", exclude_none=True)["spm_config"] == { + **SELECTION.model_dump(mode="json"), + "geography_id": None, + "as_of": None, + } with pytest.raises(ValueError, match="scenario"): build_spm_calculation_provenance( @@ -160,6 +170,28 @@ def test_calculation_builder_requires_complete_matching_config() -> None: ) +@pytest.mark.parametrize( + "field", + [ + "forecast_content_sha256", + "scenario", + "geography_kind", + "geography_id", + "county_vintage", + "as_of", + ], +) +def test_completed_calculation_config_requires_every_resolved_field(field: str) -> None: + calculated = build_spm_calculation_provenance( + config=SELECTION, + receipt=receipt(), + ).model_dump(mode="json", by_alias=True) + del calculated["spm_config"][field] + + with pytest.raises(ValidationError): + SPMCalculationProvenance.model_validate(calculated) + + def test_comparison_collapses_identical_children_and_counts_them() -> None: comparison = build_spm_comparison_provenance( baseline_receipts=[receipt()] * 20, diff --git a/projects/policyengine-apis-integ/tests/test_spm_generated_client.py b/projects/policyengine-apis-integ/tests/test_spm_generated_client.py index 2cd685864..396d2e125 100644 --- a/projects/policyengine-apis-integ/tests/test_spm_generated_client.py +++ b/projects/policyengine-apis-integ/tests/test_spm_generated_client.py @@ -7,6 +7,7 @@ from policyengine_api_simulation_client.models import ( SPMCapability, SPMComparisonProvenance, + SPMResolvedConfiguration, SPMSelection, SimulationRequest, ) @@ -68,6 +69,33 @@ def test_generated_receipts_use_only_compact_typed_metadata(): assert SPMComparisonProvenance.from_dict(provenance).to_dict() == provenance +@pytest.mark.parametrize( + "field", + [ + "forecast_content_sha256", + "scenario", + "geography_kind", + "geography_id", + "county_vintage", + "as_of", + ], +) +def test_generated_completed_config_requires_all_resolved_fields(field): + resolved = { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + } + assert SPMResolvedConfiguration.from_dict(resolved).to_dict() == resolved + del resolved[field] + + with pytest.raises(KeyError): + SPMResolvedConfiguration.from_dict(resolved) + + def test_generated_comparison_rejects_old_receipt_lists(): with pytest.raises((KeyError, TypeError, ValueError)): SPMComparisonProvenance.from_dict({"baseline": [], "reform": []}) diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index c7cc77054..05e9f2cf7 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -88,6 +88,23 @@ def test_canonical_spm_extensions_are_public(): assert "spm" in schemas["BudgetWindowBatchRequest"]["properties"] assert "spm_provenance" in schemas["SingleYearMacroOutput"]["properties"] assert "spm_provenance" in schemas["BudgetWindowAnnualImpact"]["properties"] + assert schemas["SPMResolvedConfiguration"]["required"] == [ + "forecast_content_sha256", + "scenario", + "geography_kind", + "geography_id", + "county_vintage", + "as_of", + ] + assert all( + schemas["SPMRuntimeVersions"]["properties"][field]["minLength"] == 1 + for field in ( + "policyengine", + "policyengine-core", + "policyengine-us", + "spm-calculator", + ) + ) for path in ("/jobs/{job_id}", "/budget-window-jobs/{batch_job_id}"): assert "400" in spec["paths"][path]["get"]["responses"] diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py index cedff79b8..79425ab66 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py @@ -3,9 +3,13 @@ from __future__ import annotations from decimal import Decimal -from policyengine_simulation_contract.spm import SPMSelection, validate_spm_result from typing import Any +from policyengine_simulation_contract.spm import ( + SPMResolvedConfiguration, + validate_spm_result, +) + from policyengine_simulation_contract.gateway_models import ( BudgetWindowAnnualImpact, BudgetWindowResult, @@ -64,7 +68,9 @@ def extract_annual_impact( state_tax_revenue_impact = 0.0 return BudgetWindowAnnualImpact( - spm_config=SPMSelection.model_validate(spm) if spm is not None else None, + spm_config=( + SPMResolvedConfiguration.model_validate(spm) if spm is not None else None + ), spm_provenance=receipt, year=simulation_year, taxRevenueImpact=tax_revenue_impact, diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py index 85a26df60..27028ab99 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py @@ -235,24 +235,23 @@ def test_annual_child_and_segment_receipts_must_cover_requested_year(): combine_spm_results([result(), wrong_year], SELECTION, expected_year=2026) -def test_result_transport_may_omit_nulls_but_not_resolved_options(): - transported = result() - transported["spm_config"] = { - key: value - for key, value in transported["spm_config"].items() - if value is not None - } - assert combine_spm_results([transported], SELECTION)["spm_config"] == SELECTION - for field in ( +@pytest.mark.parametrize( + "field", + [ "forecast_content_sha256", "scenario", "geography_kind", + "geography_id", "county_vintage", - ): - malformed = deepcopy(transported) - del malformed["spm_config"][field] - with pytest.raises(SPMInputError, match="complete resolved"): - combine_spm_results([malformed], SELECTION) + "as_of", + ], +) +def test_result_transport_requires_every_resolved_option(field): + transported = result() + del transported["spm_config"][field] + + with pytest.raises(SPMInputError, match="Field required"): + combine_spm_results([transported], SELECTION) def test_sync_compatibility_endpoint_preserves_explicit_null( diff --git a/projects/policyengine-simulation-gateway/tests/golden/openapi.json b/projects/policyengine-simulation-gateway/tests/golden/openapi.json index 00a61112a..72055e359 100644 --- a/projects/policyengine-simulation-gateway/tests/golden/openapi.json +++ b/projects/policyengine-simulation-gateway/tests/golden/openapi.json @@ -409,7 +409,7 @@ "spm_config": { "anyOf": [ { - "$ref": "#/components/schemas/SPMSelection" + "$ref": "#/components/schemas/SPMResolvedConfiguration" }, { "type": "null" @@ -1908,22 +1908,90 @@ "title": "SPMProvenance", "description": "Compact detached receipt for one SPM calculation configuration." }, + "SPMResolvedConfiguration": { + "properties": { + "forecast_content_sha256": { + "type": "string", + "pattern": "^[0-9a-f]{64}$", + "title": "Forecast Content Sha256" + }, + "scenario": { + "type": "string", + "minLength": 1, + "pattern": "^\\S+$", + "title": "Scenario" + }, + "geography_kind": { + "type": "string", + "enum": [ + "county", + "national", + "metro" + ], + "title": "Geography Kind" + }, + "geography_id": { + "anyOf": [ + { + "type": "string", + "minLength": 1 + }, + { + "type": "null" + } + ], + "title": "Geography Id" + }, + "county_vintage": { + "type": "string", + "pattern": "^[0-9]{4}$", + "title": "County Vintage" + }, + "as_of": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "As Of" + } + }, + "additionalProperties": false, + "type": "object", + "required": [ + "forecast_content_sha256", + "scenario", + "geography_kind", + "geography_id", + "county_vintage", + "as_of" + ], + "title": "SPMResolvedConfiguration", + "description": "The complete six-field SPM selection recorded beside a receipt." + }, "SPMRuntimeVersions": { "properties": { "policyengine": { "type": "string", + "minLength": 1, "title": "Policyengine" }, "policyengine-core": { "type": "string", + "minLength": 1, "title": "Policyengine-Core" }, "policyengine-us": { "type": "string", + "minLength": 1, "title": "Policyengine-Us" }, "spm-calculator": { "type": "string", + "minLength": 1, "title": "Spm-Calculator" } }, @@ -1936,7 +2004,7 @@ "spm-calculator" ], "title": "SPMRuntimeVersions", - "description": "Package versions needed to reproduce an SPM calculation." + "description": "Package versions required to reproduce one certified SPM execution." }, "SPMSelection": { "properties": { @@ -2181,7 +2249,7 @@ "spm_config": { "anyOf": [ { - "$ref": "#/components/schemas/SPMSelection" + "$ref": "#/components/schemas/SPMResolvedConfiguration" }, { "type": "null" From 03d35a0e8638820e34d7b0d304b566bd618e42ac Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 1 Oct 2026 20:03:21 +0400 Subject: [PATCH 6/9] Stabilize incomplete SPM result errors --- .../policyengine_simulation_contract/spm.py | 20 +++++++++++++++++-- .../tests/test_canonical_spm.py | 3 ++- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py index bd455aa10..6955546ca 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -1,7 +1,7 @@ """Dependency-light public models for US SPM selection and calculation receipts.""" import re -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import date from typing import Any, Literal, Optional @@ -254,6 +254,22 @@ def _resolved_spm_configuration( ) +def _validate_completed_spm_configuration( + value: object, +) -> SPMResolvedConfiguration: + """Validate child metadata behind a stable typed-error boundary.""" + + if not isinstance(value, Mapping): + raise ValueError("Result has no complete resolved SPM selection") + missing = set(SPMResolvedConfiguration.model_fields).difference(value) + if missing: + raise ValueError( + "Result has no complete resolved SPM selection; missing " + + ", ".join(sorted(missing)) + ) + return SPMResolvedConfiguration.model_validate(value) + + def build_spm_provenance( *, forecast_id: str, @@ -504,7 +520,7 @@ def validate_spm_result( return None try: chosen = _resolved_spm_configuration(SPMSelection.model_validate(selection)) - result_selection = SPMResolvedConfiguration.model_validate( + result_selection = _validate_completed_spm_configuration( result.get("spm_config") ) if result_selection != chosen: diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py index 27028ab99..4a2594fe9 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py @@ -250,8 +250,9 @@ def test_result_transport_requires_every_resolved_option(field): transported = result() del transported["spm_config"][field] - with pytest.raises(SPMInputError, match="Field required"): + with pytest.raises(SPMInputError, match="complete resolved") as error: combine_spm_results([transported], SELECTION) + assert error.value.code == "SPM_CONFIGURATION_UNAVAILABLE" def test_sync_compatibility_endpoint_preserves_explicit_null( From 194eef8f42c756286ee1d184595315e358dd8786 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 1 Oct 2026 20:57:16 +0400 Subject: [PATCH 7/9] Use SPM receipts as completed result metadata --- .../stage-12-modal-worker-foundation.md | 9 ++ .../gateway_models.py | 4 +- .../macro_output.py | 6 +- .../policyengine_simulation_contract/spm.py | 114 +++++++++--------- .../tests/test_gateway_models.py | 14 +++ .../tests/test_spm_provenance.py | 25 ++-- .../tests/test_spm_generated_client.py | 33 ++--- .../tests/test_openapi.py | 11 +- .../fixtures/test_simulation_api_contracts.py | 2 - .../src/modal/budget_window_results.py | 8 +- .../policyengine_simulation_executor/spm.py | 1 - .../stage12_runtime/aggregation.py | 1 - .../tests/test_budget_window_scheduler.py | 1 - .../tests/test_canonical_spm.py | 15 ++- .../tests/test_canonical_spm_native.py | 7 +- .../tests/test_segmented_national.py | 14 +-- .../tests/test_simulation_output_builder.py | 3 +- .../tests/test_stage12_runtime.py | 5 +- .../tests/golden/openapi.json | 85 +------------ .../tests/test_spm_routes.py | 31 ++--- 20 files changed, 156 insertions(+), 233 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 1a9681d2c..f5a802179 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -208,6 +208,15 @@ metadata so PolicyEngine can correctly recognize conditional labor-supply analysis. `include_cliffs=true` is supported by this path and is carried through both output planning and aggregation. +For US SPM calculations, each child artifact stores one compact +`spm_provenance` receipt inside its calculation metadata. The coordinator +validates the requested selection against the forecast, scenario, geography, +county-vintage, and date fields in that receipt, then publishes only the +combined `spm_provenance` object in the completed report. Completed results do +not contain a sibling `spm_config`; the receipt is the sole source of resolved +SPM configuration. Request payloads still accept the partial `spm` selection +and resolve omitted options from the certified bundle before execution. + ### Temporary UK local-authority display metadata Stage 12 detects the local-authority boundary configuration from the complete diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py index a5c8bf202..97c16a11b 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py @@ -15,7 +15,6 @@ ) from policyengine_simulation_contract.spm import ( - SPMResolvedConfiguration, SPMSelection, SPMCapability, SPMComparisonProvenance, @@ -238,7 +237,8 @@ def validate_end_year(self) -> "BudgetWindowBatchRequest": class BudgetWindowAnnualImpact(BaseModel): """Annual budget-window impact row.""" - spm_config: Optional[SPMResolvedConfiguration] = None + model_config = ConfigDict(extra="forbid") + spm_provenance: Optional[SPMComparisonProvenance] = None year: str taxRevenueImpact: float diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py index 8986c0dbc..a46ca9203 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py @@ -6,10 +6,7 @@ from pydantic import BaseModel, ConfigDict, RootModel -from policyengine_simulation_contract.spm import ( - SPMComparisonProvenance, - SPMResolvedConfiguration, -) +from policyengine_simulation_contract.spm import SPMComparisonProvenance T = TypeVar("T") @@ -192,7 +189,6 @@ class CongressionalDistrictImpactOutput(MacroOutputModel): class SingleYearMacroOutput(MacroOutputModel): """Completed response returned by a single-year macro simulation.""" - spm_config: SPMResolvedConfiguration | None = None spm_provenance: SPMComparisonProvenance | None = None model_version: str diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py index 6955546ca..2261679e9 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -1,7 +1,7 @@ """Dependency-light public models for US SPM selection and calculation receipts.""" import re -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from datetime import date from typing import Any, Literal, Optional @@ -11,6 +11,7 @@ Field, SerializationInfo, SerializerFunctionWrapHandler, + ValidationError, field_validator, model_serializer, model_validator, @@ -158,6 +159,23 @@ class SPMProvenance(BaseModel): years: tuple[str, ...] runtime_versions: "SPMRuntimeVersions" + @model_serializer(mode="wrap") + def serialize_provenance( + self, handler: SerializerFunctionWrapHandler, info: SerializationInfo + ): + """Keep required nullable configuration fields in response bodies.""" + + dumped = handler(self) + if info.exclude_none: + for name in ("geography_id", "as_of"): + if ( + (info.include is None or name in info.include) + and (info.exclude is None or name not in info.exclude) + and getattr(self, name) is None + ): + dumped[name] = None + return dumped + @field_validator("years") @classmethod def validate_years(cls, value: tuple[str, ...]) -> tuple[str, ...]: @@ -194,33 +212,17 @@ class SPMRuntimeVersions(BaseModel): class SPMCalculationProvenance(BaseModel): - """Resolved SPM selection paired with its compact calculation receipt.""" + """One compact receipt stored with an internal calculation artifact.""" model_config = ConfigDict(extra="forbid", frozen=True) - spm_config: SPMResolvedConfiguration spm_provenance: SPMProvenance - @model_validator(mode="after") - def require_matching_receipt(self) -> "SPMCalculationProvenance": - expected = { - "forecast_sha256": self.spm_config.forecast_content_sha256, - "scenario": self.spm_config.scenario, - "geography_kind": self.spm_config.geography_kind, - "geography_id": self.spm_config.geography_id, - "county_vintage": self.spm_config.county_vintage, - "as_of": ( - date.fromisoformat(self.spm_config.as_of) - if self.spm_config.as_of is not None - else None - ), - } - for field_name, expected_value in expected.items(): - if getattr(self.spm_provenance, field_name) != expected_value: - raise ValueError( - f"SPM receipt {field_name} differs from the resolved selection" - ) - return self + @property + def spm_config(self) -> SPMResolvedConfiguration: + """Derive internal resolved configuration from the canonical receipt.""" + + return _resolved_spm_configuration_from_provenance(self.spm_provenance) def _require_complete_selection(selection: SPMSelection) -> None: @@ -254,20 +256,19 @@ def _resolved_spm_configuration( ) -def _validate_completed_spm_configuration( - value: object, +def _resolved_spm_configuration_from_provenance( + receipt: SPMProvenance, ) -> SPMResolvedConfiguration: - """Validate child metadata behind a stable typed-error boundary.""" + """Derive internal configuration from a receipt without a sibling object.""" - if not isinstance(value, Mapping): - raise ValueError("Result has no complete resolved SPM selection") - missing = set(SPMResolvedConfiguration.model_fields).difference(value) - if missing: - raise ValueError( - "Result has no complete resolved SPM selection; missing " - + ", ".join(sorted(missing)) - ) - return SPMResolvedConfiguration.model_validate(value) + return SPMResolvedConfiguration( + forecast_content_sha256=receipt.forecast_sha256, + scenario=receipt.scenario, + geography_kind=receipt.geography_kind, + geography_id=receipt.geography_id, + county_vintage=receipt.county_vintage, + as_of=receipt.as_of.isoformat() if receipt.as_of is not None else None, + ) def build_spm_provenance( @@ -308,10 +309,14 @@ def build_spm_calculation_provenance( ) -> SPMCalculationProvenance: """Pair a resolved selection with a matching compact receipt.""" - return SPMCalculationProvenance( - spm_config=_resolved_spm_configuration(config), - spm_provenance=receipt, - ) + resolved = _resolved_spm_configuration(config) + receipt_configuration = _resolved_spm_configuration_from_provenance(receipt) + for field_name in SPMResolvedConfiguration.model_fields: + if getattr(receipt_configuration, field_name) != getattr(resolved, field_name): + raise ValueError( + f"SPM receipt {field_name} differs from the resolved selection" + ) + return SPMCalculationProvenance(spm_provenance=receipt) SPM_CONTRACT_VERSION = "canonical-spm-v1" @@ -508,11 +513,13 @@ def validate_spm_result( result: dict, selection: Any, *, expected_year: int | str | None = None ): """Do not accept incomplete or mixed-method child/cached output.""" + if "spm_config" in result: + raise SPMInputError( + "SPM_CONFIGURATION_UNAVAILABLE", + "Completed results must not contain legacy spm_config metadata", + ) if selection is None: - if ( - result.get("spm_config") is not None - or result.get("spm_provenance") is not None - ): + if result.get("spm_provenance") is not None: raise SPMInputError( "SPM_CONFIGURATION_UNAVAILABLE", "Unexpected SPM receipt for a historical result", @@ -520,14 +527,16 @@ def validate_spm_result( return None try: chosen = _resolved_spm_configuration(SPMSelection.model_validate(selection)) - result_selection = _validate_completed_spm_configuration( - result.get("spm_config") - ) - if result_selection != chosen: - raise ValueError("Result SPM selection differs from the request") - provenance = SPMComparisonProvenance.model_validate( - result.get("spm_provenance") - ) + try: + provenance = SPMComparisonProvenance.model_validate( + result.get("spm_provenance") + ) + except ValidationError as exc: + if any(error["type"] == "missing" for error in exc.errors()): + raise ValueError( + "Result has no complete resolved SPM provenance" + ) from exc + raise for execution in (provenance.baseline, provenance.reform): build_spm_calculation_provenance( config=SPMSelection.model_validate(chosen.model_dump(mode="json")), @@ -577,8 +586,5 @@ def combine_spm_results( ], ) return { - "spm_config": _resolved_spm_configuration( - SPMSelection.model_validate(selection) - ).model_dump(mode="json"), "spm_provenance": combined.model_dump(mode="json", by_alias=True), } diff --git a/libs/policyengine-simulation-contract/tests/test_gateway_models.py b/libs/policyengine-simulation-contract/tests/test_gateway_models.py index c1f6dd2e3..07bfb51a0 100644 --- a/libs/policyengine-simulation-contract/tests/test_gateway_models.py +++ b/libs/policyengine-simulation-contract/tests/test_gateway_models.py @@ -622,6 +622,20 @@ def test_budget_window_batch_submit_response_serializes_correctly(self): class TestBudgetWindowBatchStatusResponse: """Tests for budget-window batch status responses.""" + def test_annual_impact_rejects_legacy_sibling_spm_config(self): + with pytest.raises(ValidationError): + BudgetWindowAnnualImpact.model_validate( + { + "spm_config": {"scenario": "ce_trend"}, + "year": "2026", + "taxRevenueImpact": 10, + "federalTaxRevenueImpact": 8, + "stateTaxRevenueImpact": 2, + "benefitSpendingImpact": -3, + "budgetaryImpact": 13, + } + ) + def test_budget_window_batch_status_response_accepts_child_jobs_and_result(self): response = BudgetWindowBatchStatusResponse( status="complete", diff --git a/libs/policyengine-simulation-contract/tests/test_spm_provenance.py b/libs/policyengine-simulation-contract/tests/test_spm_provenance.py index 121b27cea..e04e10085 100644 --- a/libs/policyengine-simulation-contract/tests/test_spm_provenance.py +++ b/libs/policyengine-simulation-contract/tests/test_spm_provenance.py @@ -48,6 +48,9 @@ def receipt(*, year: str = "2026", forecast_id: str = "forecast-2026"): def test_compact_receipt_has_exact_public_shape() -> None: value = receipt().model_dump(mode="json", by_alias=True) + without_optional_nulls = receipt().model_dump( + mode="json", by_alias=True, exclude_none=True + ) assert value == { "schema_version": "canonical-spm-provenance-v2", @@ -66,6 +69,8 @@ def test_compact_receipt_has_exact_public_shape() -> None: "spm-calculator": "1.0.0", }, } + assert without_optional_nulls["geography_id"] is None + assert without_optional_nulls["as_of"] is None assert len(json.dumps(value, separators=(",", ":")).encode()) < 1_024 @@ -157,11 +162,14 @@ def test_calculation_builder_requires_complete_matching_config() -> None: ) assert isinstance(calculated, SPMCalculationProvenance) assert isinstance(calculated.spm_config, SPMResolvedConfiguration) - assert calculated.model_dump(mode="json", exclude_none=True)["spm_config"] == { + assert calculated.spm_config.model_dump(mode="json") == { **SELECTION.model_dump(mode="json"), "geography_id": None, "as_of": None, } + assert calculated.model_dump(mode="json", by_alias=True) == { + "spm_provenance": receipt().model_dump(mode="json", by_alias=True) + } with pytest.raises(ValueError, match="scenario"): build_spm_calculation_provenance( @@ -170,23 +178,12 @@ def test_calculation_builder_requires_complete_matching_config() -> None: ) -@pytest.mark.parametrize( - "field", - [ - "forecast_content_sha256", - "scenario", - "geography_kind", - "geography_id", - "county_vintage", - "as_of", - ], -) -def test_completed_calculation_config_requires_every_resolved_field(field: str) -> None: +def test_calculation_provenance_rejects_legacy_sibling_config() -> None: calculated = build_spm_calculation_provenance( config=SELECTION, receipt=receipt(), ).model_dump(mode="json", by_alias=True) - del calculated["spm_config"][field] + calculated["spm_config"] = SELECTION.model_dump(mode="json") with pytest.raises(ValidationError): SPMCalculationProvenance.model_validate(calculated) diff --git a/projects/policyengine-apis-integ/tests/test_spm_generated_client.py b/projects/policyengine-apis-integ/tests/test_spm_generated_client.py index 396d2e125..7e37009ac 100644 --- a/projects/policyengine-apis-integ/tests/test_spm_generated_client.py +++ b/projects/policyengine-apis-integ/tests/test_spm_generated_client.py @@ -4,11 +4,13 @@ import pytest +from policyengine_api_simulation_client import models as generated_models from policyengine_api_simulation_client.models import ( + BudgetWindowAnnualImpact, SPMCapability, SPMComparisonProvenance, - SPMResolvedConfiguration, SPMSelection, + SingleYearMacroOutput, SimulationRequest, ) @@ -69,31 +71,10 @@ def test_generated_receipts_use_only_compact_typed_metadata(): assert SPMComparisonProvenance.from_dict(provenance).to_dict() == provenance -@pytest.mark.parametrize( - "field", - [ - "forecast_content_sha256", - "scenario", - "geography_kind", - "geography_id", - "county_vintage", - "as_of", - ], -) -def test_generated_completed_config_requires_all_resolved_fields(field): - resolved = { - "forecast_content_sha256": "a" * 64, - "scenario": "ce_trend", - "geography_kind": "national", - "geography_id": None, - "county_vintage": "2020", - "as_of": None, - } - assert SPMResolvedConfiguration.from_dict(resolved).to_dict() == resolved - del resolved[field] - - with pytest.raises(KeyError): - SPMResolvedConfiguration.from_dict(resolved) +def test_generated_completed_results_have_no_sibling_spm_config(): + assert "spm_config" not in SingleYearMacroOutput.__annotations__ + assert "spm_config" not in BudgetWindowAnnualImpact.__annotations__ + assert not hasattr(generated_models, "SPMResolvedConfiguration") def test_generated_comparison_rejects_old_receipt_lists(): diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index 05e9f2cf7..2a86571f1 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -88,14 +88,9 @@ def test_canonical_spm_extensions_are_public(): assert "spm" in schemas["BudgetWindowBatchRequest"]["properties"] assert "spm_provenance" in schemas["SingleYearMacroOutput"]["properties"] assert "spm_provenance" in schemas["BudgetWindowAnnualImpact"]["properties"] - assert schemas["SPMResolvedConfiguration"]["required"] == [ - "forecast_content_sha256", - "scenario", - "geography_kind", - "geography_id", - "county_vintage", - "as_of", - ] + assert "spm_config" not in schemas["SingleYearMacroOutput"]["properties"] + assert "spm_config" not in schemas["BudgetWindowAnnualImpact"]["properties"] + assert "SPMResolvedConfiguration" not in schemas assert all( schemas["SPMRuntimeVersions"]["properties"][field]["minLength"] == 1 for field in ( diff --git a/projects/policyengine-simulation-executor/fixtures/test_simulation_api_contracts.py b/projects/policyengine-simulation-executor/fixtures/test_simulation_api_contracts.py index f296f91dc..a86d9b830 100644 --- a/projects/policyengine-simulation-executor/fixtures/test_simulation_api_contracts.py +++ b/projects/policyengine-simulation-executor/fixtures/test_simulation_api_contracts.py @@ -1,7 +1,6 @@ """Fixtures for simulation API contract tests.""" CURRENT_SINGLE_YEAR_MACRO_KEYS = { - "spm_config", "spm_provenance", "model_version", "data_version", @@ -32,7 +31,6 @@ } CURRENT_SINGLE_YEAR_MACRO_RESULT = { - "spm_config": None, "spm_provenance": None, "model_version": "1.715.2", "data_version": "1.115.5", diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py index 79425ab66..8eabb3801 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py @@ -5,10 +5,7 @@ from decimal import Decimal from typing import Any -from policyengine_simulation_contract.spm import ( - SPMResolvedConfiguration, - validate_spm_result, -) +from policyengine_simulation_contract.spm import validate_spm_result from policyengine_simulation_contract.gateway_models import ( BudgetWindowAnnualImpact, @@ -68,9 +65,6 @@ def extract_annual_impact( state_tax_revenue_impact = 0.0 return BudgetWindowAnnualImpact( - spm_config=( - SPMResolvedConfiguration.model_validate(spm) if spm is not None else None - ), spm_provenance=receipt, year=simulation_year, taxRevenueImpact=tax_revenue_impact, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py index 90b735d49..078983954 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py @@ -295,6 +295,5 @@ def simulation_spm_result(baseline, reform, selection, *, expected_year=None): "Result SPM provenance does not cover the requested year", ) return { - "spm_config": baseline_provenance.spm_config.model_dump(mode="json"), "spm_provenance": comparison.model_dump(mode="json", by_alias=True), } diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py index 7e6d5d87c..7b56428f1 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py @@ -217,6 +217,5 @@ def build_spm_result( reform_receipts=[reform_provenance.spm_provenance], ) return { - "spm_config": selection.model_dump(mode="json"), "spm_provenance": comparison.model_dump(mode="json", by_alias=True), } diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py b/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py index abfe10e5a..1c0a0e032 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py @@ -58,7 +58,6 @@ def spm_child_result(runtime, simulation_year, *, receipt_year): }, } return { - "spm_config": dict(SPM_SELECTION), "spm_provenance": { "schema_version": "canonical-spm-comparison-v2", "baseline": {"receipt": receipt, "execution_count": 1}, diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py index 4a2594fe9..b091a9a23 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py @@ -94,7 +94,6 @@ def result(selection=SELECTION, year="2026"): }, ) return { - "spm_config": deepcopy(selection), "spm_provenance": { "schema_version": "canonical-spm-comparison-v2", "baseline": {"receipt": receipt, "execution_count": 1}, @@ -248,13 +247,25 @@ def test_annual_child_and_segment_receipts_must_cover_requested_year(): ) def test_result_transport_requires_every_resolved_option(field): transported = result() - del transported["spm_config"][field] + receipt_field = "forecast_sha256" if field == "forecast_content_sha256" else field + del transported["spm_provenance"]["baseline"]["receipt"][receipt_field] with pytest.raises(SPMInputError, match="complete resolved") as error: combine_spm_results([transported], SELECTION) assert error.value.code == "SPM_CONFIGURATION_UNAVAILABLE" +def test_result_transport_rejects_legacy_sibling_config(): + transported = result() + transported["spm_config"] = deepcopy(SELECTION) + + with pytest.raises( + SPMInputError, match="must not contain legacy spm_config" + ) as error: + combine_spm_results([transported], SELECTION) + assert error.value.code == "SPM_CONFIGURATION_UNAVAILABLE" + + def test_sync_compatibility_endpoint_preserves_explicit_null( monkeypatch, observability_runtime ): diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py index b03de5121..6b0e732e0 100644 --- a/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py @@ -63,7 +63,12 @@ def test_worker_baseline_reform_national_local_cache_and_receipts( ).all() result = simulation_spm_result(baseline, reform, baseline.spm_config) dumped = json.loads(json.dumps(result)) - assert dumped["spm_config"] == baseline.spm_config == reform.spm_config + assert "spm_config" not in dumped + assert baseline.spm_config == reform.spm_config + assert ( + dumped["spm_provenance"]["baseline"]["receipt"]["scenario"] + == baseline.spm_config["scenario"] + ) assert dumped["spm_provenance"]["reform"]["receipt"]["years"] == ["2024"] original = baseline.spm_provenance() original["years"].clear() diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national.py b/projects/policyengine-simulation-executor/tests/test_segmented_national.py index 9328b078e..547c202e3 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national.py @@ -423,7 +423,6 @@ def _spm_child(index, *, selection=SPM_SELECTION, year="2026"): """One child's result carrying that segment's own SPM receipts.""" return { "child": index, - "spm_config": dict(selection), "spm_provenance": { "schema_version": "canonical-spm-comparison-v2", "baseline": { @@ -439,9 +438,9 @@ def _spm_child(index, *, selection=SPM_SELECTION, year="2026"): def _spm_child_missing(index, field): - """A child whose transported selection lost a resolved option.""" + """A child whose canonical receipt lost a resolved option.""" child = _spm_child(index) - del child["spm_config"][field] + del child["spm_provenance"]["baseline"]["receipt"][field] return child @@ -463,16 +462,17 @@ def test__child_receipts_combine_into_one_national_receipt( output = runner.run() - # The resolved selection rides to every child unchanged and comes - # back on the parent as the one national selection. + # The resolved selection rides to every child unchanged. Completed + # output publishes that configuration only inside the receipt. assert all(p["spm"] == SPM_SELECTION for p in fake.spawned_payloads) - assert output["spm_config"] == SPM_SELECTION + assert "spm_config" not in output # One receipt describes all identical child executions. The counts # prove that every segment contributed without duplicating metadata. provenance = output["spm_provenance"] assert provenance["baseline"]["execution_count"] == 20 assert provenance["reform"]["execution_count"] == 20 assert provenance["baseline"]["receipt"]["forecast_id"] == "test-only" + assert provenance["baseline"]["receipt"]["county_vintage"] == "2020" assert output["budget"] == {} @pytest.mark.parametrize( @@ -480,7 +480,7 @@ def test__child_receipts_combine_into_one_national_receipt( [ ( _spm_child(7, selection={**SPM_SELECTION, "scenario": "zero_real"}), - "selection differs from the request", + "differs from the resolved selection", ), (_spm_child(7, year="2025"), "does not cover the requested year"), (_spm_child_missing(7, "county_vintage"), "complete resolved"), diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py index 490c24194..869489048 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -92,7 +92,7 @@ def _params_with_spm(params, selection): def _assert_macro_result(result, selection): """Check the ordinary macro result plus any installed SPM record.""" - spm_keys = ("spm_config", "spm_provenance") + spm_keys = ("spm_provenance",) assert {key: value for key, value in result.items() if key not in spm_keys} == { key: value for key, value in CURRENT_SINGLE_YEAR_MACRO_RESULT.items() @@ -101,7 +101,6 @@ def _assert_macro_result(result, selection): if selection is None: assert all(result[key] is None for key in spm_keys) return - assert result["spm_config"] == selection assert set(result["spm_provenance"]) == { "schema_version", "baseline", diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 3059df0a3..dab2a16d7 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -610,6 +610,7 @@ def test_single_worker_retains_detached_calculation_provenance() -> None: assert result["calculation_provenance"] == provenance.model_dump( mode="json", by_alias=True ) + assert "spm_config" not in result["calculation_provenance"] def test_single_worker_exposes_only_a_bounded_failure() -> None: @@ -663,7 +664,6 @@ def test_single_worker_rejects_frames_that_do_not_satisfy_the_output_plan() -> N def test_aggregate_combines_detached_spm_receipts() -> None: provenance = _spm_calculation_provenance() - selection = provenance.spm_config.model_dump(mode="json") report = _report() report = report.model_copy( @@ -693,9 +693,10 @@ def test_aggregate_combines_detached_spm_receipts() -> None: reform_descriptor=reform, ) - assert result["spm_config"] == selection + assert "spm_config" not in result assert result["spm_provenance"]["baseline"]["execution_count"] == 1 assert result["spm_provenance"]["reform"]["execution_count"] == 1 + assert result["spm_provenance"]["baseline"]["receipt"]["scenario"] == "official" def test_aggregate_stand_ins_preserve_policy_and_cliff_options(monkeypatch) -> None: diff --git a/projects/policyengine-simulation-gateway/tests/golden/openapi.json b/projects/policyengine-simulation-gateway/tests/golden/openapi.json index 72055e359..346a301cb 100644 --- a/projects/policyengine-simulation-gateway/tests/golden/openapi.json +++ b/projects/policyengine-simulation-gateway/tests/golden/openapi.json @@ -406,16 +406,6 @@ }, "BudgetWindowAnnualImpact": { "properties": { - "spm_config": { - "anyOf": [ - { - "$ref": "#/components/schemas/SPMResolvedConfiguration" - }, - { - "type": "null" - } - ] - }, "spm_provenance": { "anyOf": [ { @@ -451,6 +441,7 @@ "title": "Budgetaryimpact" } }, + "additionalProperties": false, "type": "object", "required": [ "year", @@ -1908,70 +1899,6 @@ "title": "SPMProvenance", "description": "Compact detached receipt for one SPM calculation configuration." }, - "SPMResolvedConfiguration": { - "properties": { - "forecast_content_sha256": { - "type": "string", - "pattern": "^[0-9a-f]{64}$", - "title": "Forecast Content Sha256" - }, - "scenario": { - "type": "string", - "minLength": 1, - "pattern": "^\\S+$", - "title": "Scenario" - }, - "geography_kind": { - "type": "string", - "enum": [ - "county", - "national", - "metro" - ], - "title": "Geography Kind" - }, - "geography_id": { - "anyOf": [ - { - "type": "string", - "minLength": 1 - }, - { - "type": "null" - } - ], - "title": "Geography Id" - }, - "county_vintage": { - "type": "string", - "pattern": "^[0-9]{4}$", - "title": "County Vintage" - }, - "as_of": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "As Of" - } - }, - "additionalProperties": false, - "type": "object", - "required": [ - "forecast_content_sha256", - "scenario", - "geography_kind", - "geography_id", - "county_vintage", - "as_of" - ], - "title": "SPMResolvedConfiguration", - "description": "The complete six-field SPM selection recorded beside a receipt." - }, "SPMRuntimeVersions": { "properties": { "policyengine": { @@ -2246,16 +2173,6 @@ }, "SingleYearMacroOutput": { "properties": { - "spm_config": { - "anyOf": [ - { - "$ref": "#/components/schemas/SPMResolvedConfiguration" - }, - { - "type": "null" - } - ] - }, "spm_provenance": { "anyOf": [ { diff --git a/projects/policyengine-simulation-gateway/tests/test_spm_routes.py b/projects/policyengine-simulation-gateway/tests/test_spm_routes.py index a500b1310..53d214880 100644 --- a/projects/policyengine-simulation-gateway/tests/test_spm_routes.py +++ b/projects/policyengine-simulation-gateway/tests/test_spm_routes.py @@ -264,14 +264,10 @@ def spm_provenance(year): } -def test_completed_result_body_keeps_resolved_nulls(mock_modal, client): - """A client must be able to replay the selection it was handed back. - - ``response_model_exclude_none`` applies to every poll body, and an - omitted option inherits the bundle default. Dropping a resolved null - would silently re-resolve ``as_of`` and ``geography_id`` on the next - request, changing the baseline key and the receipts. - """ +def test_completed_result_publishes_configuration_only_in_provenance( + mock_modal, client +): + """The compact receipt is the sole completed-result configuration.""" shared_app_state(mock_modal, sibling_model="1.824.7") submitted = client.post( "/simulate/economy/comparison", @@ -288,17 +284,22 @@ def test_completed_result_body_keeps_resolved_nulls(mock_modal, client): call = mock_modal["function_call"].registry[job_id] call.result = { **call.result, - "spm_config": RESOLVED_SELECTION, "spm_provenance": spm_provenance(2026), } polled = client.get(f"/jobs/{job_id}") assert polled.status_code == 200, polled.text - assert polled.json()["result"]["spm_config"] == RESOLVED_SELECTION + result = polled.json()["result"] + assert "spm_config" not in result + receipt = result["spm_provenance"]["baseline"]["receipt"] + assert receipt["geography_id"] is None + assert receipt["as_of"] is None -def test_completed_budget_window_rows_keep_resolved_nulls(mock_modal, client): - """Each ``annualImpacts`` row carries the same replayable selection.""" +def test_completed_budget_window_configuration_is_only_in_provenance( + mock_modal, client +): + """Each annual row publishes configuration only in its receipt.""" from policyengine_simulation_contract.budget_window_state import ( put_batch_job_state, ) @@ -312,7 +313,6 @@ def test_completed_budget_window_rows_keep_resolved_nulls(mock_modal, client): shared_app_state(mock_modal, sibling_model="1.824.7") impact = BudgetWindowAnnualImpact( - spm_config=RESOLVED_SELECTION, spm_provenance=spm_provenance(2026), year="2026", taxRevenueImpact=10, @@ -365,7 +365,10 @@ def test_completed_budget_window_rows_keep_resolved_nulls(mock_modal, client): polled = client.get("/budget-window-jobs/mock-batch-job-id-123") assert polled.status_code == 200, polled.text row = polled.json()["result"]["annualImpacts"][0] - assert row["spm_config"] == RESOLVED_SELECTION + assert "spm_config" not in row + receipt = row["spm_provenance"]["baseline"]["receipt"] + assert receipt["geography_id"] is None + assert receipt["as_of"] is None # Every optional field the pre-SPM gateway emitted in a raw 202/500 body. From 800fe4abb00ea012e42e44871dde83eb7ca2d501 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 1 Oct 2026 21:23:39 +0400 Subject: [PATCH 8/9] Align Stage 12 provenance contracts --- .../policyengine_simulation_contract/spm.py | 32 +++--- .../stage12_execution.py | 20 +++- .../tests/test_spm_provenance.py | 38 +++---- .../tests/test_stage12_execution.py | 101 ++++++++++++++++++ .../policyengine_simulation_executor/spm.py | 11 +- .../stage12_artifacts.py | 10 +- .../stage12_qualification.py | 7 +- .../stage12_runtime/aggregation.py | 12 +-- .../stage12_runtime/simulation.py | 4 +- .../tests/test_stage12_artifacts.py | 4 +- .../tests/test_stage12_runtime.py | 4 +- 11 files changed, 172 insertions(+), 71 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py index 2261679e9..1df9228a0 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -98,7 +98,7 @@ def validate_location(self): class SPMResolvedConfiguration(BaseModel): - """The complete six-field SPM selection recorded beside a receipt.""" + """An internal complete SPM selection derived from a compact receipt.""" model_config = ConfigDict(frozen=True, extra="forbid") @@ -211,20 +211,6 @@ class SPMRuntimeVersions(BaseModel): spm_calculator: str = Field(min_length=1, alias="spm-calculator") -class SPMCalculationProvenance(BaseModel): - """One compact receipt stored with an internal calculation artifact.""" - - model_config = ConfigDict(extra="forbid", frozen=True) - - spm_provenance: SPMProvenance - - @property - def spm_config(self) -> SPMResolvedConfiguration: - """Derive internal resolved configuration from the canonical receipt.""" - - return _resolved_spm_configuration_from_provenance(self.spm_provenance) - - def _require_complete_selection(selection: SPMSelection) -> None: required = set(SPMSelection.model_fields) missing = required.difference(selection.model_fields_set) @@ -302,12 +288,12 @@ def build_spm_provenance( ) -def build_spm_calculation_provenance( +def validate_spm_calculation_provenance( *, config: SPMSelection, receipt: SPMProvenance, -) -> SPMCalculationProvenance: - """Pair a resolved selection with a matching compact receipt.""" +) -> SPMProvenance: + """Validate and return the sole compact receipt for a calculation.""" resolved = _resolved_spm_configuration(config) receipt_configuration = _resolved_spm_configuration_from_provenance(receipt) @@ -316,7 +302,7 @@ def build_spm_calculation_provenance( raise ValueError( f"SPM receipt {field_name} differs from the resolved selection" ) - return SPMCalculationProvenance(spm_provenance=receipt) + return receipt SPM_CONTRACT_VERSION = "canonical-spm-v1" @@ -401,6 +387,12 @@ class SPMComparisonProvenance(BaseModel): baseline: SPMExecutionProvenance reform: SPMExecutionProvenance + @model_validator(mode="after") + def validate_compatible_receipts(self) -> "SPMComparisonProvenance": + if self.baseline.receipt != self.reform.receipt: + raise ValueError("baseline and reform SPM receipts differ") + return self + def _collapse_spm_receipts( side: str, @@ -538,7 +530,7 @@ def validate_spm_result( ) from exc raise for execution in (provenance.baseline, provenance.reform): - build_spm_calculation_provenance( + validate_spm_calculation_provenance( config=SPMSelection.model_validate(chosen.model_dump(mode="json")), receipt=execution.receipt, ) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py index 37069da46..4241c4082 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py @@ -18,9 +18,15 @@ field_validator, model_validator, ) -from policyengine_simulation_contract.spm import SPMCalculationProvenance + +from policyengine_simulation_contract.spm import ( + SPMComparisonProvenance, + SPMProvenance, +) from policyengine_simulation_contract.uk_geography import ( UKLocalAuthorityBoundaryVersion as UKLocalAuthorityBoundaryVersion, +) +from policyengine_simulation_contract.uk_geography import ( UKLocalAuthorityMetadata as UKLocalAuthorityMetadata, ) @@ -333,7 +339,7 @@ class SimulationArtifactDescriptor(StrictContractModel): output_plan_sha256: Sha256Digest row_identity: RowIdentity bundle: BundleProvenance - calculation_provenance: SPMCalculationProvenance | None = None + calculation_provenance: SPMProvenance | None = None class ReportExecutionInput(StrictContractModel): @@ -401,6 +407,16 @@ def require_unique_aggregates( raise ValueError("requested aggregates must be unique") return value + @model_validator(mode="after") + def validate_spm_provenance(self) -> AggregateReportArtifactPayload: + """Reject duplicate or noncanonical SPM data in aggregate artifacts.""" + + if "spm_config" in self.result: + raise ValueError("aggregate result must not include legacy spm_config") + if "spm_provenance" in self.result: + SPMComparisonProvenance.model_validate(self.result["spm_provenance"]) + return self + class ResultDifference(StrictContractModel): """One exact aggregate-result difference identified by JSON Pointer.""" diff --git a/libs/policyengine-simulation-contract/tests/test_spm_provenance.py b/libs/policyengine-simulation-contract/tests/test_spm_provenance.py index e04e10085..896524cf4 100644 --- a/libs/policyengine-simulation-contract/tests/test_spm_provenance.py +++ b/libs/policyengine-simulation-contract/tests/test_spm_provenance.py @@ -6,18 +6,15 @@ from pydantic import ValidationError from policyengine_simulation_contract.spm import ( - SPMCalculationProvenance, SPMComparisonProvenance, SPMProvenance, - SPMResolvedConfiguration, SPMRuntimeVersions, SPMSelection, - build_spm_calculation_provenance, build_spm_comparison_provenance, build_spm_provenance, + validate_spm_calculation_provenance, ) - SELECTION = SPMSelection( forecast_content_sha256="a" * 64, scenario="ce_trend", @@ -156,37 +153,31 @@ def test_builder_requires_receipt_to_match_resolved_selection() -> None: def test_calculation_builder_requires_complete_matching_config() -> None: - calculated = build_spm_calculation_provenance( + calculated = validate_spm_calculation_provenance( config=SELECTION, receipt=receipt(), ) - assert isinstance(calculated, SPMCalculationProvenance) - assert isinstance(calculated.spm_config, SPMResolvedConfiguration) - assert calculated.spm_config.model_dump(mode="json") == { - **SELECTION.model_dump(mode="json"), - "geography_id": None, - "as_of": None, - } - assert calculated.model_dump(mode="json", by_alias=True) == { - "spm_provenance": receipt().model_dump(mode="json", by_alias=True) - } + assert calculated == receipt() + assert calculated.model_dump(mode="json", by_alias=True) == receipt().model_dump( + mode="json", by_alias=True + ) with pytest.raises(ValueError, match="scenario"): - build_spm_calculation_provenance( + validate_spm_calculation_provenance( config=SELECTION.model_copy(update={"scenario": "zero_real"}), receipt=receipt(), ) def test_calculation_provenance_rejects_legacy_sibling_config() -> None: - calculated = build_spm_calculation_provenance( + calculated = validate_spm_calculation_provenance( config=SELECTION, receipt=receipt(), ).model_dump(mode="json", by_alias=True) calculated["spm_config"] = SELECTION.model_dump(mode="json") with pytest.raises(ValidationError): - SPMCalculationProvenance.model_validate(calculated) + SPMProvenance.model_validate(calculated) def test_comparison_collapses_identical_children_and_counts_them() -> None: @@ -220,6 +211,17 @@ def test_comparison_rejects_missing_or_mismatched_children() -> None: ) +def test_comparison_model_rejects_mismatched_receipts() -> None: + comparison = build_spm_comparison_provenance( + baseline_receipts=[receipt()], + reform_receipts=[receipt()], + ).model_dump(mode="json", by_alias=True) + comparison["reform"]["receipt"]["years"] = ["2027"] + + with pytest.raises(ValidationError, match="baseline and reform"): + SPMComparisonProvenance.model_validate(comparison) + + def test_comparison_contract_rejects_old_receipt_lists() -> None: with pytest.raises(ValidationError): SPMComparisonProvenance.model_validate( diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_execution.py b/libs/policyengine-simulation-contract/tests/test_stage12_execution.py index 62b5ba87b..7c8a326a0 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_execution.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_execution.py @@ -2,13 +2,27 @@ from __future__ import annotations +from typing import cast +from uuid import UUID + import pytest +from pydantic import JsonValue +from policyengine_simulation_contract.spm import ( + SPMRuntimeVersions, + SPMSelection, + build_spm_comparison_provenance, + build_spm_provenance, +) from policyengine_simulation_contract.stage12_execution import ( SIMULATION_PARQUET_PAYLOAD_CONTRACT, + AggregateReportArtifactPayload, + BundleProvenance, + DatasetProvenance, EntityOutputPlan, ReportAggregate, ReportOutputRequirements, + SimulationArtifactDescriptor, Stage12OutputPlan, stage12_output_plan_sha256, ) @@ -41,6 +55,54 @@ def _plan() -> Stage12OutputPlan: ) +def _bundle() -> BundleProvenance: + return BundleProvenance( + policyengine_version="6.2.1", + country_package_name="policyengine-us", + country_package_version="2.2.1", + dataset=DatasetProvenance( + identity="populace_us_2024", + uri="hf://policyengine/data/populace_us_2024.h5@revision", + artifact_revision="revision", + data_package_name="populace-data", + data_package_version="0.1.0", + ), + bundle_manifest_sha256="b" * 64, + ) + + +def _spm_comparison() -> dict[str, JsonValue]: + selection = SPMSelection( + forecast_content_sha256="a" * 64, + scenario="official", + geography_kind="national", + geography_id=None, + county_vintage="2020", + as_of=None, + ) + receipt = build_spm_provenance( + forecast_id="forecast-2026", + forecast_sha256="a" * 64, + selection=selection, + years=("2026",), + runtime_versions=SPMRuntimeVersions.model_validate( + { + "policyengine": "6.2.1", + "policyengine-core": "3.32.10", + "policyengine-us": "2.2.1", + "spm-calculator": "1.0.0", + } + ), + ) + return cast( + dict[str, JsonValue], + build_spm_comparison_provenance( + baseline_receipts=(receipt,), + reform_receipts=(receipt,), + ).model_dump(mode="json", by_alias=True), + ) + + def test_output_plan_has_a_deterministic_digest() -> None: plan = _plan() @@ -141,3 +203,42 @@ def test_parquet_contract_names_uk_local_authority_metadata() -> None: SIMULATION_PARQUET_PAYLOAD_CONTRACT.uk_local_authority_metadata_key == "policyengine.stage12.uk_local_authority_metadata" ) + + +def test_simulation_artifact_schema_uses_the_receipt_without_a_wrapper() -> None: + schema_text = str(SimulationArtifactDescriptor.model_json_schema()) + + assert "SPMProvenance" in schema_text + assert "SPMCalculationProvenance" not in schema_text + + +def test_aggregate_artifact_accepts_only_canonical_spm_provenance() -> None: + comparison = _spm_comparison() + artifact = AggregateReportArtifactPayload( + evaluation_id=UUID("00000000-0000-0000-0000-000000000001"), + requested_aggregates=(ReportAggregate.BUDGET,), + bundle=_bundle(), + result={"spm_provenance": comparison}, + ) + + assert artifact.result["spm_provenance"] == comparison + + with pytest.raises(ValueError, match="legacy spm_config"): + AggregateReportArtifactPayload( + evaluation_id=artifact.evaluation_id, + requested_aggregates=artifact.requested_aggregates, + bundle=artifact.bundle, + result={"spm_config": {}}, + ) + + mismatched = _spm_comparison() + reform = cast(dict[str, JsonValue], mismatched["reform"]) + receipt = cast(dict[str, JsonValue], reform["receipt"]) + receipt["years"] = ["2027"] + with pytest.raises(ValueError, match="baseline and reform"): + AggregateReportArtifactPayload( + evaluation_id=artifact.evaluation_id, + requested_aggregates=artifact.requested_aggregates, + bundle=artifact.bundle, + result={"spm_provenance": mismatched}, + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py index 078983954..45daa3a0c 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py @@ -6,12 +6,11 @@ from policyengine_simulation_contract.spm import ( SPMCapability, - SPMCalculationProvenance, SPMInputError, SPMProvenance, SPMRuntimeVersions, SPMSelection, - build_spm_calculation_provenance, + validate_spm_calculation_provenance, build_spm_comparison_provenance, build_spm_provenance, resolve_spm_selection, @@ -251,7 +250,7 @@ def compact_spm_provenance( def simulation_spm_calculation_provenance( simulation: object, selection: object, -) -> SPMCalculationProvenance: +) -> SPMProvenance: """Read and compact one simulation's resolved SPM receipt.""" resolved = SPMSelection.model_validate(selection) @@ -267,7 +266,7 @@ def simulation_spm_calculation_provenance( "Simulation has no valid SPM calculation receipt", ) try: - return build_spm_calculation_provenance( + return validate_spm_calculation_provenance( config=resolved, receipt=compact_spm_provenance(receipt(), resolved), ) @@ -284,8 +283,8 @@ def simulation_spm_result(baseline, reform, selection, *, expected_year=None): baseline_provenance = simulation_spm_calculation_provenance(baseline, selection) reform_provenance = simulation_spm_calculation_provenance(reform, selection) comparison = build_spm_comparison_provenance( - baseline_receipts=[baseline_provenance.spm_provenance], - reform_receipts=[reform_provenance.spm_provenance], + baseline_receipts=[baseline_provenance], + reform_receipts=[reform_provenance], ) if expected_year is not None and str(expected_year) not in ( comparison.baseline.receipt.years diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py index 7be267fd9..ff24ea350 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py @@ -24,7 +24,7 @@ SimulationArtifactDescriptor, stage12_output_plan_sha256, ) -from policyengine_simulation_contract.spm import SPMCalculationProvenance +from policyengine_simulation_contract.spm import SPMProvenance from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata from policyengine_simulation_executor.artifact_store import ArtifactStore @@ -118,7 +118,7 @@ def _frame_payload( def serialize_simulation_frames( frames: Mapping[str, pd.DataFrame], *, - calculation_provenance: SPMCalculationProvenance | None = None, + calculation_provenance: SPMProvenance | None = None, uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, ) -> tuple[bytes, RowIdentity]: combined, dtypes, identifier_values = _frame_payload(frames) @@ -214,13 +214,13 @@ def _parquet_schema_metadata(payload: bytes) -> Mapping[bytes, bytes]: def deserialize_calculation_provenance( payload: bytes, -) -> SPMCalculationProvenance | None: +) -> SPMProvenance | None: raw = _parquet_schema_metadata(payload).get( PARQUET_CONTRACT.calculation_provenance_metadata_key.encode() ) if raw is None: return None - return SPMCalculationProvenance.model_validate_json(raw) + return SPMProvenance.model_validate_json(raw) def deserialize_uk_local_authority_metadata( @@ -289,7 +289,7 @@ def write_simulation( prefix: str, simulation: PlannedSimulationExecutionInput, frames: Mapping[str, pd.DataFrame], - calculation_provenance: SPMCalculationProvenance | None = None, + calculation_provenance: SPMProvenance | None = None, uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, ) -> SimulationArtifactDescriptor: payload, row_identity = serialize_simulation_frames( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py index 8ffd4561f..fc2ec3c38 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py @@ -5,7 +5,7 @@ from collections.abc import Callable, Mapping from concurrent.futures import ThreadPoolExecutor from hashlib import sha256 -from typing import Any, Literal, cast +from typing import Any, Literal import pandas as pd from policyengine_observability import ( @@ -208,10 +208,7 @@ def _descriptor( output_plan_sha256=stage12_output_plan_sha256(simulation.output_plan), row_identity=row_identity, bundle=simulation.bundle, - calculation_provenance=cast( - dict[str, Any] | None, - calculation.calculation_provenance, - ), + calculation_provenance=calculation.calculation_provenance, ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py index 7b56428f1..b60db9243 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py @@ -203,18 +203,12 @@ def build_spm_result( reform_provenance = reform_descriptor.calculation_provenance if baseline_provenance is None or reform_provenance is None: raise ValueError("SPM calculation provenance is missing") - selection = baseline_provenance.spm_config - if selection != reform_provenance.spm_config: - raise ValueError("SPM calculation selections do not match") - for receipt in ( - baseline_provenance.spm_provenance, - reform_provenance.spm_provenance, - ): + for receipt in (baseline_provenance, reform_provenance): if str(report.baseline.year) not in receipt.years: raise ValueError("SPM calculation provenance does not cover report year") comparison = build_spm_comparison_provenance( - baseline_receipts=[baseline_provenance.spm_provenance], - reform_receipts=[reform_provenance.spm_provenance], + baseline_receipts=[baseline_provenance], + reform_receipts=[reform_provenance], ) return { "spm_provenance": comparison.model_dump(mode="json", by_alias=True), diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 7bf43c91b..46c499f1a 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -21,7 +21,7 @@ Stage12InvocationContext, stage12_output_plan_sha256, ) -from policyengine_simulation_contract.spm import SPMCalculationProvenance +from policyengine_simulation_contract.spm import SPMProvenance from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata from policyengine_simulation_observability.stages import ( STAGE12_SIMULATION_STAGES, @@ -44,7 +44,7 @@ @dataclass(frozen=True) class SimulationCalculation: frames: Mapping[str, pd.DataFrame] - calculation_provenance: SPMCalculationProvenance | None = None + calculation_provenance: SPMProvenance | None = None uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py b/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py index b54b1a938..cacc5c332 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py @@ -19,8 +19,8 @@ from policyengine_simulation_contract.spm import ( SPMRuntimeVersions, SPMSelection, - build_spm_calculation_provenance, build_spm_provenance, + validate_spm_calculation_provenance, ) from pydantic import ValidationError @@ -75,7 +75,7 @@ def _calculation_provenance(): } ), ) - return build_spm_calculation_provenance(config=selection, receipt=receipt) + return validate_spm_calculation_provenance(config=selection, receipt=receipt) def test_parquet_encoding_is_deterministic_and_preserves_rows_and_dtypes() -> None: diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index dab2a16d7..d291a555e 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -43,8 +43,8 @@ from policyengine_simulation_contract.spm import ( SPMRuntimeVersions, SPMSelection, - build_spm_calculation_provenance, build_spm_provenance, + validate_spm_calculation_provenance, ) from policyengine_simulation_contract.uk_geography import ( UKLocalAuthorityBoundaryVersion, @@ -99,7 +99,7 @@ def _spm_calculation_provenance(): } ), ) - return build_spm_calculation_provenance(config=selection, receipt=receipt) + return validate_spm_calculation_provenance(config=selection, receipt=receipt) def _bundle() -> BundleProvenance: From 52bfe14a8171ed48caa0ddca402b6da82604866a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 1 Oct 2026 21:50:10 +0400 Subject: [PATCH 9/9] Clarify synthetic SPM version fixture --- .../tests/test_spm_provenance.py | 31 ++++++++++--------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/libs/policyengine-simulation-contract/tests/test_spm_provenance.py b/libs/policyengine-simulation-contract/tests/test_spm_provenance.py index 896524cf4..a3eb99d24 100644 --- a/libs/policyengine-simulation-contract/tests/test_spm_provenance.py +++ b/libs/policyengine-simulation-contract/tests/test_spm_provenance.py @@ -3,8 +3,6 @@ import json import pytest -from pydantic import ValidationError - from policyengine_simulation_contract.spm import ( SPMComparisonProvenance, SPMProvenance, @@ -14,6 +12,7 @@ build_spm_provenance, validate_spm_calculation_provenance, ) +from pydantic import ValidationError SELECTION = SPMSelection( forecast_content_sha256="a" * 64, @@ -23,12 +22,14 @@ county_vintage="2020", as_of=None, ) -VERSIONS = SPMRuntimeVersions.model_validate( +# Synthetic fixture values used only to exercise the public receipt contract. +# Deployed workers obtain real package versions from the PolicyEngine.py bundle. +TEST_RUNTIME_VERSIONS = SPMRuntimeVersions.model_validate( { - "policyengine": "6.2.1", - "policyengine-core": "3.32.10", - "policyengine-us": "2.2.1", - "spm-calculator": "1.0.0", + "policyengine": "0.0.0-test-policyengine", + "policyengine-core": "0.0.0-test-policyengine-core", + "policyengine-us": "0.0.0-test-policyengine-us", + "spm-calculator": "0.0.0-test-spm-calculator", } ) @@ -39,7 +40,7 @@ def receipt(*, year: str = "2026", forecast_id: str = "forecast-2026"): forecast_sha256="a" * 64, selection=SELECTION, years=(year,), - runtime_versions=VERSIONS, + runtime_versions=TEST_RUNTIME_VERSIONS, ) @@ -60,10 +61,10 @@ def test_compact_receipt_has_exact_public_shape() -> None: "as_of": None, "years": ["2026"], "runtime_versions": { - "policyengine": "6.2.1", - "policyengine-core": "3.32.10", - "policyengine-us": "2.2.1", - "spm-calculator": "1.0.0", + "policyengine": "0.0.0-test-policyengine", + "policyengine-core": "0.0.0-test-policyengine-core", + "policyengine-us": "0.0.0-test-policyengine-us", + "spm-calculator": "0.0.0-test-spm-calculator", }, } assert without_optional_nulls["geography_id"] is None @@ -120,7 +121,7 @@ def test_compact_receipt_rejects_old_rich_shape() -> None: def test_runtime_versions_reject_invalid_package_versions( package: str, invalid_value: str | None ) -> None: - versions = VERSIONS.model_dump(mode="json", by_alias=True) + versions = TEST_RUNTIME_VERSIONS.model_dump(mode="json", by_alias=True) invalid_versions: dict[str, str | None] = {**versions, package: invalid_value} with pytest.raises(ValidationError): @@ -135,7 +136,7 @@ def test_compact_receipt_serializes_as_of_as_an_iso_date() -> None: update={"as_of": "2026-09-09"}, ), years=("2026",), - runtime_versions=VERSIONS, + runtime_versions=TEST_RUNTIME_VERSIONS, ) assert dated.model_dump(mode="json", by_alias=True)["as_of"] == "2026-09-09" @@ -148,7 +149,7 @@ def test_builder_requires_receipt_to_match_resolved_selection() -> None: forecast_sha256="b" * 64, selection=SELECTION, years=("2026",), - runtime_versions=VERSIONS, + runtime_versions=TEST_RUNTIME_VERSIONS, )