learningcircuit--local-deep-research/tests/web/auth/test_connection_cleanup.py

"""
Tests for automatic database connection cleanup.
"""

import datetime
from datetime import UTC
from unittest.mock import MagicMock, patch

import pytest

from local_deep_research.web.auth.connection_cleanup import (
    cleanup_idle_connections,
    start_connection_cleanup_scheduler,
)
from local_deep_research.web.auth.session_manager import SessionManager


@pytest.fixture
def sm():
    """Create a fresh SessionManager with short timeouts for testing."""
    with patch(
        "local_deep_research.web.auth.session_manager.get_security_default",
        return_value=1,
    ):
        mgr = SessionManager()
    # Use a very short timeout for testing
    mgr.session_timeout = datetime.timedelta(seconds=1)
    mgr.remember_me_timeout = datetime.timedelta(seconds=2)
    return mgr


@pytest.fixture
def db():
    """Create a mock DatabaseManager."""
    mock = MagicMock()
    mock.get_connected_usernames.return_value = set()
    return mock


class TestCleanupIdleConnections:
    """Tests for cleanup_idle_connections()."""

    @patch(
        "local_deep_research.scheduler.background.get_background_job_scheduler",
    )
    @patch(
        "local_deep_research.web.auth.connection_cleanup.get_usernames_with_active_research",
        return_value=set(),
    )
    def test_closes_connection_no_sessions_no_research(
        self, _mock_research, _mock_sched, sm, db
    ):
        """Connection closed when user has no active sessions and no research."""
        db.get_connected_usernames.return_value = {"alice"}

        cleanup_idle_connections(sm, db)

        db.close_user_database.assert_called_once_with("alice")

    @patch(
        "local_deep_research.web.auth.connection_cleanup.get_usernames_with_active_research",
        return_value=set(),
    )
    def test_keeps_connection_with_active_session(self, _mock, sm, db):
        """Connection NOT closed when user still has an active session."""
        sm.create_session("bob")
        db.get_connected_usernames.return_value = {"bob"}

        cleanup_idle_connections(sm, db)

        db.close_user_database.assert_not_called()

    @patch(
        "local_deep_research.web.auth.connection_cleanup.get_usernames_with_active_research",
        return_value={"carol"},
    )
    def test_keeps_connection_with_active_research(self, _mock, sm, db):
        """Connection NOT closed when user has active research."""
        db.get_connected_usernames.return_value = {"carol"}

        cleanup_idle_connections(sm, db)

        db.close_user_database.assert_not_called()

    @patch(
        "local_deep_research.scheduler.background.get_background_job_scheduler",
    )
    @patch(
        "local_deep_research.web.auth.connection_cleanup.get_usernames_with_active_research",
        return_value=set(),
    )
    def test_multiple_users_handled_independently(
        self, _mock_research, _mock_sched, sm, db
    ):
        """Each user is evaluated independently."""
        sm.create_session("dave")  # active session
        db.get_connected_usernames.return_value = {"dave", "eve"}

        cleanup_idle_connections(sm, db)

        # eve has no session, should be closed; dave should not
        db.close_user_database.assert_called_once_with("eve")

    @patch(
        "local_deep_research.web.auth.connection_cleanup.get_usernames_with_active_research",
        return_value=set(),
    )
    def test_double_check_prevents_race(self, _mock, sm, db):
        """If user logs in between snapshot and close, connection is kept."""
        db.get_connected_usernames.return_value = {"frank"}

        # Simulate: frank has no session at snapshot time, but gains one
        # during the candidate iteration (via has_active_sessions_for).
        original_has = sm.has_active_sessions_for

        def fake_has(username):
            if username == "frank":
                # Simulate login between snapshot and close
                sm.create_session("frank")
                return original_has("frank")
            return original_has(username)

        sm.has_active_sessions_for = fake_has

        cleanup_idle_connections(sm, db)

        db.close_user_database.assert_not_called()

    def test_double_check_research_prevents_race(self, sm, db):
        """If user starts research between snapshot and close, connection is kept."""
        db.get_connected_usernames.return_value = {"heidi"}

        call_count = 0

        def research_side_effect():
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return set()  # First call: no research (snapshot phase)
            return {
                "heidi"
            }  # Second call: research started (double-check phase)

        with patch(
            "local_deep_research.web.auth.connection_cleanup.get_usernames_with_active_research",
            side_effect=research_side_effect,
        ):
            cleanup_idle_connections(sm, db)

        db.close_user_database.assert_not_called()

    @patch(
        "local_deep_research.scheduler.background.get_background_job_scheduler",
    )
    @patch(
        "local_deep_research.web.auth.connection_cleanup.get_usernames_with_active_research",
        return_value=set(),
    )
    def test_expired_sessions_purged_before_check(
        self, _mock_research, _mock_sched, sm, db
    ):
        """Expired sessions are cleaned up before determining active users."""
        # Create session, then expire it
        sid = sm.create_session("grace")
        with sm._lock:
            sm.sessions[sid]["last_access"] = datetime.datetime.now(
                UTC
            ) - datetime.timedelta(hours=5)

        db.get_connected_usernames.return_value = {"grace"}

        cleanup_idle_connections(sm, db)

        # Session expired, so connection should be closed
        db.close_user_database.assert_called_once_with("grace")

    @patch(
        "local_deep_research.web.auth.connection_cleanup.get_usernames_with_active_research",
        return_value=set(),
    )
    def test_no_connections_is_noop(self, _mock, sm, db):
        """No-op when there are no open connections."""
        db.get_connected_usernames.return_value = set()

        cleanup_idle_connections(sm, db)

        db.close_user_database.assert_not_called()

    def test_close_failure_does_not_abort_loop(self, sm, db):
        """If close_user_database raises for one user, others are still closed."""
        db.get_connected_usernames.return_value = {"alice", "bob"}

        def selective_raise(username):
            if username == "alice":
                raise RuntimeError("simulated failure")

        db.close_user_database.side_effect = selective_raise

        with (
            patch(
                "local_deep_research.web.auth.connection_cleanup.get_usernames_with_active_research",
                return_value=set(),
            ),
            patch(
                "local_deep_research.scheduler.background.get_background_job_scheduler",
            ),
        ):
            cleanup_idle_connections(sm, db)

        assert db.close_user_database.call_count == 2
        db.close_user_database.assert_any_call("alice")
        db.close_user_database.assert_any_call("bob")

    @patch(
        "local_deep_research.scheduler.background.get_background_job_scheduler",
    )
    @patch(
        "local_deep_research.web.auth.connection_cleanup.get_usernames_with_active_research",
        return_value=set(),
    )
    def test_unregister_user_called_on_idle_close(
        self, _mock_research, mock_get_sched, sm, db
    ):
        """Scheduler unregister_user is called before closing idle connection."""
        mock_scheduler = MagicMock()
        mock_scheduler.is_running = True
        mock_get_sched.return_value = mock_scheduler

        db.get_connected_usernames.return_value = {"alice"}

        cleanup_idle_connections(sm, db)

        mock_scheduler.unregister_user.assert_called_once_with("alice")
        db.close_user_database.assert_called_once_with("alice")

    @patch(
        "local_deep_research.scheduler.background.get_background_job_scheduler",
    )
    @patch(
        "local_deep_research.web.auth.connection_cleanup.get_usernames_with_active_research",
        return_value=set(),
    )
    def test_scheduler_failure_does_not_block_close(
        self, _mock_research, mock_get_sched, sm, db
    ):
        """If scheduler unregister raises, db close still proceeds."""
        mock_scheduler = MagicMock()
        mock_scheduler.is_running = True
        mock_scheduler.unregister_user.side_effect = RuntimeError(
            "scheduler down"
        )
        mock_get_sched.return_value = mock_scheduler

        db.get_connected_usernames.return_value = {"alice"}

        cleanup_idle_connections(sm, db)

        db.close_user_database.assert_called_once_with("alice")

    @patch(
        "local_deep_research.web.auth.connection_cleanup.session_password_store.clear_all_for_user"
    )
    @patch(
        "local_deep_research.scheduler.background.get_background_job_scheduler",
    )
    @patch(
        "local_deep_research.web.auth.connection_cleanup.get_usernames_with_active_research",
        return_value=set(),
    )
    def test_clear_all_for_user_called_on_idle_close(
        self, _mock_research, mock_get_sched, mock_clear_pwd, sm, db
    ):
        """Session password store is cleared when closing idle connection."""
        mock_scheduler = MagicMock()
        mock_scheduler.is_running = True
        mock_get_sched.return_value = mock_scheduler

        db.get_connected_usernames.return_value = {"alice"}

        cleanup_idle_connections(sm, db)

        mock_clear_pwd.assert_called_once_with("alice")


