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

"""Coverage tests for token_counter.py targeting specific logic paths."""

import json
import threading
from unittest.mock import MagicMock, Mock, patch


from local_deep_research.metrics.token_counter import (
    TokenCounter,
    TokenCountingCallback,
)


def _make_callback(**overrides):
    """Create a TokenCountingCallback with sensible defaults."""
    ctx = overrides.pop(
        "research_context",
        {"research_query": "test query", "research_mode": "quick"},
    )
    cb = TokenCountingCallback(
        research_id=overrides.pop("research_id", "r-100"),
        research_context=ctx,
    )
    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", 200)
    cb.success_status = overrides.pop("success_status", "success")
    cb.error_type = overrides.pop("error_type", None)
    cb.calling_file = overrides.pop("calling_file", "runner.py")
    cb.calling_function = overrides.pop("calling_function", "execute")
    cb.call_stack = overrides.pop("call_stack", None)
    for k, v in overrides.items():
        setattr(cb, k, v)
    return cb


def _patch_main_thread():
    t = MagicMock()
    t.name = "MainThread"
    return patch.object(threading, "current_thread", return_value=t)


def _patch_worker_thread():
    t = MagicMock()
    t.name = "WorkerThread-42"
    return patch.object(threading, "current_thread", return_value=t)


def _mock_db_session():
    s = MagicMock()
    s.query.return_value.filter_by.return_value.first.return_value = None
    return s


def _patch_get_user_db_session(mock_session):
    ctx = MagicMock()
    ctx.__enter__ = Mock(return_value=mock_session)
    ctx.__exit__ = Mock(return_value=False)
    return patch(
        "local_deep_research.database.session_context.get_user_db_session",
        return_value=ctx,
    )


class TestSaveToDbMainThreadCoverage:
    def test_saves_token_usage_and_new_model_usage(self):
        cb = _make_callback()
        mock_session = _mock_db_session()
        added = []
        mock_session.add.side_effect = lambda obj: added.append(obj)
        with _patch_main_thread():
            with patch("flask.session", {"username": "alice"}, create=True):
                with _patch_get_user_db_session(mock_session):
                    cb._save_to_db(80, 40)
        assert len(added) == 2
        assert added[0].prompt_tokens == 80
        assert added[0].total_tokens == 120
        assert added[1].model_name == "gpt-4"
        assert added[1].total_calls == 1
        mock_session.commit.assert_called_once()

    def test_updates_existing_model_usage(self):
        cb = _make_callback()
        existing = MagicMock()
        existing.total_tokens = 500
        existing.total_calls = 3
        mock_session = _mock_db_session()
        mock_session.query.return_value.filter_by.return_value.first.return_value = existing
        with _patch_main_thread():
            with patch("flask.session", {"username": "alice"}, create=True):
                with _patch_get_user_db_session(mock_session):
                    cb._save_to_db(100, 50)
        assert existing.total_tokens == 650
        assert existing.total_calls == 4
        assert mock_session.add.call_count == 1

    def test_skips_without_username(self):
        cb = _make_callback()
        with _patch_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(10, 5)
                    mock_gs.assert_not_called()

    def test_handles_db_exception(self):
        cb = _make_callback()
        with _patch_main_thread():
            with patch("flask.session", {"username": "alice"}, create=True):
                with patch(
                    "local_deep_research.database.session_context.get_user_db_session",
                    side_effect=RuntimeError("connection lost"),
                ):
                    cb._save_to_db(10, 5)

    def test_list_search_engines_planned_converted_to_json(self):
        cb = _make_callback(
            research_context={
                "research_query": "q",
                "search_engines_planned": ["google", "brave"],
            }
        )
        mock_session = _mock_db_session()
        added = []
        mock_session.add.side_effect = lambda obj: added.append(obj)
        with _patch_main_thread():
            with patch("flask.session", {"username": "alice"}, create=True):
                with _patch_get_user_db_session(mock_session):
                    cb._save_to_db(50, 25)
        assert isinstance(added[0].search_engines_planned, str)
        assert json.loads(added[0].search_engines_planned) == [
            "google",
            "brave",
        ]

    def test_string_search_engines_planned_unchanged(self):
        cb = _make_callback(
            research_context={
                "research_query": "q",
                "search_engines_planned": '["already_json"]',
            }
        )
        mock_session = _mock_db_session()
        added = []
        mock_session.add.side_effect = lambda obj: added.append(obj)
        with _patch_main_thread():
            with patch("flask.session", {"username": "alice"}, create=True):
                with _patch_get_user_db_session(mock_session):
                    cb._save_to_db(10, 5)
        assert added[0].search_engines_planned == '["already_json"]'


