refactor(tests): reuse import-state helper in session tests

Test-only refactor continuing #2523. Reuses the shared import-state helper in session-related tests, removes duplicated local save/restore logic, and preserves existing test behavior.

tests/test_session_owner_attribution.py

@@ -16,50 +16,15 @@ from unittest.mock import MagicMock
 
 import pytest
 
+from tests.helpers.import_state import clear_module, preserve_import_state
+
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 # Stub heavy ORM modules so routes.session_routes can be imported under
-# conftest's MagicMock sqlalchemy shim. Both the stubs and the cached route
-# module — including the parent `routes` package attribute — are restored in the
-# finally block to prevent poisoning later tests via `import routes.session_routes`.
-_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)
+# conftest's MagicMock sqlalchemy shim. preserve_import_state restores both the
+# stubs and the cached route module — including the parent `routes`/`core`
+# package attributes — on exit, preventing poisoning of later tests via
+# `import routes.session_routes`.
 
 
 def _set_module_and_parent_attr(dotted_name, module):
@@ -69,7 +34,7 @@ def _set_module_and_parent_attr(dotted_name, module):
     pointing at the previous (real) module, so a later import resolving through
     the parent would bypass the stub — and, symmetrically, a stub left on the
     parent attribute would poison later tests. Controlling both keeps the two
-    views consistent so the finally block can fully undo them.
+    views consistent so preserve_import_state can fully undo them.
     """
     sys.modules[dotted_name] = module
     pkg_name, _, attr = dotted_name.rpartition(".")
@@ -81,25 +46,22 @@ def _set_module_and_parent_attr(dotted_name, module):
 # Modules whose import-time effects leak through both sys.modules and the parent
 # `core`/`routes` package attributes. core.database/core.models are stubbed so
 # routes.session_routes imports under conftest's MagicMock sqlalchemy shim;
-# core.session_manager and routes.session_routes are (re)imported fresh. Each is
-# captured at both levels and restored in the finally block so this file cannot
-# poison later tests via `import core.<...>` / `import routes.session_routes`.
+# core.session_manager and routes.session_routes are (re)imported fresh.
+# preserve_import_state captures each at both levels and restores them on exit so
+# this file cannot poison later tests via `import core.<...>` /
+# `import routes.session_routes`.
 _TEMP_STUBS = ("core.database", "core.models")
 _MANAGED = _TEMP_STUBS + ("core.session_manager", "routes.session_routes")
-_saved = {name: _save_module_and_parent_attr(name) for name in _MANAGED}
-try:
+with preserve_import_state(*_MANAGED):
     for _name in _TEMP_STUBS:
         _set_module_and_parent_attr(_name, MagicMock(name=_name))
     # Clear sys.modules AND the parent package attribute for the modules we
     # re-import so the stubbed import below yields fresh modules with no stale
     # binding reachable behind them.
-    _restore_module_and_parent_attr("core.session_manager", _ABSENT, _ABSENT)
-    _restore_module_and_parent_attr("routes.session_routes", _ABSENT, _ABSENT)
+    clear_module("core.session_manager")
+    clear_module("routes.session_routes")
     importlib.import_module("core.session_manager")
     import routes.session_routes as SR  # noqa: E402
-finally:
-    for _name, _save in _saved.items():
-        _restore_module_and_parent_attr(_name, *_save)
 
 from fastapi import HTTPException  # noqa: E402
 from src.auth_helpers import effective_user  # noqa: E402