diff --git a/tests/test_webhook_ssrf_resilience.py b/tests/test_webhook_ssrf_resilience.py index c7f93b9..7678941 100644 --- a/tests/test_webhook_ssrf_resilience.py +++ b/tests/test_webhook_ssrf_resilience.py @@ -3,9 +3,53 @@ 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. -if "src.database" in sys.modules: - del sys.modules["src.database"] +# 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() + + +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 ( @@ -26,6 +70,11 @@ if ( 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(): # SSRF bypasses that must be rejected, including IPv6 unspecified and