shaguabijia-app-server/tests/test_comparison_llm_backfill.py

from __future__ import annotations

from datetime import UTC, datetime, timedelta

from app.core.rewards import CN_TZ
from app.db.session import SessionLocal
from app.models.app_config import AppConfig
from app.models.comparison import ComparisonRecord
from app.services import comparison_llm_backfill


def _record(
    trace_id: str,
    *,
    status: str = "success",
    created_at: datetime | None = None,
) -> int:
    with SessionLocal() as db:
        rec = ComparisonRecord(
            trace_id=trace_id,
            status=status,
            created_at=created_at or datetime.now(),
        )
        db.add(rec)
        db.commit()
        return rec.id


def _delete(record_id: int) -> None:
    with SessionLocal() as db:
        rec = db.get(ComparisonRecord, record_id)
        if rec is not None:
            db.delete(rec)
            db.commit()


def test_backfill_retries_then_persists_cost(monkeypatch):
    record_id = _record("llm-retry-1")
    calls = [
        {
            "model": "unknown-model",
            "error": None,
            "usage": {"prompt_tokens": 1000, "completion_tokens": 500},
        }
    ]
    responses = iter([[], calls])
    monkeypatch.setattr(
        comparison_llm_backfill,
        "fetch_llm_calls",
        lambda trace_id: next(responses),
    )
    sleeps: list[float] = []
    monkeypatch.setattr(comparison_llm_backfill.time, "sleep", sleeps.append)
    monkeypatch.setattr(
        comparison_llm_backfill.settings, "INTERNAL_API_SECRET", "test-secret"
    )

    try:
        assert comparison_llm_backfill.backfill_comparison_llm_cost(
            record_id, "llm-retry-1", attempts=2, retry_delays=(0.25,)
        )
        assert sleeps == [0.25]
        with SessionLocal() as db:
            rec = db.get(ComparisonRecord, record_id)
            assert rec.input_tokens == 1000
            assert rec.output_tokens == 500
            assert rec.llm_cost_yuan is not None
            assert rec.llm_calls == calls
    finally:
        _delete(record_id)


def test_repair_batch_only_targets_terminal_missing_rows(monkeypatch):
    missing_id = _record("llm-repair-missing")
    running_id = _record("llm-repair-running", status="running")
    calls = [
        {
            "model": "unknown-model",
            "error": None,
            "usage": {"prompt_tokens": 100, "completion_tokens": 20},
        }
    ]
    seen: list[str] = []

    def fetch(trace_id: str) -> list[dict]:
        seen.append(trace_id)
        return calls

    monkeypatch.setattr(comparison_llm_backfill, "fetch_llm_calls", fetch)
    monkeypatch.setattr(
        comparison_llm_backfill.settings, "INTERNAL_API_SECRET", "test-secret"
    )
    try:
        result = comparison_llm_backfill.repair_missing_comparison_llm_costs(
            limit=10, lookback_days=1
        )
        assert result["repaired"] >= 1
        assert "llm-repair-missing" in seen
        assert "llm-repair-running" not in seen
        with SessionLocal() as db:
            assert db.get(ComparisonRecord, missing_id).llm_cost_yuan is not None
            assert db.get(ComparisonRecord, running_id).llm_cost_yuan is None
    finally:
        _delete(missing_id)
        _delete(running_id)


def test_repair_excludes_records_before_current_price_config(monkeypatch):
    price_changed_at = datetime.now(UTC) - timedelta(hours=1)
    before_change = (price_changed_at - timedelta(minutes=30)).astimezone(CN_TZ)
    after_change = (price_changed_at + timedelta(minutes=30)).astimezone(CN_TZ)
    before_id = _record(
        "llm-before-price-change",
        created_at=before_change.replace(tzinfo=None),
    )
    after_id = _record(
        "llm-after-price-change",
        created_at=after_change.replace(tzinfo=None),
    )
    with SessionLocal() as db:
        existing = db.get(AppConfig, "llm_token_price")
        if existing is not None:
            db.delete(existing)
            db.flush()
        db.add(
            AppConfig(
                key="llm_token_price",
                value={"default": {"input_per_1m": 1, "output_per_1m": 1}},
                updated_at=price_changed_at,
            )
        )
        db.commit()

    seen: list[str] = []

    def backfill(record_id: int, trace_id: str, **kwargs) -> bool:
        seen.append(trace_id)
        return True

    monkeypatch.setattr(
        comparison_llm_backfill,
        "backfill_comparison_llm_cost",
        backfill,
    )
    try:
        result = comparison_llm_backfill.repair_missing_comparison_llm_costs(
            limit=10_000,
            lookback_days=1,
        )
        assert "llm-after-price-change" in seen
        assert "llm-before-price-change" not in seen
        assert result["repaired"] == len(seen)
        assert result["unresolved"] == 0
    finally:
        _delete(before_id)
        _delete(after_id)
        with SessionLocal() as db:
            config = db.get(AppConfig, "llm_token_price")
            if config is not None:
                db.delete(config)
                db.commit()