refactor(tests): add import-state isolation helper

Test-only refactor continuing #2523. Adds a shared import-state isolation helper with focused coverage and migrates two pilot tests that manually preserved sys.modules and parent package attributes.

tests/helpers/import_state.py

@@ -0,0 +1,87 @@
+"""Shared helper for saving and restoring Python import state in tests.
+
+Use ``preserve_import_state`` as a context manager around any block that needs
+to mutate ``sys.modules`` or parent-package attributes temporarily. On exit
+(normal or exception), every named module is restored to exactly the state it
+had before the block — present, absent, or carrying a parent-package attribute.
+
+Use ``clear_module`` to drop a single module from both ``sys.modules`` and its
+parent-package attribute (e.g. before forcing a fresh import inside the block).
+
+Background: importing ``routes.session_routes`` also sets ``session_routes`` on
+the parent ``routes`` package object. A ``from routes import session_routes``
+or ``import routes.session_routes as X`` statement resolves through that parent
+attribute, so restoring ``sys.modules`` alone is not sufficient — the parent
+attribute must be restored too. This helper handles both.
+
+Restoration in ``preserve_import_state`` is two-phased: all ``sys.modules``
+entries are written back first, then all parent-package attributes. This means
+parent-attr restoration always resolves the parent through the already-restored
+``sys.modules``, so results are deterministic regardless of argument order —
+safe for callers that pass both a parent package and a child module.
+"""
+
+import sys
+from contextlib import contextmanager
+
+_ABSENT = object()
+
+
+def _save_one(dotted_name):
+    saved_mod = sys.modules.get(dotted_name, _ABSENT)
+    pkg_name, _, attr = dotted_name.rpartition(".")
+    pkg = sys.modules.get(pkg_name)
+    saved_attr = getattr(pkg, attr, _ABSENT) if pkg is not None else _ABSENT
+    return saved_mod, saved_attr
+
+
+def _restore_parent_attr(dotted_name, saved_attr):
+    pkg_name, _, attr = dotted_name.rpartition(".")
+    pkg = sys.modules.get(pkg_name)
+    if pkg is None:
+        return
+    if saved_attr is _ABSENT:
+        if hasattr(pkg, attr):
+            delattr(pkg, attr)
+    else:
+        setattr(pkg, attr, saved_attr)
+
+
+def _restore_one(dotted_name, saved_mod, saved_attr):
+    if saved_mod is _ABSENT:
+        sys.modules.pop(dotted_name, None)
+    else:
+        sys.modules[dotted_name] = saved_mod
+    _restore_parent_attr(dotted_name, saved_attr)
+
+
+def clear_module(dotted_name):
+    """Remove a module from sys.modules and its parent-package attribute."""
+    _restore_one(dotted_name, _ABSENT, _ABSENT)
+
+
+@contextmanager
+def preserve_import_state(*module_names):
+    """Save and restore sys.modules entries and parent-package attributes.
+
+    Restoration is two-phased: sys.modules entries are written back first,
+    then parent-package attributes. This ensures parent-attr restoration always
+    sees the correctly restored parent in sys.modules, regardless of argument
+    order — safe for callers that pass both a parent and a child module.
+
+    On exit (normal or exception), each named module is restored to its state
+    before the block — whether present, absent, or carrying a parent attribute.
+    """
+    saved = {name: _save_one(name) for name in module_names}
+    try:
+        yield
+    finally:
+        # Phase 1: restore all sys.modules entries.
+        for name, (saved_mod, _) in saved.items():
+            if saved_mod is _ABSENT:
+                sys.modules.pop(name, None)
+            else:
+                sys.modules[name] = saved_mod
+        # Phase 2: restore all parent-package attributes.
+        for name, (_, saved_attr) in saved.items():
+            _restore_parent_attr(name, saved_attr)

tests/test_blind_compare_redaction.py

@@ -22,77 +22,23 @@ import importlib
 from pathlib import Path
 from unittest.mock import MagicMock
 
+from tests.helpers.import_state import clear_module, preserve_import_state
+
 _REPO = Path(__file__).resolve().parent.parent
 
