From 311c68d6a7d97789312e9f078d19625647d9376e Mon Sep 17 00:00:00 2001
From: Max Ghenis <mghenis@gmail.com>
Date: Sun, 11 Oct 2026 17:38:08 -0400
Subject: [PATCH 1/4] Release cached test systems and simulations at module end

Five test modules cached tax-benefit systems or simulations in module-level
lru_caches that were never cleared, so a pytest worker kept them for the rest
of its run. The two-child limit file kept one fresh system per parameter value
drawn (13 on CI, +0.54 GiB RSS that stayed).

- two-child limit and cliff property tests: build reforms over one shared
  baseline system (a reform clones only the parameters), keep the last two
  reforms (two-child) or the four mtr_adults values (cliff), and clear the
  caches in a module-scoped autouse fixture.
- test_state_pension_age, test_state_pension_credit_qualifying_age and
  test_uc_work_related_groups_properties: clear their cached simulations in a
  module-scoped autouse fixture.
- code_health/test_test_module_caches_released.py: every module-level
  functools cache a test module defines or imports must be cleared after the
  yield of a module-scoped autouse fixture.

Test functions' source is unchanged, so Hypothesis's derandomized examples are
the same.

Fixes #2357

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 changelog.d/2357.fixed.md                     |   1 +
 .../test_test_module_caches_released.py       | 159 ++++++++++++++++++
 .../tests/test_cliff_evaluated_properties.py  |  21 ++-
 .../tests/test_state_pension_age.py           |  10 ++
 ...est_state_pension_credit_qualifying_age.py |  12 ++
 ...wo_child_limit_age_exemption_properties.py |  36 +++-
 .../test_uc_work_related_groups_properties.py |   8 +
 7 files changed, 238 insertions(+), 9 deletions(-)
 create mode 100644 changelog.d/2357.fixed.md
 create mode 100644 policyengine_uk/tests/code_health/test_test_module_caches_released.py

diff --git a/changelog.d/2357.fixed.md b/changelog.d/2357.fixed.md
new file mode 100644
index 0000000000..c4906fa91c
--- /dev/null
+++ b/changelog.d/2357.fixed.md
@@ -0,0 +1 @@
+- Release the tax-benefit systems and simulations that five test modules cache once each module's tests finish, build the two-child limit and cliff property tests' reforms over one shared baseline system, and check that every test module cache is released (#2357).
diff --git a/policyengine_uk/tests/code_health/test_test_module_caches_released.py b/policyengine_uk/tests/code_health/test_test_module_caches_released.py
new file mode 100644
index 0000000000..71e123d79e
--- /dev/null
+++ b/policyengine_uk/tests/code_health/test_test_module_caches_released.py
@@ -0,0 +1,159 @@
+"""Test modules must release their module-level caches when they finish.
+
+Property tests cache tax-benefit systems and simulations with
+``functools.lru_cache`` so that each Hypothesis example does not build them
+again. A module-level cache outlives the module's tests: it keeps everything
+it holds for the rest of the test process, and a system also keeps the last
+simulation built from it (``Simulation.__init__`` sets
+``tax_benefit_system.simulation``). One such cache kept 13 systems, about
+0.5 GiB, in a CI worker for the rest of its run.
+
+So every function a test module caches with ``functools.lru_cache`` or
+``functools.cache``, whether it defines it or imports it from another test
+module, must have its ``cache_clear()`` called after the ``yield`` of a
+module-scoped autouse fixture in that module. An imported cache counts
+separately: pytest imports test files by bare name, so a module that imports
+from another test file loads a second copy of it, with its own caches.
+"""
+
+import ast
+from pathlib import Path
+
+TESTS = Path(__file__).resolve().parents[1]
+CACHE_DECORATORS = {"lru_cache", "cache"}
+
+
+def _name(node: ast.AST) -> str | None:
+    if isinstance(node, ast.Call):
+        node = node.func
+    if isinstance(node, ast.Attribute):
+        return node.attr
+    if isinstance(node, ast.Name):
+        return node.id
+    return None
+
+
+def cached_functions(tree: ast.Module) -> set[str]:
+    """Module-level functions decorated with a functools cache."""
+    return {
+        node.name
+        for node in tree.body
+        if isinstance(node, ast.FunctionDef)
+        and any(_name(d) in CACHE_DECORATORS for d in node.decorator_list)
+    }
+
+
+def _keyword(call: ast.Call, name: str):
+    for keyword in call.keywords:
+        if keyword.arg == name and isinstance(keyword.value, ast.Constant):
+            return keyword.value.value
+    return None
+
+
+def _is_module_autouse_fixture(function: ast.FunctionDef) -> bool:
+    for decorator in function.decorator_list:
+        if (
+            isinstance(decorator, ast.Call)
+            and _name(decorator) == "fixture"
+            and _keyword(decorator, "scope") == "module"
+            and _keyword(decorator, "autouse") is True
+        ):
+            return True
+    return False
+
+
+def released_caches(tree: ast.Module) -> set[str]:
+    """Names whose cache_clear() runs after the yield of a module-scoped
+    autouse fixture."""
+    released = set()
+    for function in tree.body:
+        if not (
+            isinstance(function, ast.FunctionDef)
+            and _is_module_autouse_fixture(function)
+        ):
+            continue
+        after_yield = False
+        for statement in function.body:
+            if isinstance(statement, ast.Expr) and isinstance(
+                statement.value, ast.Yield
+            ):
+                after_yield = True
+                continue
+            if not after_yield:
+                continue
+            for node in ast.walk(statement):
+                if (
+                    isinstance(node, ast.Call)
+                    and isinstance(node.func, ast.Attribute)
+                    and node.func.attr == "cache_clear"
+                    and isinstance(node.func.value, ast.Name)
+                ):
+                    released.add(node.func.value.id)
+    return released
+
+
+def imported_test_modules(tree: ast.Module) -> dict[Path, list[str]]:
+    """Other test files this module imports names from."""
+    imports = {}
+    for node in tree.body:
+        if isinstance(node, ast.ImportFrom) and node.module:
+            module = node.module.removeprefix("policyengine_uk.tests.")
+            path = TESTS / (module.replace(".", "/") + ".py")
+            if path.exists():
+                imports[path] = [alias.name for alias in node.names]
+    return imports
+
+
+def parse(path: Path) -> ast.Module:
+    return ast.parse(path.read_text(), filename=str(path))
+
+
+def unreleased_caches(path: Path) -> list[str]:
+    tree = parse(path)
+    caches = cached_functions(tree)
+    for imported, names in imported_test_modules(tree).items():
+        caches |= cached_functions(parse(imported)) & set(names)
+    return sorted(caches - released_caches(tree))
+
+
+def test_every_test_module_cache_is_released():
+    findings = {
+        str(path.relative_to(TESTS)): missing
+        for path in sorted(TESTS.rglob("*.py"))
+        if (missing := unreleased_caches(path))
+    }
+    assert not findings, (
+        "Clear these caches after the yield of a module-scoped autouse "
+        f"fixture in the module that uses them: {findings}"
+    )
+
+
+def test_the_scan_finds_caches_and_releases(tmp_path):
+    source = """
+from functools import cache, lru_cache
+
+import pytest
+
+
+@lru_cache(maxsize=None)
+def systems(value):
+    return value
+
+
+@cache
+def simulation():
+    return 1
+
+
+@pytest.fixture(scope="module", autouse=True)
+def release():
+    systems.cache_clear()  # before the yield: does not count
+    yield
+    simulation.cache_clear()
+"""
+    tree = ast.parse(source)
+    assert cached_functions(tree) == {"systems", "simulation"}
+    assert released_caches(tree) == {"simulation"}
+
+    function_scoped = source.replace('scope="module", ', "")
+    assert released_caches(ast.parse(function_scoped)) == set()
diff --git a/policyengine_uk/tests/test_cliff_evaluated_properties.py b/policyengine_uk/tests/test_cliff_evaluated_properties.py
index 0e2144d3e8..14c46b230c 100644
--- a/policyengine_uk/tests/test_cliff_evaluated_properties.py
+++ b/policyengine_uk/tests/test_cliff_evaluated_properties.py
@@ -16,6 +16,7 @@
 from functools import lru_cache
 
 import numpy as np
+import pytest
 from hypothesis import HealthCheck, given, settings
 from hypothesis import strategies as st
 from policyengine_core.reforms import Reform, set_parameter
@@ -32,7 +33,14 @@
 )
 
 
