From 17e7bd41586bb3ccdef56fe8642d628f3905c5ba Mon Sep 17 00:00:00 2001 From: Ziming Date: Wed, 7 Oct 2026 20:21:57 -0400 Subject: [PATCH] Reduce repeated setup in expensive Rest tests --- changelog.d/reduce-rest-test-setup.changed.md | 1 + .../core/test_marginal_tax_rate_coverage.py | 58 ++++++++++++------- .../tests/test_federal_credit_limit_order.py | 7 ++- 3 files changed, 44 insertions(+), 22 deletions(-) create mode 100644 changelog.d/reduce-rest-test-setup.changed.md diff --git a/changelog.d/reduce-rest-test-setup.changed.md b/changelog.d/reduce-rest-test-setup.changed.md new file mode 100644 index 00000000000..4449190dc93 --- /dev/null +++ b/changelog.d/reduce-rest-test-setup.changed.md @@ -0,0 +1 @@ +Reduce repeated simulation and model setup in Rest Python tests while preserving their existing cases and assertions. diff --git a/policyengine_us/tests/core/test_marginal_tax_rate_coverage.py b/policyengine_us/tests/core/test_marginal_tax_rate_coverage.py index bb75598951f..a15eb5199b7 100644 --- a/policyengine_us/tests/core/test_marginal_tax_rate_coverage.py +++ b/policyengine_us/tests/core/test_marginal_tax_rate_coverage.py @@ -21,6 +21,7 @@ changes measured marginal tax rates and so earnings and ranks. """ +from functools import cache from itertools import product import numpy as np @@ -102,20 +103,44 @@ def _calc(simulation, variable): @pytest.fixture(scope="module") -def households(): - simulation = _simulation(0) - return { - "household_id": simulation.populations["household"].members_entity_id, - "is_adult": _calc(simulation, "is_adult").astype(bool), - "market_income": _calc(simulation, "market_income"), - } +def results_for_cap(): + # The flag and rate checks use identical situations. Keep only read-only + # output snapshots; each cap still gets its own simulation and reform. + @cache + def calculate(cap): + simulation = _simulation(cap) + results = { + variable: _calc(simulation, variable) + for variable in ("marginal_tax_rate_computed", "cliff_evaluated") + } + if cap == 0: + results.update( + household_id=simulation.populations["household"].members_entity_id, + is_adult=_calc(simulation, "is_adult").astype(bool), + market_income=_calc(simulation, "market_income"), + ) + if cap in (2, 3): + results.update( + (variable, _calc(simulation, variable)) for variable in MTR_VARIABLES + ) + for variable, values in results.items(): + results[variable] = values.copy() + results[variable].setflags(write=False) + return results + + return calculate + + +@pytest.fixture(scope="module") +def households(results_for_cap): + return results_for_cap(0) @pytest.mark.parametrize("cap", [0, 1, 2, 3, 5]) -def test_flag_selects_the_top_earning_adults(households, cap): - simulation = _simulation(cap) - flag = _calc(simulation, "marginal_tax_rate_computed").astype(bool) - assert np.array_equal(_calc(simulation, "cliff_evaluated").astype(bool), flag) +def test_flag_selects_the_top_earning_adults(households, results_for_cap, cap): + results = results_for_cap(cap) + flag = results["marginal_tax_rate_computed"].astype(bool) + assert np.array_equal(results["cliff_evaluated"].astype(bool), flag) assert not flag[~households["is_adult"]].any() for household_id in np.unique(households["household_id"]): in_household = households["household_id"] == household_id @@ -129,15 +154,8 @@ def test_flag_selects_the_top_earning_adults(households, cap): @pytest.fixture(scope="module") -def mtr_by_cap(): - results = {} - for cap in (2, 3): - simulation = _simulation(cap) - results[cap] = { - variable: _calc(simulation, variable) - for variable in ("marginal_tax_rate_computed",) + MTR_VARIABLES - } - return results +def mtr_by_cap(results_for_cap): + return {cap: results_for_cap(cap) for cap in (2, 3)} @pytest.mark.parametrize("variable", MTR_VARIABLES) diff --git a/policyengine_us/tests/test_federal_credit_limit_order.py b/policyengine_us/tests/test_federal_credit_limit_order.py index 0444e200574..39e829c0f4e 100644 --- a/policyengine_us/tests/test_federal_credit_limit_order.py +++ b/policyengine_us/tests/test_federal_credit_limit_order.py @@ -133,8 +133,11 @@ def test_2020_irc_pin_keeps_the_order_in_2021(baseline_system): (non_refundable_ss_credit_reform(), "non_refundable_ctc"), ], ) -def test_reforms_order_the_ctc_and_25d_lists(reform, credit_before_25d): - credits = reform(CountryTaxBenefitSystem()).parameters.gov.irs.credits +def test_reforms_order_the_ctc_and_25d_lists( + baseline_system, reform, credit_before_25d +): + # Each reform clones the baseline's parameters and variable registry. + credits = reform(baseline_system).parameters.gov.irs.credits date = "2026-06-01" non_refundable = list(credits.non_refundable(date)) assert non_refundable[0] == "foreign_tax_credit"