From 5656209f638dc1e66bfd9cafb57bf39fa704d769 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 19 Feb 2026 23:01:48 +0100 Subject: [PATCH 1/8] feat: Add poverty_type field to Poverty output class Harmonize poverty type tracking between policyengine.py and the API by adding poverty_type directly to the Poverty class. Convenience functions now iterate .items() on the poverty variable dicts to capture both the type enum and variable name, and include poverty_type in DataFrame output. Co-Authored-By: Claude Opus 4.6 --- src/policyengine/outputs/poverty.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/policyengine/outputs/poverty.py b/src/policyengine/outputs/poverty.py index 4955b814..25ed9c79 100644 --- a/src/policyengine/outputs/poverty.py +++ b/src/policyengine/outputs/poverty.py @@ -51,6 +51,7 @@ class Poverty(Output): simulation: Simulation poverty_variable: str + poverty_type: str | None = None entity: str = "person" # Optional demographic filters @@ -151,10 +152,11 @@ def calculate_uk_poverty_rates( """ results = [] - for poverty_variable in UK_POVERTY_VARIABLES.values(): + for poverty_type, poverty_variable in UK_POVERTY_VARIABLES.items(): poverty = Poverty( simulation=simulation, poverty_variable=poverty_variable, + poverty_type=str(poverty_type), entity="person", filter_variable=filter_variable, filter_variable_eq=filter_variable_eq, @@ -168,6 +170,7 @@ def calculate_uk_poverty_rates( [ { "simulation_id": r.simulation.id, + "poverty_type": r.poverty_type, "poverty_variable": r.poverty_variable, "filter_variable": r.filter_variable, "filter_variable_eq": r.filter_variable_eq, @@ -205,10 +208,11 @@ def calculate_us_poverty_rates( """ results = [] - for poverty_variable in US_POVERTY_VARIABLES.values(): + for poverty_type, poverty_variable in US_POVERTY_VARIABLES.items(): poverty = Poverty( simulation=simulation, poverty_variable=poverty_variable, + poverty_type=str(poverty_type), entity="person", filter_variable=filter_variable, filter_variable_eq=filter_variable_eq, @@ -222,6 +226,7 @@ def calculate_us_poverty_rates( [ { "simulation_id": r.simulation.id, + "poverty_type": r.poverty_type, "poverty_variable": r.poverty_variable, "filter_variable": r.filter_variable, "filter_variable_eq": r.filter_variable_eq, From 0c51d3418f063f6424cc405a331efee27fb79343 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 20 Feb 2026 00:18:07 +0100 Subject: [PATCH 2/8] feat: Add poverty by age group computation functions Add AGE_GROUPS dict and calculate_uk/us_poverty_by_age() convenience functions that compute poverty rates for child (<18), adult (18-64), and senior (65+) age groups across all poverty types. Co-Authored-By: Claude Opus 4.6 --- src/policyengine/outputs/poverty.py | 88 +++++++++++++++++++++++++++++ 1 file changed, 88 insertions(+) diff --git a/src/policyengine/outputs/poverty.py b/src/policyengine/outputs/poverty.py index 25ed9c79..a2090921 100644 --- a/src/policyengine/outputs/poverty.py +++ b/src/policyengine/outputs/poverty.py @@ -241,3 +241,91 @@ def calculate_us_poverty_rates( ) return OutputCollection(outputs=results, dataframe=df) + + +# Age group definitions (same for UK and US) +AGE_GROUPS = { + "child": {"filter_variable": "age", "filter_variable_leq": 17}, + "adult": { + "filter_variable": "age", + "filter_variable_geq": 18, + "filter_variable_leq": 64, + }, + "senior": {"filter_variable": "age", "filter_variable_geq": 65}, +} + + +def calculate_uk_poverty_by_age( + simulation: Simulation, +) -> OutputCollection[Poverty]: + """Calculate UK poverty rates broken down by age group. + + Computes poverty rates for child (< 18), adult (18-64), and + senior (65+) groups across all UK poverty types. + + Returns: + OutputCollection containing Poverty objects for each + age group x poverty type combination (3 x 4 = 12 records). + """ + results = [] + + for group_name, filters in AGE_GROUPS.items(): + group_results = calculate_uk_poverty_rates(simulation, **filters) + for pov in group_results.outputs: + pov.filter_variable = group_name + results.append(pov) + + df = pd.DataFrame( + [ + { + "simulation_id": r.simulation.id, + "poverty_type": r.poverty_type, + "poverty_variable": r.poverty_variable, + "filter_variable": r.filter_variable, + "headcount": r.headcount, + "total_population": r.total_population, + "rate": r.rate, + } + for r in results + ] + ) + + return OutputCollection(outputs=results, dataframe=df) + + +def calculate_us_poverty_by_age( + simulation: Simulation, +) -> OutputCollection[Poverty]: + """Calculate US poverty rates broken down by age group. + + Computes poverty rates for child (< 18), adult (18-64), and + senior (65+) groups across all US poverty types. + + Returns: + OutputCollection containing Poverty objects for each + age group x poverty type combination (3 x 2 = 6 records). + """ + results = [] + + for group_name, filters in AGE_GROUPS.items(): + group_results = calculate_us_poverty_rates(simulation, **filters) + for pov in group_results.outputs: + pov.filter_variable = group_name + results.append(pov) + + df = pd.DataFrame( + [ + { + "simulation_id": r.simulation.id, + "poverty_type": r.poverty_type, + "poverty_variable": r.poverty_variable, + "filter_variable": r.filter_variable, + "headcount": r.headcount, + "total_population": r.total_population, + "rate": r.rate, + } + for r in results + ] + ) + + return OutputCollection(outputs=results, dataframe=df) From 6e6f25ba6d52f4ac6a163db62d8accbc09efe587 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 20 Feb 2026 00:45:07 +0100 Subject: [PATCH 3/8] feat: Add gender poverty computation functions Add is_male to both UK and US entity_variables so it's available in simulation output datasets. Add GENDER_GROUPS dict and calculate_uk/us_poverty_by_gender() convenience functions that compute poverty rates for male and female groups across all poverty types. Co-Authored-By: Claude Opus 4.6 --- src/policyengine/outputs/poverty.py | 82 +++++++++++++++++++ .../tax_benefit_models/uk/model.py | 1 + .../tax_benefit_models/us/model.py | 1 + 3 files changed, 84 insertions(+) diff --git a/src/policyengine/outputs/poverty.py b/src/policyengine/outputs/poverty.py index a2090921..5e0fe300 100644 --- a/src/policyengine/outputs/poverty.py +++ b/src/policyengine/outputs/poverty.py @@ -243,6 +243,12 @@ def calculate_us_poverty_rates( return OutputCollection(outputs=results, dataframe=df) +# Gender group definitions (same for UK and US — both use is_male boolean) +GENDER_GROUPS = { + "male": {"filter_variable": "is_male", "filter_variable_eq": True}, + "female": {"filter_variable": "is_male", "filter_variable_eq": False}, +} + # Age group definitions (same for UK and US) AGE_GROUPS = { "child": {"filter_variable": "age", "filter_variable_leq": 17}, @@ -329,3 +335,79 @@ def calculate_us_poverty_by_age( ) return OutputCollection(outputs=results, dataframe=df) + + +def calculate_uk_poverty_by_gender( + simulation: Simulation, +) -> OutputCollection[Poverty]: + """Calculate UK poverty rates broken down by gender. + + Computes poverty rates for male and female groups across + all UK poverty types using the is_male boolean variable. + + Returns: + OutputCollection containing Poverty objects for each + gender x poverty type combination (2 x 4 = 8 records). + """ + results = [] + + for group_name, filters in GENDER_GROUPS.items(): + group_results = calculate_uk_poverty_rates(simulation, **filters) + for pov in group_results.outputs: + pov.filter_variable = group_name + results.append(pov) + + df = pd.DataFrame( + [ + { + "simulation_id": r.simulation.id, + "poverty_type": r.poverty_type, + "poverty_variable": r.poverty_variable, + "filter_variable": r.filter_variable, + "headcount": r.headcount, + "total_population": r.total_population, + "rate": r.rate, + } + for r in results + ] + ) + + return OutputCollection(outputs=results, dataframe=df) + + +def calculate_us_poverty_by_gender( + simulation: Simulation, +) -> OutputCollection[Poverty]: + """Calculate US poverty rates broken down by gender. + + Computes poverty rates for male and female groups across + all US poverty types using the is_male boolean variable. + + Returns: + OutputCollection containing Poverty objects for each + gender x poverty type combination (2 x 2 = 4 records). + """ + results = [] + + for group_name, filters in GENDER_GROUPS.items(): + group_results = calculate_us_poverty_rates(simulation, **filters) + for pov in group_results.outputs: + pov.filter_variable = group_name + results.append(pov) + + df = pd.DataFrame( + [ + { + "simulation_id": r.simulation.id, + "poverty_type": r.poverty_type, + "poverty_variable": r.poverty_variable, + "filter_variable": r.filter_variable, + "headcount": r.headcount, + "total_population": r.total_population, + "rate": r.rate, + } + for r in results + ] + ) + + return OutputCollection(outputs=results, dataframe=df) diff --git a/src/policyengine/tax_benefit_models/uk/model.py b/src/policyengine/tax_benefit_models/uk/model.py index 88ead217..042a5916 100644 --- a/src/policyengine/tax_benefit_models/uk/model.py +++ b/src/policyengine/tax_benefit_models/uk/model.py @@ -62,6 +62,7 @@ class PolicyEngineUKLatest(TaxBenefitModelVersion): # Demographics "age", "gender", + "is_male", "is_adult", "is_SP_age", "is_child", diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index 3c8a5aae..30b6fedb 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -65,6 +65,7 @@ class PolicyEngineUSLatest(TaxBenefitModelVersion): "person_weight", # Demographics "age", + "is_male", "is_child", "is_adult", # Income From ce28ee93593bb19b5368382381afa66aa7a7aeb5 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 20 Feb 2026 00:56:39 +0100 Subject: [PATCH 4/8] feat: Add racial poverty computation function (US only) Add race to US entity_variables so it's available in simulation output datasets. Add RACE_GROUPS dict and calculate_us_poverty_by_race() that computes poverty rates for white, black, hispanic, and other racial groups across all US poverty types. Co-Authored-By: Claude Opus 4.6 --- src/policyengine/outputs/poverty.py | 49 +++++++++++++++++++ .../tax_benefit_models/us/model.py | 1 + 2 files changed, 50 insertions(+) diff --git a/src/policyengine/outputs/poverty.py b/src/policyengine/outputs/poverty.py index 5e0fe300..c874d1c1 100644 --- a/src/policyengine/outputs/poverty.py +++ b/src/policyengine/outputs/poverty.py @@ -243,6 +243,14 @@ def calculate_us_poverty_rates( return OutputCollection(outputs=results, dataframe=df) +# Race group definitions (US only — race Enum stored as string names) +RACE_GROUPS = { + "white": {"filter_variable": "race", "filter_variable_eq": "WHITE"}, + "black": {"filter_variable": "race", "filter_variable_eq": "BLACK"}, + "hispanic": {"filter_variable": "race", "filter_variable_eq": "HISPANIC"}, + "other": {"filter_variable": "race", "filter_variable_eq": "OTHER"}, +} + # Gender group definitions (same for UK and US — both use is_male boolean) GENDER_GROUPS = { "male": {"filter_variable": "is_male", "filter_variable_eq": True}, @@ -411,3 +419,44 @@ def calculate_us_poverty_by_gender( ) return OutputCollection(outputs=results, dataframe=df) + + +def calculate_us_poverty_by_race( + simulation: Simulation, +) -> OutputCollection[Poverty]: + """Calculate US poverty rates broken down by race. + + Computes poverty rates for white, black, hispanic, and other + racial groups across all US poverty types using the race Enum + variable (stored as string names in the output dataset). + + US-only — the UK does not have a race variable. + + Returns: + OutputCollection containing Poverty objects for each + race x poverty type combination (4 x 2 = 8 records). + """ + results = [] + + for group_name, filters in RACE_GROUPS.items(): + group_results = calculate_us_poverty_rates(simulation, **filters) + for pov in group_results.outputs: + pov.filter_variable = group_name + results.append(pov) + + df = pd.DataFrame( + [ + { + "simulation_id": r.simulation.id, + "poverty_type": r.poverty_type, + "poverty_variable": r.poverty_variable, + "filter_variable": r.filter_variable, + "headcount": r.headcount, + "total_population": r.total_population, + "rate": r.rate, + } + for r in results + ] + ) + + return OutputCollection(outputs=results, dataframe=df) diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index 30b6fedb..d75c760f 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -66,6 +66,7 @@ class PolicyEngineUSLatest(TaxBenefitModelVersion): # Demographics "age", "is_male", + "race", "is_child", "is_adult", # Income From d5aa56cdcc4ab7baab418ef567214fd4e39258d3 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 20 Feb 2026 02:07:05 +0100 Subject: [PATCH 5/8] feat: Add household_state_income_tax to US tax_unit entity_variables MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Needed for budget summary computation — captures state-level tax revenue impact (matching V1's state_tax_revenue_impact field). Co-Authored-By: Claude Opus 4.6 --- src/policyengine/tax_benefit_models/us/model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index d75c760f..d76cddef 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -100,6 +100,7 @@ class PolicyEngineUSLatest(TaxBenefitModelVersion): "tax_unit_weight", "income_tax", "employee_payroll_tax", + "household_state_income_tax", "eitc", "ctc", ], From 9651de863d1ec79b5148ed4ee716442df78ebf93 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 20 Feb 2026 16:47:20 +0100 Subject: [PATCH 6/8] feat: Add household_income_decile and household_count_people to entity_variables Both UK and US models now include these household-level variables, needed for intra-decile income change distribution computation. Co-Authored-By: Claude Opus 4.6 --- src/policyengine/tax_benefit_models/uk/model.py | 2 ++ src/policyengine/tax_benefit_models/us/model.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/src/policyengine/tax_benefit_models/uk/model.py b/src/policyengine/tax_benefit_models/uk/model.py index 042a5916..cb2337e0 100644 --- a/src/policyengine/tax_benefit_models/uk/model.py +++ b/src/policyengine/tax_benefit_models/uk/model.py @@ -103,8 +103,10 @@ class PolicyEngineUKLatest(TaxBenefitModelVersion): # IDs and weights "household_id", "household_weight", + "household_count_people", # Income measures "household_net_income", + "household_income_decile", "hbai_household_net_income", "equiv_hbai_household_net_income", "household_market_income", diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index d76cddef..7eddf544 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -107,7 +107,9 @@ class PolicyEngineUSLatest(TaxBenefitModelVersion): "household": [ "household_id", "household_weight", + "household_count_people", "household_net_income", + "household_income_decile", "household_benefits", "household_tax", "household_market_income", From 46601a2a7463c09c0784ac5086f059e121e72ff6 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 20 Feb 2026 19:17:38 +0100 Subject: [PATCH 7/8] test: Add tests for poverty-by-demographics convenience functions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 16 tests covering calculate_*_poverty_by_age, calculate_*_poverty_by_gender, and calculate_us_poverty_by_race — verifying delegation, record counts, filter_variable assignment, and correct filter kwargs. Co-Authored-By: Claude Opus 4.6 --- .../poverty_by_demographics_fixtures.py | 110 +++++++ tests/test_poverty_by_demographics.py | 294 ++++++++++++++++++ 2 files changed, 404 insertions(+) create mode 100644 tests/fixtures/poverty_by_demographics_fixtures.py create mode 100644 tests/test_poverty_by_demographics.py diff --git a/tests/fixtures/poverty_by_demographics_fixtures.py b/tests/fixtures/poverty_by_demographics_fixtures.py new file mode 100644 index 00000000..c19d453d --- /dev/null +++ b/tests/fixtures/poverty_by_demographics_fixtures.py @@ -0,0 +1,110 @@ +"""Fixtures for poverty-by-demographics convenience function tests. + +These tests verify the delegation logic in calculate_*_poverty_by_age, +calculate_*_poverty_by_gender, and calculate_us_poverty_by_race. The +heavy lifting (actual poverty computation) is covered by test_poverty.py; +these fixtures mock the base poverty functions so we only test the +iteration, grouping, and filter_variable assignment. +""" + +from unittest.mock import MagicMock + +import pandas as pd + +from policyengine.core import OutputCollection +from policyengine.outputs.poverty import ( + AGE_GROUPS, + GENDER_GROUPS, + RACE_GROUPS, + UK_POVERTY_VARIABLES, + US_POVERTY_VARIABLES, + Poverty, +) + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +UK_POVERTY_TYPE_COUNT = len(UK_POVERTY_VARIABLES) # 4 +US_POVERTY_TYPE_COUNT = len(US_POVERTY_VARIABLES) # 2 + +AGE_GROUP_COUNT = len(AGE_GROUPS) # 3 +GENDER_GROUP_COUNT = len(GENDER_GROUPS) # 2 +RACE_GROUP_COUNT = len(RACE_GROUPS) # 4 + +AGE_GROUP_NAMES = list(AGE_GROUPS.keys()) # ["child", "adult", "senior"] +GENDER_GROUP_NAMES = list(GENDER_GROUPS.keys()) # ["male", "female"] +RACE_GROUP_NAMES = list(RACE_GROUPS.keys()) # ["white", "black", "hispanic", "other"] + +EXPECTED_UK_BY_AGE_COUNT = AGE_GROUP_COUNT * UK_POVERTY_TYPE_COUNT # 12 +EXPECTED_US_BY_AGE_COUNT = AGE_GROUP_COUNT * US_POVERTY_TYPE_COUNT # 6 +EXPECTED_UK_BY_GENDER_COUNT = GENDER_GROUP_COUNT * UK_POVERTY_TYPE_COUNT # 8 +EXPECTED_US_BY_GENDER_COUNT = GENDER_GROUP_COUNT * US_POVERTY_TYPE_COUNT # 4 +EXPECTED_US_BY_RACE_COUNT = RACE_GROUP_COUNT * US_POVERTY_TYPE_COUNT # 8 + + +# --------------------------------------------------------------------------- +# Factory functions +# --------------------------------------------------------------------------- + + +def make_mock_simulation() -> MagicMock: + """Create a minimal mock Simulation for Poverty constructor.""" + sim = MagicMock() + sim.id = "mock-sim-id" + return sim + + +def make_mock_poverty_outputs( + simulation: MagicMock, + poverty_variables: dict, +) -> list[Poverty]: + """Create mock Poverty outputs matching the given poverty variable mapping. + + Each Poverty object has headcount, total_population, and rate set to + predictable values so tests can verify the outputs are passed through + correctly. + """ + results = [] + for ptype, pvar in poverty_variables.items(): + p = MagicMock(spec=Poverty) + p.simulation = simulation + p.poverty_type = str(ptype) + p.poverty_variable = pvar + p.entity = "person" + p.filter_variable = None + p.headcount = 100.0 + p.total_population = 1000.0 + p.rate = 0.1 + results.append(p) + return results + + +def make_mock_output_collection( + outputs: list, +) -> OutputCollection: + """Wrap mock outputs in an OutputCollection with a dummy DataFrame.""" + df = pd.DataFrame( + [ + { + "poverty_type": o.poverty_type, + "headcount": o.headcount, + "total_population": o.total_population, + "rate": o.rate, + } + for o in outputs + ] + ) + return OutputCollection(outputs=outputs, dataframe=df) + + +def make_uk_mock_collection(simulation: MagicMock) -> OutputCollection: + """Create an OutputCollection mimicking calculate_uk_poverty_rates output.""" + outputs = make_mock_poverty_outputs(simulation, UK_POVERTY_VARIABLES) + return make_mock_output_collection(outputs) + + +def make_us_mock_collection(simulation: MagicMock) -> OutputCollection: + """Create an OutputCollection mimicking calculate_us_poverty_rates output.""" + outputs = make_mock_poverty_outputs(simulation, US_POVERTY_VARIABLES) + return make_mock_output_collection(outputs) diff --git a/tests/test_poverty_by_demographics.py b/tests/test_poverty_by_demographics.py new file mode 100644 index 00000000..859326c4 --- /dev/null +++ b/tests/test_poverty_by_demographics.py @@ -0,0 +1,294 @@ +"""Tests for poverty-by-demographics convenience functions. + +Tests calculate_*_poverty_by_age, calculate_*_poverty_by_gender, and +calculate_us_poverty_by_race. These are thin wrappers that iterate over +demographic groups and delegate to the base poverty rate functions, so +we mock the base functions and verify delegation logic. +""" + +from unittest.mock import patch + +from policyengine.outputs.poverty import ( + AGE_GROUPS, + GENDER_GROUPS, + RACE_GROUPS, + calculate_uk_poverty_by_age, + calculate_uk_poverty_by_gender, + calculate_us_poverty_by_age, + calculate_us_poverty_by_gender, + calculate_us_poverty_by_race, +) +from tests.fixtures.poverty_by_demographics_fixtures import ( + AGE_GROUP_NAMES, + EXPECTED_UK_BY_AGE_COUNT, + EXPECTED_UK_BY_GENDER_COUNT, + EXPECTED_US_BY_AGE_COUNT, + EXPECTED_US_BY_GENDER_COUNT, + EXPECTED_US_BY_RACE_COUNT, + GENDER_GROUP_NAMES, + RACE_GROUP_NAMES, + make_mock_simulation, + make_uk_mock_collection, + make_us_mock_collection, +) + + +# --------------------------------------------------------------------------- +# UK poverty by age +# --------------------------------------------------------------------------- + + +class TestCalculateUkPovertyByAge: + """Tests for calculate_uk_poverty_by_age.""" + + @patch("policyengine.outputs.poverty.calculate_uk_poverty_rates") + def test__given_simulation__then_returns_12_records(self, mock_rates): + # Given + sim = make_mock_simulation() + mock_rates.return_value = make_uk_mock_collection(sim) + + # When + result = calculate_uk_poverty_by_age(sim) + + # Then + assert len(result.outputs) == EXPECTED_UK_BY_AGE_COUNT + + @patch("policyengine.outputs.poverty.calculate_uk_poverty_rates") + def test__given_simulation__then_calls_base_once_per_age_group(self, mock_rates): + # Given + sim = make_mock_simulation() + mock_rates.return_value = make_uk_mock_collection(sim) + + # When + calculate_uk_poverty_by_age(sim) + + # Then + assert mock_rates.call_count == len(AGE_GROUPS) + + @patch("policyengine.outputs.poverty.calculate_uk_poverty_rates") + def test__given_simulation__then_filter_variable_set_to_group_name( + self, mock_rates + ): + # Given + sim = make_mock_simulation() + mock_rates.side_effect = lambda *a, **kw: make_uk_mock_collection(sim) + + # When + result = calculate_uk_poverty_by_age(sim) + + # Then + filter_vars = {o.filter_variable for o in result.outputs} + assert filter_vars == set(AGE_GROUP_NAMES) + + @patch("policyengine.outputs.poverty.calculate_uk_poverty_rates") + def test__given_simulation__then_passes_correct_filter_kwargs(self, mock_rates): + # Given + sim = make_mock_simulation() + mock_rates.return_value = make_uk_mock_collection(sim) + + # When + calculate_uk_poverty_by_age(sim) + + # Then — verify that the child group passes age <= 17 + calls = mock_rates.call_args_list + child_call = calls[0] # "child" is first in AGE_GROUPS + assert child_call.kwargs["filter_variable"] == "age" + assert child_call.kwargs["filter_variable_leq"] == 17 + + +# --------------------------------------------------------------------------- +# US poverty by age +# --------------------------------------------------------------------------- + + +class TestCalculateUsPovertyByAge: + """Tests for calculate_us_poverty_by_age.""" + + @patch("policyengine.outputs.poverty.calculate_us_poverty_rates") + def test__given_simulation__then_returns_6_records(self, mock_rates): + # Given + sim = make_mock_simulation() + mock_rates.return_value = make_us_mock_collection(sim) + + # When + result = calculate_us_poverty_by_age(sim) + + # Then + assert len(result.outputs) == EXPECTED_US_BY_AGE_COUNT + + @patch("policyengine.outputs.poverty.calculate_us_poverty_rates") + def test__given_simulation__then_filter_variable_set_to_group_name( + self, mock_rates + ): + # Given + sim = make_mock_simulation() + mock_rates.side_effect = lambda *a, **kw: make_us_mock_collection(sim) + + # When + result = calculate_us_poverty_by_age(sim) + + # Then + filter_vars = {o.filter_variable for o in result.outputs} + assert filter_vars == set(AGE_GROUP_NAMES) + + +# --------------------------------------------------------------------------- +# UK poverty by gender +# --------------------------------------------------------------------------- + + +class TestCalculateUkPovertyByGender: + """Tests for calculate_uk_poverty_by_gender.""" + + @patch("policyengine.outputs.poverty.calculate_uk_poverty_rates") + def test__given_simulation__then_returns_8_records(self, mock_rates): + # Given + sim = make_mock_simulation() + mock_rates.return_value = make_uk_mock_collection(sim) + + # When + result = calculate_uk_poverty_by_gender(sim) + + # Then + assert len(result.outputs) == EXPECTED_UK_BY_GENDER_COUNT + + @patch("policyengine.outputs.poverty.calculate_uk_poverty_rates") + def test__given_simulation__then_filter_variable_set_to_gender_names( + self, mock_rates + ): + # Given + sim = make_mock_simulation() + mock_rates.side_effect = lambda *a, **kw: make_uk_mock_collection(sim) + + # When + result = calculate_uk_poverty_by_gender(sim) + + # Then + filter_vars = {o.filter_variable for o in result.outputs} + assert filter_vars == set(GENDER_GROUP_NAMES) + + @patch("policyengine.outputs.poverty.calculate_uk_poverty_rates") + def test__given_simulation__then_passes_is_male_filter(self, mock_rates): + # Given + sim = make_mock_simulation() + mock_rates.return_value = make_uk_mock_collection(sim) + + # When + calculate_uk_poverty_by_gender(sim) + + # Then — first call is "male" group + male_call = mock_rates.call_args_list[0] + assert male_call.kwargs["filter_variable"] == "is_male" + assert male_call.kwargs["filter_variable_eq"] is True + + +# --------------------------------------------------------------------------- +# US poverty by gender +# --------------------------------------------------------------------------- + + +class TestCalculateUsPovertyByGender: + """Tests for calculate_us_poverty_by_gender.""" + + @patch("policyengine.outputs.poverty.calculate_us_poverty_rates") + def test__given_simulation__then_returns_4_records(self, mock_rates): + # Given + sim = make_mock_simulation() + mock_rates.return_value = make_us_mock_collection(sim) + + # When + result = calculate_us_poverty_by_gender(sim) + + # Then + assert len(result.outputs) == EXPECTED_US_BY_GENDER_COUNT + + @patch("policyengine.outputs.poverty.calculate_us_poverty_rates") + def test__given_simulation__then_filter_variable_set_to_gender_names( + self, mock_rates + ): + # Given + sim = make_mock_simulation() + mock_rates.side_effect = lambda *a, **kw: make_us_mock_collection(sim) + + # When + result = calculate_us_poverty_by_gender(sim) + + # Then + filter_vars = {o.filter_variable for o in result.outputs} + assert filter_vars == set(GENDER_GROUP_NAMES) + + +# --------------------------------------------------------------------------- +# US poverty by race +# --------------------------------------------------------------------------- + + +class TestCalculateUsPovertyByRace: + """Tests for calculate_us_poverty_by_race.""" + + @patch("policyengine.outputs.poverty.calculate_us_poverty_rates") + def test__given_simulation__then_returns_8_records(self, mock_rates): + # Given + sim = make_mock_simulation() + mock_rates.return_value = make_us_mock_collection(sim) + + # When + result = calculate_us_poverty_by_race(sim) + + # Then + assert len(result.outputs) == EXPECTED_US_BY_RACE_COUNT + + @patch("policyengine.outputs.poverty.calculate_us_poverty_rates") + def test__given_simulation__then_calls_base_once_per_race_group(self, mock_rates): + # Given + sim = make_mock_simulation() + mock_rates.return_value = make_us_mock_collection(sim) + + # When + calculate_us_poverty_by_race(sim) + + # Then + assert mock_rates.call_count == len(RACE_GROUPS) + + @patch("policyengine.outputs.poverty.calculate_us_poverty_rates") + def test__given_simulation__then_filter_variable_set_to_race_names( + self, mock_rates + ): + # Given + sim = make_mock_simulation() + mock_rates.side_effect = lambda *a, **kw: make_us_mock_collection(sim) + + # When + result = calculate_us_poverty_by_race(sim) + + # Then + filter_vars = {o.filter_variable for o in result.outputs} + assert filter_vars == set(RACE_GROUP_NAMES) + + @patch("policyengine.outputs.poverty.calculate_us_poverty_rates") + def test__given_simulation__then_passes_race_filter_with_correct_eq_value( + self, mock_rates + ): + # Given + sim = make_mock_simulation() + mock_rates.return_value = make_us_mock_collection(sim) + + # When + calculate_us_poverty_by_race(sim) + + # Then — first call is "white" group + white_call = mock_rates.call_args_list[0] + assert white_call.kwargs["filter_variable"] == "race" + assert white_call.kwargs["filter_variable_eq"] == "WHITE" + + @patch("policyengine.outputs.poverty.calculate_us_poverty_rates") + def test__given_simulation__then_dataframe_has_correct_row_count(self, mock_rates): + # Given + sim = make_mock_simulation() + mock_rates.return_value = make_us_mock_collection(sim) + + # When + result = calculate_us_poverty_by_race(sim) + + # Then + assert len(result.dataframe) == EXPECTED_US_BY_RACE_COUNT From e26e4007377d9b4c58375786a751f81230beba85 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 20 Feb 2026 19:25:41 +0100 Subject: [PATCH 8/8] style: Apply ruff and black formatting to test files Co-Authored-By: Claude Opus 4.6 --- .../poverty_by_demographics_fixtures.py | 4 +++- tests/test_poverty_by_demographics.py | 18 ++++++++++++------ 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/tests/fixtures/poverty_by_demographics_fixtures.py b/tests/fixtures/poverty_by_demographics_fixtures.py index c19d453d..aca3ceca 100644 --- a/tests/fixtures/poverty_by_demographics_fixtures.py +++ b/tests/fixtures/poverty_by_demographics_fixtures.py @@ -34,7 +34,9 @@ AGE_GROUP_NAMES = list(AGE_GROUPS.keys()) # ["child", "adult", "senior"] GENDER_GROUP_NAMES = list(GENDER_GROUPS.keys()) # ["male", "female"] -RACE_GROUP_NAMES = list(RACE_GROUPS.keys()) # ["white", "black", "hispanic", "other"] +RACE_GROUP_NAMES = list( + RACE_GROUPS.keys() +) # ["white", "black", "hispanic", "other"] EXPECTED_UK_BY_AGE_COUNT = AGE_GROUP_COUNT * UK_POVERTY_TYPE_COUNT # 12 EXPECTED_US_BY_AGE_COUNT = AGE_GROUP_COUNT * US_POVERTY_TYPE_COUNT # 6 diff --git a/tests/test_poverty_by_demographics.py b/tests/test_poverty_by_demographics.py index 859326c4..783a8085 100644 --- a/tests/test_poverty_by_demographics.py +++ b/tests/test_poverty_by_demographics.py @@ -10,7 +10,6 @@ from policyengine.outputs.poverty import ( AGE_GROUPS, - GENDER_GROUPS, RACE_GROUPS, calculate_uk_poverty_by_age, calculate_uk_poverty_by_gender, @@ -32,7 +31,6 @@ make_us_mock_collection, ) - # --------------------------------------------------------------------------- # UK poverty by age # --------------------------------------------------------------------------- @@ -54,7 +52,9 @@ def test__given_simulation__then_returns_12_records(self, mock_rates): assert len(result.outputs) == EXPECTED_UK_BY_AGE_COUNT @patch("policyengine.outputs.poverty.calculate_uk_poverty_rates") - def test__given_simulation__then_calls_base_once_per_age_group(self, mock_rates): + def test__given_simulation__then_calls_base_once_per_age_group( + self, mock_rates + ): # Given sim = make_mock_simulation() mock_rates.return_value = make_uk_mock_collection(sim) @@ -81,7 +81,9 @@ def test__given_simulation__then_filter_variable_set_to_group_name( assert filter_vars == set(AGE_GROUP_NAMES) @patch("policyengine.outputs.poverty.calculate_uk_poverty_rates") - def test__given_simulation__then_passes_correct_filter_kwargs(self, mock_rates): + def test__given_simulation__then_passes_correct_filter_kwargs( + self, mock_rates + ): # Given sim = make_mock_simulation() mock_rates.return_value = make_uk_mock_collection(sim) @@ -239,7 +241,9 @@ def test__given_simulation__then_returns_8_records(self, mock_rates): assert len(result.outputs) == EXPECTED_US_BY_RACE_COUNT @patch("policyengine.outputs.poverty.calculate_us_poverty_rates") - def test__given_simulation__then_calls_base_once_per_race_group(self, mock_rates): + def test__given_simulation__then_calls_base_once_per_race_group( + self, mock_rates + ): # Given sim = make_mock_simulation() mock_rates.return_value = make_us_mock_collection(sim) @@ -282,7 +286,9 @@ def test__given_simulation__then_passes_race_filter_with_correct_eq_value( assert white_call.kwargs["filter_variable_eq"] == "WHITE" @patch("policyengine.outputs.poverty.calculate_us_poverty_rates") - def test__given_simulation__then_dataframe_has_correct_row_count(self, mock_rates): + def test__given_simulation__then_dataframe_has_correct_row_count( + self, mock_rates + ): # Given sim = make_mock_simulation() mock_rates.return_value = make_us_mock_collection(sim)