-@lru_cache(maxsize=None)
+@lru_cache(maxsize=1)
+def baseline_system():
+    return CountryTaxBenefitSystem()
+
+
+# One reform per value scenarios() draws for mtr_adults. A reform clones only
+# the baseline's parameters.
+@lru_cache(maxsize=4)
 def tax_benefit_system(mtr_adults):
     modifier = set_parameter(
         MTR_ADULTS, mtr_adults, return_modifier=True, period="year:2000:100"
@@ -42,7 +50,16 @@ class override(Reform):
         def apply(self):
             self.parameters = modifier(self.parameters)
 
-    return override(CountryTaxBenefitSystem())
+    return override(baseline_system())
+
+
+@pytest.fixture(scope="module", autouse=True)
+def release_systems():
+    # Each system also keeps its last simulation alive. Drop them once this
+    # module's tests are done, instead of for the rest of the test process.
+    yield
+    tax_benefit_system.cache_clear()
+    baseline_system.cache_clear()
 
 
 earnings = st.sampled_from([0, 5_000, 20_000, 60_000, 120_000])
diff --git a/policyengine_uk/tests/test_state_pension_age.py b/policyengine_uk/tests/test_state_pension_age.py
index a8cb7d948c..459c5af90e 100644
--- a/policyengine_uk/tests/test_state_pension_age.py
+++ b/policyengine_uk/tests/test_state_pension_age.py
@@ -256,6 +256,16 @@ def differential_simulation() -> tuple:
     return Simulation(situation=situation(people)), keys
 
 
+@pytest.fixture(scope="module", autouse=True)
+def release_cached_simulations():
+    # The differential simulation keeps every year it has calculated. Drop
+    # the cached simulations once this module's tests are done, instead of
+    # for the rest of the test process.
+    yield
+    differential_simulation.cache_clear()
+    situation_template.cache_clear()
+
+
 @pytest.mark.parametrize("year", YEARS)
 def test_status_matches_statute_for_every_birth_date(year):
     """State Pension age status, type, the Savings Credit age test and the day
diff --git a/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py b/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py
index a1eea4e249..3ae1020499 100644
--- a/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py
+++ b/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py
@@ -81,6 +81,18 @@ def property_simulation(people: dict):
     return sim
 
 
+@pytest.fixture(scope="module", autouse=True)
+def release_cached_simulations():
+    # Pytest imports test files here by bare name (the directory has no
+    # __init__.py), so the import above loads test_state_pension_age a second
+    # time and this module fills its own differential_simulation cache. Drop
+    # the cached simulations once this module's tests are done, instead of for
+    # the rest of the test process.
+    yield
+    differential_simulation.cache_clear()
+    one_person_units_template.cache_clear()
+
+
 @pytest.mark.parametrize("year", YEARS)
 def test_qualifying_age_matches_statute_for_every_birth_date(year):
     """For births every third day from 1945 to 1984 and on both sides of every
diff --git a/policyengine_uk/tests/test_two_child_limit_age_exemption_properties.py b/policyengine_uk/tests/test_two_child_limit_age_exemption_properties.py
index 315dc29705..6c2835d2d8 100644
--- a/policyengine_uk/tests/test_two_child_limit_age_exemption_properties.py
+++ b/policyengine_uk/tests/test_two_child_limit_age_exemption_properties.py
@@ -17,6 +17,7 @@
 from functools import lru_cache
 
 import numpy as np
+import pytest
 from hypothesis import HealthCheck, given, settings
 from hypothesis import strategies as st
 from policyengine_core.reforms import Reform, set_parameter
@@ -38,12 +39,17 @@
 TOLERANCE = 0.01
 
 
-@lru_cache(maxsize=None)
-def tax_benefit_system(parameter, value):
-    """The UK system, with one parameter overridden for all years if given."""
-    system = CountryTaxBenefitSystem()
-    if parameter is None:
-        return system
+@lru_cache(maxsize=1)
+def baseline_system():
+    return CountryTaxBenefitSystem()
+
+
+# Each example uses the repealed limit and one threshold. A reform clones only
+# the baseline's parameters, and Hypothesis draws the same threshold several
+# examples in a row, so keeping the last two reforms builds 16 over the 60
+# examples instead of holding one system per threshold drawn.
+@lru_cache(maxsize=2)
+def reformed_system(parameter, value):
     modifier = set_parameter(
         parameter, value, return_modifier=True, period="year:2000:100"
     )
@@ -52,7 +58,23 @@ class override(Reform):
         def apply(self):
             self.parameters = modifier(self.parameters)
 
-    return override(system)
+    return override(baseline_system())
+
+
+def tax_benefit_system(parameter, value):
+    """The UK system, with one parameter overridden for all years if given."""
+    if parameter is None:
+        return baseline_system()
+    return reformed_system(parameter, value)
+
+
+@pytest.fixture(scope="module", autouse=True)
+def release_systems():
+    # Each system also keeps its last simulation alive. Drop them once this
+    # module's tests are done, instead of for the rest of the test process.
+    yield
+    reformed_system.cache_clear()
+    baseline_system.cache_clear()
 
 
 @st.composite
diff --git a/policyengine_uk/tests/test_uc_work_related_groups_properties.py b/policyengine_uk/tests/test_uc_work_related_groups_properties.py
index 2b3c643ae7..573e820fe5 100644
--- a/policyengine_uk/tests/test_uc_work_related_groups_properties.py
+++ b/policyengine_uk/tests/test_uc_work_related_groups_properties.py
@@ -499,6 +499,14 @@ def qualifying_age_couples():
     return Simulation(situation=situation), keys
 
 
+@pytest.fixture(scope="module", autouse=True)
+def release_cached_simulation():
+    # Drop the cached simulation, and the system it built, once this module's
+    # tests are done, instead of for the rest of the test process.
+    yield
+    qualifying_age_couples.cache_clear()
+
+
 @pytest.mark.parametrize("year", QUALIFYING_AGE_YEARS)
 def test_regulation_89_1_a_follows_the_statute_for_every_birth_date(year):
     """The elder member is in section 19 exactly when the statute puts a woman

From a30aa90471d0331066ef5046f11577344a47f7a0 Mon Sep 17 00:00:00 2001
From: Max Ghenis <mghenis@gmail.com>
Date: Sun, 11 Oct 2026 17:42:41 -0400
Subject: [PATCH 2/4] Count cache_clear after a yield inside with or try in the
 guard

The guard only saw a yield written directly in the fixture body, so a
fixture that yields inside a with or try block (as reuse_country_model
does) was reported as not releasing its caches. Count any cache_clear()
after the fixture's first yield, and test decorator forms, that case and
imports from other test files. Drop the leading "- " from the changelog
fragment: towncrier's default template adds the bullet (#2162).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 changelog.d/2357.fixed.md                     |  2 +-
 .../test_test_module_caches_released.py       | 77 +++++++++++++------
 2 files changed, 56 insertions(+), 23 deletions(-)

diff --git a/changelog.d/2357.fixed.md b/changelog.d/2357.fixed.md
index c4906fa91c..cf545ad608 100644
--- a/changelog.d/2357.fixed.md
+++ b/changelog.d/2357.fixed.md
@@ -1 +1 @@
-- Release the tax-benefit systems and simulations that five test modules cache once each module's tests finish, build the two-child limit and cliff property tests' reforms over one shared baseline system, and check that every test module cache is released (#2357).
+Release the tax-benefit systems and simulations that five test modules cache once each module's tests finish, build the two-child limit and cliff property tests' reforms over one shared baseline system, and check that every test module cache is released (#2357).
diff --git a/policyengine_uk/tests/code_health/test_test_module_caches_released.py b/policyengine_uk/tests/code_health/test_test_module_caches_released.py
index 71e123d79e..c66dacc334 100644
--- a/policyengine_uk/tests/code_health/test_test_module_caches_released.py
+++ b/policyengine_uk/tests/code_health/test_test_module_caches_released.py
@@ -63,8 +63,8 @@ def _is_module_autouse_fixture(function: ast.FunctionDef) -> bool:
 
 
 def released_caches(tree: ast.Module) -> set[str]:
-    """Names whose cache_clear() runs after the yield of a module-scoped
-    autouse fixture."""
+    """Names whose cache_clear() comes after the yield of a module-scoped
+    autouse fixture, including a yield inside ``with`` or ``try``."""
     released = set()
     for function in tree.body:
         if not (
@@ -72,23 +72,20 @@ def released_caches(tree: ast.Module) -> set[str]:
             and _is_module_autouse_fixture(function)
         ):
             continue
-        after_yield = False
-        for statement in function.body:
-            if isinstance(statement, ast.Expr) and isinstance(
-                statement.value, ast.Yield
+        yields = [
+            node.lineno for node in ast.walk(function) if isinstance(node, ast.Yield)
+        ]
+        if not yields:
+            continue
+        for node in ast.walk(function):
+            if (
+                isinstance(node, ast.Call)
+                and isinstance(node.func, ast.Attribute)
+                and node.func.attr == "cache_clear"
+                and isinstance(node.func.value, ast.Name)
+                and node.lineno > min(yields)
             ):
-                after_yield = True
-                continue
-            if not after_yield:
-                continue
-            for node in ast.walk(statement):
-                if (
-                    isinstance(node, ast.Call)
-                    and isinstance(node.func, ast.Attribute)
-                    and node.func.attr == "cache_clear"
-                    and isinstance(node.func.value, ast.Name)
-                ):
-                    released.add(node.func.value.id)
+                released.add(node.func.value.id)
     return released
 
 
@@ -128,8 +125,9 @@ def test_every_test_module_cache_is_released():
     )
 
 
-def test_the_scan_finds_caches_and_releases(tmp_path):
+def test_the_scan_finds_caches_and_releases():
     source = """
+import functools
 from functools import cache, lru_cache
 
 import pytest
@@ -145,15 +143,50 @@ def simulation():
     return 1
 
 
+@functools.lru_cache
+def template():
+    return 2
+
+
+def not_cached():
+    return 3
+
+
 @pytest.fixture(scope="module", autouse=True)
 def release():
     systems.cache_clear()  # before the yield: does not count
     yield
     simulation.cache_clear()
+
+
+@pytest.fixture(autouse=True, scope="module")
+def release_after_a_yield_inside_with():
+    with open(__file__):
+        try:
+            yield
+        finally:
+            template.cache_clear()
 """
     tree = ast.parse(source)
-    assert cached_functions(tree) == {"systems", "simulation"}
-    assert released_caches(tree) == {"simulation"}
+    assert cached_functions(tree) == {"systems", "simulation", "template"}
+    assert released_caches(tree) == {"simulation", "template"}
 
-    function_scoped = source.replace('scope="module", ', "")
+    function_scoped = source.replace('scope="module"', 'scope="function"')
     assert released_caches(ast.parse(function_scoped)) == set()
+    not_autouse = source.replace("autouse=True", "autouse=False")
+    assert released_caches(ast.parse(not_autouse)) == set()
+
+
+def test_the_scan_follows_imports_from_other_test_files():
+    tree = ast.parse(
+        "from policyengine_uk.tests.test_state_pension_age import (\n"
+        "    YEARS,\n"
+        "    differential_simulation,\n"
+        ")\n"
+    )
+    imported = imported_test_modules(tree)
+    assert imported == {
+        TESTS / "test_state_pension_age.py": ["YEARS", "differential_simulation"]
+    }
+    caches = cached_functions(parse(TESTS / "test_state_pension_age.py"))
+    assert "differential_simulation" in caches

From da5a9c8807a28766b37a855f69413ff6c3c7f153 Mon Sep 17 00:00:00 2001
From: Max Ghenis <mghenis@gmail.com>
Date: Sun, 11 Oct 2026 18:55:10 -0400
Subject: [PATCH 3/4] Follow imports and helper functions in the cache guard;
 reword comments

From the two reviews of #2358 (both approve, no blocking findings):

- The guard now also sees async cached functions and
  `name = lru_cache(...)(function)`, follows aliased and relative imports
  from other test files, and requires a module to clear a cache it fills
  through a helper imported from another test file. It ignores a yield or
  cache_clear inside a function nested in the fixture, and only treats
  `lru_cache`/`cache` and `functools.*` as caches. Its docstring says what a
  static scan cannot see.
- The two-child comment states why a reform is built only when the
  threshold changes instead of quoting a count that depends on the
  Hypothesis version.
- The SPC comment says which module the import loads.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../test_test_module_caches_released.py       | 272 +++++++++++++-----
 ...est_state_pension_credit_qualifying_age.py |  11 +-
 ...wo_child_limit_age_exemption_properties.py |   8 +-
 3 files changed, 215 insertions(+), 76 deletions(-)

diff --git a/policyengine_uk/tests/code_health/test_test_module_caches_released.py b/policyengine_uk/tests/code_health/test_test_module_caches_released.py
index c66dacc334..425c1b7585 100644
--- a/policyengine_uk/tests/code_health/test_test_module_caches_released.py
+++ b/policyengine_uk/tests/code_health/test_test_module_caches_released.py
@@ -6,57 +6,122 @@
 it holds for the rest of the test process, and a system also keeps the last
 simulation built from it (``Simulation.__init__`` sets
 ``tax_benefit_system.simulation``). One such cache kept 13 systems, about
-0.5 GiB, in a CI worker for the rest of its run.
-
-So every function a test module caches with ``functools.lru_cache`` or
-``functools.cache``, whether it defines it or imports it from another test
-module, must have its ``cache_clear()`` called after the ``yield`` of a
-module-scoped autouse fixture in that module. An imported cache counts
-separately: pytest imports test files by bare name, so a module that imports
-from another test file loads a second copy of it, with its own caches.
+0.5 GiB, in a CI worker for the rest of its run (run 38165445604, #2357).
+
+So a test module must call ``cache_clear()``, after the ``yield`` of a
+module-scoped autouse fixture, on every functools cache that it:
+
+- defines at module level (a decorated function, or
+  ``name = lru_cache(...)(function)``);
+- imports from another test file; or
+- fills through a function it imports from another test file, where that
+  function reaches the cache through that file's own module-level functions.
+
+An imported cache counts separately because pytest collects the files in this
+directory by bare name (it has no ``__init__.py``), so
+``from policyengine_uk.tests.X import f`` loads a second copy of ``X`` with its
+own caches, which ``X``'s fixtures never clear.
+
+This is a static scan. It does not see cached methods, caches reached through
+``import module`` rather than ``from module import name``, or whether a
+``cache_clear()`` written after the ``yield`` runs (one in an ``except`` branch
+counts).
 """
 
 import ast
+from collections.abc import Iterator
 from pathlib import Path
 
 TESTS = Path(__file__).resolve().parents[1]
-CACHE_DECORATORS = {"lru_cache", "cache"}
+CACHES = {"lru_cache", "cache"}
+FUNCTIONS = (ast.FunctionDef, ast.AsyncFunctionDef)
 
 
-def _name(node: ast.AST) -> str | None:
+def _is_functools_cache(node: ast.AST) -> bool:
+    """``lru_cache``, ``cache``, ``functools.lru_cache``, called or not."""
     if isinstance(node, ast.Call):
         node = node.func
     if isinstance(node, ast.Attribute):
-        return node.attr
-    if isinstance(node, ast.Name):
-        return node.id
-    return None
+        return (
+            node.attr in CACHES
+            and isinstance(node.value, ast.Name)
+            and node.value.id == "functools"
+        )
+    return isinstance(node, ast.Name) and node.id in CACHES
 
 
-def cached_functions(tree: ast.Module) -> set[str]:
-    """Module-level functions decorated with a functools cache."""
-    return {
-        node.name
-        for node in tree.body
-        if isinstance(node, ast.FunctionDef)
-        and any(_name(d) in CACHE_DECORATORS for d in node.decorator_list)
-    }
+def module_functions(tree: ast.Module) -> dict[str, ast.AST]:
+    return {node.name: node for node in tree.body if isinstance(node, FUNCTIONS)}
 
 
-def _keyword(call: ast.Call, name: str):
-    for keyword in call.keywords:
-        if keyword.arg == name and isinstance(keyword.value, ast.Constant):
-            return keyword.value.value
-    return None
+def cached_functions(tree: ast.Module) -> set[str]:
+    """Module-level names bound to a functools cache."""
+    cached = {
+        name
+        for name, function in module_functions(tree).items()
+        if any(_is_functools_cache(d) for d in function.decorator_list)
+    }
+    for node in tree.body:
+        # name = lru_cache(maxsize=1)(function) or name = cache(function)
+        if (
+            isinstance(node, ast.Assign)
+            and isinstance(node.value, ast.Call)
+            and (
+                _is_functools_cache(node.value.func) or _is_functools_cache(node.value)
+            )
+        ):
+            cached |= {t.id for t in node.targets if isinstance(t, ast.Name)}
+    return cached
 
 
-def _is_module_autouse_fixture(function: ast.FunctionDef) -> bool:
-    for decorator in function.decorator_list:
+def reached_caches(tree: ast.Module, names: set[str]) -> set[str]:
+    """Caches of this module that calling the named functions can fill: the
+    names themselves, and any cache their bodies reach through the module's
+    own functions."""
+    functions = module_functions(tree)
+    caches = cached_functions(tree)
+    seen: set[str] = set()
+    pending = [name for name in names if name in functions or name in caches]
+    while pending:
+        name = pending.pop()
+        if name in seen:
+            continue
+        seen.add(name)
+        if name in functions:
+            pending += [
+                node.id
+                for node in ast.walk(functions[name])
+                if isinstance(node, ast.Name)
+                and (node.id in functions or node.id in caches)
+            ]
+    return seen & caches
+
+
+def _own_nodes(function: ast.AST) -> Iterator[ast.AST]:
+    """The function's nodes, without the bodies of functions nested in it."""
+    pending = list(ast.iter_child_nodes(function))
+    while pending:
+        node = pending.pop()
+        yield node
+        if not isinstance(node, (*FUNCTIONS, ast.Lambda, ast.ClassDef)):
+            pending.extend(ast.iter_child_nodes(node))
+
+
+def _is_module_autouse_fixture(function: ast.AST) -> bool:
+    for decorator in getattr(function, "decorator_list", []):
+        if not isinstance(decorator, ast.Call):
+            continue
+        func = decorator.func
+        name = func.attr if isinstance(func, ast.Attribute) else getattr(func, "id", "")
+        keywords = {
+            keyword.arg: keyword.value.value
+            for keyword in decorator.keywords
+            if isinstance(keyword.value, ast.Constant)
+        }
         if (
-            isinstance(decorator, ast.Call)
-            and _name(decorator) == "fixture"
-            and _keyword(decorator, "scope") == "module"
-            and _keyword(decorator, "autouse") is True
+            name == "fixture"
+            and keywords.get("scope") == "module"
+            and keywords.get("autouse") is True
         ):
             return True
     return False
@@ -66,18 +131,14 @@ def released_caches(tree: ast.Module) -> set[str]:
     """Names whose cache_clear() comes after the yield of a module-scoped
     autouse fixture, including a yield inside ``with`` or ``try``."""
     released = set()
-    for function in tree.body:
-        if not (
-            isinstance(function, ast.FunctionDef)
-            and _is_module_autouse_fixture(function)
-        ):
+    for function in module_functions(tree).values():
+        if not _is_module_autouse_fixture(function):
             continue
-        yields = [
-            node.lineno for node in ast.walk(function) if isinstance(node, ast.Yield)
-        ]
+        nodes = list(_own_nodes(function))
+        yields = [node.lineno for node in nodes if isinstance(node, ast.Yield)]
         if not yields:
             continue
-        for node in ast.walk(function):
+        for node in nodes:
             if (
                 isinstance(node, ast.Call)
                 and isinstance(node.func, ast.Attribute)
@@ -89,15 +150,24 @@ def released_caches(tree: ast.Module) -> set[str]:
     return released
 
 
-def imported_test_modules(tree: ast.Module) -> dict[Path, list[str]]:
-    """Other test files this module imports names from."""
-    imports = {}
+def imported_test_files(tree: ast.Module, path: Path) -> dict[Path, dict[str, str]]:
+    """Other test files this module imports names from: for each file, the
+    imported names and what this module calls them."""
+    imports: dict[Path, dict[str, str]] = {}
     for node in tree.body:
-        if isinstance(node, ast.ImportFrom) and node.module:
-            module = node.module.removeprefix("policyengine_uk.tests.")
-            path = TESTS / (module.replace(".", "/") + ".py")
-            if path.exists():
-                imports[path] = [alias.name for alias in node.names]
+        if not isinstance(node, ast.ImportFrom):
+            continue
+        module = (node.module or "").removeprefix("policyengine_uk.tests.")
+        if not module:
+            continue
+        relative = Path(*module.split(".")).with_suffix(".py")
+        bases = [path.parents[node.level - 1]] if node.level else [TESTS, path.parent]
+        for base in bases:
+            if (base / relative).is_file():
+                imports.setdefault(base / relative, {}).update(
+                    {alias.name: alias.asname or alias.name for alias in node.names}
+                )
+                break
     return imports
 
 
@@ -108,8 +178,15 @@ def parse(path: Path) -> ast.Module:
 def unreleased_caches(path: Path) -> list[str]:
     tree = parse(path)
     caches = cached_functions(tree)
-    for imported, names in imported_test_modules(tree).items():
-        caches |= cached_functions(parse(imported)) & set(names)
+    for imported, names in imported_test_files(tree, path).items():
+        if imported == path:
+            continue
+        # A cache this module fills without importing it is reported under
+        # its own name: import it to clear it.
+        caches |= {
+            names.get(cache, cache)
+            for cache in reached_caches(parse(imported), set(names))
+        }
     return sorted(caches - released_caches(tree))
 
 
@@ -125,8 +202,7 @@ def test_every_test_module_cache_is_released():
     )
 
 
-def test_the_scan_finds_caches_and_releases():
-    source = """
+SCANNED = """
 import functools
 from functools import cache, lru_cache
 
@@ -144,18 +220,39 @@ def simulation():
 
 
 @functools.lru_cache
-def template():
+async def template():
     return 2
 
 
-def not_cached():
+def build():
     return 3
 
 
+assigned = lru_cache(maxsize=1)(build)
+wrapped = cache(build)
+not_a_functools_cache = registry.cache(build)
+
+
+def helper():
+    return indirect()
+
+
+def indirect():
+    return simulation()
+
+
+def pure():
+    return build()
+
+
 @pytest.fixture(scope="module", autouse=True)
 def release():
     systems.cache_clear()  # before the yield: does not count
     yield
+
+    def never_called():
+        wrapped.cache_clear()  # in a nested function: does not count
+
     simulation.cache_clear()
 
 
@@ -167,26 +264,67 @@ def release_after_a_yield_inside_with():
         finally:
             template.cache_clear()
 """
-    tree = ast.parse(source)
-    assert cached_functions(tree) == {"systems", "simulation", "template"}
+
+
+def test_the_scan_finds_caches_and_releases():
+    tree = ast.parse(SCANNED)
+    assert cached_functions(tree) == {
+        "systems",
+        "simulation",
+        "template",
+        "assigned",
+        "wrapped",
+    }
     assert released_caches(tree) == {"simulation", "template"}
 
-    function_scoped = source.replace('scope="module"', 'scope="function"')
+    function_scoped = SCANNED.replace('scope="module"', 'scope="function"')
     assert released_caches(ast.parse(function_scoped)) == set()
-    not_autouse = source.replace("autouse=True", "autouse=False")
+    not_autouse = SCANNED.replace("autouse=True", "autouse=False")
     assert released_caches(ast.parse(not_autouse)) == set()
 
 
-def test_the_scan_follows_imports_from_other_test_files():
+def test_the_scan_follows_functions_to_the_caches_they_fill():
+    tree = ast.parse(SCANNED)
+    assert reached_caches(tree, {"helper"}) == {"simulation"}
+    assert reached_caches(tree, {"systems", "pure"}) == {"systems"}
+    assert reached_caches(tree, {"pure", "not_defined_here"}) == set()
+
+
+def test_the_scan_follows_imports_from_other_test_files(tmp_path):
+    (tmp_path / "test_source.py").write_text(SCANNED)
+    user = tmp_path / "test_user.py"
+    for imports, expected in [
+        # An imported cache, under the name this module gives it.
+        ("from test_source import systems as cached_systems", ["cached_systems"]),
+        # A cache filled through an imported helper, which is not imported.
+        ("from test_source import helper", ["simulation"]),
+        ("from .test_source import helper, simulation", ["simulation"]),
+        ("from test_source import pure", []),
+    ]:
+        user.write_text(imports + "\n")
+        assert unreleased_caches(user) == expected, imports
+
+    user.write_text(
+        "import pytest\n"
+        "from test_source import helper, simulation as shared\n\n\n"
+        '@pytest.fixture(scope="module", autouse=True)\n'
+        "def release():\n"
+        "    yield\n"
+        "    shared.cache_clear()\n"
+    )
+    assert unreleased_caches(user) == []
+
+
+def test_the_scan_resolves_this_directorys_imports():
     tree = ast.parse(
         "from policyengine_uk.tests.test_state_pension_age import (\n"
         "    YEARS,\n"
-        "    differential_simulation,\n"
+        "    differential_simulation as cached,\n"
         ")\n"
     )
-    imported = imported_test_modules(tree)
-    assert imported == {
-        TESTS / "test_state_pension_age.py": ["YEARS", "differential_simulation"]
+    assert imported_test_files(tree, TESTS / "test_x.py") == {
+        TESTS / "test_state_pension_age.py": {
+            "YEARS": "YEARS",
+            "differential_simulation": "cached",
+        }
     }
-    caches = cached_functions(parse(TESTS / "test_state_pension_age.py"))
-    assert "differential_simulation" in caches
diff --git a/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py b/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py
index 3ae1020499..6664f53b19 100644
--- a/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py
+++ b/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py
@@ -83,11 +83,12 @@ def property_simulation(people: dict):
 
 @pytest.fixture(scope="module", autouse=True)
 def release_cached_simulations():
-    # Pytest imports test files here by bare name (the directory has no
-    # __init__.py), so the import above loads test_state_pension_age a second
-    # time and this module fills its own differential_simulation cache. Drop
-    # the cached simulations once this module's tests are done, instead of for
-    # the rest of the test process.
+    # The import above loads policyengine_uk.tests.test_state_pension_age, a
+    # separate module from the test_state_pension_age that pytest collects by
+    # bare name (this directory has no __init__.py). Its differential_simulation
+    # cache is filled here, so it is cleared here. Drop the cached simulations
+    # once this module's tests are done, instead of for the rest of the test
+    # process.
     yield
     differential_simulation.cache_clear()
     one_person_units_template.cache_clear()
diff --git a/policyengine_uk/tests/test_two_child_limit_age_exemption_properties.py b/policyengine_uk/tests/test_two_child_limit_age_exemption_properties.py
index 6c2835d2d8..cc636ab030 100644
--- a/policyengine_uk/tests/test_two_child_limit_age_exemption_properties.py
+++ b/policyengine_uk/tests/test_two_child_limit_age_exemption_properties.py
@@ -44,10 +44,10 @@ def baseline_system():
     return CountryTaxBenefitSystem()
 
 
-# Each example uses the repealed limit and one threshold. A reform clones only
-# the baseline's parameters, and Hypothesis draws the same threshold several
-# examples in a row, so keeping the last two reforms builds 16 over the 60
-# examples instead of holding one system per threshold drawn.
+# Each example uses the repealed limit, then one threshold, so the last two
+# reforms are the repealed limit's and the latest threshold's: a reform is
+# built only when the threshold changes, instead of one system being held per
+# threshold drawn. A reform clones only the baseline's parameters.
 @lru_cache(maxsize=2)
 def reformed_system(parameter, value):
     modifier = set_parameter(

From 5420c0656be60e4954d79dd8370c6948c41028e0 Mon Sep 17 00:00:00 2001
From: Max Ghenis <mghenis@gmail.com>
Date: Sun, 11 Oct 2026 19:46:57 -0400
Subject: [PATCH 4/4] Tighten the cache guard after the re-review

From the GPT-6.1 Sol re-review of #2358 at da5a9c880 (approve, nits only):

- A cache made by assignment (`name = cache(function)`) is followed into
  the function it wraps, and `name = lru_cache(maxsize=1)` is not taken for
  a cache.
- A cache imported under several names is released by clearing any of them.
- The self-tests cover a relative import from a subdirectory, which only
  the import's level resolves. Seven mutations of the scan each fail a
  self-test.
- The docstring and the SPC comment say the second module copy arises under
  pytest's default import mode, and the docstring lists re-exports through a
  third test file among what the scan does not see.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../test_test_module_caches_released.py       | 132 ++++++++++++------
 ...est_state_pension_credit_qualifying_age.py |  12 +-
 2 files changed, 99 insertions(+), 45 deletions(-)

diff --git a/policyengine_uk/tests/code_health/test_test_module_caches_released.py b/policyengine_uk/tests/code_health/test_test_module_caches_released.py
index 425c1b7585..d3059879f5 100644
--- a/policyengine_uk/tests/code_health/test_test_module_caches_released.py
+++ b/policyengine_uk/tests/code_health/test_test_module_caches_released.py
@@ -17,13 +17,14 @@
 - fills through a function it imports from another test file, where that
   function reaches the cache through that file's own module-level functions.
 
-An imported cache counts separately because pytest collects the files in this
-directory by bare name (it has no ``__init__.py``), so
-``from policyengine_uk.tests.X import f`` loads a second copy of ``X`` with its
-own caches, which ``X``'s fixtures never clear.
+An imported cache counts separately because, under pytest's default (prepend)
+import mode, pytest collects the files in this directory by bare name (it has
+no ``__init__.py``), so ``from policyengine_uk.tests.X import f`` loads a second
+copy of ``X`` with its own caches, which ``X``'s fixtures never clear.
 
 This is a static scan. It does not see cached methods, caches reached through
-``import module`` rather than ``from module import name``, or whether a
+``import module`` rather than ``from module import name``, a cache re-exported
+through a third test file or reached through a helper in one, or whether a
 ``cache_clear()`` written after the ``yield`` runs (one in an ``except`` branch
 counts).
 """
@@ -54,47 +55,63 @@ def module_functions(tree: ast.Module) -> dict[str, ast.AST]:
     return {node.name: node for node in tree.body if isinstance(node, FUNCTIONS)}
 
 
-def cached_functions(tree: ast.Module) -> set[str]:
-    """Module-level names bound to a functools cache."""
-    cached = {
-        name
-        for name, function in module_functions(tree).items()
+def cache_bodies(tree: ast.Module) -> dict[str, ast.AST | None]:
+    """Module-level names bound to a functools cache, each with the module
+    function it wraps (None when it wraps something else)."""
+    functions = module_functions(tree)
+    caches: dict[str, ast.AST | None] = {
+        name: function
+        for name, function in functions.items()
         if any(_is_functools_cache(d) for d in function.decorator_list)
     }
     for node in tree.body:
-        # name = lru_cache(maxsize=1)(function) or name = cache(function)
-        if (
-            isinstance(node, ast.Assign)
-            and isinstance(node.value, ast.Call)
-            and (
-                _is_functools_cache(node.value.func) or _is_functools_cache(node.value)
+        if not (isinstance(node, ast.Assign) and isinstance(node.value, ast.Call)):
+            continue
+        call = node.value
+        # name = lru_cache(maxsize=1)(function) or name = cache(function), but
+        # not name = lru_cache(maxsize=1), which is a decorator, not a cache.
+        if isinstance(call.func, ast.Call):
+            applied = _is_functools_cache(call.func)
+        else:
+            applied = (
+                _is_functools_cache(call)
+                and bool(call.args)
+                and not isinstance(call.args[0], ast.Constant)
             )
-        ):
-            cached |= {t.id for t in node.targets if isinstance(t, ast.Name)}
-    return cached
+        if applied:
+            wrapped = call.args[0] if call.args else None
+            body = functions.get(wrapped.id) if isinstance(wrapped, ast.Name) else None
+            caches.update({t.id: body for t in node.targets if isinstance(t, ast.Name)})
+    return caches
+
+
+def cached_functions(tree: ast.Module) -> set[str]:
+    return set(cache_bodies(tree))
 
 
 def reached_caches(tree: ast.Module, names: set[str]) -> set[str]:
     """Caches of this module that calling the named functions can fill: the
     names themselves, and any cache their bodies reach through the module's
     own functions."""
-    functions = module_functions(tree)
-    caches = cached_functions(tree)
+    caches = cache_bodies(tree)
+    bodies = module_functions(tree) | {
+        name: body for name, body in caches.items() if body is not None
+    }
     seen: set[str] = set()
-    pending = [name for name in names if name in functions or name in caches]
+    pending = [name for name in names if name in bodies or name in caches]
     while pending:
         name = pending.pop()
         if name in seen:
             continue
         seen.add(name)
-        if name in functions:
+        if name in bodies:
             pending += [
                 node.id
-                for node in ast.walk(functions[name])
+                for node in ast.walk(bodies[name])
                 if isinstance(node, ast.Name)
-                and (node.id in functions or node.id in caches)
+                and (node.id in bodies or node.id in caches)
             ]
-    return seen & caches
+    return seen & set(caches)
 
 
 def _own_nodes(function: ast.AST) -> Iterator[ast.AST]:
@@ -150,10 +167,12 @@ def released_caches(tree: ast.Module) -> set[str]:
     return released
 
 
-def imported_test_files(tree: ast.Module, path: Path) -> dict[Path, dict[str, str]]:
+def imported_test_files(
+    tree: ast.Module, path: Path
+) -> dict[Path, dict[str, set[str]]]:
     """Other test files this module imports names from: for each file, the
-    imported names and what this module calls them."""
-    imports: dict[Path, dict[str, str]] = {}
+    imported names and everything this module calls each of them."""
+    imports: dict[Path, dict[str, set[str]]] = {}
     for node in tree.body:
         if not isinstance(node, ast.ImportFrom):
             continue
@@ -164,9 +183,9 @@ def imported_test_files(tree: ast.Module, path: Path) -> dict[Path, dict[str, st
         bases = [path.parents[node.level - 1]] if node.level else [TESTS, path.parent]
         for base in bases:
             if (base / relative).is_file():
-                imports.setdefault(base / relative, {}).update(
-                    {alias.name: alias.asname or alias.name for alias in node.names}
-                )
+                names = imports.setdefault(base / relative, {})
+                for alias in node.names:
+                    names.setdefault(alias.name, set()).add(alias.asname or alias.name)
                 break
     return imports
 
@@ -177,17 +196,19 @@ def parse(path: Path) -> ast.Module:
 
 def unreleased_caches(path: Path) -> list[str]:
     tree = parse(path)
-    caches = cached_functions(tree)
+    # Each cache, as the names this module can clear it by.
+    caches = [{name} for name in cached_functions(tree)]
     for imported, names in imported_test_files(tree, path).items():
         if imported == path:
             continue
         # A cache this module fills without importing it is reported under
         # its own name: import it to clear it.
-        caches |= {
-            names.get(cache, cache)
+        caches += [
+            names.get(cache, {cache})
             for cache in reached_caches(parse(imported), set(names))
-        }
-    return sorted(caches - released_caches(tree))
+        ]
+    released = released_caches(tree)
+    return sorted({min(names) for names in caches if not names & released})
 
 
 def test_every_test_module_cache_is_released():
@@ -245,6 +266,14 @@ def pure():
     return build()
 
 
+def fills_simulation():
+    return simulation()
+
+
+through_a_wrapper = cache(fills_simulation)
+a_decorator = lru_cache(maxsize=1)
+
+
 @pytest.fixture(scope="module", autouse=True)
 def release():
     systems.cache_clear()  # before the yield: does not count
@@ -274,6 +303,7 @@ def test_the_scan_finds_caches_and_releases():
         "template",
         "assigned",
         "wrapped",
+        "through_a_wrapper",
     }
     assert released_caches(tree) == {"simulation", "template"}
 
@@ -287,6 +317,11 @@ def test_the_scan_follows_functions_to_the_caches_they_fill():
     tree = ast.parse(SCANNED)
     assert reached_caches(tree, {"helper"}) == {"simulation"}
     assert reached_caches(tree, {"systems", "pure"}) == {"systems"}
+    # An assigned cache also fills what the function it wraps reaches.
+    assert reached_caches(tree, {"through_a_wrapper"}) == {
+        "through_a_wrapper",
+        "simulation",
+    }
     assert reached_caches(tree, {"pure", "not_defined_here"}) == set()
 
 
@@ -304,6 +339,25 @@ def test_the_scan_follows_imports_from_other_test_files(tmp_path):
         user.write_text(imports + "\n")
         assert unreleased_caches(user) == expected, imports
 
+    # A relative import from a subdirectory, which only the import's level
+    # resolves.
+    nested = tmp_path / "nested" / "test_user.py"
+    nested.parent.mkdir()
+    nested.write_text("from ..test_source import helper\n")
+    assert unreleased_caches(nested) == ["simulation"]
+
+    # A cache is released under any of the names this module imports it by.
+    user.write_text(
+        "import pytest\n"
+        "from test_source import systems as a\n"
+        "from test_source import systems as b\n\n\n"
+        '@pytest.fixture(scope="module", autouse=True)\n'
+        "def release():\n"
+        "    yield\n"
+        "    a.cache_clear()\n"
+    )
+    assert unreleased_caches(user) == []
+
     user.write_text(
         "import pytest\n"
         "from test_source import helper, simulation as shared\n\n\n"
@@ -324,7 +378,7 @@ def test_the_scan_resolves_this_directorys_imports():
     )
     assert imported_test_files(tree, TESTS / "test_x.py") == {
         TESTS / "test_state_pension_age.py": {
-            "YEARS": "YEARS",
-            "differential_simulation": "cached",
+            "YEARS": {"YEARS"},
+            "differential_simulation": {"cached"},
         }
     }
diff --git a/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py b/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py
index 6664f53b19..3dc872d96a 100644
--- a/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py
+++ b/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py
@@ -83,12 +83,12 @@ def property_simulation(people: dict):
 
 @pytest.fixture(scope="module", autouse=True)
 def release_cached_simulations():
-    # The import above loads policyengine_uk.tests.test_state_pension_age, a
-    # separate module from the test_state_pension_age that pytest collects by
-    # bare name (this directory has no __init__.py). Its differential_simulation
-    # cache is filled here, so it is cleared here. Drop the cached simulations
-    # once this module's tests are done, instead of for the rest of the test
-    # process.
+    # Under pytest's default import mode, the import above loads
+    # policyengine_uk.tests.test_state_pension_age, a separate module from the
+    # test_state_pension_age that pytest collects by bare name (this directory
+    # has no __init__.py). Its differential_simulation cache is filled here, so
+    # it is cleared here. Drop the cached simulations once this module's tests
+    # are done, instead of for the rest of the test process.
     yield
     differential_simulation.cache_clear()
     one_person_units_template.cache_clear()