"""Coverage tests for thread_context.py. Focuses on untested paths: - preserve_research_context: cleanup_current_thread called on success, cleanup_current_thread exception suppressed """ import sys from types import ModuleType from unittest.mock import MagicMock, patch import pytest from local_deep_research.utilities.thread_context import ( clear_search_context, get_search_context, preserve_research_context, set_search_context, ) def _cleanup(): """Remove any leftover context on the current thread.""" clear_search_context() # --------------------------------------------------------------------------- # preserve_research_context — cleanup paths # --------------------------------------------------------------------------- class TestPreserveResearchContextCleanup: """Tests for the cleanup logic inside preserve_research_context wrapper.""" def setup_method(self): _cleanup() # Evict the module so each test exercises a fresh lazy import; keep # the original so teardown can put it back instead of leaving the # NEXT importer to build a new singleton distinct from the one # earlier-bound consumers hold. self._saved_tls_module = sys.modules.pop( "local_deep_research.database.thread_local_session", None ) def teardown_method(self): _cleanup() sys.modules.pop( "local_deep_research.database.thread_local_session", None ) db_pkg = sys.modules.get("local_deep_research.database") if self._saved_tls_module is not None: sys.modules["local_deep_research.database.thread_local_session"] = ( self._saved_tls_module ) if db_pkg is not None: db_pkg.thread_local_session = self._saved_tls_module elif db_pkg is not None and hasattr(db_pkg, "thread_local_session"): # The test's fresh import bound a throwaway module on the parent # package; drop it so attribute lookups can't resolve to it. delattr(db_pkg, "thread_local_session") def test_cleanup_current_thread_called_on_success(self): """When context is set and the function succeeds, cleanup_current_thread is called in the finally block.""" set_search_context({"research_id": "cleanup-success"}) @preserve_research_context def task(): return 42 clear_search_context() mock_cleanup_fn = MagicMock() fake_module = ModuleType( "local_deep_research.database.thread_local_session" ) fake_module.cleanup_current_thread = mock_cleanup_fn with patch.dict( sys.modules, { "local_deep_research.database.thread_local_session": fake_module, }, ): result = task() assert result == 42 mock_cleanup_fn.assert_called_once() def test_cleanup_current_thread_called_on_exception(self): """When context is set and the function raises, cleanup_current_thread is still called in the finally block.""" set_search_context({"research_id": "cleanup-exception"}) @preserve_research_context def failing_task(): raise ValueError("boom") clear_search_context() mock_cleanup_fn = MagicMock() fake_module = ModuleType( "local_deep_research.database.thread_local_session" ) fake_module.cleanup_current_thread = mock_cleanup_fn with patch.dict( sys.modules, { "local_deep_research.database.thread_local_session": fake_module, }, ): with pytest.raises(ValueError, match="boom"): failing_task() mock_cleanup_fn.assert_called_once() def test_cleanup_exception_is_suppressed(self): """If cleanup_current_thread raises, the exception is suppressed and does not propagate to the caller.""" set_search_context({"research_id": "cleanup-suppressed"}) @preserve_research_context def task(): return "ok" clear_search_context() fake_module = ModuleType( "local_deep_research.database.thread_local_session" ) fake_module.cleanup_current_thread = MagicMock( side_effect=RuntimeError("db engine error") ) with patch.dict( sys.modules, { "local_deep_research.database.thread_local_session": fake_module, }, ): result = task() assert result == "ok" assert get_search_context() is None def test_cleanup_import_failure_is_suppressed(self): """If importing thread_local_session itself fails, the exception is suppressed (caught by the bare except).""" set_search_context({"research_id": "import-fail"}) @preserve_research_context def task(): return "fine" clear_search_context() with patch.dict( sys.modules, {"local_deep_research.database.thread_local_session": None}, ): result = task() assert result == "fine" assert get_search_context() is None def test_no_cleanup_when_no_context(self): """When no context was captured at decoration time, the cleanup block is never entered (context is None branch).""" assert get_search_context() is None @preserve_research_context def task(): return "no-context" with patch.dict( sys.modules, {"local_deep_research.database.thread_local_session": None}, ): result = task() assert result == "no-context"