From 14a662414c1f2b17b5dbe5b3baa24db5b650a2a3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:13:54 +0400 Subject: [PATCH 01/37] Derive every v2 country list from one list in code The Public API listed its countries in several places: the v2 routes' `CountryId` type (`Literal["us", "uk"]`), an unused `SUPPORTED_COUNTRY_IDS` set in the query models, catalog extraction's `SUPPORTED_COUNTRIES` and its required packages and installed models, the set of countries the v1 routes copy into v2, and the household copy command's options. `policyengine_api/supported_countries.py` now holds the one list, `SupportedCountryId`, and `SUPPORTED_COUNTRY_IDS` derived from it; each of those places derives from it. The list lives in its own module, not in the query models, because the query models import the catalog package, which must import the list too. Tests: in a fresh interpreter, adding "zz" to the one list adds it to every derived list (this failed before the change: `CountryId` rejected "zz"); the catalog's per-country settings must cover exactly the listed countries. Co-Authored-By: Claude Opus 5.5 --- .../data/v2/catalog/catalog_selection.py | 3 +- .../data/v2/catalog/extraction.py | 10 +- policyengine_api/query_parameters.py | 4 +- policyengine_api/supported_countries.py | 19 +++ scripts/process_v1_household_mirror_event.py | 3 +- tests/unit/test_supported_countries.py | 139 ++++++++++++++++++ 6 files changed, 169 insertions(+), 9 deletions(-) create mode 100644 policyengine_api/supported_countries.py create mode 100644 tests/unit/test_supported_countries.py diff --git a/policyengine_api/data/v2/catalog/catalog_selection.py b/policyengine_api/data/v2/catalog/catalog_selection.py index d9b50fc76..1e97b6e3c 100644 --- a/policyengine_api/data/v2/catalog/catalog_selection.py +++ b/policyengine_api/data/v2/catalog/catalog_selection.py @@ -12,8 +12,9 @@ TaxBenefitModel, TaxBenefitModelVersion, ) +from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS -SUPPORTED_V2_COUNTRY_IDS = frozenset({"us", "uk"}) +SUPPORTED_V2_COUNTRY_IDS = frozenset(SUPPORTED_COUNTRY_IDS) MAXIMUM_POLICYENGINE_VERSION_LENGTH = 128 diff --git a/policyengine_api/data/v2/catalog/extraction.py b/policyengine_api/data/v2/catalog/extraction.py index 0872eea09..0817e41ae 100644 --- a/policyengine_api/data/v2/catalog/extraction.py +++ b/policyengine_api/data/v2/catalog/extraction.py @@ -27,15 +27,15 @@ RegionRecord, VariableRecord, ) +from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS if TYPE_CHECKING: from policyengine_api.json_types import JSONValue -SUPPORTED_COUNTRIES = ("us", "uk") +SUPPORTED_COUNTRIES = SUPPORTED_COUNTRY_IDS REQUIRED_DEPENDENCIES = ( "policyengine-core", - "policyengine-us", - "policyengine-uk", + *(f"policyengine-{country_id}" for country_id in SUPPORTED_COUNTRY_IDS), ) REVIEWED_DEFAULT_DATASETS = { "us": "populace_us_2024", @@ -928,7 +928,7 @@ def extract_installed_catalog() -> NormalizedCatalog: bundle=get_current_bundle(), policyengine_version=policyengine_version, models={ - "us": policyengine.us.model, - "uk": policyengine.uk.model, + country_id: getattr(policyengine, country_id).model + for country_id in SUPPORTED_COUNTRIES }, ) diff --git a/policyengine_api/query_parameters.py b/policyengine_api/query_parameters.py index 06e38cf01..619c6924b 100644 --- a/policyengine_api/query_parameters.py +++ b/policyengine_api/query_parameters.py @@ -29,6 +29,7 @@ validate_policyengine_version, ) from policyengine_api.spm import SPMSelection +from policyengine_api.supported_countries import SupportedCountryId from policyengine_api.utils.budget_window import ( BUDGET_WINDOW_MAX_END_YEAR, BUDGET_WINDOW_MAX_YEARS, @@ -40,7 +41,6 @@ DEFAULT_QUERY_LIMIT = 100 MAXIMUM_QUERY_LIMIT = 500 MAXIMUM_USER_ID_LENGTH = 255 -SUPPORTED_COUNTRY_IDS = frozenset({"us", "uk"}) def normalize_country_id(value: object) -> object: @@ -57,7 +57,7 @@ def validate_legacy_user_id(value: str) -> str: CountryId = Annotated[ - Literal["us", "uk"], + SupportedCountryId, BeforeValidator(normalize_country_id), Field(description="Supported PolicyEngine country ID"), ] diff --git a/policyengine_api/supported_countries.py b/policyengine_api/supported_countries.py new file mode 100644 index 000000000..3756e2bef --- /dev/null +++ b/policyengine_api/supported_countries.py @@ -0,0 +1,19 @@ +"""The countries the Public API serves, listed once in code. + +``SupportedCountryId`` is the one country list. Every other list derives from +it: the v2 routes' ``CountryId`` type, the countries catalog extraction +publishes, and the countries the v1 routes copy into v2 (dual-write). Adding +a country starts here. + +The module imports nothing from the package, so modules that the query +models import, such as the catalog package, can import it too. +""" + +from __future__ import annotations + +from typing import Final, Literal, get_args + +SupportedCountryId = Literal["us", "uk"] +SUPPORTED_COUNTRY_IDS: Final[tuple[SupportedCountryId, ...]] = get_args( + SupportedCountryId +) diff --git a/scripts/process_v1_household_mirror_event.py b/scripts/process_v1_household_mirror_event.py index f7af31072..49717019b 100755 --- a/scripts/process_v1_household_mirror_event.py +++ b/scripts/process_v1_household_mirror_event.py @@ -20,6 +20,7 @@ process_household_event_after_commit, ) from policyengine_api.services.household_service import HouseholdService +from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS if TYPE_CHECKING: from collections.abc import Mapping, Sequence @@ -149,7 +150,7 @@ def _parser() -> argparse.ArgumentParser: required=True, choices=("staging", "production"), ) - parser.add_argument("--country-id", required=True, choices=("us", "uk")) + parser.add_argument("--country-id", required=True, choices=SUPPORTED_COUNTRY_IDS) parser.add_argument("--legacy-household-id", required=True, type=int) return parser diff --git a/tests/unit/test_supported_countries.py b/tests/unit/test_supported_countries.py new file mode 100644 index 000000000..429f27657 --- /dev/null +++ b/tests/unit/test_supported_countries.py @@ -0,0 +1,139 @@ +"""The Public API's one country list, and the lists that derive from it.""" + +from __future__ import annotations + +import json +import subprocess +import sys +from typing import get_args + +import pytest +from pydantic import TypeAdapter, ValidationError + +from policyengine_api.constants import REPO +from policyengine_api.data.v2.catalog import extraction +from policyengine_api.query_parameters import CountryId +from policyengine_api.supported_countries import ( + SUPPORTED_COUNTRY_IDS, + SupportedCountryId, +) + +# Runs in a fresh interpreter: it adds a country to the one list before any +# module that derives from it is imported, then reports what each derived +# list contains. PolicyEngine.py is replaced by a stand-in, because importing +# the real package builds every country model. +EXTENDED_LIST_SCRIPT = """ +import json +import sys +from types import ModuleType, SimpleNamespace +from typing import Literal, get_args + +import policyengine_api.supported_countries as supported_countries + +supported_countries.SupportedCountryId = Literal["us", "uk", "zz"] +supported_countries.SUPPORTED_COUNTRY_IDS = get_args( + supported_countries.SupportedCountryId +) + +installed_policyengine = ModuleType("policyengine") +for country_id in ("us", "uk", "zz", "ca"): + setattr( + installed_policyengine, + country_id, + SimpleNamespace(model=f"{country_id} model"), + ) +installed_bundle = ModuleType("policyengine.bundle") +installed_bundle.get_current_bundle = dict +sys.modules["policyengine"] = installed_policyengine +sys.modules["policyengine.bundle"] = installed_bundle + +from pydantic import TypeAdapter + +from policyengine_api.data.v2.catalog import catalog_selection, extraction +from policyengine_api.query_parameters import CountryId +from scripts import process_v1_household_mirror_event + +captured_arguments = {} + + +def capture_extraction(**arguments): + captured_arguments.update(arguments) + + +extraction.extract_catalog = capture_extraction +extraction.extract_installed_catalog() +country_option = next( + action + for action in process_v1_household_mirror_event._parser()._actions + if action.dest == "country_id" +) +print( + json.dumps( + { + "country_type": TypeAdapter(CountryId).validate_python("ZZ"), + "extracted_countries": list(extraction.SUPPORTED_COUNTRIES), + "required_dependencies": list(extraction.REQUIRED_DEPENDENCIES), + "installed_models": captured_arguments["models"], + "dual_write_countries": sorted(catalog_selection.SUPPORTED_V2_COUNTRY_IDS), + "household_copy_countries": list(country_option.choices), + } + ) +) +""" + + +def test_the_country_list_is_the_us_and_the_uk_in_catalog_order() -> None: + assert SUPPORTED_COUNTRY_IDS == ("us", "uk") + assert get_args(SupportedCountryId) == SUPPORTED_COUNTRY_IDS + + +def test_the_v2_country_type_accepts_exactly_the_listed_countries() -> None: + country_id_adapter = TypeAdapter(CountryId) + + for country_id in SUPPORTED_COUNTRY_IDS: + assert country_id_adapter.validate_python(country_id) == country_id + assert country_id_adapter.validate_python(country_id.upper()) == country_id + for unlisted_country_id in ("ca", "ng", "il", ""): + with pytest.raises(ValidationError): + country_id_adapter.validate_python(unlisted_country_id) + + +def test_every_country_list_follows_a_country_added_to_the_one_list() -> None: + result = subprocess.run( + [sys.executable, "-c", EXTENDED_LIST_SCRIPT], + cwd=REPO, + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, result.stderr + derived_lists = json.loads(result.stdout.strip().splitlines()[-1]) + assert derived_lists == { + "country_type": "zz", + "extracted_countries": ["us", "uk", "zz"], + "required_dependencies": [ + "policyengine-core", + "policyengine-us", + "policyengine-uk", + "policyengine-zz", + ], + # The installed package's other countries are not extracted. + "installed_models": { + "us": "us model", + "uk": "uk model", + "zz": "zz model", + }, + "dual_write_countries": ["uk", "us", "zz"], + "household_copy_countries": ["us", "uk", "zz"], + } + + +def test_catalog_settings_cover_exactly_the_listed_countries() -> None: + # Each listed country needs reviewed catalog settings, so adding a country + # without them fails here rather than during publication. + listed_countries = set(SUPPORTED_COUNTRY_IDS) + + assert set(extraction.REVIEWED_DEFAULT_DATASETS) == listed_countries + assert set(extraction.CURRENT_LAW_IDS) == listed_countries + assert set(extraction.METADATA_TIME_PERIODS) == listed_countries From 2729c37f4c212e4c2787281bb03d8f015231a99a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:29:39 +0400 Subject: [PATCH 02/37] Record each served country in a countries table that country columns reference Adds the v2 `countries` table: one row per country the Public API serves, with `household_enabled` and `society_wide_enabled` saying which calculations it supports (a check requires society-wide countries to be household capable). `users.primary_country` and the `country_id` of households, policies and the four user association tables reference it, with `ON DELETE RESTRICT`, instead of each checking `IN ('us', 'uk')`. The tables that pair v1 records with their v2 copies keep their checks. Revision `fb1a3e40d3d6` was autogenerated. Its one hand-written operation inserts the US and UK rows, both capable of both, between the table and the foreign keys, which otherwise fail on every database whose rows already name either country. This insert is an exception to the rule that revisions contain only generated operations, approved for this revision alone; the revision's comment and the Alembic skill say so. The generated downgrade restores the checks and drops the table with its rows. The privilege manifest gives the countries table no role access yet. Tests (failed first because the `countries` table and `Country` model did not exist): model structure and foreign keys; SQLite persistence (an unlisted country is refused, a listed country in use cannot be deleted, a household-only row lets records name that country); the revision's operations; and the Postgres lifecycle (upgrade keeps existing US and UK rows, unlisted countries are refused, downgrade and re-upgrade, a downgrade is refused while records name another country, and the previous release's table usage still works after the upgrade). Co-Authored-By: Claude Opus 5.5 --- changelog.d/v2-countries-table.added.md | 1 + docs/engineering/skills/alembic-migrations.md | 24 ++ .../fb1a3e40d3d6_add_countries_table.py | 195 ++++++++++ policyengine_api/data/v2/migration_target.py | 2 +- policyengine_api/data/v2/models/__init__.py | 2 + .../data/v2/models/associations.py | 40 +- policyengine_api/data/v2/models/countries.py | 37 ++ policyengine_api/data/v2/models/households.py | 10 +- policyengine_api/data/v2/models/policies.py | 10 +- policyengine_api/data/v2/models/users.py | 11 +- policyengine_api/data/v2/privileges.py | 2 + policyengine_api/supported_countries.py | 4 + tests/fixtures/v2_countries.py | 28 ++ .../integration/test_alembic_v2_lifecycle.py | 349 ++++++++++++++++-- tests/unit/v2/test_alembic_v2.py | 91 ++++- tests/unit/v2/test_household_models.py | 3 + tests/unit/v2/test_import_side_effects.py | 2 +- tests/unit/v2/test_model_persistence.py | 78 +++- tests/unit/v2/test_models.py | 90 ++++- tests/unit/v2/user_service_support.py | 3 + 20 files changed, 909 insertions(+), 73 deletions(-) create mode 100644 changelog.d/v2-countries-table.added.md create mode 100644 migrations/v2/versions/fb1a3e40d3d6_add_countries_table.py create mode 100644 policyengine_api/data/v2/models/countries.py create mode 100644 tests/fixtures/v2_countries.py diff --git a/changelog.d/v2-countries-table.added.md b/changelog.d/v2-countries-table.added.md new file mode 100644 index 000000000..5c7e480ac --- /dev/null +++ b/changelog.d/v2-countries-table.added.md @@ -0,0 +1 @@ +Added a v2 `countries` table that records, for each country the Public API serves, whether it supports household and society-wide calculations. Its revision lists the US and the UK, both with both. Users' primary country and the country of every v2 household, policy and user association now reference this table instead of a fixed US-and-UK check. diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index d4a043d91..7a5819fa9 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -294,3 +294,27 @@ last dependent table is dropped. Its only post-generation correction drops the nine generated `v2_*` enum types at the end of the baseline downgrade. The v2 Postgres lifecycle test covers empty upgrade, downgrade to base, and re-upgrade so stale enum types cannot make the generated baseline non-reversible. + +### API v2 countries rows + +Revision `fb1a3e40d3d6` was autogenerated after adding the `countries` table +and replacing the US-and-UK checks on `users.primary_country` and on the +`country_id` of households, policies and the four user association tables +with foreign keys to it. It contains one hand-written operation: it inserts +the `uk` and `us` rows, both household and society-wide capable, between the +generated table creation and the generated foreign keys, which otherwise fail +on any database whose rows already name either country. The rows are exactly +the values the replaced checks allowed. The generated downgrade removes them +with the table, and its restored checks fail if any record names another +country. + +Inserting rows is a data operation, which this skill otherwise forbids. The +maintainers approved this insert as a one-off exception, because the new +foreign keys need the rows in the same revision. Unit tests pin the +operation, and the v2 Postgres lifecycle test covers an upgrade that keeps +existing rows, the downgrade, a downgrade refused by another country's +records, and the re-upgrade. + +The table must list exactly the countries in +`policyengine_api/supported_countries.py`. Listing another country later +needs a revision that inserts its row, and the same human migration decision. diff --git a/migrations/v2/versions/fb1a3e40d3d6_add_countries_table.py b/migrations/v2/versions/fb1a3e40d3d6_add_countries_table.py new file mode 100644 index 000000000..1b02ec819 --- /dev/null +++ b/migrations/v2/versions/fb1a3e40d3d6_add_countries_table.py @@ -0,0 +1,195 @@ +"""add countries table + +Revision ID: fb1a3e40d3d6 +Revises: fbb7488ebe8a +Create Date: 2026-10-22 02:16:50.126715 +Generation: uv run alembic -c alembic-v2.ini revision --autogenerate +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import sqlmodel + + +revision: str = "fb1a3e40d3d6" +down_revision: Union[str, None] = "fbb7488ebe8a" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + countries = op.create_table( + "countries", + sa.Column("id", sqlmodel.sql.sqltypes.AutoString(length=2), nullable=False), + sa.Column("household_enabled", sa.Boolean(), nullable=False), + sa.Column("society_wide_enabled", sa.Boolean(), nullable=False), + sa.CheckConstraint( + "household_enabled OR NOT society_wide_enabled", + name=op.f("ck_countries_society_wide_requires_household"), + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_countries")), + ) + # Approved hand-written data operation: the one exception to generated-only + # revisions, decided by the maintainers for this table alone. The foreign + # keys below fail on any database whose rows already name the US or the UK + # unless both countries are listed first, so their rows are inserted here, + # in the same revision. They are the countries the replaced checks + # allowed, and both support household and society-wide calculations. The + # generated downgrade removes them with the table. Listing another + # country later needs its own revision and the same kind of decision. + op.bulk_insert( + countries, + [ + {"id": "uk", "household_enabled": True, "society_wide_enabled": True}, + {"id": "us", "household_enabled": True, "society_wide_enabled": True}, + ], + ) + op.create_foreign_key( + op.f("fk_households_country_id_countries"), + "households", + "countries", + ["country_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint(op.f("ck_households_country"), "households", type_="check") + op.create_foreign_key( + op.f("fk_policies_country_id_countries"), + "policies", + "countries", + ["country_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint(op.f("ck_policies_country"), "policies", type_="check") + op.create_foreign_key( + op.f("fk_user_household_associations_country_id_countries"), + "user_household_associations", + "countries", + ["country_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint( + op.f("ck_user_household_associations_country"), + "user_household_associations", + type_="check", + ) + op.create_foreign_key( + op.f("fk_user_policies_country_id_countries"), + "user_policies", + "countries", + ["country_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint(op.f("ck_user_policies_country"), "user_policies", type_="check") + op.create_foreign_key( + op.f("fk_user_report_associations_country_id_countries"), + "user_report_associations", + "countries", + ["country_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint( + op.f("ck_user_report_associations_country"), + "user_report_associations", + type_="check", + ) + op.create_foreign_key( + op.f("fk_user_simulation_associations_country_id_countries"), + "user_simulation_associations", + "countries", + ["country_id"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint( + op.f("ck_user_simulation_associations_country"), + "user_simulation_associations", + type_="check", + ) + op.create_foreign_key( + op.f("fk_users_primary_country_countries"), + "users", + "countries", + ["primary_country"], + ["id"], + ondelete="RESTRICT", + ) + op.drop_constraint(op.f("ck_users_primary_country"), "users", type_="check") + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_check_constraint( + op.f("ck_users_primary_country"), + "users", + "primary_country::text = ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])", + ) + op.drop_constraint( + op.f("fk_users_primary_country_countries"), "users", type_="foreignkey" + ) + op.create_check_constraint( + op.f("ck_user_simulation_associations_country"), + "user_simulation_associations", + "country_id::text = ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])", + ) + op.drop_constraint( + op.f("fk_user_simulation_associations_country_id_countries"), + "user_simulation_associations", + type_="foreignkey", + ) + op.create_check_constraint( + op.f("ck_user_report_associations_country"), + "user_report_associations", + "country_id::text = ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])", + ) + op.drop_constraint( + op.f("fk_user_report_associations_country_id_countries"), + "user_report_associations", + type_="foreignkey", + ) + op.create_check_constraint( + op.f("ck_user_policies_country"), + "user_policies", + "country_id::text = ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])", + ) + op.drop_constraint( + op.f("fk_user_policies_country_id_countries"), + "user_policies", + type_="foreignkey", + ) + op.create_check_constraint( + op.f("ck_user_household_associations_country"), + "user_household_associations", + "country_id::text = ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])", + ) + op.drop_constraint( + op.f("fk_user_household_associations_country_id_countries"), + "user_household_associations", + type_="foreignkey", + ) + op.create_check_constraint( + op.f("ck_policies_country"), + "policies", + "country_id::text = ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])", + ) + op.drop_constraint( + op.f("fk_policies_country_id_countries"), "policies", type_="foreignkey" + ) + op.create_check_constraint( + op.f("ck_households_country"), + "households", + "country_id::text = ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])", + ) + op.drop_constraint( + op.f("fk_households_country_id_countries"), "households", type_="foreignkey" + ) + op.drop_table("countries") + # ### end Alembic commands ### diff --git a/policyengine_api/data/v2/migration_target.py b/policyengine_api/data/v2/migration_target.py index fd4c05fdf..2c7b6986d 100644 --- a/policyengine_api/data/v2/migration_target.py +++ b/policyengine_api/data/v2/migration_target.py @@ -24,7 +24,7 @@ from collections.abc import Mapping V2_ALEMBIC_DISPOSABLE_TEST = "V2_ALEMBIC_DISPOSABLE_TEST" -V2_ALEMBIC_HEAD_REVISION = "fbb7488ebe8a" +V2_ALEMBIC_HEAD_REVISION = "fb1a3e40d3d6" DISPOSABLE_DATABASE_NAME = "policyengine_v2_alembic_test" DISPOSABLE_HOSTS = frozenset({"localhost", "127.0.0.1", "::1", "postgres"}) MIGRATION_ROLE = "policyengine_v2_migrator" diff --git a/policyengine_api/data/v2/models/__init__.py b/policyengine_api/data/v2/models/__init__.py index c63655bb6..a3aedae09 100644 --- a/policyengine_api/data/v2/models/__init__.py +++ b/policyengine_api/data/v2/models/__init__.py @@ -31,6 +31,7 @@ Stage12RunStatus, Stage12SimulationRole, ) +from policyengine_api.data.v2.models.countries import Country # noqa: E402 from policyengine_api.data.v2.models.household_mappings import ( # noqa: E402 LegacyHouseholdMapping, ) @@ -102,6 +103,7 @@ "ChangeAggregate", "CongressionalDistrictImpact", "ConstituencyImpact", + "Country", "Dataset", "DatasetVersion", "DecileImpact", diff --git a/policyengine_api/data/v2/models/associations.py b/policyengine_api/data/v2/models/associations.py index 2c5cd21b1..dc7e7d757 100644 --- a/policyengine_api/data/v2/models/associations.py +++ b/policyengine_api/data/v2/models/associations.py @@ -51,10 +51,6 @@ class UserHouseholdAssociation(TimestampedModel, table=True): "country_id", name="uq_user_household_associations_id_country", ), - sa.CheckConstraint( - "country_id IN ('us', 'uk')", - name="ck_user_household_associations_country", - ), sa.CheckConstraint( "household_id IS NOT NULL OR legacy_household_id IS NOT NULL", name="ck_user_household_associations_household_ref", @@ -91,7 +87,11 @@ class UserHouseholdAssociation(TimestampedModel, table=True): ) household_id: UUID | None = Field(default=None, index=True) legacy_household_id: int | None = Field(default=None, sa_type=sa.BigInteger) - country_id: str = Field(max_length=2) + country_id: str = Field( + max_length=2, + foreign_key="countries.id", + ondelete="RESTRICT", + ) name: str | None = Field(default=None, max_length=255) description: str | None = Field(default=None, sa_type=sa.Text) @@ -113,10 +113,6 @@ class UserPolicy(TimestampedModel, table=True): "country_id", name="uq_user_policies_id_country", ), - sa.CheckConstraint( - "country_id IN ('us', 'uk')", - name="ck_user_policies_country", - ), sa.CheckConstraint( "policy_id IS NOT NULL OR legacy_policy_id IS NOT NULL", name="ck_user_policies_policy_ref", @@ -157,7 +153,11 @@ class UserPolicy(TimestampedModel, table=True): ) policy_id: UUID | None = Field(default=None, index=True) legacy_policy_id: int | None = Field(default=None, sa_type=sa.BigInteger) - country_id: str = Field(max_length=2) + country_id: str = Field( + max_length=2, + foreign_key="countries.id", + ondelete="RESTRICT", + ) name: str | None = Field(default=None, max_length=255) description: str | None = Field(default=None, sa_type=sa.Text) baseline: str | None = Field(default=None, max_length=16) @@ -185,10 +185,6 @@ class UserSimulationAssociation(TimestampedModel, table=True): "country_id", name="uq_user_simulation_associations_id_country", ), - sa.CheckConstraint( - "country_id IN ('us', 'uk')", - name="ck_user_simulation_associations_country", - ), sa.CheckConstraint( "simulation_id IS NOT NULL OR legacy_simulation_id IS NOT NULL", name="ck_user_simulation_associations_simulation_ref", @@ -214,7 +210,11 @@ class UserSimulationAssociation(TimestampedModel, table=True): ) simulation_id: UUID | None = Field(default=None, index=True) legacy_simulation_id: int | None = Field(default=None, sa_type=sa.BigInteger) - country_id: str = Field(max_length=2) + country_id: str = Field( + max_length=2, + foreign_key="countries.id", + ondelete="RESTRICT", + ) name: str | None = Field(default=None, max_length=255) description: str | None = Field(default=None, sa_type=sa.Text) # Retired; see the module docstring. @@ -239,10 +239,6 @@ class UserReportAssociation(TimestampedModel, table=True): "country_id", name="uq_user_report_associations_id_country", ), - sa.CheckConstraint( - "country_id IN ('us', 'uk')", - name="ck_user_report_associations_country", - ), sa.CheckConstraint( "report_id IS NOT NULL OR legacy_report_id IS NOT NULL", name="ck_user_report_associations_report_ref", @@ -273,7 +269,11 @@ class UserReportAssociation(TimestampedModel, table=True): ) report_id: UUID | None = Field(default=None, index=True) legacy_report_id: int | None = Field(default=None, sa_type=sa.BigInteger) - country_id: str = Field(max_length=2) + country_id: str = Field( + max_length=2, + foreign_key="countries.id", + ondelete="RESTRICT", + ) name: str | None = Field(default=None, max_length=255) description: str | None = Field(default=None, sa_type=sa.Text) user_policy_id: UUID | None = Field( diff --git a/policyengine_api/data/v2/models/countries.py b/policyengine_api/data/v2/models/countries.py new file mode 100644 index 000000000..fc3ae8677 --- /dev/null +++ b/policyengine_api/data/v2/models/countries.py @@ -0,0 +1,37 @@ +"""Canonical SQLModel table recording what each served country supports. + +Which countries exist is decided in code: ``policyengine_api.supported_countries`` +lists them, and this table must hold exactly one row per listed country. A +row records the calculations its country supports: a household-only country +has ``household_enabled`` set and ``society_wide_enabled`` unset. + +``users.primary_country`` and the ``country_id`` of households, policies and +the four user association tables reference ``countries.id``. Other v2 country +columns do not: the mapping tables that pair v1 records with their v2 copies +keep their US-and-UK checks, and the simulation-domain tables, the comparison +run tables and the retired association ``country`` columns store unconstrained +text. +""" + +import sqlalchemy as sa +from sqlmodel import Field, SQLModel + + +class Country(SQLModel, table=True): + """One served country and the calculations it supports. + + A society-wide country must also be household capable. A country that a + user, household, policy or user association names cannot be deleted. + """ + + __tablename__ = "countries" + __table_args__ = ( + sa.CheckConstraint( + "household_enabled OR NOT society_wide_enabled", + name="ck_countries_society_wide_requires_household", + ), + ) + + id: str = Field(primary_key=True, max_length=2) + household_enabled: bool + society_wide_enabled: bool diff --git a/policyengine_api/data/v2/models/households.py b/policyengine_api/data/v2/models/households.py index 78e76b2ab..181266d9d 100644 --- a/policyengine_api/data/v2/models/households.py +++ b/policyengine_api/data/v2/models/households.py @@ -53,10 +53,6 @@ class Household(TimestampedModel, table=True): "content_hash", name="uq_households_canonicalization_content_hash", ), - sa.CheckConstraint( - "country_id IN ('us', 'uk')", - name="ck_households_country", - ), sa.CheckConstraint( "default_year IS NULL OR default_year BETWEEN 1900 AND 2200", name="ck_households_default_year", @@ -78,7 +74,11 @@ class Household(TimestampedModel, table=True): ), ) - country_id: str = Field(max_length=2) + country_id: str = Field( + max_length=2, + foreign_key="countries.id", + ondelete="RESTRICT", + ) default_year: int | None = None household_data: dict[str, Any] = Field( sa_type=sa.JSON().with_variant(sa.dialects.postgresql.JSONB(), "postgresql") diff --git a/policyengine_api/data/v2/models/policies.py b/policyengine_api/data/v2/models/policies.py index 97e4a619e..f0da6e8cc 100644 --- a/policyengine_api/data/v2/models/policies.py +++ b/policyengine_api/data/v2/models/policies.py @@ -40,10 +40,6 @@ class Policy(TimestampedModel, table=True): "content_hash", name="uq_policies_canonicalization_content_hash", ), - sa.CheckConstraint( - "country_id IN ('us', 'uk')", - name="ck_policies_country", - ), sa.CheckConstraint( "canonicalization_version > 0", name="ck_policies_canonicalization_version", @@ -64,7 +60,11 @@ class Policy(TimestampedModel, table=True): ), ) - country_id: str = Field(max_length=2) + country_id: str = Field( + max_length=2, + foreign_key="countries.id", + ondelete="RESTRICT", + ) tax_benefit_model_id: UUID = Field( foreign_key="tax_benefit_models.id", ondelete="RESTRICT", diff --git a/policyengine_api/data/v2/models/users.py b/policyengine_api/data/v2/models/users.py index 85c0ea19a..f3b987123 100644 --- a/policyengine_api/data/v2/models/users.py +++ b/policyengine_api/data/v2/models/users.py @@ -44,10 +44,6 @@ class User(IdentifiedModel, table=True): __tablename__ = "users" __table_args__ = ( sa.UniqueConstraint("email", name="uq_users_email"), - sa.CheckConstraint( - "primary_country IN ('us', 'uk')", - name="ck_users_primary_country", - ), sa.CheckConstraint( "user_type IN ('unauthenticated', 'standard')", name="ck_users_user_type", @@ -61,7 +57,12 @@ class User(IdentifiedModel, table=True): first_name: str | None = Field(default=None, max_length=255) last_name: str | None = Field(default=None, max_length=255) email: str | None = Field(default=None, max_length=320, index=True) - primary_country: str | None = Field(default=None, max_length=2) + primary_country: str | None = Field( + default=None, + max_length=2, + foreign_key="countries.id", + ondelete="RESTRICT", + ) # The server default keeps inserts by earlier Public API revisions, which # do not map this column, creating standard users. user_type: str = Field( diff --git a/policyengine_api/data/v2/privileges.py b/policyengine_api/data/v2/privileges.py index 8f8c34515..6fd0cc750 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -132,6 +132,8 @@ def _access( # Stage 10 saved-policy mirror. "legacy_user_mappings": _access(runtime=READ_WRITE), "legacy_user_policy_mappings": _access(runtime=READ_WRITE), + # Served countries; the migration that creates the table inserts its rows. + "countries": _access(), # Users and their saved work (Stage 12A). "user_household_associations": _access(runtime=READ_WRITE), "user_identities": _access(runtime=READ_WRITE), diff --git a/policyengine_api/supported_countries.py b/policyengine_api/supported_countries.py index 3756e2bef..51a495dd4 100644 --- a/policyengine_api/supported_countries.py +++ b/policyengine_api/supported_countries.py @@ -5,6 +5,10 @@ publishes, and the countries the v1 routes copy into v2 (dual-write). Adding a country starts here. +The v2 ``countries`` table records what each country supports, and must hold +one row for each listed country and no other: a v2 Alembic revision inserts +the row when a country is added. + The module imports nothing from the package, so modules that the query models import, such as the catalog package, can import it too. """ diff --git a/tests/fixtures/v2_countries.py b/tests/fixtures/v2_countries.py new file mode 100644 index 000000000..36d4b6604 --- /dev/null +++ b/tests/fixtures/v2_countries.py @@ -0,0 +1,28 @@ +"""Countries rows for SQLite tests that enforce v2 foreign keys. + +``V2_METADATA.create_all`` creates an empty countries table, while the v2 +Alembic revision that adds the table inserts the US and UK rows. Tests that +turn on SQLite foreign-key enforcement add the same rows before writing +records that name a country. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_api.data.v2.models import Country + +if TYPE_CHECKING: + from sqlmodel import Session + +# The rows the countries revision inserts, ordered by country ID. +LISTED_COUNTRY_ROWS = ( + {"id": "uk", "household_enabled": True, "society_wide_enabled": True}, + {"id": "us", "household_enabled": True, "society_wide_enabled": True}, +) + + +def add_listed_countries(session: Session) -> None: + """Insert and commit the US and UK countries rows.""" + session.add_all(Country(**row) for row in LISTED_COUNTRY_ROWS) + session.commit() diff --git a/tests/integration/test_alembic_v2_lifecycle.py b/tests/integration/test_alembic_v2_lifecycle.py index 814971236..2df169511 100644 --- a/tests/integration/test_alembic_v2_lifecycle.py +++ b/tests/integration/test_alembic_v2_lifecycle.py @@ -3,6 +3,7 @@ import itertools import os from collections.abc import Iterator +from dataclasses import dataclass from datetime import UTC, datetime, timedelta from uuid import UUID, uuid4 @@ -30,6 +31,7 @@ STAGE_12_PREVIOUS_REVISION = "724b1b11a33e" STAGE_12_COMPARISON_PREVIOUS_REVISION = "439303be14fe" STAGE_12A_PREVIOUS_REVISION = "91af4fc8d8c1" +COUNTRIES_PREVIOUS_REVISION = "fbb7488ebe8a" V2_TABLE_NAMES = frozenset(table.name for table in V2_METADATA.tables.values()) @@ -187,7 +189,6 @@ def _assert_stage_12a_association_shape(engine) -> None: for constraint in inspect(engine).get_check_constraints("users") } assert { - "ck_users_primary_country", "ck_users_user_type", "ck_users_standard_primary_country", } <= user_checks @@ -352,6 +353,75 @@ def _assert_base(engine: Engine) -> None: ) +# Every column that names a record's country, with the check it replaced. +COUNTRY_REFERENCE_COLUMNS = { + ("users", "primary_country"): "ck_users_primary_country", + ("households", "country_id"): "ck_households_country", + ("policies", "country_id"): "ck_policies_country", + ("user_household_associations", "country_id"): ( + "ck_user_household_associations_country" + ), + ("user_policies", "country_id"): "ck_user_policies_country", + ("user_simulation_associations", "country_id"): ( + "ck_user_simulation_associations_country" + ), + ("user_report_associations", "country_id"): "ck_user_report_associations_country", +} +# The rows the countries revision inserts: the countries the replaced checks +# allowed, both household and society-wide capable. +LISTED_COUNTRIES = [("uk", True, True), ("us", True, True)] + + +def _country_foreign_keys(engine, table_name: str) -> list[dict]: + return [ + foreign_key + for foreign_key in inspect(engine).get_foreign_keys(table_name) + if foreign_key["referred_table"] == "countries" + ] + + +def _check_names(engine, table_name: str) -> set[str]: + return { + constraint["name"] + for constraint in inspect(engine).get_check_constraints(table_name) + } + + +def _listed_countries(engine) -> list[tuple]: + with engine.connect() as connection: + return [ + tuple(row) + for row in connection.execute( + text( + "SELECT id, household_enabled, society_wide_enabled " + "FROM countries ORDER BY id" + ) + ) + ] + + +def _assert_countries_shape(engine) -> None: + assert _listed_countries(engine) == LISTED_COUNTRIES + assert "ck_countries_society_wide_requires_household" in _check_names( + engine, "countries" + ) + for (table_name, column_name), check_name in COUNTRY_REFERENCE_COLUMNS.items(): + country_keys = _country_foreign_keys(engine, table_name) + assert len(country_keys) == 1, table_name + assert country_keys[0]["name"] == f"fk_{table_name}_{column_name}_countries" + assert country_keys[0]["constrained_columns"] == [column_name] + assert country_keys[0]["referred_columns"] == ["id"] + assert country_keys[0]["options"]["ondelete"] == "RESTRICT" + assert check_name not in _check_names(engine, table_name) + + +def _assert_pre_countries_shape(engine) -> None: + assert "countries" not in inspect(engine).get_table_names(schema="public") + for (table_name, _column_name), check_name in COUNTRY_REFERENCE_COLUMNS.items(): + assert _country_foreign_keys(engine, table_name) == [] + assert check_name in _check_names(engine, table_name) + + def test_empty_upgrade_check_base_downgrade_and_reupgrade() -> None: database_url = _disposable_url() config = _config() @@ -364,6 +434,13 @@ def test_empty_upgrade_check_base_downgrade_and_reupgrade() -> None: } _upgrade_and_assert_head(config, engine) + _assert_countries_shape(engine) + + command.downgrade(config, COUNTRIES_PREVIOUS_REVISION) + _assert_revision(engine, COUNTRIES_PREVIOUS_REVISION) + _assert_pre_countries_shape(engine) + _upgrade_and_assert_head(config, engine) + _assert_countries_shape(engine) command.downgrade(config, STAGE_12A_PREVIOUS_REVISION) _assert_revision(engine, STAGE_12A_PREVIOUS_REVISION) @@ -396,10 +473,10 @@ def test_empty_upgrade_check_base_downgrade_and_reupgrade() -> None: engine.dispose() -# Dormant at the parent revision: its models map these tables, and deleting a -# user cascades through them, but no parent-revision code inserts into them. -# They are the only tables Stage 12A may give a new required column. -PARENT_REVISION_READ_ONLY_TABLES = frozenset( +# Dormant at the Stage 12A parent revision: its models map these tables, and +# deleting a user cascades through them, but no parent-revision code inserts +# into them. They are the only tables Stage 12A may give a new required column. +STAGE_12A_PARENT_READ_ONLY_TABLES = frozenset( {"user_simulation_associations", "user_report_associations"} ) USER_ASSOCIATION_TABLES = ( @@ -409,8 +486,18 @@ def test_empty_upgrade_check_base_downgrade_and_reupgrade() -> None: "user_report_associations", ) COUNTRY_COLUMN_NAMES = frozenset({"country", "country_id", "primary_country"}) -# Nullable references that a parent-revision check constraint requires. -REFERENCES_REQUIRED_BY_CHECKS = frozenset({("simulations", "household_id")}) +# Nullable references that a parent-revision check constraint requires: a +# standard user's country, and one of each association's two references. +REFERENCES_REQUIRED_BY_CHECKS = frozenset( + { + ("simulations", "household_id"), + ("users", "primary_country"), + ("user_policies", "policy_id"), + ("user_household_associations", "household_id"), + ("user_simulation_associations", "simulation_id"), + ("user_report_associations", "report_id"), + } +) # Required values that a parent-revision check constraint restricts. CHECKED_VALUES = {("datasets", "year"): 2026} @@ -547,14 +634,33 @@ def _rewrite_mapped_columns(connection, table: sa.Table, row: dict) -> None: assert result.rowcount == 1, table.name +@dataclass(frozen=True) +class PreviousRevisionSchemas: + """A previous revision's reflected schema and head's, after the upgrade. + + ``read_only_tables`` are the previous revision's tables that only new code + writes, so their test rows are inserted through head's table. + """ + + parent: sa.MetaData + head: sa.MetaData + read_only_tables: frozenset[str] + + def insert_table(self, table_name: str) -> sa.Table | None: + """Head's table for a read-only table, otherwise ``None``.""" + if table_name in self.read_only_tables: + return self.head.tables[table_name] + return None + + def _assert_user_deletion_cascades( connection, - parent: sa.MetaData, - head: sa.MetaData, + schemas: PreviousRevisionSchemas, stored_rows: dict[str, dict], sequence: Iterator[int], ) -> None: - """Delete, through the parent's users table, a user who owns Stage 12A rows.""" + """Delete, through the parent's users table, a user who owns association rows.""" + parent, head = schemas.parent, schemas.head users = parent.tables["users"] owner = _insert_and_read_back( connection, users, _row_values(users, stored_rows, sequence) @@ -562,11 +668,7 @@ def _assert_user_deletion_cascades( owner_rows = {**stored_rows, "users": owner} for table_name in USER_ASSOCIATION_TABLES: table = parent.tables[table_name] - head_table = ( - head.tables[table_name] - if table_name in PARENT_REVISION_READ_ONLY_TABLES - else None - ) + head_table = schemas.insert_table(table_name) _insert_and_read_back( connection, table, @@ -596,22 +698,26 @@ def _assert_user_deletion_cascades( assert remaining == 0, table_name -def test_stage_12a_schema_serves_the_previous_revision() -> None: - """Run the previous Public API revision's table usage against head. +def _assert_head_serves_previous_revision( + previous_revision: str, + read_only_tables: frozenset[str], +) -> None: + """Run a previous Public API revision's table usage against head. The parent revision's schema is reflected and then used, after upgrading, to insert a row into every table it maps (supplying only what it requires), read the row back with every mapped column, rewrite those - columns, delete it, and delete a user who owns Stage 12A rows so the + columns, delete it, and delete a user who owns association rows so the cascades run. Head may not give any table the parent revision writes a - new required column without a server default. + new required column without a server default; ``read_only_tables`` are the + parent's tables that only new code writes. """ database_url = _disposable_url() config = _config() engine = create_engine(database_url) try: - command.downgrade(config, STAGE_12A_PREVIOUS_REVISION) + command.downgrade(config, previous_revision) parent = sa.MetaData() parent.reflect(engine) command.upgrade(config, "head") @@ -634,21 +740,18 @@ def test_stage_12a_schema_serves_the_previous_revision() -> None: for table in parent_tables if _new_required_columns(table, head.tables[table.name]) } - assert set(tables_with_new_required_columns) <= ( - PARENT_REVISION_READ_ONLY_TABLES - ), tables_with_new_required_columns + assert set(tables_with_new_required_columns) <= read_only_tables, ( + tables_with_new_required_columns + ) + schemas = PreviousRevisionSchemas(parent, head, read_only_tables) sequence = itertools.count() with engine.connect() as connection: transaction = connection.begin() try: stored_rows: dict[str, dict] = {} for table in parent_tables: - head_table = ( - head.tables[table.name] - if table.name in PARENT_REVISION_READ_ONLY_TABLES - else None - ) + head_table = schemas.insert_table(table.name) stored_rows[table.name] = _insert_and_read_back( connection, table, @@ -660,7 +763,7 @@ def test_stage_12a_schema_serves_the_previous_revision() -> None: _rewrite_mapped_columns(connection, table, stored_rows[table.name]) _assert_user_deletion_cascades( - connection, parent, head, stored_rows, sequence + connection, schemas, stored_rows, sequence ) for table in reversed(parent_tables): @@ -677,6 +780,21 @@ def test_stage_12a_schema_serves_the_previous_revision() -> None: engine.dispose() +def test_stage_12a_schema_serves_the_previous_revision() -> None: + _assert_head_serves_previous_revision( + STAGE_12A_PREVIOUS_REVISION, + STAGE_12A_PARENT_READ_ONLY_TABLES, + ) + + +def test_countries_schema_serves_the_previous_revision() -> None: + """The release before the countries table keeps working during a deploy.""" + _assert_head_serves_previous_revision( + COUNTRIES_PREVIOUS_REVISION, + frozenset(), + ) + + def test_upgrade_to_head_validates_the_resulting_schema_against_metadata() -> None: database_url = _disposable_url() config = _config() @@ -957,3 +1075,176 @@ def reflected(name: str) -> sa.Table: connection.execute(sa.delete(table).where(table.c.id == row_id)) command.upgrade(config, "head") engine.dispose() + + +def test_countries_revision_keeps_existing_rows_and_rejects_unlisted_countries() -> ( + None +): + """Rows saved before the revision survive it, and only listed countries pass.""" + database_url = _disposable_url() + config = _config() + engine = create_engine(database_url) + user_id = uuid4() + + def reflected(name: str) -> sa.Table: + return sa.Table(name, sa.MetaData(), autoload_with=engine) + + try: + command.downgrade(config, COUNTRIES_PREVIOUS_REVISION) + with engine.begin() as connection: + connection.execute( + sa.insert(reflected("users")).values(id=user_id, primary_country="uk") + ) + connection.execute( + sa.insert(reflected("user_report_associations")).values( + id=uuid4(), + user_id=user_id, + country_id="uk", + legacy_report_id=7, + ) + ) + + command.upgrade(config, "head") + command.check(config) + _assert_countries_shape(engine) + + with engine.connect() as connection: + stored_countries = connection.execute( + text( + "SELECT users.primary_country, " + "user_report_associations.country_id FROM users " + "JOIN user_report_associations " + "ON user_report_associations.user_id = users.id " + "WHERE users.id = :user_id" + ), + {"user_id": user_id}, + ).one() + assert tuple(stored_countries) == ("uk", "uk") + + with engine.connect() as connection, pytest.raises(sa.exc.IntegrityError): + connection.execute( + sa.insert(reflected("user_report_associations")).values( + id=uuid4(), + user_id=user_id, + country_id="ca", + legacy_report_id=8, + ) + ) + + with engine.connect() as connection, pytest.raises(sa.exc.IntegrityError): + connection.execute(text("DELETE FROM countries WHERE id = 'uk'")) + finally: + with engine.begin() as connection: + users = reflected("users") + connection.execute(sa.delete(users).where(users.c.id == user_id)) + command.upgrade(config, "head") + engine.dispose() + + +def test_countries_downgrade_refuses_records_of_an_unlisted_country() -> None: + """The restored US-and-UK checks stop a downgrade that would orphan records.""" + database_url = _disposable_url() + config = _config() + engine = create_engine(database_url) + user_id = uuid4() + + try: + command.upgrade(config, "head") + with engine.begin() as connection: + connection.execute( + text( + "INSERT INTO countries " + "(id, household_enabled, society_wide_enabled) " + "VALUES ('be', true, false)" + ) + ) + connection.execute( + text("INSERT INTO users (id, primary_country) VALUES (:user_id, 'be')"), + {"user_id": user_id}, + ) + + with pytest.raises(sa.exc.IntegrityError): + command.downgrade(config, COUNTRIES_PREVIOUS_REVISION) + + with engine.connect() as connection: + context = MigrationContext.configure(connection) + assert context.get_current_revision() == V2_ALEMBIC_HEAD_REVISION + finally: + with engine.begin() as connection: + connection.execute( + text("DELETE FROM users WHERE id = :user_id"), {"user_id": user_id} + ) + connection.execute(text("DELETE FROM countries WHERE id = 'be'")) + command.upgrade(config, "head") + engine.dispose() + + +def test_a_household_only_country_row_lets_records_name_that_country() -> None: + """Once a household-only country has a row, its records can be stored.""" + database_url = _disposable_url() + config = _config() + engine = create_engine(database_url) + user_id = uuid4() + household_id = uuid4() + + try: + command.upgrade(config, "head") + with engine.connect() as connection: + transaction = connection.begin() + try: + with pytest.raises(sa.exc.IntegrityError), connection.begin_nested(): + connection.execute( + text( + "INSERT INTO users (id, primary_country) " + "VALUES (:user_id, 'be')" + ), + {"user_id": user_id}, + ) + + connection.execute( + text( + "INSERT INTO countries " + "(id, household_enabled, society_wide_enabled) " + "VALUES ('be', true, false)" + ) + ) + connection.execute( + text( + "INSERT INTO users (id, primary_country) " + "VALUES (:user_id, 'be')" + ), + {"user_id": user_id}, + ) + connection.execute( + text( + "INSERT INTO households (id, country_id, household_data, " + "canonicalization_version, content_hash) VALUES " + "(:household_id, 'be', '{}'::jsonb, 1, :content_hash)" + ), + {"household_id": household_id, "content_hash": "c" * 64}, + ) + connection.execute( + text( + "INSERT INTO user_household_associations " + "(id, user_id, household_id, country_id) VALUES " + "(:association_id, :user_id, :household_id, 'be')" + ), + { + "association_id": uuid4(), + "user_id": user_id, + "household_id": household_id, + }, + ) + + with pytest.raises(sa.exc.IntegrityError), connection.begin_nested(): + connection.execute( + text( + "INSERT INTO countries " + "(id, household_enabled, society_wide_enabled) " + "VALUES ('xx', false, true)" + ) + ) + finally: + transaction.rollback() + finally: + engine.dispose() diff --git a/tests/unit/v2/test_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index 4c0e6ee09..56e31410b 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -40,6 +40,7 @@ V2_TABLE_NAMES = frozenset(table.name for table in V2_METADATA.tables.values()) POST_BASELINE_TABLE_NAMES = frozenset( { + "countries", "legacy_household_mappings", "legacy_policy_mappings", "legacy_user_mappings", @@ -235,6 +236,7 @@ def test_v2_revision_chain_has_generated_policy_and_user_identity_changes() -> N assert script.get_heads() == [V2_ALEMBIC_HEAD_REVISION] assert [revision.revision for revision in script.walk_revisions()] == [ V2_ALEMBIC_HEAD_REVISION, + "fbb7488ebe8a", "91af4fc8d8c1", "60d6518b6a98", "439303be14fe", @@ -495,7 +497,6 @@ def test_stage_12a_user_identity_and_association_revision_is_generated() -> None / "migrations/v2/versions/fbb7488ebe8a_link_signed_in_users_to_v1_anchored_.py" ).read_text(encoding="utf-8") - assert V2_ALEMBIC_HEAD_REVISION == "fbb7488ebe8a" assert ( "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" in revision @@ -565,6 +566,94 @@ def test_stage_12a_user_identity_and_association_revision_is_generated() -> None ) in upgrade +COUNTRY_REFERENCE_COLUMNS = ( + ("users", "primary_country"), + ("households", "country_id"), + ("policies", "country_id"), + ("user_household_associations", "country_id"), + ("user_policies", "country_id"), + ("user_simulation_associations", "country_id"), + ("user_report_associations", "country_id"), +) +# The generated downgrade restores each check in PostgreSQL's reflected form +# of `` IN ('us', 'uk')``. +US_AND_UK = "ANY (ARRAY['us'::character varying, 'uk'::character varying]::text[])" +REPLACED_COUNTRY_CHECKS = { + "ck_users_primary_country": ("users", f"primary_country::text = {US_AND_UK}"), + "ck_households_country": ("households", f"country_id::text = {US_AND_UK}"), + "ck_policies_country": ("policies", f"country_id::text = {US_AND_UK}"), + "ck_user_household_associations_country": ( + "user_household_associations", + f"country_id::text = {US_AND_UK}", + ), + "ck_user_policies_country": ("user_policies", f"country_id::text = {US_AND_UK}"), + "ck_user_simulation_associations_country": ( + "user_simulation_associations", + f"country_id::text = {US_AND_UK}", + ), + "ck_user_report_associations_country": ( + "user_report_associations", + f"country_id::text = {US_AND_UK}", + ), +} + + +def _countries_revision_source() -> str: + script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) + countries_revision = next( + revision + for revision in script.walk_revisions() + if revision.down_revision == "fbb7488ebe8a" + ) + assert countries_revision.revision == V2_ALEMBIC_HEAD_REVISION + assert countries_revision.path is not None + return Path(countries_revision.path).read_text(encoding="utf-8") + + +def test_countries_revision_inserts_us_and_uk_and_replaces_country_checks() -> None: + revision = _countries_revision_source() + upgrade = revision[ + revision.index("def upgrade()") : revision.index("def downgrade()") + ] + downgrade = revision[revision.index("def downgrade()") :] + + assert ( + "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" + in revision + ) + assert revision.count("op.create_table(") == 1 + assert 'op.create_table(\n "countries"' in upgrade + assert "ck_countries_society_wide_requires_household" in upgrade + assert 'op.drop_table("countries")' in downgrade + # The one hand-written operation, an approved exception to generated-only + # revisions: the US and UK rows, inserted before the foreign keys that + # existing rows must satisfy. The generated downgrade drops them with the + # table. + assert revision.count("Approved hand-written data operation") == 1 + assert revision.count("Post-generation correction") == 0 + assert revision.count("op.bulk_insert(") == 1 + assert "op.execute(" not in revision + inserted_rows = upgrade[ + upgrade.index("op.bulk_insert(") : upgrade.index("op.create_foreign_key(") + ] + for country_id in ("uk", "us"): + assert f'"id": "{country_id}"' in inserted_rows + assert inserted_rows.count('"household_enabled": True') == 2 + assert inserted_rows.count('"society_wide_enabled": True') == 2 + for table_name, column_name in COUNTRY_REFERENCE_COLUMNS: + key_name = f"fk_{table_name}_{column_name}_countries" + assert key_name in upgrade + assert key_name in downgrade + assert upgrade.count('ondelete="RESTRICT"') == len(COUNTRY_REFERENCE_COLUMNS) + for check_name, (table_name, expression) in REPLACED_COUNTRY_CHECKS.items(): + assert check_name in upgrade + assert f'"{table_name}"' in upgrade + assert check_name in downgrade + assert expression in downgrade + assert "op.drop_column(" not in upgrade + assert "op.add_column(" not in upgrade + + def test_alembic_rejects_unknown_missing_and_divergent_history(tmp_path: Path) -> None: original = REPO / "migrations/v2" missing = tmp_path / "missing" diff --git a/tests/unit/v2/test_household_models.py b/tests/unit/v2/test_household_models.py index f0798e35e..bef59c843 100644 --- a/tests/unit/v2/test_household_models.py +++ b/tests/unit/v2/test_household_models.py @@ -18,6 +18,7 @@ User, UserHouseholdAssociation, ) +from tests.fixtures.v2_countries import add_listed_countries def _engine(): @@ -30,6 +31,8 @@ def enable_foreign_keys(dbapi_connection, _connection_record) -> None: cursor.close() V2_METADATA.create_all(engine) + with Session(engine) as session: + add_listed_countries(session) return engine diff --git a/tests/unit/v2/test_import_side_effects.py b/tests/unit/v2/test_import_side_effects.py index 7b753004d..af6913897 100644 --- a/tests/unit/v2/test_import_side_effects.py +++ b/tests/unit/v2/test_import_side_effects.py @@ -69,7 +69,7 @@ def reject_ddl(*args, **kwargs): import sys after = set(pathlib.Path.cwd().iterdir()) assert before == after -assert len(V2_METADATA.tables) == 39 +assert len(V2_METADATA.tables) == 40 assert "policyengine_api.data.v2.catalog.initialization" not in sys.modules assert "policyengine_api.data.v2.policy_migration_qualification" not in sys.modules assert "policyengine_api.data.v2.household_migration_qualification" not in sys.modules diff --git a/tests/unit/v2/test_model_persistence.py b/tests/unit/v2/test_model_persistence.py index be7879424..732fa7d98 100644 --- a/tests/unit/v2/test_model_persistence.py +++ b/tests/unit/v2/test_model_persistence.py @@ -10,7 +10,9 @@ from policyengine_api.data.v2.models import ( V2_METADATA, + Country, Dynamic, + Household, LegacyPolicyMapping, LegacyUserMapping, LegacyUserPolicyMapping, @@ -20,9 +22,11 @@ TaxBenefitModel, TaxBenefitModelVersion, User, + UserHouseholdAssociation, UserPolicy, ) from policyengine_api.data.v2.models.base import DIRECT_SQLALCHEMY_EXCEPTIONS +from tests.fixtures.v2_countries import add_listed_countries def _relational_sqlite_engine(): @@ -35,6 +39,8 @@ def enable_foreign_keys(dbapi_connection, _connection_record) -> None: cursor.close() V2_METADATA.create_all(engine) + with Session(engine) as session: + add_listed_countries(session) return engine @@ -99,9 +105,8 @@ def test_direct_sqlalchemy_categories_are_complete_and_documented() -> None: assert all(DIRECT_SQLALCHEMY_EXCEPTIONS.values()) -def test_user_primary_country_can_change_between_us_and_uk_only() -> None: - engine = create_engine("sqlite://") - V2_METADATA.create_all(engine) +def test_user_primary_country_must_be_a_listed_country() -> None: + engine = _relational_sqlite_engine() with Session(engine) as session: user = User( @@ -126,6 +131,73 @@ def test_user_primary_country_can_change_between_us_and_uk_only() -> None: engine.dispose() +def test_a_country_row_lets_records_use_a_household_only_country() -> None: + engine = _relational_sqlite_engine() + household_only_country = Country( + id="be", + household_enabled=True, + society_wide_enabled=False, + ) + + with Session(engine) as session: + session.add(User(primary_country="be")) + with pytest.raises(IntegrityError): + session.commit() + session.rollback() + + session.add(household_only_country) + session.commit() + user = User(primary_country="be") + household = Household( + country_id="be", + household_data={"people": {}}, + canonicalization_version=1, + content_hash="b" * 64, + ) + session.add_all([user, household]) + session.flush() + session.add( + UserHouseholdAssociation( + user_id=user.id, + household_id=household.id, + country_id="be", + ) + ) + session.commit() + + stored_country = session.exec(select(UserHouseholdAssociation.country_id)).one() + assert stored_country == "be" + + engine.dispose() + + +def test_a_listed_country_cannot_be_deleted_while_records_use_it() -> None: + engine = _relational_sqlite_engine() + + with Session(engine) as session: + session.add(User(primary_country="uk")) + session.commit() + + session.delete(session.get(Country, "uk")) + with pytest.raises(IntegrityError): + session.commit() + + engine.dispose() + + +def test_a_society_wide_country_must_also_be_household_capable() -> None: + engine = _relational_sqlite_engine() + + with Session(engine) as session: + session.add( + Country(id="be", household_enabled=False, society_wide_enabled=True) + ) + with pytest.raises(IntegrityError): + session.commit() + + engine.dispose() + + def test_only_unauthenticated_users_may_omit_a_primary_country() -> None: engine = create_engine("sqlite://") V2_METADATA.create_all(engine) diff --git a/tests/unit/v2/test_models.py b/tests/unit/v2/test_models.py index 20947ffcb..baa0e2590 100644 --- a/tests/unit/v2/test_models.py +++ b/tests/unit/v2/test_models.py @@ -8,6 +8,7 @@ from policyengine_api.data.v1_models import V1Base from policyengine_api.data.v2.models import ( V2_METADATA, + Country, DatasetVersion, Dynamic, Household, @@ -29,6 +30,18 @@ UserType, ) +# Every column that names the country a record belongs to. Each references the +# countries table instead of checking a fixed list of country IDs. +COUNTRY_REFERENCE_COLUMNS = ( + ("users", "primary_country"), + ("households", "country_id"), + ("policies", "country_id"), + ("user_household_associations", "country_id"), + ("user_policies", "country_id"), + ("user_simulation_associations", "country_id"), + ("user_report_associations", "country_id"), +) + RUN_OUTPUT_TABLES = frozenset( { "aggregates", @@ -48,6 +61,7 @@ def test_domain_models_are_grouped_into_topic_scoped_modules() -> None: expected_modules = { + Country: "countries", User: "users", UserIdentity: "users", Policy: "policies", @@ -354,8 +368,8 @@ def test_standard_users_require_a_supported_primary_country() -> None: for constraint in users.constraints if isinstance(constraint, sa.CheckConstraint) } + # The countries table, not a check, lists the countries a user may choose. assert checks == { - "ck_users_primary_country": "primary_country IN ('us', 'uk')", "ck_users_user_type": "user_type IN ('unauthenticated', 'standard')", "ck_users_standard_primary_country": ( "user_type = 'unauthenticated' OR primary_country IS NOT NULL" @@ -474,7 +488,8 @@ def test_named_checks_and_required_indexes_cover_core_invariants() -> None: assert { "ck_regions_required_filter_values", "ck_simulations_type_input", - "ck_users_primary_country", + "ck_users_standard_primary_country", + "ck_countries_society_wide_requires_household", "ck_report_runs_terminal_completion", "ck_parameter_values_single_owner", }.issubset(constraint_names) @@ -607,10 +622,11 @@ def test_simulation_and_report_associations_have_no_cross_service_keys() -> None assert table.c.label.nullable assert { foreign_key.target_fullname for foreign_key in simulations.foreign_keys - } == {"users.id"} + } == {"users.id", "countries.id"} assert {foreign_key.target_fullname for foreign_key in reports.foreign_keys} == { "users.id", "user_policies.id", + "countries.id", } user_policy_key = next(iter(reports.c.user_policy_id.foreign_keys)) assert user_policy_key.ondelete == "SET NULL" @@ -635,3 +651,71 @@ def test_user_policy_records_v3_reform_fields() -> None: assert "current-policy" in str(baseline_check.sqltext) postgres_type = policies.c.provenance.type.dialect_impl(postgresql.dialect()) assert isinstance(postgres_type, postgresql.JSONB) + + +def test_countries_record_household_and_society_wide_capability() -> None: + countries = V2_METADATA.tables["countries"] + + assert countries.c.keys() == ["id", "household_enabled", "society_wide_enabled"] + assert countries.primary_key.columns.keys() == ["id"] + assert countries.primary_key.name == "pk_countries" + assert countries.c.id.type.length == 2 + for flag_name in ("household_enabled", "society_wide_enabled"): + assert isinstance(countries.c[flag_name].type, sa.Boolean) + assert not countries.c[flag_name].nullable + assert countries.c[flag_name].server_default is None + checks = { + constraint.name: str(constraint.sqltext) + for constraint in countries.constraints + if isinstance(constraint, sa.CheckConstraint) + } + # Society-wide reports need household calculations, so a society-wide + # country is always household capable too. + assert checks == { + "ck_countries_society_wide_requires_household": ( + "household_enabled OR NOT society_wide_enabled" + ), + } + + +def test_country_columns_reference_the_countries_table_instead_of_a_fixed_list() -> ( + None +): + for table_name, column_name in COUNTRY_REFERENCE_COLUMNS: + table = V2_METADATA.tables[table_name] + country_keys = [ + foreign_key + for foreign_key in table.c[column_name].foreign_keys + if foreign_key.column.table.name == "countries" + ] + assert len(country_keys) == 1, table_name + assert country_keys[0].target_fullname == "countries.id" + assert country_keys[0].ondelete == "RESTRICT" + assert country_keys[0].constraint.name == ( + f"fk_{table_name}_{column_name}_countries" + ) + fixed_country_checks = [ + constraint.name + for constraint in table.constraints + if isinstance(constraint, sa.CheckConstraint) + and "'us', 'uk'" in str(constraint.sqltext) + ] + assert fixed_country_checks == [], table_name + + +def test_v1_copy_mappings_keep_their_us_and_uk_checks() -> None: + # These tables pair v1 records with their v2 copies. Each already + # references a v2 record of the same country, and they are removed with + # dual-write, so they keep their fixed checks. + for table_name in ( + "legacy_household_mappings", + "legacy_policy_mappings", + "legacy_user_policy_mappings", + ): + table = V2_METADATA.tables[table_name] + assert f"ck_{table_name}_country" in { + constraint.name for constraint in table.constraints + } + assert "countries.id" not in { + foreign_key.target_fullname for foreign_key in table.foreign_keys + } diff --git a/tests/unit/v2/user_service_support.py b/tests/unit/v2/user_service_support.py index ce54237e3..209a63dc5 100644 --- a/tests/unit/v2/user_service_support.py +++ b/tests/unit/v2/user_service_support.py @@ -31,6 +31,7 @@ from policyengine_api.services.v2.legacy_refs.validators import ( require_legacy_references, ) +from tests.fixtures.v2_countries import add_listed_countries if TYPE_CHECKING: from collections.abc import Collection, Iterator, Mapping @@ -128,6 +129,8 @@ def begin_transaction(connection) -> None: V2_METADATA.create_all(engine) sessions = sessionmaker(engine, class_=Session, expire_on_commit=False) + with sessions() as session: + add_listed_countries(session) with sessions.begin() as session: model = TaxBenefitModel(name="policyengine-us") version = TaxBenefitModelVersion( From b7cb97c31752bc1c8f54e3069379e654b13c51cc Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:10:51 +0400 Subject: [PATCH 03/37] Add a check that a deployment calls its own environment's simulation service The new script requests /health from the Simulation Entrypoint URL that SIM_ENTRYPOINT selects, reads the Cloud Run revision from the x-policyengine-simulation-revision header, and removes the revision's number and suffix to get the service name. Staging must be served by policyengine-simulation-entry-staging and production by policyengine-simulation-entry; a missing header, a name that is not a Cloud Run revision, a failed request or an unknown environment fails. The shared old gateway reports no environment and is not checked. Co-Authored-By: Claude Opus 5.5 --- ...check_simulation_entrypoint_environment.sh | 86 +++++++ ...imulation_entrypoint_environment_script.py | 216 ++++++++++++++++++ 2 files changed, 302 insertions(+) create mode 100755 .github/scripts/check_simulation_entrypoint_environment.sh create mode 100644 tests/unit/test_simulation_entrypoint_environment_script.py diff --git a/.github/scripts/check_simulation_entrypoint_environment.sh b/.github/scripts/check_simulation_entrypoint_environment.sh new file mode 100755 index 000000000..7b6555949 --- /dev/null +++ b/.github/scripts/check_simulation_entrypoint_environment.sh @@ -0,0 +1,86 @@ +#!/usr/bin/env bash + +# Check that a deployment environment calls its own Simulation Entrypoint: +# staging must call the staging service and production the production +# service. The entrypoint names the Cloud Run revision that served a request +# in its x-policyengine-simulation-revision response header, and a Cloud Run +# revision name is its service name followed by a five-digit number and a +# three-character suffix. +# +# The old gateway (SIM_ENTRYPOINT=old_gateway_direct) is one shared Modal app +# that reports no environment, so it is not checked. +# +# Required env vars: DEPLOYMENT_ENVIRONMENT (staging or production), +# SIM_ENTRYPOINT, and the URL variable that SIM_ENTRYPOINT selects. +# Exit codes: 0 when the environments match, 1 when they do not or the +# response names no revision, 2 for an unknown DEPLOYMENT_ENVIRONMENT, and +# curl's own exit code when the request fails. + +set -euo pipefail + +source .github/scripts/simulation_entrypoint_env.sh + +deployment_environment="${DEPLOYMENT_ENVIRONMENT:-}" +case "${deployment_environment}" in + staging) + expected_service="policyengine-simulation-entry-staging" + ;; + production) + expected_service="policyengine-simulation-entry" + ;; + *) + printf 'DEPLOYMENT_ENVIRONMENT must be staging or production, not %q\n' \ + "${deployment_environment}" >&2 + exit 2 + ;; +esac + +if [[ "${SIM_ENTRYPOINT:-}" == "old_gateway_direct" ]]; then + printf 'SIM_ENTRYPOINT=old_gateway_direct: the shared old gateway reports no environment, so there is nothing to check\n' + exit 0 +fi + +simulation_url="$(simulation_entrypoint_selected_url "${SIM_ENTRYPOINT:-}")" +simulation_url="${simulation_url%/}" + +health_response="$( + curl --fail --silent --show-error --include \ + --connect-timeout 10 --max-time 60 \ + "${simulation_url}/health" +)" + +# With --include, curl prints every response's headers; the last revision +# header belongs to the final response. +served_revision="$( + awk ' + tolower($1) == "x-policyengine-simulation-revision:" { + gsub("\r", "", $2) + print $2 + } + ' <<<"${health_response}" | + tail -n 1 +)" + +if [[ -z "${served_revision}" ]]; then + printf '%s/health did not name a Simulation Entrypoint revision; the %s deployment must call the stable URL of %s\n' \ + "${simulation_url}" "${deployment_environment}" "${expected_service}" >&2 + exit 1 +fi + +if ! [[ "${served_revision}" =~ ^(.+)-[0-9]{5}-[a-z0-9]{3}$ ]]; then + printf '%s, served by %s, is not a Cloud Run revision name\n' \ + "${served_revision}" "${simulation_url}" >&2 + exit 1 +fi +served_service="${BASH_REMATCH[1]}" + +if [[ "${served_service}" != "${expected_service}" ]]; then + printf 'The %s deployment calls %s, which revision %s of %s serves; it must call %s\n' \ + "${deployment_environment}" "${simulation_url}" "${served_revision}" \ + "${served_service}" "${expected_service}" >&2 + exit 1 +fi + +printf 'The %s deployment calls %s, served by revision %s of %s\n' \ + "${deployment_environment}" "${simulation_url}" "${served_revision}" \ + "${served_service}" diff --git a/tests/unit/test_simulation_entrypoint_environment_script.py b/tests/unit/test_simulation_entrypoint_environment_script.py new file mode 100644 index 000000000..7efe842aa --- /dev/null +++ b/tests/unit/test_simulation_entrypoint_environment_script.py @@ -0,0 +1,216 @@ +"""Tests for .github/scripts/check_simulation_entrypoint_environment.sh.""" + +from __future__ import annotations + +import os +import subprocess +from dataclasses import dataclass +from pathlib import Path + +import pytest + +REPO = Path(__file__).resolve().parents[2] +SCRIPT = REPO / ".github/scripts/check_simulation_entrypoint_environment.sh" +SIMULATION_URL = "https://simulation.example.test" +STAGING_REVISION = "policyengine-simulation-entry-staging-00070-abc" +PRODUCTION_REVISION = "policyengine-simulation-entry-00122-low" +HEALTHY_BODY = '{"status":"healthy"}' + + +@dataclass(frozen=True) +class FakeCurlReply: + """What the fake curl prints, and the exit code it returns.""" + + output: str + exit_code: int = 0 + + +# Cloud Run's sample container answers with an HTML page and no revision header. +SAMPLE_CONTAINER_RESPONSE = FakeCurlReply( + "HTTP/2 200\r\ncontent-type: text/html; charset=utf-8\r\n\r\n" +) +# curl --fail exits with 22 when the server answers with an HTTP error. +FAILED_REQUEST = FakeCurlReply("", exit_code=22) + + +def _entrypoint_response(revision: str) -> FakeCurlReply: + return FakeCurlReply( + "HTTP/2 200\r\n" + "content-type: application/json\r\n" + f"x-policyengine-simulation-revision: {revision}\r\n" + "\r\n" + f"{HEALTHY_BODY}" + ) + + +def _install_fake_curl(tmp_path: Path, reply: FakeCurlReply) -> Path: + """Install a fake curl that records its arguments and gives ``reply``.""" + fake_bin_directory = tmp_path / "bin" + fake_bin_directory.mkdir() + response_file = tmp_path / "response.txt" + response_file.write_text(reply.output, encoding="utf-8") + calls_file = tmp_path / "curl-calls.log" + fake_curl = fake_bin_directory / "curl" + fake_curl.write_text( + "#!/usr/bin/env bash\n" + f'echo "$*" >> "{calls_file}"\n' + f'cat "{response_file}"\n' + f"exit {reply.exit_code}\n", + encoding="utf-8", + ) + fake_curl.chmod(0o755) + return fake_bin_directory + + +def _run_check( + tmp_path: Path, + curl_reply: FakeCurlReply, + *, + deployment_environment: str, + sim_entrypoint: str = "cloud_run_simulation_entrypoint", + simulation_url: str = SIMULATION_URL, +) -> tuple[subprocess.CompletedProcess[str], list[str]]: + fake_bin_directory = _install_fake_curl(tmp_path, curl_reply) + result = subprocess.run( + ["bash", str(SCRIPT)], + cwd=REPO, + env={ + "PATH": f"{fake_bin_directory}{os.pathsep}{os.environ['PATH']}", + "DEPLOYMENT_ENVIRONMENT": deployment_environment, + "SIM_ENTRYPOINT": sim_entrypoint, + "SIMULATION_ENTRYPOINT_URL": simulation_url, + "OLD_SIMULATION_GATEWAY_URL": "https://old-gateway.example.test", + }, + text=True, + capture_output=True, + check=False, + ) + calls_file = tmp_path / "curl-calls.log" + curl_calls = calls_file.read_text().splitlines() if calls_file.exists() else [] + return result, curl_calls + + +@pytest.mark.parametrize( + ("deployment_environment", "revision"), + [("staging", STAGING_REVISION), ("production", PRODUCTION_REVISION)], +) +def test_accepts_the_simulation_service_of_the_same_environment( + tmp_path: Path, deployment_environment: str, revision: str +) -> None: + result, curl_calls = _run_check( + tmp_path, + _entrypoint_response(revision), + deployment_environment=deployment_environment, + ) + + assert result.returncode == 0, result.stderr + assert len(curl_calls) == 1 + assert curl_calls[0].endswith(f"{SIMULATION_URL}/health") + assert revision in result.stdout + + +def test_requests_health_once_without_a_doubled_slash(tmp_path: Path) -> None: + result, curl_calls = _run_check( + tmp_path, + _entrypoint_response(STAGING_REVISION), + deployment_environment="staging", + simulation_url=f"{SIMULATION_URL}/", + ) + + assert result.returncode == 0, result.stderr + assert curl_calls[0].endswith(f" {SIMULATION_URL}/health") + + +def test_rejects_staging_calling_the_production_simulation_service( + tmp_path: Path, +) -> None: + result, _ = _run_check( + tmp_path, + _entrypoint_response(PRODUCTION_REVISION), + deployment_environment="staging", + ) + + assert result.returncode == 1 + assert "it must call policyengine-simulation-entry-staging" in result.stderr + assert PRODUCTION_REVISION in result.stderr + + +def test_rejects_production_calling_the_staging_simulation_service( + tmp_path: Path, +) -> None: + result, _ = _run_check( + tmp_path, + _entrypoint_response(STAGING_REVISION), + deployment_environment="production", + ) + + assert result.returncode == 1 + assert "it must call policyengine-simulation-entry\n" in result.stderr + assert STAGING_REVISION in result.stderr + + +def test_rejects_a_url_that_names_no_entrypoint_revision(tmp_path: Path) -> None: + result, _ = _run_check( + tmp_path, SAMPLE_CONTAINER_RESPONSE, deployment_environment="staging" + ) + + assert result.returncode == 1 + assert "did not name a Simulation Entrypoint revision" in result.stderr + + +def test_rejects_a_revision_name_without_a_cloud_run_suffix(tmp_path: Path) -> None: + result, _ = _run_check( + tmp_path, + _entrypoint_response("policyengine-simulation-entry-staging"), + deployment_environment="staging", + ) + + assert result.returncode == 1 + assert "is not a Cloud Run revision name" in result.stderr + + +def test_fails_when_the_health_request_fails(tmp_path: Path) -> None: + result, curl_calls = _run_check( + tmp_path, FAILED_REQUEST, deployment_environment="staging" + ) + + assert result.returncode == 22 + assert len(curl_calls) == 1 + + +def test_rejects_an_unknown_deployment_environment(tmp_path: Path) -> None: + result, curl_calls = _run_check( + tmp_path, + _entrypoint_response(STAGING_REVISION), + deployment_environment="preview", + ) + + assert result.returncode == 2 + assert "DEPLOYMENT_ENVIRONMENT must be staging or production" in result.stderr + assert curl_calls == [] + + +def test_rejects_a_missing_simulation_url(tmp_path: Path) -> None: + result, curl_calls = _run_check( + tmp_path, + _entrypoint_response(STAGING_REVISION), + deployment_environment="staging", + simulation_url="", + ) + + assert result.returncode == 1 + assert "SIMULATION_ENTRYPOINT_URL is required" in result.stderr + assert curl_calls == [] + + +def test_does_not_check_the_shared_old_gateway(tmp_path: Path) -> None: + result, curl_calls = _run_check( + tmp_path, + _entrypoint_response(PRODUCTION_REVISION), + deployment_environment="staging", + sim_entrypoint="old_gateway_direct", + ) + + assert result.returncode == 0, result.stderr + assert curl_calls == [] + assert "old_gateway_direct" in result.stdout From ac9d7a0e7629661eec7210ea9b367b4f416d0da7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:12:44 +0400 Subject: [PATCH 04/37] Check each release calls its own environment's simulation service The staging and production model-version alignment jobs now set DEPLOYMENT_ENVIRONMENT and run check_simulation_entrypoint_environment.sh right after checkout, before the bundle compatibility check. Every staging and production deploy job needs these jobs, so a release whose staging SIMULATION_ENTRYPOINT_URL is served by the production Simulation Entrypoint, or the reverse, stops before it deploys. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/push.yml | 6 +++ ...tion-entrypoint-environment-check.added.md | 1 + docs/migration/cloud-run-operations.md | 15 +++++++ ...imulation_entrypoint_environment_script.py | 40 +++++++++++++++++++ 4 files changed, 62 insertions(+) create mode 100644 changelog.d/simulation-entrypoint-environment-check.added.md diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index af02524d3..e8007f472 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -73,11 +73,14 @@ jobs: && (github.event.head_commit.message == 'Update PolicyEngine API') environment: staging env: + DEPLOYMENT_ENVIRONMENT: staging SIM_ENTRYPOINT: ${{ vars.SIM_ENTRYPOINT }} SIMULATION_ENTRYPOINT_URL: ${{ secrets.SIMULATION_ENTRYPOINT_URL }} steps: - name: Checkout repo uses: actions/checkout@v4 + - name: Check the simulation service belongs to staging + run: bash .github/scripts/check_simulation_entrypoint_environment.sh - name: Install jq run: sudo apt-get install -y jq - name: Set up uv for installed bundle validation @@ -438,11 +441,14 @@ jobs: && (github.event.head_commit.message == 'Update PolicyEngine API') environment: production env: + DEPLOYMENT_ENVIRONMENT: production SIM_ENTRYPOINT: ${{ vars.SIM_ENTRYPOINT }} SIMULATION_ENTRYPOINT_URL: ${{ secrets.SIMULATION_ENTRYPOINT_URL }} steps: - name: Checkout repo uses: actions/checkout@v4 + - name: Check the simulation service belongs to production + run: bash .github/scripts/check_simulation_entrypoint_environment.sh - name: Install jq run: sudo apt-get install -y jq - name: Set up uv for installed bundle validation diff --git a/changelog.d/simulation-entrypoint-environment-check.added.md b/changelog.d/simulation-entrypoint-environment-check.added.md new file mode 100644 index 000000000..2e592c45b --- /dev/null +++ b/changelog.d/simulation-entrypoint-environment-check.added.md @@ -0,0 +1 @@ +Each release now checks, in its staging and production model-version alignment jobs, that the environment's simulation URL is served by that environment's own Simulation Entrypoint service, and fails otherwise; staging must call the staging simulation service, and production the production one. diff --git a/docs/migration/cloud-run-operations.md b/docs/migration/cloud-run-operations.md index a69209d24..4815630e6 100644 --- a/docs/migration/cloud-run-operations.md +++ b/docs/migration/cloud-run-operations.md @@ -259,6 +259,21 @@ not block a direct-Modal deployment. The legacy `SIMULATION_API_URL` secret is intentionally unsupported: its value is opaque, cannot be audited, and must not silently choose a deployment upstream. +Each environment calls its own Simulation Entrypoint. The `staging` secret +holds the stable URL of the `policyengine-simulation-entry-staging` Cloud Run +service, which calls the staging Modal environment; the `production` secret +holds `https://simulation.api.policyengine.org`, served by +`policyengine-simulation-entry`. In each release, the staging and production +model-version alignment jobs first run +`.github/scripts/check_simulation_entrypoint_environment.sh`. It reads the +`x-policyengine-simulation-revision` header from the selected URL's `/health` +response and fails the release unless the revision belongs to that +environment's service, so a release whose `staging` secret points at +production, or the reverse, stops before it deploys. Staging's live simulation +tests therefore run against the staging simulation service and staging Modal +workers. The old gateway is one shared Modal app that reports no environment, +so `old_gateway_direct` is not checked. + ## IAM and bootstrap constraints - The GitHub deploy service account holds `roles/run.developer`: it can deploy to diff --git a/tests/unit/test_simulation_entrypoint_environment_script.py b/tests/unit/test_simulation_entrypoint_environment_script.py index 7efe842aa..3f5ca8e04 100644 --- a/tests/unit/test_simulation_entrypoint_environment_script.py +++ b/tests/unit/test_simulation_entrypoint_environment_script.py @@ -8,9 +8,12 @@ from pathlib import Path import pytest +import yaml REPO = Path(__file__).resolve().parents[2] SCRIPT = REPO / ".github/scripts/check_simulation_entrypoint_environment.sh" +PUSH_WORKFLOW = REPO / ".github/workflows/push.yml" +CHECK_COMMAND = "bash .github/scripts/check_simulation_entrypoint_environment.sh" SIMULATION_URL = "https://simulation.example.test" STAGING_REVISION = "policyengine-simulation-entry-staging-00070-abc" PRODUCTION_REVISION = "policyengine-simulation-entry-00122-low" @@ -214,3 +217,40 @@ def test_does_not_check_the_shared_old_gateway(tmp_path: Path) -> None: assert result.returncode == 0, result.stderr assert curl_calls == [] assert "old_gateway_direct" in result.stdout + + +@pytest.mark.parametrize( + ("job_name", "deployment_environment"), + [ + ("ensure-staging-model-version-aligns-with-sim-api", "staging"), + ("ensure-production-model-version-aligns-with-sim-api", "production"), + ], +) +def test_each_release_checks_its_own_simulation_service_first( + job_name: str, deployment_environment: str +) -> None: + """Each version-alignment job fails if it reaches the other environment. + + The check runs right after checkout, before the job asks the simulation + service which bundle it supports, so a crossed environment fails the + release instead of validating against the wrong service. + """ + job = yaml.safe_load(PUSH_WORKFLOW.read_text(encoding="utf-8"))["jobs"][job_name] + + assert job["environment"] == deployment_environment + assert job["env"]["DEPLOYMENT_ENVIRONMENT"] == deployment_environment + assert job["env"]["SIM_ENTRYPOINT"] == "${{ vars.SIM_ENTRYPOINT }}" + assert ( + job["env"]["SIMULATION_ENTRYPOINT_URL"] + == "${{ secrets.SIMULATION_ENTRYPOINT_URL }}" + ) + step_commands = [step.get("run", "") for step in job["steps"]] + assert step_commands.count(CHECK_COMMAND) == 1 + check_position = step_commands.index(CHECK_COMMAND) + bundle_check_position = step_commands.index( + "bash .github/check-policyengine-bundle-supported.sh" + ) + assert ( + job["steps"][check_position - 1].get("uses", "").startswith("actions/checkout@") + ) + assert check_position < bundle_check_position From 0d30696ebab2e0350454cea3429d98945bbbb58b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:21:36 +0400 Subject: [PATCH 05/37] Serve the countries table at GET /v2/countries `GET /v2/countries` lists every row of the countries table, ordered by ID: each country's `id`, `household_enabled` and `society_wide_enabled`. It takes no query parameters (any is refused with a 422), needs no user credential, and answers 503 when the v2 database is not configured or cannot be reached, and a secret-safe 500 for anything else, through the shared v2 error mapping. Requests are logged under the `metadata` route group. The route follows the v2 layout: a `countries` route package, and a service package with a read connector, a session container, a transformation and a read type. The Public API's runtime role gets SELECT on `countries` in the privilege manifest, which the deploy grants after the schema upgrade. The migration contract registry gains the `country_selection_v2` workflow, and the cross-database integration workflow runs the new integration test. Tests (failed first: the route and service tests could not import `policyengine_api.services.v2.countries`, the OpenAPI test did not find `/v2/countries`, the generated contract files were stale, and the integration workflow did not name the new test): route responses, the query refusal, 405 for other methods, 503 and 500 mapping and the OpenAPI document; the service against SQLite; and, against migrated disposable Postgres as the runtime role, the route returning the two rows the revision inserted. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/v2-integration-check.yml | 4 +- changelog.d/v2-countries-route.added.md | 1 + docs/engineering/migration-contracts.md | 15 +- docs/generated/migration_contracts.json | 24 ++- policyengine_api/data/v2/privileges.py | 5 +- .../fastapi_routes/dependencies.py | 23 ++ .../fastapi_routes/v2/countries/__init__.py | 1 + .../v2/countries/response_models.py | 71 +++++++ .../fastapi_routes/v2/countries/routes.py | 84 ++++++++ policyengine_api/fastapi_routes/v2/routes.py | 4 + policyengine_api/migration_registry.py | 1 + policyengine_api/query_parameters.py | 4 + .../services/v2/countries/__init__.py | 1 + .../countries/database_connectors/__init__.py | 1 + .../v2/countries/database_connectors/reads.py | 12 ++ .../services/v2/countries/database_session.py | 26 +++ .../services/v2/countries/services.py | 29 +++ .../services/v2/countries/transformations.py | 19 ++ .../services/v2/countries/types.py | 19 ++ pyproject.toml | 2 + tests/contract/registry.py | 18 ++ .../test_app_v2_workflow_contracts.py | 2 + tests/integration/test_v2_country_routes.py | 71 +++++++ tests/unit/test_alembic_workflows.py | 1 + .../unit/test_migration_contract_artifacts.py | 5 +- tests/unit/v2/test_country_routes.py | 198 ++++++++++++++++++ tests/unit/v2/test_country_service.py | 49 +++++ tests/unit/v2/test_metadata_routes.py | 1 + 28 files changed, 680 insertions(+), 11 deletions(-) create mode 100644 changelog.d/v2-countries-route.added.md create mode 100644 policyengine_api/fastapi_routes/v2/countries/__init__.py create mode 100644 policyengine_api/fastapi_routes/v2/countries/response_models.py create mode 100644 policyengine_api/fastapi_routes/v2/countries/routes.py create mode 100644 policyengine_api/services/v2/countries/__init__.py create mode 100644 policyengine_api/services/v2/countries/database_connectors/__init__.py create mode 100644 policyengine_api/services/v2/countries/database_connectors/reads.py create mode 100644 policyengine_api/services/v2/countries/database_session.py create mode 100644 policyengine_api/services/v2/countries/services.py create mode 100644 policyengine_api/services/v2/countries/transformations.py create mode 100644 policyengine_api/services/v2/countries/types.py create mode 100644 tests/integration/test_v2_country_routes.py create mode 100644 tests/unit/v2/test_country_routes.py create mode 100644 tests/unit/v2/test_country_service.py diff --git a/.github/workflows/v2-integration-check.yml b/.github/workflows/v2-integration-check.yml index 3f402ae80..79d264eee 100644 --- a/.github/workflows/v2-integration-check.yml +++ b/.github/workflows/v2-integration-check.yml @@ -77,8 +77,8 @@ jobs: run: uv run coverage run --branch -m pytest -q tests/integration/test_v2_catalog_installed.py env: RUN_V2_CATALOG_COMPATIBILITY: "1" - - name: Test v2 metadata publication and resource routes - run: uv run coverage run -a --branch -m pytest -q tests/integration/test_v2_catalog_publication.py tests/integration/test_v2_metadata_routes.py + - name: Test v2 metadata publication, resource routes and the countries list + run: uv run coverage run -a --branch -m pytest -q tests/integration/test_v2_catalog_publication.py tests/integration/test_v2_metadata_routes.py tests/integration/test_v2_country_routes.py - name: Test v2 policy persistence and immediate v1 mirroring run: uv run coverage run -a --branch -m pytest -q tests/integration/test_v2_policy_persistence.py tests/integration/test_v1_policy_dual_write.py tests/integration/test_mysql_policy_dual_write.py tests/integration/test_v2_user_policy_mirroring.py tests/integration/test_v1_user_policy_dual_write.py - name: Test signed-in identity provisioning and user associations diff --git a/changelog.d/v2-countries-route.added.md b/changelog.d/v2-countries-route.added.md new file mode 100644 index 000000000..66b8d8e46 --- /dev/null +++ b/changelog.d/v2-countries-route.added.md @@ -0,0 +1 @@ +Added `GET /v2/countries`, which lists every row of the countries table: each served country's ID and whether it supports household and society-wide calculations. diff --git a/docs/engineering/migration-contracts.md b/docs/engineering/migration-contracts.md index f39a5e369..f7975a813 100644 --- a/docs/engineering/migration-contracts.md +++ b/docs/engineering/migration-contracts.md @@ -7,8 +7,8 @@ Generated from `policyengine_api/migration_registry.py` and `tests/contract/regi | Metric | Count | | --- | ---: | | route group count | 9 | -| workflow count | 16 | -| request count | 89 | +| workflow count | 17 | +| request count | 90 | | db entity count | 6 | | sim flow count | 3 | @@ -18,7 +18,7 @@ Generated from `policyengine_api/migration_registry.py` and `tests/contract/regi | --- | --- | --- | --- | | `health` | `health`, `simulation-gateway-check`, `liveness-check`, `readiness-check` | `none` | `none` | | `specification` | `specification` | `none` | `none` | -| `metadata` | `metadata`, `datasets`, `economy-options`, `parameter-values`, `parameters`, `regions`, `tax-benefit-model-versions`, `tax-benefit-models`, `variables` | `metadata` | `none` | +| `metadata` | `metadata`, `countries`, `datasets`, `economy-options`, `parameter-values`, `parameters`, `regions`, `tax-benefit-model-versions`, `tax-benefit-models`, `variables` | `metadata` | `none` | | `policy` | `policy`, `policies`, `user-policy`, `user-policies` | `policy` | `none` | | `household` | `household`, `households`, `user-household`, `user-households`, `calculate`, `calculate-full` | `household` | `household` | | `economy` | `economy` | `simulation` | `economy` | @@ -190,6 +190,15 @@ Generated from `policyengine_api/migration_registry.py` and `tests/contract/regi | `GET` | `/us/metadata` | 200 | `metadata` | none | `status`, `result.current_law_id`, `result.economy_options.region`, `result.economy_options.time_period` | | `GET` | `/uk/metadata` | 200 | `metadata` | none | `status`, `result.current_law_id`, `result.economy_options.region`, `result.economy_options.time_period` | +### `country_selection_v2` + +- Current contract: `typed_v2_resources` +- Future owner: Countries table and one country list in code + +| Method | Path | Status | Route group | Required headers | Stable response fields | +| --- | --- | ---: | --- | --- | --- | +| `GET` | `/v2/countries` | 200 | `metadata` | none | `status`, `message`, `result.items` | + ### `metadata_resources_v2_preview` - Current contract: `typed_v2_resources` diff --git a/docs/generated/migration_contracts.json b/docs/generated/migration_contracts.json index 7a0f1b67d..9ee0774f5 100644 --- a/docs/generated/migration_contracts.json +++ b/docs/generated/migration_contracts.json @@ -1,10 +1,10 @@ { "metadata": { "db_entity_count": 6, - "request_count": 89, + "request_count": 90, "route_group_count": 9, "sim_flow_count": 3, - "workflow_count": 16 + "workflow_count": 17 }, "route_groups": [ { @@ -31,6 +31,7 @@ "name": "metadata", "path_segments": [ "metadata", + "countries", "datasets", "economy-options", "parameter-values", @@ -1348,6 +1349,25 @@ } ] }, + { + "current_contract": "typed_v2_resources", + "future_owner_pr": "Countries table and one country list in code", + "name": "country_selection_v2", + "requests": [ + { + "expected_status": 200, + "headers": [], + "method": "GET", + "path": "/v2/countries", + "route_group": "metadata", + "stable_response_fields": [ + "status", + "message", + "result.items" + ] + } + ] + }, { "current_contract": "typed_v2_resources", "future_owner_pr": "Later metadata read cutover and v2 route-prefix removal", diff --git a/policyengine_api/data/v2/privileges.py b/policyengine_api/data/v2/privileges.py index 6fd0cc750..cec2fa66e 100644 --- a/policyengine_api/data/v2/privileges.py +++ b/policyengine_api/data/v2/privileges.py @@ -132,8 +132,9 @@ def _access( # Stage 10 saved-policy mirror. "legacy_user_mappings": _access(runtime=READ_WRITE), "legacy_user_policy_mappings": _access(runtime=READ_WRITE), - # Served countries; the migration that creates the table inserts its rows. - "countries": _access(), + # Served countries; the migration that creates the table inserts its rows, + # and the Public API reads them. + "countries": _access(runtime=READ), # Users and their saved work (Stage 12A). "user_household_associations": _access(runtime=READ_WRITE), "user_identities": _access(runtime=READ_WRITE), diff --git a/policyengine_api/fastapi_routes/dependencies.py b/policyengine_api/fastapi_routes/dependencies.py index 5d5a1a38d..b2fcbc0ff 100644 --- a/policyengine_api/fastapi_routes/dependencies.py +++ b/policyengine_api/fastapi_routes/dependencies.py @@ -22,6 +22,7 @@ from policyengine_api.services.v2.associations.types import ( LegacyReferenceChecker, ) + from policyengine_api.services.v2.countries.types import CountryRead from policyengine_api.services.v2.households.types import ( HouseholdPage, HouseholdRead, @@ -312,6 +313,14 @@ def get_economy_options( ... +class V2CountryResourceService(Protocol): + """Route-independent countries list read for one request.""" + + def list_countries(self) -> list[CountryRead]: + """Return every served country and its capability, ordered by ID.""" + ... + + class V2PolicyResourceService(Protocol): """Route-independent native policy operations for one request.""" @@ -742,6 +751,18 @@ def _default_v2_metadata_reader_factory() -> V2MetadataResourceReader: ) +def _default_v2_country_service_factory() -> V2CountryResourceService: + from policyengine_api.data.v2.database import ( + get_v2_session_factory, + ) + from policyengine_api.services.v2.countries.database_session import ( + CountryDatabaseSession, + ) + from policyengine_api.services.v2.countries.services import V2CountryService + + return V2CountryService(CountryDatabaseSession(get_v2_session_factory())) + + def _default_v2_policy_service_factory() -> V2PolicyResourceService: from policyengine_api.data.v2.database import ( get_v2_session_factory, @@ -917,6 +938,7 @@ class NativeRouteDependencies: metadata_reader_factory: Callable[[], MetadataReader] specification_provider: Callable[[], JSONObject] v2_metadata_reader_factory: Callable[[], V2MetadataResourceReader] | None = None + v2_country_service_factory: Callable[[], V2CountryResourceService] | None = None v2_policy_service_factory: Callable[[], V2PolicyResourceService] | None = None v2_household_service_factory: Callable[[], V2HouseholdResourceService] | None = None v2_user_policy_service_factory: Callable[[], V2UserPolicyResourceService] | None = ( @@ -947,6 +969,7 @@ def defaults(cls) -> NativeRouteDependencies: metadata_reader_factory=_default_metadata_reader_factory, specification_provider=_default_specification_provider, v2_metadata_reader_factory=_default_v2_metadata_reader_factory, + v2_country_service_factory=_default_v2_country_service_factory, v2_policy_service_factory=_default_v2_policy_service_factory, v2_household_service_factory=_default_v2_household_service_factory, v2_user_policy_service_factory=(_default_v2_user_policy_service_factory), diff --git a/policyengine_api/fastapi_routes/v2/countries/__init__.py b/policyengine_api/fastapi_routes/v2/countries/__init__.py new file mode 100644 index 000000000..3d7f8eaf0 --- /dev/null +++ b/policyengine_api/fastapi_routes/v2/countries/__init__.py @@ -0,0 +1 @@ +"""FastAPI adapters for the v2 countries list.""" diff --git a/policyengine_api/fastapi_routes/v2/countries/response_models.py b/policyengine_api/fastapi_routes/v2/countries/response_models.py new file mode 100644 index 000000000..3ea3c7ee8 --- /dev/null +++ b/policyengine_api/fastapi_routes/v2/countries/response_models.py @@ -0,0 +1,71 @@ +"""Strict HTTP response models for the native v2 countries list.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal + +from pydantic import BaseModel, ConfigDict + +from policyengine_api.fastapi_routes.v2.errors import V2ErrorResponse + +if TYPE_CHECKING: + from policyengine_api.services.v2.countries.types import CountryRead + + +class StrictCountryAPIModel(BaseModel): + """Strict response base with dataclass conversion support.""" + + model_config = ConfigDict(extra="forbid", from_attributes=True) + + +class CountryItem(StrictCountryAPIModel): + """One served country and the calculations it supports. + + ``id`` is the value other v2 resources take as ``country_id``. + ``household_enabled`` means household calculations and saved households + are available; ``society_wide_enabled`` means society-wide reports are, + and is only set for a country that is also household enabled. + """ + + id: str + household_enabled: bool + society_wide_enabled: bool + + @classmethod + def from_read(cls, country: CountryRead) -> CountryItem: + """Build the response item from the service read record.""" + return cls.model_validate(country) + + +class CountryListResult(StrictCountryAPIModel): + """Every served country, ordered by ``id``.""" + + items: list[CountryItem] + + +class CountryListResponse(StrictCountryAPIModel): + """Success envelope: ``status`` "ok", a null ``message`` and the countries.""" + + status: Literal["ok"] = "ok" + message: None = None + result: CountryListResult + + +COUNTRY_ERROR_RESPONSES: dict[int | str, dict[str, Any]] = { + 405: { + "model": V2ErrorResponse, + "description": "The countries list supports GET only.", + }, + 422: { + "model": V2ErrorResponse, + "description": "The request has query parameters; the list takes none.", + }, + 500: { + "model": V2ErrorResponse, + "description": "The countries query failed internally.", + }, + 503: { + "model": V2ErrorResponse, + "description": "The v2 database is not configured or cannot be reached.", + }, +} diff --git a/policyengine_api/fastapi_routes/v2/countries/routes.py b/policyengine_api/fastapi_routes/v2/countries/routes.py new file mode 100644 index 000000000..279aaf54f --- /dev/null +++ b/policyengine_api/fastapi_routes/v2/countries/routes.py @@ -0,0 +1,84 @@ +"""Native FastAPI route for the v2 countries list.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Annotated + +from fastapi import APIRouter, Depends +from sqlalchemy.exc import SQLAlchemyError +from starlette.responses import JSONResponse + +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.fastapi_routes.dependencies import ( + _default_v2_country_service_factory, +) +from policyengine_api.fastapi_routes.query_parameters import query_dependency +from policyengine_api.fastapi_routes.v2.countries.response_models import ( + COUNTRY_ERROR_RESPONSES, + CountryItem, + CountryListResponse, + CountryListResult, +) +from policyengine_api.fastapi_routes.v2.errors import V2ErrorMapping, run_v2_operation +from policyengine_api.query_parameters import CountryListQuery + +if TYPE_CHECKING: + from collections.abc import Callable + + from policyengine_api.fastapi_routes.dependencies import ( + NativeRouteDependencies, + V2CountryResourceService, + ) + + +CountryListQueryDependency = Annotated[ + CountryListQuery, Depends(query_dependency(CountryListQuery)) +] + +COUNTRY_ERROR_MAPPINGS = ( + V2ErrorMapping( + (V2ConfigurationError, SQLAlchemyError), + 503, + "V2 countries are unavailable", + ), +) + + +def _service_factory( + dependencies: NativeRouteDependencies, +) -> Callable[[], V2CountryResourceService]: + if dependencies.v2_country_service_factory is not None: + return dependencies.v2_country_service_factory + return _default_v2_country_service_factory + + +def build_v2_country_router( + dependencies: NativeRouteDependencies, +) -> APIRouter: + """Build the countries route without opening a database connection.""" + router = APIRouter(prefix="/v2", responses=COUNTRY_ERROR_RESPONSES) + service_factory = _service_factory(dependencies) + + @router.get( + "/countries", + response_model=CountryListResponse, + summary="List the served countries and the calculations each supports", + ) + def list_countries( + _query: CountryListQueryDependency, + ) -> CountryListResponse | JSONResponse: + def read() -> CountryListResponse: + countries = service_factory().list_countries() + return CountryListResponse( + result=CountryListResult( + items=[CountryItem.from_read(country) for country in countries] + ) + ) + + return run_v2_operation( + read, + COUNTRY_ERROR_MAPPINGS, + fallback_message="V2 countries query failed", + ) + + return router diff --git a/policyengine_api/fastapi_routes/v2/routes.py b/policyengine_api/fastapi_routes/v2/routes.py index 0521719be..fa39d7783 100644 --- a/policyengine_api/fastapi_routes/v2/routes.py +++ b/policyengine_api/fastapi_routes/v2/routes.py @@ -7,6 +7,9 @@ from fastapi import APIRouter, Request from starlette.responses import JSONResponse +from policyengine_api.fastapi_routes.v2.countries.routes import ( + build_v2_country_router, +) from policyengine_api.fastapi_routes.v2.errors import V2ErrorResponse from policyengine_api.fastapi_routes.v2.households.routes import ( build_v2_household_router, @@ -53,6 +56,7 @@ def build_v2_router( router.include_router(build_v2_user_policy_router(dependencies)) router.include_router(build_v2_household_router(dependencies)) router.include_router(build_v2_policy_router(dependencies)) + router.include_router(build_v2_country_router(dependencies)) router.include_router(build_v2_metadata_model_router(dependencies)) router.include_router(build_v2_metadata_parameter_router(dependencies)) router.include_router(build_v2_metadata_geography_router(dependencies)) diff --git a/policyengine_api/migration_registry.py b/policyengine_api/migration_registry.py index e49de6986..5daadd05a 100644 --- a/policyengine_api/migration_registry.py +++ b/policyengine_api/migration_registry.py @@ -37,6 +37,7 @@ class RouteGroupConfig: name="metadata", path_segments=( "metadata", + "countries", "datasets", "economy-options", "parameter-values", diff --git a/policyengine_api/query_parameters.py b/policyengine_api/query_parameters.py index 619c6924b..086d33131 100644 --- a/policyengine_api/query_parameters.py +++ b/policyengine_api/query_parameters.py @@ -214,6 +214,10 @@ class PaginationQuery(StrictQueryParameters): limit: Limit = DEFAULT_QUERY_LIMIT +class CountryListQuery(StrictQueryParameters): + """Query contract for the countries list, which takes no parameters.""" + + class PolicyCreateQuery(CatalogQuery): """Query contract for native immutable policy creation.""" diff --git a/policyengine_api/services/v2/countries/__init__.py b/policyengine_api/services/v2/countries/__init__.py new file mode 100644 index 000000000..f0f7b972b --- /dev/null +++ b/policyengine_api/services/v2/countries/__init__.py @@ -0,0 +1 @@ +"""Application services for the v2 countries list.""" diff --git a/policyengine_api/services/v2/countries/database_connectors/__init__.py b/policyengine_api/services/v2/countries/database_connectors/__init__.py new file mode 100644 index 000000000..01c1fa018 --- /dev/null +++ b/policyengine_api/services/v2/countries/database_connectors/__init__.py @@ -0,0 +1 @@ +"""SQL-facing connector functions for the v2 countries service.""" diff --git a/policyengine_api/services/v2/countries/database_connectors/reads.py b/policyengine_api/services/v2/countries/database_connectors/reads.py new file mode 100644 index 000000000..5741285d5 --- /dev/null +++ b/policyengine_api/services/v2/countries/database_connectors/reads.py @@ -0,0 +1,12 @@ +"""Database selections used by the v2 countries list.""" + +from __future__ import annotations + +from sqlmodel import Session, col, select + +from policyengine_api.data.v2.models import Country + + +def read_countries(session: Session) -> list[Country]: + """Return every countries row, ordered by country ID.""" + return list(session.exec(select(Country).order_by(col(Country.id))).all()) diff --git a/policyengine_api/services/v2/countries/database_session.py b/policyengine_api/services/v2/countries/database_session.py new file mode 100644 index 000000000..455d4aa39 --- /dev/null +++ b/policyengine_api/services/v2/countries/database_session.py @@ -0,0 +1,26 @@ +"""Database session lifetime management for the v2 countries service.""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy.orm import sessionmaker + from sqlmodel import Session + + +class CountryDatabaseSession: + """Open the read sessions the countries service uses.""" + + def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory`` to open every session.""" + self._session_factory = session_factory + + @contextmanager + def read(self) -> Iterator[Session]: + """Yield a session for reads; nothing is committed.""" + with self._session_factory() as session: + yield session diff --git a/policyengine_api/services/v2/countries/services.py b/policyengine_api/services/v2/countries/services.py new file mode 100644 index 000000000..a7373636d --- /dev/null +++ b/policyengine_api/services/v2/countries/services.py @@ -0,0 +1,29 @@ +"""Router-facing service for the v2 countries list.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_api.services.v2.countries.database_connectors.reads import ( + read_countries, +) +from policyengine_api.services.v2.countries.transformations import country_read + +if TYPE_CHECKING: + from policyengine_api.services.v2.countries.database_session import ( + CountryDatabaseSession, + ) + from policyengine_api.services.v2.countries.types import CountryRead + + +class V2CountryService: + """List the served countries with the calculations each supports.""" + + def __init__(self, database_session: CountryDatabaseSession) -> None: + """Read through ``database_session``, which opens one session per call.""" + self._database_session = database_session + + def list_countries(self) -> list[CountryRead]: + """Return every countries row, ordered by country ID.""" + with self._database_session.read() as session: + return [country_read(country) for country in read_countries(session)] diff --git a/policyengine_api/services/v2/countries/transformations.py b/policyengine_api/services/v2/countries/transformations.py new file mode 100644 index 000000000..fd1d65fe2 --- /dev/null +++ b/policyengine_api/services/v2/countries/transformations.py @@ -0,0 +1,19 @@ +"""Conversions between stored countries rows and service read records.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_api.services.v2.countries.types import CountryRead + +if TYPE_CHECKING: + from policyengine_api.data.v2.models import Country + + +def country_read(country: Country) -> CountryRead: + """Convert one stored countries row to the service's read record.""" + return CountryRead( + id=country.id, + household_enabled=country.household_enabled, + society_wide_enabled=country.society_wide_enabled, + ) diff --git a/policyengine_api/services/v2/countries/types.py b/policyengine_api/services/v2/countries/types.py new file mode 100644 index 000000000..b4a1e74fc --- /dev/null +++ b/policyengine_api/services/v2/countries/types.py @@ -0,0 +1,19 @@ +"""Framework-independent data exchanged by the v2 countries layers.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class CountryRead: + """One served country and the calculations it supports. + + ``household_enabled`` allows household calculations and saved households; + ``society_wide_enabled`` allows society-wide reports, and implies + ``household_enabled``. + """ + + id: str + household_enabled: bool + society_wide_enabled: bool diff --git a/pyproject.toml b/pyproject.toml index ed101271e..10a5cd6f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -212,6 +212,8 @@ banned-module-level-imports = [ "policyengine", "policyengine_api.data.v2.database", "policyengine_api.libs.auth0_management", + "policyengine_api.services.v2.countries.database_session", + "policyengine_api.services.v2.countries.services", "policyengine_api.services.v2.households.database_session", "policyengine_api.services.v2.households.services", "policyengine_api.services.v2.legacy_refs.database_session", diff --git a/tests/contract/registry.py b/tests/contract/registry.py index 778c3da26..a04522cbf 100644 --- a/tests/contract/registry.py +++ b/tests/contract/registry.py @@ -587,6 +587,24 @@ class WorkflowContract: ), ), ), + WorkflowContract( + name="country_selection_v2", + current_contract="typed_v2_resources", + future_owner_pr="Countries table and one country list in code", + requests=( + ContractRequest( + method="GET", + path="/v2/countries", + expected_status=200, + stable_response_fields=( + "status", + "message", + "result.items", + ), + route_group="metadata", + ), + ), + ), WorkflowContract( name="metadata_resources_v2_preview", current_contract="typed_v2_resources", diff --git a/tests/contract/test_app_v2_workflow_contracts.py b/tests/contract/test_app_v2_workflow_contracts.py index ff861db14..f74640678 100644 --- a/tests/contract/test_app_v2_workflow_contracts.py +++ b/tests/contract/test_app_v2_workflow_contracts.py @@ -20,6 +20,7 @@ def test_app_v2_workflow_contract_registry_is_complete(): "user_account_v2", "household_calculate", "region_selection", + "country_selection_v2", "metadata_resources_v2_preview", "simulation_submit_poll", "report_create_poll", @@ -31,6 +32,7 @@ def test_app_v2_workflow_contract_registry_is_complete(): "typed_v2_resources" if workflow.name in { + "country_selection_v2", "metadata_resources_v2_preview", "policy_resources_v2", "user_policy_associations_v2", diff --git a/tests/integration/test_v2_country_routes.py b/tests/integration/test_v2_country_routes.py new file mode 100644 index 000000000..f31c28b5d --- /dev/null +++ b/tests/integration/test_v2_country_routes.py @@ -0,0 +1,71 @@ +"""Serve the v2 countries list from a migrated disposable Postgres database. + +The route runs as the Public API's runtime role, so the test proves the role's +declared privileges cover it (rule +``api.v2.runtime-privileges-tested-as-runtime-role``). +""" + +from __future__ import annotations + +from alembic import command +from alembic.config import Config +from fastapi.testclient import TestClient +from flask import Flask +from sqlalchemy.orm import sessionmaker +from sqlalchemy.pool import NullPool +from sqlmodel import Session + +from policyengine_api.asgi_factory import create_asgi_app +from policyengine_api.constants import REPO +from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies +from policyengine_api.migration_flags import ( + RouteImplementation, + RouteImplementationSettings, +) +from policyengine_api.services.v2.countries.database_session import ( + CountryDatabaseSession, +) +from policyengine_api.services.v2.countries.services import V2CountryService +from tests.fixtures.v2_countries import LISTED_COUNTRY_ROWS + + +def test_postgres_countries_list_the_rows_the_revision_inserted( + managed_v2_roles, + v2_engine_as_role, +) -> None: + command.upgrade(Config(str(REPO / "alembic-v2.ini")), "head") + runtime_engine = v2_engine_as_role(managed_v2_roles.runtime, poolclass=NullPool) + sessions = sessionmaker(runtime_engine, class_=Session) + dependencies = NativeRouteDependencies( + readiness_probe=lambda: True, + gateway_client_factory=lambda: None, + metadata_reader_factory=lambda: None, + specification_provider=dict, + v2_country_service_factory=lambda: V2CountryService( + CountryDatabaseSession(sessions) + ), + ) + settings = RouteImplementationSettings( + health=RouteImplementation.FLASK_FALLBACK, + specification=RouteImplementation.FLASK_FALLBACK, + metadata=RouteImplementation.FLASK_FALLBACK, + ) + client = TestClient( + create_asgi_app( + Flask(__name__), + dependencies=dependencies, + route_settings=settings, + ) + ) + + try: + response = client.get("/v2/countries") + finally: + runtime_engine.dispose() + + assert response.status_code == 200 + assert response.json() == { + "status": "ok", + "message": None, + "result": {"items": list(LISTED_COUNTRY_ROWS)}, + } diff --git a/tests/unit/test_alembic_workflows.py b/tests/unit/test_alembic_workflows.py index 797c82848..fa69a8378 100644 --- a/tests/unit/test_alembic_workflows.py +++ b/tests/unit/test_alembic_workflows.py @@ -221,6 +221,7 @@ def test_reusable_v2_integration_check_uses_disposable_databases_and_coverage(): assert "RUN_V2_CATALOG_COMPATIBILITY" in workflow assert "test_v2_catalog_publication.py" in workflow assert "test_v2_metadata_routes.py" in workflow + assert "test_v2_country_routes.py" in workflow assert "test_v2_policy_persistence.py" in workflow assert "test_v1_policy_dual_write.py" in workflow assert "test_mysql_policy_dual_write.py" in workflow diff --git a/tests/unit/test_migration_contract_artifacts.py b/tests/unit/test_migration_contract_artifacts.py index d08090feb..62ddabd71 100644 --- a/tests/unit/test_migration_contract_artifacts.py +++ b/tests/unit/test_migration_contract_artifacts.py @@ -13,8 +13,8 @@ def test_migration_contract_payload_summarizes_route_contracts(): assert payload["version"] == 1 assert payload["metadata"] == { "route_group_count": 9, - "workflow_count": 16, - "request_count": 89, + "workflow_count": 17, + "request_count": 90, "db_entity_count": 6, "sim_flow_count": 3, } @@ -31,6 +31,7 @@ def test_migration_contract_payload_summarizes_route_contracts(): "user_account_v2", "household_calculate", "region_selection", + "country_selection_v2", "metadata_resources_v2_preview", "simulation_submit_poll", "report_create_poll", diff --git a/tests/unit/v2/test_country_routes.py b/tests/unit/v2/test_country_routes.py new file mode 100644 index 000000000..7fd2e2b91 --- /dev/null +++ b/tests/unit/v2/test_country_routes.py @@ -0,0 +1,198 @@ +"""Route coverage for the native v2 countries list.""" + +from __future__ import annotations + +import pytest +from fastapi.testclient import TestClient +from flask import Flask, jsonify +from sqlalchemy.exc import OperationalError + +from policyengine_api.asgi_factory import create_asgi_app +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies +from policyengine_api.migration_flags import ( + RouteImplementation, + RouteImplementationSettings, + infer_route_group, +) +from policyengine_api.services.v2.countries.types import CountryRead + +UNITED_KINGDOM = CountryRead(id="uk", household_enabled=True, society_wide_enabled=True) +UNITED_STATES = CountryRead(id="us", household_enabled=True, society_wide_enabled=True) +HOUSEHOLD_ONLY_COUNTRY = CountryRead( + id="be", household_enabled=True, society_wide_enabled=False +) + + +class FakeCountryService: + """Return fixed countries, or raise ``error``, and count the calls.""" + + def __init__( + self, + countries: list[CountryRead], + *, + error: Exception | None = None, + ) -> None: + self.countries = countries + self.error = error + self.list_calls = 0 + + def list_countries(self) -> list[CountryRead]: + self.list_calls += 1 + if self.error is not None: + raise self.error + return self.countries + + +def _client(service: FakeCountryService) -> tuple[TestClient, dict[str, int]]: + flask_calls = {"count": 0} + flask_app = Flask(__name__) + + @flask_app.route("/", methods=["GET", "POST"]) + def fallback(resource: str): + flask_calls["count"] += 1 + return jsonify({"source": "flask", "resource": resource}) + + dependencies = NativeRouteDependencies( + readiness_probe=lambda: True, + gateway_client_factory=lambda: None, + metadata_reader_factory=lambda: None, + specification_provider=dict, + v2_country_service_factory=lambda: service, + ) + settings = RouteImplementationSettings( + health=RouteImplementation.FLASK_FALLBACK, + specification=RouteImplementation.FLASK_FALLBACK, + metadata=RouteImplementation.FLASK_FALLBACK, + ) + client = TestClient( + create_asgi_app(flask_app, dependencies=dependencies, route_settings=settings), + raise_server_exceptions=False, + ) + return client, flask_calls + + +def test_countries_list_every_country_with_its_capability() -> None: + service = FakeCountryService( + [HOUSEHOLD_ONLY_COUNTRY, UNITED_KINGDOM, UNITED_STATES] + ) + client, flask_calls = _client(service) + + response = client.get("/v2/countries") + + assert response.status_code == 200 + assert response.json() == { + "status": "ok", + "message": None, + "result": { + "items": [ + { + "id": "be", + "household_enabled": True, + "society_wide_enabled": False, + }, + { + "id": "uk", + "household_enabled": True, + "society_wide_enabled": True, + }, + { + "id": "us", + "household_enabled": True, + "society_wide_enabled": True, + }, + ] + }, + } + assert service.list_calls == 1 + assert flask_calls["count"] == 0 + + +def test_countries_list_needs_no_user_credential() -> None: + client, _flask_calls = _client(FakeCountryService([UNITED_STATES])) + + response = client.get("/v2/countries", headers={"Authorization": ""}) + + assert response.status_code == 200 + + +def test_countries_list_rejects_query_parameters() -> None: + service = FakeCountryService([UNITED_STATES]) + client, _flask_calls = _client(service) + + response = client.get("/v2/countries?country_id=us") + + assert response.status_code == 422 + assert response.json()["status"] == "error" + assert service.list_calls == 0 + + +@pytest.mark.parametrize( + "error", + [ + V2ConfigurationError("V2_RUNTIME_DATABASE_URL is required"), + OperationalError("SELECT", {}, Exception("connection refused")), + ], + ids=["configuration", "database"], +) +def test_countries_list_reports_an_unavailable_database_as_503( + error: Exception, +) -> None: + client, _flask_calls = _client(FakeCountryService([], error=error)) + + response = client.get("/v2/countries") + + assert response.status_code == 503 + assert response.json() == { + "status": "error", + "message": "V2 countries are unavailable", + } + + +def test_countries_list_hides_unexpected_failures_behind_a_500() -> None: + client, _flask_calls = _client( + FakeCountryService([], error=RuntimeError("secret detail")) + ) + + response = client.get("/v2/countries") + + assert response.status_code == 500 + assert response.json() == { + "status": "error", + "message": "V2 countries query failed", + } + + +def test_countries_list_supports_get_only() -> None: + service = FakeCountryService([UNITED_STATES]) + client, flask_calls = _client(service) + + response = client.post("/v2/countries", json={}) + + assert response.status_code == 405 + assert service.list_calls == 0 + assert flask_calls["count"] == 0 + + +def test_countries_list_is_documented_in_the_v2_openapi_document() -> None: + client, _flask_calls = _client(FakeCountryService([])) + + document = client.get("/v2/openapi.json").json() + operation = document["paths"]["/v2/countries"]["get"] + + assert operation.get("parameters", []) == [] + assert operation["responses"]["200"]["content"]["application/json"]["schema"] == { + "$ref": "#/components/schemas/CountryListResponse" + } + assert {"422", "500", "503"} <= operation["responses"].keys() + item_schema = document["components"]["schemas"]["CountryItem"] + assert item_schema["required"] == [ + "id", + "household_enabled", + "society_wide_enabled", + ] + assert item_schema["additionalProperties"] is False + + +def test_countries_requests_are_logged_with_the_metadata_route_group() -> None: + assert infer_route_group("/v2/countries") == "metadata" diff --git a/tests/unit/v2/test_country_service.py b/tests/unit/v2/test_country_service.py new file mode 100644 index 000000000..aae0ff982 --- /dev/null +++ b/tests/unit/v2/test_country_service.py @@ -0,0 +1,49 @@ +"""SQLite-backed coverage for the v2 countries read service.""" + +from __future__ import annotations + +from sqlalchemy.orm import sessionmaker +from sqlmodel import Session, create_engine + +from policyengine_api.data.v2.models import V2_METADATA, Country +from policyengine_api.services.v2.countries.database_session import ( + CountryDatabaseSession, +) +from policyengine_api.services.v2.countries.services import V2CountryService +from policyengine_api.services.v2.countries.types import CountryRead +from tests.fixtures.v2_countries import add_listed_countries + + +def _service_with_countries(*extra_countries: Country) -> V2CountryService: + engine = create_engine("sqlite://") + V2_METADATA.create_all(engine) + sessions = sessionmaker(engine, class_=Session) + with sessions() as session: + add_listed_countries(session) + session.add_all(extra_countries) + session.commit() + return V2CountryService(CountryDatabaseSession(sessions)) + + +def test_countries_are_listed_by_id_with_their_capability_flags() -> None: + service = _service_with_countries( + Country(id="be", household_enabled=True, society_wide_enabled=False) + ) + + countries = service.list_countries() + + assert countries == [ + CountryRead(id="be", household_enabled=True, society_wide_enabled=False), + CountryRead(id="uk", household_enabled=True, society_wide_enabled=True), + CountryRead(id="us", household_enabled=True, society_wide_enabled=True), + ] + + +def test_an_empty_countries_table_lists_no_countries() -> None: + engine = create_engine("sqlite://") + V2_METADATA.create_all(engine) + service = V2CountryService( + CountryDatabaseSession(sessionmaker(engine, class_=Session)) + ) + + assert service.list_countries() == [] diff --git a/tests/unit/v2/test_metadata_routes.py b/tests/unit/v2/test_metadata_routes.py index 79ec7c1ed..1c2e542a5 100644 --- a/tests/unit/v2/test_metadata_routes.py +++ b/tests/unit/v2/test_metadata_routes.py @@ -538,6 +538,7 @@ def test_openapi_references_explicit_resource_response_schemas() -> None: "/v2/variables/{variable_id}", } native_paths = { + "/v2/countries", "/v2/households", "/v2/households/{household_id}", "/v2/policies", From f9a19c11e663215740ea7437ca168e58aca2c583 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:30:54 +0400 Subject: [PATCH 06/37] Refuse to start when the countries table and the country list differ `policyengine_api.asgi` now checks, while the application loads, that the v2 `countries` table holds one row for each country in `policyengine_api/supported_countries.py` and no other. A worker that finds a difference raises `CountryListMismatchError`, naming the listed countries without a row and the rows not listed, so it never serves a request; on Cloud Run the revision then fails its health checks. The check runs on Cloud Run (`K_SERVICE` is set) and wherever a v2 runtime database URL or its secret resource is configured. A local or test startup without v2 configuration opens no connection. A missing table, a missing privilege and unusable v2 configuration also stop startup. A database that cannot be reached is logged, by error type only, and startup continues: the v1 routes do not read the v2 database, and their copies into it already continue without it. The countries service gains `verify_countries_match_code_list`, which reads the row IDs through a new read connector and compares them in a validator. Tests (failed first: `policyengine_api.services.v2.countries.validators` did not exist): the comparison in every direction; the service against SQLite; the startup function's selection, each failure that stops startup, and the logged skip without the error's detail; importing the ASGI module with an unusable v2 runtime URL failing without printing the password; and, against migrated disposable Postgres as the runtime role, the check passing on the revision's rows and failing once an extra row exists. Co-Authored-By: Claude Opus 5.5 --- .../v2-countries-startup-check.added.md | 1 + docs/engineering/skills/alembic-migrations.md | 6 +- policyengine_api/asgi.py | 5 + .../v2/countries/database_connectors/reads.py | 7 +- .../services/v2/countries/services.py | 24 ++- .../services/v2/countries/validators.py | 48 +++++ policyengine_api/startup_country_check.py | 96 +++++++++ policyengine_api/supported_countries.py | 3 +- tests/integration/test_v2_country_routes.py | 46 +++- tests/unit/test_startup_country_check.py | 196 ++++++++++++++++++ tests/unit/v2/test_country_service.py | 38 +++- tests/unit/v2/test_country_validators.py | 68 ++++++ 12 files changed, 528 insertions(+), 10 deletions(-) create mode 100644 changelog.d/v2-countries-startup-check.added.md create mode 100644 policyengine_api/services/v2/countries/validators.py create mode 100644 policyengine_api/startup_country_check.py create mode 100644 tests/unit/test_startup_country_check.py create mode 100644 tests/unit/v2/test_country_validators.py diff --git a/changelog.d/v2-countries-startup-check.added.md b/changelog.d/v2-countries-startup-check.added.md new file mode 100644 index 000000000..66f42aad2 --- /dev/null +++ b/changelog.d/v2-countries-startup-check.added.md @@ -0,0 +1 @@ +The Public API refuses to start when the v2 `countries` table and its country list in code name different countries. The check runs on Cloud Run and wherever a v2 runtime database is configured; an unreachable database is logged and skipped. diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 7a5819fa9..2111ace59 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -316,5 +316,7 @@ existing rows, the downgrade, a downgrade refused by another country's records, and the re-upgrade. The table must list exactly the countries in -`policyengine_api/supported_countries.py`. Listing another country later -needs a revision that inserts its row, and the same human migration decision. +`policyengine_api/supported_countries.py`: the Public API refuses to start +when they differ (`policyengine_api/startup_country_check.py`). Listing +another country later needs a revision that inserts its row, and the same +human migration decision. diff --git a/policyengine_api/asgi.py b/policyengine_api/asgi.py index 1f981c56a..fc5809cf1 100644 --- a/policyengine_api/asgi.py +++ b/policyengine_api/asgi.py @@ -10,6 +10,7 @@ from policyengine_api.observability import get_runtime from policyengine_api.readiness import mark_not_ready, mark_ready from policyengine_api.runtime_cache.client import close_runtime_cache_clients +from policyengine_api.startup_country_check import check_countries_table_at_startup from policyengine_api.warmup import run_startup_warmup @@ -24,6 +25,10 @@ def _close_runtime_resources() -> None: shutdown_callback=_close_runtime_resources, ) +# Refuse to start when the v2 countries table and the country list in code +# disagree (see policyengine_api.startup_country_check). +check_countries_table_at_startup() + # Warm the simulation machinery before serving (see policyengine_api.warmup). # POLICYENGINE_API_STARTUP_WARMUP=0 skips it. if os.environ.get("POLICYENGINE_API_STARTUP_WARMUP", "1") != "0": diff --git a/policyengine_api/services/v2/countries/database_connectors/reads.py b/policyengine_api/services/v2/countries/database_connectors/reads.py index 5741285d5..a13a8f3e4 100644 --- a/policyengine_api/services/v2/countries/database_connectors/reads.py +++ b/policyengine_api/services/v2/countries/database_connectors/reads.py @@ -1,4 +1,4 @@ -"""Database selections used by the v2 countries list.""" +"""Database selections used by the v2 countries list and the startup check.""" from __future__ import annotations @@ -10,3 +10,8 @@ def read_countries(session: Session) -> list[Country]: """Return every countries row, ordered by country ID.""" return list(session.exec(select(Country).order_by(col(Country.id))).all()) + + +def read_country_ids(session: Session) -> list[str]: + """Return the ID of every countries row, ordered.""" + return list(session.exec(select(Country.id).order_by(col(Country.id))).all()) diff --git a/policyengine_api/services/v2/countries/services.py b/policyengine_api/services/v2/countries/services.py index a7373636d..96e0ad601 100644 --- a/policyengine_api/services/v2/countries/services.py +++ b/policyengine_api/services/v2/countries/services.py @@ -1,4 +1,4 @@ -"""Router-facing service for the v2 countries list.""" +"""Router-facing and startup service for the v2 countries table.""" from __future__ import annotations @@ -6,10 +6,16 @@ from policyengine_api.services.v2.countries.database_connectors.reads import ( read_countries, + read_country_ids, ) from policyengine_api.services.v2.countries.transformations import country_read +from policyengine_api.services.v2.countries.validators import ( + validate_countries_match_code_list, +) if TYPE_CHECKING: + from collections.abc import Iterable + from policyengine_api.services.v2.countries.database_session import ( CountryDatabaseSession, ) @@ -17,7 +23,7 @@ class V2CountryService: - """List the served countries with the calculations each supports.""" + """List the served countries, and check them against the list in code.""" def __init__(self, database_session: CountryDatabaseSession) -> None: """Read through ``database_session``, which opens one session per call.""" @@ -27,3 +33,17 @@ def list_countries(self) -> list[CountryRead]: """Return every countries row, ordered by country ID.""" with self._database_session.read() as session: return [country_read(country) for country in read_countries(session)] + + def verify_countries_match_code_list( + self, listed_country_ids: Iterable[str] + ) -> None: + """Raise ``CountryListMismatchError`` unless the rows name exactly these. + + The application runs this at startup with the country list in code. + """ + with self._database_session.read() as session: + stored_country_ids = read_country_ids(session) + validate_countries_match_code_list( + stored_country_ids=stored_country_ids, + listed_country_ids=listed_country_ids, + ) diff --git a/policyengine_api/services/v2/countries/validators.py b/policyengine_api/services/v2/countries/validators.py new file mode 100644 index 000000000..fbb30ca5b --- /dev/null +++ b/policyengine_api/services/v2/countries/validators.py @@ -0,0 +1,48 @@ +"""Checks of the stored countries rows against the country list in code.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterable + + +class CountryListMismatchError(RuntimeError): + """The countries table and the country list in code name different countries. + + ``listed_without_row`` holds the listed countries that have no row, and + ``rows_not_listed`` the rows for countries the list does not name, each + sorted by ID. + """ + + def __init__( + self, + *, + listed_without_row: tuple[str, ...], + rows_not_listed: tuple[str, ...], + ) -> None: + """Record both differences and describe them in the message.""" + self.listed_without_row = listed_without_row + self.rows_not_listed = rows_not_listed + super().__init__( + "the countries table does not match the country list in code: " + f"listed without a row: {', '.join(listed_without_row) or 'none'}; " + f"rows not listed: {', '.join(rows_not_listed) or 'none'}" + ) + + +def validate_countries_match_code_list( + *, + stored_country_ids: Iterable[str], + listed_country_ids: Iterable[str], +) -> None: + """Raise ``CountryListMismatchError`` unless both name the same countries.""" + stored = set(stored_country_ids) + listed = set(listed_country_ids) + if stored == listed: + return + raise CountryListMismatchError( + listed_without_row=tuple(sorted(listed - stored)), + rows_not_listed=tuple(sorted(stored - listed)), + ) diff --git a/policyengine_api/startup_country_check.py b/policyengine_api/startup_country_check.py new file mode 100644 index 000000000..43963d077 --- /dev/null +++ b/policyengine_api/startup_country_check.py @@ -0,0 +1,96 @@ +"""Refuse to start when the v2 countries table and the country list differ. + +Which countries the Public API serves is decided in code +(``policyengine_api.supported_countries``). The v2 ``countries`` table records +what each supports, and must hold one row for each listed country and no +other. ``policyengine_api.asgi`` runs this check while the application loads, +so a disagreement stops the worker before it serves a request. + +The check runs only where a v2 runtime database is selected: on Cloud Run +(``K_SERVICE`` is set), or when a v2 runtime database URL or its secret +resource is configured. A local or test startup without v2 configuration +opens no connection. + +A disagreement, a missing table, a missing privilege or unusable v2 +configuration stops startup. A database that cannot be reached is logged and +startup continues: the v1 routes do not read the v2 database, and their +copies into it already continue without it. +""" + +from __future__ import annotations + +import logging +import os +from typing import TYPE_CHECKING, Protocol + +from sqlalchemy.exc import OperationalError + +from policyengine_api.data.v2.settings import ( + V2_RUNTIME_DATABASE_URL, + V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE, +) +from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable, Mapping + +logger = logging.getLogger(__name__) + +CLOUD_RUN_SERVICE = "K_SERVICE" +V2_RUNTIME_SELECTION_VARIABLES = ( + CLOUD_RUN_SERVICE, + V2_RUNTIME_DATABASE_URL, + V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE, +) + + +class CountryListVerifier(Protocol): + """Compare the stored countries rows with a list of country IDs.""" + + def verify_countries_match_code_list( + self, listed_country_ids: Iterable[str] + ) -> None: + """Raise ``CountryListMismatchError`` unless the rows name exactly these.""" + ... + + +def _default_country_list_verifier() -> CountryListVerifier: + from policyengine_api.data.v2.database import get_v2_session_factory + from policyengine_api.services.v2.countries.database_session import ( + CountryDatabaseSession, + ) + from policyengine_api.services.v2.countries.services import V2CountryService + + return V2CountryService(CountryDatabaseSession(get_v2_session_factory())) + + +def v2_runtime_database_selected(environ: Mapping[str, str]) -> bool: + """Whether this process runs on Cloud Run or has a v2 runtime database.""" + return any(environ.get(name, "").strip() for name in V2_RUNTIME_SELECTION_VARIABLES) + + +def check_countries_table_at_startup( + environ: Mapping[str, str] | None = None, + verifier_factory: Callable[[], CountryListVerifier] | None = None, +) -> None: + """Raise unless the countries table names exactly the countries in code. + + Raises ``CountryListMismatchError`` on a disagreement, SQLAlchemy's + ``ProgrammingError`` when the table or the privilege to read it is + missing, and ``V2ConfigurationError`` for unusable v2 configuration. An + unreachable database is logged, without its error detail, and skipped. + """ + values = os.environ if environ is None else environ + if not v2_runtime_database_selected(values): + return + factory = verifier_factory or _default_country_list_verifier + try: + factory().verify_countries_match_code_list(SUPPORTED_COUNTRY_IDS) + except OperationalError as error: + # The error's detail can name the database host, so only its type is + # logged. + logger.warning( + "Skipped the countries table check at startup: the v2 database " + "could not be reached (%s)", + type(error).__name__, + ) diff --git a/policyengine_api/supported_countries.py b/policyengine_api/supported_countries.py index 51a495dd4..15b9d1b60 100644 --- a/policyengine_api/supported_countries.py +++ b/policyengine_api/supported_countries.py @@ -7,7 +7,8 @@ The v2 ``countries`` table records what each country supports, and must hold one row for each listed country and no other: a v2 Alembic revision inserts -the row when a country is added. +the row when a country is added. The application refuses to start when they +differ (``policyengine_api.startup_country_check``). The module imports nothing from the package, so modules that the query models import, such as the catalog package, can import it too. diff --git a/tests/integration/test_v2_country_routes.py b/tests/integration/test_v2_country_routes.py index f31c28b5d..ee6408ddc 100644 --- a/tests/integration/test_v2_country_routes.py +++ b/tests/integration/test_v2_country_routes.py @@ -1,16 +1,18 @@ -"""Serve the v2 countries list from a migrated disposable Postgres database. +"""Serve and check the v2 countries from a migrated disposable Postgres database. -The route runs as the Public API's runtime role, so the test proves the role's -declared privileges cover it (rule +The route and the startup check run as the Public API's runtime role, so the +tests prove the role's declared privileges cover them (rule ``api.v2.runtime-privileges-tested-as-runtime-role``). """ from __future__ import annotations +import pytest from alembic import command from alembic.config import Config from fastapi.testclient import TestClient from flask import Flask +from sqlalchemy import create_engine, text from sqlalchemy.orm import sessionmaker from sqlalchemy.pool import NullPool from sqlmodel import Session @@ -26,6 +28,10 @@ CountryDatabaseSession, ) from policyengine_api.services.v2.countries.services import V2CountryService +from policyengine_api.services.v2.countries.validators import ( + CountryListMismatchError, +) +from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS from tests.fixtures.v2_countries import LISTED_COUNTRY_ROWS @@ -69,3 +75,37 @@ def test_postgres_countries_list_the_rows_the_revision_inserted( "message": None, "result": {"items": list(LISTED_COUNTRY_ROWS)}, } + + +def test_postgres_countries_match_the_code_list_as_the_runtime_role( + managed_v2_roles, + v2_engine_as_role, + disposable_v2_database_url: str, +) -> None: + """The startup check passes on the revision's rows and fails on an extra row.""" + command.upgrade(Config(str(REPO / "alembic-v2.ini")), "head") + runtime_engine = v2_engine_as_role(managed_v2_roles.runtime, poolclass=NullPool) + service = V2CountryService( + CountryDatabaseSession(sessionmaker(runtime_engine, class_=Session)) + ) + owner_engine = create_engine(disposable_v2_database_url, poolclass=NullPool) + + try: + service.verify_countries_match_code_list(SUPPORTED_COUNTRY_IDS) + + with owner_engine.begin() as connection: + connection.execute( + text( + "INSERT INTO countries " + "(id, household_enabled, society_wide_enabled) " + "VALUES ('ca', true, false)" + ) + ) + with pytest.raises(CountryListMismatchError) as raised: + service.verify_countries_match_code_list(SUPPORTED_COUNTRY_IDS) + assert raised.value.rows_not_listed == ("ca",) + finally: + with owner_engine.begin() as connection: + connection.execute(text("DELETE FROM countries WHERE id = 'ca'")) + owner_engine.dispose() + runtime_engine.dispose() diff --git a/tests/unit/test_startup_country_check.py b/tests/unit/test_startup_country_check.py new file mode 100644 index 000000000..53cefa172 --- /dev/null +++ b/tests/unit/test_startup_country_check.py @@ -0,0 +1,196 @@ +"""The startup check that the v2 countries table matches the country list in code.""" + +from __future__ import annotations + +import logging +import os +import subprocess +import sys +from typing import TYPE_CHECKING + +import pytest +from sqlalchemy.exc import OperationalError, ProgrammingError + +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.services.v2.countries.validators import ( + CountryListMismatchError, +) +from policyengine_api.startup_country_check import check_countries_table_at_startup +from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS + +if TYPE_CHECKING: + from collections.abc import Iterable + from pathlib import Path + +V2_RUNTIME_SECRET_RESOURCE = "projects/example/secrets/v2-runtime/versions/latest" + + +class FakeCountryListVerifier: + """Record each verification, then raise ``error`` if one is set.""" + + def __init__(self, error: Exception | None = None) -> None: + self.error = error + self.verified_country_lists: list[tuple[str, ...]] = [] + + def verify_countries_match_code_list( + self, listed_country_ids: Iterable[str] + ) -> None: + self.verified_country_lists.append(tuple(listed_country_ids)) + if self.error is not None: + raise self.error + + +class CountingVerifierFactory: + """Return one verifier and count how often startup asks for it.""" + + def __init__(self, verifier: FakeCountryListVerifier) -> None: + self.verifier = verifier + self.calls = 0 + + def __call__(self) -> FakeCountryListVerifier: + self.calls += 1 + return self.verifier + + +def test_startup_without_a_v2_runtime_database_skips_the_check() -> None: + factory = CountingVerifierFactory(FakeCountryListVerifier()) + + check_countries_table_at_startup( + environ={"FLASK_DEBUG": "1"}, verifier_factory=factory + ) + + assert factory.calls == 0 + + +def test_blank_selection_variables_skip_the_check() -> None: + factory = CountingVerifierFactory(FakeCountryListVerifier()) + + check_countries_table_at_startup( + environ={ + "K_SERVICE": "", + "V2_RUNTIME_DATABASE_URL": " ", + "V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE": "", + }, + verifier_factory=factory, + ) + + assert factory.calls == 0 + + +@pytest.mark.parametrize( + "environ", + [ + {"K_SERVICE": "policyengine-api"}, + {"V2_RUNTIME_DATABASE_URL": "postgresql+psycopg://configured"}, + {"V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE": V2_RUNTIME_SECRET_RESOURCE}, + ], + ids=["cloud-run", "runtime-url", "runtime-secret"], +) +def test_a_selected_v2_runtime_database_is_checked_against_the_code_list( + environ: dict[str, str], +) -> None: + verifier = FakeCountryListVerifier() + + check_countries_table_at_startup( + environ=environ, verifier_factory=CountingVerifierFactory(verifier) + ) + + assert verifier.verified_country_lists == [SUPPORTED_COUNTRY_IDS] + + +@pytest.mark.parametrize( + "error", + [ + CountryListMismatchError(listed_without_row=("us",), rows_not_listed=()), + ProgrammingError( + "SELECT countries.id FROM countries", + {}, + Exception('relation "countries" does not exist'), + ), + ProgrammingError( + "SELECT countries.id FROM countries", + {}, + Exception("permission denied for table countries"), + ), + ], + ids=["mismatch", "missing-table", "missing-privilege"], +) +def test_a_failed_check_stops_startup(error: Exception) -> None: + verifier = FakeCountryListVerifier(error=error) + + with pytest.raises(type(error)): + check_countries_table_at_startup( + environ={"K_SERVICE": "policyengine-api"}, + verifier_factory=CountingVerifierFactory(verifier), + ) + + +def test_unusable_v2_configuration_stops_startup() -> None: + def unusable_configuration() -> FakeCountryListVerifier: + msg = "V2_RUNTIME_DATABASE_URL must name a non-local Postgres host" + raise V2ConfigurationError(msg) + + with pytest.raises(V2ConfigurationError): + check_countries_table_at_startup( + environ={"K_SERVICE": "policyengine-api"}, + verifier_factory=unusable_configuration, + ) + + +def test_an_unreachable_database_is_logged_and_startup_continues( + caplog: pytest.LogCaptureFixture, +) -> None: + verifier = FakeCountryListVerifier( + error=OperationalError( + "SELECT countries.id FROM countries", + {}, + Exception("connection to db.secret-host failed: password=hunter2"), + ) + ) + + with caplog.at_level(logging.WARNING): + check_countries_table_at_startup( + environ={"K_SERVICE": "policyengine-api"}, + verifier_factory=CountingVerifierFactory(verifier), + ) + + assert verifier.verified_country_lists == [SUPPORTED_COUNTRY_IDS] + assert [record.levelno for record in caplog.records] == [logging.WARNING] + message = caplog.records[0].getMessage() + assert "countries table" in message + assert "OperationalError" in message + assert "secret-host" not in message + assert "hunter2" not in message + + +def test_importing_the_asgi_application_runs_the_check(tmp_path: Path) -> None: + """An unusable v2 runtime database URL stops the application from loading.""" + environment = os.environ.copy() + environment.update( + { + "POLICYENGINE_API_STARTUP_WARMUP": "0", + "V2_RUNTIME_DATABASE_URL": ( + "postgresql+psycopg://runtime:hunter2@localhost:5432/postgres" + "?sslmode=require" + ), + "V2_SUPABASE_PROJECT_REF": "abcdefghijklmnopqrst", + "V2_SUPABASE_ENVIRONMENT": "test-startup", + } + ) + for name in ("V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE", "K_SERVICE"): + environment.pop(name, None) + + result = subprocess.run( + [sys.executable, "-c", "import policyengine_api.asgi"], + cwd=tmp_path, + env=environment, + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode != 0 + assert "V2_RUNTIME_DATABASE_URL must name a non-local Postgres host" in ( + result.stderr + ) + assert "hunter2" not in result.stderr diff --git a/tests/unit/v2/test_country_service.py b/tests/unit/v2/test_country_service.py index aae0ff982..253278599 100644 --- a/tests/unit/v2/test_country_service.py +++ b/tests/unit/v2/test_country_service.py @@ -1,7 +1,8 @@ -"""SQLite-backed coverage for the v2 countries read service.""" +"""SQLite-backed coverage for the v2 countries service: the list and the check.""" from __future__ import annotations +import pytest from sqlalchemy.orm import sessionmaker from sqlmodel import Session, create_engine @@ -11,6 +12,10 @@ ) from policyengine_api.services.v2.countries.services import V2CountryService from policyengine_api.services.v2.countries.types import CountryRead +from policyengine_api.services.v2.countries.validators import ( + CountryListMismatchError, +) +from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS from tests.fixtures.v2_countries import add_listed_countries @@ -47,3 +52,34 @@ def test_an_empty_countries_table_lists_no_countries() -> None: ) assert service.list_countries() == [] + + +def test_rows_for_exactly_the_listed_countries_pass_verification() -> None: + service = _service_with_countries() + + service.verify_countries_match_code_list(SUPPORTED_COUNTRY_IDS) + + +def test_a_row_for_an_unlisted_country_fails_verification() -> None: + service = _service_with_countries( + Country(id="ca", household_enabled=True, society_wide_enabled=False) + ) + + with pytest.raises(CountryListMismatchError) as raised: + service.verify_countries_match_code_list(SUPPORTED_COUNTRY_IDS) + + assert raised.value.rows_not_listed == ("ca",) + assert raised.value.listed_without_row == () + + +def test_an_empty_countries_table_fails_verification() -> None: + engine = create_engine("sqlite://") + V2_METADATA.create_all(engine) + service = V2CountryService( + CountryDatabaseSession(sessionmaker(engine, class_=Session)) + ) + + with pytest.raises(CountryListMismatchError) as raised: + service.verify_countries_match_code_list(SUPPORTED_COUNTRY_IDS) + + assert raised.value.listed_without_row == tuple(sorted(SUPPORTED_COUNTRY_IDS)) diff --git a/tests/unit/v2/test_country_validators.py b/tests/unit/v2/test_country_validators.py new file mode 100644 index 000000000..d56db5c00 --- /dev/null +++ b/tests/unit/v2/test_country_validators.py @@ -0,0 +1,68 @@ +"""The comparison of stored countries rows with the country list in code.""" + +from __future__ import annotations + +import pytest + +from policyengine_api.services.v2.countries.validators import ( + CountryListMismatchError, + validate_countries_match_code_list, +) + + +def test_the_same_countries_in_any_order_match() -> None: + validate_countries_match_code_list( + stored_country_ids=["us", "uk"], + listed_country_ids=("uk", "us"), + ) + + +def test_a_listed_country_without_a_row_is_a_mismatch() -> None: + with pytest.raises(CountryListMismatchError) as raised: + validate_countries_match_code_list( + stored_country_ids=["uk"], + listed_country_ids=("uk", "us"), + ) + + assert raised.value.listed_without_row == ("us",) + assert raised.value.rows_not_listed == () + assert str(raised.value) == ( + "the countries table does not match the country list in code: " + "listed without a row: us; rows not listed: none" + ) + + +def test_a_row_for_an_unlisted_country_is_a_mismatch() -> None: + with pytest.raises(CountryListMismatchError) as raised: + validate_countries_match_code_list( + stored_country_ids=["ca", "uk", "us"], + listed_country_ids=("uk", "us"), + ) + + assert raised.value.listed_without_row == () + assert raised.value.rows_not_listed == ("ca",) + + +def test_every_difference_is_reported_in_order() -> None: + with pytest.raises(CountryListMismatchError) as raised: + validate_countries_match_code_list( + stored_country_ids=["ng", "ca"], + listed_country_ids=("us", "uk"), + ) + + assert raised.value.listed_without_row == ("uk", "us") + assert raised.value.rows_not_listed == ("ca", "ng") + assert str(raised.value) == ( + "the countries table does not match the country list in code: " + "listed without a row: uk, us; rows not listed: ca, ng" + ) + + +def test_an_empty_table_misses_every_listed_country() -> None: + with pytest.raises(CountryListMismatchError) as raised: + validate_countries_match_code_list( + stored_country_ids=[], + listed_country_ids=("uk", "us"), + ) + + assert raised.value.listed_without_row == ("uk", "us") From 6b8cbd2bd1dfdeb09534e8b33a39c42dc034fc49 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:55:08 +0400 Subject: [PATCH 07/37] Stop startup only when the countries table answers with a disagreement Review found that the startup check sorted its failures by exception type in a way that stopped startup for a brief Secret Manager outage (reported as unusable v2 configuration) while only logging rejected credentials (a connection error). The check now stops startup only when the database answers and the answer disagrees with the country list: different countries, a missing table, or no privilege to read it. When it cannot get an answer, because the v2 configuration cannot be used, the secret cannot be resolved, or the database cannot be reached or queried, it logs why and startup continues; the readiness check already reports unusable v2 configuration as not ready and recovers once the secret resolves. The rule `api.startup.no-import-side-effects` now says what deployed startup does: it opens no network connections unless the process runs on Cloud Run or has a v2 runtime database configured. The country list and the startup check join the mypy file list and rule `api.quality.mypy-v2`. The country list's docstring and the Alembic skill now say that a deploy runs a country-adding revision before the new release starts, so the previous release refuses to start afterwards, and list the US-and-UK checks a new country must also widen. Tests (failed first: configuration errors and `InterfaceError` were raised, and importing the ASGI module with an unusable v2 URL exited with an error): configuration and secret failures are logged with their message and startup continues, connection and interface errors are logged by type only, and importing the ASGI module with an unusable v2 URL logs the skipped check without the password. Co-Authored-By: Claude Opus 5.5 --- .../v2-countries-startup-check.added.md | 2 +- docs/engineering/skills/alembic-migrations.md | 6 ++- docs/generated/rules.md | 4 +- policyengine_api/startup_country_check.py | 34 ++++++++----- policyengine_api/supported_countries.py | 11 ++++- pyproject.toml | 2 + rules/registry.toml | 8 ++-- tests/unit/test_startup_country_check.py | 48 ++++++++++++++----- 8 files changed, 82 insertions(+), 33 deletions(-) diff --git a/changelog.d/v2-countries-startup-check.added.md b/changelog.d/v2-countries-startup-check.added.md index 66f42aad2..8e30ae192 100644 --- a/changelog.d/v2-countries-startup-check.added.md +++ b/changelog.d/v2-countries-startup-check.added.md @@ -1 +1 @@ -The Public API refuses to start when the v2 `countries` table and its country list in code name different countries. The check runs on Cloud Run and wherever a v2 runtime database is configured; an unreachable database is logged and skipped. +The Public API refuses to start when the v2 `countries` table and its country list in code name different countries, or when the table or the privilege to read it is missing. The check runs on Cloud Run and wherever a v2 runtime database is configured; when it cannot query the table, it logs why and startup continues. diff --git a/docs/engineering/skills/alembic-migrations.md b/docs/engineering/skills/alembic-migrations.md index 2111ace59..0b5c165b4 100644 --- a/docs/engineering/skills/alembic-migrations.md +++ b/docs/engineering/skills/alembic-migrations.md @@ -319,4 +319,8 @@ The table must list exactly the countries in `policyengine_api/supported_countries.py`: the Public API refuses to start when they differ (`policyengine_api/startup_country_check.py`). Listing another country later needs a revision that inserts its row, and the same -human migration decision. +human migration decision. The deploy runs that revision before it starts +the new release, so instances of the previous release that start afterwards, +and a rollback to it, refuse to start until the revision is downgraded. The +module docstring of `supported_countries.py` lists the other checks a new +country must widen. diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 8c3cc5f5d..e7f4bfedb 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -164,7 +164,7 @@ flowchart LR | --- | --- | --- | --- | --- | --- | --- | --- | | `api.catalog.initializer-not-at-startup` | Application startup modules never import the metadata catalog initializer. | mechanical | enforced | yes | pytest `tests/unit/v2/test_catalog_initialization.py` | `pr.yml:test` | **unproven** | | `api.startup.fail-closed-configuration` | Default Cloud SQL startup needs no Supabase configuration, and a selected v2 runtime without its database URL fails closed. | mechanical | enforced | yes | pytest `tests/unit/v2/test_import_side_effects.py` | `pr.yml:test` | **unproven** | -| `api.startup.no-import-side-effects` | Importing v2 modules, starting up and serving a request open no network connections and create no files or runtime SQLite databases. | mechanical | enforced | yes | pytest `tests/unit/v2/test_import_side_effects.py` | `pr.yml:test` | **unproven** | +| `api.startup.no-import-side-effects` | Importing v2 modules opens no network connections, starting up and serving a request open none unless the process runs on Cloud Run or has a v2 runtime database configured, and none of them creates files or runtime SQLite databases. | mechanical | enforced | yes | pytest `tests/unit/v2/test_import_side_effects.py` | `pr.yml:test` | **unproven** | | `api.startup.no-runtime-sqlite-modules` | Runtime SQLite modules are absent from the production package. | mechanical | enforced | yes | pytest `tests/unit/v2/test_import_side_effects.py` | `pr.yml:test` | **unproven** | ```mermaid @@ -294,7 +294,7 @@ flowchart LR | `api.quality.command-line-modules-ruff-exemptions` | Application modules that also run as command-line programs (if __name__ == "__main__") may print their output (T201); every other ruff rule applies to them. | mechanical | enforced | yes | ruff per-file ignore of T201 for `policyengine_api/{data/v2/catalog/initialization,data/v2/household_migration_qualification,data/v2/policy_migration_qualification,data/v2/privilege_reconciliation,data/v2/user_association_migration_qualification,worker_spm_release}.py` | `pr.yml:lint` | — | | `api.quality.command-line-scripts-ruff-exemptions` | Command-line scripts under scripts/ may print their output, run other programs, fetch fixed URLs, stamp today's date and pick random samples (T201, S603, S607, S310, DTZ011, S311); every other ruff rule applies to them. | mechanical | enforced | yes | ruff per-file ignore of DTZ011, S310, S311, S603, S607, T201 for `scripts/**` | `pr.yml:lint` | — | | `api.quality.environment-variables-load` | Required environment variables and credentials load in a deployed-like environment. | mechanical | enforced | yes | CI job `test_env_vars` | `pr.yml:test_env_vars` | — | -| `api.quality.mypy-v2` | Shared query models, FastAPI dependencies, v2 routes and services, and the user sign-in and account-deletion libraries pass mypy with untyped definitions disallowed. | mechanical | enforced | yes | mypy | `pr.yml:typecheck-v2` | — | +| `api.quality.mypy-v2` | Shared query models, FastAPI dependencies, v2 routes and services, the user sign-in and account-deletion libraries, the country list and its startup check pass mypy with untyped definitions disallowed. | mechanical | enforced | yes | mypy | `pr.yml:typecheck-v2` | — | | `api.quality.ruff-format` | All Python files are formatted with ruff format. | mechanical | enforced | yes | ruff format | `pr.yml:lint` | — | | `api.quality.ruff-lint` | All Python code passes ruff check with every ruff rule selected (ALL), with no ratchet, except the ignores in pyproject.toml: COM812 and ISC001, which conflict with ruff format, and D203 and D213, which conflict with the kept D211 and D212. | mechanical | enforced | yes | ruff check (ALL) | `pr.yml:lint` | — | | `api.quality.runtime-evaluated-models-resolve` | Every package model is one ruff treats as runtime-evaluated (a base resolves, through the import in its own module, to a name in runtime-evaluated-base-classes, or is a same-module class that does), and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. | mechanical | enforced | yes | pytest `tests/unit/test_runtime_evaluated_annotations.py` | `pr.yml:test` | **unproven** | diff --git a/policyengine_api/startup_country_check.py b/policyengine_api/startup_country_check.py index 43963d077..69f233988 100644 --- a/policyengine_api/startup_country_check.py +++ b/policyengine_api/startup_country_check.py @@ -11,10 +11,14 @@ resource is configured. A local or test startup without v2 configuration opens no connection. -A disagreement, a missing table, a missing privilege or unusable v2 -configuration stops startup. A database that cannot be reached is logged and -startup continues: the v1 routes do not read the v2 database, and their -copies into it already continue without it. +Startup stops when the database answers and its answer disagrees with the +list: different countries, a missing table, or no privilege to read it. +When the check cannot get an answer, it logs why and startup continues: the +v1 routes do not read the v2 database, and their copies into it already +continue without it. That covers a database that cannot be reached or +refuses the credentials, and v2 configuration that cannot be used, including +a secret that cannot be resolved; the readiness check already reports the +last as not ready, and recovers once the secret resolves. """ from __future__ import annotations @@ -23,11 +27,12 @@ import os from typing import TYPE_CHECKING, Protocol -from sqlalchemy.exc import OperationalError +from sqlalchemy.exc import InterfaceError, OperationalError from policyengine_api.data.v2.settings import ( V2_RUNTIME_DATABASE_URL, V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE, + V2ConfigurationError, ) from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS @@ -75,10 +80,11 @@ def check_countries_table_at_startup( ) -> None: """Raise unless the countries table names exactly the countries in code. - Raises ``CountryListMismatchError`` on a disagreement, SQLAlchemy's + Raises ``CountryListMismatchError`` on a disagreement and SQLAlchemy's ``ProgrammingError`` when the table or the privilege to read it is - missing, and ``V2ConfigurationError`` for unusable v2 configuration. An - unreachable database is logged, without its error detail, and skipped. + missing. Unusable v2 configuration is logged with its message, which + never contains a secret, and a database connection failure by its error + type only, as its detail can name the host; both skip the check. """ values = os.environ if environ is None else environ if not v2_runtime_database_selected(values): @@ -86,11 +92,15 @@ def check_countries_table_at_startup( factory = verifier_factory or _default_country_list_verifier try: factory().verify_countries_match_code_list(SUPPORTED_COUNTRY_IDS) - except OperationalError as error: - # The error's detail can name the database host, so only its type is - # logged. + except V2ConfigurationError as error: + logger.warning( + "Skipped the countries table check at startup: the v2 runtime " + "configuration cannot be used (%s)", + error, + ) + except (OperationalError, InterfaceError) as error: logger.warning( "Skipped the countries table check at startup: the v2 database " - "could not be reached (%s)", + "could not be queried (%s)", type(error).__name__, ) diff --git a/policyengine_api/supported_countries.py b/policyengine_api/supported_countries.py index 15b9d1b60..0c5ad01a1 100644 --- a/policyengine_api/supported_countries.py +++ b/policyengine_api/supported_countries.py @@ -8,7 +8,16 @@ The v2 ``countries`` table records what each country supports, and must hold one row for each listed country and no other: a v2 Alembic revision inserts the row when a country is added. The application refuses to start when they -differ (``policyengine_api.startup_country_check``). +differ (``policyengine_api.startup_country_check``). Because a deploy runs +that revision before it starts the new release, instances of the previous +release that start afterwards, or a rollback to it, refuse to start until +the revision is downgraded. + +A new country also needs the US-and-UK checks widened on the tables that +pair v1 records with their v2 copies (``legacy_policy_mappings``, +``legacy_user_policy_mappings``, ``legacy_household_mappings``) and on the +v1 ``household_mirror_events`` table, which the v1 household insert writes +in its own transaction. The module imports nothing from the package, so modules that the query models import, such as the catalog package, can import it too. diff --git a/pyproject.toml b/pyproject.toml index 10a5cd6f3..ae1ce0e70 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -158,6 +158,8 @@ files = [ "policyengine_api/services/v2", "policyengine_api/libs/user_auth.py", "policyengine_api/libs/auth0_management.py", + "policyengine_api/startup_country_check.py", + "policyengine_api/supported_countries.py", ] explicit_package_bases = true follow_imports = "silent" diff --git a/rules/registry.toml b/rules/registry.toml index e517ec50a..46718196b 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -396,7 +396,7 @@ invocation = "make test" [[rule]] id = "api.startup.no-import-side-effects" category = "startup" -statement = "Importing v2 modules, starting up and serving a request open no network connections and create no files or runtime SQLite databases." +statement = "Importing v2 modules opens no network connections, starting up and serving a request open none unless the process runs on Cloud Run or has a v2 runtime database configured, and none of them creates files or runtime SQLite databases." rationale = "docs/engineering/skills/testing.md#focused-commands" scope = ["policyengine_api/**/*.py"] tier = "mechanical" @@ -919,9 +919,9 @@ invocation = "ruff check ." [[rule]] id = "api.quality.mypy-v2" category = "quality" -statement = "Shared query models, FastAPI dependencies, v2 routes and services, and the user sign-in and account-deletion libraries pass mypy with untyped definitions disallowed." +statement = "Shared query models, FastAPI dependencies, v2 routes and services, the user sign-in and account-deletion libraries, the country list and its startup check pass mypy with untyped definitions disallowed." rationale = "docs/engineering/skills/v2-code-organization.md" -scope = ["policyengine_api/fastapi_routes/v2/**/*.py", "policyengine_api/services/v2/**/*.py", "policyengine_api/libs/user_auth.py", "policyengine_api/libs/auth0_management.py"] +scope = ["policyengine_api/fastapi_routes/v2/**/*.py", "policyengine_api/services/v2/**/*.py", "policyengine_api/libs/user_auth.py", "policyengine_api/libs/auth0_management.py", "policyengine_api/startup_country_check.py", "policyengine_api/supported_countries.py"] tier = "mechanical" status = "enforced" autofix = false @@ -935,6 +935,8 @@ paths = [ "policyengine_api/services/v2", "policyengine_api/libs/user_auth.py", "policyengine_api/libs/auth0_management.py", + "policyengine_api/startup_country_check.py", + "policyengine_api/supported_countries.py", ] [[rule.ci]] workflow = ".github/workflows/pr.yml" diff --git a/tests/unit/test_startup_country_check.py b/tests/unit/test_startup_country_check.py index 53cefa172..bf479c868 100644 --- a/tests/unit/test_startup_country_check.py +++ b/tests/unit/test_startup_country_check.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING import pytest -from sqlalchemy.exc import OperationalError, ProgrammingError +from sqlalchemy.exc import InterfaceError, OperationalError, ProgrammingError from policyengine_api.data.v2.settings import V2ConfigurationError from policyengine_api.services.v2.countries.validators import ( @@ -125,23 +125,44 @@ def test_a_failed_check_stops_startup(error: Exception) -> None: ) -def test_unusable_v2_configuration_stops_startup() -> None: +@pytest.mark.parametrize( + "message", + [ + "V2_RUNTIME_DATABASE_URL must name a non-local Postgres host", + "V2_RUNTIME_DATABASE_URL_SECRET_RESOURCE could not be resolved", + ], + ids=["invalid-url", "secret-unavailable"], +) +def test_unusable_v2_configuration_is_logged_and_left_to_readiness( + message: str, + caplog: pytest.LogCaptureFixture, +) -> None: + """Readiness already reports not ready for it, and recovers when it does.""" + def unusable_configuration() -> FakeCountryListVerifier: - msg = "V2_RUNTIME_DATABASE_URL must name a non-local Postgres host" - raise V2ConfigurationError(msg) + raise V2ConfigurationError(message) - with pytest.raises(V2ConfigurationError): + with caplog.at_level(logging.WARNING): check_countries_table_at_startup( environ={"K_SERVICE": "policyengine-api"}, verifier_factory=unusable_configuration, ) + assert [record.levelno for record in caplog.records] == [logging.WARNING] + assert message in caplog.records[0].getMessage() + +@pytest.mark.parametrize( + "error_type", + [OperationalError, InterfaceError], + ids=["operational", "interface"], +) def test_an_unreachable_database_is_logged_and_startup_continues( + error_type: type[OperationalError | InterfaceError], caplog: pytest.LogCaptureFixture, ) -> None: verifier = FakeCountryListVerifier( - error=OperationalError( + error=error_type( "SELECT countries.id FROM countries", {}, Exception("connection to db.secret-host failed: password=hunter2"), @@ -158,13 +179,13 @@ def test_an_unreachable_database_is_logged_and_startup_continues( assert [record.levelno for record in caplog.records] == [logging.WARNING] message = caplog.records[0].getMessage() assert "countries table" in message - assert "OperationalError" in message + assert error_type.__name__ in message assert "secret-host" not in message assert "hunter2" not in message def test_importing_the_asgi_application_runs_the_check(tmp_path: Path) -> None: - """An unusable v2 runtime database URL stops the application from loading.""" + """The check runs at import; an unusable v2 runtime URL is only logged.""" environment = os.environ.copy() environment.update( { @@ -189,8 +210,9 @@ def test_importing_the_asgi_application_runs_the_check(tmp_path: Path) -> None: check=False, ) - assert result.returncode != 0 - assert "V2_RUNTIME_DATABASE_URL must name a non-local Postgres host" in ( - result.stderr - ) - assert "hunter2" not in result.stderr + # The application's log handler writes structured records to stdout. + output = result.stdout + result.stderr + assert result.returncode == 0, result.stderr + assert "Skipped the countries table check" in output + assert "V2_RUNTIME_DATABASE_URL must name a non-local Postgres host" in output + assert "hunter2" not in output From 975c80fc12df75301bcfa25c3c01240cbf0e014f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:13:15 +0400 Subject: [PATCH 08/37] Keep the startup check's log tests independent of earlier Alembic tests In a full test run, the tests that run Alembic commands load Alembic's logging configuration, which disables every logger that already exists, including the startup check's. Its four log tests then captured nothing and failed, though they passed on their own. An autouse fixture now re-enables the check's logger for each of its tests. Tests (failed first: running the Alembic unit tests before the startup check tests failed all four log tests); with the fixture, both files pass together (46 passed). Co-Authored-By: Claude Opus 5.5 --- tests/unit/test_startup_country_check.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/unit/test_startup_country_check.py b/tests/unit/test_startup_country_check.py index bf479c868..64bd0f219 100644 --- a/tests/unit/test_startup_country_check.py +++ b/tests/unit/test_startup_country_check.py @@ -11,6 +11,7 @@ import pytest from sqlalchemy.exc import InterfaceError, OperationalError, ProgrammingError +from policyengine_api import startup_country_check from policyengine_api.data.v2.settings import V2ConfigurationError from policyengine_api.services.v2.countries.validators import ( CountryListMismatchError, @@ -25,6 +26,16 @@ V2_RUNTIME_SECRET_RESOURCE = "projects/example/secrets/v2-runtime/versions/latest" +@pytest.fixture(autouse=True) +def enabled_startup_check_logger(monkeypatch: pytest.MonkeyPatch) -> None: + """Re-enable the check's logger, which Alembic's logging setup disables. + + Tests that run Alembic commands load its logging configuration, which + disables every logger that already exists, including this module's. + """ + monkeypatch.setattr(startup_country_check.logger, "disabled", False) + + class FakeCountryListVerifier: """Record each verification, then raise ``error`` if one is set.""" From 8417d1c6b3edbd60e1f3d34c9cd304ec01536f3f Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 20:41:45 +0400 Subject: [PATCH 09/37] Record each catalog version's SPM availability from the simulation service Catalog versions gain spm_availability: whether the version offers SPM settings, with their JSON schema and the certified defaults, or {"available": false}. After publishing, catalog initialization reads the simulation service's public versions answer and records it on each published version that has none yet; a US version offers SPM when the service lists a canonical capability with pinned defaults for its PolicyEngine.py version, and other countries never do. A bundle pins its SPM release, so a recorded value never changes, and versions published earlier get it on the next run. The publication step now receives SIMULATION_ENTRYPOINT_URL. The three database steps of initialization are grouped as PublicationSteps. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/seed-v2-database.yml | 3 + .github/workflows/v2-integration-check.yml | 2 +- changelog.d/catalog-spm-availability.added.md | 1 + ...f3e8f662_store_catalog_spm_availability.py | 32 ++++ .../data/v2/catalog/initialization.py | 44 ++++- .../data/v2/catalog/spm_availability.py | 159 ++++++++++++++++++ policyengine_api/data/v2/migration_target.py | 2 +- policyengine_api/data/v2/models/metadata.py | 5 + .../test_v2_catalog_spm_availability.py | 80 +++++++++ tests/unit/v2/test_alembic_v2.py | 3 +- tests/unit/v2/test_catalog_initialization.py | 17 +- .../unit/v2/test_catalog_spm_availability.py | 134 +++++++++++++++ tests/unit/v2/test_metadata_deployment.py | 6 + tests/unit/v2/test_models.py | 9 + 14 files changed, 481 insertions(+), 16 deletions(-) create mode 100644 changelog.d/catalog-spm-availability.added.md create mode 100644 migrations/v2/versions/aeb5f3e8f662_store_catalog_spm_availability.py create mode 100644 policyengine_api/data/v2/catalog/spm_availability.py create mode 100644 tests/integration/test_v2_catalog_spm_availability.py create mode 100644 tests/unit/v2/test_catalog_spm_availability.py diff --git a/.github/workflows/seed-v2-database.yml b/.github/workflows/seed-v2-database.yml index 61af8978c..723b0bcb2 100644 --- a/.github/workflows/seed-v2-database.yml +++ b/.github/workflows/seed-v2-database.yml @@ -69,3 +69,6 @@ jobs: # Reads the certified UK data release, whose repository is private, # to check each new catalog version's region filters. HUGGING_FACE_TOKEN: ${{ secrets.PE_UK_PRIVATE_HF_READ_TOKEN }} + # Each version's SPM availability comes from this environment's + # simulation service's public versions answer. + SIMULATION_ENTRYPOINT_URL: ${{ secrets.SIMULATION_ENTRYPOINT_URL }} diff --git a/.github/workflows/v2-integration-check.yml b/.github/workflows/v2-integration-check.yml index 6bf58aae2..599726735 100644 --- a/.github/workflows/v2-integration-check.yml +++ b/.github/workflows/v2-integration-check.yml @@ -78,7 +78,7 @@ jobs: env: RUN_V2_CATALOG_COMPATIBILITY: "1" - name: Test v2 metadata publication, resource routes and the countries list - run: uv run coverage run -a --branch -m pytest -q tests/integration/test_v2_catalog_publication.py tests/integration/test_v2_metadata_routes.py tests/integration/test_v2_country_routes.py + run: uv run coverage run -a --branch -m pytest -q tests/integration/test_v2_catalog_publication.py tests/integration/test_v2_catalog_spm_availability.py tests/integration/test_v2_metadata_routes.py tests/integration/test_v2_country_routes.py - name: Test v2 policy persistence and immediate v1 mirroring run: uv run coverage run -a --branch -m pytest -q tests/integration/test_v2_policy_persistence.py tests/integration/test_v1_policy_dual_write.py tests/integration/test_mysql_policy_dual_write.py tests/integration/test_v2_user_policy_mirroring.py tests/integration/test_v1_user_policy_dual_write.py - name: Test signed-in identity provisioning and user associations diff --git a/changelog.d/catalog-spm-availability.added.md b/changelog.d/catalog-spm-availability.added.md new file mode 100644 index 000000000..bdf508050 --- /dev/null +++ b/changelog.d/catalog-spm-availability.added.md @@ -0,0 +1 @@ +Each published catalog version now records whether it offers SPM settings, with their schema and certified defaults, from the simulation service's public versions answer; the catalog publication step reads `SIMULATION_ENTRYPOINT_URL`, and versions published earlier get it on the next publication run. diff --git a/migrations/v2/versions/aeb5f3e8f662_store_catalog_spm_availability.py b/migrations/v2/versions/aeb5f3e8f662_store_catalog_spm_availability.py new file mode 100644 index 000000000..a8722ee1f --- /dev/null +++ b/migrations/v2/versions/aeb5f3e8f662_store_catalog_spm_availability.py @@ -0,0 +1,32 @@ +"""store catalog spm availability + +Revision ID: aeb5f3e8f662 +Revises: fb1a3e40d3d6 +Create Date: 2026-10-27 20:34:25.487436 +Generation: uv run alembic -c alembic-v2.ini revision --autogenerate +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +revision: str = "aeb5f3e8f662" +down_revision: str | None = "fb1a3e40d3d6" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "tax_benefit_model_versions", + sa.Column("spm_availability", sa.JSON(), nullable=True), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("tax_benefit_model_versions", "spm_availability") + # ### end Alembic commands ### diff --git a/policyengine_api/data/v2/catalog/initialization.py b/policyengine_api/data/v2/catalog/initialization.py index 978e40e47..036b055d9 100644 --- a/policyengine_api/data/v2/catalog/initialization.py +++ b/policyengine_api/data/v2/catalog/initialization.py @@ -5,6 +5,7 @@ import json import os import sys +from dataclasses import dataclass from typing import TYPE_CHECKING from sqlalchemy import Engine, create_engine @@ -27,6 +28,9 @@ RegionDataVerificationError, verify_region_data, ) +from policyengine_api.data.v2.catalog.spm_availability import ( + record_published_spm_availability, +) from policyengine_api.data.v2.settings import ( V2ConfigurationError, V2DatabaseSettings, @@ -96,25 +100,47 @@ def verify_unpublished_region_data( ) +@dataclass(frozen=True) +class PublicationSteps: + """What initialization does with the database, in order, after extraction. + + ``verify_region_data`` checks the regions' certified data before anything + is written, ``publish`` writes the catalog, and ``record_spm_availability`` + then records each version's SPM availability from the simulation service + (``spm_availability``). + """ + + verify_region_data: Callable[ + [Engine, NormalizedCatalog, Mapping[str, str] | None], None + ] = verify_unpublished_region_data + publish: Callable[[Engine, NormalizedCatalog], PublicationEvidence] = ( + publish_catalog + ) + record_spm_availability: Callable[ + [Engine, NormalizedCatalog, Mapping[str, str] | None], None + ] = record_published_spm_availability + + def initialize_catalog( environ: Mapping[str, str] | None = None, *, extractor: Callable[[], NormalizedCatalog] = extract_installed_catalog, engine_builder: Callable[[V2DatabaseSettings], Engine] = build_data_write_engine, - publisher: Callable[[Engine, NormalizedCatalog], PublicationEvidence] = ( - publish_catalog - ), - region_data_verifier: Callable[ - [Engine, NormalizedCatalog, Mapping[str, str] | None], None - ] = verify_unpublished_region_data, + steps: PublicationSteps | None = None, ) -> PublicationEvidence: - """Load only row-write settings, extract fully, verify, then publish explicitly.""" + """Load only row-write settings, extract fully, verify, then publish explicitly. + + After publishing, record each version's SPM availability. + """ + selected_steps = steps or PublicationSteps() settings = load_v2_data_write_database_settings(environ) catalog = extractor() engine = engine_builder(settings) try: - region_data_verifier(engine, catalog, environ) - return publisher(engine, catalog) + selected_steps.verify_region_data(engine, catalog, environ) + evidence = selected_steps.publish(engine, catalog) + selected_steps.record_spm_availability(engine, catalog, environ) + return evidence finally: engine.dispose() diff --git a/policyengine_api/data/v2/catalog/spm_availability.py b/policyengine_api/data/v2/catalog/spm_availability.py new file mode 100644 index 000000000..673759971 --- /dev/null +++ b/policyengine_api/data/v2/catalog/spm_availability.py @@ -0,0 +1,159 @@ +"""Record whether each published catalog version offers SPM settings. + +The simulation service's public ``GET /versions`` lists, under +``spm_capabilities``, the certified SPM capability of each PolicyEngine.py +bundle it runs: its contract version and pinned defaults. The Public API +already checks economy requests' SPM selections against it. A US catalog +version offers SPM when the service reports a canonical capability for its +PolicyEngine.py version; other countries never do. + +Publication records it after the catalog itself, on each version that has +none yet, so versions published before it existed get it on the next run. A +bundle pins its SPM release, so a recorded value never changes. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING, Any + +import httpx +import sqlalchemy as sa +from pydantic import ValidationError + +from policyengine_api.data.v2.catalog.publication_types import ( + CatalogPublicationError, +) +from policyengine_api.data.v2.models import TaxBenefitModel, TaxBenefitModelVersion +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.spm import SPMSelection + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + + from sqlalchemy import Connection, Engine + + from policyengine_api.data.v2.catalog.records import NormalizedCatalog + +SIMULATION_ENTRYPOINT_URL = "SIMULATION_ENTRYPOINT_URL" +SPM_COUNTRY_ID = "us" +SPM_CONTRACT_VERSION = "canonical-spm-v1" +VERSIONS_TIMEOUT_SECONDS = 30.0 +UNAVAILABLE: dict[str, Any] = {"available": False} +MODELS = TaxBenefitModel.__table__ +MODEL_VERSIONS = TaxBenefitModelVersion.__table__ + + +def spm_availability( + country_id: str, policyengine_version: str, versions: object +) -> dict[str, Any]: + """Return ``{available, settings_schema, defaults}`` for one catalog version. + + Only a US version whose PolicyEngine.py version has a canonical capability + with pinned defaults (forecast artifact and scenario) offers SPM; anything + else is ``{"available": False}``. + """ + if country_id != SPM_COUNTRY_ID or not isinstance(versions, dict): + return dict(UNAVAILABLE) + capabilities = versions.get("spm_capabilities") + capability = ( + capabilities.get(policyengine_version) + if isinstance(capabilities, dict) + else None + ) + if ( + not isinstance(capability, dict) + or set(capability) != {"contract_version", "defaults"} + or capability["contract_version"] != SPM_CONTRACT_VERSION + ): + return dict(UNAVAILABLE) + try: + defaults = SPMSelection.model_validate(capability["defaults"]) + except ValidationError: + return dict(UNAVAILABLE) + if not defaults.forecast_content_sha256 or not defaults.scenario: + return dict(UNAVAILABLE) + return { + "available": True, + "settings_schema": SPMSelection.model_json_schema(), + "defaults": { + name: getattr(defaults, name) for name in SPMSelection.model_fields + }, + } + + +def simulation_entrypoint_url(environ: Mapping[str, str]) -> str: + """Return the simulation service's base URL; raise when it is not set.""" + url = environ.get(SIMULATION_ENTRYPOINT_URL, "").strip().rstrip("/") + if not url: + msg = f"{SIMULATION_ENTRYPOINT_URL} is required to record SPM availability" + raise V2ConfigurationError(msg) + return url + + +def fetch_simulation_versions( + base_url: str, *, client: httpx.Client | None = None +) -> object: + """Read the simulation service's public ``GET /versions`` answer.""" + url = f"{base_url.rstrip('/')}/versions" + try: + if client is None: + response = httpx.get(url, timeout=VERSIONS_TIMEOUT_SECONDS) + else: + response = client.get(url, timeout=VERSIONS_TIMEOUT_SECONDS) + response.raise_for_status() + return response.json() + except (httpx.HTTPError, ValueError) as error: + msg = "the simulation service's versions could not be read" + raise CatalogPublicationError(msg) from error + + +def record_spm_availability( + connection: Connection, + *, + model_name: str, + version: str, + availability: Mapping[str, Any], +) -> bool: + """Store a catalog version's SPM availability unless it already has one. + + Return whether the version was changed. + """ + model_ids = sa.select(MODELS.c.id).where(MODELS.c.name == model_name) + result = connection.execute( + sa.update(MODEL_VERSIONS) + .where( + MODEL_VERSIONS.c.model_id.in_(model_ids), + MODEL_VERSIONS.c.version == version, + MODEL_VERSIONS.c.spm_availability.is_(None), + ) + .values(spm_availability=dict(availability)) + ) + return bool(result.rowcount) + + +def record_published_spm_availability( + engine: Engine, + catalog: NormalizedCatalog, + environ: Mapping[str, str] | None = None, + *, + fetch_versions: Callable[[str], object] = fetch_simulation_versions, +) -> None: + """Record each published country's SPM availability for the catalog version. + + Raise V2ConfigurationError without ``SIMULATION_ENTRYPOINT_URL``, and + CatalogPublicationError when the service's versions cannot be read. + """ + versions = fetch_versions( + simulation_entrypoint_url(os.environ if environ is None else environ) + ) + with engine.begin() as connection: + for country in catalog.countries: + record_spm_availability( + connection, + model_name=country.model.name, + version=country.model_version.version, + availability=spm_availability( + country.country_id, catalog.policyengine_version, versions + ), + ) diff --git a/policyengine_api/data/v2/migration_target.py b/policyengine_api/data/v2/migration_target.py index 2c7b6986d..77e019ec6 100644 --- a/policyengine_api/data/v2/migration_target.py +++ b/policyengine_api/data/v2/migration_target.py @@ -24,7 +24,7 @@ from collections.abc import Mapping V2_ALEMBIC_DISPOSABLE_TEST = "V2_ALEMBIC_DISPOSABLE_TEST" -V2_ALEMBIC_HEAD_REVISION = "fb1a3e40d3d6" +V2_ALEMBIC_HEAD_REVISION = "aeb5f3e8f662" DISPOSABLE_DATABASE_NAME = "policyengine_v2_alembic_test" DISPOSABLE_HOSTS = frozenset({"localhost", "127.0.0.1", "::1", "postgres"}) MIGRATION_ROLE = "policyengine_v2_migrator" diff --git a/policyengine_api/data/v2/models/metadata.py b/policyengine_api/data/v2/models/metadata.py index 5431a6abc..5779d65b9 100644 --- a/policyengine_api/data/v2/models/metadata.py +++ b/policyengine_api/data/v2/models/metadata.py @@ -113,6 +113,11 @@ class TaxBenefitModelVersion(IdentifiedModel, table=True): catalog_format: int | None = None entities: list[CatalogEntity] | None = Field(default=None, sa_type=sa.JSON) basic_inputs: list[str] | None = Field(default=None, sa_type=sa.JSON) + # Whether this version offers SPM settings, with their schema and the + # certified defaults, as the simulation service reports its bundle's SPM + # support: {available, settings_schema, defaults}. Recorded after the + # catalog is published, so a version briefly has none. + spm_availability: dict[str, Any] | None = Field(default=None, sa_type=sa.JSON) model: TaxBenefitModel = Relationship(back_populates="versions") variables: list["Variable"] = Relationship( diff --git a/tests/integration/test_v2_catalog_spm_availability.py b/tests/integration/test_v2_catalog_spm_availability.py new file mode 100644 index 000000000..9df0b3564 --- /dev/null +++ b/tests/integration/test_v2_catalog_spm_availability.py @@ -0,0 +1,80 @@ +"""Record SPM availability on published catalog versions in disposable Postgres.""" + +from __future__ import annotations + +from uuid import uuid4 + +import sqlalchemy as sa +from alembic import command +from alembic.config import Config +from sqlalchemy import create_engine +from sqlalchemy.pool import NullPool + +from policyengine_api.constants import REPO +from policyengine_api.data.v2.catalog.spm_availability import ( + MODEL_VERSIONS, + MODELS, + record_spm_availability, +) + +AVAILABLE = {"available": True, "settings_schema": {}, "defaults": {"scenario": "x"}} +UNAVAILABLE = {"available": False} + + +def test_availability_fills_only_the_named_version_and_only_once( + disposable_v2_database_url: str, +) -> None: + command.upgrade(Config(str(REPO / "alembic-v2.ini")), "head") + engine = create_engine(disposable_v2_database_url, poolclass=NullPool) + model_id, version_id, other_version_id = uuid4(), uuid4(), uuid4() + model_name = f"policyengine-us-{model_id.hex[:8]}" + try: + with engine.begin() as connection: + connection.execute(sa.insert(MODELS).values(id=model_id, name=model_name)) + connection.execute( + sa.insert(MODEL_VERSIONS), + [ + { + "id": identifier, + "model_id": model_id, + "version": version, + "current_law_id": 2, + "metadata_time_periods": [2026], + } + for identifier, version in ( + (version_id, "1.450.0"), + (other_version_id, "1.449.0"), + ) + ], + ) + first = record_spm_availability( + connection, + model_name=model_name, + version="1.450.0", + availability=AVAILABLE, + ) + again = record_spm_availability( + connection, + model_name=model_name, + version="1.450.0", + availability=UNAVAILABLE, + ) + with engine.connect() as connection: + stored = dict( + connection.execute( + sa.select( + MODEL_VERSIONS.c.id, MODEL_VERSIONS.c.spm_availability + ).where(MODEL_VERSIONS.c.model_id == model_id) + ).all() + ) + finally: + with engine.begin() as connection: + connection.execute( + sa.delete(MODEL_VERSIONS).where(MODEL_VERSIONS.c.model_id == model_id) + ) + connection.execute(sa.delete(MODELS).where(MODELS.c.id == model_id)) + engine.dispose() + + assert first + assert not again + assert stored == {version_id: AVAILABLE, other_version_id: None} diff --git a/tests/unit/v2/test_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index b3dbb17f9..fbea9c3ef 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -248,6 +248,7 @@ def test_v2_revision_chain_has_generated_policy_and_user_identity_changes() -> N assert script.get_heads() == [V2_ALEMBIC_HEAD_REVISION] assert [revision.revision for revision in script.walk_revisions()] == [ V2_ALEMBIC_HEAD_REVISION, + "fb1a3e40d3d6", "f5c1d9b4685b", SAVED_POLICY_RENAME_REVISION, SAVED_POLICY_RENAME_PARENT_REVISION, @@ -622,7 +623,7 @@ def _countries_revision_source() -> str: for revision in script.walk_revisions() if revision.down_revision == "f5c1d9b4685b" ) - assert countries_revision.revision == V2_ALEMBIC_HEAD_REVISION + assert countries_revision.revision == "fb1a3e40d3d6" assert countries_revision.path is not None return Path(countries_revision.path).read_text(encoding="utf-8") diff --git a/tests/unit/v2/test_catalog_initialization.py b/tests/unit/v2/test_catalog_initialization.py index d4d6c22c7..aee1342b5 100644 --- a/tests/unit/v2/test_catalog_initialization.py +++ b/tests/unit/v2/test_catalog_initialization.py @@ -78,12 +78,19 @@ def verify(selected_engine, catalog, environ): environ, ) + def record_spm(selected_engine, catalog, environ): + assert "engine" in observed + observed["spm"] = (selected_engine, catalog.policyengine_version, environ) + result = initialization.initialize_catalog( ENVIRONMENT, extractor=normalized_catalog, engine_builder=build, - publisher=publish, - region_data_verifier=verify, + steps=initialization.PublicationSteps( + verify_region_data=verify, + publish=publish, + record_spm_availability=record_spm, + ), ) assert result == _evidence() @@ -92,6 +99,7 @@ def verify(selected_engine, catalog, environ): "verified": (engine, POLICYENGINE_VERSION, ENVIRONMENT), "engine": engine, "version": POLICYENGINE_VERSION, + "spm": (engine, POLICYENGINE_VERSION, ENVIRONMENT), } assert engine.disposed @@ -112,8 +120,9 @@ def refuse_publication(_engine, _catalog): ENVIRONMENT, extractor=normalized_catalog, engine_builder=lambda _settings: engine, - publisher=refuse_publication, - region_data_verifier=reject, + steps=initialization.PublicationSteps( + verify_region_data=reject, publish=refuse_publication + ), ) assert engine.disposed diff --git a/tests/unit/v2/test_catalog_spm_availability.py b/tests/unit/v2/test_catalog_spm_availability.py new file mode 100644 index 000000000..b07f3c6d8 --- /dev/null +++ b/tests/unit/v2/test_catalog_spm_availability.py @@ -0,0 +1,134 @@ +"""Each catalog version's SPM availability, from the simulation service's versions.""" + +from __future__ import annotations + +import httpx +import pytest + +from policyengine_api.data.v2.catalog import spm_availability +from policyengine_api.data.v2.catalog.publication_types import CatalogPublicationError +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.spm import SPMSelection + +POLICYENGINE_VERSION = "6.10.0" +DEFAULTS = { + "forecast_content_sha256": "3d" * 32, + "scenario": "ce_trend", + "geography_kind": "county", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, +} +VERSIONS = { + "policyengine": {"latest": POLICYENGINE_VERSION}, + "spm_capabilities": { + POLICYENGINE_VERSION: { + "contract_version": "canonical-spm-v1", + "defaults": DEFAULTS, + } + }, +} + + +def test_a_us_version_with_a_certified_capability_offers_spm() -> None: + availability = spm_availability.spm_availability( + "us", POLICYENGINE_VERSION, VERSIONS + ) + + assert availability == { + "available": True, + "settings_schema": SPMSelection.model_json_schema(), + "defaults": DEFAULTS, + } + + +@pytest.mark.parametrize( + ("country_id", "policyengine_version", "versions"), + [ + ("uk", POLICYENGINE_VERSION, VERSIONS), + ("us", "6.11.0", VERSIONS), + ("us", POLICYENGINE_VERSION, {}), + ("us", POLICYENGINE_VERSION, {"spm_capabilities": []}), + ( + "us", + POLICYENGINE_VERSION, + { + "spm_capabilities": { + POLICYENGINE_VERSION: { + "contract_version": "canonical-spm-v0", + "defaults": DEFAULTS, + } + } + }, + ), + ( + "us", + POLICYENGINE_VERSION, + { + "spm_capabilities": { + POLICYENGINE_VERSION: { + "contract_version": "canonical-spm-v1", + "defaults": {**DEFAULTS, "scenario": None}, + } + } + }, + ), + ], + ids=[ + "uk", + "unlisted-version", + "no-capabilities", + "malformed-capabilities", + "other-contract", + "unpinned-defaults", + ], +) +def test_anything_else_offers_no_spm( + country_id: str, policyengine_version: str, versions: object +) -> None: + assert spm_availability.spm_availability( + country_id, policyengine_version, versions + ) == {"available": False} + + +def test_the_versions_come_from_the_simulation_services_public_route() -> None: + requests: list[httpx.Request] = [] + + def answer(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, json=VERSIONS) + + with httpx.Client(transport=httpx.MockTransport(answer)) as client: + versions = spm_availability.fetch_simulation_versions( + "https://simulation.example/", client=client + ) + + assert versions == VERSIONS + assert [str(request.url) for request in requests] == [ + "https://simulation.example/versions" + ] + assert "authorization" not in requests[0].headers + + +def test_an_unreadable_versions_answer_stops_the_recording() -> None: + with ( + httpx.Client( + transport=httpx.MockTransport(lambda _request: httpx.Response(503)) + ) as client, + pytest.raises(CatalogPublicationError, match="versions"), + ): + spm_availability.fetch_simulation_versions( + "https://simulation.example", client=client + ) + + +def test_recording_needs_the_simulation_service_url() -> None: + with pytest.raises(V2ConfigurationError, match="SIMULATION_ENTRYPOINT_URL"): + spm_availability.simulation_entrypoint_url({}) + + assert ( + spm_availability.simulation_entrypoint_url( + {"SIMULATION_ENTRYPOINT_URL": " https://simulation.example/ "} + ) + == "https://simulation.example" + ) diff --git a/tests/unit/v2/test_metadata_deployment.py b/tests/unit/v2/test_metadata_deployment.py index a5e51e864..1784a108a 100644 --- a/tests/unit/v2/test_metadata_deployment.py +++ b/tests/unit/v2/test_metadata_deployment.py @@ -54,6 +54,12 @@ def test_reusable_seeding_workflow_separates_database_credentials() -> None: assert "V2_MIGRATION_DATABASE_URL" not in publication assert "V2_RUNTIME_DATABASE_URL" not in workflow assert "scripts/publish_v2_metadata_catalog.py" in publication + # Publication records each version's SPM availability from the + # environment's own simulation service. + assert ( + "SIMULATION_ENTRYPOINT_URL: ${{ secrets.SIMULATION_ENTRYPOINT_URL }}" + in publication + ) assert (REPO / "scripts/publish_v2_metadata_catalog.py").is_file() assert not (REPO / "scripts/initialize_v2_metadata.py").exists() diff --git a/tests/unit/v2/test_models.py b/tests/unit/v2/test_models.py index 66ac8b49b..88ed2db5e 100644 --- a/tests/unit/v2/test_models.py +++ b/tests/unit/v2/test_models.py @@ -695,6 +695,15 @@ def test_saved_policy_table_is_named_like_the_other_association_tables() -> None } +def test_spm_availability_is_nullable_for_versions_not_yet_recorded() -> None: + """Publication records it after the catalog, so a version may lack it briefly.""" + column = V2_METADATA.tables["tax_benefit_model_versions"].c["spm_availability"] + + assert column.nullable + assert column.server_default is None + assert isinstance(column.type, sa.JSON) + + def test_catalog_format_2_columns_are_nullable_for_earlier_versions() -> None: """Catalog format 2 stores v3's metadata; versions published before it keep null. From fcad8fd6fe3c91c79210cee0417eced3cbbfa256 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 20:49:17 +0400 Subject: [PATCH 10/37] Serve each catalog version's SPM availability at /v2/model-catalog The model catalog result gains spm: available, and when it is, the JSON schema of a household's SPM selection and the certified defaults. A version whose availability is not recorded yet, or that does not offer SPM, is not available. The route contract and the stage 9 guide name the field. Co-Authored-By: Claude Opus 5.5 --- .../model-catalog-spm-availability.added.md | 1 + docs/engineering/migration-contracts.md | 2 +- docs/generated/migration_contracts.json | 3 +- docs/migration/stage-9-v2-metadata.md | 9 +++++- .../services/v2/metadata/transformations.py | 7 +++- .../services/v2/metadata/types.py | 20 ++++++++++-- tests/contract/registry.py | 1 + tests/unit/v2/test_metadata_routes.py | 11 +++++++ tests/unit/v2/test_metadata_service.py | 32 +++++++++++++++++++ 9 files changed, 80 insertions(+), 6 deletions(-) create mode 100644 changelog.d/model-catalog-spm-availability.added.md diff --git a/changelog.d/model-catalog-spm-availability.added.md b/changelog.d/model-catalog-spm-availability.added.md new file mode 100644 index 000000000..b8515016b --- /dev/null +++ b/changelog.d/model-catalog-spm-availability.added.md @@ -0,0 +1 @@ +`GET /v2/model-catalog` now returns `spm`: whether the catalog version offers SPM settings for a household, with their JSON schema and certified defaults when it does. diff --git a/docs/engineering/migration-contracts.md b/docs/engineering/migration-contracts.md index 67d63c2b1..31614d7d4 100644 --- a/docs/engineering/migration-contracts.md +++ b/docs/engineering/migration-contracts.md @@ -224,7 +224,7 @@ Generated from `policyengine_api/migration_registry.py` and `tests/contract/regi | `GET` | `/v2/regions/by-code/state/ca?country_id=us` | 200 | `metadata` | none | `status`, `message`, `result.policyengine_version`, `result.item` | | `GET` | `/v2/tax-benefit-models/by-country/us` | 200 | `metadata` | none | `status`, `message`, `result.policyengine_version`, `result.model`, `result.model_version` | | `GET` | `/v2/economy-options?country_id=us` | 200 | `metadata` | none | `status`, `message`, `result.policyengine_version`, `result.current_law_id`, `result.region`, `result.time_period`, `result.datasets` | -| `GET` | `/v2/model-catalog?country_id=us` | 200 | `metadata` | none | `status`, `message`, `result.policyengine_version`, `result.catalog_format`, `result.entities`, `result.basic_inputs`, `result.regions`, `result.datasets`, `result.economy_options` | +| `GET` | `/v2/model-catalog?country_id=us` | 200 | `metadata` | none | `status`, `message`, `result.policyengine_version`, `result.catalog_format`, `result.entities`, `result.basic_inputs`, `result.regions`, `result.datasets`, `result.economy_options`, `result.spm` | | `GET` | `/v2/model-catalog/variables?country_id=us` | 200 | `metadata` | none | `status`, `message`, `result.policyengine_version`, `result.catalog_format`, `result.variables`, `result.part_parameter_values` | | `GET` | `/v2/model-catalog/parameters?country_id=us&name=gov.irs.income.bracket.rates.1` | 200 | `metadata` | none | `status`, `message`, `result.policyengine_version`, `result.catalog_format`, `result.parameters`, `result.node_labels`, `result.missing_names` | diff --git a/docs/generated/migration_contracts.json b/docs/generated/migration_contracts.json index b6a956745..928b956b3 100644 --- a/docs/generated/migration_contracts.json +++ b/docs/generated/migration_contracts.json @@ -1651,7 +1651,8 @@ "result.basic_inputs", "result.regions", "result.datasets", - "result.economy_options" + "result.economy_options", + "result.spm" ] }, { diff --git a/docs/migration/stage-9-v2-metadata.md b/docs/migration/stage-9-v2-metadata.md index 8e9ea0123..d02cbd0b4 100644 --- a/docs/migration/stage-9-v2-metadata.md +++ b/docs/migration/stage-9-v2-metadata.md @@ -36,7 +36,14 @@ version (with the current-law policy id and the years offered), the entities input dataset, and the economy choices. A catalog without regions, such as a household-only country's, has empty region, dataset and economy-choice lists rather than an error. A version published before catalog format 2 has a null -`catalog_format` and empty entities and basic inputs. The result names the +`catalog_format` and empty entities and basic inputs. `spm` says whether the +version offers SPM settings for a household: `available`, and when it does, +`settings_schema` (the JSON schema of a household's `spm` selection) and +`defaults` (the certified selection). Catalog publication records it from the +simulation service's public `GET /versions`: a US version offers SPM when the +service lists a canonical capability with pinned defaults for its +PolicyEngine.py version, and other countries never do. A version whose +availability is not recorded yet offers none. The result names the selected PolicyEngine.py version, so a client that starts without a version learns the running one and requests every other part for that exact version. The query uses the shared `CatalogQuery` model: unknown and repeated diff --git a/policyengine_api/services/v2/metadata/transformations.py b/policyengine_api/services/v2/metadata/transformations.py index 00c197cac..67d860a56 100644 --- a/policyengine_api/services/v2/metadata/transformations.py +++ b/policyengine_api/services/v2/metadata/transformations.py @@ -29,6 +29,7 @@ MetadataRegion, MetadataRegionOption, MetadataRegionType, + MetadataSPMAvailability, MetadataTimePeriodOption, MetadataVariable, ParameterNodeLabelRow, @@ -402,7 +403,8 @@ def metadata_model_catalog( Without a national dataset (a catalog without regions) the economy choices are empty. Entities and basic inputs a version published before catalog - format 2 lacks are empty lists. + format 2 lacks are empty lists, and a version whose SPM availability is + not recorded yet offers none. """ model_version = selected.model_version return MetadataModelCatalogResult( @@ -424,4 +426,7 @@ def metadata_model_catalog( if national_dataset is not None else MetadataEconomyOptions(region=[], time_period=[], datasets=[]) ), + spm=MetadataSPMAvailability.model_validate( + model_version.spm_availability or {"available": False} + ), ) diff --git a/policyengine_api/services/v2/metadata/types.py b/policyengine_api/services/v2/metadata/types.py index 63a77944e..eb15bca55 100644 --- a/policyengine_api/services/v2/metadata/types.py +++ b/policyengine_api/services/v2/metadata/types.py @@ -326,14 +326,29 @@ class MetadataEconomyOptions(StrictResponseModel): datasets: list[MetadataDatasetOption] +class MetadataSPMAvailability(StrictResponseModel): + """Whether a catalog version offers SPM settings for a household. + + When ``available``, ``settings_schema`` is the JSON schema of a + household's ``spm`` selection and ``defaults`` the certified selection + that applies to every field a household leaves out. A version without + them (any country but the US, or a bundle without certified SPM support) + is not available, and both are null. + """ + + available: bool + settings_schema: dict[str, JsonValue] | None = None + defaults: dict[str, JsonValue] | None = None + + class MetadataModelCatalogResult(StrictResponseModel): """The small parts of one catalog version, read when a client starts. policyengine_version names the selected catalog version, so a client can request its other parts for that exact version. catalog_format is null, and entities and basic_inputs are empty, for a version published before - catalog format 2. Variables, parameters and their values are served by - their own routes. + catalog format 2. ``spm`` says whether the version offers SPM settings. + Variables, parameters and their values are served by their own routes. """ policyengine_version: str @@ -345,6 +360,7 @@ class MetadataModelCatalogResult(StrictResponseModel): regions: list[MetadataRegion] datasets: list[MetadataDataset] economy_options: MetadataEconomyOptions + spm: MetadataSPMAvailability class MetadataDatedValue(StrictResponseModel): diff --git a/tests/contract/registry.py b/tests/contract/registry.py index 02a53761f..ac187802d 100644 --- a/tests/contract/registry.py +++ b/tests/contract/registry.py @@ -703,6 +703,7 @@ class WorkflowContract: "result.regions", "result.datasets", "result.economy_options", + "result.spm", ), route_group="metadata", ), diff --git a/tests/unit/v2/test_metadata_routes.py b/tests/unit/v2/test_metadata_routes.py index 2ba7580ef..a1e4708c1 100644 --- a/tests/unit/v2/test_metadata_routes.py +++ b/tests/unit/v2/test_metadata_routes.py @@ -51,6 +51,7 @@ MetadataParameterSummary, MetadataRegion, MetadataRegionOption, + MetadataSPMAvailability, MetadataTimePeriodOption, MetadataVariable, ParameterFilter, @@ -273,6 +274,11 @@ def detail(item: object) -> MetadataDetailResult: ) ], ), + spm=MetadataSPMAvailability( + available=True, + settings_schema={"type": "object"}, + defaults={"scenario": "ce_trend"}, + ), ), "get_economy_options": MetadataEconomyOptionsResult( policyengine_version=version, @@ -789,6 +795,11 @@ def test_model_catalog_route_serves_the_small_parts_of_the_selected_version() -> assert result["basic_inputs"] == ["employment_income"] assert [region["code"] for region in result["regions"]] == ["us"] assert result["economy_options"]["datasets"][0]["default"] is True + assert result["spm"] == { + "available": True, + "settings_schema": {"type": "object"}, + "defaults": {"scenario": "ce_trend"}, + } assert reader.calls == [("get_model_catalog", ("us", None), {})] assert reader.closed diff --git a/tests/unit/v2/test_metadata_service.py b/tests/unit/v2/test_metadata_service.py index b2f35d7be..6b6a19f32 100644 --- a/tests/unit/v2/test_metadata_service.py +++ b/tests/unit/v2/test_metadata_service.py @@ -702,6 +702,38 @@ def test_model_catalog_published_before_format_2_has_empty_new_fields( assert result.basic_inputs == [] +SPM_AVAILABLE = { + "available": True, + "settings_schema": {"title": "SPMSelection", "type": "object"}, + "defaults": {"scenario": "ce_trend", "geography_kind": "county"}, +} + + +def test_model_catalog_serves_the_versions_recorded_spm_availability( + catalog_session: Session, +) -> None: + model_version = _model_version(catalog_session, "policyengine-us") + model_version.spm_availability = SPM_AVAILABLE + catalog_session.add(model_version) + catalog_session.commit() + + result = _service(catalog_session).get_model_catalog("us") + + assert result.spm.model_dump(mode="json") == SPM_AVAILABLE + + +def test_model_catalog_without_recorded_spm_availability_offers_none( + catalog_session: Session, +) -> None: + result = _service(catalog_session).get_model_catalog("uk") + + assert result.spm.model_dump(mode="json") == { + "available": False, + "settings_schema": None, + "defaults": None, + } + + def test_model_catalog_with_regions_requires_the_national_region( catalog_session: Session, ) -> None: From 54071cdf7d252a21c04068fbaa80f8ad273d7604 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 21:25:46 +0400 Subject: [PATCH 11/37] Convert stored households and policies to what the simulation service takes A stored v2 household becomes the situation shape the simulation service's household functions take: people and each group collection by its plural, each group listing the people whose memberships name it, every record named by its ID so a calculated value maps back to its record. The saved SPM choice is read separately. A v2 policy's values become parameter values keyed by YYYY-MM-DD.YYYY-MM-DD period ranges, with an open end as 2100-12-31. Co-Authored-By: Claude Opus 5.5 --- .../v2/household_calculations/__init__.py | 1 + .../household_calculations/transformations.py | 79 ++++++++ ...t_household_calculation_transformations.py | 172 ++++++++++++++++++ 3 files changed, 252 insertions(+) create mode 100644 policyengine_api/services/v2/household_calculations/__init__.py create mode 100644 policyengine_api/services/v2/household_calculations/transformations.py create mode 100644 tests/unit/v2/test_household_calculation_transformations.py diff --git a/policyengine_api/services/v2/household_calculations/__init__.py b/policyengine_api/services/v2/household_calculations/__init__.py new file mode 100644 index 000000000..0501ab630 --- /dev/null +++ b/policyengine_api/services/v2/household_calculations/__init__.py @@ -0,0 +1 @@ +"""Household calculations on v2, run by the Simulation Entrypoint.""" diff --git a/policyengine_api/services/v2/household_calculations/transformations.py b/policyengine_api/services/v2/household_calculations/transformations.py new file mode 100644 index 000000000..53911db0f --- /dev/null +++ b/policyengine_api/services/v2/household_calculations/transformations.py @@ -0,0 +1,79 @@ +"""Stored households and policies, in the shapes the Simulation Entrypoint takes. + +A stored v2 household lists its people and each group entity as records, +each with an ``id`` and its ``values``; each person names, for every group +collection, the ``id`` of the group it belongs to. The Simulation Entrypoint's +household functions take the Public API's situation shape instead: ``people`` +maps each person to their values, and each group collection, by its plural, +maps each group to its ``members`` and values. Records keep their ``id`` as +their name, so a calculated value maps back to its record. + +A v2 policy stores each parameter value with its start date and an optional +end date; the household functions take the Public API's policy JSON, each +parameter's values keyed by ``YYYY-MM-DD.YYYY-MM-DD`` period ranges, where an +open end is ``2100-12-31``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Mapping + + from policyengine_api.services.v2.policies.types import PolicyRead + +# Each country's stored group collections and their plural in the situation. +GROUP_COLLECTIONS: dict[str, dict[str, str]] = { + "us": { + "household": "households", + "family": "families", + "tax_unit": "tax_units", + "spm_unit": "spm_units", + "marital_unit": "marital_units", + }, + "uk": {"household": "households", "benunit": "benunits"}, +} +OPEN_END_DATE = "2100-12-31" + + +def household_situation( + country_id: str, household_data: Mapping[str, Any] +) -> dict[str, Any]: + """Return a stored household's people and groups in the situation shape. + + The saved SPM choice is not part of the situation; ``saved_spm_choice`` + reads it. + """ + people = household_data["people"] + situation: dict[str, Any] = { + "people": {person["id"]: dict(person["values"]) for person in people} + } + for collection, plural in GROUP_COLLECTIONS[country_id].items(): + situation[plural] = { + group["id"]: { + "members": [ + person["id"] + for person in people + if person["memberships"].get(collection) == group["id"] + ], + **group["values"], + } + for group in household_data.get(collection, []) + } + return situation + + +def saved_spm_choice(household_data: Mapping[str, Any]) -> object: + """Return the household's saved SPM selection, or None when it has none.""" + return household_data.get("spm") + + +def policy_parameters(policy: PolicyRead) -> dict[str, dict[str, Any]]: + """Return the policy's values by parameter and period range.""" + parameters: dict[str, dict[str, Any]] = {} + for value in policy.parameter_values: + end = OPEN_END_DATE if value.end_date is None else value.end_date.date() + period = f"{value.start_date.date().isoformat()}.{end}" + parameters.setdefault(value.parameter_name, {})[period] = value.value + return parameters diff --git a/tests/unit/v2/test_household_calculation_transformations.py b/tests/unit/v2/test_household_calculation_transformations.py new file mode 100644 index 000000000..6efd6e8f7 --- /dev/null +++ b/tests/unit/v2/test_household_calculation_transformations.py @@ -0,0 +1,172 @@ +"""Stored households and policies, as the Simulation Entrypoint takes them.""" + +from __future__ import annotations + +from datetime import UTC, datetime +from uuid import uuid4 + +from policyengine_api.services.v2.household_calculations.transformations import ( + household_situation, + policy_parameters, + saved_spm_choice, +) +from policyengine_api.services.v2.policies.types import ( + PolicyParameterValueRead, + PolicyRead, +) + +US_HOUSEHOLD = { + "people": [ + { + "id": "person-1", + "source_name": "you", + "values": {"age": {"2026": 40}, "employment_income": {"2026": 30000}}, + "memberships": { + "household": "household-1", + "family": "family-1", + "tax_unit": "tax_unit-1", + "spm_unit": "spm_unit-1", + "marital_unit": "marital_unit-1", + }, + }, + { + "id": "person-2", + "values": {"age": {"2026": 8}}, + "memberships": { + "household": "household-1", + "family": "family-1", + "tax_unit": "tax_unit-1", + "spm_unit": "spm_unit-1", + "marital_unit": "marital_unit-2", + }, + }, + ], + "household": [{"id": "household-1", "values": {"state_name": {"2026": "CA"}}}], + "family": [{"id": "family-1", "values": {}}], + "tax_unit": [{"id": "tax_unit-1", "values": {"eitc": {"2026": None}}}], + "spm_unit": [{"id": "spm_unit-1", "values": {}}], + "marital_unit": [ + {"id": "marital_unit-1", "values": {}}, + {"id": "marital_unit-2", "values": {}}, + ], + "spm": {"scenario": "ce_trend", "geography_kind": "national"}, +} +UK_HOUSEHOLD = { + "people": [ + { + "id": "person-1", + "values": {"age": {"2026": 30}}, + "memberships": {"household": "household-1", "benunit": "benunit-1"}, + } + ], + "household": [{"id": "household-1", "values": {"region": {"2026": "WALES"}}}], + "benunit": [{"id": "benunit-1", "values": {"universal_credit": {"2026": None}}}], +} + + +def test_a_us_household_becomes_a_situation_keyed_by_record_id() -> None: + assert household_situation("us", US_HOUSEHOLD) == { + "people": { + "person-1": {"age": {"2026": 40}, "employment_income": {"2026": 30000}}, + "person-2": {"age": {"2026": 8}}, + }, + "households": { + "household-1": { + "members": ["person-1", "person-2"], + "state_name": {"2026": "CA"}, + } + }, + "families": {"family-1": {"members": ["person-1", "person-2"]}}, + "tax_units": { + "tax_unit-1": {"members": ["person-1", "person-2"], "eitc": {"2026": None}} + }, + "spm_units": {"spm_unit-1": {"members": ["person-1", "person-2"]}}, + "marital_units": { + "marital_unit-1": {"members": ["person-1"]}, + "marital_unit-2": {"members": ["person-2"]}, + }, + } + + +def test_a_uk_household_has_households_and_benefit_units() -> None: + assert household_situation("uk", UK_HOUSEHOLD) == { + "people": {"person-1": {"age": {"2026": 30}}}, + "households": { + "household-1": {"members": ["person-1"], "region": {"2026": "WALES"}} + }, + "benunits": { + "benunit-1": { + "members": ["person-1"], + "universal_credit": {"2026": None}, + } + }, + } + + +def test_a_group_without_members_keeps_an_empty_member_list() -> None: + household = {**UK_HOUSEHOLD, "benunit": [*UK_HOUSEHOLD["benunit"]]} + household["benunit"].append({"id": "benunit-2", "values": {}}) + + situation = household_situation("uk", household) + + assert situation["benunits"]["benunit-2"] == {"members": []} + + +def test_the_saved_spm_choice_is_the_documents_selection_or_none() -> None: + assert saved_spm_choice(US_HOUSEHOLD) == { + "scenario": "ce_trend", + "geography_kind": "national", + } + assert saved_spm_choice(UK_HOUSEHOLD) is None + + +def _value(*, name: str, value: object, start: datetime, end: datetime | None): + return PolicyParameterValueRead( + id=uuid4(), + parameter_id=uuid4(), + parameter_name=name, + value=value, + start_date=start, + end_date=end, + ) + + +def test_a_policy_becomes_parameter_values_by_period_range() -> None: + policy = PolicyRead( + id=uuid4(), + country_id="us", + tax_benefit_model_id=uuid4(), + tax_benefit_model_version_id=uuid4(), + created_at=datetime(2026, 1, 1, tzinfo=UTC), + updated_at=datetime(2026, 1, 1, tzinfo=UTC), + parameter_values=( + _value( + name="gov.irs.credits.ctc.amount.base[0].amount", + value=3000, + start=datetime(2026, 1, 1, tzinfo=UTC), + end=None, + ), + _value( + name="gov.irs.credits.ctc.refundable.fully_refundable", + value=True, + start=datetime(2026, 1, 1, tzinfo=UTC), + end=datetime(2027, 12, 31, tzinfo=UTC), + ), + _value( + name="gov.irs.credits.ctc.amount.base[0].amount", + value=3500, + start=datetime(2028, 1, 1, tzinfo=UTC), + end=None, + ), + ), + ) + + assert policy_parameters(policy) == { + "gov.irs.credits.ctc.amount.base[0].amount": { + "2026-01-01.2100-12-31": 3000, + "2028-01-01.2100-12-31": 3500, + }, + "gov.irs.credits.ctc.refundable.fully_refundable": { + "2026-01-01.2027-12-31": True + }, + } From d12b80451cc87ffa4b6d7dc393902f41db557c37 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 21:49:03 +0400 Subject: [PATCH 12/37] Calculate stored households through /v2/household-calculations POST /v2/household-calculations checks the country's countries row allows household calculations (422 naming the country otherwise), reads the stored household and optional policy through the households and policies services, selects the country's running catalog version and sends the situation, saved SPM choice, policy parameters and options (axes, add_missing_variables) to the Simulation Entrypoint. GET /v2/household-calculations/{calculation_id} reads the calculation back with its household, warnings, spm_config, spm_provenance and error. Neither route needs an account. The simulation service is called through the generated simulation client, one per process, with the gateway credentials and correlation headers; a missing or partial configuration, an unreachable service, an error answer or a refusal is a 503. Co-Authored-By: Claude Opus 5.5 --- .../v2-household-calculations.added.md | 1 + docs/migration/stage-11-v2-households.md | 36 +++ .../fastapi_routes/dependencies.py | 52 ++++ .../v2/household_calculations/__init__.py | 1 + .../household_calculations/request_models.py | 27 ++ .../household_calculations/response_models.py | 112 ++++++++ .../v2/household_calculations/routes.py | 147 +++++++++++ policyengine_api/fastapi_routes/v2/routes.py | 4 + .../v2/household_calculations/services.py | 136 ++++++++++ .../simulation_client.py | 171 ++++++++++++ .../v2/household_calculations/types.py | 205 +++++++++++++++ .../v2/household_calculations/validators.py | 47 ++++ pyproject.toml | 4 + .../unit/v2/household_calculation_support.py | 248 ++++++++++++++++++ .../v2/test_household_calculation_routes.py | 224 ++++++++++++++++ .../v2/test_household_calculation_service.py | 246 +++++++++++++++++ ...household_calculation_simulation_client.py | 166 ++++++++++++ tests/unit/v2/test_metadata_routes.py | 2 + uv.lock | 15 ++ 19 files changed, 1844 insertions(+) create mode 100644 changelog.d/v2-household-calculations.added.md create mode 100644 policyengine_api/fastapi_routes/v2/household_calculations/__init__.py create mode 100644 policyengine_api/fastapi_routes/v2/household_calculations/request_models.py create mode 100644 policyengine_api/fastapi_routes/v2/household_calculations/response_models.py create mode 100644 policyengine_api/fastapi_routes/v2/household_calculations/routes.py create mode 100644 policyengine_api/services/v2/household_calculations/services.py create mode 100644 policyengine_api/services/v2/household_calculations/simulation_client.py create mode 100644 policyengine_api/services/v2/household_calculations/types.py create mode 100644 policyengine_api/services/v2/household_calculations/validators.py create mode 100644 tests/unit/v2/household_calculation_support.py create mode 100644 tests/unit/v2/test_household_calculation_routes.py create mode 100644 tests/unit/v2/test_household_calculation_service.py create mode 100644 tests/unit/v2/test_household_calculation_simulation_client.py diff --git a/changelog.d/v2-household-calculations.added.md b/changelog.d/v2-household-calculations.added.md new file mode 100644 index 000000000..9a4a83ccc --- /dev/null +++ b/changelog.d/v2-household-calculations.added.md @@ -0,0 +1 @@ +`POST /v2/household-calculations` sends a stored household, under a stored policy or current law, with optional axes, to the simulation service for calculation, and `GET /v2/household-calculations/{calculation_id}` returns its status and, once it succeeded, the calculated household with its warnings and, in the US, the SPM selection applied and its receipt. Neither needs an account; a country whose `countries` row does not allow household calculations is refused with 422. diff --git a/docs/migration/stage-11-v2-households.md b/docs/migration/stage-11-v2-households.md index c9a68eb1c..2a58f7dd1 100644 --- a/docs/migration/stage-11-v2-households.md +++ b/docs/migration/stage-11-v2-households.md @@ -256,3 +256,39 @@ Apply the reviewed schema revisions and application image only to independently verified production targets. Enable native v2 use and v1 create copying as separate operational decisions and retain the preceding application revision for immediate rollback. + +## Household Calculations + +`POST /v2/household-calculations` calculates a household stored through +`/v2/households`, under a policy stored through `/v2/policies` or, without +`policy_id`, under current law. The body may add `year` (the household's +`default_year` otherwise), `axes`, `add_missing_variables` and `force`. The +Public API checks the country's `countries` row allows household calculations +(422 otherwise; a household-only country is accepted), reads the household and +policy through the same services as their routes (404 when either is not +stored in that country), selects the country's running catalog version, and +sends the household in the situation shape, its saved SPM choice and the +policy's parameter values to the simulation service, which cannot read these +tables. The answer is 202 with the calculation's `calculation_id` and +`status`. Poll `GET /v2/household-calculations/{calculation_id}` until +`status` is `succeeded`, with `result`, or `failed`, with `error`. A US result +carries `spm_config` and `spm_provenance`. + +The simulation service stores every calculation as a simulation run, so a +calculation no account saved is still stored, with no report. Neither route +needs an account. + +The routes always call the Cloud Run Simulation Entrypoint, whatever +`SIM_ENTRYPOINT` selects, because only it serves v2 calculations. They read: + +- `SIMULATION_ENTRYPOINT_URL`: the `staging` and `production` environment + secret of the same name, passed to the deployment job's `env:` and written + to the Cloud Run revision by `deploy_cloud_run_candidate.sh` when set. +- `GATEWAY_AUTH_ISSUER`, `GATEWAY_AUTH_AUDIENCE`, `GATEWAY_AUTH_CLIENT_ID` and + `GATEWAY_AUTH_CLIENT_SECRET_RESOURCE`: the organization secrets already used + for every simulation call, with `GATEWAY_AUTH_REQUIRED=1`. + +Without the URL, or with gateway auth partly configured, both routes answer +503. They also answer 503 when the simulation service cannot be reached, +answers with an error, or refuses the calculation (for example when no v2 +worker is deployed for the catalog version's PolicyEngine.py version). diff --git a/policyengine_api/fastapi_routes/dependencies.py b/policyengine_api/fastapi_routes/dependencies.py index e6336d5e5..5d0bff862 100644 --- a/policyengine_api/fastapi_routes/dependencies.py +++ b/policyengine_api/fastapi_routes/dependencies.py @@ -29,6 +29,10 @@ LegacyReferenceChecker, ) from policyengine_api.services.v2.countries.types import CountryRead + from policyengine_api.services.v2.household_calculations.types import ( + HouseholdCalculationInput, + HouseholdCalculationRead, + ) from policyengine_api.services.v2.households.types import ( HouseholdPage, HouseholdRead, @@ -444,6 +448,29 @@ def list_households( ... +class V2HouseholdCalculationResourceService(Protocol): + """Route-independent household calculation operations for one request.""" + + def submit_calculation( + self, calculation: HouseholdCalculationInput + ) -> HouseholdCalculationRead: + """Send a stored household's calculation to the simulation service. + + Raise ``HouseholdCalculationCountryError`` for a country without + household calculations, ``HouseholdNotFoundError`` or + ``PolicyNotFoundError`` for a record not stored in that country, and + the simulation service errors when it cannot calculate. + """ + ... + + def get_calculation(self, calculation_id: UUID) -> HouseholdCalculationRead: + """Return the calculation as it stands on the simulation service. + + Raise ``HouseholdCalculationNotFoundError`` when it has none. + """ + ... + + class V2UserPolicyResourceService(Protocol): """Route-independent user-policy operations for one request.""" @@ -857,6 +884,25 @@ def _default_v2_household_service_factory() -> V2HouseholdResourceService: ) +def _default_v2_household_calculation_service_factory() -> ( + V2HouseholdCalculationResourceService +): + from policyengine_api.services.v2.household_calculations.services import ( # noqa: PLC0415 # loaded on the first calculation request only + V2HouseholdCalculationService, + ) + from policyengine_api.services.v2.household_calculations.simulation_client import ( # noqa: PLC0415 # imports the generated simulation client on first use only + process_simulation_household_client, + ) + + return V2HouseholdCalculationService( + countries=_default_v2_country_service_factory(), + households=_default_v2_household_service_factory(), + policies=_default_v2_policy_service_factory(), + catalogs=_default_v2_metadata_reader_factory(), + simulation=process_simulation_household_client(), + ) + + _user_bearer_scheme = HTTPBearer( auto_error=False, scheme_name="UserAccessToken", @@ -1006,6 +1052,9 @@ class NativeRouteDependencies: v2_country_service_factory: Callable[[], V2CountryResourceService] | None = None v2_policy_service_factory: Callable[[], V2PolicyResourceService] | None = None v2_household_service_factory: Callable[[], V2HouseholdResourceService] | None = None + v2_household_calculation_service_factory: ( + Callable[[], V2HouseholdCalculationResourceService] | None + ) = None v2_user_policy_service_factory: Callable[[], V2UserPolicyResourceService] | None = ( None ) @@ -1041,6 +1090,9 @@ def defaults(cls) -> NativeRouteDependencies: v2_country_service_factory=_default_v2_country_service_factory, v2_policy_service_factory=_default_v2_policy_service_factory, v2_household_service_factory=_default_v2_household_service_factory, + v2_household_calculation_service_factory=( + _default_v2_household_calculation_service_factory + ), v2_user_policy_service_factory=(_default_v2_user_policy_service_factory), v2_user_household_service_factory=( _default_v2_user_household_service_factory diff --git a/policyengine_api/fastapi_routes/v2/household_calculations/__init__.py b/policyengine_api/fastapi_routes/v2/household_calculations/__init__.py new file mode 100644 index 000000000..59e3798f7 --- /dev/null +++ b/policyengine_api/fastapi_routes/v2/household_calculations/__init__.py @@ -0,0 +1 @@ +"""Native FastAPI routes calculating stored v2 households.""" diff --git a/policyengine_api/fastapi_routes/v2/household_calculations/request_models.py b/policyengine_api/fastapi_routes/v2/household_calculations/request_models.py new file mode 100644 index 000000000..c8b0e4427 --- /dev/null +++ b/policyengine_api/fastapi_routes/v2/household_calculations/request_models.py @@ -0,0 +1,27 @@ +"""Strict HTTP request models for v2 household calculations.""" + +from __future__ import annotations + +from policyengine_api.services.v2.household_calculations.types import ( + HouseholdCalculationInput, +) + + +class HouseholdCalculationRequest(HouseholdCalculationInput): + """A stored household to calculate under a stored policy, or current law. + + ``household_id`` and ``policy_id`` name records of ``country_id`` stored + through ``/v2/households`` and ``/v2/policies``; without ``policy_id`` the + household is calculated under current law. ``year`` defaults to the + household's ``default_year``. Each group in ``axes`` varies its inputs + together across ``count`` points; several groups form a grid. + ``add_missing_variables`` adds every variable the household lacks before + calculating, and ``force`` calculates again even when a recent identical + calculation exists. + """ + + def to_service_input(self) -> HouseholdCalculationInput: + """Convert the request to the calculation service's input.""" + return HouseholdCalculationInput.model_validate( + self.model_dump(mode="python", exclude_unset=True) + ) diff --git a/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py b/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py new file mode 100644 index 000000000..b2b1b160d --- /dev/null +++ b/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py @@ -0,0 +1,112 @@ +"""Strict HTTP response models for v2 household calculations.""" + +from __future__ import annotations + +from datetime import datetime +from typing import TYPE_CHECKING, Any, Literal +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, JsonValue + +from policyengine_api.fastapi_routes.v2.errors import V2ErrorResponse + +if TYPE_CHECKING: + from policyengine_api.services.v2.household_calculations.types import ( + HouseholdCalculationRead, + ) + + +class StrictHouseholdCalculationAPIModel(BaseModel): + """Base for calculation responses; rejects unknown fields and NaN or infinity.""" + + model_config = ConfigDict(extra="forbid", allow_inf_nan=False, from_attributes=True) + + +class HouseholdCalculationOutput(StrictHouseholdCalculationAPIModel): + """A succeeded calculation's household. + + ``household`` is the stored household in the situation shape (people and + groups keyed by their record IDs) with each calculated value filled in by + period. ``warnings`` lists inputs the country model ignored. In the US, + ``spm_config`` is the SPM selection the calculation applied and + ``spm_provenance`` its receipt; both are null elsewhere. + ``policyengine_version`` is the PolicyEngine.py version that calculated it. + """ + + household: dict[str, JsonValue] + warnings: list[str] + spm_config: dict[str, JsonValue] | None + spm_provenance: dict[str, JsonValue] | None + policyengine_version: str | None + + +class HouseholdCalculationError(StrictHouseholdCalculationAPIModel): + """Why a calculation failed; ``detail`` holds the calculation's own error.""" + + code: str + message: str + detail: dict[str, JsonValue] | None + + +class HouseholdCalculationItem(StrictHouseholdCalculationAPIModel): + """A household calculation and its outcome. + + ``status`` moves from ``pending`` through ``dispatching`` and ``running`` + to ``succeeded``, with ``result``, or ``failed``, with ``error``. Poll + ``GET /v2/household-calculations/{calculation_id}`` until it is one of + those two. ``catalog_version_id`` is the catalog version calculated. + """ + + calculation_id: UUID + country_id: str + household_id: UUID + policy_id: UUID | None + year: int + catalog_version_id: UUID + status: Literal["pending", "dispatching", "running", "succeeded", "failed"] + created_at: datetime + completed_at: datetime | None + result: HouseholdCalculationOutput | None + error: HouseholdCalculationError | None + + @classmethod + def from_read( + cls, calculation: HouseholdCalculationRead + ) -> HouseholdCalculationItem: + """Build the response item from the service's calculation.""" + return cls.model_validate(calculation) + + +class HouseholdCalculationResponse(StrictHouseholdCalculationAPIModel): + """Success envelope: ``status`` "ok", a null ``message`` and the calculation.""" + + status: Literal["ok"] = "ok" + message: None = None + result: HouseholdCalculationItem + + +HOUSEHOLD_CALCULATION_ERROR_RESPONSES: dict[int | str, dict[str, Any]] = { + 404: { + "model": V2ErrorResponse, + "description": "The household, policy or calculation does not exist.", + }, + 422: { + "model": V2ErrorResponse, + "description": ( + "The request does not match the schema, the country does not offer " + "household calculations, or no year can be chosen." + ), + }, + 500: { + "model": V2ErrorResponse, + "description": "The household calculation request failed unexpectedly.", + }, + 503: { + "model": V2ErrorResponse, + "description": ( + "The simulation service, the country's catalog version or v2 " + "storage is unavailable, or the simulation service cannot " + "calculate this household." + ), + }, +} diff --git a/policyengine_api/fastapi_routes/v2/household_calculations/routes.py b/policyengine_api/fastapi_routes/v2/household_calculations/routes.py new file mode 100644 index 000000000..c88ce2289 --- /dev/null +++ b/policyengine_api/fastapi_routes/v2/household_calculations/routes.py @@ -0,0 +1,147 @@ +"""Native FastAPI routes submitting stored households for calculation, and polling them. + +The calculation itself runs on the simulation service, which stores every +calculation as a simulation run; a calculation not saved to an account has no +report. The routes need no account. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, TypeVar +from uuid import UUID + +from fastapi import APIRouter +from sqlalchemy.exc import SQLAlchemyError +from starlette.responses import JSONResponse + +from policyengine_api.data.v2.catalog.catalog_selection import ( + MetadataCatalogUnavailableError, + MetadataCatalogVersionNotFoundError, + UnsupportedPreviewCountryError, +) +from policyengine_api.data.v2.settings import V2ConfigurationError +from policyengine_api.fastapi_routes.dependencies import ( + _default_v2_household_calculation_service_factory, +) +from policyengine_api.fastapi_routes.v2.errors import V2ErrorMapping, run_v2_operation +from policyengine_api.fastapi_routes.v2.household_calculations.request_models import ( + HouseholdCalculationRequest, +) +from policyengine_api.fastapi_routes.v2.household_calculations.response_models import ( + HOUSEHOLD_CALCULATION_ERROR_RESPONSES, + HouseholdCalculationItem, + HouseholdCalculationResponse, +) +from policyengine_api.services.v2.household_calculations.types import ( + SimulationServiceRefusedError, + SimulationServiceUnavailableError, +) +from policyengine_api.services.v2.household_calculations.validators import ( + HouseholdCalculationCountryError, + HouseholdCalculationNotFoundError, + HouseholdCalculationYearError, +) +from policyengine_api.services.v2.households.validators import HouseholdNotFoundError +from policyengine_api.services.v2.policies.validators import PolicyNotFoundError + +if TYPE_CHECKING: + from collections.abc import Callable + + from policyengine_api.fastapi_routes.dependencies import ( + NativeRouteDependencies, + V2HouseholdCalculationResourceService, + ) + +OperationT = TypeVar("OperationT") + +HOUSEHOLD_CALCULATION_ERROR_MAPPINGS = ( + V2ErrorMapping( + (HouseholdCalculationCountryError, HouseholdCalculationYearError), 422 + ), + V2ErrorMapping( + ( + HouseholdNotFoundError, + PolicyNotFoundError, + HouseholdCalculationNotFoundError, + ), + 404, + ), + V2ErrorMapping( + (SimulationServiceUnavailableError, SimulationServiceRefusedError), 503 + ), + V2ErrorMapping( + ( + MetadataCatalogUnavailableError, + MetadataCatalogVersionNotFoundError, + UnsupportedPreviewCountryError, + ), + 503, + "The country's catalog version is unavailable", + ), + V2ErrorMapping( + (V2ConfigurationError, SQLAlchemyError), + 503, + "V2 household calculation storage is unavailable", + ), +) + + +def _service_factory( + dependencies: NativeRouteDependencies, +) -> Callable[[], V2HouseholdCalculationResourceService]: + if dependencies.v2_household_calculation_service_factory is not None: + return dependencies.v2_household_calculation_service_factory + return _default_v2_household_calculation_service_factory + + +def _calculation_operation( + operation: Callable[[], OperationT], +) -> OperationT | JSONResponse: + return run_v2_operation( + operation, + HOUSEHOLD_CALCULATION_ERROR_MAPPINGS, + fallback_message="V2 household calculation failed", + ) + + +def build_v2_household_calculation_router( + dependencies: NativeRouteDependencies, +) -> APIRouter: + """Build the household calculation routes without opening any connection.""" + router = APIRouter(prefix="/v2", responses=HOUSEHOLD_CALCULATION_ERROR_RESPONSES) + service_factory = _service_factory(dependencies) + + @router.post( + "/household-calculations", + response_model=HouseholdCalculationResponse, + status_code=202, + summary="Calculate a stored household under a stored policy", + ) + def submit_household_calculation( + body: HouseholdCalculationRequest, + ) -> HouseholdCalculationResponse | JSONResponse: + def submit() -> HouseholdCalculationResponse: + calculation = service_factory().submit_calculation(body.to_service_input()) + return HouseholdCalculationResponse( + result=HouseholdCalculationItem.from_read(calculation) + ) + + return _calculation_operation(submit) + + @router.get( + "/household-calculations/{calculation_id}", + response_model=HouseholdCalculationResponse, + summary="Read a household calculation and its result", + ) + def get_household_calculation( + calculation_id: UUID, + ) -> HouseholdCalculationResponse | JSONResponse: + def read() -> HouseholdCalculationResponse: + calculation = service_factory().get_calculation(calculation_id) + return HouseholdCalculationResponse( + result=HouseholdCalculationItem.from_read(calculation) + ) + + return _calculation_operation(read) + + return router diff --git a/policyengine_api/fastapi_routes/v2/routes.py b/policyengine_api/fastapi_routes/v2/routes.py index d3b412252..1aa2032b3 100644 --- a/policyengine_api/fastapi_routes/v2/routes.py +++ b/policyengine_api/fastapi_routes/v2/routes.py @@ -11,6 +11,9 @@ build_v2_country_router, ) from policyengine_api.fastapi_routes.v2.errors import V2ErrorResponse +from policyengine_api.fastapi_routes.v2.household_calculations.routes import ( + build_v2_household_calculation_router, +) from policyengine_api.fastapi_routes.v2.households.routes import ( build_v2_household_router, ) @@ -58,6 +61,7 @@ def build_v2_router( router.include_router(build_v2_user_household_router(dependencies)) router.include_router(build_v2_user_policy_router(dependencies)) router.include_router(build_v2_household_router(dependencies)) + router.include_router(build_v2_household_calculation_router(dependencies)) router.include_router(build_v2_policy_router(dependencies)) router.include_router(build_v2_country_router(dependencies)) router.include_router(build_v2_metadata_model_router(dependencies)) diff --git a/policyengine_api/services/v2/household_calculations/services.py b/policyengine_api/services/v2/household_calculations/services.py new file mode 100644 index 000000000..3eafaeb71 --- /dev/null +++ b/policyengine_api/services/v2/household_calculations/services.py @@ -0,0 +1,136 @@ +"""Calculate stored households through the simulation service, and read them. + +Submitting a calculation checks the country's ``countries`` row allows +household calculations (a household-only country does), reads the stored +household and policy through the same services as ``/v2/households`` and +``/v2/policies`` (so a household or policy of another country, or none, is +not found), selects the country's running catalog version, and sends the +simulation service everything it needs: it cannot read these tables. The +simulation service records every calculation, so one not saved to an +account is still stored, as a simulation run with no report. + +Reading a calculation asks the simulation service for it as it stands. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from policyengine_api.services.v2.household_calculations.transformations import ( + household_situation, + policy_parameters, + saved_spm_choice, +) +from policyengine_api.services.v2.household_calculations.types import ( + HouseholdCalculationRead, +) +from policyengine_api.services.v2.household_calculations.validators import ( + HouseholdCalculationNotFoundError, + calculation_year, + require_household_country, +) + +if TYPE_CHECKING: + from uuid import UUID + + from policyengine_api.services.v2.household_calculations.types import ( + CatalogSource, + CountrySource, + HouseholdCalculationInput, + HouseholdSource, + PolicySource, + SimulationHouseholdCalculations, + ) + + +class V2HouseholdCalculationService: + """Submit household calculations to the simulation service and read them.""" + + def __init__( + self, + *, + countries: CountrySource, + households: HouseholdSource, + policies: PolicySource, + catalogs: CatalogSource, + simulation: SimulationHouseholdCalculations, + ) -> None: + """Read through these services for one request.""" + self._countries = countries + self._households = households + self._policies = policies + self._catalogs = catalogs + self._simulation = simulation + + def submit_calculation( + self, calculation: HouseholdCalculationInput + ) -> HouseholdCalculationRead: + """Send the calculation to the simulation service and return it. + + Raise HouseholdCalculationCountryError, HouseholdNotFoundError, + PolicyNotFoundError or HouseholdCalculationYearError before anything + is sent, and the simulation errors when the service cannot calculate. + """ + country_id = calculation.country_id + require_household_country(self._countries.list_countries(), country_id) + household = self._households.get_household( + country_id=country_id, household_id=calculation.household_id + ) + parameters: dict[str, dict[str, Any]] = {} + if calculation.policy_id is not None: + parameters = policy_parameters( + self._policies.get_policy( + country_id=country_id, policy_id=calculation.policy_id + ) + ) + year = calculation_year(calculation.year, household.default_year) + try: + catalog = self._catalogs.get_model_by_country(country_id) + finally: + self._catalogs.close() + submission: dict[str, Any] = { + "country": country_id, + "catalog_version_id": str(catalog.model_version.id), + "policyengine_version": catalog.policyengine_version, + "year": year, + "household_id": str(household.id), + "household": household_situation(country_id, household.household_data), + "policy": { + "policy_id": ( + None + if calculation.policy_id is None + else str(calculation.policy_id) + ), + "parameters": parameters, + }, + "options": { + "axes": ( + None + if calculation.axes is None + else [ + [ + axis.model_dump(mode="json", exclude_unset=True) + for axis in group + ] + for group in calculation.axes + ] + ), + "add_missing_variables": calculation.add_missing_variables, + }, + "force": calculation.force, + "trigger": "initial", + } + spm = saved_spm_choice(household.household_data) + if spm is not None: + submission["spm"] = spm + return HouseholdCalculationRead.from_simulation( + self._simulation.submit(submission) + ) + + def get_calculation(self, calculation_id: UUID) -> HouseholdCalculationRead: + """Return the calculation, or raise HouseholdCalculationNotFoundError.""" + calculation = self._simulation.read(calculation_id) + if calculation is None: + msg = f"Household calculation {calculation_id} not found." + raise HouseholdCalculationNotFoundError(msg) + return HouseholdCalculationRead.from_simulation(calculation) diff --git a/policyengine_api/services/v2/household_calculations/simulation_client.py b/policyengine_api/services/v2/household_calculations/simulation_client.py new file mode 100644 index 000000000..1474bd462 --- /dev/null +++ b/policyengine_api/services/v2/household_calculations/simulation_client.py @@ -0,0 +1,171 @@ +"""The simulation service's household calculation routes, through its client. + +Calls go to the Cloud Run Simulation Entrypoint (``SIMULATION_ENTRYPOINT_URL``), +the only simulation service with these routes, with the gateway credentials +the Public API already uses (``GATEWAY_AUTH_*``) and the request's +correlation headers. A 400 is a refusal, whose reason the service gives; any +other failure, including a 503, no answer or missing configuration, means the +service is unavailable. + +One client serves the process, so its connections and its gateway token are +reused across requests. +""" + +from __future__ import annotations + +from functools import lru_cache +from http import HTTPStatus +from typing import TYPE_CHECKING, Any + +import httpx +from policyengine_api_simulation_client import AuthenticatedClient +from policyengine_api_simulation_client.api.default import ( + get_household_calculation, + submit_household_calculation, +) +from policyengine_api_simulation_client.models.household_calculation_submission import ( + HouseholdCalculationSubmission, +) +from policyengine_observability import ( # type: ignore[import-untyped] # the package has no type information + instrument_httpx, +) + +from policyengine_api.libs.gateway_auth import ( + GatewayAuthError, + GatewayAuthTokenProvider, + GatewayBearerAuth, + _require_all_or_none_gateway_auth_env, + gateway_auth_required, +) +from policyengine_api.libs.simulation_entrypoint import ( + _attach_current_request_id, + resolve_simulation_entrypoint_url, +) +from policyengine_api.observability import get_runtime +from policyengine_api.services.v2.household_calculations.types import ( + SimulationServiceRefusedError, + SimulationServiceUnavailableError, +) + +if TYPE_CHECKING: + from uuid import UUID + +REQUEST_TIMEOUT_SECONDS = 30.0 +CLOUD_RUN_SIMULATION_ENTRYPOINT = "cloud_run_simulation_entrypoint" + + +def _answer(response: Any) -> tuple[int, object]: # noqa: ANN401 # the generated client's Response + try: + body = httpx.Response(response.status_code, content=response.content).json() + except ValueError: + body = None + return int(response.status_code), body + + +def _gateway_auth() -> GatewayBearerAuth | None: + """Return the gateway credentials, or None when none are set and none required.""" + _require_all_or_none_gateway_auth_env() + token_provider = GatewayAuthTokenProvider() + if token_provider.configured: + return GatewayBearerAuth(token_provider) + if gateway_auth_required(): + msg = "Gateway auth is required but GATEWAY_AUTH_* is not set." + raise GatewayAuthError(msg) + return None + + +def _refusal_reason(body: object) -> str: + reason = body.get("detail") if isinstance(body, dict) else None + if isinstance(reason, str): + return reason + return "The simulation service refused the household calculation." + + +class SimulationHouseholdClient: + """Submit and read household calculations on the simulation service.""" + + def __init__( + self, base_url: str, *, httpx_args: dict[str, Any] | None = None + ) -> None: + """Call the service at ``base_url`` with these httpx client arguments.""" + http_client = httpx.Client( + base_url=base_url, + timeout=httpx.Timeout(REQUEST_TIMEOUT_SECONDS), + event_hooks={"request": [_attach_current_request_id]}, + **(httpx_args or {}), + ) + instrument_httpx(http_client, get_runtime()) + self._base_url = base_url.rstrip("/") + # The generated routes take an AuthenticatedClient. Its fixed token is + # sent only from an httpx client it builds itself; the one set here + # authenticates through ``httpx_args`` (GatewayBearerAuth, which + # fetches and renews the gateway token), so the empty token is unused. + self._client = AuthenticatedClient( + base_url=base_url, token="" + ).set_httpx_client(http_client) + + @property + def base_url(self) -> str: + """Return the simulation service's base URL, without a trailing slash.""" + return self._base_url + + @classmethod + def from_environment(cls) -> SimulationHouseholdClient: + """Call the configured Cloud Run Simulation Entrypoint with gateway auth. + + Raise SimulationServiceUnavailableError when the URL is missing or + invalid, or gateway auth is partly configured or required but unset. + """ + try: + base_url = resolve_simulation_entrypoint_url( + CLOUD_RUN_SIMULATION_ENTRYPOINT + ) + auth = _gateway_auth() + except (ValueError, GatewayAuthError) as error: + msg = "The simulation service is not configured." + raise SimulationServiceUnavailableError(msg) from error + return cls(base_url, httpx_args={} if auth is None else {"auth": auth}) + + def submit(self, submission: dict[str, Any]) -> dict[str, Any]: + """Submit a calculation and return the service's calculation answer.""" + try: + response = submit_household_calculation.sync_detailed( + client=self._client, + body=HouseholdCalculationSubmission.from_dict(submission), + ) + except httpx.HTTPError as error: + msg = "The simulation service could not be reached." + raise SimulationServiceUnavailableError(msg) from error + status, body = _answer(response) + if status == HTTPStatus.BAD_REQUEST: + raise SimulationServiceRefusedError(_refusal_reason(body)) + if status != HTTPStatus.ACCEPTED or not isinstance(body, dict): + msg = f"The simulation service answered {status}." + raise SimulationServiceUnavailableError(msg) + return body + + def read(self, calculation_id: UUID) -> dict[str, Any] | None: + """Return the service's calculation answer, or None when it has none.""" + try: + response = get_household_calculation.sync_detailed( + calculation_id, client=self._client + ) + except httpx.HTTPError as error: + msg = "The simulation service could not be reached." + raise SimulationServiceUnavailableError(msg) from error + status, body = _answer(response) + if status == HTTPStatus.NOT_FOUND: + return None + if status != HTTPStatus.OK or not isinstance(body, dict): + msg = f"The simulation service answered {status}." + raise SimulationServiceUnavailableError(msg) + return body + + +@lru_cache(maxsize=1) +def process_simulation_household_client() -> SimulationHouseholdClient: + """Return the process's client, built from the environment on first use. + + A configuration error is raised again on the next call, not kept. + """ + return SimulationHouseholdClient.from_environment() diff --git a/policyengine_api/services/v2/household_calculations/types.py b/policyengine_api/services/v2/household_calculations/types.py new file mode 100644 index 000000000..3cc20ceba --- /dev/null +++ b/policyengine_api/services/v2/household_calculations/types.py @@ -0,0 +1,205 @@ +"""Framework-independent data exchanged by the household calculation layers.""" + +from __future__ import annotations + +from datetime import datetime +from typing import TYPE_CHECKING, Annotated, Any, Literal, Protocol, Self +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +from policyengine_api.query_parameters import CountryId, DefaultYear + +if TYPE_CHECKING: + from collections.abc import Mapping + + from policyengine_api.services.v2.countries.types import CountryRead + from policyengine_api.services.v2.households.types import HouseholdRead + from policyengine_api.services.v2.metadata.types import ( + MetadataModelSelectionResult, + ) + from policyengine_api.services.v2.policies.types import PolicyRead + +# The simulation service's limits on axes. +MAXIMUM_AXIS_POINTS = 10_000 +MAXIMUM_AXIS_GROUPS = 8 +MAXIMUM_PARALLEL_AXES = 32 +CalculationStatus = Literal["pending", "dispatching", "running", "succeeded", "failed"] + + +class StrictCalculationModel(BaseModel): + """Base for household calculation data: no unknown fields, finite numbers.""" + + model_config = ConfigDict(extra="forbid", allow_inf_nan=False, frozen=True) + + +class HouseholdCalculationAxis(StrictCalculationModel): + """One input varied across points. + + The variable ``name`` takes ``count`` evenly spaced values from ``min`` to + ``max`` in ``period`` (the household's default period when omitted), for + the entity member at position ``index`` (the first when omitted). + """ + + name: Annotated[str, Field(min_length=1, max_length=255)] + count: Annotated[int, Field(ge=1, le=MAXIMUM_AXIS_POINTS)] + min: int | float + max: int | float + period: Annotated[str, Field(min_length=1, max_length=64)] | int | None = None + index: Annotated[int, Field(ge=0)] | None = None + + +AxisGroup = Annotated[ + list[HouseholdCalculationAxis], + Field(min_length=1, max_length=MAXIMUM_PARALLEL_AXES), +] + + +class HouseholdCalculationInput(StrictCalculationModel): + """A stored household to calculate under a stored policy, or current law. + + ``year`` defaults to the household's default year. ``axes`` vary inputs + across points; ``add_missing_variables`` adds every yearly, monthly and + eternal variable the household lacks, as ``/calculate-full`` does. + ``force`` calculates again even when a recent calculation can be reused. + """ + + country_id: CountryId + household_id: UUID + policy_id: UUID | None = None + year: DefaultYear | None = None + axes: Annotated[list[AxisGroup], Field(max_length=MAXIMUM_AXIS_GROUPS)] | None = ( + None + ) + add_missing_variables: bool = False + force: bool = False + + +class HouseholdCalculationResult(StrictCalculationModel): + """A succeeded calculation's household and its notes. + + ``household`` is the stored household in the situation shape, its people + and groups named by their record IDs, with each requested value filled + in. ``spm_config`` is the SPM selection the calculation applied and + ``spm_provenance`` its receipt, both null for a country without SPM. + """ + + household: dict[str, JsonValue] + warnings: list[str] + spm_config: dict[str, JsonValue] | None + spm_provenance: dict[str, JsonValue] | None + policyengine_version: str | None + + +class HouseholdCalculationFailure(StrictCalculationModel): + """Why a calculation failed, with the simulation service's error detail.""" + + code: str + message: str + detail: dict[str, JsonValue] | None + + +class HouseholdCalculationRead(StrictCalculationModel): + """A household calculation as it stands: what it calculates and its outcome.""" + + calculation_id: UUID + country_id: str + household_id: UUID + policy_id: UUID | None + year: int + catalog_version_id: UUID + status: CalculationStatus + created_at: datetime + completed_at: datetime | None + result: HouseholdCalculationResult | None + error: HouseholdCalculationFailure | None + + @classmethod + def from_simulation(cls, calculation: Mapping[str, Any]) -> Self: + """Build the read from the simulation service's calculation answer.""" + result = calculation.get("result") + error = calculation.get("error") + bundle = result.get("bundle") if isinstance(result, dict) else None + return cls( + calculation_id=calculation["calculation_id"], + country_id=calculation["country"], + household_id=calculation["household_id"], + policy_id=calculation["policy_id"], + year=calculation["year"], + catalog_version_id=calculation["catalog_version_id"], + status=calculation["status"], + created_at=calculation["created_at"], + completed_at=calculation["completed_at"], + result=( + None + if result is None + else HouseholdCalculationResult( + household=result["household"], + warnings=result["warnings"], + spm_config=result["spm_config"], + spm_provenance=result["spm_provenance"], + policyengine_version=( + bundle.get("policyengine_version") + if isinstance(bundle, dict) + else None + ), + ) + ), + error=None if error is None else HouseholdCalculationFailure(**error), + ) + + +class SimulationServiceUnavailableError(RuntimeError): + """The simulation service could not be reached, or could not answer.""" + + +class SimulationServiceRefusedError(RuntimeError): + """The simulation service refused the calculation; the message says why.""" + + +class CountrySource(Protocol): + """The countries table (``V2CountryService``).""" + + def list_countries(self) -> list[CountryRead]: + """Return every served country and its capability.""" + ... + + +class HouseholdSource(Protocol): + """Stored households (``V2HouseholdService``).""" + + def get_household(self, *, country_id: str, household_id: UUID) -> HouseholdRead: + """Return one household of this country, or raise HouseholdNotFoundError.""" + ... + + +class PolicySource(Protocol): + """Stored policies (``V2PolicyService``).""" + + def get_policy(self, *, country_id: str, policy_id: UUID) -> PolicyRead: + """Return one policy of this country, or raise PolicyNotFoundError.""" + ... + + +class CatalogSource(Protocol): + """The country's catalog versions (a v2 metadata reader).""" + + def get_model_by_country(self, country_id: str) -> MetadataModelSelectionResult: + """Return the country's model and the running catalog version.""" + ... + + def close(self) -> None: + """Close the reader's database session.""" + ... + + +class SimulationHouseholdCalculations(Protocol): + """The simulation service's household calculation routes.""" + + def submit(self, submission: dict[str, Any]) -> dict[str, Any]: + """Submit a calculation and return it as the service answers.""" + ... + + def read(self, calculation_id: UUID) -> dict[str, Any] | None: + """Return the calculation as the service answers, or None if unknown.""" + ... diff --git a/policyengine_api/services/v2/household_calculations/validators.py b/policyengine_api/services/v2/household_calculations/validators.py new file mode 100644 index 000000000..507590131 --- /dev/null +++ b/policyengine_api/services/v2/household_calculations/validators.py @@ -0,0 +1,47 @@ +"""Checks of a household calculation against what is already loaded.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterable + + from policyengine_api.services.v2.countries.types import CountryRead + + +class HouseholdCalculationCountryError(ValueError): + """The country does not offer household calculations.""" + + +class HouseholdCalculationYearError(ValueError): + """The request names no year and the household has no default year.""" + + +class HouseholdCalculationNotFoundError(LookupError): + """The simulation service has no such household calculation.""" + + +def require_household_country( + countries: Iterable[CountryRead], country_id: str +) -> None: + """Raise HouseholdCalculationCountryError unless the country calculates households. + + A household-only country (household calculations without society-wide + reports) passes. + """ + for country in countries: + if country.id == country_id and country.household_enabled: + return + msg = f"Household calculations are not available for country {country_id}." + raise HouseholdCalculationCountryError(msg) + + +def calculation_year(requested_year: int | None, default_year: int | None) -> int: + """Return the requested year, or else the household's default year.""" + if requested_year is not None: + return requested_year + if default_year is None: + msg = "This household has no default year; send the year to calculate." + raise HouseholdCalculationYearError(msg) + return default_year diff --git a/pyproject.toml b/pyproject.toml index 4d225b491..234fc0e6b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,7 @@ dependencies = [ "psycopg[binary]>=3.3,<4", "pyjwt[crypto]>=2.10,<3", "typing-extensions>=4.4,<5", + "policyengine-api-simulation-client>=0.20261027.288", ] [project.optional-dependencies] @@ -249,6 +250,7 @@ runtime-evaluated-base-classes = [ "policyengine_api.data.v2.models.base.TimestampedModel", "policyengine_api.data.v2.models.comparison_runs.Stage12RunTimestamps", "policyengine_api.data.v2.models.reports.GeographicImpactBase", + "policyengine_api.fastapi_routes.v2.household_calculations.response_models.StrictHouseholdCalculationAPIModel", "policyengine_api.fastapi_routes.v2.households.document_models.HouseholdEntityRecord", "policyengine_api.fastapi_routes.v2.households.document_models.StrictHouseholdDocumentModel", "policyengine_api.fastapi_routes.v2.households.response_models.HouseholdSuccessResponse", @@ -273,6 +275,8 @@ runtime-evaluated-base-classes = [ "policyengine_api.query_parameters.PaginationQuery", "policyengine_api.query_parameters.StrictQueryParameters", "policyengine_api.services.v2.associations.types.StrictAssociationModel", + "policyengine_api.services.v2.household_calculations.types.HouseholdCalculationInput", + "policyengine_api.services.v2.household_calculations.types.StrictCalculationModel", "policyengine_api.services.v2.households.types.HouseholdCreationInput", "policyengine_api.services.v2.households.types.StrictHouseholdInput", "policyengine_api.services.v2.metadata.types.MetadataModelVersion", diff --git a/tests/unit/v2/household_calculation_support.py b/tests/unit/v2/household_calculation_support.py new file mode 100644 index 000000000..b28ab2ba0 --- /dev/null +++ b/tests/unit/v2/household_calculation_support.py @@ -0,0 +1,248 @@ +"""Fakes and records shared by the household calculation tests.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import UTC, datetime +from types import SimpleNamespace +from typing import Any +from uuid import UUID + +from policyengine_api.services.v2.countries.types import CountryRead +from policyengine_api.services.v2.households.types import HouseholdRead +from policyengine_api.services.v2.households.validators import HouseholdNotFoundError +from policyengine_api.services.v2.policies.types import ( + PolicyParameterValueRead, + PolicyRead, +) +from policyengine_api.services.v2.policies.validators import PolicyNotFoundError + +HOUSEHOLD_ID = UUID("00000000-0000-4000-8000-0000000000a1") +POLICY_ID = UUID("00000000-0000-4000-8000-0000000000b1") +CALCULATION_ID = UUID("00000000-0000-4000-8000-0000000000c1") +SIMULATION_ID = UUID("00000000-0000-4000-8000-0000000000d1") +CATALOG_VERSION_ID = UUID("00000000-0000-4000-8000-0000000000e1") +POLICYENGINE_VERSION = "6.10.0" +CREATED_AT = datetime(2026, 10, 27, 12, 0, tzinfo=UTC) +SPM_CHOICE = {"scenario": "ce_trend", "geography_kind": "national"} +HOUSEHOLD_DATA = { + "people": [ + { + "id": "person-1", + "values": {"age": {"2026": 40}, "employment_income": {"2026": 30000}}, + "memberships": { + "household": "household-1", + "family": "family-1", + "tax_unit": "tax_unit-1", + "spm_unit": "spm_unit-1", + "marital_unit": "marital_unit-1", + }, + } + ], + "household": [{"id": "household-1", "values": {"state_name": {"2026": "CA"}}}], + "family": [{"id": "family-1", "values": {}}], + "tax_unit": [{"id": "tax_unit-1", "values": {}}], + "spm_unit": [{"id": "spm_unit-1", "values": {}}], + "marital_unit": [{"id": "marital_unit-1", "values": {}}], + "spm": SPM_CHOICE, +} +SITUATION = { + "people": {"person-1": {"age": {"2026": 40}, "employment_income": {"2026": 30000}}}, + "households": { + "household-1": {"members": ["person-1"], "state_name": {"2026": "CA"}} + }, + "families": {"family-1": {"members": ["person-1"]}}, + "tax_units": {"tax_unit-1": {"members": ["person-1"]}}, + "spm_units": {"spm_unit-1": {"members": ["person-1"]}}, + "marital_units": {"marital_unit-1": {"members": ["person-1"]}}, +} +SPM_CONFIG = { + "forecast_content_sha256": "3d" * 32, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, +} +SPM_PROVENANCE = {"schema_version": "canonical-spm-provenance-v2", "years": ["2026"]} +BUNDLE = { + "policyengine_version": POLICYENGINE_VERSION, + "core_package_version": "3.33.0", + "country_package_name": "policyengine-us", + "country_package_version": "1.450.0", + "bundle_manifest_sha256": "a" * 64, +} + + +def simulation_calculation(**changes: object) -> dict[str, Any]: + """Return the simulation service's answer for a running calculation.""" + answer: dict[str, Any] = { + "calculation_id": str(CALCULATION_ID), + "simulation_id": str(SIMULATION_ID), + "country": "us", + "household_id": str(HOUSEHOLD_ID), + "policy_id": str(POLICY_ID), + "options": {}, + "catalog_version_id": str(CATALOG_VERSION_ID), + "year": 2026, + "attempt": 1, + "trigger": "initial", + "status": "running", + "created_at": CREATED_AT.isoformat(), + "dispatched_at": CREATED_AT.isoformat(), + "completed_at": None, + "result": None, + "error": None, + } + answer.update(changes) + return answer + + +def succeeded_calculation() -> dict[str, Any]: + """Return the simulation service's answer once the calculation succeeded.""" + return simulation_calculation( + status="succeeded", + completed_at=CREATED_AT.isoformat(), + result={ + "household": { + **SITUATION, + "households": { + "household-1": { + "members": ["person-1"], + "state_name": {"2026": "CA"}, + "household_net_income": {"2026": 27000.5}, + } + }, + }, + "warnings": ["The input `age_group` is deprecated and was ignored."], + "spm_config": SPM_CONFIG, + "spm_provenance": SPM_PROVENANCE, + "bundle": BUNDLE, + }, + ) + + +def household_read(*, default_year: int | None = 2026) -> HouseholdRead: + """Return the stored US household.""" + return HouseholdRead( + id=HOUSEHOLD_ID, + country_id="us", + default_year=default_year, + household_data=HOUSEHOLD_DATA, + created_at=CREATED_AT, + updated_at=CREATED_AT, + ) + + +def policy_read() -> PolicyRead: + """Return the stored US reform with one parameter value.""" + return PolicyRead( + id=POLICY_ID, + country_id="us", + tax_benefit_model_id=CATALOG_VERSION_ID, + tax_benefit_model_version_id=CATALOG_VERSION_ID, + created_at=CREATED_AT, + updated_at=CREATED_AT, + parameter_values=( + PolicyParameterValueRead( + id=POLICY_ID, + parameter_id=POLICY_ID, + parameter_name="gov.irs.credits.ctc.amount.base[0].amount", + value=3000, + start_date=datetime(2026, 1, 1, tzinfo=UTC), + end_date=None, + ), + ), + ) + + +@dataclass +class FakeCountries: + """The countries table: the US with every calculation, the UK households only.""" + + rows: tuple[CountryRead, ...] = ( + CountryRead(id="uk", household_enabled=True, society_wide_enabled=False), + CountryRead(id="us", household_enabled=True, society_wide_enabled=True), + ) + + def list_countries(self) -> list[CountryRead]: + """Return every row.""" + return list(self.rows) + + +class FakeHouseholds: + """Stored households by country and ID.""" + + def __init__(self) -> None: + """Hold the stored US household.""" + self.households = {("us", HOUSEHOLD_ID): household_read()} + self.calls: list[tuple[str, UUID]] = [] + + def get_household(self, *, country_id: str, household_id: UUID) -> HouseholdRead: + """Return the household or raise HouseholdNotFoundError.""" + self.calls.append((country_id, household_id)) + household = self.households.get((country_id, household_id)) + if household is None: + msg = f"Household {household_id} not found." + raise HouseholdNotFoundError(msg) + return household + + +class FakePolicies: + """Stored policies by country and ID.""" + + def __init__(self) -> None: + """Hold the stored US reform.""" + self.policies = {("us", POLICY_ID): policy_read()} + + def get_policy(self, *, country_id: str, policy_id: UUID) -> PolicyRead: + """Return the policy or raise PolicyNotFoundError.""" + policy = self.policies.get((country_id, policy_id)) + if policy is None: + msg = f"Policy {policy_id} not found." + raise PolicyNotFoundError(msg) + return policy + + +class FakeCatalogs: + """Each country's current catalog version, and whether it was closed.""" + + def __init__(self) -> None: + """Start open.""" + self.closed = False + + def get_model_by_country(self, country_id: str) -> SimpleNamespace: + """Return the running catalog version's ID and PolicyEngine.py version.""" + return SimpleNamespace( + policyengine_version=POLICYENGINE_VERSION, + model_version=SimpleNamespace(id=CATALOG_VERSION_ID), + country_id=country_id, + ) + + def close(self) -> None: + """Record the close.""" + self.closed = True + + +class FakeSimulation: + """The simulation service's household routes, as scripted answers.""" + + def __init__(self) -> None: + """Answer submissions with a running calculation.""" + self.submissions: list[dict[str, Any]] = [] + self.answer: dict[str, Any] = simulation_calculation() + self.read_answers: dict[UUID, dict[str, Any]] = {} + self.error: Exception | None = None + + def submit(self, submission: dict[str, Any]) -> dict[str, Any]: + """Record the submission and answer it.""" + self.submissions.append(submission) + if self.error is not None: + raise self.error + return self.answer + + def read(self, calculation_id: UUID) -> dict[str, Any] | None: + """Return the scripted calculation, or None.""" + if self.error is not None: + raise self.error + return self.read_answers.get(calculation_id) diff --git a/tests/unit/v2/test_household_calculation_routes.py b/tests/unit/v2/test_household_calculation_routes.py new file mode 100644 index 000000000..a6f396bd7 --- /dev/null +++ b/tests/unit/v2/test_household_calculation_routes.py @@ -0,0 +1,224 @@ +"""The /v2/household-calculations routes: submit a stored household, then poll it.""" + +from __future__ import annotations + +from uuid import uuid4 + +import pytest +from fastapi.testclient import TestClient +from flask import Flask + +from policyengine_api.asgi_factory import create_asgi_app +from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies +from policyengine_api.migration_flags import ( + RouteImplementation, + RouteImplementationSettings, +) +from policyengine_api.services.v2.countries.types import CountryRead +from policyengine_api.services.v2.household_calculations.services import ( + V2HouseholdCalculationService, +) +from policyengine_api.services.v2.household_calculations.types import ( + SimulationServiceRefusedError, + SimulationServiceUnavailableError, +) +from tests.unit.v2.household_calculation_support import ( + CALCULATION_ID, + HOUSEHOLD_ID, + POLICY_ID, + SPM_CONFIG, + SPM_PROVENANCE, + FakeCatalogs, + FakeCountries, + FakeHouseholds, + FakePolicies, + FakeSimulation, + succeeded_calculation, +) + +PATH = "/v2/household-calculations" + + +class App: + """A client of the native app serving household calculations over fakes.""" + + def __init__(self, countries: FakeCountries | None = None) -> None: + """Build the fakes and the app.""" + self.simulation = FakeSimulation() + self.households = FakeHouseholds() + self.service = V2HouseholdCalculationService( + countries=countries or FakeCountries(), + households=self.households, + policies=FakePolicies(), + catalogs=FakeCatalogs(), + simulation=self.simulation, + ) + dependencies = NativeRouteDependencies( + readiness_probe=lambda: True, + gateway_client_factory=lambda: None, + metadata_reader_factory=lambda: None, + specification_provider=dict, + v2_household_calculation_service_factory=lambda: self.service, + ) + settings = RouteImplementationSettings( + health=RouteImplementation.FLASK_FALLBACK, + specification=RouteImplementation.FLASK_FALLBACK, + metadata=RouteImplementation.FLASK_FALLBACK, + ) + self.client = TestClient( + create_asgi_app( + Flask(__name__), dependencies=dependencies, route_settings=settings + ), + raise_server_exceptions=False, + ) + + +def body(**changes: object) -> dict[str, object]: + """Return a request calculating the stored household under the reform.""" + request: dict[str, object] = { + "country_id": "us", + "household_id": str(HOUSEHOLD_ID), + "policy_id": str(POLICY_ID), + } + request.update(changes) + return request + + +def test_a_calculation_is_accepted_without_an_account() -> None: + app = App() + + response = app.client.post(PATH, json=body()) + + assert response.status_code == 202 + payload = response.json() + assert payload["status"] == "ok" + assert payload["result"]["calculation_id"] == str(CALCULATION_ID) + assert payload["result"]["status"] == "running" + assert payload["result"]["household_id"] == str(HOUSEHOLD_ID) + assert payload["result"]["result"] is None + + +def test_a_calculation_with_axes_sends_them() -> None: + app = App() + axes = [[{"name": "employment_income", "count": 3, "min": 0, "max": 50000}]] + + response = app.client.post(PATH, json=body(axes=axes)) + + assert response.status_code == 202 + assert app.simulation.submissions[0]["options"]["axes"] == axes + + +def test_a_succeeded_us_calculation_reads_back_its_spm_choice() -> None: + app = App() + app.simulation.read_answers[CALCULATION_ID] = succeeded_calculation() + + response = app.client.get(f"{PATH}/{CALCULATION_ID}") + + assert response.status_code == 200 + result = response.json()["result"] + assert result["status"] == "succeeded" + assert result["result"]["spm_config"] == SPM_CONFIG + assert result["result"]["spm_provenance"] == SPM_PROVENANCE + assert result["result"]["warnings"] == [ + "The input `age_group` is deprecated and was ignored." + ] + + +def test_a_country_without_household_calculations_gets_422_naming_it() -> None: + app = App( + FakeCountries( + rows=( + CountryRead( + id="us", household_enabled=False, society_wide_enabled=False + ), + ) + ) + ) + + response = app.client.post(PATH, json=body()) + + assert response.status_code == 422 + assert "us" in response.json()["message"] + assert app.simulation.submissions == [] + + +def test_a_household_only_country_is_accepted() -> None: + app = App() + stored = app.households.households[("us", HOUSEHOLD_ID)] + app.households.households[("uk", HOUSEHOLD_ID)] = stored.__class__( + **{ + **stored.__dict__, + "country_id": "uk", + "household_data": { + "people": [ + { + "id": "person-1", + "values": {}, + "memberships": { + "household": "household-1", + "benunit": "benunit-1", + }, + } + ], + "household": [{"id": "household-1", "values": {}}], + "benunit": [{"id": "benunit-1", "values": {}}], + }, + } + ) + + response = app.client.post(PATH, json=body(country_id="uk", policy_id=None)) + + assert response.status_code == 202 + + +@pytest.mark.parametrize( + "changes", + [{"household_id": str(uuid4())}, {"policy_id": str(uuid4())}], + ids=["household", "policy"], +) +def test_an_unknown_household_or_policy_is_not_found(changes) -> None: + response = App().client.post(PATH, json=body(**changes)) + + assert response.status_code == 404 + + +def test_an_unknown_calculation_is_not_found() -> None: + assert App().client.get(f"{PATH}/{uuid4()}").status_code == 404 + + +@pytest.mark.parametrize( + ("error", "status_code"), + [ + (SimulationServiceUnavailableError("unreachable"), 503), + (SimulationServiceRefusedError("no household function"), 503), + ], + ids=["unreachable", "refused"], +) +def test_a_simulation_service_that_cannot_calculate_answers_503( + error, status_code +) -> None: + app = App() + app.simulation.error = error + + response = app.client.post(PATH, json=body()) + + assert response.status_code == status_code + + +@pytest.mark.parametrize( + "changes", + [ + {"axes": [[{"name": "employment_income", "count": 0, "min": 0, "max": 1}]]}, + {"country_id": "ca"}, + {"unexpected": True}, + {"year": 1800}, + ], + ids=["axis-count", "country", "unknown-field", "year"], +) +def test_a_malformed_request_is_refused_before_anything_is_read(changes) -> None: + app = App() + + response = app.client.post(PATH, json=body(**changes)) + + assert response.status_code == 422 + assert app.simulation.submissions == [] diff --git a/tests/unit/v2/test_household_calculation_service.py b/tests/unit/v2/test_household_calculation_service.py new file mode 100644 index 000000000..b2b16bbc1 --- /dev/null +++ b/tests/unit/v2/test_household_calculation_service.py @@ -0,0 +1,246 @@ +"""Sending a stored household's calculation to the simulation service; reading it.""" + +from __future__ import annotations + +from uuid import uuid4 + +import pytest + +from policyengine_api.services.v2.countries.types import CountryRead +from policyengine_api.services.v2.household_calculations.services import ( + V2HouseholdCalculationService, +) +from policyengine_api.services.v2.household_calculations.types import ( + HouseholdCalculationInput, + SimulationServiceUnavailableError, +) +from policyengine_api.services.v2.household_calculations.validators import ( + HouseholdCalculationCountryError, + HouseholdCalculationNotFoundError, + HouseholdCalculationYearError, +) +from policyengine_api.services.v2.households.validators import HouseholdNotFoundError +from policyengine_api.services.v2.policies.validators import PolicyNotFoundError +from tests.unit.v2.household_calculation_support import ( + CALCULATION_ID, + CATALOG_VERSION_ID, + HOUSEHOLD_ID, + POLICY_ID, + POLICYENGINE_VERSION, + SITUATION, + SPM_CHOICE, + SPM_CONFIG, + SPM_PROVENANCE, + FakeCatalogs, + FakeCountries, + FakeHouseholds, + FakePolicies, + FakeSimulation, + household_read, + simulation_calculation, + succeeded_calculation, +) + +AXES = [[{"name": "employment_income", "count": 5, "min": 0, "max": 80000}]] + + +class Parts: + """The service's collaborators, kept for the assertions.""" + + def __init__(self, countries: FakeCountries | None = None) -> None: + """Build every fake.""" + self.countries = countries or FakeCountries() + self.households = FakeHouseholds() + self.policies = FakePolicies() + self.catalogs = FakeCatalogs() + self.simulation = FakeSimulation() + self.service = V2HouseholdCalculationService( + countries=self.countries, + households=self.households, + policies=self.policies, + catalogs=self.catalogs, + simulation=self.simulation, + ) + + +def calculation_input(**changes: object) -> HouseholdCalculationInput: + """Return a calculation of the stored US household under the stored reform.""" + values: dict[str, object] = { + "country_id": "us", + "household_id": HOUSEHOLD_ID, + "policy_id": POLICY_ID, + } + values.update(changes) + return HouseholdCalculationInput.model_validate(values) + + +def test_a_calculation_sends_the_stored_household_policy_and_catalog() -> None: + parts = Parts() + + calculation = parts.service.submit_calculation( + calculation_input(axes=AXES, add_missing_variables=True, force=True) + ) + + (submission,) = parts.simulation.submissions + assert submission == { + "country": "us", + "catalog_version_id": str(CATALOG_VERSION_ID), + "policyengine_version": POLICYENGINE_VERSION, + "year": 2026, + "household_id": str(HOUSEHOLD_ID), + "household": SITUATION, + "spm": SPM_CHOICE, + "policy": { + "policy_id": str(POLICY_ID), + "parameters": { + "gov.irs.credits.ctc.amount.base[0].amount": { + "2026-01-01.2100-12-31": 3000 + } + }, + }, + "options": {"axes": AXES, "add_missing_variables": True}, + "force": True, + "trigger": "initial", + } + assert calculation.calculation_id == CALCULATION_ID + assert calculation.status == "running" + assert calculation.result is None + assert parts.catalogs.closed + + +def test_current_law_and_no_saved_choice_send_neither() -> None: + parts = Parts() + stored = household_read() + data = {key: value for key, value in stored.household_data.items() if key != "spm"} + parts.households.households[("us", HOUSEHOLD_ID)] = stored.__class__( + **{**stored.__dict__, "household_data": data} + ) + + parts.service.submit_calculation(calculation_input(policy_id=None)) + + (submission,) = parts.simulation.submissions + assert "spm" not in submission + assert submission["policy"] == {"policy_id": None, "parameters": {}} + assert submission["options"] == {"axes": None, "add_missing_variables": False} + + +def test_the_year_is_the_requested_one_or_the_households_default() -> None: + parts = Parts() + + parts.service.submit_calculation(calculation_input(year=2027)) + parts.households.households[("us", HOUSEHOLD_ID)] = household_read( + default_year=None + ) + + with pytest.raises(HouseholdCalculationYearError): + parts.service.submit_calculation(calculation_input()) + assert [submission["year"] for submission in parts.simulation.submissions] == [2027] + + +def test_a_country_without_household_calculations_is_refused_by_name() -> None: + parts = Parts( + FakeCountries( + rows=( + CountryRead( + id="us", household_enabled=False, society_wide_enabled=False + ), + ) + ) + ) + + with pytest.raises(HouseholdCalculationCountryError, match="us"): + parts.service.submit_calculation(calculation_input()) + + assert parts.households.calls == [] + assert parts.simulation.submissions == [] + + +def test_a_household_only_country_is_accepted() -> None: + parts = Parts() + uk_household = household_read() + parts.households.households[("uk", HOUSEHOLD_ID)] = uk_household.__class__( + **{ + **uk_household.__dict__, + "country_id": "uk", + "household_data": { + "people": [ + { + "id": "person-1", + "values": {"age": {"2026": 30}}, + "memberships": { + "household": "household-1", + "benunit": "benunit-1", + }, + } + ], + "household": [{"id": "household-1", "values": {}}], + "benunit": [{"id": "benunit-1", "values": {}}], + }, + } + ) + + parts.service.submit_calculation(calculation_input(country_id="uk", policy_id=None)) + + (submission,) = parts.simulation.submissions + assert submission["country"] == "uk" + assert set(submission["household"]) == {"people", "households", "benunits"} + + +def test_an_unknown_household_or_policy_sends_nothing() -> None: + parts = Parts() + + with pytest.raises(HouseholdNotFoundError): + parts.service.submit_calculation(calculation_input(household_id=uuid4())) + with pytest.raises(PolicyNotFoundError): + parts.service.submit_calculation(calculation_input(policy_id=uuid4())) + + assert parts.simulation.submissions == [] + + +def test_a_us_result_includes_the_spm_choice_applied_and_its_receipt() -> None: + parts = Parts() + parts.simulation.read_answers[CALCULATION_ID] = succeeded_calculation() + + calculation = parts.service.get_calculation(CALCULATION_ID) + + assert calculation.status == "succeeded" + assert calculation.result is not None + assert calculation.result.spm_config == SPM_CONFIG + assert calculation.result.spm_provenance == SPM_PROVENANCE + assert calculation.result.policyengine_version == POLICYENGINE_VERSION + assert calculation.result.household["households"]["household-1"][ + "household_net_income" + ] == {"2026": 27000.5} + + +def test_a_failed_calculation_reads_back_its_error() -> None: + parts = Parts() + detail = {"code": "invalid_inputs", "message": "Unrecognized inputs: incme"} + parts.simulation.read_answers[CALCULATION_ID] = simulation_calculation( + status="failed", + error={ + "code": "household_calculation_failed", + "message": "The household calculation failed: invalid_inputs.", + "detail": detail, + }, + ) + + calculation = parts.service.get_calculation(CALCULATION_ID) + + assert calculation.status == "failed" + assert calculation.error is not None + assert calculation.error.code == "household_calculation_failed" + assert calculation.error.detail == detail + + +def test_an_unknown_calculation_is_not_found() -> None: + with pytest.raises(HouseholdCalculationNotFoundError): + Parts().service.get_calculation(uuid4()) + + +def test_an_unreachable_simulation_service_is_reported() -> None: + parts = Parts() + parts.simulation.error = SimulationServiceUnavailableError("unreachable") + + with pytest.raises(SimulationServiceUnavailableError): + parts.service.submit_calculation(calculation_input()) diff --git a/tests/unit/v2/test_household_calculation_simulation_client.py b/tests/unit/v2/test_household_calculation_simulation_client.py new file mode 100644 index 000000000..2f73b5a82 --- /dev/null +++ b/tests/unit/v2/test_household_calculation_simulation_client.py @@ -0,0 +1,166 @@ +"""The simulation service's household calculation routes, through its client. + +These tests need a simulation client with the household calculation routes. +""" + +from __future__ import annotations + +import json +from uuid import uuid4 + +import httpx +import pytest + +from policyengine_api.services.v2.household_calculations.simulation_client import ( + SimulationHouseholdClient, +) +from policyengine_api.services.v2.household_calculations.types import ( + SimulationServiceRefusedError, + SimulationServiceUnavailableError, +) +from tests.unit.v2.household_calculation_support import ( + CALCULATION_ID, + CATALOG_VERSION_ID, + HOUSEHOLD_ID, + POLICYENGINE_VERSION, + SITUATION, + simulation_calculation, + succeeded_calculation, +) + +BASE_URL = "https://simulation.example" +SUBMISSION = { + "country": "us", + "catalog_version_id": str(CATALOG_VERSION_ID), + "policyengine_version": POLICYENGINE_VERSION, + "year": 2026, + "household_id": str(HOUSEHOLD_ID), + "household": SITUATION, + "spm": {"scenario": "ce_trend"}, + "policy": {"policy_id": None, "parameters": {}}, + "options": {"axes": None, "add_missing_variables": False}, + "force": False, + "trigger": "initial", +} + + +def client_answering( + status_code: int, payload: object, requests: list +) -> SimulationHouseholdClient: + """Return a client whose requests get one fixed answer.""" + + def answer(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(status_code, json=payload) + + return SimulationHouseholdClient( + BASE_URL, httpx_args={"transport": httpx.MockTransport(answer)} + ) + + +def test_a_submission_is_posted_and_its_calculation_returned() -> None: + requests: list[httpx.Request] = [] + client = client_answering(202, simulation_calculation(), requests) + + calculation = client.submit(SUBMISSION) + + assert calculation == simulation_calculation() + (request,) = requests + assert request.method == "POST" + assert str(request.url) == f"{BASE_URL}/v2/household-calculations" + assert json.loads(request.content) == SUBMISSION + + +def test_a_calculation_is_read_by_id() -> None: + requests: list[httpx.Request] = [] + client = client_answering(200, succeeded_calculation(), requests) + + assert client.read(CALCULATION_ID) == succeeded_calculation() + assert str(requests[0].url) == ( + f"{BASE_URL}/v2/household-calculations/{CALCULATION_ID}" + ) + + +def test_an_unknown_calculation_reads_as_none() -> None: + client = client_answering(404, {"detail": "No such household calculation."}, []) + + assert client.read(uuid4()) is None + + +def test_a_refusal_carries_the_services_reason() -> None: + client = client_answering( + 400, {"detail": "PolicyEngine.py 9.9.9 has no deployed v2 worker"}, [] + ) + + with pytest.raises(SimulationServiceRefusedError, match=r"9\.9\.9"): + client.submit(SUBMISSION) + + +@pytest.mark.parametrize("status_code", [500, 503]) +def test_a_service_that_cannot_answer_is_unavailable(status_code: int) -> None: + client = client_answering(status_code, {"detail": "unavailable"}, []) + + with pytest.raises(SimulationServiceUnavailableError): + client.submit(SUBMISSION) + with pytest.raises(SimulationServiceUnavailableError): + client.read(CALCULATION_ID) + + +def test_an_unreachable_service_is_unavailable() -> None: + def refuse(request: httpx.Request) -> httpx.Response: + message = "refused" + raise httpx.ConnectError(message, request=request) + + client = SimulationHouseholdClient( + BASE_URL, httpx_args={"transport": httpx.MockTransport(refuse)} + ) + + with pytest.raises(SimulationServiceUnavailableError): + client.read(CALCULATION_ID) + + +GATEWAY_AUTH_VARIABLES = ( + "GATEWAY_AUTH_ISSUER", + "GATEWAY_AUTH_AUDIENCE", + "GATEWAY_AUTH_CLIENT_ID", + "GATEWAY_AUTH_CLIENT_SECRET", + "GATEWAY_AUTH_CLIENT_SECRET_RESOURCE", + "GATEWAY_AUTH_REQUIRED", +) + + +@pytest.mark.parametrize( + "environment", + [ + {}, + {"SIMULATION_ENTRYPOINT_URL": "not a url"}, + { + "SIMULATION_ENTRYPOINT_URL": BASE_URL, + "GATEWAY_AUTH_ISSUER": "https://issuer.example/", + }, + {"SIMULATION_ENTRYPOINT_URL": BASE_URL, "GATEWAY_AUTH_REQUIRED": "true"}, + ], + ids=["no-url", "invalid-url", "partial-auth", "auth-required"], +) +def test_a_misconfigured_service_is_unavailable( + monkeypatch: pytest.MonkeyPatch, environment: dict[str, str] +) -> None: + for name in ("SIMULATION_ENTRYPOINT_URL", *GATEWAY_AUTH_VARIABLES): + monkeypatch.delenv(name, raising=False) + for name, value in environment.items(): + monkeypatch.setenv(name, value) + + with pytest.raises(SimulationServiceUnavailableError, match="not configured"): + SimulationHouseholdClient.from_environment() + + +def test_the_configured_service_is_called_without_auth_when_none_is_set( + monkeypatch: pytest.MonkeyPatch, +) -> None: + for name in GATEWAY_AUTH_VARIABLES: + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("SIMULATION_ENTRYPOINT_URL", f"{BASE_URL}/") + + client = SimulationHouseholdClient.from_environment() + + assert client.base_url == BASE_URL diff --git a/tests/unit/v2/test_metadata_routes.py b/tests/unit/v2/test_metadata_routes.py index a1e4708c1..701e2ff2d 100644 --- a/tests/unit/v2/test_metadata_routes.py +++ b/tests/unit/v2/test_metadata_routes.py @@ -642,6 +642,8 @@ def test_openapi_references_explicit_resource_response_schemas() -> None: "/v2/countries", "/v2/households", "/v2/households/{household_id}", + "/v2/household-calculations", + "/v2/household-calculations/{calculation_id}", "/v2/policies", "/v2/policies/{policy_id}", "/v2/user-households", diff --git a/uv.lock b/uv.lock index 14d4b7346..f9bdfaf54 100644 --- a/uv.lock +++ b/uv.lock @@ -2794,6 +2794,7 @@ dependencies = [ { name = "openai" }, { name = "packaging" }, { name = "policyengine", extra = ["models"] }, + { name = "policyengine-api-simulation-client" }, { name = "policyengine-canada" }, { name = "policyengine-il" }, { name = "policyengine-ng" }, @@ -2851,6 +2852,7 @@ requires-dist = [ { name = "openai" }, { name = "packaging", specifier = ">=24,<27" }, { name = "policyengine", extras = ["models"], specifier = "==6.10.0" }, + { name = "policyengine-api-simulation-client", specifier = ">=0.20261027.288" }, { name = "policyengine-canada", specifier = "==0.96.3" }, { name = "policyengine-il", specifier = "==0.1.0" }, { name = "policyengine-ng", specifier = "==0.5.1" }, @@ -2876,6 +2878,19 @@ requires-dist = [ ] provides-extras = ["dev"] +[[package]] +name = "policyengine-api-simulation-client" +version = "0.20261027.288" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "httpx" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/46/21/145e238dec20e79a1ce7fb4818271e08603fab88656ce4d7e5c7307ab90e/policyengine_api_simulation_client-0.20261027.288.tar.gz", hash = "sha256:5129724a3e5ba4b0dc257d421c32dd1be7f5c4a1ae3c1a9b60b20d7fab58f9f5", size = 44777, upload-time = "2026-10-27T17:02:39.052Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/27/85e9da0f0561d0548cbe23ec46e4bf63a4b814d7adbe1f0932c514b52191/policyengine_api_simulation_client-0.20261027.288-py3-none-any.whl", hash = "sha256:82b4aacb954de0b8c747b313c1b92ea1d6a028402618f3d819c8298b87977486", size = 99950, upload-time = "2026-10-27T17:02:37.398Z" }, +] + [[package]] name = "policyengine-canada" version = "0.96.3" From f6d1d4e8703dd4a1e28658365ae253c771330f5c Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 21:51:08 +0400 Subject: [PATCH 13/37] Check household calculations and their parity with /calculate in staging The parallel live staging job now runs a submit, poll and result cycle for a new US household; six calculations submitted at once against submit and finish time limits; a household across 401 axis points, which must finish; and a parity set of two US and two UK households (one of each with axes) compared value by value with API v1's /calculate on the same candidate. The job's time limit rises from 20 to 30 minutes for the Modal calculations. Co-Authored-By: Claude Opus 5.5 --- .github/scripts/run_live_staging_tests.sh | 2 + .github/workflows/push.yml | 2 +- docs/migration/stage-11-v2-households.md | 25 ++ .../live_household_calculations.py | 222 ++++++++++++++++++ .../test_live_v2_household_calculations.py | 169 +++++++++++++ .../test_live_v2_household_parity.py | 154 ++++++++++++ tests/unit/test_cloud_run_deploy_scripts.py | 6 + 7 files changed, 579 insertions(+), 1 deletion(-) create mode 100644 tests/fixtures/integration/live_household_calculations.py create mode 100644 tests/integration/test_live_v2_household_calculations.py create mode 100644 tests/integration/test_live_v2_household_parity.py diff --git a/.github/scripts/run_live_staging_tests.sh b/.github/scripts/run_live_staging_tests.sh index 39bc7a3cb..d7b5f6731 100755 --- a/.github/scripts/run_live_staging_tests.sh +++ b/.github/scripts/run_live_staging_tests.sh @@ -8,4 +8,6 @@ python -m pytest \ tests/integration/test_live_budget_window_cache.py \ tests/integration/test_live_calculate.py \ tests/integration/test_live_economy.py \ + tests/integration/test_live_v2_household_calculations.py \ + tests/integration/test_live_v2_household_parity.py \ -v diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index e8007f472..d18fa4f53 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -342,7 +342,7 @@ jobs: parallel-live-simulation-tests-staging-cloud-run: name: Run live simulations against the staging candidate runs-on: ubuntu-latest - timeout-minutes: 20 + timeout-minutes: 30 needs: deploy-cloud-run-staging if: | (github.repository == 'PolicyEngine/policyengine-api') diff --git a/docs/migration/stage-11-v2-households.md b/docs/migration/stage-11-v2-households.md index 2a58f7dd1..46c8730d3 100644 --- a/docs/migration/stage-11-v2-households.md +++ b/docs/migration/stage-11-v2-households.md @@ -292,3 +292,28 @@ Without the URL, or with gateway auth partly configured, both routes answer 503. They also answer 503 when the simulation service cannot be reached, answers with an error, or refuses the calculation (for example when no v2 worker is deployed for the catalog version's PolicyEngine.py version). + +### Staging checks + +The parallel live staging job (`run_live_staging_tests.sh`) runs two modules +against each tagged staging candidate, through staging API v1, the staging +Simulation Entrypoint, the v2 household functions on Modal and the staging v2 +database: + +- `test_live_v2_household_calculations.py`: a submit, poll and result cycle + for a new US household; six calculations submitted at once, each accepted + within 10 seconds and finished within 300 seconds; and one household varied + across 401 points of employment income, which must finish. The measured + times appear in the job log as `HouseholdCalculationMeasurement` warnings. +- `test_live_v2_household_parity.py`: two US and two UK households, one of + each with axes, calculated through `/v2/household-calculations` and through + API v1's `/calculate`; every calculated value and warning must match. + +`STAGING_HOUSEHOLD_PARALLEL_COUNT`, `STAGING_HOUSEHOLD_SUBMIT_LIMIT_SECONDS`, +`STAGING_HOUSEHOLD_FINISH_LIMIT_SECONDS` and +`STAGING_HOUSEHOLD_CALCULATION_TIMEOUT_SECONDS` (600 seconds per calculation) +change those limits; none is set in the workflow. The checks see neither the +number of Simulation Entrypoint instances that served the calculations nor +the household functions' memory use: a calculation that runs out of memory +fails, and so fails the check. Both modules fail until the staging Simulation +Entrypoint serves `/v2/household-calculations`. diff --git a/tests/fixtures/integration/live_household_calculations.py b/tests/fixtures/integration/live_household_calculations.py new file mode 100644 index 000000000..6fbf92647 --- /dev/null +++ b/tests/fixtures/integration/live_household_calculations.py @@ -0,0 +1,222 @@ +"""Households and helpers for the live staging household calculation checks. + +They need only httpx and the standard library, because the parallel live test +job installs nothing else. +""" + +from __future__ import annotations + +import math +import os +import time +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import httpx + +CALCULATION_PATH = "/v2/household-calculations" +CALCULATION_TIMEOUT_SECONDS = float( + os.environ.get("STAGING_HOUSEHOLD_CALCULATION_TIMEOUT_SECONDS", "600") +) +POLL_INTERVAL_SECONDS = float( + os.environ.get("STAGING_API_TEST_POLL_INTERVAL_SECONDS", "5") +) +TRANSIENT_STATUS_CODES = {500, 502, 503, 504} +FINISHED_STATUSES = {"succeeded", "failed"} +YEAR = "2026" +COLLECTIONS = { + "us": { + "household": "households", + "family": "families", + "tax_unit": "tax_units", + "spm_unit": "spm_units", + "marital_unit": "marital_units", + }, + "uk": {"household": "households", "benunit": "benunits"}, +} + + +@dataclass(frozen=True) +class LiveHousehold: + """A household whose people all share one group of each kind. + + ``people`` and ``groups`` map variable names to values by period; a null + value asks for the variable to be calculated. ``axes`` vary inputs across + points, as in a ``/calculate`` situation. + """ + + name: str + country_id: str + people: dict[str, dict[str, Any]] + groups: dict[str, dict[str, Any]] = field(default_factory=dict) + axes: list[list[dict[str, Any]]] | None = None + + def document(self) -> dict[str, Any]: + """Return the household as ``/v2/households`` stores it.""" + collections = COLLECTIONS[self.country_id] + document: dict[str, Any] = { + "people": [ + { + "id": person_id, + "values": values, + "memberships": { + collection: f"{collection}-1" for collection in collections + }, + } + for person_id, values in self.people.items() + ] + } + for collection in collections: + document[collection] = [ + {"id": f"{collection}-1", "values": self.groups.get(collection, {})} + ] + return document + + def situation(self) -> dict[str, Any]: + """Return the same household as a ``/calculate`` situation.""" + situation: dict[str, Any] = {"people": self.people} + for collection, plural in COLLECTIONS[self.country_id].items(): + situation[plural] = { + f"{collection}-1": { + "members": list(self.people), + **self.groups.get(collection, {}), + } + } + if self.axes is not None: + situation["axes"] = self.axes + return situation + + +@dataclass(frozen=True) +class FinishedCalculation: + """A calculation's final answer and how long each step took.""" + + item: dict[str, Any] + submit_seconds: float + total_seconds: float + + +def store_household(client: httpx.Client, household: LiveHousehold) -> str: + """Store the household through ``/v2/households`` and return its ID.""" + response = client.post( + "/v2/households", + params={"country_id": household.country_id}, + json={ + "country_id": household.country_id, + "default_year": int(YEAR), + "household_data": household.document(), + }, + ) + assert response.status_code in {200, 201}, response.text + return response.json()["result"]["item"]["id"] + + +def calculate( + client: httpx.Client, + household: LiveHousehold, + household_id: str, + *, + force: bool = False, +) -> FinishedCalculation: + """Submit the stored household's calculation and poll until it finishes.""" + started = time.monotonic() + body: dict[str, Any] = { + "country_id": household.country_id, + "household_id": household_id, + "year": int(YEAR), + "force": force, + } + if household.axes is not None: + body["axes"] = household.axes + response = client.post(CALCULATION_PATH, json=body) + submit_seconds = time.monotonic() - started + assert response.status_code == 202, response.text + item = response.json()["result"] + deadline = started + CALCULATION_TIMEOUT_SECONDS + while item["status"] not in FINISHED_STATUSES: + if time.monotonic() >= deadline: + msg = ( + f"Calculation {item['calculation_id']} did not finish within " + f"{CALCULATION_TIMEOUT_SECONDS:.0f} s; last status {item['status']}" + ) + raise AssertionError(msg) + time.sleep(POLL_INTERVAL_SECONDS) + polled = client.get(f"{CALCULATION_PATH}/{item['calculation_id']}") + if polled.status_code in TRANSIENT_STATUS_CODES: + continue + assert polled.status_code == 200, polled.text + item = polled.json()["result"] + return FinishedCalculation( + item=item, + submit_seconds=submit_seconds, + total_seconds=time.monotonic() - started, + ) + + +def calculate_with_v1(client: httpx.Client, household: LiveHousehold) -> dict: + """Return API v1's ``/calculate`` answer for the same household.""" + response = client.post( + f"/{household.country_id}/calculate", + json={"household": household.situation()}, + ) + assert response.status_code == 200, response.text + payload = response.json() + assert payload["status"] == "ok", payload + return payload + + +def _same_value(v1_value: object, v2_value: object) -> bool: + if isinstance(v1_value, list) and isinstance(v2_value, list): + return len(v1_value) == len(v2_value) and all( + _same_value(v1_item, v2_item) + for v1_item, v2_item in zip(v1_value, v2_value, strict=True) + ) + # API v1 writes an undefined number as a bare NaN; v2 writes "NaN". + if isinstance(v1_value, float) and math.isnan(v1_value): + return v2_value == "NaN" + if isinstance(v1_value, bool) or isinstance(v2_value, bool): + return v1_value == v2_value + if isinstance(v1_value, int | float) and isinstance(v2_value, int | float): + return math.isclose(v1_value, v2_value, rel_tol=1e-9, abs_tol=1e-6) + return v1_value == v2_value + + +def household_differences( + v1_household: dict[str, Any], v2_household: dict[str, Any] +) -> list[str]: + """Return every value that differs between the two calculated households. + + Each entity, variable and period in either household is compared, so a + value only one of them has is a difference too. + """ + differences: list[str] = [] + for plural in sorted(set(v1_household) | set(v2_household)): + v1_entities = v1_household.get(plural) + v2_entities = v2_household.get(plural) + if not isinstance(v1_entities, dict) or not isinstance(v2_entities, dict): + if v1_entities != v2_entities: + differences.append(f"{plural}: {v1_entities!r} != {v2_entities!r}") + continue + for entity_id in sorted(set(v1_entities) | set(v2_entities)): + v1_variables = v1_entities.get(entity_id) or {} + v2_variables = v2_entities.get(entity_id) or {} + for variable in sorted(set(v1_variables) | set(v2_variables)): + v1_value = v1_variables.get(variable) + v2_value = v2_variables.get(variable) + if variable == "members" or not ( + isinstance(v1_value, dict) and isinstance(v2_value, dict) + ): + if not _same_value(v1_value, v2_value): + differences.append( + f"{plural}.{entity_id}.{variable}: " + f"{v1_value!r} != {v2_value!r}" + ) + continue + differences.extend( + f"{plural}.{entity_id}.{variable}.{period}: " + f"{v1_value.get(period)!r} != {v2_value.get(period)!r}" + for period in sorted(set(v1_value) | set(v2_value)) + if not _same_value(v1_value.get(period), v2_value.get(period)) + ) + return differences diff --git a/tests/integration/test_live_v2_household_calculations.py b/tests/integration/test_live_v2_household_calculations.py new file mode 100644 index 000000000..26a63fc49 --- /dev/null +++ b/tests/integration/test_live_v2_household_calculations.py @@ -0,0 +1,169 @@ +"""Live staging checks of ``/v2/household-calculations``. + +These run against the tagged staging candidate, through staging API v1, the +staging Simulation Entrypoint, the v2 household functions on Modal and the +staging v2 database: + +- a submit, poll and result cycle for a new US household; +- a latency check: several calculations submitted at once, so the Simulation + Entrypoint serves them in parallel, each accepted within + ``STAGING_HOUSEHOLD_SUBMIT_LIMIT_SECONDS`` and finished within + ``STAGING_HOUSEHOLD_FINISH_LIMIT_SECONDS``; +- a memory check: a household varied across 401 points of employment income, + every variable calculated at each point, which must finish without failing. + +The measured times are reported as a warning, so the job log keeps them. +""" + +from __future__ import annotations + +import os +import statistics +import warnings +from concurrent.futures import ThreadPoolExecutor + +import httpx + +from tests.fixtures.integration.live_household_calculations import ( + YEAR, + FinishedCalculation, + LiveHousehold, + calculate, + store_household, +) + +PARALLEL_CALCULATIONS = int(os.environ.get("STAGING_HOUSEHOLD_PARALLEL_COUNT", "6")) +SUBMIT_LIMIT_SECONDS = float( + os.environ.get("STAGING_HOUSEHOLD_SUBMIT_LIMIT_SECONDS", "10") +) +FINISH_LIMIT_SECONDS = float( + os.environ.get("STAGING_HOUSEHOLD_FINISH_LIMIT_SECONDS", "300") +) +MEMORY_CHECK_POINTS = 401 + + +class HouseholdCalculationMeasurement(UserWarning): + """Times measured by the live household calculation checks.""" + + +def _us_household(name: str, employment_income: int) -> LiveHousehold: + return LiveHousehold( + name=name, + country_id="us", + people={ + "adult": { + "age": {YEAR: 40}, + "employment_income": {YEAR: employment_income}, + } + }, + groups={ + "household": { + "state_name": {YEAR: "NY"}, + "household_net_income": {YEAR: None}, + } + }, + ) + + +def _probe_amount(integration_probe_id: str) -> int: + return int(integration_probe_id.rsplit("-", 1)[-1], 16) % 997 + + +def test_a_us_household_is_submitted_polled_and_read(api_client, integration_probe_id): + household = _us_household("cycle", 51_000 + _probe_amount(integration_probe_id)) + household_id = store_household(api_client, household) + + finished = calculate(api_client, household, household_id) + + item = finished.item + assert item["status"] == "succeeded", item["error"] + assert item["household_id"] == household_id + assert item["completed_at"] is not None + result = item["result"] + net_income = result["household"]["households"]["household-1"][ + "household_net_income" + ][YEAR] + assert isinstance(net_income, int | float) + assert result["spm_config"] is not None + assert result["spm_provenance"] is not None + assert result["policyengine_version"] + read_again = api_client.get( + f"/v2/household-calculations/{item['calculation_id']}" + ).json()["result"] + assert read_again == item + + +def test_parallel_calculations_finish_within_the_latency_limits( + api_base_url, integration_probe_id +): + probe_amount = _probe_amount(integration_probe_id) + households = [ + _us_household(f"parallel-{index}", 20_000 + 1_000 * index + probe_amount) + for index in range(PARALLEL_CALCULATIONS) + ] + + def run(household: LiveHousehold) -> FinishedCalculation: + with httpx.Client(base_url=api_base_url, timeout=90.0) as client: + return calculate(client, household, store_household(client, household)) + + with ThreadPoolExecutor(max_workers=PARALLEL_CALCULATIONS) as pool: + outcomes = list(pool.map(run, households)) + + submit_times = [outcome.submit_seconds for outcome in outcomes] + finish_times = [outcome.total_seconds for outcome in outcomes] + warnings.warn( + HouseholdCalculationMeasurement( + f"{PARALLEL_CALCULATIONS} parallel household calculations: submit " + f"median {statistics.median(submit_times):.2f} s, longest " + f"{max(submit_times):.2f} s; finish median " + f"{statistics.median(finish_times):.1f} s, longest " + f"{max(finish_times):.1f} s" + ), + stacklevel=1, + ) + assert [outcome.item["status"] for outcome in outcomes] == ["succeeded"] * len( + outcomes + ), [outcome.item["error"] for outcome in outcomes] + assert max(submit_times) <= SUBMIT_LIMIT_SECONDS, submit_times + assert max(finish_times) <= FINISH_LIMIT_SECONDS, finish_times + + +def test_a_household_across_many_points_finishes(api_client, integration_probe_id): + household = LiveHousehold( + name="memory", + country_id="us", + people={"adult": {"age": {YEAR: 45}}}, + groups={ + "household": { + "state_name": {YEAR: "CA"}, + "household_net_income": {YEAR: None}, + } + }, + axes=[ + [ + { + "name": "employment_income", + "count": MEMORY_CHECK_POINTS, + "min": 0, + "max": 400_000 + _probe_amount(integration_probe_id), + "period": YEAR, + } + ] + ], + ) + household_id = store_household(api_client, household) + + finished = calculate(api_client, household, household_id) + + warnings.warn( + HouseholdCalculationMeasurement( + f"Household across {MEMORY_CHECK_POINTS} points: submit " + f"{finished.submit_seconds:.2f} s, finish {finished.total_seconds:.1f} s" + ), + stacklevel=1, + ) + assert finished.item["status"] == "succeeded", finished.item["error"] + net_income = finished.item["result"]["household"]["households"]["household-1"][ + "household_net_income" + ][YEAR] + assert len(net_income) == MEMORY_CHECK_POINTS diff --git a/tests/integration/test_live_v2_household_parity.py b/tests/integration/test_live_v2_household_parity.py new file mode 100644 index 000000000..0cebaa53e --- /dev/null +++ b/tests/integration/test_live_v2_household_parity.py @@ -0,0 +1,154 @@ +"""Live staging parity of v2 household calculations with API v1's ``/calculate``. + +Each household is stored through ``/v2/households``, calculated through +``/v2/household-calculations`` (the v2 household function on Modal) and, as +the same situation, through API v1's in-process ``/calculate`` on the same +candidate. Every calculated value must match, axes included. A run-specific +input makes each household new, so no earlier calculation is reused. +""" + +from __future__ import annotations + +import pytest + +from tests.fixtures.integration.live_household_calculations import ( + YEAR, + LiveHousehold, + calculate, + calculate_with_v1, + household_differences, + store_household, +) + +CALCULATED = {YEAR: None} + + +def _probe_amount(integration_probe_id: str) -> int: + """Return a small run-specific amount that makes each household new.""" + return int(integration_probe_id.rsplit("-", 1)[-1], 16) % 997 + + +def _parity_households(probe_amount: int) -> list[LiveHousehold]: + return [ + LiveHousehold( + name="us-single-california", + country_id="us", + people={ + "adult": { + "age": {YEAR: 35}, + "employment_income": {YEAR: 42_000 + probe_amount}, + "income_tax": CALCULATED, + } + }, + groups={ + "household": { + "state_name": {YEAR: "CA"}, + "household_net_income": CALCULATED, + }, + "tax_unit": {"eitc": CALCULATED, "ctc": CALCULATED}, + "spm_unit": {"snap": CALCULATED}, + }, + ), + LiveHousehold( + name="us-couple-two-children-texas-axes", + country_id="us", + people={ + "first-adult": {"age": {YEAR: 38}}, + "second-adult": { + "age": {YEAR: 36}, + "employment_income": {YEAR: 15_000 + probe_amount}, + }, + "first-child": {"age": {YEAR: 7}}, + "second-child": {"age": {YEAR: 3}}, + }, + groups={ + "household": { + "state_name": {YEAR: "TX"}, + "household_net_income": CALCULATED, + }, + "tax_unit": {"eitc": CALCULATED, "ctc": CALCULATED}, + "spm_unit": {"snap": CALCULATED}, + }, + axes=[ + [ + { + "name": "employment_income", + "count": 5, + "min": 0, + "max": 100_000, + "period": YEAR, + } + ] + ], + ), + LiveHousehold( + name="uk-single", + country_id="uk", + people={ + "adult": { + "age": {YEAR: 30}, + "employment_income": {YEAR: 28_000 + probe_amount}, + "income_tax": CALCULATED, + "national_insurance": CALCULATED, + } + }, + groups={"household": {"household_net_income": CALCULATED}}, + ), + LiveHousehold( + name="uk-couple-child-axes", + country_id="uk", + people={ + "first-adult": {"age": {YEAR: 34}}, + "second-adult": { + "age": {YEAR: 33}, + "employment_income": {YEAR: 9_000 + probe_amount}, + }, + "child": {"age": {YEAR: 4}}, + }, + groups={ + "household": {"household_net_income": CALCULATED}, + "benunit": { + "universal_credit": CALCULATED, + "child_benefit": CALCULATED, + }, + }, + axes=[ + [ + { + "name": "employment_income", + "count": 5, + "min": 0, + "max": 60_000, + "period": YEAR, + } + ] + ], + ), + ] + + +PARITY_NAMES = [household.name for household in _parity_households(0)] + + +@pytest.mark.parametrize("household_name", PARITY_NAMES) +def test_a_v2_household_calculation_matches_api_v1( + api_client, integration_probe_id, household_name +): + (household,) = [ + candidate + for candidate in _parity_households(_probe_amount(integration_probe_id)) + if candidate.name == household_name + ] + household_id = store_household(api_client, household) + + finished = calculate(api_client, household, household_id) + v1_answer = calculate_with_v1(api_client, household) + + assert finished.item["status"] == "succeeded", finished.item["error"] + result = finished.item["result"] + differences = household_differences(v1_answer["result"], result["household"]) + assert differences == [], "\n".join(differences[:50]) + assert result["warnings"] == v1_answer.get("warnings", []) + if household.country_id == "us": + assert result["spm_config"] is not None + assert result["spm_provenance"] is not None diff --git a/tests/unit/test_cloud_run_deploy_scripts.py b/tests/unit/test_cloud_run_deploy_scripts.py index ee7a48b9f..a9d0dc878 100644 --- a/tests/unit/test_cloud_run_deploy_scripts.py +++ b/tests/unit/test_cloud_run_deploy_scripts.py @@ -1894,6 +1894,12 @@ def _assert_live_tests_run_in_parallel_after_readiness( assert "tests/integration/test_live_budget_window_cache.py" in live_test_script assert "tests/integration/test_live_calculate.py" in live_test_script assert "tests/integration/test_live_economy.py" in live_test_script + assert "tests/integration/test_live_v2_household_calculations.py" in ( + live_test_script + ) + assert "tests/integration/test_live_v2_household_parity.py" in live_test_script + # The household calculation checks wait on Modal; 30 minutes leaves room. + assert "timeout-minutes: 30" in parallel_live_tests assert "tests/integration/test_live_v2_policies.py" not in parallel_live_tests assert ( "API_BASE_URL: ${{ needs.deploy-cloud-run-staging.outputs.url }}" From c66b1d0b0c6cc721297fbd78c9402f71b0248346 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 22:12:21 +0400 Subject: [PATCH 14/37] Read simulation service answers through the generated models A 422 from the simulation service (a household or options it finds invalid, such as a group no person belongs to) now answers 422 with each invalid field and why, instead of 503 without a reason. A body that is not JSON or not of the documented shape, and a gateway token that cannot be fetched, now answer 503 instead of 500. Each answer is decoded once, by the generated client. Co-Authored-By: Claude Opus 5.5 --- docs/migration/stage-11-v2-households.md | 9 +- .../household_calculations/response_models.py | 4 +- .../v2/household_calculations/routes.py | 8 +- .../simulation_client.py | 117 +++++++++++------- .../v2/household_calculations/types.py | 7 ++ .../v2/test_household_calculation_routes.py | 14 +++ ...household_calculation_simulation_client.py | 73 +++++++++++ 7 files changed, 185 insertions(+), 47 deletions(-) diff --git a/docs/migration/stage-11-v2-households.md b/docs/migration/stage-11-v2-households.md index 46c8730d3..7aee10c96 100644 --- a/docs/migration/stage-11-v2-households.md +++ b/docs/migration/stage-11-v2-households.md @@ -289,9 +289,12 @@ The routes always call the Cloud Run Simulation Entrypoint, whatever for every simulation call, with `GATEWAY_AUTH_REQUIRED=1`. Without the URL, or with gateway auth partly configured, both routes answer -503. They also answer 503 when the simulation service cannot be reached, -answers with an error, or refuses the calculation (for example when no v2 -worker is deployed for the catalog version's PolicyEngine.py version). +503. They also answer 503 when the simulation service cannot be reached, its +gateway token cannot be fetched, it answers with an error or with a body of +another shape, or it refuses the calculation (for example when no v2 worker is +deployed for the catalog version's PolicyEngine.py version). A household or +options the simulation service finds invalid, such as a group no person +belongs to, get 422 with each invalid field and why. ### Staging checks diff --git a/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py b/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py index b2b1b160d..d21b25968 100644 --- a/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py +++ b/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py @@ -94,7 +94,9 @@ class HouseholdCalculationResponse(StrictHouseholdCalculationAPIModel): "model": V2ErrorResponse, "description": ( "The request does not match the schema, the country does not offer " - "household calculations, or no year can be chosen." + "household calculations, the household has no default year, or the " + "simulation service finds the household or its options invalid " + "(the message lists why)." ), }, 500: { diff --git a/policyengine_api/fastapi_routes/v2/household_calculations/routes.py b/policyengine_api/fastapi_routes/v2/household_calculations/routes.py index c88ce2289..019ba1c41 100644 --- a/policyengine_api/fastapi_routes/v2/household_calculations/routes.py +++ b/policyengine_api/fastapi_routes/v2/household_calculations/routes.py @@ -33,6 +33,7 @@ HouseholdCalculationResponse, ) from policyengine_api.services.v2.household_calculations.types import ( + SimulationServiceInvalidRequestError, SimulationServiceRefusedError, SimulationServiceUnavailableError, ) @@ -56,7 +57,12 @@ HOUSEHOLD_CALCULATION_ERROR_MAPPINGS = ( V2ErrorMapping( - (HouseholdCalculationCountryError, HouseholdCalculationYearError), 422 + ( + HouseholdCalculationCountryError, + HouseholdCalculationYearError, + SimulationServiceInvalidRequestError, + ), + 422, ), V2ErrorMapping( ( diff --git a/policyengine_api/services/v2/household_calculations/simulation_client.py b/policyengine_api/services/v2/household_calculations/simulation_client.py index 1474bd462..498d23614 100644 --- a/policyengine_api/services/v2/household_calculations/simulation_client.py +++ b/policyengine_api/services/v2/household_calculations/simulation_client.py @@ -3,9 +3,11 @@ Calls go to the Cloud Run Simulation Entrypoint (``SIMULATION_ENTRYPOINT_URL``), the only simulation service with these routes, with the gateway credentials the Public API already uses (``GATEWAY_AUTH_*``) and the request's -correlation headers. A 400 is a refusal, whose reason the service gives; any -other failure, including a 503, no answer or missing configuration, means the -service is unavailable. +correlation headers. Answers are read through the client's generated models: +a 400 is a refusal, whose reason the service gives; a 422 lists what the +service finds invalid in the household or its options; anything else, +including a 503, no answer, an answer of another shape, a gateway token that +cannot be fetched or missing configuration, means the service is unavailable. One client serves the process, so its connections and its gateway token are reused across requests. @@ -14,8 +16,7 @@ from __future__ import annotations from functools import lru_cache -from http import HTTPStatus -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, TypeVar import httpx from policyengine_api_simulation_client import AuthenticatedClient @@ -23,9 +24,18 @@ get_household_calculation, submit_household_calculation, ) +from policyengine_api_simulation_client.models.household_calculation_response import ( + HouseholdCalculationResponse, +) from policyengine_api_simulation_client.models.household_calculation_submission import ( HouseholdCalculationSubmission, ) +from policyengine_api_simulation_client.models.household_route_error import ( + HouseholdRouteError, +) +from policyengine_api_simulation_client.models.http_validation_error import ( + HTTPValidationError, +) from policyengine_observability import ( # type: ignore[import-untyped] # the package has no type information instrument_httpx, ) @@ -43,23 +53,57 @@ ) from policyengine_api.observability import get_runtime from policyengine_api.services.v2.household_calculations.types import ( + SimulationServiceInvalidRequestError, SimulationServiceRefusedError, SimulationServiceUnavailableError, ) if TYPE_CHECKING: + from collections.abc import Callable from uuid import UUID + from policyengine_api_simulation_client.types import Response + +ParsedT = TypeVar("ParsedT") + REQUEST_TIMEOUT_SECONDS = 30.0 CLOUD_RUN_SIMULATION_ENTRYPOINT = "cloud_run_simulation_entrypoint" -def _answer(response: Any) -> tuple[int, object]: # noqa: ANN401 # the generated client's Response +def _call(send: Callable[[], Response[ParsedT]]) -> Response[ParsedT]: + """Send one request; raise SimulationServiceUnavailableError if it cannot be read. + + The generated routes decode the body for each status they document, so a + body that is not JSON or not of the documented shape raises while + decoding. + """ try: - body = httpx.Response(response.status_code, content=response.content).json() - except ValueError: - body = None - return int(response.status_code), body + return send() + except (httpx.HTTPError, GatewayAuthError) as error: + msg = "The simulation service could not be reached." + raise SimulationServiceUnavailableError(msg) from error + except (ValueError, KeyError, TypeError) as error: + msg = "The simulation service's answer could not be read." + raise SimulationServiceUnavailableError(msg) from error + + +def _invalid_request_reason(error: HTTPValidationError) -> str: + """Return each invalid field and why, without the request's ``body`` prefix.""" + reasons = [] + for item in error.detail if isinstance(error.detail, list) else []: + location = [str(part) for part in item.loc] + if location[:1] == ["body"]: + location = location[1:] + reasons.append(f"{'.'.join(location)}: {item.msg}") + return "The simulation service cannot calculate this household: " + ( + "; ".join(reasons) or "the request is invalid" + ) + + +def _unavailable(status_code: int) -> SimulationServiceUnavailableError: + return SimulationServiceUnavailableError( + f"The simulation service answered {status_code}." + ) def _gateway_auth() -> GatewayBearerAuth | None: @@ -74,13 +118,6 @@ def _gateway_auth() -> GatewayBearerAuth | None: return None -def _refusal_reason(body: object) -> str: - reason = body.get("detail") if isinstance(body, dict) else None - if isinstance(reason, str): - return reason - return "The simulation service refused the household calculation." - - class SimulationHouseholdClient: """Submit and read household calculations on the simulation service.""" @@ -128,38 +165,34 @@ def from_environment(cls) -> SimulationHouseholdClient: def submit(self, submission: dict[str, Any]) -> dict[str, Any]: """Submit a calculation and return the service's calculation answer.""" - try: - response = submit_household_calculation.sync_detailed( - client=self._client, - body=HouseholdCalculationSubmission.from_dict(submission), + body = HouseholdCalculationSubmission.from_dict(submission) + response = _call( + lambda: submit_household_calculation.sync_detailed( + client=self._client, body=body ) - except httpx.HTTPError as error: - msg = "The simulation service could not be reached." - raise SimulationServiceUnavailableError(msg) from error - status, body = _answer(response) - if status == HTTPStatus.BAD_REQUEST: - raise SimulationServiceRefusedError(_refusal_reason(body)) - if status != HTTPStatus.ACCEPTED or not isinstance(body, dict): - msg = f"The simulation service answered {status}." - raise SimulationServiceUnavailableError(msg) - return body + ) + answer = response.parsed + if isinstance(answer, HouseholdCalculationResponse): + return answer.to_dict() + if isinstance(answer, HTTPValidationError): + raise SimulationServiceInvalidRequestError(_invalid_request_reason(answer)) + if isinstance(answer, HouseholdRouteError) and response.status_code == 400: # noqa: PLR2004 # the service's refusal status + raise SimulationServiceRefusedError(answer.detail) + raise _unavailable(response.status_code) def read(self, calculation_id: UUID) -> dict[str, Any] | None: """Return the service's calculation answer, or None when it has none.""" - try: - response = get_household_calculation.sync_detailed( + response = _call( + lambda: get_household_calculation.sync_detailed( calculation_id, client=self._client ) - except httpx.HTTPError as error: - msg = "The simulation service could not be reached." - raise SimulationServiceUnavailableError(msg) from error - status, body = _answer(response) - if status == HTTPStatus.NOT_FOUND: + ) + answer = response.parsed + if isinstance(answer, HouseholdCalculationResponse): + return answer.to_dict() + if isinstance(answer, HouseholdRouteError) and response.status_code == 404: # noqa: PLR2004 # the service's unknown-calculation status return None - if status != HTTPStatus.OK or not isinstance(body, dict): - msg = f"The simulation service answered {status}." - raise SimulationServiceUnavailableError(msg) - return body + raise _unavailable(response.status_code) @lru_cache(maxsize=1) diff --git a/policyengine_api/services/v2/household_calculations/types.py b/policyengine_api/services/v2/household_calculations/types.py index 3cc20ceba..b94a61323 100644 --- a/policyengine_api/services/v2/household_calculations/types.py +++ b/policyengine_api/services/v2/household_calculations/types.py @@ -157,6 +157,13 @@ class SimulationServiceRefusedError(RuntimeError): """The simulation service refused the calculation; the message says why.""" +class SimulationServiceInvalidRequestError(ValueError): + """The simulation service finds the household or its options invalid. + + The message lists each invalid field and why, as the service reports it. + """ + + class CountrySource(Protocol): """The countries table (``V2CountryService``).""" diff --git a/tests/unit/v2/test_household_calculation_routes.py b/tests/unit/v2/test_household_calculation_routes.py index a6f396bd7..5cef47c7e 100644 --- a/tests/unit/v2/test_household_calculation_routes.py +++ b/tests/unit/v2/test_household_calculation_routes.py @@ -19,6 +19,7 @@ V2HouseholdCalculationService, ) from policyengine_api.services.v2.household_calculations.types import ( + SimulationServiceInvalidRequestError, SimulationServiceRefusedError, SimulationServiceUnavailableError, ) @@ -222,3 +223,16 @@ def test_a_malformed_request_is_refused_before_anything_is_read(changes) -> None assert response.status_code == 422 assert app.simulation.submissions == [] + + +def test_a_household_the_simulation_service_finds_invalid_gets_422_with_why() -> None: + app = App() + app.simulation.error = SimulationServiceInvalidRequestError( + "The simulation service cannot calculate this household: " + "household.households.household-1.members: List should have at least 1 item" + ) + + response = app.client.post(PATH, json=body()) + + assert response.status_code == 422 + assert "members" in response.json()["message"] diff --git a/tests/unit/v2/test_household_calculation_simulation_client.py b/tests/unit/v2/test_household_calculation_simulation_client.py index 2f73b5a82..e99eb96ca 100644 --- a/tests/unit/v2/test_household_calculation_simulation_client.py +++ b/tests/unit/v2/test_household_calculation_simulation_client.py @@ -11,10 +11,12 @@ import httpx import pytest +from policyengine_api.libs.gateway_auth import GatewayAuthError from policyengine_api.services.v2.household_calculations.simulation_client import ( SimulationHouseholdClient, ) from policyengine_api.services.v2.household_calculations.types import ( + SimulationServiceInvalidRequestError, SimulationServiceRefusedError, SimulationServiceUnavailableError, ) @@ -164,3 +166,74 @@ def test_the_configured_service_is_called_without_auth_when_none_is_set( client = SimulationHouseholdClient.from_environment() assert client.base_url == BASE_URL + + +def test_a_household_the_service_finds_invalid_is_reported_with_its_reasons() -> None: + client = client_answering( + 422, + { + "detail": [ + { + "loc": ["body", "household", "households", "h", "members"], + "msg": "List should have at least 1 item after validation", + "type": "too_short", + } + ] + }, + [], + ) + + with pytest.raises( + SimulationServiceInvalidRequestError, + match=r"household\.households\.h\.members: List should have at least 1 item", + ): + client.submit(SUBMISSION) + + +@pytest.mark.parametrize("status_code", [400, 503, 202, 200]) +def test_an_answer_that_is_not_json_is_unavailable(status_code: int) -> None: + def answer(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code, + content=b"Service Unavailable", + headers={"Content-Type": "text/html"}, + ) + + client = SimulationHouseholdClient( + BASE_URL, httpx_args={"transport": httpx.MockTransport(answer)} + ) + + with pytest.raises(SimulationServiceUnavailableError): + client.submit(SUBMISSION) + with pytest.raises(SimulationServiceUnavailableError): + client.read(CALCULATION_ID) + + +def test_an_answer_of_the_wrong_shape_is_unavailable() -> None: + client = client_answering(202, {"calculation_id": "not a calculation"}, []) + + with pytest.raises(SimulationServiceUnavailableError): + client.submit(SUBMISSION) + + +def test_a_gateway_token_that_cannot_be_fetched_is_unavailable() -> None: + class FailingTokenAuth(httpx.Auth): + def auth_flow(self, request: httpx.Request): # httpx's generator protocol + message = "The gateway token request failed" + raise GatewayAuthError(message) + yield request + + client = SimulationHouseholdClient( + BASE_URL, + httpx_args={ + "auth": FailingTokenAuth(), + "transport": httpx.MockTransport( + lambda request: httpx.Response(202, json=simulation_calculation()) + ), + }, + ) + + with pytest.raises(SimulationServiceUnavailableError): + client.submit(SUBMISSION) + with pytest.raises(SimulationServiceUnavailableError): + client.read(CALCULATION_ID) From 0f9bf9fbd6eb2be0cc041ca883c10b8998583a9f Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 22:15:32 +0400 Subject: [PATCH 15/37] Refuse axes the simulation service would refuse, before sending them The axes of one group must share one point count, and the groups' point counts may multiply to at most 10,000, as the simulation service requires. Such requests now get 422 from the Public API before anything is read or sent. Co-Authored-By: Claude Opus 5.5 --- .../v2/household_calculations/types.py | 27 +++++++++++-- .../v2/test_household_calculation_routes.py | 38 ++++++++++++++++++- 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/policyengine_api/services/v2/household_calculations/types.py b/policyengine_api/services/v2/household_calculations/types.py index b94a61323..11712acfa 100644 --- a/policyengine_api/services/v2/household_calculations/types.py +++ b/policyengine_api/services/v2/household_calculations/types.py @@ -2,11 +2,12 @@ from __future__ import annotations +import math from datetime import datetime from typing import TYPE_CHECKING, Annotated, Any, Literal, Protocol, Self from uuid import UUID -from pydantic import BaseModel, ConfigDict, Field, JsonValue +from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator from policyengine_api.query_parameters import CountryId, DefaultYear @@ -20,7 +21,8 @@ ) from policyengine_api.services.v2.policies.types import PolicyRead -# The simulation service's limits on axes. +# The simulation service's limits on axes: the product of the groups' point +# counts is at most MAXIMUM_AXIS_POINTS. MAXIMUM_AXIS_POINTS = 10_000 MAXIMUM_AXIS_GROUPS = 8 MAXIMUM_PARALLEL_AXES = 32 @@ -59,7 +61,9 @@ class HouseholdCalculationInput(StrictCalculationModel): """A stored household to calculate under a stored policy, or current law. ``year`` defaults to the household's default year. ``axes`` vary inputs - across points; ``add_missing_variables`` adds every yearly, monthly and + across points: the axes of one group vary together, so they share one + point count, and the groups form a grid of at most 10,000 points. + ``add_missing_variables`` adds every yearly, monthly and eternal variable the household lacks, as ``/calculate-full`` does. ``force`` calculates again even when a recent calculation can be reused. """ @@ -74,6 +78,23 @@ class HouseholdCalculationInput(StrictCalculationModel): add_missing_variables: bool = False force: bool = False + @field_validator("axes") + @classmethod + def limit_axis_points( + cls, axes: list[list[HouseholdCalculationAxis]] | None + ) -> list[list[HouseholdCalculationAxis]] | None: + """Refuse groups of uneven point counts and grids of too many points.""" + if axes is None: + return None + for group in axes: + if len({axis.count for axis in group}) != 1: + msg = "the axes of one group must share one point count" + raise ValueError(msg) + if math.prod(group[0].count for group in axes) > MAXIMUM_AXIS_POINTS: + msg = f"the axis groups may vary at most {MAXIMUM_AXIS_POINTS} points" + raise ValueError(msg) + return axes + class HouseholdCalculationResult(StrictCalculationModel): """A succeeded calculation's household and its notes. diff --git a/tests/unit/v2/test_household_calculation_routes.py b/tests/unit/v2/test_household_calculation_routes.py index 5cef47c7e..836de39bd 100644 --- a/tests/unit/v2/test_household_calculation_routes.py +++ b/tests/unit/v2/test_household_calculation_routes.py @@ -210,11 +210,32 @@ def test_a_simulation_service_that_cannot_calculate_answers_503( "changes", [ {"axes": [[{"name": "employment_income", "count": 0, "min": 0, "max": 1}]]}, + { + "axes": [ + [ + {"name": "employment_income", "count": 3, "min": 0, "max": 1}, + {"name": "age", "count": 4, "min": 20, "max": 60}, + ] + ] + }, + { + "axes": [ + [{"name": "employment_income", "count": 101, "min": 0, "max": 1}], + [{"name": "age", "count": 100, "min": 20, "max": 60}], + ] + }, {"country_id": "ca"}, {"unexpected": True}, {"year": 1800}, ], - ids=["axis-count", "country", "unknown-field", "year"], + ids=[ + "axis-count", + "parallel-axes-of-different-counts", + "too-many-axis-points", + "country", + "unknown-field", + "year", + ], ) def test_a_malformed_request_is_refused_before_anything_is_read(changes) -> None: app = App() @@ -236,3 +257,18 @@ def test_a_household_the_simulation_service_finds_invalid_gets_422_with_why() -> assert response.status_code == 422 assert "members" in response.json()["message"] + + +def test_axis_groups_within_the_point_limit_are_accepted() -> None: + app = App() + axes = [ + [ + {"name": "employment_income", "count": 100, "min": 0, "max": 1}, + {"name": "self_employment_income", "count": 100, "min": 0, "max": 1}, + ], + [{"name": "age", "count": 100, "min": 20, "max": 60}], + ] + + response = app.client.post(PATH, json=body(axes=axes)) + + assert response.status_code == 202 From 0c9b65c63f44163dd32db58bdee5c338988638a5 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 22:18:40 +0400 Subject: [PATCH 16/37] Calculate each household for its default year, and take no force The request no longer takes `year`: the calculation's year is the household's default year, the year its values name, so a calculation is never labelled with a year it did not calculate. It no longer takes `force` either, so callers without an account cannot skip the 28-day reuse of a succeeded calculation and start a new Modal call for each request. The stage 11 guide notes the year the household function uses when it adds missing variables. Co-Authored-By: Claude Opus 5.5 --- docs/migration/stage-11-v2-households.md | 13 +++++++++++-- .../household_calculations/request_models.py | 10 +++++----- .../v2/household_calculations/services.py | 6 ++++-- .../v2/household_calculations/types.py | 8 +++----- .../v2/household_calculations/validators.py | 10 ++++------ .../live_household_calculations.py | 4 ---- .../v2/test_household_calculation_routes.py | 4 +++- .../v2/test_household_calculation_service.py | 19 ++++++++++++++----- 8 files changed, 44 insertions(+), 30 deletions(-) diff --git a/docs/migration/stage-11-v2-households.md b/docs/migration/stage-11-v2-households.md index 7aee10c96..c43665420 100644 --- a/docs/migration/stage-11-v2-households.md +++ b/docs/migration/stage-11-v2-households.md @@ -261,8 +261,10 @@ for immediate rollback. `POST /v2/household-calculations` calculates a household stored through `/v2/households`, under a policy stored through `/v2/policies` or, without -`policy_id`, under current law. The body may add `year` (the household's -`default_year` otherwise), `axes`, `add_missing_variables` and `force`. The +`policy_id`, under current law, for the household's `default_year`. The body +may add `axes` and `add_missing_variables`. An identical calculation that +succeeded in the last 28 days is returned again; callers cannot force a new +one. The Public API checks the country's `countries` row allows household calculations (422 otherwise; a household-only country is accepted), reads the household and policy through the same services as their routes (404 when either is not @@ -274,6 +276,13 @@ tables. The answer is 202 with the calculation's `calculation_id` and `status` is `succeeded`, with `result`, or `failed`, with `error`. A US result carries `spm_config` and `spm_provenance`. +With `add_missing_variables`, the v2 household function adds each missing +variable for the year it would choose for `/calculate-full`: the first period +of the age of a person named `you`, else the current calendar year (UTC). The +Public API names people by their record IDs, so for a household whose +`default_year` is not the current year the added variables are requested for +the current year until the household function takes the calculation's year. + The simulation service stores every calculation as a simulation run, so a calculation no account saved is still stored, with no report. Neither route needs an account. diff --git a/policyengine_api/fastapi_routes/v2/household_calculations/request_models.py b/policyengine_api/fastapi_routes/v2/household_calculations/request_models.py index c8b0e4427..eb78c9236 100644 --- a/policyengine_api/fastapi_routes/v2/household_calculations/request_models.py +++ b/policyengine_api/fastapi_routes/v2/household_calculations/request_models.py @@ -12,12 +12,12 @@ class HouseholdCalculationRequest(HouseholdCalculationInput): ``household_id`` and ``policy_id`` name records of ``country_id`` stored through ``/v2/households`` and ``/v2/policies``; without ``policy_id`` the - household is calculated under current law. ``year`` defaults to the - household's ``default_year``. Each group in ``axes`` varies its inputs - together across ``count`` points; several groups form a grid. + household is calculated under current law, for the household's + ``default_year``. Each group in ``axes`` varies its inputs together across + ``count`` points; several groups form a grid of at most 10,000 points. ``add_missing_variables`` adds every variable the household lacks before - calculating, and ``force`` calculates again even when a recent identical - calculation exists. + calculating. An identical calculation that succeeded in the last 28 days + is returned again instead of being repeated. """ def to_service_input(self) -> HouseholdCalculationInput: diff --git a/policyengine_api/services/v2/household_calculations/services.py b/policyengine_api/services/v2/household_calculations/services.py index 3eafaeb71..728ada8cc 100644 --- a/policyengine_api/services/v2/household_calculations/services.py +++ b/policyengine_api/services/v2/household_calculations/services.py @@ -83,7 +83,7 @@ def submit_calculation( country_id=country_id, policy_id=calculation.policy_id ) ) - year = calculation_year(calculation.year, household.default_year) + year = calculation_year(household.default_year) try: catalog = self._catalogs.get_model_by_country(country_id) finally: @@ -117,7 +117,9 @@ def submit_calculation( ), "add_missing_variables": calculation.add_missing_variables, }, - "force": calculation.force, + # A calculation is reused for 28 days after it succeeded; callers + # without an account must not start new Modal calls at will. + "force": False, "trigger": "initial", } spm = saved_spm_choice(household.household_data) diff --git a/policyengine_api/services/v2/household_calculations/types.py b/policyengine_api/services/v2/household_calculations/types.py index 11712acfa..78d50f8ed 100644 --- a/policyengine_api/services/v2/household_calculations/types.py +++ b/policyengine_api/services/v2/household_calculations/types.py @@ -9,7 +9,7 @@ from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator -from policyengine_api.query_parameters import CountryId, DefaultYear +from policyengine_api.query_parameters import CountryId if TYPE_CHECKING: from collections.abc import Mapping @@ -60,23 +60,21 @@ class HouseholdCalculationAxis(StrictCalculationModel): class HouseholdCalculationInput(StrictCalculationModel): """A stored household to calculate under a stored policy, or current law. - ``year`` defaults to the household's default year. ``axes`` vary inputs + The calculation's year is the household's default year, the year its + values name. ``axes`` vary inputs across points: the axes of one group vary together, so they share one point count, and the groups form a grid of at most 10,000 points. ``add_missing_variables`` adds every yearly, monthly and eternal variable the household lacks, as ``/calculate-full`` does. - ``force`` calculates again even when a recent calculation can be reused. """ country_id: CountryId household_id: UUID policy_id: UUID | None = None - year: DefaultYear | None = None axes: Annotated[list[AxisGroup], Field(max_length=MAXIMUM_AXIS_GROUPS)] | None = ( None ) add_missing_variables: bool = False - force: bool = False @field_validator("axes") @classmethod diff --git a/policyengine_api/services/v2/household_calculations/validators.py b/policyengine_api/services/v2/household_calculations/validators.py index 507590131..5475e3328 100644 --- a/policyengine_api/services/v2/household_calculations/validators.py +++ b/policyengine_api/services/v2/household_calculations/validators.py @@ -15,7 +15,7 @@ class HouseholdCalculationCountryError(ValueError): class HouseholdCalculationYearError(ValueError): - """The request names no year and the household has no default year.""" + """The household has no default year to calculate.""" class HouseholdCalculationNotFoundError(LookupError): @@ -37,11 +37,9 @@ def require_household_country( raise HouseholdCalculationCountryError(msg) -def calculation_year(requested_year: int | None, default_year: int | None) -> int: - """Return the requested year, or else the household's default year.""" - if requested_year is not None: - return requested_year +def calculation_year(default_year: int | None) -> int: + """Return the household's default year, the year its values name.""" if default_year is None: - msg = "This household has no default year; send the year to calculate." + msg = "This household has no default year to calculate." raise HouseholdCalculationYearError(msg) return default_year diff --git a/tests/fixtures/integration/live_household_calculations.py b/tests/fixtures/integration/live_household_calculations.py index 6fbf92647..cc59fbe15 100644 --- a/tests/fixtures/integration/live_household_calculations.py +++ b/tests/fixtures/integration/live_household_calculations.py @@ -116,16 +116,12 @@ def calculate( client: httpx.Client, household: LiveHousehold, household_id: str, - *, - force: bool = False, ) -> FinishedCalculation: """Submit the stored household's calculation and poll until it finishes.""" started = time.monotonic() body: dict[str, Any] = { "country_id": household.country_id, "household_id": household_id, - "year": int(YEAR), - "force": force, } if household.axes is not None: body["axes"] = household.axes diff --git a/tests/unit/v2/test_household_calculation_routes.py b/tests/unit/v2/test_household_calculation_routes.py index 836de39bd..6ab92cfc6 100644 --- a/tests/unit/v2/test_household_calculation_routes.py +++ b/tests/unit/v2/test_household_calculation_routes.py @@ -226,7 +226,8 @@ def test_a_simulation_service_that_cannot_calculate_answers_503( }, {"country_id": "ca"}, {"unexpected": True}, - {"year": 1800}, + {"year": 2026}, + {"force": True}, ], ids=[ "axis-count", @@ -235,6 +236,7 @@ def test_a_simulation_service_that_cannot_calculate_answers_503( "country", "unknown-field", "year", + "force", ], ) def test_a_malformed_request_is_refused_before_anything_is_read(changes) -> None: diff --git a/tests/unit/v2/test_household_calculation_service.py b/tests/unit/v2/test_household_calculation_service.py index b2b16bbc1..30242c867 100644 --- a/tests/unit/v2/test_household_calculation_service.py +++ b/tests/unit/v2/test_household_calculation_service.py @@ -78,7 +78,7 @@ def test_a_calculation_sends_the_stored_household_policy_and_catalog() -> None: parts = Parts() calculation = parts.service.submit_calculation( - calculation_input(axes=AXES, add_missing_variables=True, force=True) + calculation_input(axes=AXES, add_missing_variables=True) ) (submission,) = parts.simulation.submissions @@ -99,7 +99,7 @@ def test_a_calculation_sends_the_stored_household_policy_and_catalog() -> None: }, }, "options": {"axes": AXES, "add_missing_variables": True}, - "force": True, + "force": False, "trigger": "initial", } assert calculation.calculation_id == CALCULATION_ID @@ -124,17 +124,26 @@ def test_current_law_and_no_saved_choice_send_neither() -> None: assert submission["options"] == {"axes": None, "add_missing_variables": False} -def test_the_year_is_the_requested_one_or_the_households_default() -> None: +def test_the_year_is_the_households_default_year() -> None: parts = Parts() + parts.households.households[("us", HOUSEHOLD_ID)] = household_read( + default_year=2027 + ) + + parts.service.submit_calculation(calculation_input()) + + assert parts.simulation.submissions[0]["year"] == 2027 - parts.service.submit_calculation(calculation_input(year=2027)) + +def test_a_household_without_a_default_year_is_refused() -> None: + parts = Parts() parts.households.households[("us", HOUSEHOLD_ID)] = household_read( default_year=None ) with pytest.raises(HouseholdCalculationYearError): parts.service.submit_calculation(calculation_input()) - assert [submission["year"] for submission in parts.simulation.submissions] == [2027] + assert parts.simulation.submissions == [] def test_a_country_without_household_calculations_is_refused_by_name() -> None: From aaf0df33a4ffca6e972f2398554d59054407784e Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 22:21:59 +0400 Subject: [PATCH 17/37] Open the catalog reader only to submit, and reuse the household collections A household calculation's catalog reader, with its database session, is now opened only once a submission needs the catalog version, and closed after it; polls and refused submissions open none. The situation's group collections come from the households service's own lists, and the live checks share one helper for their run-specific amounts. Co-Authored-By: Claude Opus 5.5 --- .../fastapi_routes/dependencies.py | 2 +- .../v2/household_calculations/services.py | 16 +++++++++---- .../household_calculations/transformations.py | 22 ++++++++++------- .../live_household_calculations.py | 5 ++++ .../test_live_v2_household_calculations.py | 13 ++++------ .../test_live_v2_household_parity.py | 18 ++++++-------- .../v2/test_household_calculation_routes.py | 2 +- .../v2/test_household_calculation_service.py | 24 ++++++++++++++++--- 8 files changed, 64 insertions(+), 38 deletions(-) diff --git a/policyengine_api/fastapi_routes/dependencies.py b/policyengine_api/fastapi_routes/dependencies.py index 5d0bff862..9521876ae 100644 --- a/policyengine_api/fastapi_routes/dependencies.py +++ b/policyengine_api/fastapi_routes/dependencies.py @@ -898,7 +898,7 @@ def _default_v2_household_calculation_service_factory() -> ( countries=_default_v2_country_service_factory(), households=_default_v2_household_service_factory(), policies=_default_v2_policy_service_factory(), - catalogs=_default_v2_metadata_reader_factory(), + catalog_reader_factory=_default_v2_metadata_reader_factory, simulation=process_simulation_household_client(), ) diff --git a/policyengine_api/services/v2/household_calculations/services.py b/policyengine_api/services/v2/household_calculations/services.py index 728ada8cc..4c5d7fca8 100644 --- a/policyengine_api/services/v2/household_calculations/services.py +++ b/policyengine_api/services/v2/household_calculations/services.py @@ -31,6 +31,7 @@ ) if TYPE_CHECKING: + from collections.abc import Callable from uuid import UUID from policyengine_api.services.v2.household_calculations.types import ( @@ -52,14 +53,18 @@ def __init__( countries: CountrySource, households: HouseholdSource, policies: PolicySource, - catalogs: CatalogSource, + catalog_reader_factory: Callable[[], CatalogSource], simulation: SimulationHouseholdCalculations, ) -> None: - """Read through these services for one request.""" + """Read through these services for one request. + + ``catalog_reader_factory`` opens a catalog reader, with its own + database session, only once a submission needs the catalog version. + """ self._countries = countries self._households = households self._policies = policies - self._catalogs = catalogs + self._catalog_reader_factory = catalog_reader_factory self._simulation = simulation def submit_calculation( @@ -84,10 +89,11 @@ def submit_calculation( ) ) year = calculation_year(household.default_year) + catalog_reader = self._catalog_reader_factory() try: - catalog = self._catalogs.get_model_by_country(country_id) + catalog = catalog_reader.get_model_by_country(country_id) finally: - self._catalogs.close() + catalog_reader.close() submission: dict[str, Any] = { "country": country_id, "catalog_version_id": str(catalog.model_version.id), diff --git a/policyengine_api/services/v2/household_calculations/transformations.py b/policyengine_api/services/v2/household_calculations/transformations.py index 53911db0f..51de6eb10 100644 --- a/policyengine_api/services/v2/household_calculations/transformations.py +++ b/policyengine_api/services/v2/household_calculations/transformations.py @@ -18,21 +18,25 @@ from typing import TYPE_CHECKING, Any +from policyengine_api.services.v2.households.transformations import ( + LEGACY_COLLECTION_NAMES, +) +from policyengine_api.services.v2.households.validators import ( + SUPPORTED_ENTITY_COLLECTIONS, +) + if TYPE_CHECKING: from collections.abc import Mapping from policyengine_api.services.v2.policies.types import PolicyRead -# Each country's stored group collections and their plural in the situation. +# Each country's stored group collections, as the households service accepts +# them, and their plural in the situation. GROUP_COLLECTIONS: dict[str, dict[str, str]] = { - "us": { - "household": "households", - "family": "families", - "tax_unit": "tax_units", - "spm_unit": "spm_units", - "marital_unit": "marital_units", - }, - "uk": {"household": "households", "benunit": "benunits"}, + country_id: { + collection: LEGACY_COLLECTION_NAMES[collection] for collection in collections + } + for country_id, collections in SUPPORTED_ENTITY_COLLECTIONS.items() } OPEN_END_DATE = "2100-12-31" diff --git a/tests/fixtures/integration/live_household_calculations.py b/tests/fixtures/integration/live_household_calculations.py index cc59fbe15..b3dc8f936 100644 --- a/tests/fixtures/integration/live_household_calculations.py +++ b/tests/fixtures/integration/live_household_calculations.py @@ -97,6 +97,11 @@ class FinishedCalculation: total_seconds: float +def probe_amount(integration_probe_id: str) -> int: + """Return a small run-specific amount that makes each household new.""" + return int(integration_probe_id.rsplit("-", 1)[-1], 16) % 997 + + def store_household(client: httpx.Client, household: LiveHousehold) -> str: """Store the household through ``/v2/households`` and return its ID.""" response = client.post( diff --git a/tests/integration/test_live_v2_household_calculations.py b/tests/integration/test_live_v2_household_calculations.py index 26a63fc49..278575621 100644 --- a/tests/integration/test_live_v2_household_calculations.py +++ b/tests/integration/test_live_v2_household_calculations.py @@ -29,6 +29,7 @@ FinishedCalculation, LiveHousehold, calculate, + probe_amount, store_household, ) @@ -65,12 +66,8 @@ def _us_household(name: str, employment_income: int) -> LiveHousehold: ) -def _probe_amount(integration_probe_id: str) -> int: - return int(integration_probe_id.rsplit("-", 1)[-1], 16) % 997 - - def test_a_us_household_is_submitted_polled_and_read(api_client, integration_probe_id): - household = _us_household("cycle", 51_000 + _probe_amount(integration_probe_id)) + household = _us_household("cycle", 51_000 + probe_amount(integration_probe_id)) household_id = store_household(api_client, household) finished = calculate(api_client, household, household_id) @@ -96,9 +93,9 @@ def test_a_us_household_is_submitted_polled_and_read(api_client, integration_pro def test_parallel_calculations_finish_within_the_latency_limits( api_base_url, integration_probe_id ): - probe_amount = _probe_amount(integration_probe_id) + run_amount = probe_amount(integration_probe_id) households = [ - _us_household(f"parallel-{index}", 20_000 + 1_000 * index + probe_amount) + _us_household(f"parallel-{index}", 20_000 + 1_000 * index + run_amount) for index in range(PARALLEL_CALCULATIONS) ] @@ -145,7 +142,7 @@ def test_a_household_across_many_points_finishes(api_client, integration_probe_i "name": "employment_income", "count": MEMORY_CHECK_POINTS, "min": 0, - "max": 400_000 + _probe_amount(integration_probe_id), + "max": 400_000 + probe_amount(integration_probe_id), "period": YEAR, } ] diff --git a/tests/integration/test_live_v2_household_parity.py b/tests/integration/test_live_v2_household_parity.py index 0cebaa53e..f310add83 100644 --- a/tests/integration/test_live_v2_household_parity.py +++ b/tests/integration/test_live_v2_household_parity.py @@ -17,18 +17,14 @@ calculate, calculate_with_v1, household_differences, + probe_amount, store_household, ) CALCULATED = {YEAR: None} -def _probe_amount(integration_probe_id: str) -> int: - """Return a small run-specific amount that makes each household new.""" - return int(integration_probe_id.rsplit("-", 1)[-1], 16) % 997 - - -def _parity_households(probe_amount: int) -> list[LiveHousehold]: +def _parity_households(run_amount: int) -> list[LiveHousehold]: return [ LiveHousehold( name="us-single-california", @@ -36,7 +32,7 @@ def _parity_households(probe_amount: int) -> list[LiveHousehold]: people={ "adult": { "age": {YEAR: 35}, - "employment_income": {YEAR: 42_000 + probe_amount}, + "employment_income": {YEAR: 42_000 + run_amount}, "income_tax": CALCULATED, } }, @@ -56,7 +52,7 @@ def _parity_households(probe_amount: int) -> list[LiveHousehold]: "first-adult": {"age": {YEAR: 38}}, "second-adult": { "age": {YEAR: 36}, - "employment_income": {YEAR: 15_000 + probe_amount}, + "employment_income": {YEAR: 15_000 + run_amount}, }, "first-child": {"age": {YEAR: 7}}, "second-child": {"age": {YEAR: 3}}, @@ -87,7 +83,7 @@ def _parity_households(probe_amount: int) -> list[LiveHousehold]: people={ "adult": { "age": {YEAR: 30}, - "employment_income": {YEAR: 28_000 + probe_amount}, + "employment_income": {YEAR: 28_000 + run_amount}, "income_tax": CALCULATED, "national_insurance": CALCULATED, } @@ -101,7 +97,7 @@ def _parity_households(probe_amount: int) -> list[LiveHousehold]: "first-adult": {"age": {YEAR: 34}}, "second-adult": { "age": {YEAR: 33}, - "employment_income": {YEAR: 9_000 + probe_amount}, + "employment_income": {YEAR: 9_000 + run_amount}, }, "child": {"age": {YEAR: 4}}, }, @@ -136,7 +132,7 @@ def test_a_v2_household_calculation_matches_api_v1( ): (household,) = [ candidate - for candidate in _parity_households(_probe_amount(integration_probe_id)) + for candidate in _parity_households(probe_amount(integration_probe_id)) if candidate.name == household_name ] household_id = store_household(api_client, household) diff --git a/tests/unit/v2/test_household_calculation_routes.py b/tests/unit/v2/test_household_calculation_routes.py index 6ab92cfc6..e45e6cd62 100644 --- a/tests/unit/v2/test_household_calculation_routes.py +++ b/tests/unit/v2/test_household_calculation_routes.py @@ -51,7 +51,7 @@ def __init__(self, countries: FakeCountries | None = None) -> None: countries=countries or FakeCountries(), households=self.households, policies=FakePolicies(), - catalogs=FakeCatalogs(), + catalog_reader_factory=FakeCatalogs, simulation=self.simulation, ) dependencies = NativeRouteDependencies( diff --git a/tests/unit/v2/test_household_calculation_service.py b/tests/unit/v2/test_household_calculation_service.py index 30242c867..135ad18ae 100644 --- a/tests/unit/v2/test_household_calculation_service.py +++ b/tests/unit/v2/test_household_calculation_service.py @@ -52,16 +52,22 @@ def __init__(self, countries: FakeCountries | None = None) -> None: self.countries = countries or FakeCountries() self.households = FakeHouseholds() self.policies = FakePolicies() - self.catalogs = FakeCatalogs() + self.catalog_readers: list[FakeCatalogs] = [] self.simulation = FakeSimulation() self.service = V2HouseholdCalculationService( countries=self.countries, households=self.households, policies=self.policies, - catalogs=self.catalogs, + catalog_reader_factory=self.open_catalog_reader, simulation=self.simulation, ) + def open_catalog_reader(self) -> FakeCatalogs: + """Open and keep a new catalog reader.""" + reader = FakeCatalogs() + self.catalog_readers.append(reader) + return reader + def calculation_input(**changes: object) -> HouseholdCalculationInput: """Return a calculation of the stored US household under the stored reform.""" @@ -105,7 +111,8 @@ def test_a_calculation_sends_the_stored_household_policy_and_catalog() -> None: assert calculation.calculation_id == CALCULATION_ID assert calculation.status == "running" assert calculation.result is None - assert parts.catalogs.closed + (catalog_reader,) = parts.catalog_readers + assert catalog_reader.closed def test_current_law_and_no_saved_choice_send_neither() -> None: @@ -253,3 +260,14 @@ def test_an_unreachable_simulation_service_is_reported() -> None: with pytest.raises(SimulationServiceUnavailableError): parts.service.submit_calculation(calculation_input()) + + +def test_a_refused_submission_and_a_read_open_no_catalog_reader() -> None: + parts = Parts() + parts.simulation.read_answers[CALCULATION_ID] = succeeded_calculation() + + with pytest.raises(HouseholdNotFoundError): + parts.service.submit_calculation(calculation_input(household_id=uuid4())) + parts.service.get_calculation(CALCULATION_ID) + + assert parts.catalog_readers == [] From d9101539cb97907decd5d4fc7969da13769cb790 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 22:25:42 +0400 Subject: [PATCH 18/37] Leave SPM availability unrecorded until the service serves the version Catalog publication recorded {"available": false} for a US version the simulation service's /versions did not serve yet, and a recorded value is never replaced, so the version could never offer SPM settings. Such a version now stays unrecorded and a later publication run records it. The stage 9 runbook lists SIMULATION_ENTRYPOINT_URL for the publication command, and the catalog cache note says how long an earlier answer may be reused. Co-Authored-By: Claude Opus 5.5 --- docs/migration/stage-9-v2-metadata.md | 17 ++-- .../data/v2/catalog/spm_availability.py | 46 +++++++--- .../fastapi_routes/v2/metadata/common.py | 5 +- .../unit/v2/test_catalog_spm_availability.py | 89 +++++++++++++++++-- 4 files changed, 131 insertions(+), 26 deletions(-) diff --git a/docs/migration/stage-9-v2-metadata.md b/docs/migration/stage-9-v2-metadata.md index d02cbd0b4..ac447d1b1 100644 --- a/docs/migration/stage-9-v2-metadata.md +++ b/docs/migration/stage-9-v2-metadata.md @@ -42,8 +42,10 @@ version offers SPM settings for a household: `available`, and when it does, `defaults` (the certified selection). Catalog publication records it from the simulation service's public `GET /versions`: a US version offers SPM when the service lists a canonical capability with pinned defaults for its -PolicyEngine.py version, and other countries never do. A version whose -availability is not recorded yet offers none. The result names the +PolicyEngine.py version, and other countries never do. A US version the +service does not serve yet stays unrecorded until a later publication run, and +a version whose availability is not recorded yet offers none; a response read +before then, and reused for up to a day, keeps saying so. The result names the selected PolicyEngine.py version, so a client that starts without a version learns the running one and requests every other part for that exact version. The query uses the shared `CatalogQuery` model: unknown and repeated @@ -192,10 +194,15 @@ must execute in this order: Alembic chain, confirms all heads are current, and detects ungenerated model changes. 3. Remove the migration URL from the command environment. -4. Supply only `V2_DATA_WRITE_DATABASE_URL`, the same target identity and +4. Supply only `V2_DATA_WRITE_DATABASE_URL`, the same target identity, `HUGGING_FACE_TOKEN` (from the organization secret - `PE_UK_PRIVATE_HF_READ_TOKEN`), then run - `uv run python scripts/publish_v2_metadata_catalog.py`. Before publishing a + `PE_UK_PRIVATE_HF_READ_TOKEN`) and `SIMULATION_ENTRYPOINT_URL` (from the + environment secret of the same name), then run + `uv run python scripts/publish_v2_metadata_catalog.py`. After publishing, the + command reads the simulation service's public `GET /versions` to record each + version's SPM availability, and fails if the URL is unset or the answer + cannot be read; the catalog itself is already published by then, and the + next run records the availability. Before publishing a catalog version that is not stored yet, the command reads the certified data artifacts its regions' filters use and stops if one cannot select a region. It stops before reading anything if a UK version needs checking and diff --git a/policyengine_api/data/v2/catalog/spm_availability.py b/policyengine_api/data/v2/catalog/spm_availability.py index 673759971..2b5f28a2b 100644 --- a/policyengine_api/data/v2/catalog/spm_availability.py +++ b/policyengine_api/data/v2/catalog/spm_availability.py @@ -9,7 +9,11 @@ Publication records it after the catalog itself, on each version that has none yet, so versions published before it existed get it on the next run. A -bundle pins its SPM release, so a recorded value never changes. +bundle pins its SPM release, so a recorded value never changes. A US version +the service does not serve yet (its ``policyengine`` routes do not name it) +is left unrecorded rather than recorded as unavailable, and the next +publication run tries again; ``/v2/model-catalog`` offers no SPM settings for +it meanwhile. """ from __future__ import annotations @@ -44,17 +48,31 @@ MODEL_VERSIONS = TaxBenefitModelVersion.__table__ +def _served(policyengine_version: str, versions: dict[str, Any]) -> bool: + """Return whether the service's ``policyengine`` routes name the version.""" + routes = versions.get("policyengine") + if not isinstance(routes, dict): + return False + return policyengine_version != "latest" and ( + policyengine_version in routes or routes.get("latest") == policyengine_version + ) + + def spm_availability( country_id: str, policyengine_version: str, versions: object -) -> dict[str, Any]: +) -> dict[str, Any] | None: """Return ``{available, settings_schema, defaults}`` for one catalog version. Only a US version whose PolicyEngine.py version has a canonical capability - with pinned defaults (forecast artifact and scenario) offers SPM; anything - else is ``{"available": False}``. + with pinned defaults (forecast artifact and scenario) offers SPM; any other + country, and a US version the service serves without one, is + ``{"available": False}``. A US version the service does not serve yet is + None: not known. """ - if country_id != SPM_COUNTRY_ID or not isinstance(versions, dict): + if country_id != SPM_COUNTRY_ID: return dict(UNAVAILABLE) + if not isinstance(versions, dict) or not _served(policyengine_version, versions): + return None capabilities = versions.get("spm_capabilities") capability = ( capabilities.get(policyengine_version) @@ -141,6 +159,8 @@ def record_published_spm_availability( ) -> None: """Record each published country's SPM availability for the catalog version. + A US version the service does not serve yet stays unrecorded. + Raise V2ConfigurationError without ``SIMULATION_ENTRYPOINT_URL``, and CatalogPublicationError when the service's versions cannot be read. """ @@ -149,11 +169,13 @@ def record_published_spm_availability( ) with engine.begin() as connection: for country in catalog.countries: - record_spm_availability( - connection, - model_name=country.model.name, - version=country.model_version.version, - availability=spm_availability( - country.country_id, catalog.policyengine_version, versions - ), + availability = spm_availability( + country.country_id, catalog.policyengine_version, versions ) + if availability is not None: + record_spm_availability( + connection, + model_name=country.model.name, + version=country.model_version.version, + availability=availability, + ) diff --git a/policyengine_api/fastapi_routes/v2/metadata/common.py b/policyengine_api/fastapi_routes/v2/metadata/common.py index a2dab29ee..4f00dbf07 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/common.py +++ b/policyengine_api/fastapi_routes/v2/metadata/common.py @@ -140,7 +140,10 @@ def policyengine_version(self) -> str: A published catalog version never changes, so such a response stays valid until its code's response format changes; the ETag lets a client revalidate -it cheaply after that. +it cheaply after that. The one exception is a version's SPM availability in +``/v2/model-catalog``, recorded once after publication: until a later +publication run records it, a response may say the version offers no SPM +settings, and a reused response can keep saying so for up to this long. """ RUNNING_VERSION_MAX_AGE_SECONDS = 300 """How long a response for the running catalog version may be reused. diff --git a/tests/unit/v2/test_catalog_spm_availability.py b/tests/unit/v2/test_catalog_spm_availability.py index b07f3c6d8..e07e63833 100644 --- a/tests/unit/v2/test_catalog_spm_availability.py +++ b/tests/unit/v2/test_catalog_spm_availability.py @@ -2,6 +2,10 @@ from __future__ import annotations +from contextlib import contextmanager +from types import SimpleNamespace +from typing import TYPE_CHECKING + import httpx import pytest @@ -10,6 +14,9 @@ from policyengine_api.data.v2.settings import V2ConfigurationError from policyengine_api.spm import SPMSelection +if TYPE_CHECKING: + from collections.abc import Iterator + POLICYENGINE_VERSION = "6.10.0" DEFAULTS = { "forecast_content_sha256": "3d" * 32, @@ -19,8 +26,9 @@ "county_vintage": "2020", "as_of": None, } +ROUTES = {"latest": POLICYENGINE_VERSION, POLICYENGINE_VERSION: "policyengine-v6-10-0"} VERSIONS = { - "policyengine": {"latest": POLICYENGINE_VERSION}, + "policyengine": ROUTES, "spm_capabilities": { POLICYENGINE_VERSION: { "contract_version": "canonical-spm-v1", @@ -46,38 +54,38 @@ def test_a_us_version_with_a_certified_capability_offers_spm() -> None: ("country_id", "policyengine_version", "versions"), [ ("uk", POLICYENGINE_VERSION, VERSIONS), - ("us", "6.11.0", VERSIONS), - ("us", POLICYENGINE_VERSION, {}), - ("us", POLICYENGINE_VERSION, {"spm_capabilities": []}), + ("us", POLICYENGINE_VERSION, {"policyengine": ROUTES}), + ("us", POLICYENGINE_VERSION, {"policyengine": ROUTES, "spm_capabilities": []}), ( "us", POLICYENGINE_VERSION, { + "policyengine": ROUTES, "spm_capabilities": { POLICYENGINE_VERSION: { "contract_version": "canonical-spm-v0", "defaults": DEFAULTS, } - } + }, }, ), ( "us", POLICYENGINE_VERSION, { + "policyengine": ROUTES, "spm_capabilities": { POLICYENGINE_VERSION: { "contract_version": "canonical-spm-v1", "defaults": {**DEFAULTS, "scenario": None}, } - } + }, }, ), ], ids=[ "uk", - "unlisted-version", - "no-capabilities", + "no-capability", "malformed-capabilities", "other-contract", "unpinned-defaults", @@ -91,6 +99,71 @@ def test_anything_else_offers_no_spm( ) == {"available": False} +@pytest.mark.parametrize( + "versions", + [ + {**VERSIONS, "policyengine": {"latest": "6.11.0", "6.11.0": "app"}}, + {"spm_capabilities": VERSIONS["spm_capabilities"]}, + [], + ], + ids=["version-not-served", "no-route-map", "not-an-object"], +) +def test_a_us_version_the_service_does_not_serve_is_not_known_yet( + versions: object, +) -> None: + assert ( + spm_availability.spm_availability("us", POLICYENGINE_VERSION, versions) is None + ) + + +class RecordingEngine: + """An engine whose transactions are only counted.""" + + def __init__(self) -> None: + """Start with no transaction.""" + self.transactions = 0 + + @contextmanager + def begin(self) -> Iterator[str]: + """Count one transaction and hand out a placeholder connection.""" + self.transactions += 1 + yield "connection" + + +def test_publication_records_only_what_is_known( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recorded: list[tuple[str, dict]] = [] + monkeypatch.setattr( + spm_availability, + "record_spm_availability", + lambda _connection, *, model_name, version, availability: recorded.append( + (model_name, dict(availability)) + ), + ) + catalog = SimpleNamespace( + policyengine_version=POLICYENGINE_VERSION, + countries=[ + SimpleNamespace( + country_id=country_id, + model=SimpleNamespace(name=f"policyengine-{country_id}"), + model_version=SimpleNamespace(version="1.0.0"), + ) + for country_id in ("uk", "us") + ], + ) + not_served = {"policyengine": {"latest": "6.11.0", "6.11.0": "app"}} + + spm_availability.record_published_spm_availability( + RecordingEngine(), + catalog, + {"SIMULATION_ENTRYPOINT_URL": "https://simulation.example"}, + fetch_versions=lambda _url: not_served, + ) + + assert recorded == [("policyengine-uk", {"available": False})] + + def test_the_versions_come_from_the_simulation_services_public_route() -> None: requests: list[httpx.Request] = [] From f6f5f0c22a29bfb73ccc453ade1c77e5743c70dd Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 22:27:48 +0400 Subject: [PATCH 19/37] Compare reforms and added variables with API v1 in the staging parity set The parity set gains a US household under a stored reform that also adds missing variables (compared with /calculate-full) and a UK household under a stored reform, so the policy conversion and the added variables are checked against API v1 too. The memory check's description now says what it asks for. Co-Authored-By: Claude Opus 5.5 --- docs/migration/stage-11-v2-households.md | 9 ++- .../live_household_calculations.py | 70 +++++++++++++++++-- .../test_live_v2_household_calculations.py | 3 +- .../test_live_v2_household_parity.py | 42 +++++++++-- 4 files changed, 108 insertions(+), 16 deletions(-) diff --git a/docs/migration/stage-11-v2-households.md b/docs/migration/stage-11-v2-households.md index c43665420..2a350cef6 100644 --- a/docs/migration/stage-11-v2-households.md +++ b/docs/migration/stage-11-v2-households.md @@ -317,9 +317,12 @@ database: within 10 seconds and finished within 300 seconds; and one household varied across 401 points of employment income, which must finish. The measured times appear in the job log as `HouseholdCalculationMeasurement` warnings. -- `test_live_v2_household_parity.py`: two US and two UK households, one of - each with axes, calculated through `/v2/household-calculations` and through - API v1's `/calculate`; every calculated value and warning must match. +- `test_live_v2_household_parity.py`: three US and three UK households, one + of each with axes and one of each under a reform stored through + `/v2/policies` (the US one also adding missing variables), calculated + through `/v2/household-calculations` and through API v1's `/calculate` + (`/calculate-full` when adding missing variables); every calculated value + and warning must match. `STAGING_HOUSEHOLD_PARALLEL_COUNT`, `STAGING_HOUSEHOLD_SUBMIT_LIMIT_SECONDS`, `STAGING_HOUSEHOLD_FINISH_LIMIT_SECONDS` and diff --git a/tests/fixtures/integration/live_household_calculations.py b/tests/fixtures/integration/live_household_calculations.py index b3dc8f936..a640e32a6 100644 --- a/tests/fixtures/integration/live_household_calculations.py +++ b/tests/fixtures/integration/live_household_calculations.py @@ -25,6 +25,8 @@ TRANSIENT_STATUS_CODES = {500, 502, 503, 504} FINISHED_STATUSES = {"succeeded", "failed"} YEAR = "2026" +REFORM_START = f"{YEAR}-01-01" +OPEN_END = "2100-12-31" COLLECTIONS = { "us": { "household": "households", @@ -43,7 +45,10 @@ class LiveHousehold: ``people`` and ``groups`` map variable names to values by period; a null value asks for the variable to be calculated. ``axes`` vary inputs across - points, as in a ``/calculate`` situation. + points, as in a ``/calculate`` situation. ``reform`` sets each named + parameter to a value from the start of ``YEAR`` with no end, through a + policy stored with ``/v2/policies``. ``add_missing_variables`` calculates + every variable the household lacks, as ``/calculate-full`` does. """ name: str @@ -51,6 +56,8 @@ class LiveHousehold: people: dict[str, dict[str, Any]] groups: dict[str, dict[str, Any]] = field(default_factory=dict) axes: list[list[dict[str, Any]]] | None = None + reform: dict[str, int | float] | None = None + add_missing_variables: bool = False def document(self) -> dict[str, Any]: """Return the household as ``/v2/households`` stores it.""" @@ -117,16 +124,59 @@ def store_household(client: httpx.Client, household: LiveHousehold) -> str: return response.json()["result"]["item"]["id"] +def _one_catalog_item( + client: httpx.Client, resource: str, country_id: str, **filters: str +) -> dict[str, Any]: + response = client.get( + f"/v2/{resource}", params={"country_id": country_id, "limit": 1, **filters} + ) + assert response.status_code == 200, response.text[:500] + (item,) = response.json()["result"]["items"] + return item + + +def store_reform(client: httpx.Client, household: LiveHousehold) -> str | None: + """Store the household's reform through ``/v2/policies``; return its ID.""" + if household.reform is None: + return None + country_id = household.country_id + model = _one_catalog_item(client, "tax-benefit-models", country_id) + parameter_values = [ + { + "parameter_id": _one_catalog_item( + client, "parameters", country_id, name=parameter_name + )["id"], + "value": value, + "start_date": f"{REFORM_START}T00:00:00Z", + } + for parameter_name, value in household.reform.items() + ] + response = client.post( + "/v2/policies", + params={"country_id": country_id}, + json={ + "country_id": country_id, + "tax_benefit_model_id": model["id"], + "parameter_values": parameter_values, + }, + ) + assert response.status_code in {200, 201}, response.text[:500] + return response.json()["result"]["item"]["id"] + + def calculate( client: httpx.Client, household: LiveHousehold, household_id: str, + policy_id: str | None = None, ) -> FinishedCalculation: """Submit the stored household's calculation and poll until it finishes.""" started = time.monotonic() body: dict[str, Any] = { "country_id": household.country_id, "household_id": household_id, + "policy_id": policy_id, + "add_missing_variables": household.add_missing_variables, } if household.axes is not None: body["axes"] = household.axes @@ -156,11 +206,19 @@ def calculate( def calculate_with_v1(client: httpx.Client, household: LiveHousehold) -> dict: - """Return API v1's ``/calculate`` answer for the same household.""" - response = client.post( - f"/{household.country_id}/calculate", - json={"household": household.situation()}, - ) + """Return API v1's answer for the same household and reform. + + ``/calculate-full`` answers when the household adds missing variables, + ``/calculate`` otherwise. + """ + route = "calculate-full" if household.add_missing_variables else "calculate" + request: dict[str, Any] = {"household": household.situation()} + if household.reform is not None: + request["policy"] = { + parameter_name: {f"{REFORM_START}.{OPEN_END}": value} + for parameter_name, value in household.reform.items() + } + response = client.post(f"/{household.country_id}/{route}", json=request) assert response.status_code == 200, response.text payload = response.json() assert payload["status"] == "ok", payload diff --git a/tests/integration/test_live_v2_household_calculations.py b/tests/integration/test_live_v2_household_calculations.py index 278575621..c5a62fdb6 100644 --- a/tests/integration/test_live_v2_household_calculations.py +++ b/tests/integration/test_live_v2_household_calculations.py @@ -10,7 +10,8 @@ ``STAGING_HOUSEHOLD_SUBMIT_LIMIT_SECONDS`` and finished within ``STAGING_HOUSEHOLD_FINISH_LIMIT_SECONDS``; - a memory check: a household varied across 401 points of employment income, - every variable calculated at each point, which must finish without failing. + with its net income calculated at each point, which must finish without + failing; a household function that runs out of memory fails it. The measured times are reported as a warning, so the job log keeps them. """ diff --git a/tests/integration/test_live_v2_household_parity.py b/tests/integration/test_live_v2_household_parity.py index f310add83..4b9cc7be9 100644 --- a/tests/integration/test_live_v2_household_parity.py +++ b/tests/integration/test_live_v2_household_parity.py @@ -1,10 +1,12 @@ """Live staging parity of v2 household calculations with API v1's ``/calculate``. -Each household is stored through ``/v2/households``, calculated through -``/v2/household-calculations`` (the v2 household function on Modal) and, as -the same situation, through API v1's in-process ``/calculate`` on the same -candidate. Every calculated value must match, axes included. A run-specific -input makes each household new, so no earlier calculation is reused. +Each household is stored through ``/v2/households`` (and its reform through +``/v2/policies``), calculated through ``/v2/household-calculations`` (the v2 +household function on Modal) and, as the same situation and policy, through +API v1's in-process ``/calculate`` (``/calculate-full`` when it adds missing +variables) on the same candidate. Every calculated value must match, axes, +reforms and added variables included. A run-specific input makes each +household new, so no earlier calculation is reused. """ from __future__ import annotations @@ -19,6 +21,7 @@ household_differences, probe_amount, store_household, + store_reform, ) CALCULATED = {YEAR: None} @@ -77,6 +80,19 @@ def _parity_households(run_amount: int) -> list[LiveHousehold]: ] ], ), + LiveHousehold( + name="us-single-reform-added-variables", + country_id="us", + people={ + "adult": { + "age": {YEAR: 29}, + "employment_income": {YEAR: 61_000 + run_amount}, + } + }, + groups={"household": {"state_name": {YEAR: "NY"}}}, + reform={"gov.irs.deductions.standard.amount.SINGLE": 20_000}, + add_missing_variables=True, + ), LiveHousehold( name="uk-single", country_id="uk", @@ -90,6 +106,19 @@ def _parity_households(run_amount: int) -> list[LiveHousehold]: }, groups={"household": {"household_net_income": CALCULATED}}, ), + LiveHousehold( + name="uk-single-reform", + country_id="uk", + people={ + "adult": { + "age": {YEAR: 41}, + "employment_income": {YEAR: 36_000 + run_amount}, + "income_tax": CALCULATED, + } + }, + groups={"household": {"household_net_income": CALCULATED}}, + reform={"gov.hmrc.income_tax.allowances.personal_allowance.amount": 15_000}, + ), LiveHousehold( name="uk-couple-child-axes", country_id="uk", @@ -136,8 +165,9 @@ def test_a_v2_household_calculation_matches_api_v1( if candidate.name == household_name ] household_id = store_household(api_client, household) + policy_id = store_reform(api_client, household) - finished = calculate(api_client, household, household_id) + finished = calculate(api_client, household, household_id, policy_id) v1_answer = calculate_with_v1(api_client, household) assert finished.item["status"] == "succeeded", finished.item["error"] From 0a3851f9554ebf48e951029083db3e184f01e356 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 23:38:56 +0400 Subject: [PATCH 20/37] Keep the household's key order to and from the simulation service The generated client's models wrote the household with its groups before its people and each group's members after its values, in the request and in the answer. The household function lists warnings and calculated values in the household's key order, as API v1 does, so the order could differ from API v1's. The submission and the answers are now checked against the generated models but sent and returned as built and decoded. Co-Authored-By: Claude Opus 5.5 --- .../simulation_client.py | 96 +++++++++++-------- ...household_calculation_simulation_client.py | 23 +++++ 2 files changed, 77 insertions(+), 42 deletions(-) diff --git a/policyengine_api/services/v2/household_calculations/simulation_client.py b/policyengine_api/services/v2/household_calculations/simulation_client.py index 498d23614..4d482adba 100644 --- a/policyengine_api/services/v2/household_calculations/simulation_client.py +++ b/policyengine_api/services/v2/household_calculations/simulation_client.py @@ -1,11 +1,17 @@ """The simulation service's household calculation routes, through its client. Calls go to the Cloud Run Simulation Entrypoint (``SIMULATION_ENTRYPOINT_URL``), -the only simulation service with these routes, with the gateway credentials -the Public API already uses (``GATEWAY_AUTH_*``) and the request's -correlation headers. Answers are read through the client's generated models: -a 400 is a refusal, whose reason the service gives; a 422 lists what the -service finds invalid in the household or its options; anything else, +the only simulation service with these routes, through the generated client's +HTTP client, with the gateway credentials the Public API already uses +(``GATEWAY_AUTH_*``) and the request's correlation headers. + +The submission is checked against the generated submission model and sent as +built, and each answer is checked against the generated model for its status +and returned as decoded, both in their original key order: the household +function lists warnings and calculated values in the household's key order, +as API v1 does, and the generated models write their keys in their own +order. A 400 is a refusal, whose reason the service gives; a 422 lists what +the service finds invalid in the household or its options; anything else, including a 503, no answer, an answer of another shape, a gateway token that cannot be fetched or missing configuration, means the service is unavailable. @@ -16,14 +22,11 @@ from __future__ import annotations from functools import lru_cache -from typing import TYPE_CHECKING, Any, TypeVar +from http import HTTPStatus +from typing import TYPE_CHECKING, Any import httpx from policyengine_api_simulation_client import AuthenticatedClient -from policyengine_api_simulation_client.api.default import ( - get_household_calculation, - submit_household_calculation, -) from policyengine_api_simulation_client.models.household_calculation_response import ( HouseholdCalculationResponse, ) @@ -62,31 +65,43 @@ from collections.abc import Callable from uuid import UUID - from policyengine_api_simulation_client.types import Response - -ParsedT = TypeVar("ParsedT") +CALCULATIONS_PATH = "/v2/household-calculations" REQUEST_TIMEOUT_SECONDS = 30.0 CLOUD_RUN_SIMULATION_ENTRYPOINT = "cloud_run_simulation_entrypoint" -def _call(send: Callable[[], Response[ParsedT]]) -> Response[ParsedT]: - """Send one request; raise SimulationServiceUnavailableError if it cannot be read. - - The generated routes decode the body for each status they document, so a - body that is not JSON or not of the documented shape raises while - decoding. - """ +def _send(send: Callable[[], httpx.Response]) -> httpx.Response: + """Send one request; raise SimulationServiceUnavailableError if none answers.""" try: return send() except (httpx.HTTPError, GatewayAuthError) as error: msg = "The simulation service could not be reached." raise SimulationServiceUnavailableError(msg) from error - except (ValueError, KeyError, TypeError) as error: + + +def _decoded(response: httpx.Response, model: Any) -> tuple[Any, Any]: # noqa: ANN401 # a generated model class and its instance + """Return the decoded body and the generated model read from it. + + Raise SimulationServiceUnavailableError when the body is not JSON or not + of the model's shape. + """ + try: + body = response.json() + return body, model.from_dict(body) + except (ValueError, KeyError, TypeError, AttributeError) as error: msg = "The simulation service's answer could not be read." raise SimulationServiceUnavailableError(msg) from error +def _calculation(response: httpx.Response) -> dict[str, Any]: + body, _calculation_model = _decoded(response, HouseholdCalculationResponse) + if not isinstance(body, dict): + msg = "The simulation service's answer could not be read." + raise SimulationServiceUnavailableError(msg) + return body + + def _invalid_request_reason(error: HTTPValidationError) -> str: """Return each invalid field and why, without the request's ``body`` prefix.""" reasons = [] @@ -165,32 +180,29 @@ def from_environment(cls) -> SimulationHouseholdClient: def submit(self, submission: dict[str, Any]) -> dict[str, Any]: """Submit a calculation and return the service's calculation answer.""" - body = HouseholdCalculationSubmission.from_dict(submission) - response = _call( - lambda: submit_household_calculation.sync_detailed( - client=self._client, body=body - ) - ) - answer = response.parsed - if isinstance(answer, HouseholdCalculationResponse): - return answer.to_dict() - if isinstance(answer, HTTPValidationError): - raise SimulationServiceInvalidRequestError(_invalid_request_reason(answer)) - if isinstance(answer, HouseholdRouteError) and response.status_code == 400: # noqa: PLR2004 # the service's refusal status - raise SimulationServiceRefusedError(answer.detail) + HouseholdCalculationSubmission.from_dict(submission) + http_client = self._client.get_httpx_client() + response = _send(lambda: http_client.post(CALCULATIONS_PATH, json=submission)) + if response.status_code == HTTPStatus.ACCEPTED: + return _calculation(response) + if response.status_code == HTTPStatus.UNPROCESSABLE_ENTITY: + _body, invalid = _decoded(response, HTTPValidationError) + raise SimulationServiceInvalidRequestError(_invalid_request_reason(invalid)) + if response.status_code == HTTPStatus.BAD_REQUEST: + _body, refusal = _decoded(response, HouseholdRouteError) + raise SimulationServiceRefusedError(refusal.detail) raise _unavailable(response.status_code) def read(self, calculation_id: UUID) -> dict[str, Any] | None: """Return the service's calculation answer, or None when it has none.""" - response = _call( - lambda: get_household_calculation.sync_detailed( - calculation_id, client=self._client - ) + http_client = self._client.get_httpx_client() + response = _send( + lambda: http_client.get(f"{CALCULATIONS_PATH}/{calculation_id}") ) - answer = response.parsed - if isinstance(answer, HouseholdCalculationResponse): - return answer.to_dict() - if isinstance(answer, HouseholdRouteError) and response.status_code == 404: # noqa: PLR2004 # the service's unknown-calculation status + if response.status_code == HTTPStatus.OK: + return _calculation(response) + if response.status_code == HTTPStatus.NOT_FOUND: + _decoded(response, HouseholdRouteError) return None raise _unavailable(response.status_code) diff --git a/tests/unit/v2/test_household_calculation_simulation_client.py b/tests/unit/v2/test_household_calculation_simulation_client.py index e99eb96ca..747b629e9 100644 --- a/tests/unit/v2/test_household_calculation_simulation_client.py +++ b/tests/unit/v2/test_household_calculation_simulation_client.py @@ -237,3 +237,26 @@ def auth_flow(self, request: httpx.Request): # httpx's generator protocol client.submit(SUBMISSION) with pytest.raises(SimulationServiceUnavailableError): client.read(CALCULATION_ID) + + +def test_the_household_keeps_its_key_order_both_ways() -> None: + requests: list[httpx.Request] = [] + answer = succeeded_calculation() + client = client_answering(200, answer, requests) + submit_client = client_answering(202, answer, requests) + + submitted = submit_client.submit(SUBMISSION) + read = client.read(CALCULATION_ID) + + sent_household = json.loads(requests[0].content)["household"] + assert list(sent_household) == list(SITUATION) + assert list(sent_household["households"]["household-1"]) == list( + SITUATION["households"]["household-1"] + ) + for calculation in (submitted, read): + assert calculation is not None + returned_household = calculation["result"]["household"] + assert list(returned_household) == list(answer["result"]["household"]) + assert list(returned_household["households"]["household-1"]) == list( + answer["result"]["household"]["households"]["household-1"] + ) From d82e798067dd27702fd37619d44b9d902500ac51 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 23:41:40 +0400 Subject: [PATCH 21/37] Build the simulation service client only once a request reaches it The default service factory built the simulation client before any check, so with the simulation service unconfigured a country without household calculations, or a household that does not exist, got 503 instead of 422 or 404. The client is now built only when a submission or read is sent. Co-Authored-By: Claude Opus 5.5 --- .../fastapi_routes/dependencies.py | 2 +- .../v2/household_calculations/services.py | 12 ++++-- .../v2/test_household_calculation_routes.py | 2 +- .../v2/test_household_calculation_service.py | 40 ++++++++++++++++++- 4 files changed, 49 insertions(+), 7 deletions(-) diff --git a/policyengine_api/fastapi_routes/dependencies.py b/policyengine_api/fastapi_routes/dependencies.py index 9521876ae..2b9466f5a 100644 --- a/policyengine_api/fastapi_routes/dependencies.py +++ b/policyengine_api/fastapi_routes/dependencies.py @@ -899,7 +899,7 @@ def _default_v2_household_calculation_service_factory() -> ( households=_default_v2_household_service_factory(), policies=_default_v2_policy_service_factory(), catalog_reader_factory=_default_v2_metadata_reader_factory, - simulation=process_simulation_household_client(), + simulation_factory=process_simulation_household_client, ) diff --git a/policyengine_api/services/v2/household_calculations/services.py b/policyengine_api/services/v2/household_calculations/services.py index 4c5d7fca8..9250d0606 100644 --- a/policyengine_api/services/v2/household_calculations/services.py +++ b/policyengine_api/services/v2/household_calculations/services.py @@ -54,18 +54,22 @@ def __init__( households: HouseholdSource, policies: PolicySource, catalog_reader_factory: Callable[[], CatalogSource], - simulation: SimulationHouseholdCalculations, + simulation_factory: Callable[[], SimulationHouseholdCalculations], ) -> None: """Read through these services for one request. ``catalog_reader_factory`` opens a catalog reader, with its own database session, only once a submission needs the catalog version. + ``simulation_factory`` returns the simulation service's client only + once a request is sent to it, so a request the Public API refuses + gets its own answer even when the simulation service is not + configured. """ self._countries = countries self._households = households self._policies = policies self._catalog_reader_factory = catalog_reader_factory - self._simulation = simulation + self._simulation_factory = simulation_factory def submit_calculation( self, calculation: HouseholdCalculationInput @@ -132,12 +136,12 @@ def submit_calculation( if spm is not None: submission["spm"] = spm return HouseholdCalculationRead.from_simulation( - self._simulation.submit(submission) + self._simulation_factory().submit(submission) ) def get_calculation(self, calculation_id: UUID) -> HouseholdCalculationRead: """Return the calculation, or raise HouseholdCalculationNotFoundError.""" - calculation = self._simulation.read(calculation_id) + calculation = self._simulation_factory().read(calculation_id) if calculation is None: msg = f"Household calculation {calculation_id} not found." raise HouseholdCalculationNotFoundError(msg) diff --git a/tests/unit/v2/test_household_calculation_routes.py b/tests/unit/v2/test_household_calculation_routes.py index e45e6cd62..56ad0b595 100644 --- a/tests/unit/v2/test_household_calculation_routes.py +++ b/tests/unit/v2/test_household_calculation_routes.py @@ -52,7 +52,7 @@ def __init__(self, countries: FakeCountries | None = None) -> None: households=self.households, policies=FakePolicies(), catalog_reader_factory=FakeCatalogs, - simulation=self.simulation, + simulation_factory=lambda: self.simulation, ) dependencies = NativeRouteDependencies( readiness_probe=lambda: True, diff --git a/tests/unit/v2/test_household_calculation_service.py b/tests/unit/v2/test_household_calculation_service.py index 135ad18ae..aee458b1d 100644 --- a/tests/unit/v2/test_household_calculation_service.py +++ b/tests/unit/v2/test_household_calculation_service.py @@ -59,7 +59,7 @@ def __init__(self, countries: FakeCountries | None = None) -> None: households=self.households, policies=self.policies, catalog_reader_factory=self.open_catalog_reader, - simulation=self.simulation, + simulation_factory=lambda: self.simulation, ) def open_catalog_reader(self) -> FakeCatalogs: @@ -271,3 +271,41 @@ def test_a_refused_submission_and_a_read_open_no_catalog_reader() -> None: parts.service.get_calculation(CALCULATION_ID) assert parts.catalog_readers == [] + + +def _unconfigured_simulation() -> FakeSimulation: + message = "The simulation service is not configured." + raise SimulationServiceUnavailableError(message) + + +def test_requests_refused_by_the_public_api_need_no_simulation_service() -> None: + parts = Parts( + FakeCountries( + rows=( + CountryRead( + id="us", household_enabled=False, society_wide_enabled=False + ), + ) + ) + ) + unconfigured = V2HouseholdCalculationService( + countries=parts.countries, + households=parts.households, + policies=parts.policies, + catalog_reader_factory=parts.open_catalog_reader, + simulation_factory=_unconfigured_simulation, + ) + available_country = V2HouseholdCalculationService( + countries=FakeCountries(), + households=parts.households, + policies=parts.policies, + catalog_reader_factory=parts.open_catalog_reader, + simulation_factory=_unconfigured_simulation, + ) + + with pytest.raises(HouseholdCalculationCountryError): + unconfigured.submit_calculation(calculation_input()) + with pytest.raises(HouseholdNotFoundError): + available_country.submit_calculation(calculation_input(household_id=uuid4())) + with pytest.raises(SimulationServiceUnavailableError): + available_country.submit_calculation(calculation_input()) From 7ed16f2e3fb2f09dd6cc07e88aee84f0464d0b70 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 23:44:32 +0400 Subject: [PATCH 22/37] Vary the household's default year when an axis names no period An axis without a period varied the country model's default period, the current year, rather than the year the household's values name. Such an axis is now sent with the household's default year as its period. Co-Authored-By: Claude Opus 5.5 --- docs/migration/stage-11-v2-households.md | 3 +- .../v2/household_calculations/services.py | 34 +++++++++++++------ .../v2/household_calculations/types.py | 2 +- .../v2/test_household_calculation_routes.py | 4 ++- .../v2/test_household_calculation_service.py | 14 +++++++- 5 files changed, 42 insertions(+), 15 deletions(-) diff --git a/docs/migration/stage-11-v2-households.md b/docs/migration/stage-11-v2-households.md index 2a350cef6..f4e5ea674 100644 --- a/docs/migration/stage-11-v2-households.md +++ b/docs/migration/stage-11-v2-households.md @@ -262,7 +262,8 @@ for immediate rollback. `POST /v2/household-calculations` calculates a household stored through `/v2/households`, under a policy stored through `/v2/policies` or, without `policy_id`, under current law, for the household's `default_year`. The body -may add `axes` and `add_missing_variables`. An identical calculation that +may add `axes` (an axis without a `period` varies the household's +`default_year`) and `add_missing_variables`. An identical calculation that succeeded in the last 28 days is returned again; callers cannot force a new one. The Public API checks the country's `countries` row allows household calculations diff --git a/policyengine_api/services/v2/household_calculations/services.py b/policyengine_api/services/v2/household_calculations/services.py index 9250d0606..228f7461b 100644 --- a/policyengine_api/services/v2/household_calculations/services.py +++ b/policyengine_api/services/v2/household_calculations/services.py @@ -44,6 +44,28 @@ ) +def _axes_in_year( + calculation: HouseholdCalculationInput, year: int +) -> list[list[dict[str, Any]]] | None: + """Return the axes as sent, each without a period varying ``year``. + + The country model would otherwise vary its own default period, the + current year, rather than the year the household's values name. + """ + if calculation.axes is None: + return None + sent_axes = [] + for group in calculation.axes: + sent_group = [] + for axis in group: + sent_axis = axis.model_dump(mode="json", exclude_unset=True) + if sent_axis.get("period") is None: + sent_axis["period"] = str(year) + sent_group.append(sent_axis) + sent_axes.append(sent_group) + return sent_axes + + class V2HouseholdCalculationService: """Submit household calculations to the simulation service and read them.""" @@ -114,17 +136,7 @@ def submit_calculation( "parameters": parameters, }, "options": { - "axes": ( - None - if calculation.axes is None - else [ - [ - axis.model_dump(mode="json", exclude_unset=True) - for axis in group - ] - for group in calculation.axes - ] - ), + "axes": _axes_in_year(calculation, year), "add_missing_variables": calculation.add_missing_variables, }, # A calculation is reused for 28 days after it succeeded; callers diff --git a/policyengine_api/services/v2/household_calculations/types.py b/policyengine_api/services/v2/household_calculations/types.py index 78d50f8ed..f8e33e626 100644 --- a/policyengine_api/services/v2/household_calculations/types.py +++ b/policyengine_api/services/v2/household_calculations/types.py @@ -39,7 +39,7 @@ class HouseholdCalculationAxis(StrictCalculationModel): """One input varied across points. The variable ``name`` takes ``count`` evenly spaced values from ``min`` to - ``max`` in ``period`` (the household's default period when omitted), for + ``max`` in ``period`` (the household's ``default_year`` when omitted), for the entity member at position ``index`` (the first when omitted). """ diff --git a/tests/unit/v2/test_household_calculation_routes.py b/tests/unit/v2/test_household_calculation_routes.py index 56ad0b595..aca622d83 100644 --- a/tests/unit/v2/test_household_calculation_routes.py +++ b/tests/unit/v2/test_household_calculation_routes.py @@ -106,7 +106,9 @@ def test_a_calculation_with_axes_sends_them() -> None: response = app.client.post(PATH, json=body(axes=axes)) assert response.status_code == 202 - assert app.simulation.submissions[0]["options"]["axes"] == axes + assert app.simulation.submissions[0]["options"]["axes"] == [ + [{**axes[0][0], "period": "2026"}] + ] def test_a_succeeded_us_calculation_reads_back_its_spm_choice() -> None: diff --git a/tests/unit/v2/test_household_calculation_service.py b/tests/unit/v2/test_household_calculation_service.py index aee458b1d..246e119e1 100644 --- a/tests/unit/v2/test_household_calculation_service.py +++ b/tests/unit/v2/test_household_calculation_service.py @@ -104,7 +104,10 @@ def test_a_calculation_sends_the_stored_household_policy_and_catalog() -> None: } }, }, - "options": {"axes": AXES, "add_missing_variables": True}, + "options": { + "axes": [[{**AXES[0][0], "period": "2026"}]], + "add_missing_variables": True, + }, "force": False, "trigger": "initial", } @@ -309,3 +312,12 @@ def test_requests_refused_by_the_public_api_need_no_simulation_service() -> None available_country.submit_calculation(calculation_input(household_id=uuid4())) with pytest.raises(SimulationServiceUnavailableError): available_country.submit_calculation(calculation_input()) + + +def test_an_axis_keeps_its_own_period() -> None: + parts = Parts() + axes = [[{**AXES[0][0], "period": "2025-03"}]] + + parts.service.submit_calculation(calculation_input(axes=axes)) + + assert parts.simulation.submissions[0]["options"]["axes"] == axes From 538499de164b8a09762e8c44ec7ac6e2c619dfca Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 27 Oct 2026 23:57:07 +0400 Subject: [PATCH 23/37] Register the household calculation routes with the household route group /v2/household-calculations requests were logged with the route group "unknown"; they now belong to the household group. The contract registry records both routes as a typed v2 workflow, and the generated migration contract documents include them. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/migration-contracts.md | 16 ++++-- docs/generated/migration_contracts.json | 51 ++++++++++++++++++- policyengine_api/migration_registry.py | 1 + tests/contract/registry.py | 44 ++++++++++++++++ .../test_app_v2_workflow_contracts.py | 2 + .../unit/test_migration_contract_artifacts.py | 5 +- tests/unit/test_migration_flags.py | 5 ++ 7 files changed, 117 insertions(+), 7 deletions(-) diff --git a/docs/engineering/migration-contracts.md b/docs/engineering/migration-contracts.md index 31614d7d4..0398c7a94 100644 --- a/docs/engineering/migration-contracts.md +++ b/docs/engineering/migration-contracts.md @@ -7,8 +7,8 @@ Generated from `policyengine_api/migration_registry.py` and `tests/contract/regi | Metric | Count | | --- | ---: | | route group count | 9 | -| workflow count | 17 | -| request count | 93 | +| workflow count | 18 | +| request count | 95 | | db entity count | 6 | | sim flow count | 3 | @@ -20,7 +20,7 @@ Generated from `policyengine_api/migration_registry.py` and `tests/contract/regi | `specification` | `specification` | `none` | `none` | | `metadata` | `metadata`, `countries`, `datasets`, `economy-options`, `model-catalog`, `parameter-values`, `parameters`, `regions`, `tax-benefit-model-versions`, `tax-benefit-models`, `variables` | `metadata` | `none` | | `policy` | `policy`, `policies`, `user-policy`, `user-policies` | `policy` | `none` | -| `household` | `household`, `households`, `user-household`, `user-households`, `calculate`, `calculate-full` | `household` | `household` | +| `household` | `household`, `households`, `household-calculations`, `user-household`, `user-households`, `calculate`, `calculate-full` | `household` | `household` | | `economy` | `economy` | `simulation` | `economy` | | `simulation` | `simulation`, `simulations`, `user-simulations` | `simulation` | `economy` | | `report` | `report`, `user-reports`, `shared-reports` | `report` | `report` | @@ -100,6 +100,16 @@ Generated from `policyengine_api/migration_registry.py` and `tests/contract/regi | `GET` | `/v2/households/{household_id}?country_id=us` | 200 | `household` | none | `status`, `message`, `result.item.id`, `result.item.country_id`, `result.item.default_year`, `result.item.household_data`, `result.item.created_at`, `result.item.updated_at` | | `GET` | `/v2/households?country_id=us` | 200 | `household` | none | `status`, `message`, `result.items`, `result.offset`, `result.limit`, `result.has_more` | +### `household_calculations_v2` + +- Current contract: `typed_v2_resources` +- Future owner: Household calculations on v2 + +| Method | Path | Status | Route group | Required headers | Stable response fields | +| --- | --- | ---: | --- | --- | --- | +| `POST` | `/v2/household-calculations` | 202 | `household` | none | `status`, `message`, `result.calculation_id`, `result.country_id`, `result.household_id`, `result.policy_id`, `result.year`, `result.catalog_version_id`, `result.status`, `result.created_at` | +| `GET` | `/v2/household-calculations/{calculation_id}` | 200 | `household` | none | `status`, `message`, `result.calculation_id`, `result.status`, `result.completed_at`, `result.result.household`, `result.result.warnings`, `result.result.spm_config`, `result.result.spm_provenance`, `result.result.policyengine_version`, `result.error` | + ### `user_household_associations_v2` - Current contract: `typed_v2_resources` diff --git a/docs/generated/migration_contracts.json b/docs/generated/migration_contracts.json index 928b956b3..22de0e6a6 100644 --- a/docs/generated/migration_contracts.json +++ b/docs/generated/migration_contracts.json @@ -1,10 +1,10 @@ { "metadata": { "db_entity_count": 6, - "request_count": 93, + "request_count": 95, "route_group_count": 9, "sim_flow_count": 3, - "workflow_count": 17 + "workflow_count": 18 }, "route_groups": [ { @@ -61,6 +61,7 @@ "path_segments": [ "household", "households", + "household-calculations", "user-household", "user-households", "calculate", @@ -538,6 +539,52 @@ } ] }, + { + "current_contract": "typed_v2_resources", + "future_owner_pr": "Household calculations on v2", + "name": "household_calculations_v2", + "requests": [ + { + "expected_status": 202, + "headers": [], + "method": "POST", + "path": "/v2/household-calculations", + "route_group": "household", + "stable_response_fields": [ + "status", + "message", + "result.calculation_id", + "result.country_id", + "result.household_id", + "result.policy_id", + "result.year", + "result.catalog_version_id", + "result.status", + "result.created_at" + ] + }, + { + "expected_status": 200, + "headers": [], + "method": "GET", + "path": "/v2/household-calculations/{calculation_id}", + "route_group": "household", + "stable_response_fields": [ + "status", + "message", + "result.calculation_id", + "result.status", + "result.completed_at", + "result.result.household", + "result.result.warnings", + "result.result.spm_config", + "result.result.spm_provenance", + "result.result.policyengine_version", + "result.error" + ] + } + ] + }, { "current_contract": "typed_v2_resources", "future_owner_pr": "PR 12A: User Login and v1-Anchored User Associations", diff --git a/policyengine_api/migration_registry.py b/policyengine_api/migration_registry.py index 74bcba926..f46c5b8aa 100644 --- a/policyengine_api/migration_registry.py +++ b/policyengine_api/migration_registry.py @@ -60,6 +60,7 @@ class RouteGroupConfig: path_segments=( "household", "households", + "household-calculations", "user-household", "user-households", "calculate", diff --git a/tests/contract/registry.py b/tests/contract/registry.py index ac187802d..987c9dee8 100644 --- a/tests/contract/registry.py +++ b/tests/contract/registry.py @@ -355,6 +355,50 @@ class WorkflowContract: ), ), ), + WorkflowContract( + name="household_calculations_v2", + current_contract="typed_v2_resources", + future_owner_pr="Household calculations on v2", + requests=( + ContractRequest( + method="POST", + path="/v2/household-calculations", + expected_status=202, + stable_response_fields=( + "status", + "message", + "result.calculation_id", + "result.country_id", + "result.household_id", + "result.policy_id", + "result.year", + "result.catalog_version_id", + "result.status", + "result.created_at", + ), + route_group="household", + ), + ContractRequest( + method="GET", + path="/v2/household-calculations/{calculation_id}", + expected_status=200, + stable_response_fields=( + "status", + "message", + "result.calculation_id", + "result.status", + "result.completed_at", + "result.result.household", + "result.result.warnings", + "result.result.spm_config", + "result.result.spm_provenance", + "result.result.policyengine_version", + "result.error", + ), + route_group="household", + ), + ), + ), WorkflowContract( name="user_household_associations_v2", current_contract="typed_v2_resources", diff --git a/tests/contract/test_app_v2_workflow_contracts.py b/tests/contract/test_app_v2_workflow_contracts.py index f74640678..a4580d9c7 100644 --- a/tests/contract/test_app_v2_workflow_contracts.py +++ b/tests/contract/test_app_v2_workflow_contracts.py @@ -14,6 +14,7 @@ def test_app_v2_workflow_contract_registry_is_complete(): "user_policy_associations_v2", "household_save_read", "household_resources_v2", + "household_calculations_v2", "user_household_associations_v2", "user_simulation_associations_v2", "user_report_associations_v2", @@ -37,6 +38,7 @@ def test_app_v2_workflow_contract_registry_is_complete(): "policy_resources_v2", "user_policy_associations_v2", "household_resources_v2", + "household_calculations_v2", "user_household_associations_v2", "user_simulation_associations_v2", "user_report_associations_v2", diff --git a/tests/unit/test_migration_contract_artifacts.py b/tests/unit/test_migration_contract_artifacts.py index 026d0950f..90eec1780 100644 --- a/tests/unit/test_migration_contract_artifacts.py +++ b/tests/unit/test_migration_contract_artifacts.py @@ -13,8 +13,8 @@ def test_migration_contract_payload_summarizes_route_contracts(): assert payload["version"] == 1 assert payload["metadata"] == { "route_group_count": 9, - "workflow_count": 17, - "request_count": 93, + "workflow_count": 18, + "request_count": 95, "db_entity_count": 6, "sim_flow_count": 3, } @@ -25,6 +25,7 @@ def test_migration_contract_payload_summarizes_route_contracts(): "user_policy_associations_v2", "household_save_read", "household_resources_v2", + "household_calculations_v2", "user_household_associations_v2", "user_simulation_associations_v2", "user_report_associations_v2", diff --git a/tests/unit/test_migration_flags.py b/tests/unit/test_migration_flags.py index e110da879..4ca5dd68f 100644 --- a/tests/unit/test_migration_flags.py +++ b/tests/unit/test_migration_flags.py @@ -181,6 +181,11 @@ def test_explicit_migration_context_rejects_invalid_database_sources( "policy", ), ("/v2/user-households", "household"), + ("/v2/household-calculations", "household"), + ( + "/v2/household-calculations/00000000-0000-0000-0000-000000000001", + "household", + ), ("/v2/user-simulations", "simulation"), ("/v2/user-reports/00000000-0000-0000-0000-000000000001", "report"), ("/v2/me", "user_profile"), From bdbab60e44cd758fe14d3ebb0316ab3d92dde3e6 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 28 Oct 2026 00:00:27 +0400 Subject: [PATCH 24/37] Check the simulation service URL first, and never fail on its versions Catalog publication now stops before extraction when SIMULATION_ENTRYPOINT_URL is unset, instead of after the catalog is written. When the simulation service's /versions answer cannot be read, publication logs a warning and leaves SPM availability unrecorded for the next run, so a slow or unreachable service no longer stops a release. The cache note and the stage 9 guide say that an instance keeps an earlier answer until it restarts. Co-Authored-By: Claude Opus 5.5 --- docs/migration/stage-9-v2-metadata.md | 16 +++++---- .../data/v2/catalog/initialization.py | 5 ++- .../data/v2/catalog/spm_availability.py | 26 +++++++++----- .../fastapi_routes/v2/metadata/common.py | 9 +++-- tests/unit/v2/test_catalog_initialization.py | 16 +++++++++ .../unit/v2/test_catalog_spm_availability.py | 36 +++++++++++++++++++ 6 files changed, 89 insertions(+), 19 deletions(-) diff --git a/docs/migration/stage-9-v2-metadata.md b/docs/migration/stage-9-v2-metadata.md index ac447d1b1..de87f336f 100644 --- a/docs/migration/stage-9-v2-metadata.md +++ b/docs/migration/stage-9-v2-metadata.md @@ -44,8 +44,10 @@ simulation service's public `GET /versions`: a US version offers SPM when the service lists a canonical capability with pinned defaults for its PolicyEngine.py version, and other countries never do. A US version the service does not serve yet stays unrecorded until a later publication run, and -a version whose availability is not recorded yet offers none; a response read -before then, and reused for up to a day, keeps saying so. The result names the +a version whose availability is not recorded yet offers none. An API instance +that read the version before it was recorded keeps that answer until it +restarts, normally at the deployment after the next publication run, and a +client may reuse it for up to a day after that. The result names the selected PolicyEngine.py version, so a client that starts without a version learns the running one and requests every other part for that exact version. The query uses the shared `CatalogQuery` model: unknown and repeated @@ -198,11 +200,11 @@ must execute in this order: `HUGGING_FACE_TOKEN` (from the organization secret `PE_UK_PRIVATE_HF_READ_TOKEN`) and `SIMULATION_ENTRYPOINT_URL` (from the environment secret of the same name), then run - `uv run python scripts/publish_v2_metadata_catalog.py`. After publishing, the - command reads the simulation service's public `GET /versions` to record each - version's SPM availability, and fails if the URL is unset or the answer - cannot be read; the catalog itself is already published by then, and the - next run records the availability. Before publishing a + `uv run python scripts/publish_v2_metadata_catalog.py`. The command stops + before extraction if `SIMULATION_ENTRYPOINT_URL` is unset. After publishing, + it reads the simulation service's public `GET /versions` to record each + version's SPM availability; if the answer cannot be read, it logs a warning, + records none, and the next run records them. Before publishing a catalog version that is not stored yet, the command reads the certified data artifacts its regions' filters use and stops if one cannot select a region. It stops before reading anything if a UK version needs checking and diff --git a/policyengine_api/data/v2/catalog/initialization.py b/policyengine_api/data/v2/catalog/initialization.py index 036b055d9..6c974f121 100644 --- a/policyengine_api/data/v2/catalog/initialization.py +++ b/policyengine_api/data/v2/catalog/initialization.py @@ -30,6 +30,7 @@ ) from policyengine_api.data.v2.catalog.spm_availability import ( record_published_spm_availability, + simulation_entrypoint_url, ) from policyengine_api.data.v2.settings import ( V2ConfigurationError, @@ -130,10 +131,12 @@ def initialize_catalog( ) -> PublicationEvidence: """Load only row-write settings, extract fully, verify, then publish explicitly. - After publishing, record each version's SPM availability. + After publishing, record each version's SPM availability. The simulation + service's URL, which that needs, is checked before anything else. """ selected_steps = steps or PublicationSteps() settings = load_v2_data_write_database_settings(environ) + simulation_entrypoint_url(os.environ if environ is None else environ) catalog = extractor() engine = engine_builder(settings) try: diff --git a/policyengine_api/data/v2/catalog/spm_availability.py b/policyengine_api/data/v2/catalog/spm_availability.py index 2b5f28a2b..ab9c3f1c4 100644 --- a/policyengine_api/data/v2/catalog/spm_availability.py +++ b/policyengine_api/data/v2/catalog/spm_availability.py @@ -13,11 +13,13 @@ the service does not serve yet (its ``policyengine`` routes do not name it) is left unrecorded rather than recorded as unavailable, and the next publication run tries again; ``/v2/model-catalog`` offers no SPM settings for -it meanwhile. +it meanwhile. So is every version when the service's versions cannot be +read: recording never fails the publication. """ from __future__ import annotations +import logging import os from typing import TYPE_CHECKING, Any @@ -39,6 +41,7 @@ from policyengine_api.data.v2.catalog.records import NormalizedCatalog +LOGGER = logging.getLogger(__name__) SIMULATION_ENTRYPOINT_URL = "SIMULATION_ENTRYPOINT_URL" SPM_COUNTRY_ID = "us" SPM_CONTRACT_VERSION = "canonical-spm-v1" @@ -159,14 +162,21 @@ def record_published_spm_availability( ) -> None: """Record each published country's SPM availability for the catalog version. - A US version the service does not serve yet stays unrecorded. - - Raise V2ConfigurationError without ``SIMULATION_ENTRYPOINT_URL``, and - CatalogPublicationError when the service's versions cannot be read. + A US version the service does not serve yet stays unrecorded, and so does + every version when the service's versions cannot be read; the next + publication run tries again. Raise V2ConfigurationError without + ``SIMULATION_ENTRYPOINT_URL``. """ - versions = fetch_versions( - simulation_entrypoint_url(os.environ if environ is None else environ) - ) + url = simulation_entrypoint_url(os.environ if environ is None else environ) + try: + versions = fetch_versions(url) + except CatalogPublicationError as error: + LOGGER.warning( + "SPM availability was not recorded; the next publication run " + "tries again: %s", + error, + ) + return with engine.begin() as connection: for country in catalog.countries: availability = spm_availability( diff --git a/policyengine_api/fastapi_routes/v2/metadata/common.py b/policyengine_api/fastapi_routes/v2/metadata/common.py index 4f00dbf07..5d36bf50d 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/common.py +++ b/policyengine_api/fastapi_routes/v2/metadata/common.py @@ -141,9 +141,12 @@ def policyengine_version(self) -> str: A published catalog version never changes, so such a response stays valid until its code's response format changes; the ETag lets a client revalidate it cheaply after that. The one exception is a version's SPM availability in -``/v2/model-catalog``, recorded once after publication: until a later -publication run records it, a response may say the version offers no SPM -settings, and a reused response can keep saying so for up to this long. +``/v2/model-catalog``, recorded once after publication: until a publication +run records it, a response says the version offers no SPM settings. A +process keeps the selected catalog rows and built responses it has read for +its whole life, so it keeps saying so until it restarts (normally at the next +deployment, which follows the next publication run), and a client may reuse +such a response for up to this long after that. """ RUNNING_VERSION_MAX_AGE_SECONDS = 300 """How long a response for the running catalog version may be reused. diff --git a/tests/unit/v2/test_catalog_initialization.py b/tests/unit/v2/test_catalog_initialization.py index aee1342b5..22f01b72b 100644 --- a/tests/unit/v2/test_catalog_initialization.py +++ b/tests/unit/v2/test_catalog_initialization.py @@ -34,6 +34,7 @@ V2_SUPABASE_ENVIRONMENT: "test-foundation", "V2_RUNTIME_DATABASE_URL": "invalid-runtime-value", "V2_MIGRATION_DATABASE_URL": "invalid-migration-value", + "SIMULATION_ENTRYPOINT_URL": "https://simulation.example", } @@ -308,3 +309,18 @@ def verify_without_token(): assert "HUGGING_FACE_TOKEN must be set to read the certified uk data" in ( failure.err ) + + +def test_a_missing_simulation_service_url_stops_before_extraction() -> None: + def refuse_extraction(): + msg = "extraction must not start" + raise AssertionError(msg) + + environment = { + name: value + for name, value in ENVIRONMENT.items() + if name != "SIMULATION_ENTRYPOINT_URL" + } + + with pytest.raises(V2ConfigurationError, match="SIMULATION_ENTRYPOINT_URL"): + initialization.initialize_catalog(environment, extractor=refuse_extraction) diff --git a/tests/unit/v2/test_catalog_spm_availability.py b/tests/unit/v2/test_catalog_spm_availability.py index e07e63833..c8cd1ac1b 100644 --- a/tests/unit/v2/test_catalog_spm_availability.py +++ b/tests/unit/v2/test_catalog_spm_availability.py @@ -205,3 +205,39 @@ def test_recording_needs_the_simulation_service_url() -> None: ) == "https://simulation.example" ) + + +def test_an_unreadable_versions_answer_leaves_every_version_unrecorded( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recorded: list[str] = [] + monkeypatch.setattr( + spm_availability, + "record_spm_availability", + lambda _connection, *, model_name, version, availability: recorded.append( + model_name + ), + ) + catalog = SimpleNamespace( + policyengine_version=POLICYENGINE_VERSION, + countries=[ + SimpleNamespace( + country_id="us", + model=SimpleNamespace(name="policyengine-us"), + model_version=SimpleNamespace(version="1.0.0"), + ) + ], + ) + + def unreadable(_url: str) -> object: + msg = "the simulation service's versions could not be read" + raise CatalogPublicationError(msg) + + spm_availability.record_published_spm_availability( + RecordingEngine(), + catalog, + {"SIMULATION_ENTRYPOINT_URL": "https://simulation.example"}, + fetch_versions=unreadable, + ) + + assert recorded == [] From b54f7eff9ed0506cde77b7cc78cc3e3fef7f0f8c Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 28 Oct 2026 00:01:29 +0400 Subject: [PATCH 25/37] Send a live check's calculation again after a 5xx answer The simulation service answers 503 with Retry-After while its database or version manifest cannot be read. The live staging checks retried such answers when polling but not when submitting, so one short outage failed the job that staging promotion waits on. Submissions are now sent again, after the Retry-After wait, until the calculation's deadline. Co-Authored-By: Claude Opus 5.5 --- .../live_household_calculations.py | 27 +++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/tests/fixtures/integration/live_household_calculations.py b/tests/fixtures/integration/live_household_calculations.py index a640e32a6..58d0cf113 100644 --- a/tests/fixtures/integration/live_household_calculations.py +++ b/tests/fixtures/integration/live_household_calculations.py @@ -164,6 +164,29 @@ def store_reform(client: httpx.Client, household: LiveHousehold) -> str | None: return response.json()["result"]["item"]["id"] +def _retry_wait_seconds(response: httpx.Response) -> float: + retry_after = response.headers.get("Retry-After", "") + return float(retry_after) if retry_after.isdigit() else POLL_INTERVAL_SECONDS + + +def _submit( + client: httpx.Client, body: dict[str, Any], deadline: float +) -> httpx.Response: + """Post the calculation, sending it again after a 5xx answer until the deadline. + + The simulation service answers 503 with ``Retry-After`` while its database + or version manifest cannot be read; such an outage is not a failure of + the calculation. + """ + while True: + response = client.post(CALCULATION_PATH, json=body) + if response.status_code not in TRANSIENT_STATUS_CODES: + return response + if time.monotonic() >= deadline: + return response + time.sleep(_retry_wait_seconds(response)) + + def calculate( client: httpx.Client, household: LiveHousehold, @@ -180,11 +203,11 @@ def calculate( } if household.axes is not None: body["axes"] = household.axes - response = client.post(CALCULATION_PATH, json=body) + deadline = started + CALCULATION_TIMEOUT_SECONDS + response = _submit(client, body, deadline) submit_seconds = time.monotonic() - started assert response.status_code == 202, response.text item = response.json()["result"] - deadline = started + CALCULATION_TIMEOUT_SECONDS while item["status"] not in FINISHED_STATUSES: if time.monotonic() >= deadline: msg = ( From 4108a807ccd0ac699229ead50c37640b6061edee Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 28 Oct 2026 00:21:09 +0400 Subject: [PATCH 26/37] Log household calculations with the sources and compute they use Requests to /v2/household-calculations were logged with the household route group's v1 database flags and its configured compute backend. They now log a Supabase read for a submission, no database read for a poll, no write, and the Cloud Run Simulation Entrypoint's v2 workers as the compute, whatever those flags select. Co-Authored-By: Claude Opus 5.5 --- policyengine_api/migration_logging.py | 20 ++++++++ .../routes/test_migration_context_logging.py | 49 +++++++++++++++++++ 2 files changed, 69 insertions(+) diff --git a/policyengine_api/migration_logging.py b/policyengine_api/migration_logging.py index 18f5ac7a9..be94d982c 100644 --- a/policyengine_api/migration_logging.py +++ b/policyengine_api/migration_logging.py @@ -61,6 +61,15 @@ V2_USER_RESOURCE_SEGMENTS = frozenset( {"user-policies", "user-households", "user-simulations", "user-reports", "me"} ) +# /v2/household-calculations reads the stored household, policy and catalog +# from Supabase when submitting, stores nothing itself, and always computes on +# the Cloud Run Simulation Entrypoint's v2 workers, whatever the household +# route group's flags select. +V2_HOUSEHOLD_CALCULATION_SEGMENT = "household-calculations" +V2_HOUSEHOLD_CALCULATION_COMPUTE = { + "sim_entrypoint": "cloud_run_simulation_entrypoint", + "sim_compute": "v2_primary", +} V2_RESOURCE_METHODS = frozenset({"GET", "POST", "PATCH", "DELETE"}) # A v2 resource path has at least the "v2" segment and the resource segment. V2_RESOURCE_PATH_MINIMUM_SEGMENTS = 2 @@ -247,6 +256,17 @@ def _migration_context( route_impl: RouteImplementation | None, ) -> dict[str, str | None]: route_group = infer_route_group(path) + if _v2_resource_segment(path) == V2_HOUSEHOLD_CALCULATION_SEGMENT: + context = get_migration_log_context( + route_group, + route_impl=route_impl, + explicit_db_sources=ExplicitDatabaseSources( + db_read_source="supabase" if method == "POST" else None + ), + ) + if "migration_flag_error" not in context: + context.update(V2_HOUSEHOLD_CALCULATION_COMPUTE) + return context is_v2_metadata_read = _is_v2_metadata_resource_read(method, path) is_v2_supabase_resource = _is_v2_supabase_resource(method, path) uses_explicit_v2_source = is_v2_metadata_read or is_v2_supabase_resource diff --git a/tests/unit/routes/test_migration_context_logging.py b/tests/unit/routes/test_migration_context_logging.py index 641c8c34b..2be8e2f4c 100644 --- a/tests/unit/routes/test_migration_context_logging.py +++ b/tests/unit/routes/test_migration_context_logging.py @@ -717,3 +717,52 @@ def test_compute_routes_in_the_same_groups_keep_their_compute_labels(): assert context["route_group"] == "simulation" assert context["sim_flow"] == "economy" assert context["sim_compute"] is not None + + +def test_household_calculations_log_their_actual_sources_and_compute(monkeypatch): + monkeypatch.setenv("DB_READ_HOUSEHOLD", "invalid-for-native-v2") + monkeypatch.setenv("DB_WRITE_HOUSEHOLD", "invalid-for-native-v2") + monkeypatch.setenv("SIM_ENTRYPOINT", "old_gateway_direct") + monkeypatch.delenv("SIM_COMPUTE_HOUSEHOLD", raising=False) + + with patch("policyengine_api.migration_logging.logger") as mock_logger: + for method, path in ( + ("POST", "/v2/household-calculations"), + ( + "GET", + "/v2/household-calculations/00000000-0000-0000-0000-000000000001", + ), + ): + log_migration_request( + ServedRequest( + request_id=f"request-{method}", + method=method, + path=path, + status_code=202 if method == "POST" else 200, + started_at=None, + route_impl=RouteImplementation.FASTAPI_NATIVE, + ) + ) + + submit_context, read_context = ( + call.args[0]["migration"] for call in mock_logger.log_struct.call_args_list + ) + computed_on_v2 = { + "route_group": "household", + "route_impl": "fastapi_native", + "sim_flow": "household", + "sim_entrypoint": "cloud_run_simulation_entrypoint", + "sim_compute": "v2_primary", + } + assert submit_context == { + **submit_context, + **computed_on_v2, + "db_write": None, + "db_read": "supabase", + } + assert read_context == { + **read_context, + **computed_on_v2, + "db_write": None, + "db_read": None, + } From e3b5f8dc41f8a71e110840ca8dce38d7d0fadfce Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 28 Oct 2026 00:24:15 +0400 Subject: [PATCH 27/37] Drop the unused generated client wrapper and correct two descriptions The household client held its HTTP client inside the generated AuthenticatedClient with an empty token but never called a generated route; it now keeps the HTTP client itself. Its description says the household's key order is the stored household's as the v2 database returns it. The SPM availability description and changelog entry say a version gets its record on the next publication run of its own PolicyEngine.py version, and that only an unreadable versions answer never fails publication. Co-Authored-By: Claude Opus 5.5 --- changelog.d/catalog-spm-availability.added.md | 2 +- .../data/v2/catalog/spm_availability.py | 10 +++-- .../simulation_client.py | 41 ++++++++----------- 3 files changed, 25 insertions(+), 28 deletions(-) diff --git a/changelog.d/catalog-spm-availability.added.md b/changelog.d/catalog-spm-availability.added.md index bdf508050..c22bc0f26 100644 --- a/changelog.d/catalog-spm-availability.added.md +++ b/changelog.d/catalog-spm-availability.added.md @@ -1 +1 @@ -Each published catalog version now records whether it offers SPM settings, with their schema and certified defaults, from the simulation service's public versions answer; the catalog publication step reads `SIMULATION_ENTRYPOINT_URL`, and versions published earlier get it on the next publication run. +Each published catalog version now records whether it offers SPM settings, with their schema and certified defaults, from the simulation service's public versions answer; the catalog publication step reads `SIMULATION_ENTRYPOINT_URL`, and a version published earlier gets it on the next publication run of its PolicyEngine.py version. diff --git a/policyengine_api/data/v2/catalog/spm_availability.py b/policyengine_api/data/v2/catalog/spm_availability.py index ab9c3f1c4..398cb2a41 100644 --- a/policyengine_api/data/v2/catalog/spm_availability.py +++ b/policyengine_api/data/v2/catalog/spm_availability.py @@ -7,14 +7,16 @@ version offers SPM when the service reports a canonical capability for its PolicyEngine.py version; other countries never do. -Publication records it after the catalog itself, on each version that has -none yet, so versions published before it existed get it on the next run. A -bundle pins its SPM release, so a recorded value never changes. A US version +Publication records it after the catalog itself, for the installed +PolicyEngine.py version's catalog versions that have none yet, so a version +published before this existed gets it on the next publication run of that +PolicyEngine.py version; versions of other PolicyEngine.py versions keep none. +A bundle pins its SPM release, so a recorded value never changes. A US version the service does not serve yet (its ``policyengine`` routes do not name it) is left unrecorded rather than recorded as unavailable, and the next publication run tries again; ``/v2/model-catalog`` offers no SPM settings for it meanwhile. So is every version when the service's versions cannot be -read: recording never fails the publication. +read: an unreadable answer never fails the publication. """ from __future__ import annotations diff --git a/policyengine_api/services/v2/household_calculations/simulation_client.py b/policyengine_api/services/v2/household_calculations/simulation_client.py index 4d482adba..dd2c4ee6a 100644 --- a/policyengine_api/services/v2/household_calculations/simulation_client.py +++ b/policyengine_api/services/v2/household_calculations/simulation_client.py @@ -1,16 +1,18 @@ -"""The simulation service's household calculation routes, through its client. +"""The simulation service's household calculation routes, checked by its models. Calls go to the Cloud Run Simulation Entrypoint (``SIMULATION_ENTRYPOINT_URL``), -the only simulation service with these routes, through the generated client's -HTTP client, with the gateway credentials the Public API already uses -(``GATEWAY_AUTH_*``) and the request's correlation headers. - -The submission is checked against the generated submission model and sent as -built, and each answer is checked against the generated model for its status -and returned as decoded, both in their original key order: the household -function lists warnings and calculated values in the household's key order, -as API v1 does, and the generated models write their keys in their own -order. A 400 is a refusal, whose reason the service gives; a 422 lists what +the only simulation service with these routes, over HTTP with the gateway +credentials the Public API already uses (``GATEWAY_AUTH_*``) and the request's +correlation headers. + +The submission is checked against the generated client's submission model and +sent as built, and each answer is checked against the generated model for its +status and returned as decoded, both in their original key order: the +household function lists warnings and calculated values in the household's +key order, and the generated models write their keys in an order of their +own. That order is the stored household's as the v2 database returns it, +which orders each record's keys itself, so it can differ from the order a +caller sent. A 400 is a refusal, whose reason the service gives; a 422 lists what the service finds invalid in the household or its options; anything else, including a 503, no answer, an answer of another shape, a gateway token that cannot be fetched or missing configuration, means the service is unavailable. @@ -26,7 +28,6 @@ from typing import TYPE_CHECKING, Any import httpx -from policyengine_api_simulation_client import AuthenticatedClient from policyengine_api_simulation_client.models.household_calculation_response import ( HouseholdCalculationResponse, ) @@ -148,13 +149,7 @@ def __init__( ) instrument_httpx(http_client, get_runtime()) self._base_url = base_url.rstrip("/") - # The generated routes take an AuthenticatedClient. Its fixed token is - # sent only from an httpx client it builds itself; the one set here - # authenticates through ``httpx_args`` (GatewayBearerAuth, which - # fetches and renews the gateway token), so the empty token is unused. - self._client = AuthenticatedClient( - base_url=base_url, token="" - ).set_httpx_client(http_client) + self._http_client = http_client @property def base_url(self) -> str: @@ -181,8 +176,9 @@ def from_environment(cls) -> SimulationHouseholdClient: def submit(self, submission: dict[str, Any]) -> dict[str, Any]: """Submit a calculation and return the service's calculation answer.""" HouseholdCalculationSubmission.from_dict(submission) - http_client = self._client.get_httpx_client() - response = _send(lambda: http_client.post(CALCULATIONS_PATH, json=submission)) + response = _send( + lambda: self._http_client.post(CALCULATIONS_PATH, json=submission) + ) if response.status_code == HTTPStatus.ACCEPTED: return _calculation(response) if response.status_code == HTTPStatus.UNPROCESSABLE_ENTITY: @@ -195,9 +191,8 @@ def submit(self, submission: dict[str, Any]) -> dict[str, Any]: def read(self, calculation_id: UUID) -> dict[str, Any] | None: """Return the service's calculation answer, or None when it has none.""" - http_client = self._client.get_httpx_client() response = _send( - lambda: http_client.get(f"{CALCULATIONS_PATH}/{calculation_id}") + lambda: self._http_client.get(f"{CALCULATIONS_PATH}/{calculation_id}") ) if response.status_code == HTTPStatus.OK: return _calculation(response) From 243c2d9cba7d02491bf9c7aa46ced5333accc684 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 28 Oct 2026 00:25:54 +0400 Subject: [PATCH 28/37] Retry only short outages in the live checks, and compare warnings in any order The live household checks sent a request again after any 5xx answer until the calculation's 600-second deadline, so a lasting failure kept each check waiting until the job's time limit and lost the failing answer. They now retry 502, 503 and 504 for at most 90 seconds in a row and fail at once, with the answer's body, on anything else. The parity set compares warnings regardless of order, since the v2 database orders stored keys itself, and one US household now has a deprecated input so that both calculations warn. Co-Authored-By: Claude Opus 5.5 --- docs/migration/stage-11-v2-households.md | 11 +++- .../live_household_calculations.py | 56 +++++++++++-------- .../test_live_v2_household_parity.py | 8 ++- 3 files changed, 49 insertions(+), 26 deletions(-) diff --git a/docs/migration/stage-11-v2-households.md b/docs/migration/stage-11-v2-households.md index f4e5ea674..222d87e31 100644 --- a/docs/migration/stage-11-v2-households.md +++ b/docs/migration/stage-11-v2-households.md @@ -325,10 +325,17 @@ database: (`/calculate-full` when adding missing variables); every calculated value and warning must match. +Each request is sent again while it gets 502, 503 or 504 for at most 90 +seconds in a row; any other failure, or a longer outage, fails the check with +the answer's body. Warnings are compared regardless of order, because the v2 +database orders each stored record's keys itself; one US household has a +deprecated input, so both calculations warn. + `STAGING_HOUSEHOLD_PARALLEL_COUNT`, `STAGING_HOUSEHOLD_SUBMIT_LIMIT_SECONDS`, -`STAGING_HOUSEHOLD_FINISH_LIMIT_SECONDS` and +`STAGING_HOUSEHOLD_FINISH_LIMIT_SECONDS`, `STAGING_HOUSEHOLD_CALCULATION_TIMEOUT_SECONDS` (600 seconds per calculation) -change those limits; none is set in the workflow. The checks see neither the +and `STAGING_HOUSEHOLD_TRANSIENT_LIMIT_SECONDS` change those limits; none is +set in the workflow. The checks see neither the number of Simulation Entrypoint instances that served the calculations nor the household functions' memory use: a calculation that runs out of memory fails, and so fails the check. Both modules fail until the staging Simulation diff --git a/tests/fixtures/integration/live_household_calculations.py b/tests/fixtures/integration/live_household_calculations.py index 58d0cf113..580715718 100644 --- a/tests/fixtures/integration/live_household_calculations.py +++ b/tests/fixtures/integration/live_household_calculations.py @@ -10,9 +10,12 @@ import os import time from dataclasses import dataclass, field +from functools import partial from typing import TYPE_CHECKING, Any if TYPE_CHECKING: + from collections.abc import Callable + import httpx CALCULATION_PATH = "/v2/household-calculations" @@ -22,7 +25,15 @@ POLL_INTERVAL_SECONDS = float( os.environ.get("STAGING_API_TEST_POLL_INTERVAL_SECONDS", "5") ) -TRANSIENT_STATUS_CODES = {500, 502, 503, 504} +# A 502, 503 or 504 can be a short outage, such as the simulation service +# failing to read its database or version manifest for a moment; each request +# is sent again while such answers last at most TRANSIENT_LIMIT_SECONDS in a +# row. Any other failure, and a longer outage, fails the check at once with +# the answer's body. +TRANSIENT_STATUS_CODES = {502, 503, 504} +TRANSIENT_LIMIT_SECONDS = float( + os.environ.get("STAGING_HOUSEHOLD_TRANSIENT_LIMIT_SECONDS", "90") +) FINISHED_STATUSES = {"succeeded", "failed"} YEAR = "2026" REFORM_START = f"{YEAR}-01-01" @@ -49,6 +60,8 @@ class LiveHousehold: parameter to a value from the start of ``YEAR`` with no end, through a policy stored with ``/v2/policies``. ``add_missing_variables`` calculates every variable the household lacks, as ``/calculate-full`` does. + ``expects_warnings`` says the household has an input the country model + ignores with a warning. """ name: str @@ -58,6 +71,7 @@ class LiveHousehold: axes: list[list[dict[str, Any]]] | None = None reform: dict[str, int | float] | None = None add_missing_variables: bool = False + expects_warnings: bool = False def document(self) -> dict[str, Any]: """Return the household as ``/v2/households`` stores it.""" @@ -164,27 +178,21 @@ def store_reform(client: httpx.Client, household: LiveHousehold) -> str | None: return response.json()["result"]["item"]["id"] -def _retry_wait_seconds(response: httpx.Response) -> float: - retry_after = response.headers.get("Retry-After", "") - return float(retry_after) if retry_after.isdigit() else POLL_INTERVAL_SECONDS - - -def _submit( - client: httpx.Client, body: dict[str, Any], deadline: float -) -> httpx.Response: - """Post the calculation, sending it again after a 5xx answer until the deadline. - - The simulation service answers 503 with ``Retry-After`` while its database - or version manifest cannot be read; such an outage is not a failure of - the calculation. - """ +def _send_through_short_outages(send: Callable[[], httpx.Response]) -> httpx.Response: + """Send the request again while it gets 502, 503 or 504 for a short time.""" + outage_started: float | None = None while True: - response = client.post(CALCULATION_PATH, json=body) + response = send() if response.status_code not in TRANSIENT_STATUS_CODES: return response - if time.monotonic() >= deadline: - return response - time.sleep(_retry_wait_seconds(response)) + outage_started = outage_started or time.monotonic() + if time.monotonic() - outage_started >= TRANSIENT_LIMIT_SECONDS: + msg = ( + f"HTTP {response.status_code} for {TRANSIENT_LIMIT_SECONDS:.0f} s: " + f"{response.text[:500]}" + ) + raise AssertionError(msg) + time.sleep(POLL_INTERVAL_SECONDS) def calculate( @@ -204,7 +212,9 @@ def calculate( if household.axes is not None: body["axes"] = household.axes deadline = started + CALCULATION_TIMEOUT_SECONDS - response = _submit(client, body, deadline) + response = _send_through_short_outages( + lambda: client.post(CALCULATION_PATH, json=body) + ) submit_seconds = time.monotonic() - started assert response.status_code == 202, response.text item = response.json()["result"] @@ -216,9 +226,9 @@ def calculate( ) raise AssertionError(msg) time.sleep(POLL_INTERVAL_SECONDS) - polled = client.get(f"{CALCULATION_PATH}/{item['calculation_id']}") - if polled.status_code in TRANSIENT_STATUS_CODES: - continue + polled = _send_through_short_outages( + partial(client.get, f"{CALCULATION_PATH}/{item['calculation_id']}") + ) assert polled.status_code == 200, polled.text item = polled.json()["result"] return FinishedCalculation( diff --git a/tests/integration/test_live_v2_household_parity.py b/tests/integration/test_live_v2_household_parity.py index 4b9cc7be9..758482aa8 100644 --- a/tests/integration/test_live_v2_household_parity.py +++ b/tests/integration/test_live_v2_household_parity.py @@ -36,6 +36,8 @@ def _parity_households(run_amount: int) -> list[LiveHousehold]: "adult": { "age": {YEAR: 35}, "employment_income": {YEAR: 42_000 + run_amount}, + # Deprecated: dropped with a warning by both calculations. + "medical_out_of_pocket_expenses": {YEAR: 1_000}, "income_tax": CALCULATED, } }, @@ -47,6 +49,7 @@ def _parity_households(run_amount: int) -> list[LiveHousehold]: "tax_unit": {"eitc": CALCULATED, "ctc": CALCULATED}, "spm_unit": {"snap": CALCULATED}, }, + expects_warnings=True, ), LiveHousehold( name="us-couple-two-children-texas-axes", @@ -174,7 +177,10 @@ def test_a_v2_household_calculation_matches_api_v1( result = finished.item["result"] differences = household_differences(v1_answer["result"], result["household"]) assert differences == [], "\n".join(differences[:50]) - assert result["warnings"] == v1_answer.get("warnings", []) + # The v2 database orders each stored record's keys itself, so warnings, + # listed in the household's key order, may come in another order. + assert sorted(result["warnings"]) == sorted(v1_answer.get("warnings", [])) + assert bool(result["warnings"]) == household.expects_warnings if household.country_id == "us": assert result["spm_config"] is not None assert result["spm_provenance"] is not None From ba761d42ba112257721bfd801ed31a94235f55ba Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 28 Oct 2026 00:41:56 +0400 Subject: [PATCH 29/37] Log household calculations without checking flags they do not use An invalid SIM_ENTRYPOINT or SIM_COMPUTE_HOUSEHOLD, which household calculations never read, reduced their migration log record to a flag error. The record is now built from what the route uses, checking only the route implementation. The guide and the live check helper say which requests are sent again after a 502, 503 or 504, and that failures the Public API answers with 503 take about 90 seconds to fail. Co-Authored-By: Claude Opus 5.5 --- docs/migration/stage-11-v2-households.md | 9 ++-- policyengine_api/migration_logging.py | 44 +++++++++++++------ .../simulation_client.py | 14 +++--- .../live_household_calculations.py | 12 ++--- .../routes/test_migration_context_logging.py | 27 ++++++++++++ 5 files changed, 77 insertions(+), 29 deletions(-) diff --git a/docs/migration/stage-11-v2-households.md b/docs/migration/stage-11-v2-households.md index 222d87e31..55f9b75c9 100644 --- a/docs/migration/stage-11-v2-households.md +++ b/docs/migration/stage-11-v2-households.md @@ -325,9 +325,12 @@ database: (`/calculate-full` when adding missing variables); every calculated value and warning must match. -Each request is sent again while it gets 502, 503 or 504 for at most 90 -seconds in a row; any other failure, or a longer outage, fails the check with -the answer's body. Warnings are compared regardless of order, because the v2 +A calculation's submission and each poll are sent again while they get 502, +503 or 504 for at most 90 seconds in a row; any other failure, or a longer +run of such answers, fails the check with the answer's body. The Public API +also answers 503 for failures that do not pass, such as a simulation service +without these routes, so those take about 90 seconds to fail. Every other +request in the checks fails at once. Warnings are compared regardless of order, because the v2 database orders each stored record's keys itself; one US household has a deprecated input, so both calculations warn. diff --git a/policyengine_api/migration_logging.py b/policyengine_api/migration_logging.py index be94d982c..e4b2e22ac 100644 --- a/policyengine_api/migration_logging.py +++ b/policyengine_api/migration_logging.py @@ -13,10 +13,13 @@ from policyengine_api.gcp_logging import logger from policyengine_api.migration_flags import ( ExplicitDatabaseSources, + MigrationContext, RouteImplementation, get_migration_log_context, + get_route_impl, infer_route_group, ) +from policyengine_api.migration_registry import ROUTE_GROUP_CONFIG_BY_NAME from policyengine_api.observability.identifiers import ( OBSERVABILITY_ID_HEADER, normalize_observability_id, @@ -64,12 +67,11 @@ # /v2/household-calculations reads the stored household, policy and catalog # from Supabase when submitting, stores nothing itself, and always computes on # the Cloud Run Simulation Entrypoint's v2 workers, whatever the household -# route group's flags select. +# route group's database, entrypoint and compute flags select; it reads none +# of them, so none of them is checked for its log record. V2_HOUSEHOLD_CALCULATION_SEGMENT = "household-calculations" -V2_HOUSEHOLD_CALCULATION_COMPUTE = { - "sim_entrypoint": "cloud_run_simulation_entrypoint", - "sim_compute": "v2_primary", -} +V2_HOUSEHOLD_CALCULATION_ENTRYPOINT = "cloud_run_simulation_entrypoint" +V2_HOUSEHOLD_CALCULATION_COMPUTE = "v2_primary" V2_RESOURCE_METHODS = frozenset({"GET", "POST", "PATCH", "DELETE"}) # A v2 resource path has at least the "v2" segment and the resource segment. V2_RESOURCE_PATH_MINIMUM_SEGMENTS = 2 @@ -249,6 +251,27 @@ def log_migration_request(served_request: ServedRequest) -> None: ) +def _household_calculation_context( + route_group: str, *, method: str, route_impl: RouteImplementation | None +) -> dict[str, str | None]: + """Return a household calculation's log record from what it actually uses.""" + try: + selected_route_impl = route_impl or get_route_impl(route_group) + except ValueError as error: + return {"route_group": route_group, "migration_flag_error": str(error)} + route_config = ROUTE_GROUP_CONFIG_BY_NAME[route_group] + return MigrationContext( + route_group=route_group, + route_impl=selected_route_impl, + db_entity=route_config.db_entity, + db_write=None, + db_read="supabase" if method == "POST" else None, + sim_flow=route_config.sim_flow, + sim_entrypoint=V2_HOUSEHOLD_CALCULATION_ENTRYPOINT, + sim_compute=V2_HOUSEHOLD_CALCULATION_COMPUTE, + ).to_log_dict() + + def _migration_context( *, method: str, @@ -257,16 +280,9 @@ def _migration_context( ) -> dict[str, str | None]: route_group = infer_route_group(path) if _v2_resource_segment(path) == V2_HOUSEHOLD_CALCULATION_SEGMENT: - context = get_migration_log_context( - route_group, - route_impl=route_impl, - explicit_db_sources=ExplicitDatabaseSources( - db_read_source="supabase" if method == "POST" else None - ), + return _household_calculation_context( + route_group, method=method, route_impl=route_impl ) - if "migration_flag_error" not in context: - context.update(V2_HOUSEHOLD_CALCULATION_COMPUTE) - return context is_v2_metadata_read = _is_v2_metadata_resource_read(method, path) is_v2_supabase_resource = _is_v2_supabase_resource(method, path) uses_explicit_v2_source = is_v2_metadata_read or is_v2_supabase_resource diff --git a/policyengine_api/services/v2/household_calculations/simulation_client.py b/policyengine_api/services/v2/household_calculations/simulation_client.py index dd2c4ee6a..0f2d158fc 100644 --- a/policyengine_api/services/v2/household_calculations/simulation_client.py +++ b/policyengine_api/services/v2/household_calculations/simulation_client.py @@ -7,13 +7,15 @@ The submission is checked against the generated client's submission model and sent as built, and each answer is checked against the generated model for its -status and returned as decoded, both in their original key order: the +status and returned as decoded, both in their original key order, because +the generated models write their keys in an order of their own and the household function lists warnings and calculated values in the household's -key order, and the generated models write their keys in an order of their -own. That order is the stored household's as the v2 database returns it, -which orders each record's keys itself, so it can differ from the order a -caller sent. A 400 is a refusal, whose reason the service gives; a 422 lists what -the service finds invalid in the household or its options; anything else, +key order. The household's key order is the stored household's as the v2 +database returns it, which orders each record's keys itself, so it can +differ from the order a caller sent. + +A 400 is a refusal, whose reason the service gives; a 422 lists what the +service finds invalid in the household or its options; anything else, including a 503, no answer, an answer of another shape, a gateway token that cannot be fetched or missing configuration, means the service is unavailable. diff --git a/tests/fixtures/integration/live_household_calculations.py b/tests/fixtures/integration/live_household_calculations.py index 580715718..960891496 100644 --- a/tests/fixtures/integration/live_household_calculations.py +++ b/tests/fixtures/integration/live_household_calculations.py @@ -26,10 +26,10 @@ os.environ.get("STAGING_API_TEST_POLL_INTERVAL_SECONDS", "5") ) # A 502, 503 or 504 can be a short outage, such as the simulation service -# failing to read its database or version manifest for a moment; each request -# is sent again while such answers last at most TRANSIENT_LIMIT_SECONDS in a -# row. Any other failure, and a longer outage, fails the check at once with -# the answer's body. +# failing to read its database or version manifest for a moment; a +# calculation's submission and each poll are sent again while such answers +# last at most TRANSIENT_LIMIT_SECONDS in a row. Any other failure, and a +# longer run of such answers, fails the check with the answer's body. TRANSIENT_STATUS_CODES = {502, 503, 504} TRANSIENT_LIMIT_SECONDS = float( os.environ.get("STAGING_HOUSEHOLD_TRANSIENT_LIMIT_SECONDS", "90") @@ -60,8 +60,8 @@ class LiveHousehold: parameter to a value from the start of ``YEAR`` with no end, through a policy stored with ``/v2/policies``. ``add_missing_variables`` calculates every variable the household lacks, as ``/calculate-full`` does. - ``expects_warnings`` says the household has an input the country model - ignores with a warning. + ``expects_warnings`` says the household has a deprecated input, which the + calculation drops with a warning before the country model runs. """ name: str diff --git a/tests/unit/routes/test_migration_context_logging.py b/tests/unit/routes/test_migration_context_logging.py index 2be8e2f4c..efc3d27bf 100644 --- a/tests/unit/routes/test_migration_context_logging.py +++ b/tests/unit/routes/test_migration_context_logging.py @@ -766,3 +766,30 @@ def test_household_calculations_log_their_actual_sources_and_compute(monkeypatch "db_write": None, "db_read": None, } + + +def test_household_calculations_ignore_compute_flags_they_do_not_use(monkeypatch): + monkeypatch.setenv("SIM_ENTRYPOINT", "invalid-unused") + monkeypatch.setenv("SIM_COMPUTE_HOUSEHOLD", "invalid-unused") + + with patch("policyengine_api.migration_logging.logger") as mock_logger: + log_migration_request( + ServedRequest( + request_id="request-submit", + method="POST", + path="/v2/household-calculations", + status_code=202, + started_at=None, + route_impl=RouteImplementation.FASTAPI_NATIVE, + ) + ) + + context = mock_logger.log_struct.call_args.args[0]["migration"] + assert "migration_flag_error" not in context + assert context == { + **context, + "route_group": "household", + "db_read": "supabase", + "sim_entrypoint": "cloud_run_simulation_entrypoint", + "sim_compute": "v2_primary", + } From d695268a68b2b5ad2444eddae06b9d65169343c0 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 28 Oct 2026 00:46:22 +0400 Subject: [PATCH 30/37] Type the simulation answer reader without a suppression The helper that reads a simulation service answer took a generated model class as Any, with an ANN401 suppression; it now takes the model's from_dict reader, typed by the model it returns. Co-Authored-By: Claude Opus 5.5 --- .../simulation_client.py | 23 ++++++++++++------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/policyengine_api/services/v2/household_calculations/simulation_client.py b/policyengine_api/services/v2/household_calculations/simulation_client.py index 0f2d158fc..6045e8281 100644 --- a/policyengine_api/services/v2/household_calculations/simulation_client.py +++ b/policyengine_api/services/v2/household_calculations/simulation_client.py @@ -27,7 +27,7 @@ from functools import lru_cache from http import HTTPStatus -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, TypeVar import httpx from policyengine_api_simulation_client.models.household_calculation_response import ( @@ -65,7 +65,7 @@ ) if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Mapping from uuid import UUID CALCULATIONS_PATH = "/v2/household-calculations" @@ -83,7 +83,12 @@ def _send(send: Callable[[], httpx.Response]) -> httpx.Response: raise SimulationServiceUnavailableError(msg) from error -def _decoded(response: httpx.Response, model: Any) -> tuple[Any, Any]: # noqa: ANN401 # a generated model class and its instance +ModelT = TypeVar("ModelT") + + +def _decoded( + response: httpx.Response, read_model: Callable[[Mapping[str, Any]], ModelT] +) -> tuple[object, ModelT]: """Return the decoded body and the generated model read from it. Raise SimulationServiceUnavailableError when the body is not JSON or not @@ -91,14 +96,16 @@ def _decoded(response: httpx.Response, model: Any) -> tuple[Any, Any]: # noqa: """ try: body = response.json() - return body, model.from_dict(body) + return body, read_model(body) except (ValueError, KeyError, TypeError, AttributeError) as error: msg = "The simulation service's answer could not be read." raise SimulationServiceUnavailableError(msg) from error def _calculation(response: httpx.Response) -> dict[str, Any]: - body, _calculation_model = _decoded(response, HouseholdCalculationResponse) + body, _calculation_model = _decoded( + response, HouseholdCalculationResponse.from_dict + ) if not isinstance(body, dict): msg = "The simulation service's answer could not be read." raise SimulationServiceUnavailableError(msg) @@ -184,10 +191,10 @@ def submit(self, submission: dict[str, Any]) -> dict[str, Any]: if response.status_code == HTTPStatus.ACCEPTED: return _calculation(response) if response.status_code == HTTPStatus.UNPROCESSABLE_ENTITY: - _body, invalid = _decoded(response, HTTPValidationError) + _body, invalid = _decoded(response, HTTPValidationError.from_dict) raise SimulationServiceInvalidRequestError(_invalid_request_reason(invalid)) if response.status_code == HTTPStatus.BAD_REQUEST: - _body, refusal = _decoded(response, HouseholdRouteError) + _body, refusal = _decoded(response, HouseholdRouteError.from_dict) raise SimulationServiceRefusedError(refusal.detail) raise _unavailable(response.status_code) @@ -199,7 +206,7 @@ def read(self, calculation_id: UUID) -> dict[str, Any] | None: if response.status_code == HTTPStatus.OK: return _calculation(response) if response.status_code == HTTPStatus.NOT_FOUND: - _decoded(response, HouseholdRouteError) + _decoded(response, HouseholdRouteError.from_dict) return None raise _unavailable(response.status_code) From 220cf1e9401bf5cb2b63824a7779064d6dafe81d Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 28 Oct 2026 16:52:36 +0400 Subject: [PATCH 31/37] Send the simulation service household and policy IDs, not their values The simulation service now reads the stored household, the policy's parameter values and the catalog version from the Public API's tables itself. A household calculation's submission therefore carries only the country, the household's and policy's IDs, the options and the running catalog version's PolicyEngine.py version. The Public API still checks the country allows household calculations and that the household and policy exist in it, and fills an axis's missing period with the household's default year. The conversion of stored households and policies, which moved into the simulation service, is removed here. Co-Authored-By: Claude Opus 5.5 --- docs/migration/stage-11-v2-households.md | 16 +- .../v2/household_calculations/services.py | 45 ++--- .../household_calculations/transformations.py | 83 --------- .../v2/test_household_calculation_service.py | 39 ++-- ...household_calculation_simulation_client.py | 21 +-- ...t_household_calculation_transformations.py | 172 ------------------ 6 files changed, 42 insertions(+), 334 deletions(-) delete mode 100644 policyengine_api/services/v2/household_calculations/transformations.py delete mode 100644 tests/unit/v2/test_household_calculation_transformations.py diff --git a/docs/migration/stage-11-v2-households.md b/docs/migration/stage-11-v2-households.md index 55f9b75c9..0a33fa8bb 100644 --- a/docs/migration/stage-11-v2-households.md +++ b/docs/migration/stage-11-v2-households.md @@ -267,12 +267,16 @@ may add `axes` (an axis without a `period` varies the household's succeeded in the last 28 days is returned again; callers cannot force a new one. The Public API checks the country's `countries` row allows household calculations -(422 otherwise; a household-only country is accepted), reads the household and -policy through the same services as their routes (404 when either is not -stored in that country), selects the country's running catalog version, and -sends the household in the situation shape, its saved SPM choice and the -policy's parameter values to the simulation service, which cannot read these -tables. The answer is 202 with the calculation's `calculation_id` and +(422 otherwise; a household-only country is accepted), checks the household +and policy exist through the same services as their routes (404 when either +is not stored in that country), selects the country's running catalog +version, and sends the simulation service the household's and policy's IDs, +the options and that version's PolicyEngine.py version. The simulation +service reads the household, the policy's parameter values and the catalog +version from these tables itself, through its own `SELECT` grant, and +refuses a policy changing a parameter the catalog version lacks, which the +Public API answers with 503. The answer is 202 with the calculation's +`calculation_id` and `status`. Poll `GET /v2/household-calculations/{calculation_id}` until `status` is `succeeded`, with `result`, or `failed`, with `error`. A US result carries `spm_config` and `spm_provenance`. diff --git a/policyengine_api/services/v2/household_calculations/services.py b/policyengine_api/services/v2/household_calculations/services.py index 228f7461b..41dcae0fa 100644 --- a/policyengine_api/services/v2/household_calculations/services.py +++ b/policyengine_api/services/v2/household_calculations/services.py @@ -1,13 +1,15 @@ """Calculate stored households through the simulation service, and read them. Submitting a calculation checks the country's ``countries`` row allows -household calculations (a household-only country does), reads the stored -household and policy through the same services as ``/v2/households`` and -``/v2/policies`` (so a household or policy of another country, or none, is -not found), selects the country's running catalog version, and sends the -simulation service everything it needs: it cannot read these tables. The -simulation service records every calculation, so one not saved to an -account is still stored, as a simulation run with no report. +household calculations (a household-only country does), checks the stored +household and policy exist in that country through the same services as +``/v2/households`` and ``/v2/policies`` (so one of another country, or none, +is not found), selects the country's running catalog version, and sends the +simulation service their IDs and that version's PolicyEngine.py version. +The simulation service reads the household, the policy's parameter values +and the catalog version from the Public API's tables itself, and calculates +the household's default year. It records every calculation, so one not +saved to an account is still stored, as a simulation run with no report. Reading a calculation asks the simulation service for it as it stands. """ @@ -16,11 +18,6 @@ from typing import TYPE_CHECKING, Any -from policyengine_api.services.v2.household_calculations.transformations import ( - household_situation, - policy_parameters, - saved_spm_choice, -) from policyengine_api.services.v2.household_calculations.types import ( HouseholdCalculationRead, ) @@ -107,12 +104,9 @@ def submit_calculation( household = self._households.get_household( country_id=country_id, household_id=calculation.household_id ) - parameters: dict[str, dict[str, Any]] = {} if calculation.policy_id is not None: - parameters = policy_parameters( - self._policies.get_policy( - country_id=country_id, policy_id=calculation.policy_id - ) + self._policies.get_policy( + country_id=country_id, policy_id=calculation.policy_id ) year = calculation_year(household.default_year) catalog_reader = self._catalog_reader_factory() @@ -122,19 +116,11 @@ def submit_calculation( catalog_reader.close() submission: dict[str, Any] = { "country": country_id, - "catalog_version_id": str(catalog.model_version.id), "policyengine_version": catalog.policyengine_version, - "year": year, "household_id": str(household.id), - "household": household_situation(country_id, household.household_data), - "policy": { - "policy_id": ( - None - if calculation.policy_id is None - else str(calculation.policy_id) - ), - "parameters": parameters, - }, + "policy_id": ( + None if calculation.policy_id is None else str(calculation.policy_id) + ), "options": { "axes": _axes_in_year(calculation, year), "add_missing_variables": calculation.add_missing_variables, @@ -144,9 +130,6 @@ def submit_calculation( "force": False, "trigger": "initial", } - spm = saved_spm_choice(household.household_data) - if spm is not None: - submission["spm"] = spm return HouseholdCalculationRead.from_simulation( self._simulation_factory().submit(submission) ) diff --git a/policyengine_api/services/v2/household_calculations/transformations.py b/policyengine_api/services/v2/household_calculations/transformations.py deleted file mode 100644 index 51de6eb10..000000000 --- a/policyengine_api/services/v2/household_calculations/transformations.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Stored households and policies, in the shapes the Simulation Entrypoint takes. - -A stored v2 household lists its people and each group entity as records, -each with an ``id`` and its ``values``; each person names, for every group -collection, the ``id`` of the group it belongs to. The Simulation Entrypoint's -household functions take the Public API's situation shape instead: ``people`` -maps each person to their values, and each group collection, by its plural, -maps each group to its ``members`` and values. Records keep their ``id`` as -their name, so a calculated value maps back to its record. - -A v2 policy stores each parameter value with its start date and an optional -end date; the household functions take the Public API's policy JSON, each -parameter's values keyed by ``YYYY-MM-DD.YYYY-MM-DD`` period ranges, where an -open end is ``2100-12-31``. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any - -from policyengine_api.services.v2.households.transformations import ( - LEGACY_COLLECTION_NAMES, -) -from policyengine_api.services.v2.households.validators import ( - SUPPORTED_ENTITY_COLLECTIONS, -) - -if TYPE_CHECKING: - from collections.abc import Mapping - - from policyengine_api.services.v2.policies.types import PolicyRead - -# Each country's stored group collections, as the households service accepts -# them, and their plural in the situation. -GROUP_COLLECTIONS: dict[str, dict[str, str]] = { - country_id: { - collection: LEGACY_COLLECTION_NAMES[collection] for collection in collections - } - for country_id, collections in SUPPORTED_ENTITY_COLLECTIONS.items() -} -OPEN_END_DATE = "2100-12-31" - - -def household_situation( - country_id: str, household_data: Mapping[str, Any] -) -> dict[str, Any]: - """Return a stored household's people and groups in the situation shape. - - The saved SPM choice is not part of the situation; ``saved_spm_choice`` - reads it. - """ - people = household_data["people"] - situation: dict[str, Any] = { - "people": {person["id"]: dict(person["values"]) for person in people} - } - for collection, plural in GROUP_COLLECTIONS[country_id].items(): - situation[plural] = { - group["id"]: { - "members": [ - person["id"] - for person in people - if person["memberships"].get(collection) == group["id"] - ], - **group["values"], - } - for group in household_data.get(collection, []) - } - return situation - - -def saved_spm_choice(household_data: Mapping[str, Any]) -> object: - """Return the household's saved SPM selection, or None when it has none.""" - return household_data.get("spm") - - -def policy_parameters(policy: PolicyRead) -> dict[str, dict[str, Any]]: - """Return the policy's values by parameter and period range.""" - parameters: dict[str, dict[str, Any]] = {} - for value in policy.parameter_values: - end = OPEN_END_DATE if value.end_date is None else value.end_date.date() - period = f"{value.start_date.date().isoformat()}.{end}" - parameters.setdefault(value.parameter_name, {})[period] = value.value - return parameters diff --git a/tests/unit/v2/test_household_calculation_service.py b/tests/unit/v2/test_household_calculation_service.py index 246e119e1..331adb3c5 100644 --- a/tests/unit/v2/test_household_calculation_service.py +++ b/tests/unit/v2/test_household_calculation_service.py @@ -23,12 +23,9 @@ from policyengine_api.services.v2.policies.validators import PolicyNotFoundError from tests.unit.v2.household_calculation_support import ( CALCULATION_ID, - CATALOG_VERSION_ID, HOUSEHOLD_ID, POLICY_ID, POLICYENGINE_VERSION, - SITUATION, - SPM_CHOICE, SPM_CONFIG, SPM_PROVENANCE, FakeCatalogs, @@ -80,7 +77,8 @@ def calculation_input(**changes: object) -> HouseholdCalculationInput: return HouseholdCalculationInput.model_validate(values) -def test_a_calculation_sends_the_stored_household_policy_and_catalog() -> None: +def test_a_calculation_sends_the_household_and_policy_ids_and_the_version() -> None: + """The simulation service reads the household, policy and catalog version.""" parts = Parts() calculation = parts.service.submit_calculation( @@ -90,20 +88,9 @@ def test_a_calculation_sends_the_stored_household_policy_and_catalog() -> None: (submission,) = parts.simulation.submissions assert submission == { "country": "us", - "catalog_version_id": str(CATALOG_VERSION_ID), "policyengine_version": POLICYENGINE_VERSION, - "year": 2026, "household_id": str(HOUSEHOLD_ID), - "household": SITUATION, - "spm": SPM_CHOICE, - "policy": { - "policy_id": str(POLICY_ID), - "parameters": { - "gov.irs.credits.ctc.amount.base[0].amount": { - "2026-01-01.2100-12-31": 3000 - } - }, - }, + "policy_id": str(POLICY_ID), "options": { "axes": [[{**AXES[0][0], "period": "2026"}]], "add_missing_variables": True, @@ -118,31 +105,27 @@ def test_a_calculation_sends_the_stored_household_policy_and_catalog() -> None: assert catalog_reader.closed -def test_current_law_and_no_saved_choice_send_neither() -> None: +def test_current_law_sends_no_policy() -> None: parts = Parts() - stored = household_read() - data = {key: value for key, value in stored.household_data.items() if key != "spm"} - parts.households.households[("us", HOUSEHOLD_ID)] = stored.__class__( - **{**stored.__dict__, "household_data": data} - ) parts.service.submit_calculation(calculation_input(policy_id=None)) (submission,) = parts.simulation.submissions - assert "spm" not in submission - assert submission["policy"] == {"policy_id": None, "parameters": {}} + assert submission["policy_id"] is None assert submission["options"] == {"axes": None, "add_missing_variables": False} -def test_the_year_is_the_households_default_year() -> None: +def test_an_axis_without_a_period_varies_the_households_default_year() -> None: parts = Parts() parts.households.households[("us", HOUSEHOLD_ID)] = household_read( default_year=2027 ) - parts.service.submit_calculation(calculation_input()) + parts.service.submit_calculation(calculation_input(axes=AXES)) - assert parts.simulation.submissions[0]["year"] == 2027 + assert parts.simulation.submissions[0]["options"]["axes"] == [ + [{**AXES[0][0], "period": "2027"}] + ] def test_a_household_without_a_default_year_is_refused() -> None: @@ -202,7 +185,7 @@ def test_a_household_only_country_is_accepted() -> None: (submission,) = parts.simulation.submissions assert submission["country"] == "uk" - assert set(submission["household"]) == {"people", "households", "benunits"} + assert submission["policy_id"] is None def test_an_unknown_household_or_policy_sends_nothing() -> None: diff --git a/tests/unit/v2/test_household_calculation_simulation_client.py b/tests/unit/v2/test_household_calculation_simulation_client.py index 747b629e9..c1624f643 100644 --- a/tests/unit/v2/test_household_calculation_simulation_client.py +++ b/tests/unit/v2/test_household_calculation_simulation_client.py @@ -22,10 +22,8 @@ ) from tests.unit.v2.household_calculation_support import ( CALCULATION_ID, - CATALOG_VERSION_ID, HOUSEHOLD_ID, POLICYENGINE_VERSION, - SITUATION, simulation_calculation, succeeded_calculation, ) @@ -33,14 +31,13 @@ BASE_URL = "https://simulation.example" SUBMISSION = { "country": "us", - "catalog_version_id": str(CATALOG_VERSION_ID), "policyengine_version": POLICYENGINE_VERSION, - "year": 2026, "household_id": str(HOUSEHOLD_ID), - "household": SITUATION, - "spm": {"scenario": "ce_trend"}, - "policy": {"policy_id": None, "parameters": {}}, - "options": {"axes": None, "add_missing_variables": False}, + "policy_id": None, + "options": { + "axes": [[{"name": "employment_income", "count": 3, "min": 0, "max": 1}]], + "add_missing_variables": False, + }, "force": False, "trigger": "initial", } @@ -239,7 +236,7 @@ def auth_flow(self, request: httpx.Request): # httpx's generator protocol client.read(CALCULATION_ID) -def test_the_household_keeps_its_key_order_both_ways() -> None: +def test_the_household_keeps_its_key_order_as_answered() -> None: requests: list[httpx.Request] = [] answer = succeeded_calculation() client = client_answering(200, answer, requests) @@ -248,11 +245,7 @@ def test_the_household_keeps_its_key_order_both_ways() -> None: submitted = submit_client.submit(SUBMISSION) read = client.read(CALCULATION_ID) - sent_household = json.loads(requests[0].content)["household"] - assert list(sent_household) == list(SITUATION) - assert list(sent_household["households"]["household-1"]) == list( - SITUATION["households"]["household-1"] - ) + assert json.loads(requests[0].content) == SUBMISSION for calculation in (submitted, read): assert calculation is not None returned_household = calculation["result"]["household"] diff --git a/tests/unit/v2/test_household_calculation_transformations.py b/tests/unit/v2/test_household_calculation_transformations.py deleted file mode 100644 index 6efd6e8f7..000000000 --- a/tests/unit/v2/test_household_calculation_transformations.py +++ /dev/null @@ -1,172 +0,0 @@ -"""Stored households and policies, as the Simulation Entrypoint takes them.""" - -from __future__ import annotations - -from datetime import UTC, datetime -from uuid import uuid4 - -from policyengine_api.services.v2.household_calculations.transformations import ( - household_situation, - policy_parameters, - saved_spm_choice, -) -from policyengine_api.services.v2.policies.types import ( - PolicyParameterValueRead, - PolicyRead, -) - -US_HOUSEHOLD = { - "people": [ - { - "id": "person-1", - "source_name": "you", - "values": {"age": {"2026": 40}, "employment_income": {"2026": 30000}}, - "memberships": { - "household": "household-1", - "family": "family-1", - "tax_unit": "tax_unit-1", - "spm_unit": "spm_unit-1", - "marital_unit": "marital_unit-1", - }, - }, - { - "id": "person-2", - "values": {"age": {"2026": 8}}, - "memberships": { - "household": "household-1", - "family": "family-1", - "tax_unit": "tax_unit-1", - "spm_unit": "spm_unit-1", - "marital_unit": "marital_unit-2", - }, - }, - ], - "household": [{"id": "household-1", "values": {"state_name": {"2026": "CA"}}}], - "family": [{"id": "family-1", "values": {}}], - "tax_unit": [{"id": "tax_unit-1", "values": {"eitc": {"2026": None}}}], - "spm_unit": [{"id": "spm_unit-1", "values": {}}], - "marital_unit": [ - {"id": "marital_unit-1", "values": {}}, - {"id": "marital_unit-2", "values": {}}, - ], - "spm": {"scenario": "ce_trend", "geography_kind": "national"}, -} -UK_HOUSEHOLD = { - "people": [ - { - "id": "person-1", - "values": {"age": {"2026": 30}}, - "memberships": {"household": "household-1", "benunit": "benunit-1"}, - } - ], - "household": [{"id": "household-1", "values": {"region": {"2026": "WALES"}}}], - "benunit": [{"id": "benunit-1", "values": {"universal_credit": {"2026": None}}}], -} - - -def test_a_us_household_becomes_a_situation_keyed_by_record_id() -> None: - assert household_situation("us", US_HOUSEHOLD) == { - "people": { - "person-1": {"age": {"2026": 40}, "employment_income": {"2026": 30000}}, - "person-2": {"age": {"2026": 8}}, - }, - "households": { - "household-1": { - "members": ["person-1", "person-2"], - "state_name": {"2026": "CA"}, - } - }, - "families": {"family-1": {"members": ["person-1", "person-2"]}}, - "tax_units": { - "tax_unit-1": {"members": ["person-1", "person-2"], "eitc": {"2026": None}} - }, - "spm_units": {"spm_unit-1": {"members": ["person-1", "person-2"]}}, - "marital_units": { - "marital_unit-1": {"members": ["person-1"]}, - "marital_unit-2": {"members": ["person-2"]}, - }, - } - - -def test_a_uk_household_has_households_and_benefit_units() -> None: - assert household_situation("uk", UK_HOUSEHOLD) == { - "people": {"person-1": {"age": {"2026": 30}}}, - "households": { - "household-1": {"members": ["person-1"], "region": {"2026": "WALES"}} - }, - "benunits": { - "benunit-1": { - "members": ["person-1"], - "universal_credit": {"2026": None}, - } - }, - } - - -def test_a_group_without_members_keeps_an_empty_member_list() -> None: - household = {**UK_HOUSEHOLD, "benunit": [*UK_HOUSEHOLD["benunit"]]} - household["benunit"].append({"id": "benunit-2", "values": {}}) - - situation = household_situation("uk", household) - - assert situation["benunits"]["benunit-2"] == {"members": []} - - -def test_the_saved_spm_choice_is_the_documents_selection_or_none() -> None: - assert saved_spm_choice(US_HOUSEHOLD) == { - "scenario": "ce_trend", - "geography_kind": "national", - } - assert saved_spm_choice(UK_HOUSEHOLD) is None - - -def _value(*, name: str, value: object, start: datetime, end: datetime | None): - return PolicyParameterValueRead( - id=uuid4(), - parameter_id=uuid4(), - parameter_name=name, - value=value, - start_date=start, - end_date=end, - ) - - -def test_a_policy_becomes_parameter_values_by_period_range() -> None: - policy = PolicyRead( - id=uuid4(), - country_id="us", - tax_benefit_model_id=uuid4(), - tax_benefit_model_version_id=uuid4(), - created_at=datetime(2026, 1, 1, tzinfo=UTC), - updated_at=datetime(2026, 1, 1, tzinfo=UTC), - parameter_values=( - _value( - name="gov.irs.credits.ctc.amount.base[0].amount", - value=3000, - start=datetime(2026, 1, 1, tzinfo=UTC), - end=None, - ), - _value( - name="gov.irs.credits.ctc.refundable.fully_refundable", - value=True, - start=datetime(2026, 1, 1, tzinfo=UTC), - end=datetime(2027, 12, 31, tzinfo=UTC), - ), - _value( - name="gov.irs.credits.ctc.amount.base[0].amount", - value=3500, - start=datetime(2028, 1, 1, tzinfo=UTC), - end=None, - ), - ), - ) - - assert policy_parameters(policy) == { - "gov.irs.credits.ctc.amount.base[0].amount": { - "2026-01-01.2100-12-31": 3000, - "2028-01-01.2100-12-31": 3500, - }, - "gov.irs.credits.ctc.refundable.fully_refundable": { - "2026-01-01.2027-12-31": True - }, - } From 78219f38f33e88d793fc9f77a4429bf686b11366 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 28 Oct 2026 17:10:30 +0400 Subject: [PATCH 32/37] Send axes as given, now that the household function fills a missing period The household function now takes the run's year, the household's default year, from the simulation service and uses it for an axis without a period, so the Public API no longer fills that period itself. The stage 11 guide says the household's default year is also the year of the variables add_missing_variables adds. Co-Authored-By: Claude Opus 5.5 --- docs/migration/stage-11-v2-households.md | 11 +++-- .../v2/household_calculations/services.py | 40 ++++++++----------- .../v2/test_household_calculation_routes.py | 4 +- .../v2/test_household_calculation_service.py | 12 ++---- 4 files changed, 26 insertions(+), 41 deletions(-) diff --git a/docs/migration/stage-11-v2-households.md b/docs/migration/stage-11-v2-households.md index 0a33fa8bb..29bc2688a 100644 --- a/docs/migration/stage-11-v2-households.md +++ b/docs/migration/stage-11-v2-households.md @@ -281,12 +281,11 @@ Public API answers with 503. The answer is 202 with the calculation's `status` is `succeeded`, with `result`, or `failed`, with `error`. A US result carries `spm_config` and `spm_provenance`. -With `add_missing_variables`, the v2 household function adds each missing -variable for the year it would choose for `/calculate-full`: the first period -of the age of a person named `you`, else the current calendar year (UTC). The -Public API names people by their record IDs, so for a household whose -`default_year` is not the current year the added variables are requested for -the current year until the household function takes the calculation's year. +The simulation service sends the household function the household's +`default_year` as the calculation's year: the year of the variables +`add_missing_variables` adds, and of values and axes the household gives +without a period. (`/calculate-full` takes that year from the first age +period of a person named `you`, or else the current year.) The simulation service stores every calculation as a simulation run, so a calculation no account saved is still stored, with no report. Neither route diff --git a/policyengine_api/services/v2/household_calculations/services.py b/policyengine_api/services/v2/household_calculations/services.py index 41dcae0fa..cadc5ab0f 100644 --- a/policyengine_api/services/v2/household_calculations/services.py +++ b/policyengine_api/services/v2/household_calculations/services.py @@ -41,28 +41,6 @@ ) -def _axes_in_year( - calculation: HouseholdCalculationInput, year: int -) -> list[list[dict[str, Any]]] | None: - """Return the axes as sent, each without a period varying ``year``. - - The country model would otherwise vary its own default period, the - current year, rather than the year the household's values name. - """ - if calculation.axes is None: - return None - sent_axes = [] - for group in calculation.axes: - sent_group = [] - for axis in group: - sent_axis = axis.model_dump(mode="json", exclude_unset=True) - if sent_axis.get("period") is None: - sent_axis["period"] = str(year) - sent_group.append(sent_axis) - sent_axes.append(sent_group) - return sent_axes - - class V2HouseholdCalculationService: """Submit household calculations to the simulation service and read them.""" @@ -108,7 +86,9 @@ def submit_calculation( self._policies.get_policy( country_id=country_id, policy_id=calculation.policy_id ) - year = calculation_year(household.default_year) + # The simulation service calculates the household's default year; + # one without it is refused here, with 422. + calculation_year(household.default_year) catalog_reader = self._catalog_reader_factory() try: catalog = catalog_reader.get_model_by_country(country_id) @@ -122,7 +102,19 @@ def submit_calculation( None if calculation.policy_id is None else str(calculation.policy_id) ), "options": { - "axes": _axes_in_year(calculation, year), + # The household function varies the run's year, the + # household's default year, along an axis without a period. + "axes": ( + None + if calculation.axes is None + else [ + [ + axis.model_dump(mode="json", exclude_unset=True) + for axis in group + ] + for group in calculation.axes + ] + ), "add_missing_variables": calculation.add_missing_variables, }, # A calculation is reused for 28 days after it succeeded; callers diff --git a/tests/unit/v2/test_household_calculation_routes.py b/tests/unit/v2/test_household_calculation_routes.py index aca622d83..56ad0b595 100644 --- a/tests/unit/v2/test_household_calculation_routes.py +++ b/tests/unit/v2/test_household_calculation_routes.py @@ -106,9 +106,7 @@ def test_a_calculation_with_axes_sends_them() -> None: response = app.client.post(PATH, json=body(axes=axes)) assert response.status_code == 202 - assert app.simulation.submissions[0]["options"]["axes"] == [ - [{**axes[0][0], "period": "2026"}] - ] + assert app.simulation.submissions[0]["options"]["axes"] == axes def test_a_succeeded_us_calculation_reads_back_its_spm_choice() -> None: diff --git a/tests/unit/v2/test_household_calculation_service.py b/tests/unit/v2/test_household_calculation_service.py index 331adb3c5..ab45ac526 100644 --- a/tests/unit/v2/test_household_calculation_service.py +++ b/tests/unit/v2/test_household_calculation_service.py @@ -92,7 +92,7 @@ def test_a_calculation_sends_the_household_and_policy_ids_and_the_version() -> N "household_id": str(HOUSEHOLD_ID), "policy_id": str(POLICY_ID), "options": { - "axes": [[{**AXES[0][0], "period": "2026"}]], + "axes": AXES, "add_missing_variables": True, }, "force": False, @@ -115,17 +115,13 @@ def test_current_law_sends_no_policy() -> None: assert submission["options"] == {"axes": None, "add_missing_variables": False} -def test_an_axis_without_a_period_varies_the_households_default_year() -> None: +def test_an_axis_without_a_period_is_sent_as_given() -> None: + """The simulation service varies the household's default year for it.""" parts = Parts() - parts.households.households[("us", HOUSEHOLD_ID)] = household_read( - default_year=2027 - ) parts.service.submit_calculation(calculation_input(axes=AXES)) - assert parts.simulation.submissions[0]["options"]["axes"] == [ - [{**AXES[0][0], "period": "2027"}] - ] + assert parts.simulation.submissions[0]["options"]["axes"] == AXES def test_a_household_without_a_default_year_is_refused() -> None: From 25427397c515383514d6bd8ca965cfd977c0e32b Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 28 Oct 2026 17:49:54 +0400 Subject: [PATCH 33/37] Pass the simulation service's 422 refusals through with their reasons The simulation service now refuses the caller's data with 422 and its reason as text: a policy changing a parameter the catalog version lacks, or a stored policy value or household the household functions cannot take, such as a group no person belongs to. The Public API read only a list of refused fields, so such a refusal became 503. It now reads the service's refusal model and answers 422 with the reason, or with each refused field. The service's other refusals (400), for references the Public API has already checked, stay 503. The route's declared 422 answer, the adapter's description and the stage 11 guide say so. Co-Authored-By: Claude Opus 5.5 --- docs/migration/stage-11-v2-households.md | 19 ++++---- .../household_calculations/response_models.py | 6 ++- .../simulation_client.py | 44 ++++++++++++------- .../v2/household_calculations/types.py | 7 ++- .../v2/test_household_calculation_routes.py | 32 +++++++++++--- ...household_calculation_simulation_client.py | 26 ++++++++--- 6 files changed, 94 insertions(+), 40 deletions(-) diff --git a/docs/migration/stage-11-v2-households.md b/docs/migration/stage-11-v2-households.md index 29bc2688a..e9968a26e 100644 --- a/docs/migration/stage-11-v2-households.md +++ b/docs/migration/stage-11-v2-households.md @@ -273,11 +273,8 @@ is not stored in that country), selects the country's running catalog version, and sends the simulation service the household's and policy's IDs, the options and that version's PolicyEngine.py version. The simulation service reads the household, the policy's parameter values and the catalog -version from these tables itself, through its own `SELECT` grant, and -refuses a policy changing a parameter the catalog version lacks, which the -Public API answers with 503. The answer is 202 with the calculation's -`calculation_id` and -`status`. Poll `GET /v2/household-calculations/{calculation_id}` until +version from these tables itself, through its own `SELECT` grant. The +answer is 202 with the calculation's `calculation_id` and `status`. Poll `GET /v2/household-calculations/{calculation_id}` until `status` is `succeeded`, with `result`, or `failed`, with `error`. A US result carries `spm_config` and `spm_provenance`. @@ -305,9 +302,15 @@ Without the URL, or with gateway auth partly configured, both routes answer 503. They also answer 503 when the simulation service cannot be reached, its gateway token cannot be fetched, it answers with an error or with a body of another shape, or it refuses the calculation (for example when no v2 worker is -deployed for the catalog version's PolicyEngine.py version). A household or -options the simulation service finds invalid, such as a group no person -belongs to, get 422 with each invalid field and why. +deployed for the catalog version's PolicyEngine.py version, or when it +cannot find a household, policy or catalog version the Public API has just +found, which means the two services disagree). The simulation service checks +every stored household and policy value against the household functions' +contract before it writes anything, and the caller's data it refuses gets +422 with its reason: a policy changing a parameter the catalog version +lacks, a stored policy value it cannot take (such as an object), or a stored +household it cannot take (such as a group no person belongs to, which +`/v2/households` accepts). ### Staging checks diff --git a/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py b/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py index d21b25968..20cf7cc0f 100644 --- a/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py +++ b/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py @@ -95,8 +95,10 @@ class HouseholdCalculationResponse(StrictHouseholdCalculationAPIModel): "description": ( "The request does not match the schema, the country does not offer " "household calculations, the household has no default year, or the " - "simulation service finds the household or its options invalid " - "(the message lists why)." + "simulation service refuses the caller's data, with its reason in " + "the message: a policy changing a parameter the catalog version " + "lacks, a stored policy value or household it cannot calculate " + "(such as a group no person belongs to), or an invalid option." ), }, 500: { diff --git a/policyengine_api/services/v2/household_calculations/simulation_client.py b/policyengine_api/services/v2/household_calculations/simulation_client.py index 6045e8281..097747d79 100644 --- a/policyengine_api/services/v2/household_calculations/simulation_client.py +++ b/policyengine_api/services/v2/household_calculations/simulation_client.py @@ -14,8 +14,13 @@ database returns it, which orders each record's keys itself, so it can differ from the order a caller sent. -A 400 is a refusal, whose reason the service gives; a 422 lists what the -service finds invalid in the household or its options; anything else, +A 400 is a refusal, whose reason the service gives (a household, policy or +catalog version it cannot find, or a household without a default year, +which the Public API checks first, or no household function for the +version); a 422 is the caller's data the service refuses, with its reason +as text (a parameter the catalog version lacks, or a stored policy value or +household the household functions cannot take) or one entry per request +field it refuses; anything else, including a 503, no answer, an answer of another shape, a gateway token that cannot be fetched or missing configuration, means the service is unavailable. @@ -36,12 +41,12 @@ from policyengine_api_simulation_client.models.household_calculation_submission import ( HouseholdCalculationSubmission, ) +from policyengine_api_simulation_client.models.household_refusal_response import ( + HouseholdRefusalResponse, +) from policyengine_api_simulation_client.models.household_route_error import ( HouseholdRouteError, ) -from policyengine_api_simulation_client.models.http_validation_error import ( - HTTPValidationError, -) from policyengine_observability import ( # type: ignore[import-untyped] # the package has no type information instrument_httpx, ) @@ -112,17 +117,22 @@ def _calculation(response: httpx.Response) -> dict[str, Any]: return body -def _invalid_request_reason(error: HTTPValidationError) -> str: - """Return each invalid field and why, without the request's ``body`` prefix.""" - reasons = [] - for item in error.detail if isinstance(error.detail, list) else []: - location = [str(part) for part in item.loc] - if location[:1] == ["body"]: - location = location[1:] - reasons.append(f"{'.'.join(location)}: {item.msg}") - return "The simulation service cannot calculate this household: " + ( - "; ".join(reasons) or "the request is invalid" - ) +def _invalid_request_reason(refusal: HouseholdRefusalResponse) -> str: + """Return the service's reason, or each refused field and why. + + A field's location is given without the request's ``body`` prefix. + """ + if isinstance(refusal.detail, str): + reason = refusal.detail + else: + reasons = [] + for item in refusal.detail: + location = [str(part) for part in item.loc] + if location[:1] == ["body"]: + location = location[1:] + reasons.append(f"{'.'.join(location)}: {item.msg}") + reason = "; ".join(reasons) or "the request is invalid" + return f"The simulation service cannot calculate this household: {reason}" def _unavailable(status_code: int) -> SimulationServiceUnavailableError: @@ -191,7 +201,7 @@ def submit(self, submission: dict[str, Any]) -> dict[str, Any]: if response.status_code == HTTPStatus.ACCEPTED: return _calculation(response) if response.status_code == HTTPStatus.UNPROCESSABLE_ENTITY: - _body, invalid = _decoded(response, HTTPValidationError.from_dict) + _body, invalid = _decoded(response, HouseholdRefusalResponse.from_dict) raise SimulationServiceInvalidRequestError(_invalid_request_reason(invalid)) if response.status_code == HTTPStatus.BAD_REQUEST: _body, refusal = _decoded(response, HouseholdRouteError.from_dict) diff --git a/policyengine_api/services/v2/household_calculations/types.py b/policyengine_api/services/v2/household_calculations/types.py index f8e33e626..d0768b4b7 100644 --- a/policyengine_api/services/v2/household_calculations/types.py +++ b/policyengine_api/services/v2/household_calculations/types.py @@ -177,9 +177,12 @@ class SimulationServiceRefusedError(RuntimeError): class SimulationServiceInvalidRequestError(ValueError): - """The simulation service finds the household or its options invalid. + """The simulation service refuses the caller's data. - The message lists each invalid field and why, as the service reports it. + Such as a policy changing a parameter the catalog version lacks, a stored + policy value or household the household functions cannot take, or an + invalid option. The message gives the service's reason, or each refused + field and why. """ diff --git a/tests/unit/v2/test_household_calculation_routes.py b/tests/unit/v2/test_household_calculation_routes.py index 56ad0b595..667cf175a 100644 --- a/tests/unit/v2/test_household_calculation_routes.py +++ b/tests/unit/v2/test_household_calculation_routes.py @@ -4,6 +4,7 @@ from uuid import uuid4 +import httpx import pytest from fastapi.testclient import TestClient from flask import Flask @@ -18,8 +19,10 @@ from policyengine_api.services.v2.household_calculations.services import ( V2HouseholdCalculationService, ) +from policyengine_api.services.v2.household_calculations.simulation_client import ( + SimulationHouseholdClient, +) from policyengine_api.services.v2.household_calculations.types import ( - SimulationServiceInvalidRequestError, SimulationServiceRefusedError, SimulationServiceUnavailableError, ) @@ -248,17 +251,34 @@ def test_a_malformed_request_is_refused_before_anything_is_read(changes) -> None assert app.simulation.submissions == [] -def test_a_household_the_simulation_service_finds_invalid_gets_422_with_why() -> None: +def test_a_stored_household_with_a_group_no_person_belongs_to_gets_422() -> None: + """The simulation service's 422 reaches the caller with its reason.""" + refusal = ( + f"the household {HOUSEHOLD_ID} cannot be calculated: " + "benunits.benunit-2.members: List should have at least 1 item after " + "validation" + ) + + def answer(request: httpx.Request) -> httpx.Response: + del request + return httpx.Response(422, json={"detail": refusal}) + app = App() - app.simulation.error = SimulationServiceInvalidRequestError( - "The simulation service cannot calculate this household: " - "household.households.household-1.members: List should have at least 1 item" + app.service = V2HouseholdCalculationService( + countries=FakeCountries(), + households=app.households, + policies=FakePolicies(), + catalog_reader_factory=FakeCatalogs, + simulation_factory=lambda: SimulationHouseholdClient( + "https://simulation.example", + httpx_args={"transport": httpx.MockTransport(answer)}, + ), ) response = app.client.post(PATH, json=body()) assert response.status_code == 422 - assert "members" in response.json()["message"] + assert refusal in response.json()["message"] def test_axis_groups_within_the_point_limit_are_accepted() -> None: diff --git a/tests/unit/v2/test_household_calculation_simulation_client.py b/tests/unit/v2/test_household_calculation_simulation_client.py index c1624f643..d132e953b 100644 --- a/tests/unit/v2/test_household_calculation_simulation_client.py +++ b/tests/unit/v2/test_household_calculation_simulation_client.py @@ -165,15 +165,31 @@ def test_the_configured_service_is_called_without_auth_when_none_is_set( assert client.base_url == BASE_URL -def test_a_household_the_service_finds_invalid_is_reported_with_its_reasons() -> None: +STORED_HOUSEHOLD_REFUSAL = ( + "the household 00000000-0000-4000-8000-0000000000a1 cannot be calculated: " + "benunits.benunit-2.members: List should have at least 1 item after validation" +) + + +def test_caller_data_the_service_refuses_is_reported_with_its_reason() -> None: + """The service refuses a stored household or policy value with 422 and text.""" + client = client_answering(422, {"detail": STORED_HOUSEHOLD_REFUSAL}, []) + + with pytest.raises( + SimulationServiceInvalidRequestError, match=r"benunit-2\.members" + ): + client.submit(SUBMISSION) + + +def test_request_fields_the_service_refuses_are_reported_one_by_one() -> None: client = client_answering( 422, { "detail": [ { - "loc": ["body", "household", "households", "h", "members"], - "msg": "List should have at least 1 item after validation", - "type": "too_short", + "loc": ["body", "options", "axes", 0, 0, "count"], + "msg": "Input should be greater than or equal to 1", + "type": "greater_than_equal", } ] }, @@ -182,7 +198,7 @@ def test_a_household_the_service_finds_invalid_is_reported_with_its_reasons() -> with pytest.raises( SimulationServiceInvalidRequestError, - match=r"household\.households\.h\.members: List should have at least 1 item", + match=r"options\.axes\.0\.0\.count: Input should be greater", ): client.submit(SUBMISSION) From 21ad2051e7a1f60b8770acb6ca1b53e5ac079d89 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 28 Oct 2026 17:50:38 +0400 Subject: [PATCH 34/37] Name the parity household's person `you`, so both sides add one year's variables The staging parity household that adds missing variables named its person `adult`, so API v1's /calculate-full added them for the current calendar year while the v2 path adds them for the household's default year, and the two would differ from 2027. It now names its person `you`, whose age gives /calculate-full the household's year, and a test with the date fixed in 2027 checks both sides use the same year. Co-Authored-By: Claude Opus 5.5 --- .../test_live_v2_household_parity.py | 10 +++-- .../test_live_household_parity_households.py | 44 +++++++++++++++++++ 2 files changed, 50 insertions(+), 4 deletions(-) create mode 100644 tests/unit/v2/test_live_household_parity_households.py diff --git a/tests/integration/test_live_v2_household_parity.py b/tests/integration/test_live_v2_household_parity.py index 758482aa8..f6042f6ca 100644 --- a/tests/integration/test_live_v2_household_parity.py +++ b/tests/integration/test_live_v2_household_parity.py @@ -27,7 +27,7 @@ CALCULATED = {YEAR: None} -def _parity_households(run_amount: int) -> list[LiveHousehold]: +def parity_households(run_amount: int) -> list[LiveHousehold]: return [ LiveHousehold( name="us-single-california", @@ -86,8 +86,10 @@ def _parity_households(run_amount: int) -> list[LiveHousehold]: LiveHousehold( name="us-single-reform-added-variables", country_id="us", + # /calculate-full adds variables for the age year of a person + # named ``you``; the v2 path for the household's default year. people={ - "adult": { + "you": { "age": {YEAR: 29}, "employment_income": {YEAR: 61_000 + run_amount}, } @@ -155,7 +157,7 @@ def _parity_households(run_amount: int) -> list[LiveHousehold]: ] -PARITY_NAMES = [household.name for household in _parity_households(0)] +PARITY_NAMES = [household.name for household in parity_households(0)] @pytest.mark.parametrize("household_name", PARITY_NAMES) @@ -164,7 +166,7 @@ def test_a_v2_household_calculation_matches_api_v1( ): (household,) = [ candidate - for candidate in _parity_households(probe_amount(integration_probe_id)) + for candidate in parity_households(probe_amount(integration_probe_id)) if candidate.name == household_name ] household_id = store_household(api_client, household) diff --git a/tests/unit/v2/test_live_household_parity_households.py b/tests/unit/v2/test_live_household_parity_households.py new file mode 100644 index 000000000..3782c329d --- /dev/null +++ b/tests/unit/v2/test_live_household_parity_households.py @@ -0,0 +1,44 @@ +"""The staging parity set's households calculate the same year on both sides. + +API v1's ``/calculate-full`` adds missing variables for the first age period +of a person named ``you``, or else the current year; the v2 household +function adds them for the stored household's default year. A parity +household that adds missing variables must therefore name its person +``you``, or the two sides differ once the calendar year moves on. +""" + +from __future__ import annotations + +from datetime import date + +from policyengine_api.services import household_calculation_service +from tests.fixtures.integration.live_household_calculations import YEAR +from tests.integration.test_live_v2_household_parity import parity_households + + +class DateIn2027(date): + """Today is a day in 2027, a year after the households' year.""" + + @classmethod + def today(cls) -> DateIn2027: + """Return 2027-06-01.""" + return cls(2027, 6, 1) + + +def test_households_that_add_variables_use_one_year_on_both_sides( + monkeypatch, +) -> None: + monkeypatch.setattr(household_calculation_service, "date", DateIn2027) + adding = [ + household + for household in parity_households(run_amount=0) + if household.add_missing_variables + ] + + assert adding + for household in adding: + v1_year = household_calculation_service.get_household_year( + household.situation() + ) + # The v2 side calculates the stored household's default year, YEAR. + assert str(v1_year) == YEAR, household.name From 5184700273e2c05ad9e3af8d1254c29787607b55 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 28 Oct 2026 17:53:29 +0400 Subject: [PATCH 35/37] Check that axis periods are forwarded unchanged, and name the year check A test now sends an axis with an explicit period and one with a null period and checks both reach the simulation service unchanged, as they must now that the household function fills a missing period itself. The check that a household has a default year returns nothing, so it is named require_default_year. Co-Authored-By: Claude Opus 5.5 --- .../services/v2/household_calculations/services.py | 4 ++-- .../services/v2/household_calculations/validators.py | 8 +++++--- tests/unit/v2/test_household_calculation_service.py | 12 ++++++++++++ 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/policyengine_api/services/v2/household_calculations/services.py b/policyengine_api/services/v2/household_calculations/services.py index cadc5ab0f..7a20d30d6 100644 --- a/policyengine_api/services/v2/household_calculations/services.py +++ b/policyengine_api/services/v2/household_calculations/services.py @@ -23,7 +23,7 @@ ) from policyengine_api.services.v2.household_calculations.validators import ( HouseholdCalculationNotFoundError, - calculation_year, + require_default_year, require_household_country, ) @@ -88,7 +88,7 @@ def submit_calculation( ) # The simulation service calculates the household's default year; # one without it is refused here, with 422. - calculation_year(household.default_year) + require_default_year(household.default_year) catalog_reader = self._catalog_reader_factory() try: catalog = catalog_reader.get_model_by_country(country_id) diff --git a/policyengine_api/services/v2/household_calculations/validators.py b/policyengine_api/services/v2/household_calculations/validators.py index 5475e3328..9eb7dcd93 100644 --- a/policyengine_api/services/v2/household_calculations/validators.py +++ b/policyengine_api/services/v2/household_calculations/validators.py @@ -37,9 +37,11 @@ def require_household_country( raise HouseholdCalculationCountryError(msg) -def calculation_year(default_year: int | None) -> int: - """Return the household's default year, the year its values name.""" +def require_default_year(default_year: int | None) -> None: + """Raise HouseholdCalculationYearError unless the household has a default year. + + The simulation service calculates that year, the year its values name. + """ if default_year is None: msg = "This household has no default year to calculate." raise HouseholdCalculationYearError(msg) - return default_year diff --git a/tests/unit/v2/test_household_calculation_service.py b/tests/unit/v2/test_household_calculation_service.py index ab45ac526..f443cca0f 100644 --- a/tests/unit/v2/test_household_calculation_service.py +++ b/tests/unit/v2/test_household_calculation_service.py @@ -300,3 +300,15 @@ def test_an_axis_keeps_its_own_period() -> None: parts.service.submit_calculation(calculation_input(axes=axes)) assert parts.simulation.submissions[0]["options"]["axes"] == axes + + +@pytest.mark.parametrize( + "period", ["2025", None], ids=["explicit-period", "null-period"] +) +def test_an_axis_period_is_forwarded_unchanged(period) -> None: + parts = Parts() + axes = [[{**AXES[0][0], "period": period}]] + + parts.service.submit_calculation(calculation_input(axes=axes)) + + assert parts.simulation.submissions[0]["options"]["axes"] == axes From 109d62630c91cf460ad1ffc8d2e02f1d0d0b4391 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 29 Oct 2026 02:22:46 +0400 Subject: [PATCH 36/37] Answer 503 for a simulation refusal of another shape The generated refusal model keeps a 422's detail as it came when it is neither text nor a list of field errors, such as null, an object, a number, a list of strings or an entry without a type, and building the reason from it raised, so the route answered 500. Such a refusal now means the simulation service's answer could not be read, and the route answers 503 as documented. Co-Authored-By: Claude Opus 5.5 --- .../simulation_client.py | 22 ++++++++++--- .../v2/test_household_calculation_routes.py | 33 +++++++++++++++++++ ...household_calculation_simulation_client.py | 18 ++++++++++ 3 files changed, 69 insertions(+), 4 deletions(-) diff --git a/policyengine_api/services/v2/household_calculations/simulation_client.py b/policyengine_api/services/v2/household_calculations/simulation_client.py index 097747d79..6ff1689a7 100644 --- a/policyengine_api/services/v2/household_calculations/simulation_client.py +++ b/policyengine_api/services/v2/household_calculations/simulation_client.py @@ -44,6 +44,9 @@ from policyengine_api_simulation_client.models.household_refusal_response import ( HouseholdRefusalResponse, ) +from policyengine_api_simulation_client.models.household_request_field_error import ( + HouseholdRequestFieldError, +) from policyengine_api_simulation_client.models.household_route_error import ( HouseholdRouteError, ) @@ -120,13 +123,24 @@ def _calculation(response: httpx.Response) -> dict[str, Any]: def _invalid_request_reason(refusal: HouseholdRefusalResponse) -> str: """Return the service's reason, or each refused field and why. - A field's location is given without the request's ``body`` prefix. + A field's location is given without the request's ``body`` prefix. Raise + SimulationServiceUnavailableError for a refusal of another shape: the + generated model keeps a ``detail`` it cannot read as text or fields. """ - if isinstance(refusal.detail, str): - reason = refusal.detail + detail: object = refusal.detail + if isinstance(detail, str): + reason = detail else: + fields = ( + [item for item in detail if isinstance(item, HouseholdRequestFieldError)] + if isinstance(detail, list) + else [] + ) + if not isinstance(detail, list) or len(fields) != len(detail): + msg = "The simulation service's answer could not be read." + raise SimulationServiceUnavailableError(msg) reasons = [] - for item in refusal.detail: + for item in fields: location = [str(part) for part in item.loc] if location[:1] == ["body"]: location = location[1:] diff --git a/tests/unit/v2/test_household_calculation_routes.py b/tests/unit/v2/test_household_calculation_routes.py index 667cf175a..2bbd63642 100644 --- a/tests/unit/v2/test_household_calculation_routes.py +++ b/tests/unit/v2/test_household_calculation_routes.py @@ -294,3 +294,36 @@ def test_axis_groups_within_the_point_limit_are_accepted() -> None: response = app.client.post(PATH, json=body(axes=axes)) assert response.status_code == 202 + + +@pytest.mark.parametrize( + "refusal_body", + [ + {"detail": None}, + {"detail": {"reason": "x"}}, + {"detail": 7}, + {"detail": ["the household is refused"]}, + {"detail": [{"loc": ["body", "options"], "msg": "Field required"}]}, + ], + ids=["null", "object", "number", "list-of-strings", "entry-without-type"], +) +def test_a_simulation_refusal_of_another_shape_answers_503(refusal_body) -> None: + def answer(request: httpx.Request) -> httpx.Response: + del request + return httpx.Response(422, json=refusal_body) + + app = App() + app.service = V2HouseholdCalculationService( + countries=FakeCountries(), + households=app.households, + policies=FakePolicies(), + catalog_reader_factory=FakeCatalogs, + simulation_factory=lambda: SimulationHouseholdClient( + "https://simulation.example", + httpx_args={"transport": httpx.MockTransport(answer)}, + ), + ) + + response = app.client.post(PATH, json=body()) + + assert response.status_code == 503 diff --git a/tests/unit/v2/test_household_calculation_simulation_client.py b/tests/unit/v2/test_household_calculation_simulation_client.py index d132e953b..8d84b8c8f 100644 --- a/tests/unit/v2/test_household_calculation_simulation_client.py +++ b/tests/unit/v2/test_household_calculation_simulation_client.py @@ -269,3 +269,21 @@ def test_the_household_keeps_its_key_order_as_answered() -> None: assert list(returned_household["households"]["household-1"]) == list( answer["result"]["household"]["households"]["household-1"] ) + + +@pytest.mark.parametrize( + "body", + [ + {"detail": None}, + {"detail": {"reason": "x"}}, + {"detail": 7}, + {"detail": ["the household is refused"]}, + {"detail": [{"loc": ["body", "options"], "msg": "Field required"}]}, + ], + ids=["null", "object", "number", "list-of-strings", "entry-without-type"], +) +def test_a_refusal_of_another_shape_is_unavailable(body) -> None: + client = client_answering(422, body, []) + + with pytest.raises(SimulationServiceUnavailableError, match="could not be read"): + client.submit(SUBMISSION) From 4c515e2e346a51eaee65c5ee59467bd1c6d6baa6 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 29 Oct 2026 02:25:49 +0400 Subject: [PATCH 37/37] Describe the remaining 503s, the per-field 422, and test the sent default year The declared 503 answer now lists what is left for it (an unavailable dependency, an unreadable answer, a refused reference the Public API has already checked, or no v2 worker), and the stage 11 guide says a refused request field gets 422 with each field. The parity year test checks the default year the live check stores the household with, not a constant, and the axis-period tests are one parametrized test. Co-Authored-By: Claude Opus 5.5 --- docs/migration/stage-11-v2-households.md | 3 ++- .../household_calculations/response_models.py | 5 ++-- .../v2/test_household_calculation_service.py | 13 +++------- .../test_live_household_parity_households.py | 26 ++++++++++++++++--- 4 files changed, 31 insertions(+), 16 deletions(-) diff --git a/docs/migration/stage-11-v2-households.md b/docs/migration/stage-11-v2-households.md index e9968a26e..701673f0b 100644 --- a/docs/migration/stage-11-v2-households.md +++ b/docs/migration/stage-11-v2-households.md @@ -310,7 +310,8 @@ contract before it writes anything, and the caller's data it refuses gets 422 with its reason: a policy changing a parameter the catalog version lacks, a stored policy value it cannot take (such as an object), or a stored household it cannot take (such as a group no person belongs to, which -`/v2/households` accepts). +`/v2/households` accepts). A request field it refuses, such as an invalid +option, gets 422 with each refused field and why. ### Staging checks diff --git a/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py b/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py index 20cf7cc0f..d1705fab7 100644 --- a/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py +++ b/policyengine_api/fastapi_routes/v2/household_calculations/response_models.py @@ -109,8 +109,9 @@ class HouseholdCalculationResponse(StrictHouseholdCalculationAPIModel): "model": V2ErrorResponse, "description": ( "The simulation service, the country's catalog version or v2 " - "storage is unavailable, or the simulation service cannot " - "calculate this household." + "storage is unavailable; the simulation service's answer cannot " + "be read; or it refuses a reference the Public API has already " + "checked, or has no v2 worker for the catalog version." ), }, } diff --git a/tests/unit/v2/test_household_calculation_service.py b/tests/unit/v2/test_household_calculation_service.py index f443cca0f..c87c820db 100644 --- a/tests/unit/v2/test_household_calculation_service.py +++ b/tests/unit/v2/test_household_calculation_service.py @@ -293,17 +293,10 @@ def test_requests_refused_by_the_public_api_need_no_simulation_service() -> None available_country.submit_calculation(calculation_input()) -def test_an_axis_keeps_its_own_period() -> None: - parts = Parts() - axes = [[{**AXES[0][0], "period": "2025-03"}]] - - parts.service.submit_calculation(calculation_input(axes=axes)) - - assert parts.simulation.submissions[0]["options"]["axes"] == axes - - @pytest.mark.parametrize( - "period", ["2025", None], ids=["explicit-period", "null-period"] + "period", + ["2025", "2025-03", None], + ids=["explicit-year", "explicit-month", "null-period"], ) def test_an_axis_period_is_forwarded_unchanged(period) -> None: parts = Parts() diff --git a/tests/unit/v2/test_live_household_parity_households.py b/tests/unit/v2/test_live_household_parity_households.py index 3782c329d..d66ebd6ea 100644 --- a/tests/unit/v2/test_live_household_parity_households.py +++ b/tests/unit/v2/test_live_household_parity_households.py @@ -9,10 +9,13 @@ from __future__ import annotations +import json from datetime import date +import httpx + from policyengine_api.services import household_calculation_service -from tests.fixtures.integration.live_household_calculations import YEAR +from tests.fixtures.integration.live_household_calculations import store_household from tests.integration.test_live_v2_household_parity import parity_households @@ -40,5 +43,22 @@ def test_households_that_add_variables_use_one_year_on_both_sides( v1_year = household_calculation_service.get_household_year( household.situation() ) - # The v2 side calculates the stored household's default year, YEAR. - assert str(v1_year) == YEAR, household.name + # The v2 side calculates the default year the household is stored with. + assert str(v1_year) == str(_stored_default_year(household)), household.name + + +def _stored_default_year(household) -> int: + """Return the default year the live check stores the household with.""" + sent: list[dict] = [] + + def answer(request: httpx.Request) -> httpx.Response: + sent.append(json.loads(request.content)) + return httpx.Response(201, json={"result": {"item": {"id": "stored"}}}) + + with httpx.Client( + base_url="https://public-api.example", + transport=httpx.MockTransport(answer), + ) as client: + store_household(client, household) + (request,) = sent + return request["default_year"]