"""Comprehensive tests for ThreadSafeMetricsWriter in database/thread_metrics.py.""" import threading from unittest.mock import MagicMock, patch import pytest from local_deep_research.database.thread_metrics import ( ThreadSafeMetricsWriter, metrics_writer, ) # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @pytest.fixture def writer(): """Return a fresh ThreadSafeMetricsWriter instance.""" return ThreadSafeMetricsWriter() @pytest.fixture def writer_with_password(writer): """Return a writer that already has a password stored for 'testuser'.""" writer.set_user_password("testuser", "testpass") return writer @pytest.fixture def mock_db_manager(): """Patch db_manager and yield (mock_db_manager, mock_session).""" with patch( "local_deep_research.database.thread_metrics.db_manager" ) as mock_db: mock_session = MagicMock() mock_db.create_thread_safe_session_for_metrics.return_value = ( mock_session ) yield mock_db, mock_session # =========================================================================== # set_user_password # =========================================================================== class TestSetUserPassword: """Tests for ThreadSafeMetricsWriter.set_user_password.""" def test_stores_password_in_thread_local(self, writer): """set_user_password stores the password under _thread_local.passwords.""" writer.set_user_password("alice", "secret123") assert writer._thread_local.passwords["alice"] == "secret123" def test_creates_passwords_dict_if_not_exists(self, writer): """Passwords dict is lazily created on first call.""" assert not hasattr(writer._thread_local, "passwords") writer.set_user_password("alice", "secret") assert hasattr(writer._thread_local, "passwords") assert isinstance(writer._thread_local.passwords, dict) def test_can_store_multiple_users(self, writer): """Multiple users can be stored simultaneously.""" writer.set_user_password("alice", "pass_a") writer.set_user_password("bob", "pass_b") writer.set_user_password("charlie", "pass_c") assert writer._thread_local.passwords == { "alice": "pass_a", "bob": "pass_b", "charlie": "pass_c", } def test_overwrites_existing_password(self, writer): """Setting a password for the same user replaces the old value.""" writer.set_user_password("alice", "old_password") writer.set_user_password("alice", "new_password") assert writer._thread_local.passwords["alice"] == "new_password" assert len(writer._thread_local.passwords) == 1 # =========================================================================== # get_session # =========================================================================== class TestGetSession: """Tests for ThreadSafeMetricsWriter.get_session.""" # --- error paths ------------------------------------------------------- def test_raises_when_no_passwords_dict(self, writer): """Raises ValueError when passwords dict has never been created.""" with pytest.raises(ValueError, match="No password set"): with writer.get_session("testuser"): pass def test_raises_when_password_not_found_for_user(self, writer): """Raises ValueError when the requested user has no stored password.""" writer.set_user_password("other_user", "other_pass") with pytest.raises(ValueError, match="No password available"): with writer.get_session("missing_user"): pass def test_raises_when_username_none_and_no_flask_context(self, writer): """Raises ValueError when username=None and Flask context is absent.""" writer.set_user_password("alice", "pass") # The real code does `from flask import session as flask_session` # inside the function body. A RuntimeError simulates no app context. with pytest.raises((ValueError, RuntimeError)): with writer.get_session(username=None): pass def test_raises_when_db_manager_returns_none_session( self, writer_with_password ): """Raises ValueError when db_manager returns None for the session.""" with patch( "local_deep_research.database.thread_metrics.db_manager" ) as mock_db: mock_db.create_thread_safe_session_for_metrics.return_value = None with pytest.raises(ValueError, match="Failed to create session"): with writer_with_password.get_session("testuser"): pass # --- happy path -------------------------------------------------------- def test_yields_session_on_success( self, writer_with_password, mock_db_manager ): """Successfully yields a session object when password is available.""" _mock_db, mock_session = mock_db_manager with writer_with_password.get_session("testuser") as session: assert session is mock_session def test_creates_session_with_correct_args( self, writer_with_password, mock_db_manager ): """Passes correct username and password to db_manager.""" mock_db, _mock_session = mock_db_manager with writer_with_password.get_session("testuser"): pass mock_db.create_thread_safe_session_for_metrics.assert_called_once_with( "testuser", "testpass" ) def test_commits_session_on_success( self, writer_with_password, mock_db_manager ): """Session is committed when context manager exits normally.""" _mock_db, mock_session = mock_db_manager with writer_with_password.get_session("testuser"): pass mock_session.commit.assert_called_once() mock_session.rollback.assert_not_called() def test_rollback_on_exception(self, writer_with_password, mock_db_manager): """Session is rolled back when an exception occurs inside the block.""" _mock_db, mock_session = mock_db_manager with pytest.raises(RuntimeError, match="boom"): with writer_with_password.get_session("testuser"): raise RuntimeError("boom") mock_session.rollback.assert_called_once() mock_session.commit.assert_not_called() def test_closes_session_on_success( self, writer_with_password, mock_db_manager ): """Session is closed in the finally block after normal exit.""" _mock_db, mock_session = mock_db_manager with writer_with_password.get_session("testuser"): pass mock_session.close.assert_called_once() def test_closes_session_on_exception( self, writer_with_password, mock_db_manager ): """Session is closed in the finally block even after an exception.""" _mock_db, mock_session = mock_db_manager with pytest.raises(RuntimeError): with writer_with_password.get_session("testuser"): raise RuntimeError("failure") mock_session.close.assert_called_once() def test_closes_session_when_db_manager_returns_none( self, writer_with_password ): """When db_manager returns None, session is not closed (it is None).""" with patch( "local_deep_research.database.thread_metrics.db_manager" ) as mock_db: mock_db.create_thread_safe_session_for_metrics.return_value = None with pytest.raises(ValueError, match="Failed to create session"): with writer_with_password.get_session("testuser"): pass # session was None, so close should never be called on it # =========================================================================== # write_token_metrics # =========================================================================== class TestWriteTokenMetrics: """Tests for ThreadSafeMetricsWriter.write_token_metrics.""" def _run_write(self, writer, mock_db_manager, token_data, research_id=42): """Helper: call write_token_metrics with mocked session and TokenUsage.""" mock_db, mock_session = mock_db_manager mock_token_cls = MagicMock() with patch( "local_deep_research.database.models.TokenUsage", mock_token_cls, ): writer.write_token_metrics("testuser", research_id, token_data) return mock_token_cls, mock_session # --- field mapping ----------------------------------------------------- def test_creates_token_usage_with_correct_fields( self, writer_with_password, mock_db_manager ): """TokenUsage is constructed with all expected fields from token_data.""" token_data = { "model_name": "gpt-4", "provider": "openai", "prompt_tokens": 100, "completion_tokens": 50, "research_query": "test query", "research_mode": "deep", "research_phase": "search", "search_iteration": 2, "response_time_ms": 1234, "success_status": "success", "error_type": None, "search_engines_planned": ["google", "bing"], "search_engine_selected": "google", "calling_file": "main.py", "calling_function": "run_search", "call_stack": ["a", "b"], "context_limit": 4096, "context_truncated": True, "tokens_truncated": 500, "truncation_ratio": 0.12, "ollama_prompt_eval_count": 90, "ollama_eval_count": 45, "ollama_total_duration": 999999, "ollama_load_duration": 111111, "ollama_prompt_eval_duration": 222222, "ollama_eval_duration": 333333, } mock_token_cls, _ = self._run_write( writer_with_password, mock_db_manager, token_data, research_id=7 ) call_kwargs = mock_token_cls.call_args[1] assert call_kwargs["research_id"] == 7 assert call_kwargs["model_name"] == "gpt-4" assert call_kwargs["model_provider"] == "openai" assert call_kwargs["prompt_tokens"] == 100 assert call_kwargs["completion_tokens"] == 50 assert call_kwargs["total_tokens"] == 150 assert call_kwargs["research_query"] == "test query" assert call_kwargs["research_mode"] == "deep" assert call_kwargs["research_phase"] == "search" assert call_kwargs["search_iteration"] == 2 assert call_kwargs["response_time_ms"] == 1234 assert call_kwargs["success_status"] == "success" assert call_kwargs["error_type"] is None assert call_kwargs["search_engines_planned"] == ["google", "bing"] assert call_kwargs["search_engine_selected"] == "google" assert call_kwargs["calling_file"] == "main.py" assert call_kwargs["calling_function"] == "run_search" assert call_kwargs["call_stack"] == ["a", "b"] assert call_kwargs["context_limit"] == 4096 assert call_kwargs["context_truncated"] is True assert call_kwargs["tokens_truncated"] == 500 assert call_kwargs["truncation_ratio"] == 0.12 assert call_kwargs["ollama_prompt_eval_count"] == 90 assert call_kwargs["ollama_eval_count"] == 45 assert call_kwargs["ollama_total_duration"] == 999999 assert call_kwargs["ollama_load_duration"] == 111111 assert call_kwargs["ollama_prompt_eval_duration"] == 222222 assert call_kwargs["ollama_eval_duration"] == 333333 def test_adds_record_to_session( self, writer_with_password, mock_db_manager ): """The TokenUsage record is added to the session via session.add().""" token_data = { "model_name": "gpt-4", "provider": "openai", "prompt_tokens": 10, "completion_tokens": 5, } mock_token_cls, mock_session = self._run_write( writer_with_password, mock_db_manager, token_data ) mock_session.add.assert_called_once_with(mock_token_cls.return_value) def test_total_tokens_calculated_as_sum( self, writer_with_password, mock_db_manager ): """total_tokens equals prompt_tokens + completion_tokens.""" token_data = { "prompt_tokens": 200, "completion_tokens": 75, } mock_token_cls, _ = self._run_write( writer_with_password, mock_db_manager, token_data ) call_kwargs = mock_token_cls.call_args[1] assert call_kwargs["total_tokens"] == 275 assert ( call_kwargs["total_tokens"] == call_kwargs["prompt_tokens"] + call_kwargs["completion_tokens"] ) # --- defaults ---------------------------------------------------------- def test_default_prompt_tokens_zero( self, writer_with_password, mock_db_manager ): """prompt_tokens defaults to 0 when not provided.""" token_data = {"completion_tokens": 30} mock_token_cls, _ = self._run_write( writer_with_password, mock_db_manager, token_data ) call_kwargs = mock_token_cls.call_args[1] assert call_kwargs["prompt_tokens"] == 0 def test_default_completion_tokens_zero( self, writer_with_password, mock_db_manager ): """completion_tokens defaults to 0 when not provided.""" token_data = {"prompt_tokens": 50} mock_token_cls, _ = self._run_write( writer_with_password, mock_db_manager, token_data ) call_kwargs = mock_token_cls.call_args[1] assert call_kwargs["completion_tokens"] == 0 def test_default_success_status( self, writer_with_password, mock_db_manager ): """success_status defaults to 'success' when not provided.""" token_data = {} mock_token_cls, _ = self._run_write( writer_with_password, mock_db_manager, token_data ) call_kwargs = mock_token_cls.call_args[1] assert call_kwargs["success_status"] == "success" def test_default_context_truncated_false( self, writer_with_password, mock_db_manager ): """context_truncated defaults to False when not provided.""" token_data = {} mock_token_cls, _ = self._run_write( writer_with_password, mock_db_manager, token_data ) call_kwargs = mock_token_cls.call_args[1] assert call_kwargs["context_truncated"] is False def test_default_total_tokens_zero_when_both_missing( self, writer_with_password, mock_db_manager ): """total_tokens is 0 when both prompt_tokens and completion_tokens are absent.""" token_data = {} mock_token_cls, _ = self._run_write( writer_with_password, mock_db_manager, token_data ) call_kwargs = mock_token_cls.call_args[1] assert call_kwargs["total_tokens"] == 0 # --- missing optional fields ------------------------------------------- def test_missing_optional_fields_are_none( self, writer_with_password, mock_db_manager ): """Optional fields not present in token_data resolve to None via .get().""" token_data = {} mock_token_cls, _ = self._run_write( writer_with_password, mock_db_manager, token_data ) call_kwargs = mock_token_cls.call_args[1] assert call_kwargs["model_name"] is None assert call_kwargs["model_provider"] is None assert call_kwargs["research_query"] is None assert call_kwargs["research_mode"] is None assert call_kwargs["research_phase"] is None assert call_kwargs["search_iteration"] is None assert call_kwargs["response_time_ms"] is None assert call_kwargs["error_type"] is None assert call_kwargs["search_engines_planned"] is None assert call_kwargs["search_engine_selected"] is None assert call_kwargs["calling_file"] is None assert call_kwargs["calling_function"] is None assert call_kwargs["call_stack"] is None assert call_kwargs["context_limit"] is None assert call_kwargs["tokens_truncated"] is None assert call_kwargs["truncation_ratio"] is None assert call_kwargs["ollama_prompt_eval_count"] is None assert call_kwargs["ollama_eval_count"] is None assert call_kwargs["ollama_total_duration"] is None assert call_kwargs["ollama_load_duration"] is None assert call_kwargs["ollama_prompt_eval_duration"] is None assert call_kwargs["ollama_eval_duration"] is None def test_research_id_passed_directly( self, writer_with_password, mock_db_manager ): """research_id comes from the function argument, not from token_data.""" token_data = {"model_name": "test-model"} mock_token_cls, _ = self._run_write( writer_with_password, mock_db_manager, token_data, research_id=999 ) call_kwargs = mock_token_cls.call_args[1] assert call_kwargs["research_id"] == 999 def test_research_id_none(self, writer_with_password, mock_db_manager): """research_id can be None.""" token_data = {"model_name": "test-model"} mock_token_cls, _ = self._run_write( writer_with_password, mock_db_manager, token_data, research_id=None ) call_kwargs = mock_token_cls.call_args[1] assert call_kwargs["research_id"] is None # --- ModelUsage upsert ------------------------------------------------- def test_creates_model_usage_when_none_exists( self, writer_with_password, mock_db_manager ): """New ModelUsage record is created when none exists.""" token_data = { "model_name": "gpt-4", "provider": "openai", "prompt_tokens": 100, "completion_tokens": 50, } mock_db, mock_session = mock_db_manager mock_session.query.return_value.filter_by.return_value.first.return_value = None with patch( "local_deep_research.database.models.TokenUsage", MagicMock() ): with patch( "local_deep_research.database.models.ModelUsage" ) as mock_model_cls: writer_with_password.write_token_metrics( "testuser", 1, token_data ) mock_model_cls.assert_called_once_with( model_name="gpt-4", model_provider="openai", total_tokens=150, total_calls=1, ) assert mock_session.add.call_count == 2 def test_updates_existing_model_usage( self, writer_with_password, mock_db_manager ): """Existing ModelUsage record is incremented.""" token_data = { "model_name": "gpt-4", "provider": "openai", "prompt_tokens": 100, "completion_tokens": 50, } mock_db, mock_session = mock_db_manager existing = MagicMock() existing.total_tokens = 1000 existing.total_calls = 5 mock_session.query.return_value.filter_by.return_value.first.return_value = existing with patch( "local_deep_research.database.models.TokenUsage", MagicMock() ): writer_with_password.write_token_metrics("testuser", 1, token_data) assert existing.total_tokens == 1150 assert existing.total_calls == 6 # =========================================================================== # Global instance # =========================================================================== class TestGlobalInstance: """Tests for the module-level metrics_writer singleton.""" def test_metrics_writer_exists(self): """A global metrics_writer instance is importable.""" assert metrics_writer is not None def test_metrics_writer_is_correct_type(self): """The global instance is a ThreadSafeMetricsWriter.""" assert isinstance(metrics_writer, ThreadSafeMetricsWriter) def test_metrics_writer_has_thread_local(self): """The global instance has _thread_local attribute.""" assert hasattr(metrics_writer, "_thread_local") assert isinstance(metrics_writer._thread_local, threading.local) # =========================================================================== # Thread isolation # =========================================================================== class TestThreadIsolation: """Verify that thread-local storage is genuinely isolated between threads.""" def test_passwords_not_shared_across_threads(self, writer): """Password set in one thread is invisible in another thread.""" results = {} def thread_a(): writer.set_user_password("alice", "alice_pass") results["a_passwords"] = dict(writer._thread_local.passwords) def thread_b(): # Wait for thread_a to finish writing import time time.sleep(0.05) results["b_has_passwords"] = hasattr( writer._thread_local, "passwords" ) t_a = threading.Thread(target=thread_a) t_b = threading.Thread(target=thread_b) t_a.start() t_b.start() t_a.join() t_b.join() assert results["a_passwords"] == {"alice": "alice_pass"} assert results["b_has_passwords"] is False def test_each_thread_has_own_password_store(self, writer): """Two threads can independently store passwords without interference.""" results = {} barrier = threading.Barrier(2) def thread_one(): writer.set_user_password("user_one", "pass_one") barrier.wait() results["one"] = dict(writer._thread_local.passwords) def thread_two(): writer.set_user_password("user_two", "pass_two") barrier.wait() results["two"] = dict(writer._thread_local.passwords) t1 = threading.Thread(target=thread_one) t2 = threading.Thread(target=thread_two) t1.start() t2.start() t1.join() t2.join() assert results["one"] == {"user_one": "pass_one"} assert results["two"] == {"user_two": "pass_two"} # Neither thread sees the other's password assert "user_two" not in results["one"] assert "user_one" not in results["two"] def test_get_session_fails_in_different_thread(self, writer): """get_session in a thread that hasn't called set_user_password raises.""" writer.set_user_password("alice", "pass") errors = [] def other_thread(): try: with writer.get_session("alice"): pass except ValueError as e: errors.append(str(e)) t = threading.Thread(target=other_thread) t.start() t.join() assert len(errors) == 1 assert "No password set" in errors[0]