class TestPopPerUserLocks:
    """Tests for ``_pop_per_user_locks``: drops the four module-level
    per-user lock-dict entries (library-init, backup, queue-processor
    critical sections, library-RAG FAISS-write locks) on user-close.
    Without this each dict grew one entry per username over process
    lifetime.
    """

    def test_pops_all_four_lock_dicts(self):
        from local_deep_research.web.auth.connection_cleanup import (
            _pop_per_user_locks,
        )
        from local_deep_research.database.library_init import (
            _get_user_init_lock,
            _user_init_locks,
        )
        from local_deep_research.database.backup.backup_service import (
            _get_user_lock,
            _user_locks,
        )
        from local_deep_research.web.queue.processor_v2 import (
            queue_processor,
        )
        from local_deep_research.research_library.services.library_rag_service import (
            _faiss_write_locks,
            _get_faiss_write_lock,
        )

        # Populate each dict with a unique test username so we don't
        # collide with state any other test might have left behind.
        u = "test-pop-locks-user-zzz"
        _get_user_init_lock(u)
        _get_user_lock(u)
        queue_processor._get_user_critical_lock(u)
        _get_faiss_write_lock(u, "/tmp/test-pop-locks/idx.faiss")

        assert u in _user_init_locks
        assert u in _user_locks
        assert u in queue_processor._user_critical_locks
        assert any(k[0] == u for k in _faiss_write_locks)

        _pop_per_user_locks(u)

        assert u not in _user_init_locks
        assert u not in _user_locks
        assert u not in queue_processor._user_critical_locks
        assert not any(k[0] == u for k in _faiss_write_locks)

    def test_idempotent_on_missing_user(self):
        """Pop on a username that was never inserted must not raise."""
        from local_deep_research.web.auth.connection_cleanup import (
            _pop_per_user_locks,
        )

        # Should silently no-op.
        _pop_per_user_locks("never-registered-user-zzz")

    def test_pop_called_from_idle_close_path(self, sm, db):
        """Integration: ``cleanup_idle_connections`` calls
        ``_pop_per_user_locks`` for each user it closes — verified
        by asserting all four lock-dict entries are gone, so a
        regression that drops the call (or breaks one of the four
        pop hops) would be caught.
        """
        from local_deep_research.database.library_init import (
            _get_user_init_lock,
            _user_init_locks,
        )
        from local_deep_research.database.backup.backup_service import (
            _get_user_lock,
            _user_locks,
        )
        from local_deep_research.web.queue.processor_v2 import (
            queue_processor,
        )
        from local_deep_research.research_library.services.library_rag_service import (
            _faiss_write_locks,
            _get_faiss_write_lock,
        )

        # Use a dedicated test username (not "alice") to avoid colliding
        # with other tests that may also touch these module-level dicts.
        u = "test-idle-close-user-zzz"
        _get_user_init_lock(u)
        _get_user_lock(u)
        queue_processor._get_user_critical_lock(u)
        _get_faiss_write_lock(u, "/tmp/test-idle-close/idx.faiss")

        assert u in _user_init_locks
        assert u in _user_locks
        assert u in queue_processor._user_critical_locks
        assert any(k[0] == u for k in _faiss_write_locks)

        db.get_connected_usernames.return_value = {u}

        with (
            patch(
                "local_deep_research.web.auth.connection_cleanup.get_usernames_with_active_research",
                return_value=set(),
            ),
            patch(
                "local_deep_research.scheduler.background.get_background_job_scheduler",
            ),
        ):
            cleanup_idle_connections(sm, db)

        # All four lock-dict entries removed by the idle-close path.
        assert u not in _user_init_locks
        assert u not in _user_locks
        assert u not in queue_processor._user_critical_locks
        assert not any(k[0] == u for k in _faiss_write_locks)

    def test_pop_runs_even_when_close_user_database_fails(self, sm, db):
        """Regression for the original PR: ``_pop_per_user_locks`` was
        previously inside the same try/except as ``close_user_database``,
        so a DB-close failure (the very path
        ``test_close_failure_does_not_abort_loop`` exercises) would skip
        the pop and leak the lock-dict entry. Now the pop is outside
        the try; this test pins that behavior.
        """
        from local_deep_research.database.library_init import (
            _get_user_init_lock,
            _user_init_locks,
        )

        u = "test-close-fails-user-zzz"
        _get_user_init_lock(u)
        assert u in _user_init_locks

        db.get_connected_usernames.return_value = {u}
        db.close_user_database.side_effect = RuntimeError(
            "simulated DB close failure"
        )

        with (
            patch(
                "local_deep_research.web.auth.connection_cleanup.get_usernames_with_active_research",
                return_value=set(),
            ),
            patch(
                "local_deep_research.scheduler.background.get_background_job_scheduler",
            ),
        ):
            cleanup_idle_connections(sm, db)

        # Despite close_user_database raising, the lock-dict entry
        # is still popped.
        assert u not in _user_init_locks


