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

"""Tests for get_context_overflow_truncation_summary helper.

The helper unifies truncation aggregation between /metrics/api/metrics
and /metrics/api/context-overflow. These tests pin its semantics so the
two endpoints cannot drift apart.
"""

from datetime import datetime, timedelta, timezone

import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from local_deep_research.database.models import Base, TokenUsage
from local_deep_research.metrics.query_utils import (
    get_context_overflow_truncation_summary,
)


@pytest.fixture
def engine():
    eng = create_engine("sqlite:///:memory:")
    Base.metadata.create_all(eng)
    yield eng
    eng.dispose()


@pytest.fixture
def session(engine):
    Session = sessionmaker(bind=engine)
    s = Session()
    yield s
    s.close()


def _make_usage(
    *,
    minutes_ago: int = 0,
    context_limit: int | None = 8192,
    context_truncated: bool = False,
    tokens_truncated: int | None = None,
    prompt_tokens: int = 100,
    research_mode: str | None = None,
) -> TokenUsage:
    return TokenUsage(
        research_id=f"r-{minutes_ago}-{context_truncated}-{research_mode}",
        model_provider="openai",
        model_name="gpt-4",
        prompt_tokens=prompt_tokens,
        completion_tokens=50,
        total_tokens=prompt_tokens + 50,
        timestamp=datetime.now(timezone.utc) - timedelta(minutes=minutes_ago),
        context_limit=context_limit,
        context_truncated=context_truncated,
        tokens_truncated=tokens_truncated,
        research_mode=research_mode,
    )


class TestEmpty:
    def test_empty_db_returns_zeros(self, session):
        result = get_context_overflow_truncation_summary(session, "30d")
        assert result == {
            "total_requests": 0,
            "requests_with_context": 0,
            "truncated_requests": 0,
            "truncation_rate": 0.0,
            "avg_tokens_truncated": 0.0,
            "total_tokens": 0,
            "total_prompt_tokens": 0,
            "total_completion_tokens": 0,
            "avg_prompt_tokens": 0.0,
            "avg_completion_tokens": 0.0,
            "max_prompt_tokens": 0,
        }


class TestCounts:
    def test_mixed_rows_count_correctly(self, session):
        # 4 rows: one with no context_limit, two with context but not truncated,
        # one truncated.
        session.add_all(
            [
                _make_usage(context_limit=None),
                _make_usage(context_limit=4096, context_truncated=False),
                _make_usage(context_limit=8192, context_truncated=False),
                _make_usage(
                    context_limit=8192,
                    context_truncated=True,
                    tokens_truncated=200,
                ),
            ]
        )
        session.commit()

        result = get_context_overflow_truncation_summary(session, "30d")

        assert result["total_requests"] == 4
        assert result["requests_with_context"] == 3
        assert result["truncated_requests"] == 1
        # 1/3 of context-aware requests truncated
        assert result["truncation_rate"] == pytest.approx(33.333333, rel=1e-3)
        assert result["avg_tokens_truncated"] == 200.0
        # Token-summary aggregates over the same 4 rows (each: prompt=100,
        # completion=50, total=150). These must be asserted directly because
        # the test_empty_db_returns_zeros case can't catch a non-zero bug
        # (e.g. SUM accidentally returning per-row count instead of sum).
        assert result["total_tokens"] == 600
        assert result["total_prompt_tokens"] == 400
        assert result["total_completion_tokens"] == 200
        assert result["avg_prompt_tokens"] == pytest.approx(100.0)
        assert result["avg_completion_tokens"] == pytest.approx(50.0)
        assert result["max_prompt_tokens"] == 100