-# Mirror tests/test_session_ghost_delete.py exactly: stub only the ORM *class*
+# Stub only the ORM class modules and import the real core.session_manager so
-# modules and import the REAL core.session_manager + src.auth_helpers. pytest
+# the cached routes.session_routes is identical regardless of collection order.
-# caches routes.session_routes after the first import, so stubbing auth_helpers /
+# preserve_import_state restores both sys.modules and parent-package attributes
-# session_manager here would poison the shared module for the sibling session
+# after the block, preventing stub leakage into siblings.
-# tests (whichever file is collected first wins). Matching their stub set keeps
-# the cached module identical regardless of collection order. We restore both
-# sys.modules AND the parent `routes` package attribute so the stub-bound module
-# never leaks into sibling modules via `import routes.session_routes as X`.
-_ABSENT = object()
-
-
-def _save_module_and_parent_attr(dotted_name):
-    """Capture a module's sys.modules entry *and* its parent-package attribute.
-
-    Importing ``routes.session_routes`` also sets ``session_routes`` on the
-    parent ``routes`` package object, and ``import routes.session_routes as X``
-    resolves ``X`` through that parent attribute — so restoring sys.modules
-    alone leaves the stale stub-bound module reachable. Returns a (module, attr)
-    pair to hand back to _restore_module_and_parent_attr.
-    """
-    saved_module = sys.modules.get(dotted_name, _ABSENT)
-    pkg_name, _, attr = dotted_name.rpartition(".")
-    pkg = sys.modules.get(pkg_name)
-    saved_attr = getattr(pkg, attr, _ABSENT) if pkg is not None else _ABSENT
-    return saved_module, saved_attr
-
-
-def _restore_module_and_parent_attr(dotted_name, saved_module, saved_attr):
-    """Restore (or remove) both the sys.modules entry and the parent attribute.
-
-    Passing _ABSENT for both clears the cache, which is how we drop any stale
-    entry before the stubbed import.
-    """
-    if saved_module is _ABSENT:
-        sys.modules.pop(dotted_name, None)
-    else:
-        sys.modules[dotted_name] = saved_module
-    pkg_name, _, attr = dotted_name.rpartition(".")
-    pkg = sys.modules.get(pkg_name)
-    if pkg is None:
-        return
-    if saved_attr is _ABSENT:
-        if hasattr(pkg, attr):
-            delattr(pkg, attr)
-    else:
-        setattr(pkg, attr, saved_attr)
-
-
 _TEMP_STUBS = ("core.database", "core.models")
-_saved = {name: sys.modules.get(name, _ABSENT) for name in _TEMP_STUBS}
+with preserve_import_state(*_TEMP_STUBS, "core.session_manager", "routes.session_routes"):
-_saved["core.session_manager"] = sys.modules.get("core.session_manager", _ABSENT)
-_sr_saved = _save_module_and_parent_attr("routes.session_routes")
-try:
     for _name in _TEMP_STUBS:
         sys.modules[_name] = MagicMock(name=_name)
     if isinstance(sys.modules.get("core.session_manager"), MagicMock):
         del sys.modules["core.session_manager"]
-    # Clear the sys.modules entry AND the parent `routes` attribute so the
+    clear_module("routes.session_routes")
-    # stubbed import below produces a fresh module with no stale binding behind it.
-    _restore_module_and_parent_attr("routes.session_routes", _ABSENT, _ABSENT)
     importlib.import_module("core.session_manager")
     import routes.session_routes as SR  # noqa: E402
-finally:
-    for _name, _val in _saved.items():
-        if _val is _ABSENT:
-            sys.modules.pop(_name, None)
-        else:
-            sys.modules[_name] = _val
-    _restore_module_and_parent_attr("routes.session_routes", *_sr_saved)
 
 
 # ── backend: GET /api/sessions model redaction ─────────────────────────────

tests/test_helpers_import_state.py

