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