class TestPeriodFilter:
    def test_7d_excludes_older_rows(self, session):
        session.add_all(
            [
                _make_usage(minutes_ago=60),  # 1 hour ago, in window
                _make_usage(minutes_ago=10 * 24 * 60),  # 10 days ago, out
            ]
        )
        session.commit()

        result_7d = get_context_overflow_truncation_summary(session, "7d")
        result_30d = get_context_overflow_truncation_summary(session, "30d")

        assert result_7d["total_requests"] == 1
        assert result_30d["total_requests"] == 2

    def test_all_period_includes_old_rows(self, session):
        # A row 2 years old should be excluded by 1y but included by all.
        session.add_all(
            [
                _make_usage(minutes_ago=60),
                _make_usage(minutes_ago=2 * 365 * 24 * 60),
            ]
        )
        session.commit()

        result_1y = get_context_overflow_truncation_summary(session, "1y")
        result_all = get_context_overflow_truncation_summary(session, "all")

        assert result_1y["total_requests"] == 1
        assert result_all["total_requests"] == 2


class TestUnknownPeriodFallback:
    def test_unknown_period_uses_30d_window(self, session):
        # Inside 30d, outside 30d.
        session.add_all(
            [
                _make_usage(minutes_ago=60),
                _make_usage(minutes_ago=60 * 24 * 60),  # 60 days ago
            ]
        )
        session.commit()

        # "foo" is not a valid period — helper falls back to 30 days, matching
        # the route's defensive whitelist behavior end-to-end.
        result = get_context_overflow_truncation_summary(session, "foo")
        assert result["total_requests"] == 1


class TestNullTokensTruncated:
    def test_null_tokens_truncated_yields_zero_not_nan(self, session):
        # Truncated row but tokens_truncated=NULL — func.avg ignores NULL,
        # the `or 0` guard then prevents NaN/None leaking out.
        session.add(
            _make_usage(
                context_limit=8192,
                context_truncated=True,
                tokens_truncated=None,
            )
        )
        session.commit()

        result = get_context_overflow_truncation_summary(session, "30d")

        assert result["truncated_requests"] == 1
        assert result["avg_tokens_truncated"] == 0.0


class TestDivisionByZeroGuard:
    def test_no_context_data_keeps_rate_zero(self, session):
        # Rows exist, but none have a context_limit — truncation_rate must
        # still be 0.0 (not NaN) since the denominator is requests_with_context.
        session.add_all(
            [
                _make_usage(context_limit=None),
                _make_usage(context_limit=None),
            ]
        )
        session.commit()

        result = get_context_overflow_truncation_summary(session, "30d")

        assert result["total_requests"] == 2
        assert result["requests_with_context"] == 0
        assert result["truncation_rate"] == 0.0


class TestResearchModeFilter:
    def test_default_includes_all_modes(self, session):
        session.add_all(
            [
                _make_usage(research_mode="quick"),
                _make_usage(research_mode="detailed"),
                _make_usage(research_mode=None),
            ]
        )
        session.commit()

        # Default research_mode is "all" — no mode filter applied
        result = get_context_overflow_truncation_summary(session, "30d")
        assert result["total_requests"] == 3

    def test_quick_excludes_detailed(self, session):
        session.add_all(
            [
                _make_usage(research_mode="quick"),
                _make_usage(research_mode="detailed"),
                _make_usage(research_mode="detailed"),
            ]
        )
        session.commit()

        result = get_context_overflow_truncation_summary(
            session, "30d", research_mode="quick"
        )
        assert result["total_requests"] == 1

    def test_detailed_excludes_quick(self, session):
        session.add_all(
            [
                _make_usage(research_mode="quick"),
                _make_usage(research_mode="detailed"),
                _make_usage(research_mode="detailed"),
            ]
        )
        session.commit()

        result = get_context_overflow_truncation_summary(
            session, "30d", research_mode="detailed"
        )
        assert result["total_requests"] == 2

    def test_all_mode_string_includes_everything(self, session):
        session.add_all(
            [
                _make_usage(research_mode="quick"),
                _make_usage(research_mode="detailed"),
            ]
        )
        session.commit()

        # Explicit "all" matches the default (no filter)
        result = get_context_overflow_truncation_summary(
            session, "30d", research_mode="all"
        )
        assert result["total_requests"] == 2