@@ -0,0 +1,141 @@
+"""Focused tests for tests/helpers/import_state.py."""
+import sys
+import types
+
+import pytest
+
+from tests.helpers.import_state import clear_module, preserve_import_state
+
+_SENTINEL = "tests._import_state_test_sentinel"
+
+
+def test_absent_module_is_removed_after_block():
+    assert _SENTINEL not in sys.modules
+    with preserve_import_state(_SENTINEL):
+        sys.modules[_SENTINEL] = types.ModuleType(_SENTINEL)
+    assert _SENTINEL not in sys.modules
+
+
+def test_present_module_is_restored_after_block():
+    original = types.ModuleType(_SENTINEL)
+    sys.modules[_SENTINEL] = original
+    try:
+        with preserve_import_state(_SENTINEL):
+            sys.modules[_SENTINEL] = types.ModuleType(_SENTINEL)
+        assert sys.modules[_SENTINEL] is original
+    finally:
+        sys.modules.pop(_SENTINEL, None)
+
+
+def test_parent_attr_restored_when_present_before_block():
+    fake_parent = types.ModuleType("_fake_istate_parent")
+    fake_child = types.ModuleType("_fake_istate_parent.child")
+    fake_parent.child = fake_child
+    sys.modules["_fake_istate_parent"] = fake_parent
+    sys.modules["_fake_istate_parent.child"] = fake_child
+    try:
+        with preserve_import_state("_fake_istate_parent.child"):
+            replacement = types.ModuleType("_fake_istate_parent.child")
+            sys.modules["_fake_istate_parent.child"] = replacement
+            fake_parent.child = replacement
+        assert sys.modules["_fake_istate_parent.child"] is fake_child
+        assert fake_parent.child is fake_child
+    finally:
+        sys.modules.pop("_fake_istate_parent", None)
+        sys.modules.pop("_fake_istate_parent.child", None)
+
+
+def test_parent_attr_removed_when_absent_before_block():
+    fake_parent = types.ModuleType("_fake_istate_parent")
+    sys.modules["_fake_istate_parent"] = fake_parent
+    try:
+        with preserve_import_state("_fake_istate_parent.child"):
+            fake_child = types.ModuleType("_fake_istate_parent.child")
+            sys.modules["_fake_istate_parent.child"] = fake_child
+            fake_parent.child = fake_child
+        assert "_fake_istate_parent.child" not in sys.modules
+        assert not hasattr(fake_parent, "child")
+    finally:
+        sys.modules.pop("_fake_istate_parent", None)
+        sys.modules.pop("_fake_istate_parent.child", None)
+
+
+def test_state_restored_on_exception():
+    assert _SENTINEL not in sys.modules
+    with pytest.raises(RuntimeError, match="expected"):
+        with preserve_import_state(_SENTINEL):
+            sys.modules[_SENTINEL] = types.ModuleType(_SENTINEL)
+            raise RuntimeError("expected")
+    assert _SENTINEL not in sys.modules
+
+
+def test_multiple_modules_all_restored():
+    names = [f"tests._istate_multi_{i}" for i in range(3)]
+    for n in names:
+        assert n not in sys.modules
+    with preserve_import_state(*names):
+        for n in names:
+            sys.modules[n] = types.ModuleType(n)
+    for n in names:
+        assert n not in sys.modules
+
+
+def test_clear_module_removes_entry():
+    sys.modules[_SENTINEL] = types.ModuleType(_SENTINEL)
+    try:
+        clear_module(_SENTINEL)
+        assert _SENTINEL not in sys.modules
+    finally:
+        sys.modules.pop(_SENTINEL, None)
+
+
+def test_clear_module_removes_parent_attr():
+    fake_parent = types.ModuleType("_fake_istate_parent")
+    fake_child = types.ModuleType("_fake_istate_parent.child")
+    fake_parent.child = fake_child
+    sys.modules["_fake_istate_parent"] = fake_parent
+    sys.modules["_fake_istate_parent.child"] = fake_child
+    try:
+        clear_module("_fake_istate_parent.child")
+        assert "_fake_istate_parent.child" not in sys.modules
+        assert not hasattr(fake_parent, "child")
+    finally:
+        sys.modules.pop("_fake_istate_parent", None)
+        sys.modules.pop("_fake_istate_parent.child", None)
+
+
+def test_clear_module_tolerates_absent_entry():
+    assert _SENTINEL not in sys.modules
+    clear_module(_SENTINEL)  # must not raise
+
+
+def test_parent_attr_restored_correctly_when_parent_also_preserved():
+    """When a parent package and its child are both named, the child's
+    parent-attr restore must target the *saved* parent module, not the mutated
+    one. This requires phase 1 (sys.modules) to complete before phase 2 (attrs).
+    Tested with child listed before parent to trigger the failure path in a
+    naive single-pass implementation.
+    """
+    fake_parent = types.ModuleType("_fake_istate_parent")
+    fake_child = types.ModuleType("_fake_istate_parent.child")
+    fake_parent.child = fake_child
+    sys.modules["_fake_istate_parent"] = fake_parent
+    sys.modules["_fake_istate_parent.child"] = fake_child
+    try:
+        # child before parent: old single-pass restore would write the child attr
+        # onto the still-mutated parent, then replace sys.modules["_fake_istate_parent"]
+        # — leaving fake_parent.child untouched.
+        with preserve_import_state("_fake_istate_parent.child", "_fake_istate_parent"):
+            new_parent = types.ModuleType("_fake_istate_parent")
+            new_child = types.ModuleType("_fake_istate_parent.child")
+            new_parent.child = new_child
+            sys.modules["_fake_istate_parent"] = new_parent
+            sys.modules["_fake_istate_parent.child"] = new_child
+        # sys.modules entries restored
+        assert sys.modules["_fake_istate_parent"] is fake_parent
+        assert sys.modules["_fake_istate_parent.child"] is fake_child
+        # parent-attr written onto the restored (saved) parent, not the mutated one
+        assert fake_parent.child is fake_child
+    finally:
+        sys.modules.pop("_fake_istate_parent", None)
+        sys.modules.pop("_fake_istate_parent.child", None)

