"""Deep coverage tests for database/thread_local_session.py. Focuses on: - ThreadLocalSessionManager.get_session: session reuse, invalid session recovery, PendingRollbackError rollback then re-execute fails -> new session - ThreadLocalSessionManager.cleanup_all - Module-level helpers: get_metrics_session, get_current_thread_session, cleanup_current_thread - _ThreadCleanup: clear_settings_context and clear_search_context exception paths - ThreadSessionContext: None session handling """ import threading from unittest.mock import MagicMock, patch import pytest from sqlalchemy.exc import PendingRollbackError MODULE = "local_deep_research.database.thread_local_session" def _make_manager(): from local_deep_research.database.thread_local_session import ( ThreadLocalSessionManager, ) return ThreadLocalSessionManager() # --------------------------------------------------------------------------- # get_session – full coverage of branches # --------------------------------------------------------------------------- class TestGetSessionBranches: def test_reuses_valid_session(self): """When an existing valid session is in thread-local, it is reused.""" mgr = _make_manager() valid_sess = MagicMock() mgr._local.session = valid_sess mgr._local.username = "alice" result = mgr.get_session("alice", "pw") assert result is valid_sess valid_sess.execute.assert_called() def test_pending_rollback_then_re_execute_fails_creates_new(self): """PendingRollbackError, rollback succeeds but SELECT 1 again fails.""" mgr = _make_manager() broken_sess = MagicMock() broken_sess.execute.side_effect = [ PendingRollbackError("pending"), Exception("still broken"), # second execute after rollback fails ] mgr._local.session = broken_sess mgr._local.username = "alice" new_sess = MagicMock() with patch(f"{MODULE}.db_manager") as mock_db: mock_db.open_user_database.return_value = MagicMock() mock_db.create_thread_safe_session_for_metrics.return_value = ( new_sess ) result = mgr.get_session("alice", "pw") assert result is new_sess def test_generic_exception_invalidates_session(self): """Any exception from execute causes cleanup and new session creation.""" mgr = _make_manager() broken_sess = MagicMock() broken_sess.execute.side_effect = OSError("db file locked") mgr._local.session = broken_sess mgr._local.username = "alice" new_sess = MagicMock() with patch(f"{MODULE}.db_manager") as mock_db: mock_db.open_user_database.return_value = MagicMock() mock_db.create_thread_safe_session_for_metrics.return_value = ( new_sess ) result = mgr.get_session("alice", "pw") assert result is new_sess def test_credentials_tracked_after_new_session(self): mgr = _make_manager() new_sess = MagicMock() with patch(f"{MODULE}.db_manager") as mock_db: mock_db.open_user_database.return_value = MagicMock() mock_db.create_thread_safe_session_for_metrics.return_value = ( new_sess ) mgr.get_session("bob", "s3cr3t") tid = threading.get_ident() assert mgr._thread_credentials.get(tid) == ("bob", "s3cr3t") # --------------------------------------------------------------------------- # cleanup_all # --------------------------------------------------------------------------- class TestCleanupAll: def test_cleanup_all_calls_cleanup_for_all_tracked_threads(self): mgr = _make_manager() mgr._thread_credentials[111] = ("a", "pw") mgr._thread_credentials[222] = ("b", "pw") called_with = [] def track_cleanup(tid=None): called_with.append(tid) mgr.cleanup_thread = track_cleanup mgr.cleanup_all() # Both thread IDs should have been cleaned up assert len(called_with) == 2 # --------------------------------------------------------------------------- # Module-level helpers # --------------------------------------------------------------------------- class TestModuleLevelHelpers: def test_get_metrics_session_delegates_to_manager(self): from local_deep_research.database.thread_local_session import ( get_metrics_session, ) mock_sess = MagicMock() with patch(f"{MODULE}.thread_session_manager") as mock_mgr: mock_mgr.get_session.return_value = mock_sess result = get_metrics_session("alice", "pw") assert result is mock_sess mock_mgr.get_session.assert_called_once_with("alice", "pw") def test_get_current_thread_session_delegates(self): from local_deep_research.database.thread_local_session import ( get_current_thread_session, ) mock_sess = MagicMock() with patch(f"{MODULE}.thread_session_manager") as mock_mgr: mock_mgr.get_current_session.return_value = mock_sess result = get_current_thread_session() assert result is mock_sess def test_cleanup_current_thread_delegates(self): from local_deep_research.database.thread_local_session import ( cleanup_current_thread, ) with patch(f"{MODULE}.thread_session_manager") as mock_mgr: cleanup_current_thread() mock_mgr.cleanup_thread.assert_called_once_with() # --------------------------------------------------------------------------- # _ThreadCleanup – exception suppression paths # --------------------------------------------------------------------------- class TestThreadCleanupExceptionPaths: def test_clear_settings_context_exception_suppressed(self): from local_deep_research.database.thread_local_session import ( _ThreadCleanup, ) with ( patch(f"{MODULE}.cleanup_current_thread"), patch( "local_deep_research.config.thread_settings.clear_settings_context", side_effect=RuntimeError("settings boom"), create=True, ), ): with _ThreadCleanup(): pass # Should not raise def test_clear_search_context_exception_suppressed(self): from local_deep_research.database.thread_local_session import ( _ThreadCleanup, ) with ( patch(f"{MODULE}.cleanup_current_thread"), patch( "local_deep_research.utilities.thread_context.clear_search_context", side_effect=RuntimeError("search boom"), create=True, ), ): with _ThreadCleanup(): pass # Should not raise def test_returns_false_from_exit(self): """__exit__ returns False to propagate exceptions from the body.""" from local_deep_research.database.thread_local_session import ( _ThreadCleanup, ) with patch(f"{MODULE}.cleanup_current_thread"): cm = _ThreadCleanup() cm.__enter__() result = cm.__exit__(None, None, None) assert result is False # --------------------------------------------------------------------------- # ThreadSessionContext – None session # --------------------------------------------------------------------------- class TestThreadSessionContextNone: def test_none_session_is_yielded(self): from local_deep_research.database.thread_local_session import ( ThreadSessionContext, ) with patch(f"{MODULE}.get_metrics_session", return_value=None): with ThreadSessionContext("alice", "pw") as sess: assert sess is None def test_exception_in_body_propagates(self): from local_deep_research.database.thread_local_session import ( ThreadSessionContext, ) mock_sess = MagicMock() with patch(f"{MODULE}.get_metrics_session", return_value=mock_sess): with pytest.raises(ValueError, match="body error"): with ThreadSessionContext("alice", "pw"): raise ValueError("body error")