From e511914932e48d431ae2c27c42d7017de8e494dc Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 31 Dec 2025 17:58:17 +0400 Subject: [PATCH 1/4] fix: Create utils for extracting labels from brackets, breakdown params --- .../tax_benefit_models/uk/model.py | 12 +- .../tax_benefit_models/us/model.py | 12 +- src/policyengine/utils/__init__.py | 4 + src/policyengine/utils/parameter_labels.py | 107 ++++++++++++++++++ 4 files changed, 131 insertions(+), 4 deletions(-) create mode 100644 src/policyengine/utils/parameter_labels.py diff --git a/src/policyengine/tax_benefit_models/uk/model.py b/src/policyengine/tax_benefit_models/uk/model.py index f547ac59..f5eae9c5 100644 --- a/src/policyengine/tax_benefit_models/uk/model.py +++ b/src/policyengine/tax_benefit_models/uk/model.py @@ -13,6 +13,10 @@ TaxBenefitModelVersion, Variable, ) +from policyengine.utils.parameter_labels import ( + build_scale_lookup, + generate_label_for_parameter, +) from .datasets import PolicyEngineUKDataset, UKYearData @@ -146,17 +150,21 @@ def __init__(self, **kwargs: dict): from policyengine_core.parameters import Parameter as CoreParameter + scale_lookup = build_scale_lookup(system) + for param_node in system.parameters.get_descendants(): if isinstance(param_node, CoreParameter): parameter = Parameter( id=self.id + "-" + param_node.name, name=param_node.name, - label=param_node.metadata.get("label", param_node.name), + label=generate_label_for_parameter( + param_node, system, scale_lookup + ), tax_benefit_model_version=self, description=param_node.description, data_type=type(param_node(2025)), unit=param_node.metadata.get("unit"), - _core_param=param_node, # Store for lazy value loading + _core_param=param_node, ) self.add_parameter(parameter) diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index 17844feb..27760e01 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -13,6 +13,10 @@ TaxBenefitModelVersion, Variable, ) +from policyengine.utils.parameter_labels import ( + build_scale_lookup, + generate_label_for_parameter, +) from .datasets import PolicyEngineUSDataset, USYearData @@ -134,17 +138,21 @@ def __init__(self, **kwargs: dict): from policyengine_core.parameters import Parameter as CoreParameter + scale_lookup = build_scale_lookup(system) + for param_node in system.parameters.get_descendants(): if isinstance(param_node, CoreParameter): parameter = Parameter( id=self.id + "-" + param_node.name, name=param_node.name, - label=param_node.metadata.get("label"), + label=generate_label_for_parameter( + param_node, system, scale_lookup + ), tax_benefit_model_version=self, description=param_node.description, data_type=type(param_node(2025)), unit=param_node.metadata.get("unit"), - _core_param=param_node, # Store for lazy value loading + _core_param=param_node, ) self.add_parameter(parameter) diff --git a/src/policyengine/utils/__init__.py b/src/policyengine/utils/__init__.py index e73de67e..bf3cc681 100644 --- a/src/policyengine/utils/__init__.py +++ b/src/policyengine/utils/__init__.py @@ -1,3 +1,7 @@ from .dates import parse_safe_date as parse_safe_date +from .parameter_labels import build_scale_lookup as build_scale_lookup +from .parameter_labels import ( + generate_label_for_parameter as generate_label_for_parameter, +) from .plotting import COLORS as COLORS from .plotting import format_fig as format_fig diff --git a/src/policyengine/utils/parameter_labels.py b/src/policyengine/utils/parameter_labels.py new file mode 100644 index 00000000..aac76e12 --- /dev/null +++ b/src/policyengine/utils/parameter_labels.py @@ -0,0 +1,107 @@ +"""Utilities for generating human-readable labels for tax-benefit parameters.""" + +import re + + +def generate_label_for_parameter(param_node, system, scale_lookup): + """ + Generate a label for a parameter that doesn't have one. + + For breakdown parameters: Uses parent label + enum value + For bracket parameters: Uses scale label + bracket info + + Args: + param_node: The CoreParameter object + system: The tax-benefit system (has variables and parameters) + scale_lookup: Dict mapping scale names to ParameterScale objects + + Returns: + str or None: Generated label, or None if cannot generate + """ + if param_node.metadata.get("label"): + return param_node.metadata.get("label") + + param_name = param_node.name + + if "[" in param_name: + return _generate_bracket_label(param_name, scale_lookup) + + if param_node.parent and param_node.parent.metadata.get("breakdown"): + return _generate_breakdown_label(param_node, system) + + return None + + +def _generate_breakdown_label(param_node, system): + """Generate label for a breakdown parameter using enum values.""" + parent = param_node.parent + parent_label = parent.metadata.get("label") + breakdown_vars = parent.metadata.get("breakdown", []) + + if not parent_label: + return None + + child_key = param_node.name.split(".")[-1] + + for var_name in breakdown_vars: + var = system.variables.get(var_name) + if var and hasattr(var, "possible_values") and var.possible_values: + enum_class = var.possible_values + try: + enum_value = enum_class[child_key].value + return f"{parent_label} ({enum_value})" + except (KeyError, AttributeError): + continue + + return f"{parent_label} ({child_key})" + + +def _generate_bracket_label(param_name, scale_lookup): + """Generate label for a bracket parameter.""" + match = re.match(r"^(.+)\[(\d+)\]\.(\w+)$", param_name) + if not match: + return None + + scale_name = match.group(1) + bracket_index = int(match.group(2)) + field_name = match.group(3) + + scale = scale_lookup.get(scale_name) + if not scale: + return None + + scale_label = scale.metadata.get("label") + scale_type = scale.metadata.get("type", "") + + if not scale_label: + return None + + bracket_num = bracket_index + 1 + + if scale_type in ("marginal_rate", "marginal_amount"): + bracket_desc = f"bracket {bracket_num}" + elif scale_type == "single_amount": + bracket_desc = f"tier {bracket_num}" + else: + bracket_desc = f"bracket {bracket_num}" + + return f"{scale_label} ({bracket_desc} {field_name})" + + +def build_scale_lookup(system): + """ + Build a lookup dict mapping scale names to ParameterScale objects. + + Args: + system: The tax-benefit system + + Returns: + dict: Mapping of scale name -> ParameterScale object + """ + from policyengine_core.parameters import ParameterScale + + return { + p.name: p + for p in system.parameters.get_descendants() + if isinstance(p, ParameterScale) + } From adfcd135e465e10d08eeb3f659a51f8ae960d091 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 31 Dec 2025 18:20:23 +0400 Subject: [PATCH 2/4] test: Fix tests --- tests/__init__.py | 1 + tests/fixtures/__init__.py | 1 + tests/fixtures/parameter_labels_fixtures.py | 159 ++++++++ tests/test_models.py | 67 ++++ tests/test_parameter_labels.py | 417 ++++++++++++++++++++ 5 files changed, 645 insertions(+) create mode 100644 tests/__init__.py create mode 100644 tests/fixtures/__init__.py create mode 100644 tests/fixtures/parameter_labels_fixtures.py create mode 100644 tests/test_parameter_labels.py diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..2a113f28 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""PolicyEngine test suite.""" diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py new file mode 100644 index 00000000..a11f7df2 --- /dev/null +++ b/tests/fixtures/__init__.py @@ -0,0 +1 @@ +"""Test fixtures for PolicyEngine tests.""" diff --git a/tests/fixtures/parameter_labels_fixtures.py b/tests/fixtures/parameter_labels_fixtures.py new file mode 100644 index 00000000..e43263c6 --- /dev/null +++ b/tests/fixtures/parameter_labels_fixtures.py @@ -0,0 +1,159 @@ +"""Fixtures for parameter_labels utility tests.""" + +from enum import Enum +from typing import Any +from unittest.mock import MagicMock + + +class MockFilingStatus(Enum): + """Mock filing status enum for testing breakdown labels.""" + + SINGLE = "Single" + JOINT = "Joint" + HEAD_OF_HOUSEHOLD = "Head of household" + + +class MockStateCode(Enum): + """Mock state code enum where key equals value.""" + + CA = "CA" + TX = "TX" + NY = "NY" + + +def create_mock_parameter( + name: str, + label: str | None = None, + parent: Any = None, +) -> MagicMock: + """Create a mock CoreParameter object.""" + param = MagicMock() + param.name = name + param.metadata = {"label": label} if label else {} + param.parent = parent + return param + + +def create_mock_parent_node( + name: str, + label: str | None = None, + breakdown: list[str] | None = None, +) -> MagicMock: + """Create a mock parent ParameterNode with optional breakdown metadata.""" + parent = MagicMock() + parent.name = name + parent.metadata = {} + if label: + parent.metadata["label"] = label + if breakdown: + parent.metadata["breakdown"] = breakdown + return parent + + +def create_mock_scale( + name: str, + label: str | None = None, + scale_type: str | None = None, +) -> MagicMock: + """Create a mock ParameterScale object.""" + scale = MagicMock() + scale.name = name + scale.metadata = {} + if label: + scale.metadata["label"] = label + if scale_type: + scale.metadata["type"] = scale_type + return scale + + +def create_mock_variable( + name: str, + possible_values: type[Enum] | None = None, +) -> MagicMock: + """Create a mock Variable object with optional enum values.""" + var = MagicMock() + var.name = name + if possible_values: + var.possible_values = possible_values + else: + var.possible_values = None + return var + + +def create_mock_system( + variables: dict[str, MagicMock] | None = None, + scales: list[MagicMock] | None = None, +) -> MagicMock: + """Create a mock tax-benefit system.""" + system = MagicMock() + system.variables = variables or {} + + descendants = list(scales) if scales else [] + system.parameters.get_descendants.return_value = descendants + + return system + + +# Pre-built fixtures for common test scenarios + +PARAM_WITH_EXPLICIT_LABEL = create_mock_parameter( + name="gov.tax.rate", + label="Tax rate", +) + +PARAM_WITHOUT_LABEL_NO_PARENT = create_mock_parameter( + name="gov.tax.rate", + label=None, + parent=None, +) + +PARENT_WITH_BREAKDOWN_AND_LABEL = create_mock_parent_node( + name="gov.exemptions.personal", + label="Personal exemption amount", + breakdown=["filing_status"], +) + +PARENT_WITH_BREAKDOWN_NO_LABEL = create_mock_parent_node( + name="gov.exemptions.personal", + label=None, + breakdown=["filing_status"], +) + +PARENT_WITHOUT_BREAKDOWN = create_mock_parent_node( + name="gov.exemptions.personal", + label="Personal exemption amount", + breakdown=None, +) + +SCALE_WITH_LABEL_MARGINAL = create_mock_scale( + name="gov.tax.rates", + label="Income tax rate", + scale_type="marginal_rate", +) + +SCALE_WITH_LABEL_SINGLE_AMOUNT = create_mock_scale( + name="gov.tax.amounts", + label="Tax amount", + scale_type="single_amount", +) + +SCALE_WITHOUT_LABEL = create_mock_scale( + name="gov.tax.rates", + label=None, + scale_type="marginal_rate", +) + +VARIABLE_WITH_FILING_STATUS_ENUM = create_mock_variable( + name="filing_status", + possible_values=MockFilingStatus, +) + +VARIABLE_WITH_STATE_CODE_ENUM = create_mock_variable( + name="state_code", + possible_values=MockStateCode, +) + +VARIABLE_WITHOUT_ENUM = create_mock_variable( + name="age", + possible_values=None, +) diff --git a/tests/test_models.py b/tests/test_models.py index 02936f62..3a74a1b9 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,5 +1,6 @@ """Tests for UK and US tax-benefit model versions.""" +import re from policyengine.tax_benefit_models.uk import uk_latest from policyengine.tax_benefit_models.us import us_latest @@ -40,6 +41,26 @@ def test_model_version_parameter_values_aggregates_all(self): all_values = uk_latest.parameter_values assert len(all_values) >= 100 + def test__given_bracket_parameter__then_has_generated_label(self): + """Bracket parameters should have auto-generated labels.""" + bracket_params_with_labels = [ + p + for p in uk_latest.parameters + if "[" in p.name and p.label and "bracket" in p.label.lower() + ] + assert ( + len(bracket_params_with_labels) > 0 + ), "Expected some bracket parameters to have generated labels" + + def test__given_bracket_label__then_follows_expected_format(self): + """Bracket labels should follow the format 'Scale label (bracket N field)'.""" + for p in uk_latest.parameters: + if "[" in p.name and p.label and "bracket" in p.label.lower(): + assert re.search( + r"\(bracket \d+ \w+\)", p.label + ), f"Label '{p.label}' doesn't match expected bracket format" + break + class TestUSModel: """Tests for PolicyEngine US model.""" @@ -75,3 +96,49 @@ def test_model_version_parameter_values_aggregates_all(self): """model_version.parameter_values should aggregate all parameter values.""" all_values = us_latest.parameter_values assert len(all_values) >= 100 + + def test__given_breakdown_parameter__then_has_generated_label(self): + """Breakdown parameters (e.g., filing status) should have auto-generated labels.""" + breakdown_params_with_labels = [ + p + for p in us_latest.parameters + if ".SINGLE" in p.name and p.label and "Single" in p.label + ] + assert ( + len(breakdown_params_with_labels) > 0 + ), "Expected some breakdown parameters with SINGLE to have labels containing 'Single'" + + def test__given_bracket_parameter__then_has_generated_label(self): + """Bracket parameters should have auto-generated labels.""" + bracket_params_with_labels = [ + p + for p in us_latest.parameters + if "[" in p.name and p.label and "bracket" in p.label.lower() + ] + assert ( + len(bracket_params_with_labels) > 0 + ), "Expected some bracket parameters to have generated labels" + + def test__given_breakdown_label__then_includes_enum_value_in_parentheses( + self, + ): + """Generated breakdown labels should include the enum value in parentheses.""" + # Find a parameter with a generated label (contains both parent info and enum value) + found = False + for p in us_latest.parameters: + if ".SINGLE" in p.name and p.label and "(" in p.label: + assert ( + "(Single)" in p.label + ), f"Label '{p.label}' should contain '(Single)'" + found = True + break + assert found, "Expected to find at least one generated breakdown label" + + def test__given_bracket_label__then_follows_expected_format(self): + """Bracket labels should follow the format 'Scale label (bracket N field)'.""" + for p in us_latest.parameters: + if "[0].rate" in p.name and p.label and "bracket" in p.label.lower(): + assert re.search( + r"\(bracket \d+ rate\)", p.label + ), f"Label '{p.label}' doesn't match expected bracket format" + break diff --git a/tests/test_parameter_labels.py b/tests/test_parameter_labels.py new file mode 100644 index 00000000..18c74af3 --- /dev/null +++ b/tests/test_parameter_labels.py @@ -0,0 +1,417 @@ +"""Tests for policyengine.utils.parameter_labels module.""" + +from unittest.mock import MagicMock + +import pytest + +from policyengine.utils.parameter_labels import ( + _generate_bracket_label, + _generate_breakdown_label, + build_scale_lookup, + generate_label_for_parameter, +) + +from tests.fixtures.parameter_labels_fixtures import ( + PARAM_WITH_EXPLICIT_LABEL, + PARAM_WITHOUT_LABEL_NO_PARENT, + PARENT_WITH_BREAKDOWN_AND_LABEL, + PARENT_WITH_BREAKDOWN_NO_LABEL, + PARENT_WITHOUT_BREAKDOWN, + SCALE_WITH_LABEL_MARGINAL, + SCALE_WITH_LABEL_SINGLE_AMOUNT, + SCALE_WITHOUT_LABEL, + VARIABLE_WITH_FILING_STATUS_ENUM, + VARIABLE_WITH_STATE_CODE_ENUM, + MockFilingStatus, + MockStateCode, + create_mock_parameter, + create_mock_parent_node, + create_mock_scale, + create_mock_system, + create_mock_variable, +) + + +class TestGenerateLabelForParameter: + """Tests for the generate_label_for_parameter function.""" + + def test__given_parameter_has_explicit_label__then_returns_explicit_label( + self, + ): + # Given + param = PARAM_WITH_EXPLICIT_LABEL + system = create_mock_system() + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then + assert result == "Tax rate" + + def test__given_parameter_without_label_and_no_parent__then_returns_none( + self, + ): + # Given + param = PARAM_WITHOUT_LABEL_NO_PARENT + system = create_mock_system() + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then + assert result is None + + def test__given_bracket_parameter_with_scale_label__then_generates_bracket_label( + self, + ): + # Given + param = create_mock_parameter(name="gov.tax.rates[0].rate") + system = create_mock_system() + scale_lookup = {"gov.tax.rates": SCALE_WITH_LABEL_MARGINAL} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then + assert result == "Income tax rate (bracket 1 rate)" + + def test__given_breakdown_parameter_with_parent_label__then_generates_breakdown_label( + self, + ): + # Given + parent = PARENT_WITH_BREAKDOWN_AND_LABEL + param = create_mock_parameter( + name="gov.exemptions.personal.SINGLE", + parent=parent, + ) + system = create_mock_system( + variables={"filing_status": VARIABLE_WITH_FILING_STATUS_ENUM} + ) + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then + assert result == "Personal exemption amount (Single)" + + def test__given_parameter_with_parent_but_no_breakdown__then_returns_none( + self, + ): + # Given + parent = PARENT_WITHOUT_BREAKDOWN + param = create_mock_parameter( + name="gov.exemptions.personal.value", + parent=parent, + ) + system = create_mock_system() + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then + assert result is None + + +class TestGenerateBreakdownLabel: + """Tests for the _generate_breakdown_label function.""" + + def test__given_parent_has_label_and_enum_match__then_returns_label_with_enum_value( + self, + ): + # Given + parent = PARENT_WITH_BREAKDOWN_AND_LABEL + param = create_mock_parameter( + name="gov.exemptions.personal.JOINT", + parent=parent, + ) + system = create_mock_system( + variables={"filing_status": VARIABLE_WITH_FILING_STATUS_ENUM} + ) + + # When + result = _generate_breakdown_label(param, system) + + # Then + assert result == "Personal exemption amount (Joint)" + + def test__given_parent_has_no_label__then_returns_none(self): + # Given + parent = PARENT_WITH_BREAKDOWN_NO_LABEL + param = create_mock_parameter( + name="gov.exemptions.personal.SINGLE", + parent=parent, + ) + system = create_mock_system( + variables={"filing_status": VARIABLE_WITH_FILING_STATUS_ENUM} + ) + + # When + result = _generate_breakdown_label(param, system) + + # Then + assert result is None + + def test__given_enum_key_not_found__then_returns_label_with_raw_key(self): + # Given + parent = PARENT_WITH_BREAKDOWN_AND_LABEL + param = create_mock_parameter( + name="gov.exemptions.personal.UNKNOWN_STATUS", + parent=parent, + ) + system = create_mock_system( + variables={"filing_status": VARIABLE_WITH_FILING_STATUS_ENUM} + ) + + # When + result = _generate_breakdown_label(param, system) + + # Then + assert result == "Personal exemption amount (UNKNOWN_STATUS)" + + def test__given_variable_not_in_system__then_returns_label_with_raw_key( + self, + ): + # Given + parent = PARENT_WITH_BREAKDOWN_AND_LABEL + param = create_mock_parameter( + name="gov.exemptions.personal.SINGLE", + parent=parent, + ) + system = create_mock_system(variables={}) + + # When + result = _generate_breakdown_label(param, system) + + # Then + assert result == "Personal exemption amount (SINGLE)" + + def test__given_state_code_enum__then_returns_label_with_state_code(self): + # Given + parent = create_mock_parent_node( + name="gov.enrollment.by_state", + label="Enrollment by state", + breakdown=["state_code"], + ) + param = create_mock_parameter( + name="gov.enrollment.by_state.CA", + parent=parent, + ) + system = create_mock_system( + variables={"state_code": VARIABLE_WITH_STATE_CODE_ENUM} + ) + + # When + result = _generate_breakdown_label(param, system) + + # Then + assert result == "Enrollment by state (CA)" + + def test__given_multiple_breakdown_vars__then_uses_first_match(self): + # Given + parent = create_mock_parent_node( + name="gov.rates.by_status_and_state", + label="Rate by status and state", + breakdown=["filing_status", "state_code"], + ) + param = create_mock_parameter( + name="gov.rates.by_status_and_state.SINGLE", + parent=parent, + ) + system = create_mock_system( + variables={ + "filing_status": VARIABLE_WITH_FILING_STATUS_ENUM, + "state_code": VARIABLE_WITH_STATE_CODE_ENUM, + } + ) + + # When + result = _generate_breakdown_label(param, system) + + # Then + assert result == "Rate by status and state (Single)" + + +class TestGenerateBracketLabel: + """Tests for the _generate_bracket_label function.""" + + def test__given_valid_bracket_param_with_scale_label__then_returns_bracket_label( + self, + ): + # Given + param_name = "gov.tax.rates[0].rate" + scale_lookup = {"gov.tax.rates": SCALE_WITH_LABEL_MARGINAL} + + # When + result = _generate_bracket_label(param_name, scale_lookup) + + # Then + assert result == "Income tax rate (bracket 1 rate)" + + def test__given_bracket_index_greater_than_zero__then_uses_one_indexed_bracket_number( + self, + ): + # Given + param_name = "gov.tax.rates[2].threshold" + scale_lookup = {"gov.tax.rates": SCALE_WITH_LABEL_MARGINAL} + + # When + result = _generate_bracket_label(param_name, scale_lookup) + + # Then + assert result == "Income tax rate (bracket 3 threshold)" + + def test__given_single_amount_scale_type__then_uses_tier_instead_of_bracket( + self, + ): + # Given + param_name = "gov.tax.amounts[0].amount" + scale_lookup = {"gov.tax.amounts": SCALE_WITH_LABEL_SINGLE_AMOUNT} + + # When + result = _generate_bracket_label(param_name, scale_lookup) + + # Then + assert result == "Tax amount (tier 1 amount)" + + def test__given_scale_without_label__then_returns_none(self): + # Given + param_name = "gov.tax.rates[0].rate" + scale_lookup = {"gov.tax.rates": SCALE_WITHOUT_LABEL} + + # When + result = _generate_bracket_label(param_name, scale_lookup) + + # Then + assert result is None + + def test__given_scale_not_in_lookup__then_returns_none(self): + # Given + param_name = "gov.tax.rates[0].rate" + scale_lookup = {} + + # When + result = _generate_bracket_label(param_name, scale_lookup) + + # Then + assert result is None + + def test__given_invalid_bracket_format__then_returns_none(self): + # Given + param_name = "gov.tax.rates.rate" # No bracket notation + scale_lookup = {"gov.tax.rates": SCALE_WITH_LABEL_MARGINAL} + + # When + result = _generate_bracket_label(param_name, scale_lookup) + + # Then + assert result is None + + def test__given_threshold_field__then_includes_threshold_in_label(self): + # Given + param_name = "gov.tax.rates[1].threshold" + scale_lookup = {"gov.tax.rates": SCALE_WITH_LABEL_MARGINAL} + + # When + result = _generate_bracket_label(param_name, scale_lookup) + + # Then + assert result == "Income tax rate (bracket 2 threshold)" + + +class TestBuildScaleLookup: + """Tests for the build_scale_lookup function.""" + + def test__given_system_with_scales__then_returns_dict_of_scales_by_name( + self, + ): + # Given + from policyengine_core.parameters import ParameterScale + + scale1 = MagicMock(spec=ParameterScale) + scale1.name = "gov.tax.rates" + scale2 = MagicMock(spec=ParameterScale) + scale2.name = "gov.benefit.amounts" + + system = MagicMock() + system.parameters.get_descendants.return_value = [scale1, scale2] + + # When + result = build_scale_lookup(system) + + # Then + assert "gov.tax.rates" in result + assert "gov.benefit.amounts" in result + assert result["gov.tax.rates"] == scale1 + assert result["gov.benefit.amounts"] == scale2 + + def test__given_system_with_no_scales__then_returns_empty_dict(self): + # Given + from policyengine_core.parameters import Parameter as CoreParameter + + # Create a mock that is NOT a ParameterScale + param = MagicMock(spec=CoreParameter) + param.name = "gov.tax.rate" + + system = MagicMock() + system.parameters.get_descendants.return_value = [param] + + # When + result = build_scale_lookup(system) + + # Then + assert result == {} + + +class TestIntegrationWithRealEnums: + """Integration tests using real-like enum scenarios.""" + + def test__given_head_of_household_status__then_generates_readable_label( + self, + ): + # Given + parent = create_mock_parent_node( + name="gov.pr.exemptions.personal", + label="Puerto Rico personal exemption", + breakdown=["filing_status"], + ) + param = create_mock_parameter( + name="gov.pr.exemptions.personal.HEAD_OF_HOUSEHOLD", + parent=parent, + ) + system = create_mock_system( + variables={"filing_status": VARIABLE_WITH_FILING_STATUS_ENUM} + ) + scale_lookup = {} + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then + assert result == "Puerto Rico personal exemption (Head of household)" + + def test__given_complex_bracket_path__then_extracts_correct_scale_name( + self, + ): + # Given + scale = create_mock_scale( + name="gov.territories.pr.tax.income.tax_rate.amount", + label="Puerto Rico tax rate", + scale_type="marginal_rate", + ) + param = create_mock_parameter( + name="gov.territories.pr.tax.income.tax_rate.amount[0].rate" + ) + system = create_mock_system() + scale_lookup = { + "gov.territories.pr.tax.income.tax_rate.amount": scale + } + + # When + result = generate_label_for_parameter(param, system, scale_lookup) + + # Then + assert result == "Puerto Rico tax rate (bracket 1 rate)" From f0db56b145997dee53a8e1dd66aa195e1f367d85 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 31 Dec 2025 19:36:48 +0400 Subject: [PATCH 3/4] chore: Format --- tests/test_household_impact.py | 1 - tests/test_models.py | 42 +++++++++++++++++++--------------- tests/test_parameter_labels.py | 4 +--- 3 files changed, 24 insertions(+), 23 deletions(-) diff --git a/tests/test_household_impact.py b/tests/test_household_impact.py index a6616219..86439134 100644 --- a/tests/test_household_impact.py +++ b/tests/test_household_impact.py @@ -1,6 +1,5 @@ """Tests for calculate_household_impact functions.""" - from policyengine.tax_benefit_models.uk import ( UKHouseholdInput, UKHouseholdOutput, diff --git a/tests/test_models.py b/tests/test_models.py index 3a74a1b9..06045df9 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -48,17 +48,17 @@ def test__given_bracket_parameter__then_has_generated_label(self): for p in uk_latest.parameters if "[" in p.name and p.label and "bracket" in p.label.lower() ] - assert ( - len(bracket_params_with_labels) > 0 - ), "Expected some bracket parameters to have generated labels" + assert len(bracket_params_with_labels) > 0, ( + "Expected some bracket parameters to have generated labels" + ) def test__given_bracket_label__then_follows_expected_format(self): """Bracket labels should follow the format 'Scale label (bracket N field)'.""" for p in uk_latest.parameters: if "[" in p.name and p.label and "bracket" in p.label.lower(): - assert re.search( - r"\(bracket \d+ \w+\)", p.label - ), f"Label '{p.label}' doesn't match expected bracket format" + assert re.search(r"\(bracket \d+ \w+\)", p.label), ( + f"Label '{p.label}' doesn't match expected bracket format" + ) break @@ -104,9 +104,9 @@ def test__given_breakdown_parameter__then_has_generated_label(self): for p in us_latest.parameters if ".SINGLE" in p.name and p.label and "Single" in p.label ] - assert ( - len(breakdown_params_with_labels) > 0 - ), "Expected some breakdown parameters with SINGLE to have labels containing 'Single'" + assert len(breakdown_params_with_labels) > 0, ( + "Expected some breakdown parameters with SINGLE to have labels containing 'Single'" + ) def test__given_bracket_parameter__then_has_generated_label(self): """Bracket parameters should have auto-generated labels.""" @@ -115,9 +115,9 @@ def test__given_bracket_parameter__then_has_generated_label(self): for p in us_latest.parameters if "[" in p.name and p.label and "bracket" in p.label.lower() ] - assert ( - len(bracket_params_with_labels) > 0 - ), "Expected some bracket parameters to have generated labels" + assert len(bracket_params_with_labels) > 0, ( + "Expected some bracket parameters to have generated labels" + ) def test__given_breakdown_label__then_includes_enum_value_in_parentheses( self, @@ -127,9 +127,9 @@ def test__given_breakdown_label__then_includes_enum_value_in_parentheses( found = False for p in us_latest.parameters: if ".SINGLE" in p.name and p.label and "(" in p.label: - assert ( - "(Single)" in p.label - ), f"Label '{p.label}' should contain '(Single)'" + assert "(Single)" in p.label, ( + f"Label '{p.label}' should contain '(Single)'" + ) found = True break assert found, "Expected to find at least one generated breakdown label" @@ -137,8 +137,12 @@ def test__given_breakdown_label__then_includes_enum_value_in_parentheses( def test__given_bracket_label__then_follows_expected_format(self): """Bracket labels should follow the format 'Scale label (bracket N field)'.""" for p in us_latest.parameters: - if "[0].rate" in p.name and p.label and "bracket" in p.label.lower(): - assert re.search( - r"\(bracket \d+ rate\)", p.label - ), f"Label '{p.label}' doesn't match expected bracket format" + if ( + "[0].rate" in p.name + and p.label + and "bracket" in p.label.lower() + ): + assert re.search(r"\(bracket \d+ rate\)", p.label), ( + f"Label '{p.label}' doesn't match expected bracket format" + ) break diff --git a/tests/test_parameter_labels.py b/tests/test_parameter_labels.py index 18c74af3..4413c259 100644 --- a/tests/test_parameter_labels.py +++ b/tests/test_parameter_labels.py @@ -406,9 +406,7 @@ def test__given_complex_bracket_path__then_extracts_correct_scale_name( name="gov.territories.pr.tax.income.tax_rate.amount[0].rate" ) system = create_mock_system() - scale_lookup = { - "gov.territories.pr.tax.income.tax_rate.amount": scale - } + scale_lookup = {"gov.territories.pr.tax.income.tax_rate.amount": scale} # When result = generate_label_for_parameter(param, system, scale_lookup) From 872ed9534e2dc4910c802e7c5edc94eec0f8a5b7 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 31 Dec 2025 19:52:05 +0400 Subject: [PATCH 4/4] fix: Remove unused imports in test_parameter_labels MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- tests/test_parameter_labels.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/test_parameter_labels.py b/tests/test_parameter_labels.py index 4413c259..b055e77e 100644 --- a/tests/test_parameter_labels.py +++ b/tests/test_parameter_labels.py @@ -2,15 +2,12 @@ from unittest.mock import MagicMock -import pytest - from policyengine.utils.parameter_labels import ( _generate_bracket_label, _generate_breakdown_label, build_scale_lookup, generate_label_for_parameter, ) - from tests.fixtures.parameter_labels_fixtures import ( PARAM_WITH_EXPLICIT_LABEL, PARAM_WITHOUT_LABEL_NO_PARENT, @@ -22,13 +19,10 @@ SCALE_WITHOUT_LABEL, VARIABLE_WITH_FILING_STATUS_ENUM, VARIABLE_WITH_STATE_CODE_ENUM, - MockFilingStatus, - MockStateCode, create_mock_parameter, create_mock_parent_node, create_mock_scale, create_mock_system, - create_mock_variable, )