tests/test_webhook_ssrf_resilience.py

@@ -2,59 +2,20 @@ import sys
 import json
 from datetime import datetime
 
-# conftest.py stubs src.database with a fake module; webhook_manager imports
+import pytest
-# from it, so drop the stub here to load the real module under test. We RESTORE
-# both the sys.modules entry AND the parent `src` package attribute afterwards,
-# so the real src.database never leaks into sibling test modules (e.g.
-# llm_core.list_model_ids resolves `from src.database import ...` against
-# sys.modules at call time, and `import src.database as X` resolves through the
-# parent attribute). This mirrors the routes.session_routes isolation fix.
-_ABSENT = object()
 
+from tests.helpers.import_state import clear_module, preserve_import_state
 
-def _save_module_and_parent_attr(dotted_name):
+# conftest.py stubs src.database; drop the stub so webhook_manager imports the
-    """Capture a module's sys.modules entry *and* its parent-package attribute.
+# real module. preserve_import_state restores both sys.modules and the parent
-
+# src.database attribute after the block, preventing stub leakage into siblings.
-    Returns a (module, attr) pair to hand back to
+with preserve_import_state("src.database"):
-    _restore_module_and_parent_attr. Either may be _ABSENT when not present.
+    clear_module("src.database")
-    """
+    _core_database = sys.modules.get("core.database")
-    saved_module = sys.modules.get(dotted_name, _ABSENT)
+    _core_database_all = (
-    pkg_name, _, attr = dotted_name.rpartition(".")
+        getattr(_core_database, "__all__", None) if _core_database is not None else None
-    pkg = sys.modules.get(pkg_name)
+    )
-    saved_attr = getattr(pkg, attr, _ABSENT) if pkg is not None else _ABSENT
+    if _core_database is not None and (
-    return saved_module, saved_attr
-
-
-def _restore_module_and_parent_attr(dotted_name, saved_module, saved_attr):
-    """Restore (or remove) both the sys.modules entry and the parent attribute.
-
-    Passing _ABSENT for both clears the cache, which is how we drop the stub
-    before the real import below.
-    """
-    if saved_module is _ABSENT:
-        sys.modules.pop(dotted_name, None)
-    else:
-        sys.modules[dotted_name] = saved_module
-    pkg_name, _, attr = dotted_name.rpartition(".")
-    pkg = sys.modules.get(pkg_name)
-    if pkg is None:
-        return
-    if saved_attr is _ABSENT:
-        if hasattr(pkg, attr):
-            delattr(pkg, attr)
-    else:
-        setattr(pkg, attr, saved_attr)
-
-
-# Capture the stub state, then clear both bindings so webhook_manager's import
-# below produces/binds the real src.database with no stale stub behind it.
-_src_database_saved = _save_module_and_parent_attr("src.database")
-_restore_module_and_parent_attr("src.database", _ABSENT, _ABSENT)
-_core_database = sys.modules.get("core.database")
-_core_database_all = getattr(_core_database, "__all__", None) if _core_database is not None else None
-if (
-    _core_database is not None
-    and (
         not getattr(_core_database, "__file__", None)
         or (
             _core_database_all is not None
@@ -63,17 +24,9 @@ if (
                 or not all(isinstance(name, str) for name in _core_database_all)
             )
         )
-    )
+    ):
-):
+        del sys.modules["core.database"]
-    del sys.modules["core.database"]
+    from src.webhook_manager import validate_webhook_url
-
-import pytest
-from src.webhook_manager import validate_webhook_url
-
-# webhook_manager is now bound to the real src.database, so restore both the
-# sys.modules entry and the parent `src.database` attribute to their original
-# stub state to avoid polluting sibling test modules.
-_restore_module_and_parent_attr("src.database", *_src_database_saved)
 
 
 def test_webhook_url_ssrf_mitigation():