From da51cf1870f445c8b4dfb1a44ab47debb09b92ce Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 00:08:52 -0400 Subject: [PATCH] Test the UK nation filters against policyengine-uk's country formula Since #565 the nation regions filter the stored household region column, and policyengine.py lists which regions make up each nation. policyengine-uk keeps its own mapping in the formula that derives country from region. These tests take the expected nations from that formula, for every value of policyengine-uk's Region enum, so the two copies cannot drift apart unnoticed. A Hypothesis property checks region-only data with several people per household, separately ordered tables and str or bytes regions, and each nation runs end to end through the UK model. Fixes #569 Co-Authored-By: Claude Opus 5.5 --- changelog.d/569.changed.md | 1 + tests/test_uk_nation_model_agreement.py | 222 ++++++++++++++++++++++++ 2 files changed, 223 insertions(+) create mode 100644 changelog.d/569.changed.md create mode 100644 tests/test_uk_nation_model_agreement.py diff --git a/changelog.d/569.changed.md b/changelog.d/569.changed.md new file mode 100644 index 00000000..9511d1d3 --- /dev/null +++ b/changelog.d/569.changed.md @@ -0,0 +1 @@ +Test that each UK nation filter keeps exactly the households policyengine-uk's `country` formula assigns to that nation, for every value of its `Region` enum, so the registry's region lists and the model's mapping cannot drift apart unnoticed. diff --git a/tests/test_uk_nation_model_agreement.py b/tests/test_uk_nation_model_agreement.py new file mode 100644 index 00000000..d6d228d6 --- /dev/null +++ b/tests/test_uk_nation_model_agreement.py @@ -0,0 +1,222 @@ +"""UK nation filters agree with policyengine-uk's ``country`` formula. + +The nation regions filter the stored household ``region`` column, so +policyengine.py keeps its own list of the regions in each nation +(``countries/uk/regions.py``). policyengine-uk keeps another in the formula +that derives ``country`` from ``region``. These tests take the expected +nations from policyengine-uk itself, so a change on either side fails here +instead of silently moving households between nations +(PolicyEngine/policyengine.py#569). +""" + +import pandas as pd +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st +from microdf import MicroDataFrame + +from policyengine.countries.uk.regions import UK_COUNTRIES, uk_region_registry + +pytest.importorskip("policyengine_uk") + +GROUP_ENTITIES = ["benunit", "household"] +NATIONS = [code.upper() for code in UK_COUNTRIES] + + +@pytest.fixture(scope="module") +def model_country() -> dict[str, str]: + """policyengine-uk's ``country`` for every value of its ``Region`` enum.""" + from policyengine_uk import Simulation + from policyengine_uk.variables.household.demographic.geography import Region + + regions = [region.name for region in Region] + situation = { + "people": {f"p{i}": {"age": {"2026": 40}} for i in range(len(regions))}, + "benunits": {f"b{i}": {"members": [f"p{i}"]} for i in range(len(regions))}, + "households": { + f"h{i}": {"members": [f"p{i}"], "region": {"2026": region}} + for i, region in enumerate(regions) + }, + } + simulation = Simulation(situation=situation) + # The households come back in situation order; check rather than assume. + assert [str(r) for r in simulation.calculate("region", 2026)] == regions + countries = simulation.calculate("country", 2026) + return {region: str(country) for region, country in zip(regions, countries)} + + +def _nation_filter(nation: str): + return uk_region_registry.get(f"country/{nation.lower()}").scoping_strategy + + +def _entity_data( + regions, + household_ids, + people=None, + person_order=None, + benunit_order=None, + encode=False, +) -> dict[str, pd.DataFrame]: + """Region-only UK tables, with no ``country`` column. + + Household ``h`` has benefit unit ``1000 + h`` and ``people[i]`` people + (default one) with IDs ``100 * h + k``. ``person_order`` and + ``benunit_order`` reorder those rows independently of the households; + ``encode`` stores ``region`` as bytes. + """ + household_ids = list(household_ids) + people = people or [1] * len(household_ids) + person = pd.DataFrame( + { + "person_id": [ + 100 * h + k for h, n in zip(household_ids, people) for k in range(n) + ], + "person_household_id": [ + h for h, n in zip(household_ids, people) for _ in range(n) + ], + "person_benunit_id": [ + 1000 + h for h, n in zip(household_ids, people) for _ in range(n) + ], + } + ).assign(person_weight=1.0) + benunit = pd.DataFrame({"benunit_id": [1000 + h for h in household_ids]}).assign( + benunit_weight=1.0 + ) + household = pd.DataFrame( + { + "household_id": household_ids, + "household_weight": 1.0, + "region": [r.encode() if encode else r for r in regions], + } + ) + if person_order is not None: + person = person.iloc[list(person_order)].reset_index(drop=True) + if benunit_order is not None: + benunit = benunit.iloc[list(benunit_order)].reset_index(drop=True) + return {"person": person, "benunit": benunit, "household": household} + + +def _micro(tables) -> dict[str, MicroDataFrame]: + return { + entity: MicroDataFrame(frame, weights=f"{entity}_weight") + for entity, frame in tables.items() + } + + +def test_registry_nations_are_policyengine_uk_nations(model_country): + assert set(model_country.values()) - {"UNKNOWN"} == set(NATIONS) + + +@pytest.mark.parametrize("nation", NATIONS) +def test_each_nation_keeps_the_regions_policyengine_uk_puts_in_it( + model_country, nation +): + """Exhaustive over policyengine-uk's Region enum, UNKNOWN included.""" + regions = list(model_country) + household_ids = range(1, len(regions) + 1) + + result = _nation_filter(nation).apply( + _micro(_entity_data(regions, household_ids)), GROUP_ENTITIES, 2026 + ) + + expected = [ + household_id + for household_id, region in zip(household_ids, regions) + if model_country[region] == nation + ] + assert sorted(result["household"]["household_id"]) == expected + + +@settings(max_examples=100, deadline=None) +@given(data=st.data()) +def test_nation_filters_split_any_dataset_as_policyengine_uk_does(model_country, data): + """For region-only data with 1-3 people per household, separately ordered + person, benefit-unit and household rows, and str or bytes regions: each + nation filter keeps exactly the households policyengine-uk assigns to that + nation, once each, with all their people and their benefit units. The four + are disjoint and together keep every household whose region policyengine-uk + places in a nation.""" + regions = data.draw( + st.lists(st.sampled_from(sorted(model_country)), min_size=1, max_size=20) + ) + n = len(regions) + household_ids = data.draw(st.permutations(range(1, n + 1))) + people = data.draw(st.lists(st.integers(1, 3), min_size=n, max_size=n)) + tables = _entity_data( + regions, + household_ids, + people=people, + person_order=data.draw(st.permutations(range(sum(people)))), + benunit_order=data.draw(st.permutations(range(n))), + encode=data.draw(st.booleans()), + ) + nation_of = { + household_id: model_country[region] + for household_id, region in zip(household_ids, regions) + } + person = tables["person"] + + kept = {} + for nation in NATIONS: + expected = sorted(h for h, n in nation_of.items() if n == nation) + if not expected: + with pytest.raises(ValueError, match="No households found"): + _nation_filter(nation).apply(_micro(tables), GROUP_ENTITIES, 2026) + continue + result = _nation_filter(nation).apply(_micro(tables), GROUP_ENTITIES, 2026) + kept[nation] = sorted(result["household"]["household_id"]) + assert kept[nation] == expected + assert sorted(result["person"]["person_id"]) == sorted( + person.loc[person["person_household_id"].isin(expected), "person_id"] + ) + assert sorted(result["benunit"]["benunit_id"]) == [1000 + h for h in expected] + + all_kept = [h for households in kept.values() for h in households] + assert sorted(all_kept) == sorted(h for h, n in nation_of.items() if n != "UNKNOWN") + + +@pytest.mark.parametrize("nation", NATIONS) +def test_nation_simulation_runs_on_region_only_data(model_country, nation): + """End to end through the UK model on data that stores ``region`` only.""" + import policyengine as pe + from policyengine.core import Simulation + from policyengine.tax_benefit_models.uk.datasets import ( + PolicyEngineUKDataset, + UKYearData, + ) + + regions = list(model_country) + household_ids = list(range(1, len(regions) + 1)) + tables = _entity_data(regions, household_ids) + tables["person"]["age"] = 40 + tables["household"] = tables["household"].assign( + tenure_type="RENT_PRIVATELY", rent=9_000.0, council_tax=1_500.0 + ) + micro = _micro(tables) + dataset = PolicyEngineUKDataset( + name="region-only", + description="UK tables with region and no country column", + filepath=None, + year=2026, + data=UKYearData( + person=micro["person"], + benunit=micro["benunit"], + household=micro["household"], + ), + ) + + simulation = Simulation( + dataset=dataset, + tax_benefit_model_version=pe.uk.model, + scoping_strategy=_nation_filter(nation), + extra_variables={"household": ["country"]}, + ) + simulation.run() + + household = pd.DataFrame(simulation.output_dataset.data.household) + assert sorted(household["household_id"]) == [ + household_id + for household_id, region in zip(household_ids, regions) + if model_country[region] == nation + ] + assert set(household["country"]) == {nation}