learningcircuit--local-deep-research/tests/database/test_thread_local_session_coverage.py

"""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")