class TestSaveToDbBackgroundThreadCoverage:
    def test_skips_without_username_in_context(self):
        cb = _make_callback(research_context={})
        with _patch_worker_thread():
            cb._save_to_db(100, 50)

    def test_skips_without_password(self):
        cb = _make_callback(research_context={"username": "bob"})
        mock_writer = MagicMock()
        with _patch_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()

    def test_handles_metrics_writer_exception(self):
        cb = _make_callback(
            research_context={"username": "bob", "user_password": "secret"}
        )
        mock_writer = MagicMock()
        mock_writer.write_token_metrics.side_effect = RuntimeError("fail")
        with _patch_worker_thread():
            with patch(
                "local_deep_research.database.thread_metrics.metrics_writer",
                mock_writer,
            ):
                cb._save_to_db(100, 50)

    def test_writes_token_data_with_context_overflow_fields(self):
        cb = _make_callback(
            research_context={
                "username": "bob",
                "user_password": "secret",
                "research_query": "AI trends",
                "search_engines_planned": ["google"],
            }
        )
        cb.context_limit = 4096
        cb.context_truncated = True
        cb.tokens_truncated = 200
        mock_writer = MagicMock()
        with _patch_worker_thread():
            with patch(
                "local_deep_research.database.thread_metrics.metrics_writer",
                mock_writer,
            ):
                cb._save_to_db(100, 50)
        token_data = mock_writer.write_token_metrics.call_args[0][2]
        assert token_data["context_limit"] == 4096
        assert token_data["context_truncated"] is True

    def test_list_search_engines_planned_json_in_thread(self):
        cb = _make_callback(
            research_context={
                "username": "bob",
                "user_password": "secret",
                "search_engines_planned": ["bing", "duckduckgo"],
            }
        )
        mock_writer = MagicMock()
        with _patch_worker_thread():
            with patch(
                "local_deep_research.database.thread_metrics.metrics_writer",
                mock_writer,
            ):
                cb._save_to_db(50, 25)
        token_data = mock_writer.write_token_metrics.call_args[0][2]
        assert isinstance(token_data["search_engines_planned"], str)
        assert json.loads(token_data["search_engines_planned"]) == [
            "bing",
            "duckduckgo",
        ]


class TestGetResearchMetricsCoverage:
    def test_empty_without_username(self):
        counter = TokenCounter()
        with patch("flask.session", {}):
            result = counter.get_research_metrics("res-1")
        assert result["total_tokens"] == 0
        assert result["model_usage"] == []

    def test_queries_and_formats_results(self):
        counter = TokenCounter()
        row1 = MagicMock(
            model_name="gpt-4",
            model_provider="openai",
            prompt_tokens=200,
            completion_tokens=100,
            total_tokens=300,
            calls=2,
        )
        row2 = MagicMock(
            model_name="claude-3",
            model_provider="anthropic",
            prompt_tokens=150,
            completion_tokens=50,
            total_tokens=200,
            calls=1,
        )
        mock_session = MagicMock()
        chain = mock_session.query.return_value
        chain.filter_by.return_value = chain
        chain.group_by.return_value = chain
        chain.order_by.return_value = chain
        chain.all.return_value = [row1, row2]
        with patch("flask.session", {"username": "alice"}):
            with _patch_get_user_db_session(mock_session):
                result = counter.get_research_metrics("res-42")
        assert result["total_tokens"] == 500
        assert result["total_calls"] == 3
        assert len(result["model_usage"]) == 2


