From 872ea200e520d2b7095af4ea4f5fc5fc80dbd053 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 6 Nov 2026 05:30:03 -0500 Subject: [PATCH 1/3] Forward explicit branches through reported period discovery Keep calls without a branch compatible with the current Core holder API. Cover branch visibility and root fallback with the existing Core fixture. Co-Authored-By: Claude Opus 5.5 --- .../known-periods-branch-wrapper.fixed.md | 1 + .../test_social_security_known_periods.py | 79 +++++++++++++++++++ .../tools/social_security_inputs.py | 15 +++- 3 files changed, 91 insertions(+), 4 deletions(-) create mode 100644 changelog.d/known-periods-branch-wrapper.fixed.md create mode 100644 policyengine_us/tests/core/test_social_security_known_periods.py diff --git a/changelog.d/known-periods-branch-wrapper.fixed.md b/changelog.d/known-periods-branch-wrapper.fixed.md new file mode 100644 index 00000000000..8745f87102c --- /dev/null +++ b/changelog.d/known-periods-branch-wrapper.fixed.md @@ -0,0 +1 @@ +Fixed Social Security period discovery to forward explicit branches while preserving compatibility with Core versions that accept no branch argument. diff --git a/policyengine_us/tests/core/test_social_security_known_periods.py b/policyengine_us/tests/core/test_social_security_known_periods.py new file mode 100644 index 00000000000..a4a6172c4a5 --- /dev/null +++ b/policyengine_us/tests/core/test_social_security_known_periods.py @@ -0,0 +1,79 @@ +"""Keep reported source periods compatible with both Core holder APIs.""" + +import inspect + +import numpy as np +import pytest +from policyengine_core.holders import Holder +from policyengine_core.periods import period + +import test_social_security_real_core as real_core +from test_social_security_real_core import installed_hooks as installed_hooks + + +@pytest.fixture +def branch_aware_core(): + physical = real_core.hooks._PHYSICAL["get_known_periods"] + if "branch_name" not in inspect.signature(physical).parameters: + pytest.skip("Core master Holder.get_known_periods accepts no branch argument") + + +@pytest.mark.parametrize("public", real_core.COMPONENTS) +def test_known_periods_without_branch_keeps_reported_sources(public): + reported = f"{public}_reported" + simulation = real_core.fresh( + {public: {"2024": [12_000]}, reported: {"2026": [6_000]}} + ) + holder = simulation.get_holder(reported) + assert {str(known) for known in holder.get_known_periods()} == {"2024", "2026"} + assert {str(known) for known in holder.get_known_periods(None)} == {"2024", "2026"} + + +@pytest.mark.usefixtures("branch_aware_core") +@pytest.mark.parametrize("public", real_core.COMPONENTS) +def test_known_periods_lists_only_sources_readable_by_requested_branch(public): + reported = f"{public}_reported" + simulation = real_core.fresh() + simulation.get_holder(public).set_input( + period(2024), np.array([12_000]), branch_name="later" + ) + holder = simulation.get_holder(reported) + holder.set_input(period(2026), np.array([6_000]), branch_name="later") + + assert holder.get_array(period(2024), "later")[0] == 12_000 + assert holder.get_array(period(2024), "default") is None + assert {str(known) for known in holder.get_known_periods("later")} == { + "2024", + "2026", + } + assert holder.get_known_periods("default") == [] + assert holder.get_known_periods("unrelated") == [] + + +@pytest.mark.usefixtures("branch_aware_core") +def test_known_periods_forwards_branch_to_fallback_check(monkeypatch): + simulation = real_core.fresh() + holder = simulation.get_holder(f"{real_core.PUBLIC}_reported") + original = real_core.hooks._allow_fallback + branches = [] + + def allow_fallback(simulation, name, period=None, branch_name=None): + if period is None: + branches.append(branch_name) + return original(simulation, name, period, branch_name) + + monkeypatch.setattr(real_core.hooks, "_allow_fallback", allow_fallback) + assert holder.get_known_periods("later") == [] + assert branches == ["later"] + + +@pytest.mark.parametrize("public", real_core.COMPONENTS) +def test_root_fallback_calculates_missing_reported_input(public): + simulation = real_core.fresh() + assert simulation.calculate(f"{public}_reported", 2024)[0] == 0 + + +def test_root_fallback_calculates_missing_nonreported_input(): + simulation = real_core.fresh() + simulation.delete_arrays("independent_input") + assert simulation.calculate("independent_input", 2024)[0] == 0 diff --git a/policyengine_us/tools/social_security_inputs.py b/policyengine_us/tools/social_security_inputs.py index 3cbfc659f8d..249a14a6a16 100644 --- a/policyengine_us/tools/social_security_inputs.py +++ b/policyengine_us/tools/social_security_inputs.py @@ -893,10 +893,13 @@ def get_array(holder, period, branch_name="default"): Holder.get_array = get_array @wraps(_PHYSICAL["get_known_periods"]) - def get_known_periods(holder): - values = _PHYSICAL["get_known_periods"](holder) + def get_known_periods(holder, branch_name=None): + branch_kwargs = {} if branch_name is None else {"branch_name": branch_name} + values = _PHYSICAL["get_known_periods"](holder, **branch_kwargs) if holder.variable.name in _REPORTED_COMPONENTS and _enabled(holder.simulation): - if _allow_fallback(holder.simulation, holder.variable.name): + if _allow_fallback( + holder.simulation, holder.variable.name, **branch_kwargs + ): values = [ *values, *( @@ -904,9 +907,13 @@ def get_known_periods(holder): for period in _public_periods( holder.simulation, holder.variable.name.removesuffix("_reported"), + **branch_kwargs, ) if _effective_value( - holder.simulation, holder.variable.name, period + holder.simulation, + holder.variable.name, + period, + **branch_kwargs, )[0] is not None ), From 626fcc7fab8127d5e7313b12213fbbe4c4bae619 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 6 Nov 2026 07:29:46 -0500 Subject: [PATCH 2/3] Use the selected Core source in Social Security state tests Co-Authored-By: Claude Opus 5.5 --- .../core/_social_security_state_backend.py | 22 +++++++++++++------ ...est_social_security_input_state_machine.py | 2 +- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/policyengine_us/tests/core/_social_security_state_backend.py b/policyengine_us/tests/core/_social_security_state_backend.py index 04d6f7d1550..dc26774a88c 100644 --- a/policyengine_us/tests/core/_social_security_state_backend.py +++ b/policyengine_us/tests/core/_social_security_state_backend.py @@ -1,6 +1,6 @@ """Small storage backend for the Social Security input state machine. -The country is intentionally never imported. Core's installed source bodies +The country is intentionally never imported. Core's selected source bodies implement every write, builder, copy, replay, extraction and dump route below; the US module's complete function bodies implement the hooks. Entity and variable descriptors, dataset construction, and disk dump bytes are synthetic. @@ -49,7 +49,6 @@ from dataclasses import dataclass from enum import Enum from functools import lru_cache, wraps -from importlib.metadata import distribution from pathlib import Path from unittest.mock import patch from weakref import WeakKeyDictionary, ref @@ -179,7 +178,7 @@ def __init__(self, *, cache_blacklist=(), variables_to_drop=()): def _setup(self): cache_blacklist = self.cache_blacklist variables_to_drop = self.variables_to_drop - core = Path(distribution("policyengine-core").locate_file("policyengine_core")) + core = Path(core_simulation.__file__).resolve().parents[1] calculation_frames = core_simulation._calculation_frames formula_simulation = core_simulation._formula_simulation branch_clone = core_simulation._branch_clone @@ -389,7 +388,7 @@ def subsample(self, *args): def calculate(self, name, period): # The country-free fixture has no tracer or concurrent # formulas. Retain Core's public reader cache behavior while - # its installed _calculate body supplies the actual value. + # its selected _calculate body supplies the actual value. period = as_period(period) key = (name, period) if key not in self._fast_cache: @@ -397,7 +396,7 @@ def calculate(self, name, period): return self._fast_cache[key] def _calculate(self, name, period): - raise AssertionError("Installed Core calculation body was not loaded") + raise AssertionError("Selected Core calculation body was not loaded") def check_macro_cache(self, *args): return False @@ -590,12 +589,12 @@ def load(self): variable.set_input = core_helpers["set_input_divide_by_period"] # Core 3.33.12 cleans overlapping raw seeds after a direct set_input; # 3.33.2 does this only in dispatch/divide helpers. Keep the reference - # model tied to the installed behavior rather than a version string. + # model tied to the selected behavior rather than a version string. self.cleans_raw_overlaps = "_drop_calculated_over_own_period" in core_helpers simulation_source = core / "simulations/simulation.py" # _latest_input_key calls _end_order, which uses Core's calendar # table even for yearly sources. Load that literal from the same - # installed module rather than maintaining a second calendar table. + # selected module rather than maintaining a second calendar table. day_offsets = next( node.value for node in source_tree(simulation_source).body @@ -636,6 +635,15 @@ def load(self): "_input_set_meanwhile", "_drop_computed", ) + sim_methods += tuple( + name + for name in ( + "_set_fallback_cache_status", + "_fallback_is_uncached", + "_uncached_fallback", + ) + if hasattr(core_simulation.Simulation, name) + ) for name, method in execute( core / "simulations/simulation.py", sim_methods, namespace, "Simulation" ).items(): diff --git a/policyengine_us/tests/core/test_social_security_input_state_machine.py b/policyengine_us/tests/core/test_social_security_input_state_machine.py index aabc2ebaabb..151883e0028 100644 --- a/policyengine_us/tests/core/test_social_security_input_state_machine.py +++ b/policyengine_us/tests/core/test_social_security_input_state_machine.py @@ -8,7 +8,7 @@ reported amount and public formula amount, including the initial scored year and a later year whose intermediate base may already have been cached. -The backend executes installed Core source bodies and the production US hooks +The backend executes selected Core source bodies and the production US hooks without importing the country. Run this file directly for a country-free test: ``python policyengine_us/tests/core/test_social_security_input_state_machine.py``. """ From de480911d02a656d912b2973b00d5319ca602be4 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 6 Nov 2026 18:51:01 -0500 Subject: [PATCH 3/3] Store the derived entry the dataset provenance test reuses Core PR 552 can leave a missing input's default uncached, so calculating independent_input for 2023 no longer stores the derived entry that the rebuilt dataset's input then shares a key with. Store that entry in the test, so it checks the same provenance under every Core release. Co-Authored-By: Claude Opus 5.5 --- policyengine_us/tests/core/test_social_security_real_core.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/policyengine_us/tests/core/test_social_security_real_core.py b/policyengine_us/tests/core/test_social_security_real_core.py index b29fb14af6f..29e5ff44f62 100644 --- a/policyengine_us/tests/core/test_social_security_real_core.py +++ b/policyengine_us/tests/core/test_social_security_real_core.py @@ -443,6 +443,9 @@ def test_dataset_rebuild_keeps_new_input_provenance_when_old_cache_shares_key(): assert since is not None assert scalar(simulation, "independent_input", "2023") == 0 old_holder = simulation.get_holder("independent_input") + # Core may leave a missing input's default uncached. Store the derived + # entry this test needs under the key the dataset's input then reuses. + old_holder.put_in_cache(np.array([0.0]), period("2023"), derived=True) assert old_holder.is_derived(period("2023")) assert sequence_number(simulation, "independent_input", "2023") >= since