"""Coverage tests for database/thread_local_session.py targeting ~8 missing statements. Uncovered functions/branches: - ThreadLocalSessionManager.get_session: PendingRollbackError recovery path - ThreadLocalSessionManager.get_session: rollback recovery fails path - ThreadLocalSessionManager.cleanup_thread: other thread branch (line 148-152) - ThreadLocalSessionManager.cleanup_dead_threads: dead thread sweep - _ThreadCleanup.__exit__: exception during cleanup paths """ from unittest.mock import MagicMock, patch from sqlalchemy.exc import PendingRollbackError from local_deep_research.database.thread_local_session import ( ThreadLocalSessionManager, ThreadSessionContext, _ThreadCleanup, cleanup_dead_threads, thread_cleanup, ) MODULE = "local_deep_research.database.thread_local_session" class TestGetSessionPendingRollback: """Tests for PendingRollbackError recovery in get_session.""" def test_pending_rollback_recovery_succeeds(self): """PendingRollbackError is recovered via rollback. Rollback is called twice in this path: once to clear the pending-rollback state so the retry SELECT 1 can run, and once more after the retry succeeds to release the SHARED lock the validation transaction holds under DEFERRED isolation. """ mgr = ThreadLocalSessionManager() mock_session = MagicMock() # First call raises PendingRollbackError, after rollback it succeeds mock_session.execute.side_effect = [ PendingRollbackError("pending"), MagicMock(), # after rollback, SELECT 1 succeeds ] mgr._local.session = mock_session mgr._local.username = "user" with patch(f"{MODULE}.db_manager"): result = mgr.get_session("user", "pass") assert result is mock_session assert mock_session.rollback.call_count == 2 def test_pending_rollback_recovery_fails_creates_new(self): """When rollback recovery fails, creates a new session.""" mgr = ThreadLocalSessionManager() mock_session = MagicMock() mock_session.execute.side_effect = PendingRollbackError("pending") mock_session.rollback.side_effect = Exception("rollback failed") mgr._local.session = mock_session mgr._local.username = "user" new_session = 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_session ) result = mgr.get_session("user", "pass") assert result is new_session class TestCleanupThread: """Tests for cleanup_thread with different thread IDs.""" def test_cleanup_other_thread_removes_credentials(self): """Cleaning up another thread only removes credentials.""" mgr = ThreadLocalSessionManager() other_tid = 99999999 mgr._thread_credentials[other_tid] = ("user", "pass") mgr.cleanup_thread(other_tid) assert other_tid not in mgr._thread_credentials class TestCleanupDeadThreads: """Tests for cleanup_dead_threads.""" def test_sweeps_dead_thread_credentials(self): """Dead thread entries are removed.""" mgr = ThreadLocalSessionManager() dead_tid = 11111111 mgr._thread_credentials[dead_tid] = ("user", "pass") # Current thread is alive, dead_tid is not mgr.cleanup_dead_threads() assert dead_tid not in mgr._thread_credentials def test_module_level_cleanup_dead_threads(self): """Module-level cleanup_dead_threads delegates to session manager.""" with patch(f"{MODULE}.thread_session_manager") as mock_mgr: cleanup_dead_threads() mock_mgr.cleanup_dead_threads.assert_called_once() def test_module_level_cleanup_session_failure_swallowed(self): """If session sweep fails, the wrapper does not raise.""" with patch(f"{MODULE}.thread_session_manager") as mock_mgr: mock_mgr.cleanup_dead_threads.side_effect = RuntimeError("boom") # Should not raise cleanup_dead_threads() class TestThreadCleanupContextManager: """Tests for _ThreadCleanup context manager.""" def test_cleanup_exception_suppressed(self): """Exceptions during cleanup are suppressed (logged).""" with patch( f"{MODULE}.cleanup_current_thread", side_effect=Exception("boom") ): # Should not raise with _ThreadCleanup(): pass def test_thread_cleanup_as_decorator(self): """thread_cleanup works as a bare decorator.""" @thread_cleanup def worker(): return 42 with patch(f"{MODULE}.cleanup_current_thread"): assert worker() == 42 def test_thread_cleanup_as_factory(self): """thread_cleanup() works as a decorator factory.""" with patch(f"{MODULE}.cleanup_current_thread"): with thread_cleanup(): pass # Should not raise class TestThreadSessionContext: """Tests for ThreadSessionContext.""" def test_context_manager_returns_session(self): """Context manager returns a session.""" with patch( f"{MODULE}.get_metrics_session", return_value=MagicMock() ) as mock_get: with ThreadSessionContext("user", "pass") as session: assert session is not None mock_get.assert_called_once_with("user", "pass")