diff --git a/changelog_entry.yaml b/changelog_entry.yaml index ce4676bc5..2e68ba388 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -16,6 +16,7 @@ - Ran Stage 12 US national simulations as 20 region-group calls each, and built an independent, immutable Stage 12 cache of three annual US datasets and 60 current-law baselines before each deployment, publishing it only after a recorded determinism check and recording its release in the v2 version manifest under a new key, while the existing key keeps a copy that earlier readers accept - Made the budget refuse a household count outside its country's plausible range (above 200 million in the US or 60 million in the UK, negative or not finite), raising ImplausibleEntityCountError instead of publishing a count like #767's changed: + - Moved the worker's credential setup, request year and policy parsing, and region resolution out of simulation_runtime.py (1,111 lines, now 651) into gcp_credentials.py, request_policy.py and region_resolution.py; simulation_runtime.py keeps dataset selection and loading, simulation building, local-area breakdowns and the run itself, which tests replace by name there, and re-exports the moved names, so every caller still imports them from simulation_runtime; the shared helpers _parse_year, _normalise_policy, _compile_interval_policy, _runtime_policy (now wrapper_policy), _resolve_region, _resolve_dataset_reference and _require_local_area_breakdown are now public, and the credential messages still log under the simulation_runtime logger; the 3,048-line output builder test file is split into seven test files and a module of shared fakes, all named test_simulation_output_*.py so the weighted-frame rule still checks them - Split the artifact precompute (precompute.py, 931 lines) and the Stage 12 cache precompute (stage12_cache/precompute.py, 844 lines) each into planning, dataset-building and baseline modules, with the manifest and run_precompute left in the original module, which re-exports every public name, so the Modal precompute apps and their tests import from the same paths; the cache precompute's _scoped_planned_simulation and _baseline_identity are now public, and the worker validation imports scoped_planned_simulation from its new module; the 1,835-line precompute test file is split into four, and the cache precompute tests patch each name in the module whose code reads it - Split the 4,709-line Stage 12 runtime test file into thirteen test files by behaviour (dataset loading, provenance, the single worker and its failures, aggregates, the coordinator and its failures, runtime store retries, failure records, cancellation, child state errors and the wait window) and three support modules for the shared builders, fake stores and fake child calls, whose helpers are now public; the other Stage 12 tests import them from those modules, every test keeps its name and body, and two parametrized cases are named after the renamed bundle helper - Split Stage 12 report coordination (stage12_runtime/coordination.py, 1,091 lines) into modules for the report coordinator class, its child simulations' records, failures and cancellation, and the parent report's checks and claim, with coordinate_report and every public name still in coordination.py and child failures still logged under its logger; and turned the contract's stage12_execution.py (929 lines) into a stage12_execution package of modules for the shared field types and simulation input, report inputs and output plans, result artifacts, and comparison-run records, which re-exports every name from the same import path, with unchanged JSON schemas; models subclass the base models through the package path ruff.toml lists, so no ruff setting changes diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 7f8323b2c..4fc82bdaf 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -66,7 +66,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt` - `sim.microdata.weight-stripping-confined`: 0 remaining in `rules/baselines/sim-microdata-weight-stripping.txt` -- `sim.quality.python-file-length`: 4 remaining in `rules/baselines/python-file-length.txt` +- `sim.quality.python-file-length`: 2 remaining in `rules/baselines/python-file-length.txt` ### Exceptions and review dates diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py index 955b42d20..73d3a82a9 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -581,7 +581,7 @@ def _reduce( self, child_results: list[dict], *, country_module: ModuleType ) -> dict[str, Any]: from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # Modal worker code: bound per call, like this module's other simulation_runtime imports - _parse_year, + parse_year, ) # The canonical internal-key stripper, plus the opt-out knob. @@ -597,7 +597,7 @@ def _reduce( country=self.country, simulation_params=simulation_params, country_module=country_module, - year=_parse_year(simulation_params), + year=parse_year(simulation_params), resolved_data_version=None, ), runtime=self.runtime, @@ -606,7 +606,7 @@ def _reduce( combine_spm_results( child_results, simulation_params.get("spm"), - expected_year=_parse_year(simulation_params), + expected_year=parse_year(simulation_params), ) ) for key in ("model_version", "data_version"): diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/gcp_credentials.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/gcp_credentials.py new file mode 100644 index 000000000..8d614524e --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/gcp_credentials.py @@ -0,0 +1,119 @@ +"""Write the worker's Google Cloud credentials from its environment.""" + +from __future__ import annotations + +import contextlib +import json +import logging +import os +import tempfile +from typing import TYPE_CHECKING + +from policyengine_simulation_observability import scopes +from policyengine_simulation_observability.stages import ANNUAL_IMPACT_STAGES, Stage + +if TYPE_CHECKING: + from collections.abc import Iterator + + from policyengine_observability import ObservabilityRuntime + +# Moved out of simulation_runtime.py; it still logs under that module's name. +logger = logging.getLogger("policyengine_simulation_executor.simulation_runtime") + +# Credential-adjacent environment values are cut to this many characters. +_ENV_VALUE_PREVIEW_CHARS = 50 + + +def _normalize_credentials_blob(creds_json: str) -> str: + """Return the raw JSON blob, decoding the outer escape if present. + + The upstream Modal secret sometimes stores the credentials payload + double-encoded (the entire JSON object is wrapped in quotes with + backslash-escaped interior quotes). Historically we always attempted + the unescape as a fallback which could accidentally parse an already + clean blob. Only unwrap when the payload looks wrapped. + """ + try: + json.loads(creds_json) + except json.JSONDecodeError: + looks_escaped = creds_json.lstrip().startswith('"') or '\\"' in creds_json + if looks_escaped: + return json.loads(f'"{creds_json}"') + raise + return creds_json + + +@contextlib.contextmanager +def setup_gcp_credentials( + runtime: ObservabilityRuntime | None = None, +) -> Iterator[None]: + """Set up GCP credentials from environment variable. + + Modal secrets are injected as environment variables. The GCP library + expects GOOGLE_APPLICATION_CREDENTIALS to point to a file path. If + credentials JSON is provided, write it to a temp file that's deleted + on exit. This runs as a context manager to guarantee cleanup even if + the caller raises mid-simulation; the previous fire-and-forget + ``tempfile.mkstemp`` path leaked credential material on disk every + time a container served a request. + """ + previous = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") + creds_file = None + with contextlib.ExitStack() as credentials_file_stack: + try: + credential_span = ( + scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.CREDENTIAL_SETUP)) + if runtime is not None + else contextlib.nullcontext() + ) + with credential_span: + # Log available GCP-related env vars for debugging. + gcp_vars = { + k: v[:_ENV_VALUE_PREVIEW_CHARS] + "..." + if len(v) > _ENV_VALUE_PREVIEW_CHARS + else v + for k, v in os.environ.items() + if "GOOGLE" in k or "GCP" in k or "CREDENTIAL" in k + } + logger.info("GCP-related env vars: %s", list(gcp_vars.keys())) + + if previous: + logger.info("GOOGLE_APPLICATION_CREDENTIALS already set") + else: + creds_json = ( + os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON") + or os.environ.get("GCP_CREDENTIALS_JSON") + or os.environ.get("GOOGLE_CREDENTIALS") + or os.environ.get("SERVICE_ACCOUNT_JSON") + ) + + if not creds_json: + logger.warning( + "No GCP credentials found in environment variables" + ) + else: + normalized = _normalize_credentials_blob(creds_json) + # The exit stack keeps the file open across the yield, + # and closing it (``delete=True``) removes the file. + # The finally block first restores any prior value of + # ``GOOGLE_APPLICATION_CREDENTIALS`` so a retry in the same + # container doesn't silently pick up a stale path. + creds_file = credentials_file_stack.enter_context( + tempfile.NamedTemporaryFile( + mode="w", + suffix=".json", + delete=True, + ) + ) + creds_file.write(normalized) + creds_file.flush() + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = creds_file.name + logger.info("GCP credentials written to %s", creds_file.name) + + yield + finally: + if creds_file is not None: + if previous is None: + os.environ.pop("GOOGLE_APPLICATION_CREDENTIALS", None) + else: + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = previous diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index 512ec9498..03602e712 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -21,7 +21,7 @@ means deleting its object first, then re-running. Cohort baselines are built through the executor's own request path -(``_resolve_region`` + ``_build_simulation`` on a synthetic child request), +(``resolve_region`` + ``_build_simulation`` on a synthetic child request), so scoping, extra variables, and the deterministic simulation id match the production child's by construction, not by convention — and a regression test asserts the writer and reader derive identical identities. All diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_planning.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_planning.py index 6f5ceb332..bafb11cd2 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_planning.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_planning.py @@ -53,12 +53,12 @@ def cohort_resolution( """ from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers _country_module, - _resolve_region, + resolve_region, ) params = cohort_params(year, group) country_module = _country_module(PRECOMPUTE_COUNTRY) - resolution = _resolve_region( + resolution = resolve_region( country_module=country_module, country=PRECOMPUTE_COUNTRY, params=params, @@ -123,14 +123,14 @@ class StateDatasetTarget: def state_dataset_targets() -> list[StateDatasetTarget]: """Return every US state region that loads a non-default dataset, by state code. - Each state is resolved exactly as a request for it is (``_resolve_region`` + Each state is resolved exactly as a request for it is (``resolve_region`` then ``_resolve_dataset_selection``), so a bundle that sends states back to the default dataset gets no state files. """ from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers _country_module, _resolve_dataset_selection, - _resolve_region, + resolve_region, ) country_module = _country_module(PRECOMPUTE_COUNTRY) @@ -141,7 +141,7 @@ def state_dataset_targets() -> list[StateDatasetTarget]: "region": region.code, "time_period": PRECOMPUTE_YEARS[0], } - resolution = _resolve_region( + resolution = resolve_region( country_module=country_module, country=PRECOMPUTE_COUNTRY, params=params ) selection = _resolve_dataset_selection(params, region_resolution=resolution) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/region_resolution.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/region_resolution.py new file mode 100644 index 000000000..0d3535edf --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/region_resolution.py @@ -0,0 +1,283 @@ +"""Resolve a simulation request's region, and check a local-area breakdown request.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from policyengine_simulation_contract.dataset_uri import runtime_dataset_uri +from policyengine_simulation_contract.local_area_breakdown import ( + LOCAL_AREAS_BREAKDOWN, + local_area_breakdown_error, +) + +from policyengine_simulation_executor.release_bundle import ( + get_country_release_bundle, + resolve_bundle_dataset_name, +) + +if TYPE_CHECKING: + from types import ModuleType + + from policyengine.core import Region + +# A bare two-letter US region is a state code, such as "ca". +_US_STATE_CODE_LENGTH = 2 + + +@dataclass(frozen=True) +class RegionResolution: + """A request's normalized region code, the dataset to load, and how to scope it. + + ``scoping_strategy`` narrows the simulation to the region; None simulates + the whole dataset, as a national run does. ``state_code`` is the two-letter + code of the US state a state or congressional-district region lies in, and + None for every other region. + """ + + code: str + dataset_reference: str | None = None + scoping_strategy: Any | None = None + state_code: str | None = None + + +def resolve_dataset_reference(country: str) -> str: + """Return the bundle's default dataset name for ``country``.""" + return resolve_bundle_dataset_name(country, None) + + +def _normalise_region_code(country: str, region: object) -> str: + if region is None or str(region).strip() == "": + return country + + raw = str(region).strip() + if raw.lower() in {"us", "uk"}: + return raw.lower() + + if "/" not in raw: + if country == "us" and len(raw) == _US_STATE_CODE_LENGTH: + return f"state/{raw.lower()}" + if country == "uk": + return f"country/{raw.lower().replace(' ', '_')}" + return raw + + prefix, value = raw.split("/", maxsplit=1) + prefix = prefix.lower() + value = value.strip() + if prefix == "state": + value = value.lower() + elif prefix == "country": + value = value.lower().replace(" ", "_") + elif prefix in {"congressional_district", "place", "local_authority"}: + value = value.upper() + return f"{prefix}/{value}" + + +def _uk_area_regions_routed() -> bool: + from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module is not loaded at import; loaded per call + uk_area_regions_routed, + ) + + return uk_area_regions_routed() + + +def _registry_uk_area_region_code(country_module: ModuleType, region_code: str) -> str: + """Spell a UK constituency or local authority as the region registry keys it. + + When the bundle routes UK area regions, policyengine.py's registry holds + every constituency and local authority under its GSS code, labelled with + its name. A request may name the area instead, as policyengine-api's + ``constituency/{name}`` does, or write the code in lower case; both map to + the registry's code. Any other value is returned unchanged, so an area + outside the registry, such as one of the 17 districts replaced in April + 2023, stays unsupported. + """ + from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module is not loaded at import; loaded per call + UK_AREA_REGION_TYPES, + ) + + region_type, _, value = region_code.partition("/") + if region_type not in UK_AREA_REGION_TYPES or not _uk_area_regions_routed(): + return region_code + requested = value.casefold() + for region in country_module.model.region_registry.get_by_type(region_type): + code = region.code.partition("/")[2] + if requested in {code.casefold(), region.label.casefold()}: + return region.code + return region_code + + +def _region_parent_dataset_reference( + country_module: ModuleType, + country: str, + region: Region, +) -> str: + parent_code = getattr(region, "parent_code", None) + visited: set[str] = set() + + while isinstance(parent_code, str) and parent_code: + if parent_code in visited: + msg = f"Region hierarchy cycle at {parent_code}" + raise ValueError(msg) + visited.add(parent_code) + + parent_region = country_module.model.get_region(parent_code) + parent_dataset_path = getattr(parent_region, "dataset_path", None) + if isinstance(parent_dataset_path, str): + bundle = get_country_release_bundle(country) + return runtime_dataset_uri( + parent_dataset_path, + default_revision=bundle.data_package_version, + artifact_revision=bundle.data_artifact_revision, + validate_hf=False, + ) + parent_code = getattr(parent_region, "parent_code", None) + + return resolve_dataset_reference(country) + + +def _reject_unscoped_us_place_region(region_code: str, region: Region) -> None: + if not region_code.startswith("place/"): + return + if getattr(region, "dataset_path", None) is not None: + return + if getattr(region, "scoping_strategy", None) is not None: + return + msg = ( + "US place regions are not yet supported for runtime simulation because " + "policyengine.py does not expose place-level dataset scoping." + ) + raise ValueError(msg) + + +def resolve_region( + *, + country_module: ModuleType, + country: str, + params: dict[str, Any], +) -> RegionResolution: + """Resolve the request's region to its code, scoping and dataset.""" + if params.get("region_group"): + return _resolve_region_group( + country_module=country_module, country=country, params=params + ) + region_code = _normalise_region_code(country, params.get("region")) + if region_code == country: + return RegionResolution( + code=region_code, + dataset_reference=resolve_dataset_reference(country), + ) + + if country == "uk": + region_code = _registry_uk_area_region_code(country_module, region_code) + region = country_module.model.get_region(region_code) + if region is None: + msg = f"Unsupported {country.upper()} region: {region_code}" + raise ValueError(msg) + if country == "us": + _reject_unscoped_us_place_region(region_code, region) + + dataset_path = getattr(region, "dataset_path", None) + if isinstance(dataset_path, str): + bundle = get_country_release_bundle(country) + dataset_reference = runtime_dataset_uri( + dataset_path, + default_revision=bundle.data_package_version, + artifact_revision=bundle.data_artifact_revision, + validate_hf=False, + ) + else: + dataset_reference = _region_parent_dataset_reference( + country_module, + country, + region, + ) + + return RegionResolution( + code=region_code, + dataset_reference=dataset_reference, + scoping_strategy=getattr(region, "scoping_strategy", None), + state_code=getattr(region, "state_code", None) if country == "us" else None, + ) + + +def _resolve_region_group( + *, + country_module: ModuleType, + country: str, + params: dict[str, Any], +) -> RegionResolution: + """Resolve a ``region_group`` list of region codes to one ``RegionGroupStrategy``. + + The strategy scopes a single simulation to the union of the members. + Members must be row-filter regions (states/CDs) so national weights are + preserved (INV-5); the group filters the national dataset. + """ + from policyengine.core.scoping_strategy import ( # noqa: PLC0415 # this module keeps policyengine out of its module-level imports (see its docstring) + RegionGroupStrategy, + RowFilterStrategy, + ) + + codes = [ + _normalise_region_code(country, code) + for code in (params.get("region_group") or []) + ] + if not codes: + msg = "region_group must contain at least one region code" + raise ValueError(msg) + + members = [] + for code in codes: + region = country_module.model.get_region(code) + if region is None: + msg = f"Unsupported {country.upper()} region in group: {code}" + raise ValueError(msg) + scoping_strategy = getattr(region, "scoping_strategy", None) + if not isinstance(scoping_strategy, RowFilterStrategy): + msg = ( + f"Region group member '{code}' is not a row-filter region " + f"(got {type(scoping_strategy).__name__}); region groups must " + "preserve national weights." + ) + raise ValueError(msg) # noqa: TRY004 # an unsupported region in the request; run_simulation_impl handles ValueError specially + members.append(scoping_strategy) + + return RegionResolution( + code="region_group/" + "+".join(sorted(codes)), + dataset_reference=resolve_dataset_reference(country), + scoping_strategy=RegionGroupStrategy(members=members), + ) + + +def require_local_area_breakdown(params: dict[str, Any], *, country: str) -> None: + """Refuse a local-area breakdown this worker cannot run, before loading data. + + The gateway contract applies the same request rules; they are checked + again here because the worker can be called with any payload. The bundle + must also route UK area regions to a local-area dataset, which every + breakdown runs on. + """ + breakdown = params.get("breakdown") + if breakdown != LOCAL_AREAS_BREAKDOWN: + msg = ( + f"Unsupported breakdown {breakdown!r}; the only breakdown is " + f"{LOCAL_AREAS_BREAKDOWN!r}" + ) + raise ValueError(msg) + region = params.get("region") + error = local_area_breakdown_error( + country=country, + region=None if region is None else str(region), + region_group=params.get("region_group"), + include_cliffs=params.get("include_cliffs"), + segmented=params.get("segmented"), + ) + if error is not None: + raise ValueError(error) + if not _uk_area_regions_routed(): + msg = ( + "This policyengine.py bundle does not route UK constituency and " + "local-authority regions to a local-area dataset, so it cannot run " + "a local-area breakdown" + ) + raise ValueError(msg) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/request_policy.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/request_policy.py new file mode 100644 index 000000000..453397bd9 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/request_policy.py @@ -0,0 +1,133 @@ +"""Read a simulation request's year, and the policy its periods describe.""" + +from __future__ import annotations + +from datetime import datetime +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from types import ModuleType + + from policyengine.core import Policy + +DEFAULT_YEAR = 2026 + +# A policy period key is one date or a ``start.stop`` pair of dates. +_DATE_RANGE_PART_COUNT = 2 + + +def parse_year(params: dict[str, Any]) -> int: + """Return the request's simulation year, or the default year without one.""" + value = params.get("time_period") or params.get("year") or DEFAULT_YEAR + return int(value) + + +def _parse_policy_period(period_key: object) -> tuple[datetime, datetime | None]: + """Parse effective dates or inclusive ``start.stop`` policy periods.""" + text = str(period_key) + try: + return _parse_policy_period_dates(text) + except ValueError: + msg = f"Invalid policy period key: {text!r}" + raise ValueError(msg) from None + + +def _parse_policy_period_dates(text: str) -> tuple[datetime, datetime | None]: + """Parse a date or ``start.stop`` range; raise a bare ``ValueError`` if invalid.""" + parts = text.split(".") + if len(parts) not in {1, 2}: + raise ValueError + # The wrapper has always parsed dates with strptime, which also takes + # unpadded months and days ("2027-4-6"); accept what it accepted. + dates = [datetime.strptime(part, "%Y-%m-%d") for part in parts] # noqa: DTZ007 # naive like policyengine.py's own period parsing; its ParameterValue dates are naive + start = dates[0] + stop = dates[1] if len(dates) == _DATE_RANGE_PART_COUNT else None + if stop is not None and stop < start: + raise ValueError + return start, stop + + +def _normalise_period_key(period_key: object) -> str: + """Validate policy period keys without discarding their stop dates.""" + text = str(period_key) + _parse_policy_period(text) + return text + + +def normalise_policy(policy: dict[str, Any] | None) -> dict[str, Any] | None: + """Return the request's policy with every period key normalised, or None.""" + if not policy: + return None + + normalised: dict[str, Any] = {} + for parameter, value in policy.items(): + if isinstance(value, dict): + normalised[parameter] = { + _normalise_period_key(period): period_value + for period, period_value in value.items() + } + else: + normalised[parameter] = value + return normalised + + +def compile_interval_policy( + policy: dict[str, Any] | None, *, country_module: ModuleType, year: int +) -> dict[str, Any] | Policy | None: + """Use public wrapper values for intervals its dict compiler cannot parse.""" + if not policy or not any( + isinstance(value, dict) and any("." in str(key) for key in value) + for value in policy.values() + ): + return policy + + from policyengine.core import ( # noqa: PLC0415 # this module keeps policyengine out of its module-level imports (see its docstring) + ParameterValue, + Policy, + ) + from policyengine.core.parameter import ( # noqa: PLC0415 # this module keeps policyengine out of its module-level imports (see its docstring) + Parameter, + ) + from policyengine.tax_benefit_models.common.reform import ( + compile_reform, + ) + + compiled = compile_reform(policy, year=year, model_version=country_module.model) + if compiled is None: + # compile_reform returns None only for an empty policy, returned above. + return policy + parameter_values = [] + for path, periods in compiled.items(): + dated_values = [ + (*_parse_policy_period(key), value) for key, value in periods.items() + ] + for start, stop, value in sorted(dated_values, key=lambda item: item[0]): + # Match the wrapper's scalar type selection and effective dates. + data_type = type(value) if isinstance(value, (int, float, bool)) else float + parameter_values.append( + ParameterValue( + parameter=Parameter( + name=path, + tax_benefit_model_version=country_module.model, + data_type=data_type, + ), + start_date=start, + end_date=stop, + value=value, + ) + ) + return Policy(name="Reform", parameter_values=parameter_values) + + +def wrapper_policy( + policy: dict[str, Any] | None, *, country_module: ModuleType, year: int +) -> dict[str, Any] | Policy | None: + """Convert a request's policy into the form the wrapper must receive. + + Normalises the period keys, then compiles a policy with interval keys + into wrapper parameter values (see ``_compile_interval_policy``). Every + path that hands a request policy to the wrapper goes through this. + """ + return compile_interval_policy( + normalise_policy(policy), country_module=country_module, year=year + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py index 470eb271a..5088f1066 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py @@ -7,40 +7,40 @@ from __future__ import annotations -import contextlib import hashlib -import json import logging import os -import tempfile import time from dataclasses import dataclass -from datetime import datetime from importlib import import_module from typing import TYPE_CHECKING, Any from policyengine_simulation_contract.dataset_uri import runtime_dataset_uri from policyengine_simulation_contract.job_progress import JobStage -from policyengine_simulation_contract.local_area_breakdown import ( - LOCAL_AREAS_BREAKDOWN, - local_area_breakdown_error, -) from policyengine_simulation_contract.spm import SPMInputError, spm_error_detail from policyengine_simulation_observability import scopes -from policyengine_simulation_observability.stages import ( - ANNUAL_IMPACT_STAGES, - Stage, -) +from policyengine_simulation_observability.stages import ANNUAL_IMPACT_STAGES, Stage from policyengine_simulation_observability.telemetry import split_internal_payload from policyengine_simulation_executor import state_datasets +from policyengine_simulation_executor.gcp_credentials import setup_gcp_credentials from policyengine_simulation_executor.job_progress import ( NULL_JOB_PROGRESS, JobProgressReporter, ) -from policyengine_simulation_executor.release_bundle import ( - get_country_release_bundle, - resolve_bundle_dataset_name, +from policyengine_simulation_executor.region_resolution import ( + RegionResolution, + require_local_area_breakdown, + resolve_dataset_reference, + resolve_region, +) +from policyengine_simulation_executor.release_bundle import get_country_release_bundle +from policyengine_simulation_executor.request_policy import ( + DEFAULT_YEAR, + compile_interval_policy, + normalise_policy, + parse_year, + wrapper_policy, ) from policyengine_simulation_executor.simulation_output_builder import ( SimulationOutputBuilder, @@ -50,48 +50,38 @@ ) if TYPE_CHECKING: - from collections.abc import Iterator from types import ModuleType - from policyengine.core import Dataset, Policy, Region + from policyengine.core import Dataset from policyengine_observability import ObservabilityRuntime - from policyengine_simulation_contract.uk_geography import ( - UKLocalAuthorityMetadata, - ) + from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata from policyengine_simulation_observability.telemetry import TelemetryEnvelope from policyengine_simulation_executor.stage_reporting_simulation import ( StageReportingSimulation, ) +__all__ = [ + "DEFAULT_YEAR", + "DatasetSelection", + "RegionResolution", + "UKAreaSimulations", + "compile_interval_policy", + "logger", + "normalise_policy", + "parse_year", + "resolve_data_folder", + "resolve_dataset_reference", + "resolve_region", + "run_simulation_impl", + "setup_gcp_credentials", + "wrapper_policy", +] + logger = logging.getLogger(__name__) os.environ.setdefault("POLICYENGINE_SKIP_COUNTRY_IMPORTS", "1") -DEFAULT_YEAR = 2026 -# Credential-adjacent environment values are cut to this many characters. -_ENV_VALUE_PREVIEW_CHARS = 50 -# A policy period key is one date or a ``start.stop`` pair of dates. -_DATE_RANGE_PART_COUNT = 2 -# A bare two-letter US region is a state code, such as "ca". -_US_STATE_CODE_LENGTH = 2 - - -@dataclass(frozen=True) -class RegionResolution: - """A request's normalized region code, the dataset to load, and how to scope it. - - ``scoping_strategy`` narrows the simulation to the region; None simulates - the whole dataset, as a national run does. ``state_code`` is the two-letter - code of the US state a state or congressional-district region lies in, and - None for every other region. - """ - - code: str - dataset_reference: str | None = None - scoping_strategy: Any | None = None - state_code: str | None = None - @dataclass(frozen=True) class DatasetSelection: @@ -119,101 +109,6 @@ class UKAreaSimulations: dataset_name: str -def _normalize_credentials_blob(creds_json: str) -> str: - """Return the raw JSON blob, decoding the outer escape if present. - - The upstream Modal secret sometimes stores the credentials payload - double-encoded (the entire JSON object is wrapped in quotes with - backslash-escaped interior quotes). Historically we always attempted - the unescape as a fallback which could accidentally parse an already - clean blob. Only unwrap when the payload looks wrapped. - """ - try: - json.loads(creds_json) - except json.JSONDecodeError: - looks_escaped = creds_json.lstrip().startswith('"') or '\\"' in creds_json - if looks_escaped: - return json.loads(f'"{creds_json}"') - raise - return creds_json - - -@contextlib.contextmanager -def setup_gcp_credentials( - runtime: ObservabilityRuntime | None = None, -) -> Iterator[None]: - """Set up GCP credentials from environment variable. - - Modal secrets are injected as environment variables. The GCP library - expects GOOGLE_APPLICATION_CREDENTIALS to point to a file path. If - credentials JSON is provided, write it to a temp file that's deleted - on exit. This runs as a context manager to guarantee cleanup even if - the caller raises mid-simulation; the previous fire-and-forget - ``tempfile.mkstemp`` path leaked credential material on disk every - time a container served a request. - """ - previous = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") - creds_file = None - with contextlib.ExitStack() as credentials_file_stack: - try: - credential_span = ( - scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.CREDENTIAL_SETUP)) - if runtime is not None - else contextlib.nullcontext() - ) - with credential_span: - # Log available GCP-related env vars for debugging. - gcp_vars = { - k: v[:_ENV_VALUE_PREVIEW_CHARS] + "..." - if len(v) > _ENV_VALUE_PREVIEW_CHARS - else v - for k, v in os.environ.items() - if "GOOGLE" in k or "GCP" in k or "CREDENTIAL" in k - } - logger.info("GCP-related env vars: %s", list(gcp_vars.keys())) - - if previous: - logger.info("GOOGLE_APPLICATION_CREDENTIALS already set") - else: - creds_json = ( - os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON") - or os.environ.get("GCP_CREDENTIALS_JSON") - or os.environ.get("GOOGLE_CREDENTIALS") - or os.environ.get("SERVICE_ACCOUNT_JSON") - ) - - if not creds_json: - logger.warning( - "No GCP credentials found in environment variables" - ) - else: - normalized = _normalize_credentials_blob(creds_json) - # The exit stack keeps the file open across the yield, - # and closing it (``delete=True``) removes the file. - # The finally block first restores any prior value of - # ``GOOGLE_APPLICATION_CREDENTIALS`` so a retry in the same - # container doesn't silently pick up a stale path. - creds_file = credentials_file_stack.enter_context( - tempfile.NamedTemporaryFile( - mode="w", - suffix=".json", - delete=True, - ) - ) - creds_file.write(normalized) - creds_file.flush() - os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = creds_file.name - logger.info("GCP credentials written to %s", creds_file.name) - - yield - finally: - if creds_file is not None: - if previous is None: - os.environ.pop("GOOGLE_APPLICATION_CREDENTIALS", None) - else: - os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = previous - - def run_simulation_impl( params: dict, *, @@ -239,327 +134,6 @@ def run_simulation_impl( raise -def _parse_year(params: dict[str, Any]) -> int: - value = params.get("time_period") or params.get("year") or DEFAULT_YEAR - return int(value) - - -def _parse_policy_period(period_key: object) -> tuple[datetime, datetime | None]: - """Parse effective dates or inclusive ``start.stop`` policy periods.""" - text = str(period_key) - try: - return _parse_policy_period_dates(text) - except ValueError: - msg = f"Invalid policy period key: {text!r}" - raise ValueError(msg) from None - - -def _parse_policy_period_dates(text: str) -> tuple[datetime, datetime | None]: - """Parse a date or ``start.stop`` range; raise a bare ``ValueError`` if invalid.""" - parts = text.split(".") - if len(parts) not in {1, 2}: - raise ValueError - # The wrapper has always parsed dates with strptime, which also takes - # unpadded months and days ("2027-4-6"); accept what it accepted. - dates = [datetime.strptime(part, "%Y-%m-%d") for part in parts] # noqa: DTZ007 # naive like policyengine.py's own period parsing; its ParameterValue dates are naive - start = dates[0] - stop = dates[1] if len(dates) == _DATE_RANGE_PART_COUNT else None - if stop is not None and stop < start: - raise ValueError - return start, stop - - -def _normalise_period_key(period_key: object) -> str: - """Validate policy period keys without discarding their stop dates.""" - text = str(period_key) - _parse_policy_period(text) - return text - - -def _normalise_policy(policy: dict[str, Any] | None) -> dict[str, Any] | None: - if not policy: - return None - - normalised: dict[str, Any] = {} - for parameter, value in policy.items(): - if isinstance(value, dict): - normalised[parameter] = { - _normalise_period_key(period): period_value - for period, period_value in value.items() - } - else: - normalised[parameter] = value - return normalised - - -def _compile_interval_policy( - policy: dict[str, Any] | None, *, country_module: ModuleType, year: int -) -> dict[str, Any] | Policy | None: - """Use public wrapper values for intervals its dict compiler cannot parse.""" - if not policy or not any( - isinstance(value, dict) and any("." in str(key) for key in value) - for value in policy.values() - ): - return policy - - from policyengine.core import ( # noqa: PLC0415 # this module keeps policyengine out of its module-level imports (see its docstring) - ParameterValue, - Policy, - ) - from policyengine.core.parameter import ( # noqa: PLC0415 # this module keeps policyengine out of its module-level imports (see its docstring) - Parameter, - ) - from policyengine.tax_benefit_models.common.reform import ( - compile_reform, - ) - - compiled = compile_reform(policy, year=year, model_version=country_module.model) - if compiled is None: - # compile_reform returns None only for an empty policy, returned above. - return policy - parameter_values = [] - for path, periods in compiled.items(): - dated_values = [ - (*_parse_policy_period(key), value) for key, value in periods.items() - ] - for start, stop, value in sorted(dated_values, key=lambda item: item[0]): - # Match the wrapper's scalar type selection and effective dates. - data_type = type(value) if isinstance(value, (int, float, bool)) else float - parameter_values.append( - ParameterValue( - parameter=Parameter( - name=path, - tax_benefit_model_version=country_module.model, - data_type=data_type, - ), - start_date=start, - end_date=stop, - value=value, - ) - ) - return Policy(name="Reform", parameter_values=parameter_values) - - -def _runtime_policy( - policy: dict[str, Any] | None, *, country_module: ModuleType, year: int -) -> dict[str, Any] | Policy | None: - """Convert a request's policy into the form the wrapper must receive. - - Normalises the period keys, then compiles a policy with interval keys - into wrapper parameter values (see ``_compile_interval_policy``). Every - path that hands a request policy to the wrapper goes through this. - """ - return _compile_interval_policy( - _normalise_policy(policy), country_module=country_module, year=year - ) - - -def _resolve_dataset_reference(country: str) -> str: - return resolve_bundle_dataset_name(country, None) - - -def _normalise_region_code(country: str, region: object) -> str: - if region is None or str(region).strip() == "": - return country - - raw = str(region).strip() - if raw.lower() in {"us", "uk"}: - return raw.lower() - - if "/" not in raw: - if country == "us" and len(raw) == _US_STATE_CODE_LENGTH: - return f"state/{raw.lower()}" - if country == "uk": - return f"country/{raw.lower().replace(' ', '_')}" - return raw - - prefix, value = raw.split("/", maxsplit=1) - prefix = prefix.lower() - value = value.strip() - if prefix == "state": - value = value.lower() - elif prefix == "country": - value = value.lower().replace(" ", "_") - elif prefix in {"congressional_district", "place", "local_authority"}: - value = value.upper() - return f"{prefix}/{value}" - - -def _uk_area_regions_routed() -> bool: - from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module is not loaded at import; loaded per call - uk_area_regions_routed, - ) - - return uk_area_regions_routed() - - -def _registry_uk_area_region_code(country_module: ModuleType, region_code: str) -> str: - """Spell a UK constituency or local authority as the region registry keys it. - - When the bundle routes UK area regions, policyengine.py's registry holds - every constituency and local authority under its GSS code, labelled with - its name. A request may name the area instead, as policyengine-api's - ``constituency/{name}`` does, or write the code in lower case; both map to - the registry's code. Any other value is returned unchanged, so an area - outside the registry, such as one of the 17 districts replaced in April - 2023, stays unsupported. - """ - from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module is not loaded at import; loaded per call - UK_AREA_REGION_TYPES, - ) - - region_type, _, value = region_code.partition("/") - if region_type not in UK_AREA_REGION_TYPES or not _uk_area_regions_routed(): - return region_code - requested = value.casefold() - for region in country_module.model.region_registry.get_by_type(region_type): - code = region.code.partition("/")[2] - if requested in {code.casefold(), region.label.casefold()}: - return region.code - return region_code - - -def _region_parent_dataset_reference( - country_module: ModuleType, - country: str, - region: Region, -) -> str: - parent_code = getattr(region, "parent_code", None) - visited: set[str] = set() - - while isinstance(parent_code, str) and parent_code: - if parent_code in visited: - msg = f"Region hierarchy cycle at {parent_code}" - raise ValueError(msg) - visited.add(parent_code) - - parent_region = country_module.model.get_region(parent_code) - parent_dataset_path = getattr(parent_region, "dataset_path", None) - if isinstance(parent_dataset_path, str): - bundle = get_country_release_bundle(country) - return runtime_dataset_uri( - parent_dataset_path, - default_revision=bundle.data_package_version, - artifact_revision=bundle.data_artifact_revision, - validate_hf=False, - ) - parent_code = getattr(parent_region, "parent_code", None) - - return _resolve_dataset_reference(country) - - -def _reject_unscoped_us_place_region(region_code: str, region: Region) -> None: - if not region_code.startswith("place/"): - return - if getattr(region, "dataset_path", None) is not None: - return - if getattr(region, "scoping_strategy", None) is not None: - return - msg = ( - "US place regions are not yet supported for runtime simulation because " - "policyengine.py does not expose place-level dataset scoping." - ) - raise ValueError(msg) - - -def _resolve_region( - *, - country_module: ModuleType, - country: str, - params: dict[str, Any], -) -> RegionResolution: - if params.get("region_group"): - return _resolve_region_group( - country_module=country_module, country=country, params=params - ) - region_code = _normalise_region_code(country, params.get("region")) - if region_code == country: - return RegionResolution( - code=region_code, - dataset_reference=_resolve_dataset_reference(country), - ) - - if country == "uk": - region_code = _registry_uk_area_region_code(country_module, region_code) - region = country_module.model.get_region(region_code) - if region is None: - msg = f"Unsupported {country.upper()} region: {region_code}" - raise ValueError(msg) - if country == "us": - _reject_unscoped_us_place_region(region_code, region) - - dataset_path = getattr(region, "dataset_path", None) - if isinstance(dataset_path, str): - bundle = get_country_release_bundle(country) - dataset_reference = runtime_dataset_uri( - dataset_path, - default_revision=bundle.data_package_version, - artifact_revision=bundle.data_artifact_revision, - validate_hf=False, - ) - else: - dataset_reference = _region_parent_dataset_reference( - country_module, - country, - region, - ) - - return RegionResolution( - code=region_code, - dataset_reference=dataset_reference, - scoping_strategy=getattr(region, "scoping_strategy", None), - state_code=getattr(region, "state_code", None) if country == "us" else None, - ) - - -def _resolve_region_group( - *, - country_module: ModuleType, - country: str, - params: dict[str, Any], -) -> RegionResolution: - """Resolve a ``region_group`` list of region codes to one ``RegionGroupStrategy``. - - The strategy scopes a single simulation to the union of the members. - Members must be row-filter regions (states/CDs) so national weights are - preserved (INV-5); the group filters the national dataset. - """ - from policyengine.core.scoping_strategy import ( # noqa: PLC0415 # this module keeps policyengine out of its module-level imports (see its docstring) - RegionGroupStrategy, - RowFilterStrategy, - ) - - codes = [ - _normalise_region_code(country, code) - for code in (params.get("region_group") or []) - ] - if not codes: - msg = "region_group must contain at least one region code" - raise ValueError(msg) - - members = [] - for code in codes: - region = country_module.model.get_region(code) - if region is None: - msg = f"Unsupported {country.upper()} region in group: {code}" - raise ValueError(msg) - scoping_strategy = getattr(region, "scoping_strategy", None) - if not isinstance(scoping_strategy, RowFilterStrategy): - msg = ( - f"Region group member '{code}' is not a row-filter region " - f"(got {type(scoping_strategy).__name__}); region groups must " - "preserve national weights." - ) - raise ValueError(msg) # noqa: TRY004 # an unsupported region in the request; run_simulation_impl handles ValueError specially - members.append(scoping_strategy) - - return RegionResolution( - code="region_group/" + "+".join(sorted(codes)), - dataset_reference=_resolve_dataset_reference(country), - scoping_strategy=RegionGroupStrategy(members=members), - ) - - def _country_module(country: str) -> ModuleType: country = country.lower() if country not in {"us", "uk"}: @@ -661,7 +235,7 @@ def _load_dataset( other dataset. """ country = params.get("country", "us").lower() - year = _parse_year(params) + year = parse_year(params) country_module = country_module or _country_module(country) if data_folder is None: data_folder = ( @@ -715,8 +289,8 @@ def _build_simulation( spm_kwargs: dict[str, Any] = {"spm": selection} if selection is not None else {} country = params.get("country", "us") country_module = _country_module(country) - runtime_policy = _compile_interval_policy( - policy, country_module=country_module, year=_parse_year(params) + runtime_policy = compile_interval_policy( + policy, country_module=country_module, year=parse_year(params) ) scoping_strategy = region.scoping_strategy simulation_id = deterministic_baseline_id( @@ -727,7 +301,7 @@ def _build_simulation( policy=runtime_policy, region_code=region.code, scoping_strategy=scoping_strategy, - year=_parse_year(params), + year=parse_year(params), ) ) if simulation_id is not None: @@ -825,40 +399,6 @@ def _uk_area_simulations( ) -def _require_local_area_breakdown(params: dict[str, Any], *, country: str) -> None: - """Refuse a local-area breakdown this worker cannot run, before loading data. - - The gateway contract applies the same request rules; they are checked - again here because the worker can be called with any payload. The bundle - must also route UK area regions to a local-area dataset, which every - breakdown runs on. - """ - breakdown = params.get("breakdown") - if breakdown != LOCAL_AREAS_BREAKDOWN: - msg = ( - f"Unsupported breakdown {breakdown!r}; the only breakdown is " - f"{LOCAL_AREAS_BREAKDOWN!r}" - ) - raise ValueError(msg) - region = params.get("region") - error = local_area_breakdown_error( - country=country, - region=None if region is None else str(region), - region_group=params.get("region_group"), - include_cliffs=params.get("include_cliffs"), - segmented=params.get("segmented"), - ) - if error is not None: - raise ValueError(error) - if not _uk_area_regions_routed(): - msg = ( - "This policyengine.py bundle does not route UK constituency and " - "local-authority regions to a local-area dataset, so it cannot run " - "a local-area breakdown" - ) - raise ValueError(msg) - - def _run_local_area_breakdown_core( simulation_params: dict[str, Any], *, @@ -873,11 +413,11 @@ def _run_local_area_breakdown_core( never loaded. The result holds the constituency and local-authority impacts only (``LocalAreaBreakdownOutput``). """ - _require_local_area_breakdown(simulation_params, country=country) + require_local_area_breakdown(simulation_params, country=country) with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.COUNTRY_MODULE_LOAD)): country_module = _country_module(country) with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.REGION_RESOLUTION)): - region_resolution = _resolve_region( + region_resolution = resolve_region( country_module=country_module, country=country, params=simulation_params, @@ -885,8 +425,8 @@ def _run_local_area_breakdown_core( runtime.set_context(region=region_resolution.code) with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.POLICY_NORMALIZATION)): policies = ( - _normalise_policy(simulation_params.get("baseline")), - _normalise_policy(simulation_params.get("reform")), + normalise_policy(simulation_params.get("baseline")), + normalise_policy(simulation_params.get("reform")), ) logger.info("Initialising the local-area baseline and reform simulations") @@ -907,7 +447,7 @@ def _run_local_area_breakdown_core( local_authority_metadata=area.local_authority_metadata, dataset_name=area.dataset_name, data_version=_manifest_dataset_revision(country, area.dataset_name), - year=_parse_year(simulation_params), + year=parse_year(simulation_params), country_module=country_module, runtime=runtime, progress=progress, @@ -953,7 +493,7 @@ def _run_simulation_impl_core( with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.COUNTRY_MODULE_LOAD)): country_module = _country_module(country) with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.REGION_RESOLUTION)): - region_resolution = _resolve_region( + region_resolution = resolve_region( country_module=country_module, country=country, params=simulation_params, @@ -968,7 +508,7 @@ def _run_simulation_impl_core( country=country, selection=dataset_selection, state_code=region_resolution.state_code, - year=_parse_year(simulation_params), + year=parse_year(simulation_params), data_folder=resolve_data_folder(), ) if not dataset_selection.is_default: @@ -997,8 +537,8 @@ def _run_simulation_impl_core( dataset_name=dataset_selection.name, ) with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.POLICY_NORMALIZATION)): - baseline_policy = _normalise_policy(simulation_params.get("baseline")) - reform_policy = _normalise_policy(simulation_params.get("reform")) + baseline_policy = normalise_policy(simulation_params.get("baseline")) + reform_policy = normalise_policy(simulation_params.get("reform")) logger.info("Initialising baseline and reform simulations") with scopes.span( @@ -1051,7 +591,7 @@ def _run_simulation_impl_core( output = builder.serialize() output.update( simulation_spm_result( - baseline, reform, selection, expected_year=_parse_year(simulation_params) + baseline, reform, selection, expected_year=parse_year(simulation_params) ) ) # ensure() has run inside the builder by now, so the artifact outcome @@ -1095,7 +635,7 @@ def _set_runtime_attributes( runtime.set_context( country=country, scope=simulation_params.get("scope"), - simulation_year=_parse_year(simulation_params), + simulation_year=parse_year(simulation_params), submission_claim_id=getattr(telemetry, "submission_claim_id", None), geography_code=getattr(telemetry, "geography_code", None), geography_type=getattr(telemetry, "geography_type", None), 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 b4b4ade8e..a1f39c2d8 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py @@ -212,11 +212,11 @@ def normalize_runtime_spm(params: dict[str, Any]) -> dict[str, Any] | None: canonical = json.dumps(selection, sort_keys=True, separators=(",", ":")) try: from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # simulation_runtime imports this module lazily; kept lazy on both sides - _parse_year, + parse_year, ) try: - start = int(params.get("start_year") or _parse_year(params)) + start = int(params.get("start_year") or parse_year(params)) window = int(params.get("window_size", 1)) except ValueError: # Keep the order the range's old parse position gave it: the diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute_planning.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute_planning.py index f67362226..5c73a18ac 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute_planning.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute_planning.py @@ -24,7 +24,7 @@ from policyengine_simulation_executor.simulation_runtime import ( _country_module, _resolve_dataset_selection, - _resolve_region, + resolve_region, ) from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle from policyengine_simulation_executor.stage12_runtime.output_planning import ( @@ -150,7 +150,7 @@ def _scope_key(year: int, group: tuple[str, ...]) -> str: "time_period": str(year), "region_group": list(group), } - resolution = _resolve_region( + resolution = resolve_region( country_module=_country_module("us"), country="us", params=params, @@ -222,7 +222,7 @@ def state_dataset_targets() -> tuple[StateDatasetTarget, ...]: """Return every US state region that loads a non-default dataset, by state code. Each state is resolved as a Stage 12 worker resolves a request for it - (``_resolve_region`` then ``_resolve_dataset_selection``), so a bundle + (``resolve_region`` then ``_resolve_dataset_selection``), so a bundle that sends states back to the default dataset gets no state files. """ country_module = _country_module("us") @@ -233,7 +233,7 @@ def state_dataset_targets() -> tuple[StateDatasetTarget, ...]: "region": region.code, "time_period": PRECOMPUTE_YEARS[0], } - resolution = _resolve_region( + resolution = resolve_region( country_module=country_module, country="us", params=params ) selection = _resolve_dataset_selection(params, region_resolution=resolution) 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 657414357..24bde0c91 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 @@ -182,7 +182,7 @@ def build_aggregate_report( ) from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime._country_module _country_module, - _runtime_policy, + wrapper_policy, ) from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module is not loaded at import; loaded per call uk_area_breakdowns_omitted, @@ -212,7 +212,7 @@ def stand_in( tax_benefit_model_version=country_module.model, # Interval keys need compiling: the wrapper's dict compiler # rejects them. - policy=_runtime_policy( + policy=wrapper_policy( policy, country_module=country_module, year=dataset.year ), ) 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 2409ed33c..1975fc8e0 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 @@ -243,13 +243,13 @@ def _require_loaded_dataset_matches_provenance( send that until they are updated. """ from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers - _resolve_dataset_reference, + resolve_dataset_reference, ) named = simulation.bundle.dataset if (selection.name, selection.uri) == (named.identity, named.uri): return - if not selection.is_default and named.identity == _resolve_dataset_reference( + if not selection.is_default and named.identity == resolve_dataset_reference( simulation.geography.country ): return @@ -285,10 +285,10 @@ def calculate_simulation_frames( from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers _country_module, _load_dataset, - _normalise_policy, _resolve_dataset_selection, - _resolve_region, + normalise_policy, resolve_data_folder, + resolve_region, setup_gcp_credentials, ) from policyengine_simulation_executor.stage12_cache import ( # noqa: PLC0415 # resolved per call so tests can patch the baked-folder lookup @@ -302,7 +302,7 @@ def calculate_simulation_frames( with stage_span(Stage.COUNTRY_MODULE_LOAD): country_module = _country_module(country) with stage_span(Stage.REGION_RESOLUTION): - region = _resolve_region( + region = resolve_region( country_module=country_module, country=country, params=params, @@ -333,7 +333,7 @@ def calculate_simulation_frames( dataset_name=dataset_selection.name, ) with stage_span(Stage.POLICY_NORMALIZATION): - policy = _normalise_policy(simulation.policy) + policy = normalise_policy(simulation.policy) with stage_span(Stage.SIMULATION_BUILD): model = build_stage12_simulation( params, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py index 69371b11f..9ff959eec 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py @@ -14,7 +14,7 @@ resolve_local_bundle_dataset_path, ) from policyengine_simulation_executor.simulation_runtime import ( - _normalise_policy, + normalise_policy, resolve_data_folder, ) from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle @@ -96,7 +96,7 @@ def _run_non_serving_calculation(country: CountryId) -> None: ] extra_variables.append("ctc") options["spm"] = {"geography_kind": "national"} - options["reform"] = _normalise_policy(US_VALIDATION_REFORM) + options["reform"] = normalise_policy(US_VALIDATION_REFORM) result = country_module.calculate_household( people=people, year=2026, diff --git a/projects/policyengine-simulation-executor/tests/stage12_runtime_support.py b/projects/policyengine-simulation-executor/tests/stage12_runtime_support.py index d51d2a64e..81fc7ff3b 100644 --- a/projects/policyengine-simulation-executor/tests/stage12_runtime_support.py +++ b/projects/policyengine-simulation-executor/tests/stage12_runtime_support.py @@ -312,7 +312,7 @@ def stop_after_loading(*_args, **_kwargs): monkeypatch.setattr( simulation_runtime, "_country_module", lambda _country: object() ) - monkeypatch.setattr(simulation_runtime, "_resolve_region", lambda **_: region) + monkeypatch.setattr(simulation_runtime, "resolve_region", lambda **_: region) monkeypatch.setattr( simulation_runtime, "_resolve_dataset_selection", diff --git a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py index a6c8bdda7..402cb199b 100644 --- a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py @@ -91,7 +91,7 @@ def test_reforms_keep_current_law_artifact_identity_separate(self, collected, ki path = "gov.hmrc.income_tax.rates.uk[0].rate" policy = {path: {"2027-04-06.2028-04-05": 0.21}} if kind == "bounded_policy": - policy = sr._compile_interval_policy( + policy = sr.compile_interval_policy( policy, country_module=sr._country_module("uk"), year=2027 ) elif kind == "scalar": diff --git a/projects/policyengine-simulation-executor/tests/test_budget_reform_contract_route.py b/projects/policyengine-simulation-executor/tests/test_budget_reform_contract_route.py index 9b3a1005d..cadabfeba 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_reform_contract_route.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_reform_contract_route.py @@ -23,9 +23,9 @@ reform_dict_from_parameter_values, ) from policyengine_simulation_executor.simulation_runtime import ( # noqa: E402 - _compile_interval_policy, _country_module, - _normalise_policy, + compile_interval_policy, + normalise_policy, ) CASES = [ @@ -44,14 +44,14 @@ def test_the_api_hands_the_wrapper_each_case_unchanged(case): wrapper_form = contract.represent(case, "wrapper.policy") if case.expect_error is not None: with pytest.raises(ValueError, match=case.expect_error): - _compile_interval_policy( - _normalise_policy(request), country_module=uk, year=year + compile_interval_policy( + normalise_policy(request), country_module=uk, year=year ) with pytest.raises(ValueError, match=case.expect_error): compile_reform_to_policy(wrapper_form, year=year, model_version=uk.model) return - through_api = _compile_interval_policy( - _normalise_policy(request), country_module=uk, year=year + through_api = compile_interval_policy( + normalise_policy(request), country_module=uk, year=year ) through_wrapper = compile_reform_to_policy( wrapper_form, year=year, model_version=uk.model diff --git a/projects/policyengine-simulation-executor/tests/test_gcp_credentials.py b/projects/policyengine-simulation-executor/tests/test_gcp_credentials.py index 51ec20a36..1b143f04e 100644 --- a/projects/policyengine-simulation-executor/tests/test_gcp_credentials.py +++ b/projects/policyengine-simulation-executor/tests/test_gcp_credentials.py @@ -7,10 +7,10 @@ from pathlib import Path import pytest -from policyengine_simulation_executor.simulation_runtime import ( +from policyengine_simulation_executor.gcp_credentials import ( _normalize_credentials_blob, - setup_gcp_credentials, ) +from policyengine_simulation_executor.simulation_runtime import setup_gcp_credentials def test_setup_gcp_credentials_deletes_temp_file_on_exit( diff --git a/projects/policyengine-simulation-executor/tests/test_job_progress.py b/projects/policyengine-simulation-executor/tests/test_job_progress.py index 2a172362e..4153bd85e 100644 --- a/projects/policyengine-simulation-executor/tests/test_job_progress.py +++ b/projects/policyengine-simulation-executor/tests/test_job_progress.py @@ -593,7 +593,7 @@ def serialize(self): ) monkeypatch.setattr( sr, - "_resolve_region", + "resolve_region", lambda **_kwargs: RegionResolution(code="us", dataset_reference=None), ) monkeypatch.setattr( diff --git a/projects/policyengine-simulation-executor/tests/test_precompute_baselines.py b/projects/policyengine-simulation-executor/tests/test_precompute_baselines.py index a5f3c393b..0e6be4ff1 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute_baselines.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute_baselines.py @@ -98,7 +98,7 @@ def upload_file(self, path, local_path): monkeypatch.setattr(sr, "_country_module", lambda _country: SimpleNamespace()) monkeypatch.setattr( sr, - "_resolve_region", + "resolve_region", lambda **_kwargs: SimpleNamespace( code="state/ca+state/wv", scoping_strategy="scoping", diff --git a/projects/policyengine-simulation-executor/tests/test_precompute_planning.py b/projects/policyengine-simulation-executor/tests/test_precompute_planning.py index f95629a8b..f4c6c3b5a 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute_planning.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute_planning.py @@ -13,12 +13,19 @@ class TestWriterReaderContract: @pytest.fixture def identity_stubs(self, monkeypatch): + from policyengine_simulation_executor import region_resolution from policyengine_simulation_executor import simulation_runtime as sr state = install_identity_stubs(monkeypatch) + # Both modules read the release bundle: dataset selection and regions. monkeypatch.setattr( sr, "get_country_release_bundle", lambda _country: state.bundle ) + monkeypatch.setattr( + region_resolution, + "get_country_release_bundle", + lambda _country: state.bundle, + ) @pytest.fixture def fake_regions(self, monkeypatch): @@ -54,7 +61,7 @@ def test_precompute_identity_equals_runtime_id(self): writer_identity = precompute.cohort_identity(2026, group) params = precompute.cohort_params(2026, group) - resolution = sr._resolve_region( + resolution = sr.resolve_region( country_module=sr._country_module("us"), country="us", params=params ) dataset_selection = sr._resolve_dataset_selection( @@ -416,7 +423,7 @@ def test_diverged_predicates_fail_loudly(self, monkeypatch): monkeypatch.setattr(sr, "_country_module", lambda _country: SimpleNamespace()) monkeypatch.setattr( sr, - "_resolve_region", + "resolve_region", lambda **_kwargs: SimpleNamespace( code="x", dataset_reference=None, scoping_strategy=None ), diff --git a/projects/policyengine-simulation-executor/tests/test_region_group_resolution.py b/projects/policyengine-simulation-executor/tests/test_region_group_resolution.py index 9e8c022ce..3c366b068 100644 --- a/projects/policyengine-simulation-executor/tests/test_region_group_resolution.py +++ b/projects/policyengine-simulation-executor/tests/test_region_group_resolution.py @@ -8,6 +8,7 @@ RowFilterStrategy, WeightReplacementStrategy, ) +from policyengine_simulation_executor import region_resolution from policyengine_simulation_executor import simulation_runtime as sr from policyengine_simulation_executor.simulation_output_geographic import ( _should_build_us_congressional_district_impact as gate, @@ -33,12 +34,14 @@ def _state(fips): class TestResolveRegionGroup: def test__builds_region_group_strategy_from_members(self, monkeypatch): monkeypatch.setattr( - sr, "_resolve_dataset_reference", lambda _country: "us-national" + region_resolution, + "resolve_dataset_reference", + lambda _country: "us-national", ) cm = _country_module( {"state/hi": _region_with(_state(15)), "state/ia": _region_with(_state(19))} ) - res = sr._resolve_region_group( + res = region_resolution._resolve_region_group( country_module=cm, country="us", params={"region_group": ["state/hi", "state/ia"]}, @@ -49,16 +52,18 @@ def test__builds_region_group_strategy_from_members(self, monkeypatch): assert res.dataset_reference == "us-national" def test__code_is_member_order_independent(self, monkeypatch): - monkeypatch.setattr(sr, "_resolve_dataset_reference", lambda _country: "x") + monkeypatch.setattr( + region_resolution, "resolve_dataset_reference", lambda _country: "x" + ) cm = _country_module( {"state/hi": _region_with(_state(15)), "state/ia": _region_with(_state(19))} ) - a = sr._resolve_region_group( + a = region_resolution._resolve_region_group( country_module=cm, country="us", params={"region_group": ["state/hi", "state/ia"]}, ) - b = sr._resolve_region_group( + b = region_resolution._resolve_region_group( country_module=cm, country="us", params={"region_group": ["state/ia", "state/hi"]}, @@ -66,15 +71,19 @@ def test__code_is_member_order_independent(self, monkeypatch): assert a.code == b.code def test__rejects_unknown_member(self, monkeypatch): - monkeypatch.setattr(sr, "_resolve_dataset_reference", lambda _country: "x") + monkeypatch.setattr( + region_resolution, "resolve_dataset_reference", lambda _country: "x" + ) cm = _country_module({}) # get_region returns None with pytest.raises(ValueError, match="Unsupported"): - sr._resolve_region_group( + region_resolution._resolve_region_group( country_module=cm, country="us", params={"region_group": ["state/zz"]} ) def test__rejects_weight_replacement_member(self, monkeypatch): - monkeypatch.setattr(sr, "_resolve_dataset_reference", lambda _country: "x") + monkeypatch.setattr( + region_resolution, "resolve_dataset_reference", lambda _country: "x" + ) weight_replacement = WeightReplacementStrategy( weight_matrix_bucket="b", weight_matrix_key="k", @@ -85,22 +94,26 @@ def test__rejects_weight_replacement_member(self, monkeypatch): cm = MagicMock() cm.model.get_region.return_value = _region_with(weight_replacement) with pytest.raises(ValueError, match="not a row-filter"): - sr._resolve_region_group( + region_resolution._resolve_region_group( country_module=cm, country="us", params={"region_group": ["state/hi"]} ) def test__requires_nonempty_group(self, monkeypatch): - monkeypatch.setattr(sr, "_resolve_dataset_reference", lambda _country: "x") + monkeypatch.setattr( + region_resolution, "resolve_dataset_reference", lambda _country: "x" + ) cm = _country_module({}) with pytest.raises(ValueError, match="at least one"): - sr._resolve_region_group( + region_resolution._resolve_region_group( country_module=cm, country="us", params={"region_group": []} ) def test__resolve_region_dispatches_to_group(self, monkeypatch): - monkeypatch.setattr(sr, "_resolve_dataset_reference", lambda _country: "x") + monkeypatch.setattr( + region_resolution, "resolve_dataset_reference", lambda _country: "x" + ) cm = _country_module({"state/hi": _region_with(_state(15))}) - res = sr._resolve_region( + res = sr.resolve_region( country_module=cm, country="us", params={"region_group": ["state/hi"]} ) assert isinstance(res.scoping_strategy, RegionGroupStrategy) 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 44b99a136..7166450fb 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -1,11 +1,5 @@ -"""Tests for building PolicyEngine v4 outputs into simulation API macro results.""" +"""Tests for the simulation output builder's schema, regions and timings.""" -from __future__ import annotations - -import json -from collections import Counter -from contextlib import AbstractContextManager -from dataclasses import dataclass from types import SimpleNamespace import pytest @@ -13,50 +7,18 @@ CURRENT_SINGLE_YEAR_MACRO_KEYS, CURRENT_SINGLE_YEAR_MACRO_RESULT, ) -from fixtures.spm_doubles import ( - installed_spm_selection, - spm_capable_simulation, -) -from fixtures.test_simulation_output_builder import ( - BASELINE_POVERTY_BY_AGE, - BASELINE_POVERTY_BY_GENDER, - BASELINE_POVERTY_BY_RACE, - INTRA_DECILE_COLLECTION, - REFORM_POVERTY_BY_AGE, - REFORM_POVERTY_BY_GENDER, - REFORM_POVERTY_BY_RACE, - FakeCollection, - FakeModelOutput, - fake_analysis, -) -from fixtures.weighted_frames import weighted_entity_frame -from policyengine.core.scoping_strategy import RowFilterStrategy -from policyengine_simulation_contract.job_progress import JobStage -from policyengine_simulation_contract.uk_geography import ( - UKLocalAuthorityBoundaryVersion, - UKLocalAuthorityMetadata, -) -from policyengine_simulation_executor import simulation_output_geographic -from policyengine_simulation_executor.job_progress import NULL_JOB_PROGRESS -from policyengine_simulation_executor.release_bundle import ( - BUNDLE_RECEIPT_FILENAME, - get_country_release_bundle, - resolve_runtime_bundle_dataset_uri, -) +from fixtures.test_simulation_output_builder import FakeCollection from policyengine_simulation_executor.simulation_macro_output import ( AgePovertyOutput, BaselineReformValue, BudgetaryImpact, BudgetaryOutput, CongressionalDistrictImpactOutput, - ConstituencyImpactRecord, DecileOutput, DetailedBudgetOutput, GenderPovertyOutput, - GeographicImpactOutput, InequalityOutput, IntraDecileOutput, - LocalAuthorityImpactRecord, PovertyByGenderOutput, PovertyModuleOutputs, PovertyOutput, @@ -72,317 +34,18 @@ build_congressional_district_impact, build_congressional_district_impact_output, ) -from policyengine_simulation_executor.simulation_output_local_areas import ( - LocalAreaBreakdownBuilder, -) -from policyengine_simulation_executor.simulation_runtime import ( - DatasetSelection, - RegionResolution, - _nondefault_data_folder, - _normalise_policy, - _resolve_dataset_reference, - _resolve_dataset_selection, - _resolve_region, - _run_simulation_impl_core, - resolve_data_folder, - run_simulation_impl, +from policyengine_simulation_observability.stages import Stage +from test_simulation_output_fakes import ( + BuilderOptions, + build_schema_output, + macro_baseline_reform, + make_output_builder, + with_observability_timings, ) -from policyengine_simulation_executor.simulation_runtime import ( - _load_dataset as _load_selected_dataset, -) -from policyengine_simulation_observability.stages import ANNUAL_IMPACT_STAGES, Stage - - -def _params_with_spm(params, selection): - """Return request parameters after the runtime resolves SPM settings.""" - return params if selection is None else {**params, "spm": selection} - - -def _assert_macro_result(result, selection): - """Check the ordinary macro result plus any installed SPM record.""" - 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() - if key not in spm_keys - } - if selection is None: - assert all(result[key] is None for key in spm_keys) - return - assert set(result["spm_provenance"]) == { - "schema_version", - "baseline", - "reform", - } - - -def _load_dataset(params, *, country_module=None, region_resolution=None): - """Exercise dataset resolution and loading together in direct loader tests.""" - selection = _resolve_dataset_selection(params, region_resolution=region_resolution) - return _load_selected_dataset( - params, country_module=country_module, selection=selection - ) - - -class _FakeOutputDataset: - def __init__(self, household): - self.data = SimpleNamespace(household=household) - - -class _FakeSimulation: - def __init__(self, household): - self.output_dataset = _FakeOutputDataset(household) - - def ensure(self): - msg = "test data is already materialized" - raise AssertionError(msg) - - -class _TrackedSpan(AbstractContextManager): - def __init__(self, runtime, name, attributes): - self.runtime = runtime - self.node = { - "name": name, - "attrs": dict(attributes or {}), - "children": [], - } - - def __enter__(self): - target = ( - self.runtime.stack[-1]["children"] - if self.runtime.stack - else self.runtime.roots - ) - target.append(self.node) - self.runtime.stack.append(self.node) - self.runtime.names.append(self.node["name"]) - return self - - def __exit__(self, exc_type, exc_value, traceback): - self.runtime.stack.pop() - - -class _TrackingRuntime: - def __init__(self): - self.roots = [] - self.stack = [] - self.names = [] - self.context = [] - - def span(self, name, *, attributes=None): - return _TrackedSpan(self, name, attributes) - - def set_context(self, **attributes): - self.context.extend(attributes.items()) - - -def _with_observability_timings(callback): - runtime = _TrackingRuntime() - result = callback(runtime) - counts = Counter(runtime.names) - return result, set(runtime.names), counts, runtime.roots - - -def _macro_baseline_reform(): - """Build households whose weighted totals are the golden macro result's budget. - - With weights 0.5 and 1.5, the weighted totals are a tax change of 100, a - benefit change of 30, a state tax change of 20, 2 households and a baseline - net income of 1000. Unweighted sums of the same rows would give 120, 40, 24 - and 1200, and summing the weights through the weighted frame (#767) would - give 2.5 households, so the golden result only matches a correct build. - """ - baseline = _FakeSimulation( - weighted_entity_frame( - { - "household_weight": [0.5, 1.5], - "household_net_income": [800.0, 400.0], - "household_tax": [60.0, 40.0], - "household_benefits": [20.0, 30.0], - "household_state_income_tax": [5.0, 5.0], - } - ) - ) - reform = _FakeSimulation( - weighted_entity_frame( - { - "household_weight": [0.5, 1.5], - "household_net_income": [810.0, 420.0], - "household_tax": [140.0, 80.0], - "household_benefits": [50.0, 40.0], - "household_state_income_tax": [21.0, 13.0], - } - ) - ) - return baseline, reform - - -@dataclass(frozen=True) -class BuilderOptions: - """Optional inputs for a SimulationOutputBuilder built from fakes.""" - - analysis: object = None - include_cliffs: bool | None = None - resolved_region_code: str | None = "us" - uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None - omit_uk_area_breakdowns: bool = False - - -def _simulation_output_builder( - country: str, - baseline, - reform, - options: BuilderOptions | None = None, -) -> SimulationOutputBuilder: - options = options or BuilderOptions() - analysis = options.analysis or fake_analysis() - - def economic_impact_analysis( - _baseline_simulation, - _reform_simulation, - *, - include_cliff_impacts=False, - ): - return analysis - - simulation_params = { - "country": country, - "data_version": "1.115.5" if country == "us" else "1.55.10", - } - if options.include_cliffs is not None: - simulation_params["include_cliffs"] = options.include_cliffs - - country_module = SimpleNamespace( - model=SimpleNamespace(version="1.715.2" if country == "us" else "2.88.20"), - economic_impact_analysis=economic_impact_analysis, - ) - return SimulationOutputBuilder( - country=country, - simulation_params=simulation_params, - country_module=country_module, - dataset=SimpleNamespace(metadata={}), - baseline=baseline, - reform=reform, - resolved_region_code=options.resolved_region_code, - uk_local_authority_metadata=options.uk_local_authority_metadata, - omit_uk_area_breakdowns=options.omit_uk_area_breakdowns, - ) - - -def _congressional_district_output() -> CongressionalDistrictImpactOutput: - return CongressionalDistrictImpactOutput( - districts=[ - { - "district": "AL-01", - "average_household_income_change": 10.0, - "relative_household_income_change": 0.01, - "winner_percentage": 0.6, - "loser_percentage": 0.3, - "no_change_percentage": 0.1, - "population": 1000.0, - } - ] - ) - - -def _constituency_impact_record() -> ConstituencyImpactRecord: - return ConstituencyImpactRecord( - constituency_code="E14000530", - constituency_name="Birmingham, Ladywood", - x=34, - y=-42, - average_household_income_change=10.0, - relative_household_income_change=0.01, - population=1000.0, - ) - - -def _local_authority_impact_record() -> LocalAuthorityImpactRecord: - return LocalAuthorityImpactRecord( - local_authority_code="E06000001", - local_authority_name="Hartlepool", - x=63, - y=-19, - average_household_income_change=12.0, - relative_household_income_change=0.02, - population=900.0, - ) - - -def _stub_policyengine_output_calls(monkeypatch, baseline, _reform) -> None: - def fake_poverty_module_function(name): - def compute(simulation): - if "by_age" in name: - return ( - BASELINE_POVERTY_BY_AGE - if simulation is baseline - else REFORM_POVERTY_BY_AGE - ) - if "by_gender" in name: - return ( - BASELINE_POVERTY_BY_GENDER - if simulation is baseline - else REFORM_POVERTY_BY_GENDER - ) - if "by_race" in name: - return ( - BASELINE_POVERTY_BY_RACE - if simulation is baseline - else REFORM_POVERTY_BY_RACE - ) - msg = f"Unexpected poverty output: {name}" - raise AssertionError(msg) - - return compute - - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_output_poverty._poverty_module_function", - fake_poverty_module_function, - ) - monkeypatch.setattr( - SimulationOutputBuilder, - "_build_intra_decile_output", - lambda self: self._build_intra_decile_output_from_collection( - INTRA_DECILE_COLLECTION - ), - ) - monkeypatch.setattr( - SimulationOutputBuilder, - "_build_congressional_district_impact", - lambda self: _congressional_district_output() if self.country == "us" else None, - ) - monkeypatch.setattr( - SimulationOutputBuilder, - "_build_uk_constituency_impact", - lambda self: ( - self._build_geographic_impact_output( - [_constituency_impact_record().model_dump(mode="json")] - ) - if self.country == "uk" - else None - ), - ) - monkeypatch.setattr( - SimulationOutputBuilder, - "_build_uk_local_authority_impact", - lambda self: ( - self._build_geographic_impact_output( - [_local_authority_impact_record().model_dump(mode="json")] - ) - if self.country == "uk" - else None - ), - ) - - -def _build_schema_output(monkeypatch, *, country: str = "us") -> SingleYearMacroOutput: - baseline, reform = _macro_baseline_reform() - _stub_policyengine_output_calls(monkeypatch, baseline, reform) - return _simulation_output_builder(country, baseline, reform).build() def test_builder_returns_schema_modules_before_legacy_dict_dump(monkeypatch): - output = _build_schema_output(monkeypatch) + output = build_schema_output(monkeypatch) assert isinstance(output, SingleYearMacroOutput) assert isinstance(output.budget, BudgetaryOutput) @@ -399,7 +62,7 @@ def test_builder_returns_schema_modules_before_legacy_dict_dump(monkeypatch): def test_builder_returns_existing_single_year_macro_shape(monkeypatch): - output = _build_schema_output(monkeypatch).model_dump(mode="json") + output = build_schema_output(monkeypatch).model_dump(mode="json") assert set(output) == CURRENT_SINGLE_YEAR_MACRO_KEYS assert output == CURRENT_SINGLE_YEAR_MACRO_RESULT @@ -509,7 +172,7 @@ def test_congressional_district_output_skips_narrow_regions(): def test_builder_passes_resolved_region_to_congressional_district_output(monkeypatch): - baseline, reform = _macro_baseline_reform() + baseline, reform = macro_baseline_reform() calls = [] def fake_build_congressional_district_impact( @@ -523,7 +186,7 @@ def fake_build_congressional_district_impact( fake_build_congressional_district_impact, ) - builder = _simulation_output_builder( + builder = make_output_builder( "us", baseline, reform, @@ -536,218 +199,6 @@ def fake_build_congressional_district_impact( assert calls == [("us", baseline, reform, "state/ut")] -def test_run_simulation_impl_records_runtime_timings_without_real_calculation( - monkeypatch, -): - dataset = object() - country_module = SimpleNamespace(model=SimpleNamespace(version="1.715.2")) - selection = installed_spm_selection() - baseline_simulation = spm_capable_simulation(selection) - reform_simulation = spm_capable_simulation(selection) - build_calls = [] - - def fake_build_simulation( - params, - *, - dataset, - dataset_selection, - policy, - region, - ): - build_calls.append( - (params, dataset, policy, region.scoping_strategy, region.code) - ) - return baseline_simulation if len(build_calls) == 1 else reform_simulation - - class FakeSimulationOutputBuilder: - def __init__(self, **kwargs): - self.kwargs = kwargs - - def serialize(self): - return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) - - monkeypatch.delenv("GOOGLE_APPLICATION_CREDENTIALS", raising=False) - monkeypatch.delenv("GOOGLE_APPLICATION_CREDENTIALS_JSON", raising=False) - monkeypatch.delenv("GCP_CREDENTIALS_JSON", raising=False) - monkeypatch.delenv("GOOGLE_CREDENTIALS", raising=False) - monkeypatch.delenv("SERVICE_ACCOUNT_JSON", raising=False) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._country_module", - lambda _country: country_module, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._resolve_region", - lambda **_kwargs: RegionResolution( - code="us", - dataset_reference=None, - scoping_strategy="mock-scoping", - ), - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._load_dataset", - lambda _params, country_module, selection, data_folder=None: dataset, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._build_simulation", - fake_build_simulation, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", - FakeSimulationOutputBuilder, - ) - - result, timings, counts, segment_tree = _with_observability_timings( - lambda runtime: run_simulation_impl( - { - "country": "us", - "baseline": {"gov.test.parameter": {"2026-01-01": 1}}, - "reform": {"gov.test.parameter": {"2026-01-01": 2}}, - }, - runtime=runtime, - ) - ) - - _assert_macro_result(result, selection) - assert set(timings) >= { - Stage.CREDENTIAL_SETUP, - Stage.REQUEST_PARSE, - Stage.COUNTRY_MODULE_LOAD, - Stage.REGION_RESOLUTION, - Stage.DATASET_RESOLUTION, - Stage.DATASET_LOAD, - Stage.POLICY_NORMALIZATION, - Stage.SIMULATION_BUILD, - } - assert counts[Stage.SIMULATION_BUILD] == 2 - assert [node["name"] for node in segment_tree] == [ - Stage.CREDENTIAL_SETUP, - Stage.REQUEST_PARSE, - Stage.COUNTRY_MODULE_LOAD, - Stage.REGION_RESOLUTION, - Stage.DATASET_RESOLUTION, - Stage.DATASET_LOAD, - Stage.POLICY_NORMALIZATION, - Stage.SIMULATION_BUILD, - Stage.SIMULATION_BUILD, - ] - simulation_builds = [ - node for node in segment_tree if node["name"] == Stage.SIMULATION_BUILD - ] - assert [node["attrs"] for node in simulation_builds] == [ - {"simulation_kind": "baseline"}, - {"simulation_kind": "reform"}, - ] - # region_code must be the RESOLVED code: it feeds the deterministic - # baseline-id predicate, and dropping it silently disables artifact - # reuse (every baseline becomes a miss). - built_params = _params_with_spm( - { - "country": "us", - "baseline": {"gov.test.parameter": {"2026-01-01": 1}}, - "reform": {"gov.test.parameter": {"2026-01-01": 2}}, - }, - selection, - ) - assert build_calls == [ - ( - built_params, - dataset, - {"gov.test.parameter": {"2026-01-01": 1}}, - "mock-scoping", - "us", - ), - ( - built_params, - dataset, - {"gov.test.parameter": {"2026-01-01": 2}}, - "mock-scoping", - "us", - ), - ] - - -def test_run_simulation_impl_exports_baseline_artifact_outcome(monkeypatch): - """run_simulation_impl must export the baseline artifact outcome attribute. - - The hit/incomplete/miss attribute is the rollout metric for the artifact - pipeline; losing the export would blind that measurement. - """ - dataset = object() - country_module = SimpleNamespace(model=SimpleNamespace(version="1.715.2")) - selection = installed_spm_selection() - simulations = { - "baseline": spm_capable_simulation(selection, artifact_outcome="incomplete"), - "reform": spm_capable_simulation(selection), - } - build_count = [0] - - def fake_build_simulation( - _params, - *, - dataset, - dataset_selection, - policy, - region, - ): - build_count[0] += 1 - return simulations["baseline"] if build_count[0] == 1 else simulations["reform"] - - class FakeSimulationOutputBuilder: - def __init__(self, **kwargs): - self.kwargs = kwargs - - def serialize(self): - return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) - - for env in ( - "GOOGLE_APPLICATION_CREDENTIALS", - "GOOGLE_APPLICATION_CREDENTIALS_JSON", - "GCP_CREDENTIALS_JSON", - "GOOGLE_CREDENTIALS", - "SERVICE_ACCOUNT_JSON", - ): - monkeypatch.delenv(env, raising=False) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._country_module", - lambda _country: country_module, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._resolve_region", - lambda **_kwargs: RegionResolution( - code="us", - dataset_reference=None, - scoping_strategy="mock-scoping", - ), - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._load_dataset", - lambda _params, country_module, selection, data_folder=None: dataset, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._build_simulation", - fake_build_simulation, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", - FakeSimulationOutputBuilder, - ) - params = { - "country": "us", - "baseline": {"gov.test.parameter": {"2026-01-01": 1}}, - "reform": {"gov.test.parameter": {"2026-01-01": 2}}, - } - runtime = _TrackingRuntime() - run_simulation_impl(params, runtime=runtime) - assert ("baseline_artifact", "incomplete") in runtime.context - - # A plain Simulation baseline (no artifact_outcome) must emit nothing. - runtime.context.clear() - build_count[0] = 0 - simulations["baseline"] = spm_capable_simulation(selection) - run_simulation_impl(params, runtime=runtime) - assert all(key != "baseline_artifact" for key, _ in runtime.context) - - def test_builder_records_output_timings_without_real_calculation(monkeypatch): baseline = object() reform = object() @@ -884,7 +335,7 @@ def serialize(runtime): builder.runtime = runtime return builder.serialize() - result, timings, counts, segment_tree = _with_observability_timings(serialize) + result, timings, counts, segment_tree = with_observability_timings(serialize) assert result["model_version"] == "mock-model-version" assert result["data_version"] == "mock-data-version" @@ -981,7 +432,7 @@ def test_decile_relative_is_a_fraction_of_the_baseline_mean(): def test_builder_maps_uk_wealth_outputs_and_omits_us_only_race(monkeypatch): - output = _build_schema_output(monkeypatch, country="uk").model_dump(mode="json") + output = build_schema_output(monkeypatch, country="uk").model_dump(mode="json") assert output["poverty_by_race"] is None assert output["wealth_decile"] == { @@ -991,2058 +442,3 @@ def test_builder_maps_uk_wealth_outputs_and_omits_us_only_race(monkeypatch): assert output["intra_wealth_decile"]["deciles"]["Lose more than 5%"] == [0.1] assert output["constituency_impact"][0]["constituency_code"] == "E14000530" assert output["local_authority_impact"][0]["local_authority_code"] == "E06000001" - - -def test_builder_calls_policyengine_economic_impact_analysis(): - baseline, reform = _macro_baseline_reform() - analysis = fake_analysis() - calls = [] - - def economic_impact_analysis( - baseline_simulation, - reform_simulation, - *, - include_cliff_impacts=False, - ): - calls.append((baseline_simulation, reform_simulation, include_cliff_impacts)) - return analysis - - country_module = SimpleNamespace( - model=SimpleNamespace(version="1.715.2"), - economic_impact_analysis=economic_impact_analysis, - ) - builder = SimulationOutputBuilder( - country="us", - simulation_params={"country": "us", "data_version": "1.115.5"}, - country_module=country_module, - dataset=SimpleNamespace(metadata={}), - baseline=baseline, - reform=reform, - ) - - assert builder.analysis is analysis - assert builder.analysis is analysis - assert calls == [(baseline, reform, False)] - - -def test_builder_passes_include_cliffs_to_policyengine_economic_impact_analysis(): - baseline, reform = _macro_baseline_reform() - analysis = fake_analysis() - calls = [] - - def economic_impact_analysis( - _baseline_simulation, - _reform_simulation, - *, - include_cliff_impacts=False, - ): - calls.append(include_cliff_impacts) - return analysis - - country_module = SimpleNamespace( - model=SimpleNamespace(version="1.715.2"), - economic_impact_analysis=economic_impact_analysis, - ) - builder = SimulationOutputBuilder( - country="us", - simulation_params={ - "country": "us", - "data_version": "1.115.5", - "include_cliffs": True, - }, - country_module=country_module, - dataset=SimpleNamespace(metadata={}), - baseline=baseline, - reform=reform, - ) - - assert builder.analysis is analysis - assert calls == [True] - - -def test_builder_serializes_cliff_impact_when_requested(monkeypatch): - baseline, reform = _macro_baseline_reform() - _stub_policyengine_output_calls(monkeypatch, baseline, reform) - analysis = fake_analysis() - analysis.cliff_impact = FakeModelOutput( - { - "baseline": {"cliff_gap": 10.0, "cliff_share": 0.25}, - "reform": {"cliff_gap": 20.0, "cliff_share": 0.5}, - } - ) - - output = _simulation_output_builder( - "us", - baseline, - reform, - BuilderOptions(analysis=analysis, include_cliffs=True), - ).build() - - assert output.cliff_impact is not None - assert output.model_dump(mode="json")["cliff_impact"] == { - "baseline": {"cliff_gap": 10.0, "cliff_share": 0.25}, - "reform": {"cliff_gap": 20.0, "cliff_share": 0.5}, - } - - -def test_normalise_policy_preserves_legacy_period_range_keys(): - assert _normalise_policy({"gov.test.parameter": {"2026-01-01.2100-12-31": 1}}) == { - "gov.test.parameter": {"2026-01-01.2100-12-31": 1} - } - - -def test_run_simulation_impl_core_builds_and_serializes_macro_output(monkeypatch): - dataset = object() - country_module = SimpleNamespace(model=SimpleNamespace(version="1.715.2")) - selection = installed_spm_selection() - baseline_simulation = spm_capable_simulation(selection) - reform_simulation = spm_capable_simulation(selection) - build_calls = [] - builder_calls = [] - - def fake_country_module(country): - assert country == "us" - return country_module - - def fake_build_simulation( - params, - *, - dataset, - dataset_selection, - policy, - region, - ): - build_calls.append((params, dataset, policy, region.scoping_strategy)) - return baseline_simulation if len(build_calls) == 1 else reform_simulation - - class FakeSimulationOutputBuilder: - def __init__(self, **kwargs): - builder_calls.append(kwargs) - - def serialize(self): - return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) - - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._country_module", - fake_country_module, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._resolve_region", - lambda **_kwargs: RegionResolution(code="us", dataset_reference=None), - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._load_dataset", - lambda _params, country_module, selection, data_folder=None: dataset, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._build_simulation", - fake_build_simulation, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", - FakeSimulationOutputBuilder, - ) - - runtime = _TrackingRuntime() - result = _run_simulation_impl_core( - { - "country": "us", - "baseline": {"gov.test.parameter": {"2026-01-01.2100-12-31": 1}}, - "reform": {"gov.test.parameter": {"2026-01-01.2100-12-31": 2}}, - }, - runtime=runtime, - ) - - _assert_macro_result(result, selection) - assert build_calls[0][2] == {"gov.test.parameter": {"2026-01-01.2100-12-31": 1}} - assert build_calls[1][2] == {"gov.test.parameter": {"2026-01-01.2100-12-31": 2}} - assert build_calls[0][3] is None - assert build_calls[1][3] is None - assert builder_calls == [ - { - "country": "us", - "simulation_params": _params_with_spm( - { - "country": "us", - "baseline": {"gov.test.parameter": {"2026-01-01.2100-12-31": 1}}, - "reform": {"gov.test.parameter": {"2026-01-01.2100-12-31": 2}}, - }, - selection, - ), - "country_module": country_module, - "dataset": dataset, - "baseline": baseline_simulation, - "reform": reform_simulation, - "resolved_data_version": None, - "resolved_region_code": "us", - "runtime": runtime, - "uk_local_authority_metadata": None, - "omit_uk_area_breakdowns": False, - "progress": NULL_JOB_PROGRESS, - } - ] - - -def test_run_simulation_impl_core_passes_region_scoping_to_simulations(monkeypatch): - dataset = object() - country_module = SimpleNamespace(model=SimpleNamespace(version="1.715.2")) - selection = installed_spm_selection() - baseline_simulation = spm_capable_simulation(selection) - reform_simulation = spm_capable_simulation(selection) - scoping_strategy = object() - region_resolution = RegionResolution( - code="state/ut", - dataset_reference=None, - scoping_strategy=scoping_strategy, - ) - build_calls = [] - load_selections = [] - - def fake_build_simulation( - params, - *, - dataset, - dataset_selection, - policy, - region, - ): - build_calls.append( - (params, dataset, policy, region.scoping_strategy, dataset_selection) - ) - return baseline_simulation if len(build_calls) == 1 else reform_simulation - - def fake_load_dataset(_params, *, country_module, selection, data_folder=None): - load_selections.append(selection) - return dataset - - class FakeSimulationOutputBuilder: - def __init__(self, **kwargs): - pass - - def serialize(self): - return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) - - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._country_module", - lambda _country: country_module, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._resolve_region", - lambda **_kwargs: region_resolution, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._load_dataset", - fake_load_dataset, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._build_simulation", - fake_build_simulation, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", - FakeSimulationOutputBuilder, - ) - - result = _run_simulation_impl_core( - { - "country": "us", - "region": "state/ut", - "baseline": {}, - "reform": {}, - }, - runtime=_TrackingRuntime(), - ) - - _assert_macro_result(result, selection) - assert build_calls[0][3] is scoping_strategy - assert build_calls[1][3] is scoping_strategy - assert len(load_selections) == 1 - assert build_calls[0][4] is load_selections[0] - assert build_calls[1][4] is load_selections[0] - - -class _StoppedAfterLoadingError(Exception): - """Raised by a stubbed simulation build once the dataset has loaded.""" - - -def _load_through_runtime( - monkeypatch, *, region_resolution, selection, baked_folder, time_period=2026 -): - """Run the annual runtime up to its dataset load and return what it did.""" - from policyengine_simulation_executor import state_datasets - - load_calls = [] - lookups = [] - - def fake_load_dataset(_params, *, country_module, selection, data_folder=None): - load_calls.append({"selection": selection, "data_folder": data_folder}) - return object() - - def fake_baked_state_data_folder(**kwargs): - lookups.append(kwargs) - return baked_folder - - def stop_after_loading(*_args, **_kwargs): - raise _StoppedAfterLoadingError - - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._country_module", - lambda _country: SimpleNamespace(model=SimpleNamespace(version="2.2.1")), - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._resolve_region", - lambda **_kwargs: region_resolution, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._resolve_dataset_selection", - lambda _params, *, region_resolution: selection, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._load_dataset", - fake_load_dataset, - ) - monkeypatch.setattr( - state_datasets, "baked_state_data_folder", fake_baked_state_data_folder - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._build_simulation", - stop_after_loading, - ) - runtime = _TrackingRuntime() - with pytest.raises(_StoppedAfterLoadingError): - _run_simulation_impl_core( - { - "country": "us", - "region": region_resolution.code, - "time_period": time_period, - "baseline": {}, - "reform": {}, - }, - runtime=runtime, - ) - return SimpleNamespace( - load_calls=load_calls, lookups=lookups, context=dict(runtime.context) - ) - - -def _regional_selection(): - return DatasetSelection( - name="populace_us_2024_acs_local", - uri="hf://policyengine/populace-us/populace_us_2024_acs_local.h5@rev", - is_default=False, - ) - - -def test_runtime_loads_a_state_from_its_baked_folder(monkeypatch): - loaded = _load_through_runtime( - monkeypatch, - region_resolution=RegionResolution( - code="state/ut", dataset_reference="ref", state_code="UT" - ), - selection=_regional_selection(), - baked_folder="/opt/policyengine/data/state-datasets/acs/UT", - ) - - assert loaded.load_calls == [ - { - "selection": _regional_selection(), - "data_folder": "/opt/policyengine/data/state-datasets/acs/UT", - } - ] - assert loaded.lookups == [ - { - "country": "us", - "selection": _regional_selection(), - "state_code": "UT", - "year": 2026, - "data_folder": resolve_data_folder(), - } - ] - assert loaded.context["dataset_source"] == "baked_state" - - -def test_runtime_prepares_a_state_on_demand_without_a_baked_file(monkeypatch): - loaded = _load_through_runtime( - monkeypatch, - region_resolution=RegionResolution( - code="congressional_district/UT-01", - dataset_reference="ref", - state_code="UT", - ), - selection=_regional_selection(), - baked_folder=None, - time_period=2028, - ) - - assert loaded.load_calls[0]["data_folder"] is None - assert loaded.lookups[0]["year"] == 2028 - assert loaded.context["dataset_source"] == "on_demand" - - -def test_runtime_loads_a_national_run_from_the_default_folder(monkeypatch): - bundle = get_country_release_bundle("us") - loaded = _load_through_runtime( - monkeypatch, - region_resolution=RegionResolution(code="us"), - selection=DatasetSelection( - name=bundle.default_dataset, - uri=bundle.default_dataset_uri, - is_default=True, - ), - baked_folder=None, - ) - - assert loaded.load_calls[0]["data_folder"] is None - assert "dataset_source" not in loaded.context - - -def _builder_data_version(monkeypatch, *, region_resolution, dataset_selection): - """Run the annual runtime and return the data version it gives the builder.""" - from policyengine_simulation_executor import state_datasets - - selection = installed_spm_selection() - simulations = [spm_capable_simulation(selection), spm_capable_simulation(selection)] - builder_calls = [] - - class FakeSimulationOutputBuilder: - def __init__(self, **kwargs): - builder_calls.append(kwargs) - - def serialize(self): - return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) - - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._country_module", - lambda _country: SimpleNamespace(model=SimpleNamespace(version="2.2.1")), - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._resolve_region", - lambda **_kwargs: region_resolution, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._resolve_dataset_selection", - lambda _params, *, region_resolution: dataset_selection, - ) - monkeypatch.setattr( - state_datasets, "baked_state_data_folder", lambda **_kwargs: None - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._load_dataset", - lambda _params, country_module, selection, data_folder=None: object(), - ) - build_iter = iter(simulations) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._build_simulation", - lambda *_args, **_kwargs: next(build_iter), - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", - FakeSimulationOutputBuilder, - ) - - _run_simulation_impl_core( - { - "country": "us", - "region": region_resolution.code, - "baseline": {}, - "reform": {}, - }, - runtime=_TrackingRuntime(), - ) - return builder_calls[0]["resolved_data_version"] - - -def test_a_state_run_reports_the_revision_of_the_dataset_it_loads(monkeypatch): - from policyengine.provenance.manifest import get_release_manifest - - bundle = get_country_release_bundle("us") - regional = "populace_us_2024_acs_local" - - data_version = _builder_data_version( - monkeypatch, - region_resolution=RegionResolution( - code="state/ut", dataset_reference="ref", state_code="UT" - ), - dataset_selection=DatasetSelection( - name=regional, uri=bundle.dataset_uris[regional], is_default=False - ), - ) - - assert data_version == get_release_manifest("us").datasets[regional].revision - assert data_version != bundle.data_version - - -def test_a_national_run_keeps_the_release_data_version(monkeypatch): - bundle = get_country_release_bundle("us") - - data_version = _builder_data_version( - monkeypatch, - region_resolution=RegionResolution(code="us"), - dataset_selection=DatasetSelection( - name=bundle.default_dataset, - uri=bundle.default_dataset_uri, - is_default=True, - ), - ) - - assert data_version is None - - -def test_run_simulation_impl_core_passes_detected_uk_boundary_to_builder(monkeypatch): - # A bundle that does not route area regions: the national file carries - # the area codes and the run builds the area breakdowns from them. - _routes_uk_area_regions(monkeypatch, routed=False) - dataset = SimpleNamespace( - data=SimpleNamespace( - entity_data={ - "household": weighted_entity_frame( - {"la_code_oa": ["E06000001", "E07000026"]} - ) - } - ) - ) - country_module = SimpleNamespace(model=SimpleNamespace(version="2.90.0")) - builder_calls: list[dict] = [] - - class FakeSimulationOutputBuilder: - def __init__(self, **kwargs): - builder_calls.append(kwargs) - - def serialize(self): - return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) - - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._country_module", - lambda _country: country_module, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._resolve_region", - lambda **_kwargs: RegionResolution(code="uk", dataset_reference=None), - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._load_dataset", - lambda _params, country_module, selection, data_folder=None: dataset, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime._build_simulation", - lambda *_args, **_kwargs: SimpleNamespace( - on_first_ensure=lambda _callback: None - ), - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", - FakeSimulationOutputBuilder, - ) - - _run_simulation_impl_core( - {"country": "uk", "baseline": {}, "reform": {}}, - runtime=_TrackingRuntime(), - ) - - metadata = builder_calls[0]["uk_local_authority_metadata"] - assert metadata == UKLocalAuthorityMetadata( - boundary_version=UKLocalAuthorityBoundaryVersion.LAD22 - ) - assert builder_calls[0]["omit_uk_area_breakdowns"] is False - - -def _routes_uk_area_regions(monkeypatch, *, routed: bool) -> None: - monkeypatch.setattr( - "policyengine_simulation_executor.uk_local_authority_metadata." - "uk_area_regions_routed", - lambda: routed, - ) - - -def _uk_dataset(**household: list[object]) -> SimpleNamespace: - return SimpleNamespace( - data=SimpleNamespace( - entity_data={"household": weighted_entity_frame(household)} - ) - ) - - -_NATIONAL_SELECTION = DatasetSelection( - "microcosm_uk_2024_25", "hf://policyengine/uk/national.h5@cut", is_default=True -) -_LOCAL_AREA_SELECTION = DatasetSelection( - "microcosm_uk_2024_25_dense", "hf://policyengine/uk/local.h5@cut", is_default=False -) - - -_LOCAL_AREA_REVISION = "microcosm-uk-2024-25-dense-20261019T000432Z-ee5c2dbe" -_WALES = RowFilterStrategy(variable_name="region", variable_value="WALES") -_TEST_POLICIES = { - "baseline": {"gov.test.parameter": {"2026-01-01.2100-12-31": 1}}, - "reform": {"gov.test.parameter": {"2026-01-01.2100-12-31": 2}}, -} -_WALES_NATIONAL_PARAMS = {"country": "uk", "region": "country/wales", **_TEST_POLICIES} -_WALES_BREAKDOWN_PARAMS = { - "country": "uk", - "scope": "macro", - "region": "country/wales", - "time_period": "2026", - # The Public API always sends include_cliffs: false. - "include_cliffs": False, - "breakdown": "local_areas", - **_TEST_POLICIES, -} - - -def _routed_area_dataset(*, without: str | None = None) -> SimpleNamespace: - """Return a local-area dataset with every column its area outputs read. - - ``without`` names one column to leave out. - """ - household: dict[str, list[object]] = { - "region": ["WALES", "WALES"], - "constituency_code_oa": ["E14001063", "W07000112"], - "la_code_oa": ["E06000063", "W06000001"], - "region_code_oa": ["WALES", "WALES"], - "source_household_id": [1, 2], - } - household.pop(without or "", None) - return _uk_dataset(**household) - - -@dataclass -class _UKAreaRun: - """A UK run through ``_run_simulation_impl_core`` with every seam faked.""" - - builds: list[dict] - builder_calls: list[dict] - runtime: _TrackingRuntime - # Each selection _load_dataset was asked for, and each country module load. - loads: list[DatasetSelection] - country_modules: list[str] - # The builder methods the run called to produce its output, and the output. - serialized: list[str] - output: object = None - - def dataset_load_attributes(self) -> list[dict]: - return [ - node["attrs"] - for node in self.runtime.roots - if node["name"] == ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD) - ] - - -class _RoutedUKSeams: - """Fakes for each seam of a routed UK run, recording into one ``_UKAreaRun``.""" - - def __init__(self, run, *, local_area_dataset): - self.run = run - self.local_area_dataset = local_area_dataset - self.national_dataset = _uk_dataset(region=["WALES", "WALES"]) - - def country_module(self, country): - self.run.country_modules.append(country) - return SimpleNamespace(model=SimpleNamespace(version="2.132.6")) - - def resolve_region(self, *, params, **_kwargs): - region = params.get("region") - if region == "country/wales": - return RegionResolution( - code="country/wales", dataset_reference=None, scoping_strategy=_WALES - ) - region_type = str(region).partition("/")[0] - if region_type in {"constituency", "local_authority"}: - # An area region names the routed local-area dataset itself. - return RegionResolution( - code=region, - dataset_reference=_LOCAL_AREA_SELECTION.name, - scoping_strategy=RowFilterStrategy( - variable_name=f"{region_type}_code_oa", - variable_value=region.partition("/")[2], - ), - ) - return RegionResolution(code="uk", dataset_reference=None) - - def resolve_selection(self, _params, *, region_resolution): - if region_resolution.dataset_reference == _LOCAL_AREA_SELECTION.name: - return _LOCAL_AREA_SELECTION - return _NATIONAL_SELECTION - - def load_dataset(self, _params, *, selection, country_module, data_folder=None): - self.run.loads.append(selection) - if selection is _LOCAL_AREA_SELECTION: - return self.local_area_dataset - return self.national_dataset - - def build_simulation(self, _params, *, dataset, dataset_selection, policy, region): - simulation = SimpleNamespace(on_first_ensure=lambda _callback: None) - self.run.builds.append( - { - "dataset": dataset, - "selection": dataset_selection, - "policy": policy, - "scoping_strategy": region.scoping_strategy, - "simulation": simulation, - } - ) - return simulation - - def output_builder(self, **kwargs): - self.run.builder_calls.append(kwargs) - return _RecordingBuilder( - self.run, "macro", dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) - ) - - def breakdown_builder(self, **kwargs): - self.run.builder_calls.append(kwargs) - return _RecordingBuilder( - self.run, "local_area_breakdown", {"kind": "local_areas"} - ) - - -class _RecordingBuilder: - """Records which builder serialized the run's output.""" - - def __init__(self, run, kind, output): - self.run = run - self.kind = kind - self.output = output - - def serialize(self): - self.run.serialized.append(self.kind) - return self.output - - -def _install_routed_uk_fakes( - monkeypatch, *, local_area_dataset, routed: bool = True -) -> _UKAreaRun: - """Fake every seam of a UK run whose national file carries no area codes.""" - run = _UKAreaRun( - builds=[], - builder_calls=[], - runtime=_TrackingRuntime(), - loads=[], - country_modules=[], - serialized=[], - ) - seams = _RoutedUKSeams(run, local_area_dataset=local_area_dataset) - _routes_uk_area_regions(monkeypatch, routed=routed) - monkeypatch.setattr( - "policyengine_simulation_executor.uk_local_authority_metadata." - "routed_uk_area_dataset_name", - lambda: _LOCAL_AREA_SELECTION.name, - ) - for name, value in { - "_country_module": seams.country_module, - "_resolve_region": seams.resolve_region, - "_resolve_dataset_selection": seams.resolve_selection, - "_load_dataset": seams.load_dataset, - "_build_simulation": seams.build_simulation, - "SimulationOutputBuilder": seams.output_builder, - "LocalAreaBreakdownBuilder": seams.breakdown_builder, - }.items(): - monkeypatch.setattr( - f"policyengine_simulation_executor.simulation_runtime.{name}", value - ) - return run - - -def _run_routed_uk(monkeypatch, *, params, local_area_dataset) -> _UKAreaRun: - run = _install_routed_uk_fakes(monkeypatch, local_area_dataset=local_area_dataset) - run.output = _run_simulation_impl_core(dict(params), runtime=run.runtime) - return run - - -def _span_attributes(nodes) -> list[dict]: - """Return the attributes of every span in a tracked span tree.""" - attributes = [] - for node in nodes: - attributes.append(node["attrs"]) - attributes.extend(_span_attributes(node["children"])) - return attributes - - -@pytest.mark.parametrize("region", ["uk", "country/wales"]) -def test_a_routed_uk_wide_or_nation_run_builds_no_area_pair(monkeypatch, region): - from policyengine_simulation_executor import uk_local_authority_metadata - - run = _run_routed_uk( - monkeypatch, - params={**_WALES_NATIONAL_PARAMS, "region": region}, - local_area_dataset=_routed_area_dataset(), - ) - - # One dataset, the national file, and one baseline and reform on it. - assert run.loads == [_NATIONAL_SELECTION] - assert [build["selection"] for build in run.builds] == [_NATIONAL_SELECTION] * 2 - (builder_kwargs,) = run.builder_calls - assert builder_kwargs["baseline"] is run.builds[0]["simulation"] - assert builder_kwargs["reform"] is run.builds[1]["simulation"] - assert builder_kwargs["omit_uk_area_breakdowns"] is True - assert builder_kwargs["omit_uk_area_breakdowns"] is ( - uk_local_authority_metadata.uk_area_breakdowns_omitted("uk", region) - ) - assert "uk_area_simulations" not in builder_kwargs - assert run.serialized == ["macro"] - assert {"dataset_role": "uk_area"} not in _span_attributes(run.runtime.roots) - - -@pytest.mark.parametrize( - "region", ["constituency/E14001063", "local_authority/E06000063"] -) -@pytest.mark.parametrize( - "missing", - ["constituency_code_oa", "la_code_oa", "region_code_oa", "source_household_id"], -) -def test_a_routed_area_run_needs_every_column_its_outputs_read_before_simulating( - monkeypatch, region, missing -): - run = _install_routed_uk_fakes( - monkeypatch, local_area_dataset=_routed_area_dataset(without=missing) - ) - - with pytest.raises( - ValueError, - match=rf"'microcosm_uk_2024_25_dense', which the bundle routes .* has no " - rf"{missing} column", - ): - _run_simulation_impl_core( - {**_WALES_NATIONAL_PARAMS, "region": region}, runtime=run.runtime - ) - - assert run.loads == [_LOCAL_AREA_SELECTION] - assert run.builds == [] - - -def test_a_routed_area_run_with_every_column_detects_its_boundaries(monkeypatch): - run = _run_routed_uk( - monkeypatch, - params={**_WALES_NATIONAL_PARAMS, "region": "local_authority/E06000063"}, - local_area_dataset=_routed_area_dataset(), - ) - - (builder_kwargs,) = run.builder_calls - assert builder_kwargs["uk_local_authority_metadata"] == UKLocalAuthorityMetadata( - boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 - ) - assert builder_kwargs["omit_uk_area_breakdowns"] is False - - -def test_a_local_area_breakdown_runs_on_the_local_area_dataset_alone(monkeypatch): - local_area_dataset = _routed_area_dataset() - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_runtime." - "_manifest_dataset_revision", - lambda country, name: { - ("uk", _LOCAL_AREA_SELECTION.name): _LOCAL_AREA_REVISION - }[(country, name)], - ) - - run = _run_routed_uk( - monkeypatch, - params=_WALES_BREAKDOWN_PARAMS, - local_area_dataset=local_area_dataset, - ) - - # The national file is never loaded: one load, one baseline and reform. - assert run.loads == [_LOCAL_AREA_SELECTION] - assert [build["selection"] for build in run.builds] == [_LOCAL_AREA_SELECTION] * 2 - assert all(build["dataset"] is local_area_dataset for build in run.builds) - assert all(build["scoping_strategy"] is _WALES for build in run.builds) - assert [build["policy"] for build in run.builds] == [ - _normalise_policy(_TEST_POLICIES["baseline"]), - _normalise_policy(_TEST_POLICIES["reform"]), - ] - (builder_kwargs,) = run.builder_calls - assert builder_kwargs["baseline"] is run.builds[0]["simulation"] - assert builder_kwargs["reform"] is run.builds[1]["simulation"] - assert builder_kwargs["local_authority_metadata"] == UKLocalAuthorityMetadata( - boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 - ) - assert builder_kwargs["dataset_name"] == _LOCAL_AREA_SELECTION.name - assert builder_kwargs["data_version"] == _LOCAL_AREA_REVISION - assert builder_kwargs["year"] == 2026 - assert run.serialized == ["local_area_breakdown"] - assert run.output == {"kind": "local_areas"} - assert run.dataset_load_attributes() == [{"dataset_role": "uk_area"}] - - -def test_a_scotland_breakdown_resolves_against_the_installed_bundle(monkeypatch): - # Only the dataset load, the simulations and the builder are faked: the - # country module, region resolution, dataset selection, routing and the - # revision come from the installed .py bundle, with no data or network. - from policyengine.provenance.manifest import get_release_manifest - from policyengine_simulation_executor import simulation_runtime - - loads: list[DatasetSelection] = [] - builds: list[dict] = [] - builder_calls: list[dict] = [] - - def load_dataset(_params, *, selection, country_module, data_folder=None): - loads.append(selection) - return _routed_area_dataset() - - def build_simulation(_params, *, dataset, dataset_selection, policy, region): - builds.append( - { - "selection": dataset_selection, - "region_code": region.code, - "scoping_strategy": region.scoping_strategy, - } - ) - return SimpleNamespace(on_first_ensure=lambda _callback: None) - - class RecordingBreakdownBuilder: - def __init__(self, **kwargs): - builder_calls.append(kwargs) - - def serialize(self): - return {"kind": "local_areas"} - - for name, value in { - "_load_dataset": load_dataset, - "_build_simulation": build_simulation, - "LocalAreaBreakdownBuilder": RecordingBreakdownBuilder, - }.items(): - monkeypatch.setattr(simulation_runtime, name, value) - - output = _run_simulation_impl_core( - {**_WALES_BREAKDOWN_PARAMS, "region": "country/scotland"}, - runtime=_TrackingRuntime(), - ) - - (selection,) = loads - assert selection.name == "microcosm_uk_2024_25_dense" - assert selection.is_default is False - scotland = RowFilterStrategy(variable_name="region", variable_value="SCOTLAND") - assert [build["scoping_strategy"] for build in builds] == [scotland, scotland] - assert [build["selection"] for build in builds] == [selection, selection] - assert {build["region_code"] for build in builds} == {"country/scotland"} - (builder_kwargs,) = builder_calls - assert builder_kwargs["dataset_name"] == selection.name - assert builder_kwargs["data_version"] == ( - get_release_manifest("uk").datasets[selection.name].revision - ) - assert output == {"kind": "local_areas"} - - -def test_a_dataset_revision_is_read_from_the_release_manifest(monkeypatch): - from policyengine_simulation_executor import simulation_runtime - - manifests = { - "uk": SimpleNamespace( - datasets={ - _LOCAL_AREA_SELECTION.name: SimpleNamespace( - revision=_LOCAL_AREA_REVISION - ), - "unrevised": SimpleNamespace(revision=""), - } - ) - } - monkeypatch.setattr( - "policyengine.provenance.manifest.get_release_manifest", manifests.__getitem__ - ) - - revision = simulation_runtime._manifest_dataset_revision - assert revision("uk", _LOCAL_AREA_SELECTION.name) == _LOCAL_AREA_REVISION - assert revision("uk", "unrevised") is None - assert revision("uk", "absent") is None - - -def test_a_uk_wide_local_area_breakdown_simulates_the_whole_dataset(monkeypatch): - run = _run_routed_uk( - monkeypatch, - params={**_WALES_BREAKDOWN_PARAMS, "region": "uk"}, - local_area_dataset=_routed_area_dataset(), - ) - - assert run.loads == [_LOCAL_AREA_SELECTION] - assert [build["scoping_strategy"] for build in run.builds] == [None, None] - assert run.serialized == ["local_area_breakdown"] - - -@pytest.mark.parametrize( - ("change", "message"), - [ - ({"country": "us"}, "only for UK requests"), - ({"region": "constituency/E14001063"}, "'constituency/E14001063'"), - ({"region": "local_authority/E06000063"}, "'local_authority/E06000063'"), - ({"include_cliffs": True}, "include_cliffs=true"), - ({"region": None, "region_group": ["country/wales"]}, "region_group"), - ({"breakdown": "districts"}, "'districts'"), - ], -) -def test_an_invalid_local_area_breakdown_is_refused_before_any_load( - monkeypatch, change, message -): - run = _install_routed_uk_fakes( - monkeypatch, local_area_dataset=_routed_area_dataset() - ) - - with pytest.raises(ValueError, match=message): - _run_simulation_impl_core( - {**_WALES_BREAKDOWN_PARAMS, **change}, runtime=run.runtime - ) - - assert run.country_modules == [] - assert run.loads == [] - assert run.builds == [] - - -def test_a_local_area_breakdown_needs_a_bundle_that_routes_area_regions(monkeypatch): - run = _install_routed_uk_fakes( - monkeypatch, local_area_dataset=_routed_area_dataset(), routed=False - ) - - with pytest.raises( - ValueError, - match="does not route UK constituency and local-authority regions", - ): - _run_simulation_impl_core(dict(_WALES_BREAKDOWN_PARAMS), runtime=run.runtime) - - assert run.loads == [] - assert run.builds == [] - - -@pytest.mark.parametrize( - "missing", - ["constituency_code_oa", "la_code_oa", "region_code_oa", "source_household_id"], -) -def test_a_local_area_breakdown_needs_every_column_its_outputs_read( - monkeypatch, missing -): - local_area_dataset = _routed_area_dataset(without=missing) - - with pytest.raises( - ValueError, - match=rf"'microcosm_uk_2024_25_dense', which the bundle routes .* has no " - rf"{missing} column", - ): - _run_routed_uk( - monkeypatch, - params=_WALES_BREAKDOWN_PARAMS, - local_area_dataset=local_area_dataset, - ) - - -_LOCAL_AREA_URI = ( - "hf://policyengine/populace-uk-private/microcosm_uk_2024_25_dense.h5" - "@microcosm-uk-2024-25-dense-20261019T000432Z-ee5c2dbe" -) - - -def _registry_area(region_type: str, code: str, label: str) -> SimpleNamespace: - return SimpleNamespace( - code=f"{region_type}/{code}", - label=label, - region_type=region_type, - parent_code="uk", - dataset_path=_LOCAL_AREA_URI, - scoping_strategy=object(), - ) - - -def _routed_uk_country_module() -> SimpleNamespace: - """Return a UK model whose registry lists area regions, as routing gives.""" - areas = [ - _registry_area("constituency", "E14001063", "Aldershot"), - _registry_area("constituency", "W07000112", "Ynys Môn"), - _registry_area("local_authority", "E07000110", "Maidstone"), - _registry_area("local_authority", "N09000011", "Ards and North Down"), - ] - by_code = {area.code: area for area in areas} - registry = SimpleNamespace( - get_by_type=lambda region_type: [ - area for area in areas if area.region_type == region_type - ] - ) - return SimpleNamespace( - model=SimpleNamespace(get_region=by_code.get, region_registry=registry) - ) - - -@pytest.mark.parametrize( - ("requested", "expected"), - [ - ("constituency/E14001063", "constituency/E14001063"), - ("constituency/e14001063", "constituency/E14001063"), - ("constituency/Aldershot", "constituency/E14001063"), - ("constituency/ynys môn", "constituency/W07000112"), - ("local_authority/Maidstone", "local_authority/E07000110"), - ("local_authority/n09000011", "local_authority/N09000011"), - ("local_authority/Ards and North Down", "local_authority/N09000011"), - ], -) -def test_resolve_region_maps_routed_uk_area_names_and_codes_to_the_registry( - monkeypatch, requested, expected -): - _routes_uk_area_regions(monkeypatch, routed=True) - country_module = _routed_uk_country_module() - - resolution = _resolve_region( - country_module=country_module, country="uk", params={"region": requested} - ) - - region = country_module.model.get_region(expected) - assert resolution.code == expected - assert resolution.scoping_strategy is region.scoping_strategy - assert resolution.dataset_reference == _LOCAL_AREA_URI - - -@pytest.mark.parametrize( - "requested", - [ - # A district that April 2023 replaced: LAD22 only. - "local_authority/E07000026", - "local_authority/Carlisle", - "constituency/Maidstone", - ], -) -def test_resolve_region_rejects_uk_areas_outside_the_routed_registry( - monkeypatch, requested -): - _routes_uk_area_regions(monkeypatch, routed=True) - - with pytest.raises(ValueError, match="Unsupported UK region"): - _resolve_region( - country_module=_routed_uk_country_module(), - country="uk", - params={"region": requested}, - ) - - -@pytest.mark.parametrize( - "requested", ["constituency/Aldershot", "local_authority/E07000110"] -) -def test_resolve_region_rejects_uk_area_regions_on_an_unrouted_bundle( - monkeypatch, requested -): - # Without routing, policyengine.py's registry lists no area regions, and - # no weight-matrix region stands in for one: older bundles keep their own - # frozen apps. - _routes_uk_area_regions(monkeypatch, routed=False) - country_module = SimpleNamespace(model=SimpleNamespace(get_region=lambda _: None)) - - with pytest.raises(ValueError, match=f"Unsupported UK region: {requested}"): - _resolve_region( - country_module=country_module, country="uk", params={"region": requested} - ) - - -def test_resolve_dataset_reference_uses_bundle_default(): - bundle = get_country_release_bundle("us") - - assert _resolve_dataset_reference("us") == bundle.default_dataset - - -def test_dataset_selection_identifies_the_actual_bundle_default(): - bundle = get_country_release_bundle("us") - nondefault = next( - name for name in bundle.dataset_uris if name != bundle.default_dataset - ) - - omitted = _resolve_dataset_selection({"country": "us"}) - alternate = _resolve_dataset_selection( - {"country": "us"}, - region_resolution=RegionResolution( - code="state/ca", dataset_reference=bundle.dataset_uris[nondefault] - ), - ) - - assert isinstance(omitted, DatasetSelection) - assert omitted.name == bundle.default_dataset - assert omitted.is_default - assert alternate.name == nondefault - assert not alternate.is_default - - -def test_dataset_selection_rejects_a_bundled_uri_as_public_data(): - bundle = get_country_release_bundle("us") - with pytest.raises(ValueError, match="Dataset overrides are not supported"): - _resolve_dataset_selection( - {"country": "us", "data": bundle.default_dataset_uri} - ) - - -@pytest.mark.parametrize("country", ["us", "uk"]) -def test_load_dataset_delegates_default_selection_to_country_loader_with_receipt( - country, - tmp_path, - monkeypatch, -): - bundle = get_country_release_bundle(country) - installed_dataset_path = tmp_path / f"{bundle.default_dataset}.h5" - installed_dataset_path.write_bytes(b"data") - receipt_path = tmp_path / BUNDLE_RECEIPT_FILENAME - receipt_path.write_text( - json.dumps( - { - "bundle_version": "4.18.3", - "policyengine_version": "4.18.3", - "datasets": [ - { - "country": country, - "dataset": bundle.default_dataset, - "version": bundle.data_version, - "path": str(installed_dataset_path), - } - ], - } - ), - encoding="utf-8", - ) - monkeypatch.setenv("POLICYENGINE_BUNDLE_RECEIPT", str(receipt_path)) - monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) - get_country_release_bundle.cache_clear() - - dataset = SimpleNamespace() - ensure_calls = [] - - def ensure_datasets(**kwargs): - ensure_calls.append(kwargs) - return {"dataset": dataset} - - assert ( - _load_dataset( - {"country": country, "time_period": "2026"}, - country_module=SimpleNamespace(ensure_datasets=ensure_datasets), - ) - is dataset - ) - - assert ensure_calls == [ - { - "years": [2026], - "data_folder": str(tmp_path), - } - ] - - -@pytest.mark.parametrize( - "explicit_data_params", - [ - {"data_version": "custom-v1"}, - {"data": "populace_us_2024@custom-v1"}, - {"data": "hf://other-org/other-repo/populace_us_2024.h5"}, - ], -) -def test_load_dataset_rejects_unmanaged_dataset_requests( - explicit_data_params, - tmp_path, - monkeypatch, -): - monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) - - ensure_calls = [] - - def ensure_datasets(**kwargs): - ensure_calls.append(kwargs) - return {"dataset": SimpleNamespace()} - - with pytest.raises(ValueError, match="Dataset overrides are not supported"): - _load_dataset( - {"country": "us", "time_period": "2026", **explicit_data_params}, - country_module=SimpleNamespace(ensure_datasets=ensure_datasets), - ) - - assert ensure_calls == [] - - -def test_load_dataset_rejects_explicit_default(tmp_path, monkeypatch): - monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) - ensure_calls = [] - - def ensure_datasets(**kwargs): - ensure_calls.append(kwargs) - return {"dataset": SimpleNamespace()} - - with pytest.raises(ValueError, match="Dataset overrides are not supported"): - _load_dataset( - {"country": "us", "time_period": "2026", "data": "populace_us_2024"}, - country_module=SimpleNamespace(ensure_datasets=ensure_datasets), - ) - - assert ensure_calls == [] - - -def test_load_dataset_passes_region_dataset_name_to_managed_loader( - tmp_path, monkeypatch -): - monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) - bundle = get_country_release_bundle("us") - nondefault = next( - name for name in bundle.dataset_uris if name != bundle.default_dataset - ) - ensure_calls = [] - - def ensure_datasets(**kwargs): - ensure_calls.append(kwargs) - return {"dataset": SimpleNamespace()} - - _load_dataset( - {"country": "us", "time_period": "2026"}, - country_module=SimpleNamespace(ensure_datasets=ensure_datasets), - region_resolution=RegionResolution( - code="state/ca", dataset_reference=bundle.dataset_uris[nondefault] - ), - ) - - assert ensure_calls == [ - { - "datasets": [nondefault], - "years": [2026], - "data_folder": _nondefault_data_folder( - "us", nondefault, bundle.dataset_uris[nondefault] - ), - } - ] - - -@pytest.mark.parametrize("use_default", [True, False]) -def test_load_dataset_maps_region_bundle_uri_to_managed_name( - use_default, tmp_path, monkeypatch -): - monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) - bundle = get_country_release_bundle("us") - name = ( - bundle.default_dataset - if use_default - else next( - name for name in bundle.dataset_uris if name != bundle.default_dataset - ) - ) - ensure_calls = [] - - def ensure_datasets(**kwargs): - ensure_calls.append(kwargs) - return {"dataset": SimpleNamespace()} - - _load_dataset( - {"country": "us", "time_period": "2026"}, - country_module=SimpleNamespace(ensure_datasets=ensure_datasets), - region_resolution=SimpleNamespace(dataset_reference=bundle.dataset_uris[name]), - ) - - assert ensure_calls == [ - { - **({} if use_default else {"datasets": [name]}), - "years": [2026], - "data_folder": ( - str(tmp_path) - if use_default - else _nondefault_data_folder("us", name, bundle.dataset_uris[name]) - ), - } - ] - - -def test_load_dataset_rejects_explicit_dataset_for_filtered_region( - tmp_path, monkeypatch -): - monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) - bundle = get_country_release_bundle("us") - nondefault = next( - name for name in bundle.dataset_uris if name != bundle.default_dataset - ) - ensure_calls = [] - - def ensure_datasets(**kwargs): - ensure_calls.append(kwargs) - return {"dataset": SimpleNamespace()} - - with pytest.raises(ValueError, match="Dataset overrides are not supported"): - _load_dataset( - {"country": "us", "time_period": "2026", "data": nondefault}, - country_module=SimpleNamespace(ensure_datasets=ensure_datasets), - region_resolution=SimpleNamespace( - dataset_reference=bundle.default_dataset_uri, - scoping_strategy=object(), - ), - ) - - assert ensure_calls == [] - - -def test_nondefault_dataset_cache_is_distinct_for_shared_filename_stem(): - first = _nondefault_data_folder( - "us", "certified_a", "hf://policyengine/a/shared.h5@revision-a" - ) - second = _nondefault_data_folder( - "us", "certified_b", "hf://policyengine/b/shared.h5@revision-b" - ) - - assert first != second - - -def test_load_dataset_reads_a_nondefault_dataset_from_the_folder_it_is_given( - tmp_path, - monkeypatch, -): - monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path / "image-data")) - bundle = get_country_release_bundle("us") - nondefault = next( - name for name in bundle.dataset_uris if name != bundle.default_dataset - ) - selection = DatasetSelection( - name=nondefault, uri=bundle.dataset_uris[nondefault], is_default=False - ) - state_folder = str(tmp_path / "state-datasets" / nondefault / "UT") - ensure_calls = [] - - def ensure_datasets(**kwargs): - ensure_calls.append(kwargs) - return {"dataset": SimpleNamespace()} - - _load_selected_dataset( - {"country": "us", "time_period": "2026"}, - selection=selection, - country_module=SimpleNamespace(ensure_datasets=ensure_datasets), - data_folder=state_folder, - ) - - assert ensure_calls == [ - {"datasets": [nondefault], "years": [2026], "data_folder": state_folder} - ] - - -def test_load_dataset_without_a_folder_keeps_the_temporary_nondefault_folder( - tmp_path, - monkeypatch, -): - monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) - bundle = get_country_release_bundle("us") - nondefault = next( - name for name in bundle.dataset_uris if name != bundle.default_dataset - ) - selection = DatasetSelection( - name=nondefault, uri=bundle.dataset_uris[nondefault], is_default=False - ) - ensure_calls = [] - - def ensure_datasets(**kwargs): - ensure_calls.append(kwargs) - return {"dataset": SimpleNamespace()} - - _load_selected_dataset( - {"country": "us", "time_period": "2026"}, - selection=selection, - country_module=SimpleNamespace(ensure_datasets=ensure_datasets), - data_folder=None, - ) - - assert ensure_calls[0]["data_folder"] == _nondefault_data_folder( - "us", nondefault, bundle.dataset_uris[nondefault] - ) - - -def test_load_dataset_uses_baked_folder_for_default_requests( - tmp_path, - monkeypatch, -): - monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) - - ensure_calls = [] - - def ensure_datasets(**kwargs): - ensure_calls.append(kwargs) - return {"dataset": SimpleNamespace()} - - _load_dataset( - {"country": "us", "time_period": "2026"}, - country_module=SimpleNamespace(ensure_datasets=ensure_datasets), - ) - - assert ensure_calls[0]["data_folder"] == str(tmp_path) - - -def test_resolve_region_scopes_us_state_from_national_dataset(): - bundle = get_country_release_bundle("us") - scoping_strategy = object() - state = SimpleNamespace( - dataset_path=None, - scoping_strategy=scoping_strategy, - parent_code="us", - ) - national = SimpleNamespace( - dataset_path=bundle.default_dataset_uri, - scoping_strategy=None, - parent_code=None, - ) - regions = {"state/ut": state, "us": national} - country_module = SimpleNamespace(model=SimpleNamespace(get_region=regions.get)) - - resolution = _resolve_region( - country_module=country_module, - country="us", - params={"region": "state/UT"}, - ) - - assert resolution.code == "state/ut" - assert resolution.dataset_reference == resolve_runtime_bundle_dataset_uri( - "us", None, prefer_local=False - ) - assert resolution.scoping_strategy is scoping_strategy - - -def test_resolve_region_scopes_us_congressional_district_from_national_dataset(): - bundle = get_country_release_bundle("us") - scoping_strategy = object() - district = SimpleNamespace( - dataset_path=None, - scoping_strategy=scoping_strategy, - parent_code="state/ut", - ) - state = SimpleNamespace( - dataset_path=None, - scoping_strategy=object(), - parent_code="us", - ) - national = SimpleNamespace( - dataset_path=bundle.default_dataset_uri, - scoping_strategy=None, - parent_code=None, - ) - regions = { - "congressional_district/UT-01": district, - "state/ut": state, - "us": national, - } - country_module = SimpleNamespace(model=SimpleNamespace(get_region=regions.get)) - - resolution = _resolve_region( - country_module=country_module, - country="us", - params={"region": "congressional_district/ut-01"}, - ) - - assert resolution.code == "congressional_district/UT-01" - assert resolution.dataset_reference == resolve_runtime_bundle_dataset_uri( - "us", None, prefer_local=False - ) - assert resolution.scoping_strategy is scoping_strategy - - -def _us_country_module_with_regions(regions): - return SimpleNamespace(model=SimpleNamespace(get_region=regions.get)) - - -@pytest.mark.parametrize( - ("requested_region", "state_code"), - [ - ("state/ut", "UT"), - ("state/UT", "UT"), - ("ut", "UT"), - ("congressional_district/ut-01", "UT"), - ("congressional_district/WY-00", "WY"), - ], -) -def test_resolve_region_reports_the_us_state_code(requested_region, state_code): - utah = SimpleNamespace( - dataset_path=None, - scoping_strategy=object(), - parent_code="us", - state_code="UT", - ) - regions = { - "us": SimpleNamespace( - dataset_path=get_country_release_bundle("us").default_dataset_uri, - scoping_strategy=None, - parent_code=None, - state_code=None, - ), - "state/ut": utah, - "congressional_district/UT-01": SimpleNamespace( - dataset_path=None, - scoping_strategy=object(), - parent_code="state/ut", - state_code="UT", - ), - "congressional_district/WY-00": SimpleNamespace( - dataset_path=None, - scoping_strategy=object(), - parent_code="state/wy", - state_code="WY", - ), - "state/wy": SimpleNamespace( - dataset_path=None, - scoping_strategy=object(), - parent_code="us", - state_code="WY", - ), - } - - resolution = _resolve_region( - country_module=_us_country_module_with_regions(regions), - country="us", - params={"region": requested_region}, - ) - - assert resolution.state_code == state_code - - -def test_resolve_region_reports_no_state_code_for_a_national_us_run(): - resolution = _resolve_region( - country_module=_us_country_module_with_regions({}), - country="us", - params={"region": "us"}, - ) - - assert resolution.state_code is None - - -def test_resolve_region_reports_no_state_code_for_a_uk_region(): - scotland = SimpleNamespace( - dataset_path=None, - scoping_strategy=object(), - parent_code="uk", - state_code=None, - ) - national = SimpleNamespace( - dataset_path=get_country_release_bundle("uk").default_dataset_uri, - scoping_strategy=None, - parent_code=None, - state_code=None, - ) - country_module = SimpleNamespace( - model=SimpleNamespace( - get_region={"country/scotland": scotland, "uk": national}.get - ) - ) - - resolution = _resolve_region( - country_module=country_module, - country="uk", - params={"region": "country/scotland"}, - ) - - assert resolution.state_code is None - - -def test_resolve_region_reports_no_state_code_for_a_region_group(): - from policyengine.core.scoping_strategy import RowFilterStrategy - - regions = { - code: SimpleNamespace( - scoping_strategy=RowFilterStrategy( - variable_name="state_fips", variable_value=fips - ), - state_code=code.removeprefix("state/").upper(), - ) - for code, fips in (("state/ut", 49), ("state/wy", 56)) - } - - resolution = _resolve_region( - country_module=_us_country_module_with_regions(regions), - country="us", - params={"region": "us", "region_group": ["state/ut", "state/wy"]}, - ) - - assert resolution.state_code is None - - -def test_every_installed_us_state_and_district_region_names_its_state(): - from policyengine.countries.us.regions import us_region_registry - - for region_type in ("state", "congressional_district"): - regions = us_region_registry.get_by_type(region_type) - assert regions - for region in regions: - assert region.state_code - assert region.state_code == region.state_code.upper() - - -def test_resolve_region_rejects_unscoped_us_place_region(): - place = SimpleNamespace( - dataset_path=None, - scoping_strategy=None, - parent_code="state/ut", - ) - state = SimpleNamespace( - dataset_path=None, - scoping_strategy=object(), - parent_code="us", - ) - regions = {"place/UT-67000": place, "state/ut": state} - country_module = SimpleNamespace(model=SimpleNamespace(get_region=regions.get)) - - with pytest.raises(ValueError, match="US place regions are not yet supported"): - _resolve_region( - country_module=country_module, - country="us", - params={"region": "place/ut-67000"}, - ) - - -def test_resolve_region_scopes_uk_country_from_national_dataset(): - bundle = get_country_release_bundle("uk") - scoping_strategy = object() - england = SimpleNamespace( - dataset_path=None, - scoping_strategy=scoping_strategy, - parent_code="uk", - ) - uk = SimpleNamespace( - dataset_path=bundle.default_dataset_uri, - scoping_strategy=None, - parent_code=None, - ) - regions = {"country/england": england, "uk": uk} - country_module = SimpleNamespace(model=SimpleNamespace(get_region=regions.get)) - - resolution = _resolve_region( - country_module=country_module, - country="uk", - params={"region": "England"}, - ) - - assert resolution.code == "country/england" - assert resolution.dataset_reference == resolve_runtime_bundle_dataset_uri( - "uk", None, prefer_local=False - ) - assert resolution.scoping_strategy is scoping_strategy - - -def test_builder_data_version_prefers_resolved_revision_then_dataset_metadata(): - baseline, reform = _macro_baseline_reform() - - def economic_impact_analysis( - _baseline_simulation, - _reform_simulation, - *, - include_cliff_impacts=False, - ): - return fake_analysis() - - country_module = SimpleNamespace( - model=SimpleNamespace(version="1.715.2"), - economic_impact_analysis=economic_impact_analysis, - ) - - resolved_builder = SimulationOutputBuilder( - country="us", - simulation_params={"country": "us"}, - country_module=country_module, - dataset=SimpleNamespace(metadata={"version": "metadata-version"}), - baseline=baseline, - reform=reform, - resolved_data_version="1.77.0", - ) - metadata_builder = SimulationOutputBuilder( - country="us", - simulation_params={"country": "us"}, - country_module=country_module, - dataset=SimpleNamespace(metadata={"version": "metadata-version"}), - baseline=baseline, - reform=reform, - ) - - assert resolved_builder._data_version() == "1.77.0" - assert metadata_builder._data_version() == "metadata-version" - - -def test_builder_budgetary_impact_uses_materialized_columns_and_uk_state_tax_zero(): - baseline = _FakeSimulation( - weighted_entity_frame( - { - "household_weight": [1.0, 2.0], - "household_net_income": [100.0, 200.0], - "household_tax": [20.0, 40.0], - "household_benefits": [5.0, 10.0], - "household_state_income_tax": [2.0, 3.0], - } - ) - ) - reform = _FakeSimulation( - weighted_entity_frame( - { - "household_weight": [1.0, 2.0], - "household_net_income": [120.0, 210.0], - "household_tax": [25.0, 50.0], - "household_benefits": [4.0, 8.0], - "household_state_income_tax": [4.0, 6.0], - } - ) - ) - - us_budget = _simulation_output_builder( - "us", baseline, reform - )._build_budgetary_impact() - uk_budget = _simulation_output_builder( - "uk", baseline, reform - )._build_budgetary_impact() - - # Changes are weighted sums, weight * (reform - baseline), over the - # weights 1 and 2; the household count is a plain sum of the weights. - tax_revenue_impact = 1.0 * (25.0 - 20.0) + 2.0 * (50.0 - 40.0) - state_tax_revenue_impact = 1.0 * (4.0 - 2.0) + 2.0 * (6.0 - 3.0) - benefit_spending_impact = 1.0 * (4.0 - 5.0) + 2.0 * (8.0 - 10.0) - assert isinstance(us_budget, BudgetaryImpact) - assert us_budget.model_dump(mode="json") == { - "tax_revenue_impact": tax_revenue_impact, - "state_tax_revenue_impact": state_tax_revenue_impact, - "benefit_spending_impact": benefit_spending_impact, - "budgetary_impact": tax_revenue_impact - benefit_spending_impact, - "households": 1.0 + 2.0, - "baseline_net_income": 1.0 * 100.0 + 2.0 * 200.0, - } - assert uk_budget.state_tax_revenue_impact == 0.0 - - -def test_builder_budgetary_impact_propagates_required_calculation_errors(monkeypatch): - baseline, reform = _macro_baseline_reform() - - def fail_change_output_variable(*_args, **_kwargs): - msg = "household_tax missing" - raise RuntimeError(msg) - - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_output_budget._change_output_variable", - fail_change_output_variable, - ) - - with pytest.raises(RuntimeError, match="household_tax missing"): - _simulation_output_builder("us", baseline, reform)._build_budgetary_impact() - - -def test_uk_area_impacts_share_one_disclosure_release(monkeypatch): - baseline = object() - reform = object() - disclosure = simulation_output_geographic.UKGeographyDisclosureOutput() - computed: list[tuple[str, object, object, int]] = [] - built: list[tuple[str, object]] = [] - - def compute(country, baseline_simulation, reform_simulation, *, year): - computed.append((country, baseline_simulation, reform_simulation, year)) - return disclosure - - def build_constituencies(country, value): - built.append(("constituency", value)) - return GeographicImpactOutput.model_validate([]) - - def build_authorities(country, value, **kwargs): - built.append(("local_authority", value)) - return GeographicImpactOutput.model_validate([]) - - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_output_geographic." - "compute_uk_geography_disclosure_output", - compute, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_output_geographic." - "build_uk_constituency_impact", - build_constituencies, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_output_geographic." - "build_uk_local_authority_impact", - build_authorities, - ) - - builder = _simulation_output_builder("uk", baseline, reform) - builder.simulation_params["time_period"] = "2026" - builder._build_uk_constituency_impact() - builder._build_uk_local_authority_impact() - - # Both levels come from one release, computed once. - assert computed == [("uk", baseline, reform, 2026)] - assert built == [("constituency", disclosure), ("local_authority", disclosure)] - - -def test_uk_disclosure_is_not_computed_for_other_countries(monkeypatch): - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_output_geographic." - "compute_uk_geography_disclosure_output", - lambda *_args, **_kwargs: pytest.fail("US runs need no UK disclosure"), - ) - builder = _simulation_output_builder("us", object(), object()) - - assert builder._build_uk_constituency_impact() is None - assert builder._build_uk_local_authority_impact() is None - - -def test_uk_local_authority_impact_receives_detected_boundary_metadata(monkeypatch): - expected = [_local_authority_impact_record().model_dump(mode="json")] - metadata = UKLocalAuthorityMetadata( - boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 - ) - disclosure = simulation_output_geographic.UKGeographyDisclosureOutput() - calls: list[tuple[object, UKLocalAuthorityMetadata | None]] = [] - - def build(country, value, **kwargs): - if country != "uk": - return None - calls.append((value, kwargs["uk_local_authority_metadata"])) - return GeographicImpactOutput.model_validate(expected) - - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_output_geographic." - "compute_uk_geography_disclosure_output", - lambda *_args, **_kwargs: disclosure, - ) - monkeypatch.setattr( - "policyengine_simulation_executor.simulation_output_geographic." - "build_uk_local_authority_impact", - build, - ) - - result = _simulation_output_builder( - "uk", - object(), - object(), - BuilderOptions(uk_local_authority_metadata=metadata), - )._build_uk_local_authority_impact() - assert result is not None - assert result.model_dump(mode="json") == expected - assert calls == [(disclosure, metadata)] - - -def test_omitted_uk_area_breakdowns_compute_nothing(monkeypatch): - for name in ( - "compute_uk_geography_disclosure_output", - "build_uk_constituency_impact", - "build_uk_local_authority_impact", - ): - monkeypatch.setattr( - f"policyengine_simulation_executor.simulation_output_geographic.{name}", - lambda *_args, **_kwargs: pytest.fail("no area breakdown is built"), - ) - builder = _simulation_output_builder( - "uk", - object(), - object(), - BuilderOptions(resolved_region_code="uk", omit_uk_area_breakdowns=True), - ) - - assert builder._build_uk_constituency_impact() is None - assert builder._build_uk_local_authority_impact() is None - - -class _BreakdownSimulation: - """A local-area simulation that records when it is computed.""" - - def __init__(self, label, events): - self.label = label - self.events = events - - def ensure(self): - self.events.append(self.label) - - -class _RecordingProgress: - def __init__(self, events): - self.events = events - - def enter(self, stage): - self.events.append(stage) - - -def _breakdown_builder(monkeypatch, *, events, runtime=None): - """Return a UK breakdown builder whose area outputs are faked. - - The country module's impact analysis fails the test if it is called. - Each fake records into ``events``. - """ - disclosure = simulation_output_geographic.UKGeographyDisclosureOutput() - - def compute(country, baseline_simulation, reform_simulation, *, year): - events.append( - ("release", country, baseline_simulation, reform_simulation, year) - ) - return disclosure - - def constituencies(country, value): - events.append(("constituency", value)) - return GeographicImpactOutput.model_validate( - [_constituency_impact_record().model_dump(mode="json")] - ) - - def local_authorities(country, value, **kwargs): - events.append(("local_authority", value, kwargs["uk_local_authority_metadata"])) - return GeographicImpactOutput.model_validate([]) - - for name, function in { - "compute_uk_geography_disclosure_output": compute, - "build_uk_constituency_impact": constituencies, - "build_uk_local_authority_impact": local_authorities, - }.items(): - monkeypatch.setattr( - f"policyengine_simulation_executor.simulation_output_geographic.{name}", - function, - ) - - def no_impact_analysis(*_args, **_kwargs): - pytest.fail("a local-area breakdown computes no national impact analysis") - - builder = LocalAreaBreakdownBuilder( - baseline=_BreakdownSimulation("baseline", events), - reform=_BreakdownSimulation("reform", events), - local_authority_metadata=UKLocalAuthorityMetadata( - boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 - ), - dataset_name="microcosm_uk_2024_25_dense", - data_version=_LOCAL_AREA_REVISION, - year=2026, - country_module=SimpleNamespace( - model=SimpleNamespace(version="2.132.6"), - economic_impact_analysis=no_impact_analysis, - ), - runtime=runtime, - progress=_RecordingProgress(events), - ) - return builder, disclosure - - -def test_a_local_area_breakdown_publishes_both_levels_from_one_release(monkeypatch): - events: list[object] = [] - builder, disclosure = _breakdown_builder(monkeypatch, events=events) - - output = builder.serialize() - - releases = [event for event in events if event[0] == "release"] - assert releases == [("release", "uk", builder.baseline, builder.reform, 2026)] - assert [event for event in events if event[0] != "release"][-2:] == [ - ("constituency", disclosure), - ("local_authority", disclosure, builder.local_authority_metadata), - ] - assert output == { - "kind": "local_areas", - "model_version": "2.132.6", - "policyengine_version": get_country_release_bundle("uk").policyengine_version, - "dataset": "microcosm_uk_2024_25_dense", - "data_version": _LOCAL_AREA_REVISION, - "constituency_impact": [_constituency_impact_record().model_dump(mode="json")], - "local_authority_impact": [], - } - - -def test_a_local_area_breakdown_enters_outputs_after_both_simulations(monkeypatch): - events: list[object] = [] - builder, _ = _breakdown_builder(monkeypatch, events=events) - - builder.serialize() - - simulations_and_stages = [event for event in events if not isinstance(event, tuple)] - assert simulations_and_stages == ["baseline", "reform", JobStage.OUTPUTS] - assert events.index(JobStage.OUTPUTS) < next( - index for index, event in enumerate(events) if isinstance(event, tuple) - ) - - -def test_a_local_area_breakdown_opens_only_registered_spans(monkeypatch): - events: list[object] = [] - runtime = _TrackingRuntime() - builder, _ = _breakdown_builder(monkeypatch, events=events, runtime=runtime) - - builder.serialize() - - registered = {stage.value for stage in ANNUAL_IMPACT_STAGES.stages} - assert set(runtime.names) <= registered - assert ANNUAL_IMPACT_STAGES.name(Stage.ECONOMIC_IMPACT_ANALYSIS) not in ( - runtime.names - ) - (calculation,) = runtime.roots[:1] - assert calculation["name"] == ANNUAL_IMPACT_STAGES.name(Stage.CALCULATION) - (output_build,) = calculation["children"] - assert output_build["name"] == ANNUAL_IMPACT_STAGES.name( - Stage.SIMULATION_OUTPUT_BUILD - ) - assert { - ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_CONSTITUENCY), - ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_LOCAL_AUTHORITY), - } <= {child["name"] for child in output_build["children"]} diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_dataset_loading.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_dataset_loading.py new file mode 100644 index 000000000..27b498778 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_dataset_loading.py @@ -0,0 +1,242 @@ +"""Tests for the dataset a runtime run loads, and the data version it reports.""" + +from types import SimpleNamespace + +import pytest +from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT +from fixtures.spm_doubles import installed_spm_selection, spm_capable_simulation +from policyengine_simulation_executor.release_bundle import get_country_release_bundle +from policyengine_simulation_executor.simulation_runtime import ( + DatasetSelection, + RegionResolution, + _run_simulation_impl_core, + resolve_data_folder, +) +from test_simulation_output_fakes import TrackingRuntime + + +class _StoppedAfterLoadingError(Exception): + """Raised by a stubbed simulation build once the dataset has loaded.""" + + +def _load_through_runtime( + monkeypatch, *, region_resolution, selection, baked_folder, time_period=2026 +): + """Run the annual runtime up to its dataset load and return what it did.""" + from policyengine_simulation_executor import state_datasets + + load_calls = [] + lookups = [] + + def fake_load_dataset(_params, *, country_module, selection, data_folder=None): + load_calls.append({"selection": selection, "data_folder": data_folder}) + return object() + + def fake_baked_state_data_folder(**kwargs): + lookups.append(kwargs) + return baked_folder + + def stop_after_loading(*_args, **_kwargs): + raise _StoppedAfterLoadingError + + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._country_module", + lambda _country: SimpleNamespace(model=SimpleNamespace(version="2.2.1")), + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.resolve_region", + lambda **_kwargs: region_resolution, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._resolve_dataset_selection", + lambda _params, *, region_resolution: selection, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._load_dataset", + fake_load_dataset, + ) + monkeypatch.setattr( + state_datasets, "baked_state_data_folder", fake_baked_state_data_folder + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._build_simulation", + stop_after_loading, + ) + runtime = TrackingRuntime() + with pytest.raises(_StoppedAfterLoadingError): + _run_simulation_impl_core( + { + "country": "us", + "region": region_resolution.code, + "time_period": time_period, + "baseline": {}, + "reform": {}, + }, + runtime=runtime, + ) + return SimpleNamespace( + load_calls=load_calls, lookups=lookups, context=dict(runtime.context) + ) + + +def _regional_selection(): + return DatasetSelection( + name="populace_us_2024_acs_local", + uri="hf://policyengine/populace-us/populace_us_2024_acs_local.h5@rev", + is_default=False, + ) + + +def test_runtime_loads_a_state_from_its_baked_folder(monkeypatch): + loaded = _load_through_runtime( + monkeypatch, + region_resolution=RegionResolution( + code="state/ut", dataset_reference="ref", state_code="UT" + ), + selection=_regional_selection(), + baked_folder="/opt/policyengine/data/state-datasets/acs/UT", + ) + + assert loaded.load_calls == [ + { + "selection": _regional_selection(), + "data_folder": "/opt/policyengine/data/state-datasets/acs/UT", + } + ] + assert loaded.lookups == [ + { + "country": "us", + "selection": _regional_selection(), + "state_code": "UT", + "year": 2026, + "data_folder": resolve_data_folder(), + } + ] + assert loaded.context["dataset_source"] == "baked_state" + + +def test_runtime_prepares_a_state_on_demand_without_a_baked_file(monkeypatch): + loaded = _load_through_runtime( + monkeypatch, + region_resolution=RegionResolution( + code="congressional_district/UT-01", + dataset_reference="ref", + state_code="UT", + ), + selection=_regional_selection(), + baked_folder=None, + time_period=2028, + ) + + assert loaded.load_calls[0]["data_folder"] is None + assert loaded.lookups[0]["year"] == 2028 + assert loaded.context["dataset_source"] == "on_demand" + + +def test_runtime_loads_a_national_run_from_the_default_folder(monkeypatch): + bundle = get_country_release_bundle("us") + loaded = _load_through_runtime( + monkeypatch, + region_resolution=RegionResolution(code="us"), + selection=DatasetSelection( + name=bundle.default_dataset, + uri=bundle.default_dataset_uri, + is_default=True, + ), + baked_folder=None, + ) + + assert loaded.load_calls[0]["data_folder"] is None + assert "dataset_source" not in loaded.context + + +def _builder_data_version(monkeypatch, *, region_resolution, dataset_selection): + """Run the annual runtime and return the data version it gives the builder.""" + from policyengine_simulation_executor import state_datasets + + selection = installed_spm_selection() + simulations = [spm_capable_simulation(selection), spm_capable_simulation(selection)] + builder_calls = [] + + class FakeSimulationOutputBuilder: + def __init__(self, **kwargs): + builder_calls.append(kwargs) + + def serialize(self): + return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) + + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._country_module", + lambda _country: SimpleNamespace(model=SimpleNamespace(version="2.2.1")), + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.resolve_region", + lambda **_kwargs: region_resolution, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._resolve_dataset_selection", + lambda _params, *, region_resolution: dataset_selection, + ) + monkeypatch.setattr( + state_datasets, "baked_state_data_folder", lambda **_kwargs: None + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._load_dataset", + lambda _params, country_module, selection, data_folder=None: object(), + ) + build_iter = iter(simulations) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._build_simulation", + lambda *_args, **_kwargs: next(build_iter), + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", + FakeSimulationOutputBuilder, + ) + + _run_simulation_impl_core( + { + "country": "us", + "region": region_resolution.code, + "baseline": {}, + "reform": {}, + }, + runtime=TrackingRuntime(), + ) + return builder_calls[0]["resolved_data_version"] + + +def test_a_state_run_reports_the_revision_of_the_dataset_it_loads(monkeypatch): + from policyengine.provenance.manifest import get_release_manifest + + bundle = get_country_release_bundle("us") + regional = "populace_us_2024_acs_local" + + data_version = _builder_data_version( + monkeypatch, + region_resolution=RegionResolution( + code="state/ut", dataset_reference="ref", state_code="UT" + ), + dataset_selection=DatasetSelection( + name=regional, uri=bundle.dataset_uris[regional], is_default=False + ), + ) + + assert data_version == get_release_manifest("us").datasets[regional].revision + assert data_version != bundle.data_version + + +def test_a_national_run_keeps_the_release_data_version(monkeypatch): + bundle = get_country_release_bundle("us") + + data_version = _builder_data_version( + monkeypatch, + region_resolution=RegionResolution(code="us"), + dataset_selection=DatasetSelection( + name=bundle.default_dataset, + uri=bundle.default_dataset_uri, + is_default=True, + ), + ) + + assert data_version is None diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_dataset_selection.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_dataset_selection.py new file mode 100644 index 000000000..70d66cf96 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_dataset_selection.py @@ -0,0 +1,354 @@ +"""Tests for selecting and loading a request's dataset.""" + +import json +from types import SimpleNamespace + +import pytest +from policyengine_simulation_executor.release_bundle import ( + BUNDLE_RECEIPT_FILENAME, + get_country_release_bundle, +) +from policyengine_simulation_executor.simulation_runtime import ( + DatasetSelection, + RegionResolution, + _nondefault_data_folder, + _resolve_dataset_selection, + resolve_dataset_reference, +) +from policyengine_simulation_executor.simulation_runtime import ( + _load_dataset as _load_selected_dataset, +) + + +def _load_dataset(params, *, country_module=None, region_resolution=None): + """Exercise dataset resolution and loading together in direct loader tests.""" + selection = _resolve_dataset_selection(params, region_resolution=region_resolution) + return _load_selected_dataset( + params, country_module=country_module, selection=selection + ) + + +def test_resolve_dataset_reference_uses_bundle_default(): + bundle = get_country_release_bundle("us") + + assert resolve_dataset_reference("us") == bundle.default_dataset + + +def test_dataset_selection_identifies_the_actual_bundle_default(): + bundle = get_country_release_bundle("us") + nondefault = next( + name for name in bundle.dataset_uris if name != bundle.default_dataset + ) + + omitted = _resolve_dataset_selection({"country": "us"}) + alternate = _resolve_dataset_selection( + {"country": "us"}, + region_resolution=RegionResolution( + code="state/ca", dataset_reference=bundle.dataset_uris[nondefault] + ), + ) + + assert isinstance(omitted, DatasetSelection) + assert omitted.name == bundle.default_dataset + assert omitted.is_default + assert alternate.name == nondefault + assert not alternate.is_default + + +def test_dataset_selection_rejects_a_bundled_uri_as_public_data(): + bundle = get_country_release_bundle("us") + with pytest.raises(ValueError, match="Dataset overrides are not supported"): + _resolve_dataset_selection( + {"country": "us", "data": bundle.default_dataset_uri} + ) + + +@pytest.mark.parametrize("country", ["us", "uk"]) +def test_load_dataset_delegates_default_selection_to_country_loader_with_receipt( + country, + tmp_path, + monkeypatch, +): + bundle = get_country_release_bundle(country) + installed_dataset_path = tmp_path / f"{bundle.default_dataset}.h5" + installed_dataset_path.write_bytes(b"data") + receipt_path = tmp_path / BUNDLE_RECEIPT_FILENAME + receipt_path.write_text( + json.dumps( + { + "bundle_version": "4.18.3", + "policyengine_version": "4.18.3", + "datasets": [ + { + "country": country, + "dataset": bundle.default_dataset, + "version": bundle.data_version, + "path": str(installed_dataset_path), + } + ], + } + ), + encoding="utf-8", + ) + monkeypatch.setenv("POLICYENGINE_BUNDLE_RECEIPT", str(receipt_path)) + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) + get_country_release_bundle.cache_clear() + + dataset = SimpleNamespace() + ensure_calls = [] + + def ensure_datasets(**kwargs): + ensure_calls.append(kwargs) + return {"dataset": dataset} + + assert ( + _load_dataset( + {"country": country, "time_period": "2026"}, + country_module=SimpleNamespace(ensure_datasets=ensure_datasets), + ) + is dataset + ) + + assert ensure_calls == [ + { + "years": [2026], + "data_folder": str(tmp_path), + } + ] + + +@pytest.mark.parametrize( + "explicit_data_params", + [ + {"data_version": "custom-v1"}, + {"data": "populace_us_2024@custom-v1"}, + {"data": "hf://other-org/other-repo/populace_us_2024.h5"}, + ], +) +def test_load_dataset_rejects_unmanaged_dataset_requests( + explicit_data_params, + tmp_path, + monkeypatch, +): + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) + + ensure_calls = [] + + def ensure_datasets(**kwargs): + ensure_calls.append(kwargs) + return {"dataset": SimpleNamespace()} + + with pytest.raises(ValueError, match="Dataset overrides are not supported"): + _load_dataset( + {"country": "us", "time_period": "2026", **explicit_data_params}, + country_module=SimpleNamespace(ensure_datasets=ensure_datasets), + ) + + assert ensure_calls == [] + + +def test_load_dataset_rejects_explicit_default(tmp_path, monkeypatch): + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) + ensure_calls = [] + + def ensure_datasets(**kwargs): + ensure_calls.append(kwargs) + return {"dataset": SimpleNamespace()} + + with pytest.raises(ValueError, match="Dataset overrides are not supported"): + _load_dataset( + {"country": "us", "time_period": "2026", "data": "populace_us_2024"}, + country_module=SimpleNamespace(ensure_datasets=ensure_datasets), + ) + + assert ensure_calls == [] + + +def test_load_dataset_passes_region_dataset_name_to_managed_loader( + tmp_path, monkeypatch +): + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) + bundle = get_country_release_bundle("us") + nondefault = next( + name for name in bundle.dataset_uris if name != bundle.default_dataset + ) + ensure_calls = [] + + def ensure_datasets(**kwargs): + ensure_calls.append(kwargs) + return {"dataset": SimpleNamespace()} + + _load_dataset( + {"country": "us", "time_period": "2026"}, + country_module=SimpleNamespace(ensure_datasets=ensure_datasets), + region_resolution=RegionResolution( + code="state/ca", dataset_reference=bundle.dataset_uris[nondefault] + ), + ) + + assert ensure_calls == [ + { + "datasets": [nondefault], + "years": [2026], + "data_folder": _nondefault_data_folder( + "us", nondefault, bundle.dataset_uris[nondefault] + ), + } + ] + + +@pytest.mark.parametrize("use_default", [True, False]) +def test_load_dataset_maps_region_bundle_uri_to_managed_name( + use_default, tmp_path, monkeypatch +): + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) + bundle = get_country_release_bundle("us") + name = ( + bundle.default_dataset + if use_default + else next( + name for name in bundle.dataset_uris if name != bundle.default_dataset + ) + ) + ensure_calls = [] + + def ensure_datasets(**kwargs): + ensure_calls.append(kwargs) + return {"dataset": SimpleNamespace()} + + _load_dataset( + {"country": "us", "time_period": "2026"}, + country_module=SimpleNamespace(ensure_datasets=ensure_datasets), + region_resolution=SimpleNamespace(dataset_reference=bundle.dataset_uris[name]), + ) + + assert ensure_calls == [ + { + **({} if use_default else {"datasets": [name]}), + "years": [2026], + "data_folder": ( + str(tmp_path) + if use_default + else _nondefault_data_folder("us", name, bundle.dataset_uris[name]) + ), + } + ] + + +def test_load_dataset_rejects_explicit_dataset_for_filtered_region( + tmp_path, monkeypatch +): + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) + bundle = get_country_release_bundle("us") + nondefault = next( + name for name in bundle.dataset_uris if name != bundle.default_dataset + ) + ensure_calls = [] + + def ensure_datasets(**kwargs): + ensure_calls.append(kwargs) + return {"dataset": SimpleNamespace()} + + with pytest.raises(ValueError, match="Dataset overrides are not supported"): + _load_dataset( + {"country": "us", "time_period": "2026", "data": nondefault}, + country_module=SimpleNamespace(ensure_datasets=ensure_datasets), + region_resolution=SimpleNamespace( + dataset_reference=bundle.default_dataset_uri, + scoping_strategy=object(), + ), + ) + + assert ensure_calls == [] + + +def test_nondefault_dataset_cache_is_distinct_for_shared_filename_stem(): + first = _nondefault_data_folder( + "us", "certified_a", "hf://policyengine/a/shared.h5@revision-a" + ) + second = _nondefault_data_folder( + "us", "certified_b", "hf://policyengine/b/shared.h5@revision-b" + ) + + assert first != second + + +def test_load_dataset_reads_a_nondefault_dataset_from_the_folder_it_is_given( + tmp_path, + monkeypatch, +): + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path / "image-data")) + bundle = get_country_release_bundle("us") + nondefault = next( + name for name in bundle.dataset_uris if name != bundle.default_dataset + ) + selection = DatasetSelection( + name=nondefault, uri=bundle.dataset_uris[nondefault], is_default=False + ) + state_folder = str(tmp_path / "state-datasets" / nondefault / "UT") + ensure_calls = [] + + def ensure_datasets(**kwargs): + ensure_calls.append(kwargs) + return {"dataset": SimpleNamespace()} + + _load_selected_dataset( + {"country": "us", "time_period": "2026"}, + selection=selection, + country_module=SimpleNamespace(ensure_datasets=ensure_datasets), + data_folder=state_folder, + ) + + assert ensure_calls == [ + {"datasets": [nondefault], "years": [2026], "data_folder": state_folder} + ] + + +def test_load_dataset_without_a_folder_keeps_the_temporary_nondefault_folder( + tmp_path, + monkeypatch, +): + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) + bundle = get_country_release_bundle("us") + nondefault = next( + name for name in bundle.dataset_uris if name != bundle.default_dataset + ) + selection = DatasetSelection( + name=nondefault, uri=bundle.dataset_uris[nondefault], is_default=False + ) + ensure_calls = [] + + def ensure_datasets(**kwargs): + ensure_calls.append(kwargs) + return {"dataset": SimpleNamespace()} + + _load_selected_dataset( + {"country": "us", "time_period": "2026"}, + selection=selection, + country_module=SimpleNamespace(ensure_datasets=ensure_datasets), + data_folder=None, + ) + + assert ensure_calls[0]["data_folder"] == _nondefault_data_folder( + "us", nondefault, bundle.dataset_uris[nondefault] + ) + + +def test_load_dataset_uses_baked_folder_for_default_requests( + tmp_path, + monkeypatch, +): + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) + + ensure_calls = [] + + def ensure_datasets(**kwargs): + ensure_calls.append(kwargs) + return {"dataset": SimpleNamespace()} + + _load_dataset( + {"country": "us", "time_period": "2026"}, + country_module=SimpleNamespace(ensure_datasets=ensure_datasets), + ) + + assert ensure_calls[0]["data_folder"] == str(tmp_path) diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_fakes.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_fakes.py new file mode 100644 index 000000000..ad4bb8d9b --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_fakes.py @@ -0,0 +1,290 @@ +"""Fake simulations, runtimes and output builders shared by the output tests. + +It holds no tests. Its name matches tests/test_simulation_output_*.py, the files +sim.microdata.builder-tests-use-weighted-frames checks, so the entity frames it +builds stay checked. +""" + +from collections import Counter +from contextlib import AbstractContextManager +from dataclasses import dataclass +from types import SimpleNamespace + +from fixtures.test_simulation_output_builder import ( + BASELINE_POVERTY_BY_AGE, + BASELINE_POVERTY_BY_GENDER, + BASELINE_POVERTY_BY_RACE, + INTRA_DECILE_COLLECTION, + REFORM_POVERTY_BY_AGE, + REFORM_POVERTY_BY_GENDER, + REFORM_POVERTY_BY_RACE, + fake_analysis, +) +from fixtures.weighted_frames import weighted_entity_frame +from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata +from policyengine_simulation_executor.simulation_macro_output import ( + CongressionalDistrictImpactOutput, + ConstituencyImpactRecord, + LocalAuthorityImpactRecord, + SingleYearMacroOutput, +) +from policyengine_simulation_executor.simulation_output_builder import ( + SimulationOutputBuilder, +) + + +class FakeOutputDataset: + def __init__(self, household): + self.data = SimpleNamespace(household=household) + + +class FakeSimulation: + def __init__(self, household): + self.output_dataset = FakeOutputDataset(household) + + def ensure(self): + msg = "test data is already materialized" + raise AssertionError(msg) + + +class TrackedSpan(AbstractContextManager): + def __init__(self, runtime, name, attributes): + self.runtime = runtime + self.node = { + "name": name, + "attrs": dict(attributes or {}), + "children": [], + } + + def __enter__(self): + target = ( + self.runtime.stack[-1]["children"] + if self.runtime.stack + else self.runtime.roots + ) + target.append(self.node) + self.runtime.stack.append(self.node) + self.runtime.names.append(self.node["name"]) + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.runtime.stack.pop() + + +class TrackingRuntime: + def __init__(self): + self.roots = [] + self.stack = [] + self.names = [] + self.context = [] + + def span(self, name, *, attributes=None): + return TrackedSpan(self, name, attributes) + + def set_context(self, **attributes): + self.context.extend(attributes.items()) + + +def with_observability_timings(callback): + runtime = TrackingRuntime() + result = callback(runtime) + counts = Counter(runtime.names) + return result, set(runtime.names), counts, runtime.roots + + +def macro_baseline_reform(): + """Build households whose weighted totals are the golden macro result's budget. + + With weights 0.5 and 1.5, the weighted totals are a tax change of 100, a + benefit change of 30, a state tax change of 20, 2 households and a baseline + net income of 1000. Unweighted sums of the same rows would give 120, 40, 24 + and 1200, and summing the weights through the weighted frame (#767) would + give 2.5 households, so the golden result only matches a correct build. + """ + baseline = FakeSimulation( + weighted_entity_frame( + { + "household_weight": [0.5, 1.5], + "household_net_income": [800.0, 400.0], + "household_tax": [60.0, 40.0], + "household_benefits": [20.0, 30.0], + "household_state_income_tax": [5.0, 5.0], + } + ) + ) + reform = FakeSimulation( + weighted_entity_frame( + { + "household_weight": [0.5, 1.5], + "household_net_income": [810.0, 420.0], + "household_tax": [140.0, 80.0], + "household_benefits": [50.0, 40.0], + "household_state_income_tax": [21.0, 13.0], + } + ) + ) + return baseline, reform + + +@dataclass(frozen=True) +class BuilderOptions: + """Optional inputs for a SimulationOutputBuilder built from fakes.""" + + analysis: object = None + include_cliffs: bool | None = None + resolved_region_code: str | None = "us" + uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None + omit_uk_area_breakdowns: bool = False + + +def make_output_builder( + country: str, + baseline, + reform, + options: BuilderOptions | None = None, +) -> SimulationOutputBuilder: + options = options or BuilderOptions() + analysis = options.analysis or fake_analysis() + + def economic_impact_analysis( + _baseline_simulation, + _reform_simulation, + *, + include_cliff_impacts=False, + ): + return analysis + + simulation_params = { + "country": country, + "data_version": "1.115.5" if country == "us" else "1.55.10", + } + if options.include_cliffs is not None: + simulation_params["include_cliffs"] = options.include_cliffs + + country_module = SimpleNamespace( + model=SimpleNamespace(version="1.715.2" if country == "us" else "2.88.20"), + economic_impact_analysis=economic_impact_analysis, + ) + return SimulationOutputBuilder( + country=country, + simulation_params=simulation_params, + country_module=country_module, + dataset=SimpleNamespace(metadata={}), + baseline=baseline, + reform=reform, + resolved_region_code=options.resolved_region_code, + uk_local_authority_metadata=options.uk_local_authority_metadata, + omit_uk_area_breakdowns=options.omit_uk_area_breakdowns, + ) + + +def congressional_district_output() -> CongressionalDistrictImpactOutput: + return CongressionalDistrictImpactOutput( + districts=[ + { + "district": "AL-01", + "average_household_income_change": 10.0, + "relative_household_income_change": 0.01, + "winner_percentage": 0.6, + "loser_percentage": 0.3, + "no_change_percentage": 0.1, + "population": 1000.0, + } + ] + ) + + +def constituency_impact_record() -> ConstituencyImpactRecord: + return ConstituencyImpactRecord( + constituency_code="E14000530", + constituency_name="Birmingham, Ladywood", + x=34, + y=-42, + average_household_income_change=10.0, + relative_household_income_change=0.01, + population=1000.0, + ) + + +def local_authority_impact_record() -> LocalAuthorityImpactRecord: + return LocalAuthorityImpactRecord( + local_authority_code="E06000001", + local_authority_name="Hartlepool", + x=63, + y=-19, + average_household_income_change=12.0, + relative_household_income_change=0.02, + population=900.0, + ) + + +def stub_policyengine_output_calls(monkeypatch, baseline, _reform) -> None: + def fake_poverty_module_function(name): + def compute(simulation): + if "by_age" in name: + return ( + BASELINE_POVERTY_BY_AGE + if simulation is baseline + else REFORM_POVERTY_BY_AGE + ) + if "by_gender" in name: + return ( + BASELINE_POVERTY_BY_GENDER + if simulation is baseline + else REFORM_POVERTY_BY_GENDER + ) + if "by_race" in name: + return ( + BASELINE_POVERTY_BY_RACE + if simulation is baseline + else REFORM_POVERTY_BY_RACE + ) + msg = f"Unexpected poverty output: {name}" + raise AssertionError(msg) + + return compute + + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_output_poverty._poverty_module_function", + fake_poverty_module_function, + ) + monkeypatch.setattr( + SimulationOutputBuilder, + "_build_intra_decile_output", + lambda self: self._build_intra_decile_output_from_collection( + INTRA_DECILE_COLLECTION + ), + ) + monkeypatch.setattr( + SimulationOutputBuilder, + "_build_congressional_district_impact", + lambda self: congressional_district_output() if self.country == "us" else None, + ) + monkeypatch.setattr( + SimulationOutputBuilder, + "_build_uk_constituency_impact", + lambda self: ( + self._build_geographic_impact_output( + [constituency_impact_record().model_dump(mode="json")] + ) + if self.country == "uk" + else None + ), + ) + monkeypatch.setattr( + SimulationOutputBuilder, + "_build_uk_local_authority_impact", + lambda self: ( + self._build_geographic_impact_output( + [local_authority_impact_record().model_dump(mode="json")] + ) + if self.country == "uk" + else None + ), + ) + + +def build_schema_output(monkeypatch, *, country: str = "us") -> SingleYearMacroOutput: + baseline, reform = macro_baseline_reform() + stub_policyengine_output_calls(monkeypatch, baseline, reform) + return make_output_builder(country, baseline, reform).build() diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_impacts.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_impacts.py new file mode 100644 index 000000000..638687417 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_impacts.py @@ -0,0 +1,338 @@ +"""Tests for the economic, cliff, budgetary and UK area impacts the builder reports.""" + +from types import SimpleNamespace + +import pytest +from fixtures.test_simulation_output_builder import FakeModelOutput, fake_analysis +from fixtures.weighted_frames import weighted_entity_frame +from policyengine_simulation_contract.uk_geography import ( + UKLocalAuthorityBoundaryVersion, + UKLocalAuthorityMetadata, +) +from policyengine_simulation_executor import simulation_output_geographic +from policyengine_simulation_executor.simulation_macro_output import ( + BudgetaryImpact, + GeographicImpactOutput, +) +from policyengine_simulation_executor.simulation_output_builder import ( + SimulationOutputBuilder, +) +from policyengine_simulation_executor.simulation_runtime import normalise_policy +from test_simulation_output_fakes import ( + BuilderOptions, + FakeSimulation, + local_authority_impact_record, + macro_baseline_reform, + make_output_builder, + stub_policyengine_output_calls, +) + + +def test_builder_calls_policyengine_economic_impact_analysis(): + baseline, reform = macro_baseline_reform() + analysis = fake_analysis() + calls = [] + + def economic_impact_analysis( + baseline_simulation, + reform_simulation, + *, + include_cliff_impacts=False, + ): + calls.append((baseline_simulation, reform_simulation, include_cliff_impacts)) + return analysis + + country_module = SimpleNamespace( + model=SimpleNamespace(version="1.715.2"), + economic_impact_analysis=economic_impact_analysis, + ) + builder = SimulationOutputBuilder( + country="us", + simulation_params={"country": "us", "data_version": "1.115.5"}, + country_module=country_module, + dataset=SimpleNamespace(metadata={}), + baseline=baseline, + reform=reform, + ) + + assert builder.analysis is analysis + assert builder.analysis is analysis + assert calls == [(baseline, reform, False)] + + +def test_builder_passes_include_cliffs_to_policyengine_economic_impact_analysis(): + baseline, reform = macro_baseline_reform() + analysis = fake_analysis() + calls = [] + + def economic_impact_analysis( + _baseline_simulation, + _reform_simulation, + *, + include_cliff_impacts=False, + ): + calls.append(include_cliff_impacts) + return analysis + + country_module = SimpleNamespace( + model=SimpleNamespace(version="1.715.2"), + economic_impact_analysis=economic_impact_analysis, + ) + builder = SimulationOutputBuilder( + country="us", + simulation_params={ + "country": "us", + "data_version": "1.115.5", + "include_cliffs": True, + }, + country_module=country_module, + dataset=SimpleNamespace(metadata={}), + baseline=baseline, + reform=reform, + ) + + assert builder.analysis is analysis + assert calls == [True] + + +def test_builder_serializes_cliff_impact_when_requested(monkeypatch): + baseline, reform = macro_baseline_reform() + stub_policyengine_output_calls(monkeypatch, baseline, reform) + analysis = fake_analysis() + analysis.cliff_impact = FakeModelOutput( + { + "baseline": {"cliff_gap": 10.0, "cliff_share": 0.25}, + "reform": {"cliff_gap": 20.0, "cliff_share": 0.5}, + } + ) + + output = make_output_builder( + "us", + baseline, + reform, + BuilderOptions(analysis=analysis, include_cliffs=True), + ).build() + + assert output.cliff_impact is not None + assert output.model_dump(mode="json")["cliff_impact"] == { + "baseline": {"cliff_gap": 10.0, "cliff_share": 0.25}, + "reform": {"cliff_gap": 20.0, "cliff_share": 0.5}, + } + + +def test_normalise_policy_preserves_legacy_period_range_keys(): + assert normalise_policy({"gov.test.parameter": {"2026-01-01.2100-12-31": 1}}) == { + "gov.test.parameter": {"2026-01-01.2100-12-31": 1} + } + + +def test_builder_data_version_prefers_resolved_revision_then_dataset_metadata(): + baseline, reform = macro_baseline_reform() + + def economic_impact_analysis( + _baseline_simulation, + _reform_simulation, + *, + include_cliff_impacts=False, + ): + return fake_analysis() + + country_module = SimpleNamespace( + model=SimpleNamespace(version="1.715.2"), + economic_impact_analysis=economic_impact_analysis, + ) + + resolved_builder = SimulationOutputBuilder( + country="us", + simulation_params={"country": "us"}, + country_module=country_module, + dataset=SimpleNamespace(metadata={"version": "metadata-version"}), + baseline=baseline, + reform=reform, + resolved_data_version="1.77.0", + ) + metadata_builder = SimulationOutputBuilder( + country="us", + simulation_params={"country": "us"}, + country_module=country_module, + dataset=SimpleNamespace(metadata={"version": "metadata-version"}), + baseline=baseline, + reform=reform, + ) + + assert resolved_builder._data_version() == "1.77.0" + assert metadata_builder._data_version() == "metadata-version" + + +def test_builder_budgetary_impact_uses_materialized_columns_and_uk_state_tax_zero(): + baseline = FakeSimulation( + weighted_entity_frame( + { + "household_weight": [1.0, 2.0], + "household_net_income": [100.0, 200.0], + "household_tax": [20.0, 40.0], + "household_benefits": [5.0, 10.0], + "household_state_income_tax": [2.0, 3.0], + } + ) + ) + reform = FakeSimulation( + weighted_entity_frame( + { + "household_weight": [1.0, 2.0], + "household_net_income": [120.0, 210.0], + "household_tax": [25.0, 50.0], + "household_benefits": [4.0, 8.0], + "household_state_income_tax": [4.0, 6.0], + } + ) + ) + + us_budget = make_output_builder("us", baseline, reform)._build_budgetary_impact() + uk_budget = make_output_builder("uk", baseline, reform)._build_budgetary_impact() + + # Changes are weighted sums, weight * (reform - baseline), over the + # weights 1 and 2; the household count is a plain sum of the weights. + tax_revenue_impact = 1.0 * (25.0 - 20.0) + 2.0 * (50.0 - 40.0) + state_tax_revenue_impact = 1.0 * (4.0 - 2.0) + 2.0 * (6.0 - 3.0) + benefit_spending_impact = 1.0 * (4.0 - 5.0) + 2.0 * (8.0 - 10.0) + assert isinstance(us_budget, BudgetaryImpact) + assert us_budget.model_dump(mode="json") == { + "tax_revenue_impact": tax_revenue_impact, + "state_tax_revenue_impact": state_tax_revenue_impact, + "benefit_spending_impact": benefit_spending_impact, + "budgetary_impact": tax_revenue_impact - benefit_spending_impact, + "households": 1.0 + 2.0, + "baseline_net_income": 1.0 * 100.0 + 2.0 * 200.0, + } + assert uk_budget.state_tax_revenue_impact == 0.0 + + +def test_builder_budgetary_impact_propagates_required_calculation_errors(monkeypatch): + baseline, reform = macro_baseline_reform() + + def fail_change_output_variable(*_args, **_kwargs): + msg = "household_tax missing" + raise RuntimeError(msg) + + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_output_budget._change_output_variable", + fail_change_output_variable, + ) + + with pytest.raises(RuntimeError, match="household_tax missing"): + make_output_builder("us", baseline, reform)._build_budgetary_impact() + + +def test_uk_area_impacts_share_one_disclosure_release(monkeypatch): + baseline = object() + reform = object() + disclosure = simulation_output_geographic.UKGeographyDisclosureOutput() + computed: list[tuple[str, object, object, int]] = [] + built: list[tuple[str, object]] = [] + + def compute(country, baseline_simulation, reform_simulation, *, year): + computed.append((country, baseline_simulation, reform_simulation, year)) + return disclosure + + def build_constituencies(country, value): + built.append(("constituency", value)) + return GeographicImpactOutput.model_validate([]) + + def build_authorities(country, value, **kwargs): + built.append(("local_authority", value)) + return GeographicImpactOutput.model_validate([]) + + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_output_geographic." + "compute_uk_geography_disclosure_output", + compute, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_output_geographic." + "build_uk_constituency_impact", + build_constituencies, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_output_geographic." + "build_uk_local_authority_impact", + build_authorities, + ) + + builder = make_output_builder("uk", baseline, reform) + builder.simulation_params["time_period"] = "2026" + builder._build_uk_constituency_impact() + builder._build_uk_local_authority_impact() + + # Both levels come from one release, computed once. + assert computed == [("uk", baseline, reform, 2026)] + assert built == [("constituency", disclosure), ("local_authority", disclosure)] + + +def test_uk_disclosure_is_not_computed_for_other_countries(monkeypatch): + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_output_geographic." + "compute_uk_geography_disclosure_output", + lambda *_args, **_kwargs: pytest.fail("US runs need no UK disclosure"), + ) + builder = make_output_builder("us", object(), object()) + + assert builder._build_uk_constituency_impact() is None + assert builder._build_uk_local_authority_impact() is None + + +def test_uk_local_authority_impact_receives_detected_boundary_metadata(monkeypatch): + expected = [local_authority_impact_record().model_dump(mode="json")] + metadata = UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 + ) + disclosure = simulation_output_geographic.UKGeographyDisclosureOutput() + calls: list[tuple[object, UKLocalAuthorityMetadata | None]] = [] + + def build(country, value, **kwargs): + if country != "uk": + return None + calls.append((value, kwargs["uk_local_authority_metadata"])) + return GeographicImpactOutput.model_validate(expected) + + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_output_geographic." + "compute_uk_geography_disclosure_output", + lambda *_args, **_kwargs: disclosure, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_output_geographic." + "build_uk_local_authority_impact", + build, + ) + + result = make_output_builder( + "uk", + object(), + object(), + BuilderOptions(uk_local_authority_metadata=metadata), + )._build_uk_local_authority_impact() + assert result is not None + assert result.model_dump(mode="json") == expected + assert calls == [(disclosure, metadata)] + + +def test_omitted_uk_area_breakdowns_compute_nothing(monkeypatch): + for name in ( + "compute_uk_geography_disclosure_output", + "build_uk_constituency_impact", + "build_uk_local_authority_impact", + ): + monkeypatch.setattr( + f"policyengine_simulation_executor.simulation_output_geographic.{name}", + lambda *_args, **_kwargs: pytest.fail("no area breakdown is built"), + ) + builder = make_output_builder( + "uk", + object(), + object(), + BuilderOptions(resolved_region_code="uk", omit_uk_area_breakdowns=True), + ) + + assert builder._build_uk_constituency_impact() is None + assert builder._build_uk_local_authority_impact() is None diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_region_resolution.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_region_resolution.py new file mode 100644 index 000000000..8a8ab235f --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_region_resolution.py @@ -0,0 +1,257 @@ +"""Tests for how a request's US or UK region resolves to its scoping and dataset.""" + +from types import SimpleNamespace + +import pytest +from policyengine_simulation_executor.release_bundle import ( + get_country_release_bundle, + resolve_runtime_bundle_dataset_uri, +) +from policyengine_simulation_executor.simulation_runtime import resolve_region + + +def test_resolve_region_scopes_us_state_from_national_dataset(): + bundle = get_country_release_bundle("us") + scoping_strategy = object() + state = SimpleNamespace( + dataset_path=None, + scoping_strategy=scoping_strategy, + parent_code="us", + ) + national = SimpleNamespace( + dataset_path=bundle.default_dataset_uri, + scoping_strategy=None, + parent_code=None, + ) + regions = {"state/ut": state, "us": national} + country_module = SimpleNamespace(model=SimpleNamespace(get_region=regions.get)) + + resolution = resolve_region( + country_module=country_module, + country="us", + params={"region": "state/UT"}, + ) + + assert resolution.code == "state/ut" + assert resolution.dataset_reference == resolve_runtime_bundle_dataset_uri( + "us", None, prefer_local=False + ) + assert resolution.scoping_strategy is scoping_strategy + + +def test_resolve_region_scopes_us_congressional_district_from_national_dataset(): + bundle = get_country_release_bundle("us") + scoping_strategy = object() + district = SimpleNamespace( + dataset_path=None, + scoping_strategy=scoping_strategy, + parent_code="state/ut", + ) + state = SimpleNamespace( + dataset_path=None, + scoping_strategy=object(), + parent_code="us", + ) + national = SimpleNamespace( + dataset_path=bundle.default_dataset_uri, + scoping_strategy=None, + parent_code=None, + ) + regions = { + "congressional_district/UT-01": district, + "state/ut": state, + "us": national, + } + country_module = SimpleNamespace(model=SimpleNamespace(get_region=regions.get)) + + resolution = resolve_region( + country_module=country_module, + country="us", + params={"region": "congressional_district/ut-01"}, + ) + + assert resolution.code == "congressional_district/UT-01" + assert resolution.dataset_reference == resolve_runtime_bundle_dataset_uri( + "us", None, prefer_local=False + ) + assert resolution.scoping_strategy is scoping_strategy + + +def _us_country_module_with_regions(regions): + return SimpleNamespace(model=SimpleNamespace(get_region=regions.get)) + + +@pytest.mark.parametrize( + ("requested_region", "state_code"), + [ + ("state/ut", "UT"), + ("state/UT", "UT"), + ("ut", "UT"), + ("congressional_district/ut-01", "UT"), + ("congressional_district/WY-00", "WY"), + ], +) +def test_resolve_region_reports_the_us_state_code(requested_region, state_code): + utah = SimpleNamespace( + dataset_path=None, + scoping_strategy=object(), + parent_code="us", + state_code="UT", + ) + regions = { + "us": SimpleNamespace( + dataset_path=get_country_release_bundle("us").default_dataset_uri, + scoping_strategy=None, + parent_code=None, + state_code=None, + ), + "state/ut": utah, + "congressional_district/UT-01": SimpleNamespace( + dataset_path=None, + scoping_strategy=object(), + parent_code="state/ut", + state_code="UT", + ), + "congressional_district/WY-00": SimpleNamespace( + dataset_path=None, + scoping_strategy=object(), + parent_code="state/wy", + state_code="WY", + ), + "state/wy": SimpleNamespace( + dataset_path=None, + scoping_strategy=object(), + parent_code="us", + state_code="WY", + ), + } + + resolution = resolve_region( + country_module=_us_country_module_with_regions(regions), + country="us", + params={"region": requested_region}, + ) + + assert resolution.state_code == state_code + + +def test_resolve_region_reports_no_state_code_for_a_national_us_run(): + resolution = resolve_region( + country_module=_us_country_module_with_regions({}), + country="us", + params={"region": "us"}, + ) + + assert resolution.state_code is None + + +def test_resolve_region_reports_no_state_code_for_a_uk_region(): + scotland = SimpleNamespace( + dataset_path=None, + scoping_strategy=object(), + parent_code="uk", + state_code=None, + ) + national = SimpleNamespace( + dataset_path=get_country_release_bundle("uk").default_dataset_uri, + scoping_strategy=None, + parent_code=None, + state_code=None, + ) + country_module = SimpleNamespace( + model=SimpleNamespace( + get_region={"country/scotland": scotland, "uk": national}.get + ) + ) + + resolution = resolve_region( + country_module=country_module, + country="uk", + params={"region": "country/scotland"}, + ) + + assert resolution.state_code is None + + +def test_resolve_region_reports_no_state_code_for_a_region_group(): + from policyengine.core.scoping_strategy import RowFilterStrategy + + regions = { + code: SimpleNamespace( + scoping_strategy=RowFilterStrategy( + variable_name="state_fips", variable_value=fips + ), + state_code=code.removeprefix("state/").upper(), + ) + for code, fips in (("state/ut", 49), ("state/wy", 56)) + } + + resolution = resolve_region( + country_module=_us_country_module_with_regions(regions), + country="us", + params={"region": "us", "region_group": ["state/ut", "state/wy"]}, + ) + + assert resolution.state_code is None + + +def test_every_installed_us_state_and_district_region_names_its_state(): + from policyengine.countries.us.regions import us_region_registry + + for region_type in ("state", "congressional_district"): + regions = us_region_registry.get_by_type(region_type) + assert regions + for region in regions: + assert region.state_code + assert region.state_code == region.state_code.upper() + + +def test_resolve_region_rejects_unscoped_us_place_region(): + place = SimpleNamespace( + dataset_path=None, + scoping_strategy=None, + parent_code="state/ut", + ) + state = SimpleNamespace( + dataset_path=None, + scoping_strategy=object(), + parent_code="us", + ) + regions = {"place/UT-67000": place, "state/ut": state} + country_module = SimpleNamespace(model=SimpleNamespace(get_region=regions.get)) + + with pytest.raises(ValueError, match="US place regions are not yet supported"): + resolve_region( + country_module=country_module, + country="us", + params={"region": "place/ut-67000"}, + ) + + +def test_resolve_region_scopes_uk_country_from_national_dataset(): + bundle = get_country_release_bundle("uk") + scoping_strategy = object() + england = SimpleNamespace( + dataset_path=None, + scoping_strategy=scoping_strategy, + parent_code="uk", + ) + uk = SimpleNamespace( + dataset_path=bundle.default_dataset_uri, + scoping_strategy=None, + parent_code=None, + ) + regions = {"country/england": england, "uk": uk} + country_module = SimpleNamespace(model=SimpleNamespace(get_region=regions.get)) + + resolution = resolve_region( + country_module=country_module, + country="uk", + params={"region": "England"}, + ) + + assert resolution.code == "country/england" + assert resolution.dataset_reference == resolve_runtime_bundle_dataset_uri( + "uk", None, prefer_local=False + ) + assert resolution.scoping_strategy is scoping_strategy diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_runtime_runs.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_runtime_runs.py new file mode 100644 index 000000000..a0c1e3c67 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_runtime_runs.py @@ -0,0 +1,419 @@ +"""Tests for running a simulation through the runtime, with its calculation faked.""" + +from types import SimpleNamespace + +from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT +from fixtures.spm_doubles import installed_spm_selection, spm_capable_simulation +from policyengine_simulation_executor.job_progress import NULL_JOB_PROGRESS +from policyengine_simulation_executor.simulation_runtime import ( + RegionResolution, + _run_simulation_impl_core, + run_simulation_impl, +) +from policyengine_simulation_observability.stages import Stage +from test_simulation_output_fakes import TrackingRuntime, with_observability_timings + + +def _params_with_spm(params, selection): + """Return request parameters after the runtime resolves SPM settings.""" + return params if selection is None else {**params, "spm": selection} + + +def _assert_macro_result(result, selection): + """Check the ordinary macro result plus any installed SPM record.""" + 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() + if key not in spm_keys + } + if selection is None: + assert all(result[key] is None for key in spm_keys) + return + assert set(result["spm_provenance"]) == { + "schema_version", + "baseline", + "reform", + } + + +def test_run_simulation_impl_records_runtime_timings_without_real_calculation( + monkeypatch, +): + dataset = object() + country_module = SimpleNamespace(model=SimpleNamespace(version="1.715.2")) + selection = installed_spm_selection() + baseline_simulation = spm_capable_simulation(selection) + reform_simulation = spm_capable_simulation(selection) + build_calls = [] + + def fake_build_simulation( + params, + *, + dataset, + dataset_selection, + policy, + region, + ): + build_calls.append( + (params, dataset, policy, region.scoping_strategy, region.code) + ) + return baseline_simulation if len(build_calls) == 1 else reform_simulation + + class FakeSimulationOutputBuilder: + def __init__(self, **kwargs): + self.kwargs = kwargs + + def serialize(self): + return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) + + monkeypatch.delenv("GOOGLE_APPLICATION_CREDENTIALS", raising=False) + monkeypatch.delenv("GOOGLE_APPLICATION_CREDENTIALS_JSON", raising=False) + monkeypatch.delenv("GCP_CREDENTIALS_JSON", raising=False) + monkeypatch.delenv("GOOGLE_CREDENTIALS", raising=False) + monkeypatch.delenv("SERVICE_ACCOUNT_JSON", raising=False) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._country_module", + lambda _country: country_module, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.resolve_region", + lambda **_kwargs: RegionResolution( + code="us", + dataset_reference=None, + scoping_strategy="mock-scoping", + ), + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._load_dataset", + lambda _params, country_module, selection, data_folder=None: dataset, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._build_simulation", + fake_build_simulation, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", + FakeSimulationOutputBuilder, + ) + + result, timings, counts, segment_tree = with_observability_timings( + lambda runtime: run_simulation_impl( + { + "country": "us", + "baseline": {"gov.test.parameter": {"2026-01-01": 1}}, + "reform": {"gov.test.parameter": {"2026-01-01": 2}}, + }, + runtime=runtime, + ) + ) + + _assert_macro_result(result, selection) + assert set(timings) >= { + Stage.CREDENTIAL_SETUP, + Stage.REQUEST_PARSE, + Stage.COUNTRY_MODULE_LOAD, + Stage.REGION_RESOLUTION, + Stage.DATASET_RESOLUTION, + Stage.DATASET_LOAD, + Stage.POLICY_NORMALIZATION, + Stage.SIMULATION_BUILD, + } + assert counts[Stage.SIMULATION_BUILD] == 2 + assert [node["name"] for node in segment_tree] == [ + Stage.CREDENTIAL_SETUP, + Stage.REQUEST_PARSE, + Stage.COUNTRY_MODULE_LOAD, + Stage.REGION_RESOLUTION, + Stage.DATASET_RESOLUTION, + Stage.DATASET_LOAD, + Stage.POLICY_NORMALIZATION, + Stage.SIMULATION_BUILD, + Stage.SIMULATION_BUILD, + ] + simulation_builds = [ + node for node in segment_tree if node["name"] == Stage.SIMULATION_BUILD + ] + assert [node["attrs"] for node in simulation_builds] == [ + {"simulation_kind": "baseline"}, + {"simulation_kind": "reform"}, + ] + # region_code must be the RESOLVED code: it feeds the deterministic + # baseline-id predicate, and dropping it silently disables artifact + # reuse (every baseline becomes a miss). + built_params = _params_with_spm( + { + "country": "us", + "baseline": {"gov.test.parameter": {"2026-01-01": 1}}, + "reform": {"gov.test.parameter": {"2026-01-01": 2}}, + }, + selection, + ) + assert build_calls == [ + ( + built_params, + dataset, + {"gov.test.parameter": {"2026-01-01": 1}}, + "mock-scoping", + "us", + ), + ( + built_params, + dataset, + {"gov.test.parameter": {"2026-01-01": 2}}, + "mock-scoping", + "us", + ), + ] + + +def test_run_simulation_impl_exports_baseline_artifact_outcome(monkeypatch): + """run_simulation_impl must export the baseline artifact outcome attribute. + + The hit/incomplete/miss attribute is the rollout metric for the artifact + pipeline; losing the export would blind that measurement. + """ + dataset = object() + country_module = SimpleNamespace(model=SimpleNamespace(version="1.715.2")) + selection = installed_spm_selection() + simulations = { + "baseline": spm_capable_simulation(selection, artifact_outcome="incomplete"), + "reform": spm_capable_simulation(selection), + } + build_count = [0] + + def fake_build_simulation( + _params, + *, + dataset, + dataset_selection, + policy, + region, + ): + build_count[0] += 1 + return simulations["baseline"] if build_count[0] == 1 else simulations["reform"] + + class FakeSimulationOutputBuilder: + def __init__(self, **kwargs): + self.kwargs = kwargs + + def serialize(self): + return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) + + for env in ( + "GOOGLE_APPLICATION_CREDENTIALS", + "GOOGLE_APPLICATION_CREDENTIALS_JSON", + "GCP_CREDENTIALS_JSON", + "GOOGLE_CREDENTIALS", + "SERVICE_ACCOUNT_JSON", + ): + monkeypatch.delenv(env, raising=False) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._country_module", + lambda _country: country_module, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.resolve_region", + lambda **_kwargs: RegionResolution( + code="us", + dataset_reference=None, + scoping_strategy="mock-scoping", + ), + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._load_dataset", + lambda _params, country_module, selection, data_folder=None: dataset, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._build_simulation", + fake_build_simulation, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", + FakeSimulationOutputBuilder, + ) + params = { + "country": "us", + "baseline": {"gov.test.parameter": {"2026-01-01": 1}}, + "reform": {"gov.test.parameter": {"2026-01-01": 2}}, + } + runtime = TrackingRuntime() + run_simulation_impl(params, runtime=runtime) + assert ("baseline_artifact", "incomplete") in runtime.context + + # A plain Simulation baseline (no artifact_outcome) must emit nothing. + runtime.context.clear() + build_count[0] = 0 + simulations["baseline"] = spm_capable_simulation(selection) + run_simulation_impl(params, runtime=runtime) + assert all(key != "baseline_artifact" for key, _ in runtime.context) + + +def test_run_simulation_impl_core_builds_and_serializes_macro_output(monkeypatch): + dataset = object() + country_module = SimpleNamespace(model=SimpleNamespace(version="1.715.2")) + selection = installed_spm_selection() + baseline_simulation = spm_capable_simulation(selection) + reform_simulation = spm_capable_simulation(selection) + build_calls = [] + builder_calls = [] + + def fake_country_module(country): + assert country == "us" + return country_module + + def fake_build_simulation( + params, + *, + dataset, + dataset_selection, + policy, + region, + ): + build_calls.append((params, dataset, policy, region.scoping_strategy)) + return baseline_simulation if len(build_calls) == 1 else reform_simulation + + class FakeSimulationOutputBuilder: + def __init__(self, **kwargs): + builder_calls.append(kwargs) + + def serialize(self): + return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) + + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._country_module", + fake_country_module, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.resolve_region", + lambda **_kwargs: RegionResolution(code="us", dataset_reference=None), + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._load_dataset", + lambda _params, country_module, selection, data_folder=None: dataset, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._build_simulation", + fake_build_simulation, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", + FakeSimulationOutputBuilder, + ) + + runtime = TrackingRuntime() + result = _run_simulation_impl_core( + { + "country": "us", + "baseline": {"gov.test.parameter": {"2026-01-01.2100-12-31": 1}}, + "reform": {"gov.test.parameter": {"2026-01-01.2100-12-31": 2}}, + }, + runtime=runtime, + ) + + _assert_macro_result(result, selection) + assert build_calls[0][2] == {"gov.test.parameter": {"2026-01-01.2100-12-31": 1}} + assert build_calls[1][2] == {"gov.test.parameter": {"2026-01-01.2100-12-31": 2}} + assert build_calls[0][3] is None + assert build_calls[1][3] is None + assert builder_calls == [ + { + "country": "us", + "simulation_params": _params_with_spm( + { + "country": "us", + "baseline": {"gov.test.parameter": {"2026-01-01.2100-12-31": 1}}, + "reform": {"gov.test.parameter": {"2026-01-01.2100-12-31": 2}}, + }, + selection, + ), + "country_module": country_module, + "dataset": dataset, + "baseline": baseline_simulation, + "reform": reform_simulation, + "resolved_data_version": None, + "resolved_region_code": "us", + "runtime": runtime, + "uk_local_authority_metadata": None, + "omit_uk_area_breakdowns": False, + "progress": NULL_JOB_PROGRESS, + } + ] + + +def test_run_simulation_impl_core_passes_region_scoping_to_simulations(monkeypatch): + dataset = object() + country_module = SimpleNamespace(model=SimpleNamespace(version="1.715.2")) + selection = installed_spm_selection() + baseline_simulation = spm_capable_simulation(selection) + reform_simulation = spm_capable_simulation(selection) + scoping_strategy = object() + region_resolution = RegionResolution( + code="state/ut", + dataset_reference=None, + scoping_strategy=scoping_strategy, + ) + build_calls = [] + load_selections = [] + + def fake_build_simulation( + params, + *, + dataset, + dataset_selection, + policy, + region, + ): + build_calls.append( + (params, dataset, policy, region.scoping_strategy, dataset_selection) + ) + return baseline_simulation if len(build_calls) == 1 else reform_simulation + + def fake_load_dataset(_params, *, country_module, selection, data_folder=None): + load_selections.append(selection) + return dataset + + class FakeSimulationOutputBuilder: + def __init__(self, **kwargs): + pass + + def serialize(self): + return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) + + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._country_module", + lambda _country: country_module, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.resolve_region", + lambda **_kwargs: region_resolution, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._load_dataset", + fake_load_dataset, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._build_simulation", + fake_build_simulation, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", + FakeSimulationOutputBuilder, + ) + + result = _run_simulation_impl_core( + { + "country": "us", + "region": "state/ut", + "baseline": {}, + "reform": {}, + }, + runtime=TrackingRuntime(), + ) + + _assert_macro_result(result, selection) + assert build_calls[0][3] is scoping_strategy + assert build_calls[1][3] is scoping_strategy + assert len(load_selections) == 1 + assert build_calls[0][4] is load_selections[0] + assert build_calls[1][4] is load_selections[0] diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_uk_local_areas.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_uk_local_areas.py new file mode 100644 index 000000000..e4f0c4096 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_uk_local_areas.py @@ -0,0 +1,802 @@ +"""Tests for UK runs on a bundle that routes area regions, and local-area breakdowns.""" + +from dataclasses import dataclass +from types import SimpleNamespace + +import pytest +from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT +from fixtures.weighted_frames import weighted_entity_frame +from policyengine.core.scoping_strategy import RowFilterStrategy +from policyengine_simulation_contract.job_progress import JobStage +from policyengine_simulation_contract.uk_geography import ( + UKLocalAuthorityBoundaryVersion, + UKLocalAuthorityMetadata, +) +from policyengine_simulation_executor import simulation_output_geographic +from policyengine_simulation_executor.release_bundle import get_country_release_bundle +from policyengine_simulation_executor.simulation_macro_output import ( + GeographicImpactOutput, +) +from policyengine_simulation_executor.simulation_output_local_areas import ( + LocalAreaBreakdownBuilder, +) +from policyengine_simulation_executor.simulation_runtime import ( + DatasetSelection, + RegionResolution, + _run_simulation_impl_core, + normalise_policy, + resolve_region, +) +from policyengine_simulation_observability.stages import ANNUAL_IMPACT_STAGES, Stage +from test_simulation_output_fakes import TrackingRuntime, constituency_impact_record + + +def test_run_simulation_impl_core_passes_detected_uk_boundary_to_builder(monkeypatch): + # A bundle that does not route area regions: the national file carries + # the area codes and the run builds the area breakdowns from them. + _routes_uk_area_regions(monkeypatch, routed=False) + dataset = SimpleNamespace( + data=SimpleNamespace( + entity_data={ + "household": weighted_entity_frame( + {"la_code_oa": ["E06000001", "E07000026"]} + ) + } + ) + ) + country_module = SimpleNamespace(model=SimpleNamespace(version="2.90.0")) + builder_calls: list[dict] = [] + + class FakeSimulationOutputBuilder: + def __init__(self, **kwargs): + builder_calls.append(kwargs) + + def serialize(self): + return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) + + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._country_module", + lambda _country: country_module, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.resolve_region", + lambda **_kwargs: RegionResolution(code="uk", dataset_reference=None), + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._load_dataset", + lambda _params, country_module, selection, data_folder=None: dataset, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._build_simulation", + lambda *_args, **_kwargs: SimpleNamespace( + on_first_ensure=lambda _callback: None + ), + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", + FakeSimulationOutputBuilder, + ) + + _run_simulation_impl_core( + {"country": "uk", "baseline": {}, "reform": {}}, + runtime=TrackingRuntime(), + ) + + metadata = builder_calls[0]["uk_local_authority_metadata"] + assert metadata == UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD22 + ) + assert builder_calls[0]["omit_uk_area_breakdowns"] is False + + +def _routes_uk_area_regions(monkeypatch, *, routed: bool) -> None: + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "uk_area_regions_routed", + lambda: routed, + ) + + +def _uk_dataset(**household: list[object]) -> SimpleNamespace: + return SimpleNamespace( + data=SimpleNamespace( + entity_data={"household": weighted_entity_frame(household)} + ) + ) + + +_NATIONAL_SELECTION = DatasetSelection( + "microcosm_uk_2024_25", "hf://policyengine/uk/national.h5@cut", is_default=True +) +_LOCAL_AREA_SELECTION = DatasetSelection( + "microcosm_uk_2024_25_dense", "hf://policyengine/uk/local.h5@cut", is_default=False +) + + +_LOCAL_AREA_REVISION = "microcosm-uk-2024-25-dense-20261019T000432Z-ee5c2dbe" +_WALES = RowFilterStrategy(variable_name="region", variable_value="WALES") +_TEST_POLICIES = { + "baseline": {"gov.test.parameter": {"2026-01-01.2100-12-31": 1}}, + "reform": {"gov.test.parameter": {"2026-01-01.2100-12-31": 2}}, +} +_WALES_NATIONAL_PARAMS = {"country": "uk", "region": "country/wales", **_TEST_POLICIES} +_WALES_BREAKDOWN_PARAMS = { + "country": "uk", + "scope": "macro", + "region": "country/wales", + "time_period": "2026", + # The Public API always sends include_cliffs: false. + "include_cliffs": False, + "breakdown": "local_areas", + **_TEST_POLICIES, +} + + +def _routed_area_dataset(*, without: str | None = None) -> SimpleNamespace: + """Return a local-area dataset with every column its area outputs read. + + ``without`` names one column to leave out. + """ + household: dict[str, list[object]] = { + "region": ["WALES", "WALES"], + "constituency_code_oa": ["E14001063", "W07000112"], + "la_code_oa": ["E06000063", "W06000001"], + "region_code_oa": ["WALES", "WALES"], + "source_household_id": [1, 2], + } + household.pop(without or "", None) + return _uk_dataset(**household) + + +@dataclass +class _UKAreaRun: + """A UK run through ``_run_simulation_impl_core`` with every seam faked.""" + + builds: list[dict] + builder_calls: list[dict] + runtime: TrackingRuntime + # Each selection _load_dataset was asked for, and each country module load. + loads: list[DatasetSelection] + country_modules: list[str] + # The builder methods the run called to produce its output, and the output. + serialized: list[str] + output: object = None + + def dataset_load_attributes(self) -> list[dict]: + return [ + node["attrs"] + for node in self.runtime.roots + if node["name"] == ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD) + ] + + +class _RoutedUKSeams: + """Fakes for each seam of a routed UK run, recording into one ``_UKAreaRun``.""" + + def __init__(self, run, *, local_area_dataset): + self.run = run + self.local_area_dataset = local_area_dataset + self.national_dataset = _uk_dataset(region=["WALES", "WALES"]) + + def country_module(self, country): + self.run.country_modules.append(country) + return SimpleNamespace(model=SimpleNamespace(version="2.132.6")) + + def resolve_region(self, *, params, **_kwargs): + region = params.get("region") + if region == "country/wales": + return RegionResolution( + code="country/wales", dataset_reference=None, scoping_strategy=_WALES + ) + region_type = str(region).partition("/")[0] + if region_type in {"constituency", "local_authority"}: + # An area region names the routed local-area dataset itself. + return RegionResolution( + code=region, + dataset_reference=_LOCAL_AREA_SELECTION.name, + scoping_strategy=RowFilterStrategy( + variable_name=f"{region_type}_code_oa", + variable_value=region.partition("/")[2], + ), + ) + return RegionResolution(code="uk", dataset_reference=None) + + def resolve_selection(self, _params, *, region_resolution): + if region_resolution.dataset_reference == _LOCAL_AREA_SELECTION.name: + return _LOCAL_AREA_SELECTION + return _NATIONAL_SELECTION + + def load_dataset(self, _params, *, selection, country_module, data_folder=None): + self.run.loads.append(selection) + if selection is _LOCAL_AREA_SELECTION: + return self.local_area_dataset + return self.national_dataset + + def build_simulation(self, _params, *, dataset, dataset_selection, policy, region): + simulation = SimpleNamespace(on_first_ensure=lambda _callback: None) + self.run.builds.append( + { + "dataset": dataset, + "selection": dataset_selection, + "policy": policy, + "scoping_strategy": region.scoping_strategy, + "simulation": simulation, + } + ) + return simulation + + def output_builder(self, **kwargs): + self.run.builder_calls.append(kwargs) + return _RecordingBuilder( + self.run, "macro", dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) + ) + + def breakdown_builder(self, **kwargs): + self.run.builder_calls.append(kwargs) + return _RecordingBuilder( + self.run, "local_area_breakdown", {"kind": "local_areas"} + ) + + +class _RecordingBuilder: + """Records which builder serialized the run's output.""" + + def __init__(self, run, kind, output): + self.run = run + self.kind = kind + self.output = output + + def serialize(self): + self.run.serialized.append(self.kind) + return self.output + + +def _install_routed_uk_fakes( + monkeypatch, *, local_area_dataset, routed: bool = True +) -> _UKAreaRun: + """Fake every seam of a UK run whose national file carries no area codes.""" + run = _UKAreaRun( + builds=[], + builder_calls=[], + runtime=TrackingRuntime(), + loads=[], + country_modules=[], + serialized=[], + ) + seams = _RoutedUKSeams(run, local_area_dataset=local_area_dataset) + _routes_uk_area_regions(monkeypatch, routed=routed) + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "routed_uk_area_dataset_name", + lambda: _LOCAL_AREA_SELECTION.name, + ) + for name, value in { + "_country_module": seams.country_module, + "resolve_region": seams.resolve_region, + "_resolve_dataset_selection": seams.resolve_selection, + "_load_dataset": seams.load_dataset, + "_build_simulation": seams.build_simulation, + "SimulationOutputBuilder": seams.output_builder, + "LocalAreaBreakdownBuilder": seams.breakdown_builder, + }.items(): + monkeypatch.setattr( + f"policyengine_simulation_executor.simulation_runtime.{name}", value + ) + return run + + +def _run_routed_uk(monkeypatch, *, params, local_area_dataset) -> _UKAreaRun: + run = _install_routed_uk_fakes(monkeypatch, local_area_dataset=local_area_dataset) + run.output = _run_simulation_impl_core(dict(params), runtime=run.runtime) + return run + + +def _span_attributes(nodes) -> list[dict]: + """Return the attributes of every span in a tracked span tree.""" + attributes = [] + for node in nodes: + attributes.append(node["attrs"]) + attributes.extend(_span_attributes(node["children"])) + return attributes + + +@pytest.mark.parametrize("region", ["uk", "country/wales"]) +def test_a_routed_uk_wide_or_nation_run_builds_no_area_pair(monkeypatch, region): + from policyengine_simulation_executor import uk_local_authority_metadata + + run = _run_routed_uk( + monkeypatch, + params={**_WALES_NATIONAL_PARAMS, "region": region}, + local_area_dataset=_routed_area_dataset(), + ) + + # One dataset, the national file, and one baseline and reform on it. + assert run.loads == [_NATIONAL_SELECTION] + assert [build["selection"] for build in run.builds] == [_NATIONAL_SELECTION] * 2 + (builder_kwargs,) = run.builder_calls + assert builder_kwargs["baseline"] is run.builds[0]["simulation"] + assert builder_kwargs["reform"] is run.builds[1]["simulation"] + assert builder_kwargs["omit_uk_area_breakdowns"] is True + assert builder_kwargs["omit_uk_area_breakdowns"] is ( + uk_local_authority_metadata.uk_area_breakdowns_omitted("uk", region) + ) + assert "uk_area_simulations" not in builder_kwargs + assert run.serialized == ["macro"] + assert {"dataset_role": "uk_area"} not in _span_attributes(run.runtime.roots) + + +@pytest.mark.parametrize( + "region", ["constituency/E14001063", "local_authority/E06000063"] +) +@pytest.mark.parametrize( + "missing", + ["constituency_code_oa", "la_code_oa", "region_code_oa", "source_household_id"], +) +def test_a_routed_area_run_needs_every_column_its_outputs_read_before_simulating( + monkeypatch, region, missing +): + run = _install_routed_uk_fakes( + monkeypatch, local_area_dataset=_routed_area_dataset(without=missing) + ) + + with pytest.raises( + ValueError, + match=rf"'microcosm_uk_2024_25_dense', which the bundle routes .* has no " + rf"{missing} column", + ): + _run_simulation_impl_core( + {**_WALES_NATIONAL_PARAMS, "region": region}, runtime=run.runtime + ) + + assert run.loads == [_LOCAL_AREA_SELECTION] + assert run.builds == [] + + +def test_a_routed_area_run_with_every_column_detects_its_boundaries(monkeypatch): + run = _run_routed_uk( + monkeypatch, + params={**_WALES_NATIONAL_PARAMS, "region": "local_authority/E06000063"}, + local_area_dataset=_routed_area_dataset(), + ) + + (builder_kwargs,) = run.builder_calls + assert builder_kwargs["uk_local_authority_metadata"] == UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 + ) + assert builder_kwargs["omit_uk_area_breakdowns"] is False + + +def test_a_local_area_breakdown_runs_on_the_local_area_dataset_alone(monkeypatch): + local_area_dataset = _routed_area_dataset() + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime." + "_manifest_dataset_revision", + lambda country, name: { + ("uk", _LOCAL_AREA_SELECTION.name): _LOCAL_AREA_REVISION + }[(country, name)], + ) + + run = _run_routed_uk( + monkeypatch, + params=_WALES_BREAKDOWN_PARAMS, + local_area_dataset=local_area_dataset, + ) + + # The national file is never loaded: one load, one baseline and reform. + assert run.loads == [_LOCAL_AREA_SELECTION] + assert [build["selection"] for build in run.builds] == [_LOCAL_AREA_SELECTION] * 2 + assert all(build["dataset"] is local_area_dataset for build in run.builds) + assert all(build["scoping_strategy"] is _WALES for build in run.builds) + assert [build["policy"] for build in run.builds] == [ + normalise_policy(_TEST_POLICIES["baseline"]), + normalise_policy(_TEST_POLICIES["reform"]), + ] + (builder_kwargs,) = run.builder_calls + assert builder_kwargs["baseline"] is run.builds[0]["simulation"] + assert builder_kwargs["reform"] is run.builds[1]["simulation"] + assert builder_kwargs["local_authority_metadata"] == UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 + ) + assert builder_kwargs["dataset_name"] == _LOCAL_AREA_SELECTION.name + assert builder_kwargs["data_version"] == _LOCAL_AREA_REVISION + assert builder_kwargs["year"] == 2026 + assert run.serialized == ["local_area_breakdown"] + assert run.output == {"kind": "local_areas"} + assert run.dataset_load_attributes() == [{"dataset_role": "uk_area"}] + + +def test_a_scotland_breakdown_resolves_against_the_installed_bundle(monkeypatch): + # Only the dataset load, the simulations and the builder are faked: the + # country module, region resolution, dataset selection, routing and the + # revision come from the installed .py bundle, with no data or network. + from policyengine.provenance.manifest import get_release_manifest + from policyengine_simulation_executor import simulation_runtime + + loads: list[DatasetSelection] = [] + builds: list[dict] = [] + builder_calls: list[dict] = [] + + def load_dataset(_params, *, selection, country_module, data_folder=None): + loads.append(selection) + return _routed_area_dataset() + + def build_simulation(_params, *, dataset, dataset_selection, policy, region): + builds.append( + { + "selection": dataset_selection, + "region_code": region.code, + "scoping_strategy": region.scoping_strategy, + } + ) + return SimpleNamespace(on_first_ensure=lambda _callback: None) + + class RecordingBreakdownBuilder: + def __init__(self, **kwargs): + builder_calls.append(kwargs) + + def serialize(self): + return {"kind": "local_areas"} + + for name, value in { + "_load_dataset": load_dataset, + "_build_simulation": build_simulation, + "LocalAreaBreakdownBuilder": RecordingBreakdownBuilder, + }.items(): + monkeypatch.setattr(simulation_runtime, name, value) + + output = _run_simulation_impl_core( + {**_WALES_BREAKDOWN_PARAMS, "region": "country/scotland"}, + runtime=TrackingRuntime(), + ) + + (selection,) = loads + assert selection.name == "microcosm_uk_2024_25_dense" + assert selection.is_default is False + scotland = RowFilterStrategy(variable_name="region", variable_value="SCOTLAND") + assert [build["scoping_strategy"] for build in builds] == [scotland, scotland] + assert [build["selection"] for build in builds] == [selection, selection] + assert {build["region_code"] for build in builds} == {"country/scotland"} + (builder_kwargs,) = builder_calls + assert builder_kwargs["dataset_name"] == selection.name + assert builder_kwargs["data_version"] == ( + get_release_manifest("uk").datasets[selection.name].revision + ) + assert output == {"kind": "local_areas"} + + +def test_a_dataset_revision_is_read_from_the_release_manifest(monkeypatch): + from policyengine_simulation_executor import simulation_runtime + + manifests = { + "uk": SimpleNamespace( + datasets={ + _LOCAL_AREA_SELECTION.name: SimpleNamespace( + revision=_LOCAL_AREA_REVISION + ), + "unrevised": SimpleNamespace(revision=""), + } + ) + } + monkeypatch.setattr( + "policyengine.provenance.manifest.get_release_manifest", manifests.__getitem__ + ) + + revision = simulation_runtime._manifest_dataset_revision + assert revision("uk", _LOCAL_AREA_SELECTION.name) == _LOCAL_AREA_REVISION + assert revision("uk", "unrevised") is None + assert revision("uk", "absent") is None + + +def test_a_uk_wide_local_area_breakdown_simulates_the_whole_dataset(monkeypatch): + run = _run_routed_uk( + monkeypatch, + params={**_WALES_BREAKDOWN_PARAMS, "region": "uk"}, + local_area_dataset=_routed_area_dataset(), + ) + + assert run.loads == [_LOCAL_AREA_SELECTION] + assert [build["scoping_strategy"] for build in run.builds] == [None, None] + assert run.serialized == ["local_area_breakdown"] + + +@pytest.mark.parametrize( + ("change", "message"), + [ + ({"country": "us"}, "only for UK requests"), + ({"region": "constituency/E14001063"}, "'constituency/E14001063'"), + ({"region": "local_authority/E06000063"}, "'local_authority/E06000063'"), + ({"include_cliffs": True}, "include_cliffs=true"), + ({"region": None, "region_group": ["country/wales"]}, "region_group"), + ({"breakdown": "districts"}, "'districts'"), + ], +) +def test_an_invalid_local_area_breakdown_is_refused_before_any_load( + monkeypatch, change, message +): + run = _install_routed_uk_fakes( + monkeypatch, local_area_dataset=_routed_area_dataset() + ) + + with pytest.raises(ValueError, match=message): + _run_simulation_impl_core( + {**_WALES_BREAKDOWN_PARAMS, **change}, runtime=run.runtime + ) + + assert run.country_modules == [] + assert run.loads == [] + assert run.builds == [] + + +def test_a_local_area_breakdown_needs_a_bundle_that_routes_area_regions(monkeypatch): + run = _install_routed_uk_fakes( + monkeypatch, local_area_dataset=_routed_area_dataset(), routed=False + ) + + with pytest.raises( + ValueError, + match="does not route UK constituency and local-authority regions", + ): + _run_simulation_impl_core(dict(_WALES_BREAKDOWN_PARAMS), runtime=run.runtime) + + assert run.loads == [] + assert run.builds == [] + + +@pytest.mark.parametrize( + "missing", + ["constituency_code_oa", "la_code_oa", "region_code_oa", "source_household_id"], +) +def test_a_local_area_breakdown_needs_every_column_its_outputs_read( + monkeypatch, missing +): + local_area_dataset = _routed_area_dataset(without=missing) + + with pytest.raises( + ValueError, + match=rf"'microcosm_uk_2024_25_dense', which the bundle routes .* has no " + rf"{missing} column", + ): + _run_routed_uk( + monkeypatch, + params=_WALES_BREAKDOWN_PARAMS, + local_area_dataset=local_area_dataset, + ) + + +_LOCAL_AREA_URI = ( + "hf://policyengine/populace-uk-private/microcosm_uk_2024_25_dense.h5" + "@microcosm-uk-2024-25-dense-20261019T000432Z-ee5c2dbe" +) + + +def _registry_area(region_type: str, code: str, label: str) -> SimpleNamespace: + return SimpleNamespace( + code=f"{region_type}/{code}", + label=label, + region_type=region_type, + parent_code="uk", + dataset_path=_LOCAL_AREA_URI, + scoping_strategy=object(), + ) + + +def _routed_uk_country_module() -> SimpleNamespace: + """Return a UK model whose registry lists area regions, as routing gives.""" + areas = [ + _registry_area("constituency", "E14001063", "Aldershot"), + _registry_area("constituency", "W07000112", "Ynys Môn"), + _registry_area("local_authority", "E07000110", "Maidstone"), + _registry_area("local_authority", "N09000011", "Ards and North Down"), + ] + by_code = {area.code: area for area in areas} + registry = SimpleNamespace( + get_by_type=lambda region_type: [ + area for area in areas if area.region_type == region_type + ] + ) + return SimpleNamespace( + model=SimpleNamespace(get_region=by_code.get, region_registry=registry) + ) + + +@pytest.mark.parametrize( + ("requested", "expected"), + [ + ("constituency/E14001063", "constituency/E14001063"), + ("constituency/e14001063", "constituency/E14001063"), + ("constituency/Aldershot", "constituency/E14001063"), + ("constituency/ynys môn", "constituency/W07000112"), + ("local_authority/Maidstone", "local_authority/E07000110"), + ("local_authority/n09000011", "local_authority/N09000011"), + ("local_authority/Ards and North Down", "local_authority/N09000011"), + ], +) +def test_resolve_region_maps_routed_uk_area_names_and_codes_to_the_registry( + monkeypatch, requested, expected +): + _routes_uk_area_regions(monkeypatch, routed=True) + country_module = _routed_uk_country_module() + + resolution = resolve_region( + country_module=country_module, country="uk", params={"region": requested} + ) + + region = country_module.model.get_region(expected) + assert resolution.code == expected + assert resolution.scoping_strategy is region.scoping_strategy + assert resolution.dataset_reference == _LOCAL_AREA_URI + + +@pytest.mark.parametrize( + "requested", + [ + # A district that April 2023 replaced: LAD22 only. + "local_authority/E07000026", + "local_authority/Carlisle", + "constituency/Maidstone", + ], +) +def test_resolve_region_rejects_uk_areas_outside_the_routed_registry( + monkeypatch, requested +): + _routes_uk_area_regions(monkeypatch, routed=True) + + with pytest.raises(ValueError, match="Unsupported UK region"): + resolve_region( + country_module=_routed_uk_country_module(), + country="uk", + params={"region": requested}, + ) + + +@pytest.mark.parametrize( + "requested", ["constituency/Aldershot", "local_authority/E07000110"] +) +def test_resolve_region_rejects_uk_area_regions_on_an_unrouted_bundle( + monkeypatch, requested +): + # Without routing, policyengine.py's registry lists no area regions, and + # no weight-matrix region stands in for one: older bundles keep their own + # frozen apps. + _routes_uk_area_regions(monkeypatch, routed=False) + country_module = SimpleNamespace(model=SimpleNamespace(get_region=lambda _: None)) + + with pytest.raises(ValueError, match=f"Unsupported UK region: {requested}"): + resolve_region( + country_module=country_module, country="uk", params={"region": requested} + ) + + +class _BreakdownSimulation: + """A local-area simulation that records when it is computed.""" + + def __init__(self, label, events): + self.label = label + self.events = events + + def ensure(self): + self.events.append(self.label) + + +class _RecordingProgress: + def __init__(self, events): + self.events = events + + def enter(self, stage): + self.events.append(stage) + + +def _breakdown_builder(monkeypatch, *, events, runtime=None): + """Return a UK breakdown builder whose area outputs are faked. + + The country module's impact analysis fails the test if it is called. + Each fake records into ``events``. + """ + disclosure = simulation_output_geographic.UKGeographyDisclosureOutput() + + def compute(country, baseline_simulation, reform_simulation, *, year): + events.append( + ("release", country, baseline_simulation, reform_simulation, year) + ) + return disclosure + + def constituencies(country, value): + events.append(("constituency", value)) + return GeographicImpactOutput.model_validate( + [constituency_impact_record().model_dump(mode="json")] + ) + + def local_authorities(country, value, **kwargs): + events.append(("local_authority", value, kwargs["uk_local_authority_metadata"])) + return GeographicImpactOutput.model_validate([]) + + for name, function in { + "compute_uk_geography_disclosure_output": compute, + "build_uk_constituency_impact": constituencies, + "build_uk_local_authority_impact": local_authorities, + }.items(): + monkeypatch.setattr( + f"policyengine_simulation_executor.simulation_output_geographic.{name}", + function, + ) + + def no_impact_analysis(*_args, **_kwargs): + pytest.fail("a local-area breakdown computes no national impact analysis") + + builder = LocalAreaBreakdownBuilder( + baseline=_BreakdownSimulation("baseline", events), + reform=_BreakdownSimulation("reform", events), + local_authority_metadata=UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 + ), + dataset_name="microcosm_uk_2024_25_dense", + data_version=_LOCAL_AREA_REVISION, + year=2026, + country_module=SimpleNamespace( + model=SimpleNamespace(version="2.132.6"), + economic_impact_analysis=no_impact_analysis, + ), + runtime=runtime, + progress=_RecordingProgress(events), + ) + return builder, disclosure + + +def test_a_local_area_breakdown_publishes_both_levels_from_one_release(monkeypatch): + events: list[object] = [] + builder, disclosure = _breakdown_builder(monkeypatch, events=events) + + output = builder.serialize() + + releases = [event for event in events if event[0] == "release"] + assert releases == [("release", "uk", builder.baseline, builder.reform, 2026)] + assert [event for event in events if event[0] != "release"][-2:] == [ + ("constituency", disclosure), + ("local_authority", disclosure, builder.local_authority_metadata), + ] + assert output == { + "kind": "local_areas", + "model_version": "2.132.6", + "policyengine_version": get_country_release_bundle("uk").policyengine_version, + "dataset": "microcosm_uk_2024_25_dense", + "data_version": _LOCAL_AREA_REVISION, + "constituency_impact": [constituency_impact_record().model_dump(mode="json")], + "local_authority_impact": [], + } + + +def test_a_local_area_breakdown_enters_outputs_after_both_simulations(monkeypatch): + events: list[object] = [] + builder, _ = _breakdown_builder(monkeypatch, events=events) + + builder.serialize() + + simulations_and_stages = [event for event in events if not isinstance(event, tuple)] + assert simulations_and_stages == ["baseline", "reform", JobStage.OUTPUTS] + assert events.index(JobStage.OUTPUTS) < next( + index for index, event in enumerate(events) if isinstance(event, tuple) + ) + + +def test_a_local_area_breakdown_opens_only_registered_spans(monkeypatch): + events: list[object] = [] + runtime = TrackingRuntime() + builder, _ = _breakdown_builder(monkeypatch, events=events, runtime=runtime) + + builder.serialize() + + registered = {stage.value for stage in ANNUAL_IMPACT_STAGES.stages} + assert set(runtime.names) <= registered + assert ANNUAL_IMPACT_STAGES.name(Stage.ECONOMIC_IMPACT_ANALYSIS) not in ( + runtime.names + ) + (calculation,) = runtime.roots[:1] + assert calculation["name"] == ANNUAL_IMPACT_STAGES.name(Stage.CALCULATION) + (output_build,) = calculation["children"] + assert output_build["name"] == ANNUAL_IMPACT_STAGES.name( + Stage.SIMULATION_OUTPUT_BUILD + ) + assert { + ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_CONSTITUENCY), + ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_UK_LOCAL_AUTHORITY), + } <= {child["name"] for child in output_build["children"]} diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_runtime_policy.py b/projects/policyengine-simulation-executor/tests/test_simulation_runtime_policy.py index 2f032145b..9b56d4be4 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_runtime_policy.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_runtime_policy.py @@ -12,10 +12,10 @@ DatasetSelection, RegionResolution, _build_simulation, - _compile_interval_policy, _country_module, - _normalise_policy, - _parse_year, + compile_interval_policy, + normalise_policy, + parse_year, ) UK_RATE = "gov.hmrc.income_tax.rates.uk[0].rate" @@ -60,8 +60,8 @@ def test_uk_interval_policy_preserves_inclusive_stop_date(uk, stop): key = f"2027-04-06.{stop}" raw = {UK_RATE: {key: 0.21}} - policy = _compile_interval_policy( - _normalise_policy(raw), country_module=uk, year=2027 + policy = compile_interval_policy( + normalise_policy(raw), country_module=uk, year=2027 ) assert isinstance(policy, Policy) @@ -86,7 +86,7 @@ def test_mixed_effective_dates_and_intervals_are_chronological(uk, reverse): if reverse: entries.reverse() - policy = _compile_interval_policy( + policy = compile_interval_policy( {UK_RATE: dict(entries)}, country_module=uk, year=2027 ) @@ -114,7 +114,7 @@ def test_unknown_parameter_retains_wrapper_error_and_close_match_hint(uk): with pytest.raises(ValueError, match=re.escape(unknown)) as expected: wrapper_reform.compile_reform(raw, year=2027, model_version=uk.model) with pytest.raises(ValueError, match=re.escape(unknown)) as actual: - _compile_interval_policy(raw, country_module=uk, year=2027) + compile_interval_policy(raw, country_module=uk, year=2027) assert str(actual.value) == str(expected.value) assert unknown in str(actual.value) @@ -136,9 +136,9 @@ def test_invalid_policy_period_names_key_in_error(uk, key): raw = {UK_RATE: {"2026-04-06.2027-04-05": 0.20, key: 0.21}} with pytest.raises(ValueError, match=re.escape(key)) as normalise_error: - _normalise_policy(raw) + normalise_policy(raw) with pytest.raises(ValueError, match=re.escape(key)) as compile_error: - _compile_interval_policy(raw, country_module=uk, year=2027) + compile_interval_policy(raw, country_module=uk, year=2027) assert key in str(normalise_error.value) assert key in str(compile_error.value) @@ -171,7 +171,7 @@ def test_scalar_in_interval_policy_uses_request_year(period, year): region=RegionResolution(code="uk"), ) - assert _parse_year(params) == year + assert parse_year(params) == year scalar = next( value for value in simulation.policy.parameter_values @@ -192,7 +192,7 @@ def test_us_interval_policy_uses_same_public_parameter_values(): assert US_SALT_CAP in us.model.parameters_by_name raw = {US_SALT_CAP: {"2027-01-01.2027-12-31": 50000}} - policy = _compile_interval_policy(raw, country_module=us, year=2027) + policy = compile_interval_policy(raw, country_module=us, year=2027) assert isinstance(policy, Policy) assert len(policy.parameter_values) == 1 @@ -223,7 +223,7 @@ def recording_compiler(reform, *, year, model_version): return original_compiler(reform, year=year, model_version=model_version) monkeypatch.setattr(wrapper_reform, "compile_reform_to_policy", recording_compiler) - assert _compile_interval_policy(raw, country_module=uk, year=2027) is raw + assert compile_interval_policy(raw, country_module=uk, year=2027) is raw simulation = _build_simulation( {"country": "uk", "scope": "macro", "time_period": "2027"}, @@ -260,9 +260,9 @@ def recording_compiler(reform, *, year, model_version): def test_unpadded_dates_are_accepted_as_the_wrapper_always_accepted_them( key, start, stop ): - from policyengine_simulation_executor.simulation_runtime import ( + from policyengine_simulation_executor.request_policy import ( _parse_policy_period, ) - assert _normalise_policy({UK_RATE: {key: 0.21}}) == {UK_RATE: {key: 0.21}} + assert normalise_policy({UK_RATE: {key: 0.21}}) == {UK_RATE: {key: 0.21}} assert _parse_policy_period(key) == (start, stop) diff --git a/projects/policyengine-simulation-executor/tests/test_spm_runtime_native.py b/projects/policyengine-simulation-executor/tests/test_spm_runtime_native.py index acfc6fd05..a32807f2f 100644 --- a/projects/policyengine-simulation-executor/tests/test_spm_runtime_native.py +++ b/projects/policyengine-simulation-executor/tests/test_spm_runtime_native.py @@ -24,7 +24,7 @@ def forecast(): def _worker_entrypoint(params, monkeypatch): - from policyengine_simulation_executor import simulation_runtime + from policyengine_simulation_executor import region_resolution, simulation_runtime from src.modal import budget_window_batch if "start_year" in params: @@ -34,7 +34,7 @@ def _worker_entrypoint(params, monkeypatch): else: entrypoint = simulation_runtime._run_simulation_impl_core boundary = Mock(side_effect=AssertionError("Must reject before dataset lookup")) - monkeypatch.setattr(simulation_runtime, "_resolve_dataset_reference", boundary) + monkeypatch.setattr(region_resolution, "resolve_dataset_reference", boundary) return entrypoint, boundary diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py index 621481fad..95fce42ae 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py @@ -1029,7 +1029,7 @@ def test_planned_baseline_file_is_the_one_the_worker_saves_and_loads( ) params = simulation_module.simulation_runtime_params(planned) country_module = simulation_runtime._country_module("us") - region = simulation_runtime._resolve_region( + region = simulation_runtime.resolve_region( country_module=country_module, country="us", params=params ) dataset_path = "/opt/policyengine/stage12-data/populace_year_2026.h5" diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 2d5134ab2..90e205734 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -70,7 +70,7 @@ def ensure(self) -> None: ) monkeypatch.setattr( simulation_runtime, - "_resolve_region", + "resolve_region", lambda **_kwargs: region, ) @@ -183,7 +183,7 @@ def test_stage12_area_run_needs_every_column_its_outputs_read_before_simulating( ) monkeypatch.setattr( simulation_runtime, - "_resolve_region", + "resolve_region", lambda **_: RegionResolution( code=region_code, dataset_reference=_UK_LOCAL_AREA_DATASET ), @@ -270,7 +270,7 @@ def test_stage12_state_run_without_a_baked_file_loads_on_demand(monkeypatch) -> def _runtime_region_type(country: str, region: str | None) -> str | None: """Return the region type the runtime's own region normalisation implies.""" - from policyengine_simulation_executor.simulation_runtime import ( + from policyengine_simulation_executor.region_resolution import ( _normalise_region_code, ) diff --git a/rules/baselines/python-file-length.txt b/rules/baselines/python-file-length.txt index dc30e2766..0f8e078d1 100644 --- a/rules/baselines/python-file-length.txt +++ b/rules/baselines/python-file-length.txt @@ -6,5 +6,3 @@ # baseline count, in the pull request that splits the file. projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py projects/policyengine-simulation-entry/tests/test_deployment_assets.py -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py -projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py diff --git a/rules/registry.toml b/rules/registry.toml index c3fced421..3f7a6679b 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -1190,7 +1190,7 @@ job = "quality-guards" invocation = "run_quality_guards.py" [rule.baseline] file = "rules/baselines/python-file-length.txt" -count = 4 +count = 2 [rule.proof] bad_fixture = "rules/fixtures/sim.quality.python-file-length/bad_files.toml" good_fixture = "rules/fixtures/sim.quality.python-file-length/good_files.toml"