From a9c1c698b0de20643e3467bc6dca9ff6561e42b8 Mon Sep 17 00:00:00 2001 From: Alexandre Teixeira <111787685+alteixeira20@users.noreply.github.com> Date: Fri, 5 Jun 2026 07:30:14 +0100 Subject: [PATCH] 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 | 87 ++++++++++++++++ tests/test_blind_compare_redaction.py | 70 ++----------- tests/test_helpers_import_state.py | 141 ++++++++++++++++++++++++++ tests/test_webhook_ssrf_resilience.py | 77 +++----------- 4 files changed, 251 insertions(+), 124 deletions(-) create mode 100644 tests/helpers/import_state.py create mode 100644 tests/test_helpers_import_state.py diff --git a/tests/helpers/import_state.py b/tests/helpers/import_state.py new file mode 100644 index 0000000..35059bf --- /dev/null +++ b/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) diff --git a/tests/test_blind_compare_redaction.py b/tests/test_blind_compare_redaction.py index 10e0d98..c6eb462 100644 --- a/tests/test_blind_compare_redaction.py +++ b/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* -# modules and import the REAL core.session_manager + src.auth_helpers. pytest -# caches routes.session_routes after the first import, so stubbing auth_helpers / -# session_manager here would poison the shared module for the sibling session -# 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) - - +# Stub only the ORM class modules and import the real core.session_manager so +# the cached routes.session_routes is identical regardless of collection order. +# preserve_import_state restores both sys.modules and parent-package attributes +# after the block, preventing stub leakage into siblings. _TEMP_STUBS = ("core.database", "core.models") -_saved = {name: sys.modules.get(name, _ABSENT) for name in _TEMP_STUBS} -_saved["core.session_manager"] = sys.modules.get("core.session_manager", _ABSENT) -_sr_saved = _save_module_and_parent_attr("routes.session_routes") -try: +with preserve_import_state(*_TEMP_STUBS, "core.session_manager", "routes.session_routes"): 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 - # stubbed import below produces a fresh module with no stale binding behind it. - _restore_module_and_parent_attr("routes.session_routes", _ABSENT, _ABSENT) + clear_module("routes.session_routes") 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 ───────────────────────────── diff --git a/tests/test_helpers_import_state.py b/tests/test_helpers_import_state.py new file mode 100644 index 0000000..d9f9254 --- /dev/null +++ b/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) diff --git a/tests/test_webhook_ssrf_resilience.py b/tests/test_webhook_ssrf_resilience.py index 7678941..9c1ae99 100644 --- a/tests/test_webhook_ssrf_resilience.py +++ b/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 -# 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() +import pytest +from tests.helpers.import_state import clear_module, preserve_import_state -def _save_module_and_parent_attr(dotted_name): - """Capture a module's sys.modules entry *and* its parent-package attribute. - - Returns a (module, attr) pair to hand back to - _restore_module_and_parent_attr. Either may be _ABSENT when not present. - """ - 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 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 ( +# conftest.py stubs src.database; drop the stub so webhook_manager imports the +# real module. preserve_import_state restores both sys.modules and the parent +# src.database attribute after the block, preventing stub leakage into siblings. +with preserve_import_state("src.database"): + clear_module("src.database") + _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"] - -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) + ): + del sys.modules["core.database"] + from src.webhook_manager import validate_webhook_url def test_webhook_url_ssrf_mitigation():