diff --git a/changelog.d/834.fixed.md b/changelog.d/834.fixed.md new file mode 100644 index 000000000..bac592908 --- /dev/null +++ b/changelog.d/834.fixed.md @@ -0,0 +1 @@ +- Require `medicare_enrolled` to be a pure PolicyEngine US input before CPS writes Medicare enrollment values. diff --git a/policyengine_us_data/utils/policyengine.py b/policyengine_us_data/utils/policyengine.py index 869b95d9a..a292dfe92 100644 --- a/policyengine_us_data/utils/policyengine.py +++ b/policyengine_us_data/utils/policyengine.py @@ -121,10 +121,15 @@ def assert_locked_policyengine_us_version() -> PolicyEngineUSBuildInfo: @lru_cache(maxsize=1) -def _policyengine_us_variable_names() -> frozenset[str]: +def _policyengine_us_variables(): from policyengine_us import CountryTaxBenefitSystem - return frozenset(CountryTaxBenefitSystem().variables) + return CountryTaxBenefitSystem().variables + + +@lru_cache(maxsize=1) +def _policyengine_us_variable_names() -> frozenset[str]: + return frozenset(_policyengine_us_variables()) def has_policyengine_us_variables(*variables: str) -> bool: @@ -136,8 +141,31 @@ def has_policyengine_us_variables(*variables: str) -> bool: return set(variables).issubset(available_variables) +def has_policyengine_us_pure_input_variables(*variables: str) -> bool: + try: + available_variables = _policyengine_us_variables() + except Exception: + return False + + for variable in variables: + variable_definition = available_variables.get(variable) + if variable_definition is None: + return False + if not variable_definition.is_input_variable(): + return False + if getattr(variable_definition, "defined_for", None) is not None: + return False + if getattr(variable_definition, "adds", None): + return False + if getattr(variable_definition, "subtracts", None): + return False + if getattr(variable_definition, "formulas", None): + return False + return True + + def supports_medicare_enrollment_input() -> bool: - return has_policyengine_us_variables("medicare_enrolled") + return has_policyengine_us_pure_input_variables("medicare_enrolled") def supports_modeled_medicare_part_b_inputs() -> bool: diff --git a/tests/unit/test_medicare_part_b_inputs.py b/tests/unit/test_medicare_part_b_inputs.py index c69e88789..3919a6301 100644 --- a/tests/unit/test_medicare_part_b_inputs.py +++ b/tests/unit/test_medicare_part_b_inputs.py @@ -6,6 +6,21 @@ from policyengine_us_data.utils import policyengine as policyengine_utils +class _Variable: + adds = None + defined_for = None + formulas = None + subtracts = None + + def __init__(self, *, is_input=True, defined_for=None, formulas=None): + self._is_input = is_input + self.defined_for = defined_for + self.formulas = formulas + + def is_input_variable(self): + return self._is_input + + def test_medicare_part_b_clone_imputation_matches_installed_model_support(): assert ("medicare_part_b_premiums" in set(CPS_ONLY_IMPUTED_VARIABLES)) is ( not supports_modeled_medicare_part_b_inputs() @@ -22,6 +37,35 @@ def test_supports_medicare_enrollment_input_allows_partial_support(monkeypatch): "has_policyengine_us_variables", lambda *variables: variables == ("medicare_enrolled",), ) + monkeypatch.setattr( + policyengine_utils, + "has_policyengine_us_pure_input_variables", + lambda *variables: variables == ("medicare_enrolled",), + ) assert policyengine_utils.supports_medicare_enrollment_input() is True assert policyengine_utils.supports_modeled_medicare_part_b_inputs() is False + + +def test_pure_input_support_rejects_formula_and_conditional_variables(monkeypatch): + monkeypatch.setattr( + policyengine_utils, + "_policyengine_us_variables", + lambda: { + "pure_input": _Variable(), + "formula_input": _Variable(formulas={"2025": object()}), + "conditional_input": _Variable(defined_for="is_eligible"), + "formula_variable": _Variable(is_input=False), + }, + ) + + assert policyengine_utils.has_policyengine_us_pure_input_variables("pure_input") + assert not policyengine_utils.has_policyengine_us_pure_input_variables( + "formula_input" + ) + assert not policyengine_utils.has_policyengine_us_pure_input_variables( + "conditional_input" + ) + assert not policyengine_utils.has_policyengine_us_pure_input_variables( + "formula_variable" + )