class TestStartConnectionCleanupScheduler:
    """Tests for start_connection_cleanup_scheduler()."""

    @patch(
        "local_deep_research.web.auth.connection_cleanup.BackgroundScheduler"
    )
    def test_returns_running_scheduler(self, MockScheduler, sm, db):
        """Verify scheduler starts and returns a BackgroundScheduler."""
        mock_instance = MagicMock()
        MockScheduler.return_value = mock_instance

        result = start_connection_cleanup_scheduler(sm, db)

        assert result is mock_instance
        mock_instance.start.assert_called_once()

    @patch(
        "local_deep_research.web.auth.connection_cleanup.BackgroundScheduler"
    )
    def test_uses_correct_interval_and_jitter(self, MockScheduler, sm, db):
        """Verify the job is added with correct interval and jitter."""
        mock_instance = MagicMock()
        MockScheduler.return_value = mock_instance

        start_connection_cleanup_scheduler(sm, db)

        mock_instance.add_job.assert_called_once_with(
            cleanup_idle_connections,
            "interval",
            seconds=300,
            args=[sm, db],
            id="cleanup_idle_connections",
            jitter=30,
        )

    @patch(
        "local_deep_research.web.auth.connection_cleanup.BackgroundScheduler"
    )
    def test_custom_interval(self, MockScheduler, sm, db):
        """Verify custom interval_seconds parameter is respected."""
        mock_instance = MagicMock()
        MockScheduler.return_value = mock_instance

        start_connection_cleanup_scheduler(sm, db, interval_seconds=60)

        mock_instance.add_job.assert_called_once_with(
            cleanup_idle_connections,
            "interval",
            seconds=60,
            args=[sm, db],
            id="cleanup_idle_connections",
            jitter=30,
        )