class TestGetMetricsFromEncryptedDbCoverage:
    def test_empty_without_username(self):
        counter = TokenCounter()
        with patch("flask.session", {}):
            result = counter._get_metrics_from_encrypted_db("30d", "all")
        assert result["total_tokens"] == 0

    def test_exception_returns_empty_metrics(self):
        counter = TokenCounter()
        with patch("flask.session", {"username": "alice"}):
            with patch(
                "local_deep_research.database.session_context.get_user_db_session",
                side_effect=Exception("DB down"),
            ):
                result = counter._get_metrics_from_encrypted_db("30d", "all")
        assert result["total_tokens"] == 0
        assert result["by_model"] == []

    def test_complete_metrics_build_with_30d_period(self):
        counter = TokenCounter()
        mock_session = MagicMock()
        mock_query = MagicMock()
        mock_session.query.return_value = mock_query
        for attr in (
            "filter",
            "filter_by",
            "with_entities",
            "group_by",
            "order_by",
            "limit",
            "distinct",
        ):
            getattr(mock_query, attr).return_value = mock_query
        mock_query.scalar.return_value = 0
        mock_query.all.return_value = []
        breakdown = MagicMock(
            total_input_tokens=3000,
            total_output_tokens=2000,
            avg_input_tokens=150,
            avg_output_tokens=100,
            avg_total_tokens=250,
        )
        mock_query.first.return_value = breakdown
        mock_query.count.return_value = 0
        with patch("flask.session", {"username": "alice"}):
            with _patch_get_user_db_session(mock_session):
                result = counter._get_metrics_from_encrypted_db("30d", "all")
        assert "rate_limiting" in result
        assert result["rate_limiting"]["total_attempts"] == 0


class TestGetEnhancedMetricsCoverage:
    def test_empty_without_username(self):
        counter = TokenCounter()
        with patch("flask.session", {}):
            result = counter.get_enhanced_metrics()
        assert result["performance_stats"]["total_enhanced_calls"] == 0
        assert result["time_series_data"] == []

    def test_exception_returns_empty_structure(self):
        counter = TokenCounter()
        with patch("flask.session", {"username": "alice"}):
            with patch(
                "local_deep_research.database.session_context.get_user_db_session",
                side_effect=Exception("refused"),
            ):
                result = counter.get_enhanced_metrics()
        assert result["mode_breakdown"] == []

    def test_time_series_and_performance_stats(self):
        counter = TokenCounter()
        mock_session = MagicMock()
        mock_query = MagicMock()
        mock_session.query.return_value = mock_query
        for attr in (
            "filter",
            "with_entities",
            "group_by",
            "order_by",
            "limit",
        ):
            getattr(mock_query, attr).return_value = mock_query
        mock_query.all.return_value = []
        mock_query.count.return_value = 5
        mock_query.scalar.return_value = 350
        with patch("flask.session", {"username": "alice"}):
            with _patch_get_user_db_session(mock_session):
                result = counter.get_enhanced_metrics("30d", "all")
        assert "performance_stats" in result
        assert "time_series_data" in result


