learningcircuit--local-deep-research/tests/metrics/test_token_counter_db.py

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