class TestSessionManagerThreadSafety:
    """Verify SessionManager operations don't crash under concurrent access."""

    def test_concurrent_create_and_cleanup(self, sm):
        """Create and cleanup sessions concurrently without RuntimeError."""
        import threading

        errors = []

        def create_sessions():
            try:
                for i in range(50):
                    sm.create_session(f"user_{i}")
            except Exception as e:
                errors.append(e)

        def cleanup_sessions():
            try:
                for _ in range(50):
                    sm.cleanup_expired_sessions()
            except Exception as e:
                errors.append(e)

        t1 = threading.Thread(target=create_sessions)
        t2 = threading.Thread(target=cleanup_sessions)
        t1.start()
        t2.start()
        t1.join()
        t2.join()

        assert not errors, f"Concurrent access errors: {errors}"

    def test_get_active_usernames_snapshot(self, sm):
        """get_active_usernames returns a set, not a view."""
        sm.create_session("user_a")
        sm.create_session("user_b")

        result = sm.get_active_usernames()
        assert isinstance(result, set)
        assert result == {"user_a", "user_b"}

    def test_has_active_sessions_for_returns_false_when_none(self, sm):
        assert sm.has_active_sessions_for("nobody") is False

    def test_has_active_sessions_for_returns_true_when_active(self, sm):
        sm.create_session("active_user")
        assert sm.has_active_sessions_for("active_user") is True