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") cancelled_id = _record("llm-repair-cancelled", status="cancelled") 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-cancelled" 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, cancelled_id).llm_cost_yuan is not None assert db.get(ComparisonRecord, running_id).llm_cost_yuan is None finally: _delete(missing_id) _delete(cancelled_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()