class TestGetResearchTimelineMetricsCoverage:
    def test_empty_without_username(self):
        counter = TokenCounter()
        with patch("flask.session", {}):
            result = counter.get_research_timeline_metrics("res-1")
        assert result["timeline"] == []
        assert result["summary"]["total_calls"] == 0

    def test_sql_timeline_build(self):
        counter = TokenCounter()
        mock_session = MagicMock()
        row1 = (
            "2026-01-01 10:00:00",
            100,
            60,
            40,
            150,
            "success",
            None,
            "search",
            1,
            "google",
            "gpt-4",
            "runner.py",
            "execute",
            None,
        )
        row2 = (
            "2026-01-01 10:01:00",
            200,
            120,
            80,
            250,
            "success",
            None,
            "analysis",
            2,
            "brave",
            "gpt-4",
            "analyzer.py",
            "analyze",
            "stack",
        )
        research_info = (
            "AI trends",
            "quick",
            "completed",
            "2026-01-01",
            "2026-01-01",
        )
        mock_session.execute = Mock(
            side_effect=[
                MagicMock(fetchall=Mock(return_value=[row1, row2])),
                MagicMock(fetchone=Mock(return_value=research_info)),
            ]
        )
        with patch("flask.session", {"username": "alice"}):
            with _patch_get_user_db_session(mock_session):
                result = counter.get_research_timeline_metrics("res-5")
        assert len(result["timeline"]) == 2
        assert result["timeline"][1]["cumulative_tokens"] == 300
        assert result["research_details"]["query"] == "AI trends"
        assert result["summary"]["total_tokens"] == 300

    def test_no_research_info(self):
        counter = TokenCounter()
        mock_session = MagicMock()
        row1 = (
            "2026-01-01",
            50,
            30,
            20,
            100,
            "success",
            None,
            "search",
            1,
            None,
            "gpt-4",
            None,
            None,
            None,
        )
        mock_session.execute = Mock(
            side_effect=[
                MagicMock(fetchall=Mock(return_value=[row1])),
                MagicMock(fetchone=Mock(return_value=None)),
            ]
        )
        with patch("flask.session", {"username": "alice"}):
            with _patch_get_user_db_session(mock_session):
                result = counter.get_research_timeline_metrics("res-missing")
        assert result["research_details"] == {}

    def test_phase_stats_aggregation(self):
        counter = TokenCounter()
        mock_session = MagicMock()
        row1 = (
            "2026-01-01",
            100,
            60,
            40,
            200,
            "success",
            None,
            "search",
            1,
            "google",
            "gpt-4",
            None,
            None,
            None,
        )
        row2 = (
            "2026-01-01",
            150,
            90,
            60,
            400,
            "success",
            None,
            "search",
            2,
            "google",
            "gpt-4",
            None,
            None,
            None,
        )
        row3 = (
            "2026-01-01",
            200,
            120,
            80,
            300,
            "error",
            "TimeoutError",
            "analysis",
            1,
            None,
            "gpt-4",
            None,
            None,
            None,
        )
        mock_session.execute = Mock(
            side_effect=[
                MagicMock(fetchall=Mock(return_value=[row1, row2, row3])),
                MagicMock(fetchone=Mock(return_value=None)),
            ]
        )
        with patch("flask.session", {"username": "alice"}):
            with _patch_get_user_db_session(mock_session):
                result = counter.get_research_timeline_metrics("res-ps")
        assert result["phase_stats"]["search"]["count"] == 2
        assert result["phase_stats"]["search"]["avg_response_time"] == 300
        assert result["phase_stats"]["analysis"]["tokens"] == 200

    def test_success_rate_in_summary(self):
        counter = TokenCounter()
        mock_session = MagicMock()
        row_ok = (
            "2026-01-01",
            100,
            60,
            40,
            100,
            "success",
            None,
            "search",
            1,
            None,
            "gpt-4",
            None,
            None,
            None,
        )
        row_err = (
            "2026-01-01",
            50,
            30,
            20,
            200,
            "error",
            "ValueError",
            "search",
            2,
            None,
            "gpt-4",
            None,
            None,
            None,
        )
        mock_session.execute = Mock(
            side_effect=[
                MagicMock(fetchall=Mock(return_value=[row_ok, row_err])),
                MagicMock(fetchone=Mock(return_value=None)),
            ]
        )
        with patch("flask.session", {"username": "alice"}):
            with _patch_get_user_db_session(mock_session):
                result = counter.get_research_timeline_metrics("res-sr")
        assert result["summary"]["success_rate"] == 50.0

    def test_empty_timeline_data(self):
        counter = TokenCounter()
        mock_session = MagicMock()
        mock_session.execute = Mock(
            side_effect=[
                MagicMock(fetchall=Mock(return_value=[])),
                MagicMock(fetchone=Mock(return_value=None)),
            ]
        )
        with patch("flask.session", {"username": "alice"}):
            with _patch_get_user_db_session(mock_session):
                result = counter.get_research_timeline_metrics("res-empty")
        assert result["timeline"] == []
        assert result["summary"]["total_calls"] == 0
        assert result["phase_stats"] == {}