"""Tests for TokenCountingCallback._save_to_db.""" import threading from unittest.mock import MagicMock, Mock, patch from local_deep_research.metrics.token_counter import ( TokenCountingCallback, ) # ── _save_to_db MainThread path ──────────────────────────────────── class TestSaveToDbMainThread: """Tests for the MainThread path of _save_to_db (lines 484-588).""" def _make_callback(self, **overrides): """Create a callback with sane defaults.""" cb = TokenCountingCallback( research_id="r-123", research_context=overrides.pop( "research_context", {"research_query": "test"} ), ) cb.current_model = overrides.pop("current_model", "gpt-4") cb.current_provider = overrides.pop("current_provider", "openai") cb.response_time_ms = overrides.pop("response_time_ms", 150) cb.success_status = overrides.pop("success_status", "success") cb.error_type = overrides.pop("error_type", None) cb.calling_file = overrides.pop("calling_file", "test.py") cb.calling_function = overrides.pop("calling_function", "run") cb.call_stack = overrides.pop("call_stack", None) for k, v in overrides.items(): setattr(cb, k, v) return cb def _mock_main_thread(self): """Return a patch that makes current_thread().name == 'MainThread'.""" mock_thread = MagicMock() mock_thread.name = "MainThread" return patch.object( threading, "current_thread", return_value=mock_thread ) def test_creates_token_usage(self): """TokenUsage record is created and added to session.""" cb = self._make_callback() mock_session = MagicMock() mock_session.query.return_value.filter_by.return_value.first.return_value = None with self._mock_main_thread(): with patch("flask.session", {"username": "testuser"}, create=True): with patch( "local_deep_research.database.session_context.get_user_db_session" ) as mock_gs: mock_gs.return_value.__enter__ = Mock( return_value=mock_session ) mock_gs.return_value.__exit__ = Mock(return_value=False) cb._save_to_db(100, 50) assert mock_session.add.call_count >= 1 mock_session.commit.assert_called_once() def test_updates_existing_model_usage(self): """Existing ModelUsage record is incremented.""" cb = self._make_callback() existing_model_usage = MagicMock() existing_model_usage.total_tokens = 1000 existing_model_usage.total_calls = 5 mock_session = MagicMock() mock_session.query.return_value.filter_by.return_value.first.return_value = existing_model_usage with self._mock_main_thread(): with patch("flask.session", {"username": "testuser"}, create=True): with patch( "local_deep_research.database.session_context.get_user_db_session" ) as mock_gs: mock_gs.return_value.__enter__ = Mock( return_value=mock_session ) mock_gs.return_value.__exit__ = Mock(return_value=False) cb._save_to_db(100, 50) assert existing_model_usage.total_tokens == 1150 assert existing_model_usage.total_calls == 6 def test_creates_new_model_usage(self): """New ModelUsage record is created when none exists.""" cb = self._make_callback() mock_session = MagicMock() mock_session.query.return_value.filter_by.return_value.first.return_value = None with self._mock_main_thread(): with patch("flask.session", {"username": "testuser"}, create=True): with patch( "local_deep_research.database.session_context.get_user_db_session" ) as mock_gs: mock_gs.return_value.__enter__ = Mock( return_value=mock_session ) mock_gs.return_value.__exit__ = Mock(return_value=False) cb._save_to_db(200, 100) # Two session.add calls: TokenUsage + ModelUsage assert mock_session.add.call_count == 2 def test_no_username_skips(self): """No username in flask session skips the save.""" cb = self._make_callback() with self._mock_main_thread(): with patch("flask.session", {}, create=True): with patch( "local_deep_research.database.session_context.get_user_db_session" ) as mock_gs: cb._save_to_db(100, 50) mock_gs.assert_not_called() def test_database_error_caught(self): """Database exceptions are caught, not raised.""" cb = self._make_callback() with self._mock_main_thread(): with patch("flask.session", {"username": "testuser"}, create=True): with patch( "local_deep_research.database.session_context.get_user_db_session", side_effect=RuntimeError("DB exploded"), ): # Should NOT raise cb._save_to_db(100, 50) def test_enhanced_fields_included(self): """response_time_ms and calling_file are stored in TokenUsage.""" cb = self._make_callback( response_time_ms=250, calling_file="my_module.py" ) mock_session = MagicMock() mock_session.query.return_value.filter_by.return_value.first.return_value = None added_objects = [] mock_session.add.side_effect = lambda obj: added_objects.append(obj) with self._mock_main_thread(): with patch("flask.session", {"username": "testuser"}, create=True): with patch( "local_deep_research.database.session_context.get_user_db_session" ) as mock_gs: mock_gs.return_value.__enter__ = Mock( return_value=mock_session ) mock_gs.return_value.__exit__ = Mock(return_value=False) cb._save_to_db(100, 50) token_usage = added_objects[0] assert token_usage.response_time_ms == 250 assert token_usage.calling_file == "my_module.py" # ── _save_to_db thread path ─────────────────────────────────────── class TestSaveToDbThread: """Tests for the worker-thread path of _save_to_db.""" def _mock_worker_thread(self): """Return a patch that makes current_thread().name != 'MainThread'.""" mock_thread = MagicMock() mock_thread.name = "WorkerThread-1" return patch.object( threading, "current_thread", return_value=mock_thread ) def test_thread_path_converts_search_engines(self): """search_engines_planned list is JSON-serialized in thread path.""" cb = TokenCountingCallback( research_id="r-1", research_context={ "username": "testuser", "user_password": "pass", "search_engines_planned": ["google", "brave"], }, ) cb.current_model = "gpt-4" cb.current_provider = "openai" mock_writer = MagicMock() with self._mock_worker_thread(): with patch( "local_deep_research.database.thread_metrics.metrics_writer", mock_writer, ): cb._save_to_db(100, 50) mock_writer.set_user_password.assert_called_once_with( "testuser", "pass" ) mock_writer.write_token_metrics.assert_called_once() def test_thread_path_no_password_skips(self): """Thread path skips save when no password in research context.""" cb = TokenCountingCallback( research_id="r-1", research_context={"username": "testuser"}, ) cb.current_model = "gpt-4" cb.current_provider = "openai" mock_writer = MagicMock() with self._mock_worker_thread(): with patch( "local_deep_research.database.thread_metrics.metrics_writer", mock_writer, ): cb._save_to_db(100, 50) mock_writer.write_token_metrics.assert_not_called()