arena-dashboard/tests/test_probe.py

"""Путь cron-пробы (probe.py): захват usage, NULL при неизвестном, сохранение при ошибках."""
import json
import sqlite3


def _rows(db_path):
    con = sqlite3.connect(db_path)
    con.row_factory = sqlite3.Row
    rows = con.execute("SELECT * FROM checks ORDER BY id").fetchall()
    out = [dict(r) for r in rows]
    con.close()
    return out


def test_probe_captures_usage(core, omni_server):
    omni_server.set_completions(usage={"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
                                no_usage=False)
    ok, lat, err, usage = core.probe.check_model("zai/test-model-a")
    assert ok is True and err is None
    assert usage["input_tokens"] == 11
    assert usage["output_tokens"] == 4
    assert usage["total_tokens"] == 15
    assert json.loads(usage["usage_json"])["prompt_tokens"] == 11


def test_probe_missing_usage_stays_null(core, omni_server):
    omni_server.set_completions(no_usage=True)
    ok, lat, err, usage = core.probe.check_model("zai/test-model-a")
    assert ok is True  # контент есть — проверка успешна
    # но usage неизвестен: NULL, не нули
    assert usage == {"input_tokens": None, "output_tokens": None,
                     "total_tokens": None, "usage_json": None}


def test_probe_zero_usage_treated_as_unknown(core, omni_server):
    omni_server.set_completions(usage={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
                                no_usage=False)
    _, _, _, usage = core.probe.check_model("agy/test-model-b")
    assert usage["total_tokens"] is None and usage["input_tokens"] is None


def test_probe_failed_content_preserves_usage(core, omni_server):
    # ответ без choices: парсинг контента упадёт, но usage уже зафиксирован
    omni_server.set_completions(usage={"prompt_tokens": 9, "completion_tokens": 2, "total_tokens": 11},
                                omit_choices=True)
    ok, lat, err, usage = core.probe.check_model("zai/test-model-a")
    assert ok is False
    assert err and "KeyError" in err
    assert usage["input_tokens"] == 9  # токены не потеряны
    assert usage["total_tokens"] == 11


def test_probe_timeout_usage_null(core, omni_server):
    omni_server.set_completions(delay=6)  # CHECK_TIMEOUT=2 в тестовом env
    ok, lat, err, usage = core.probe.check_model("zai/test-model-a")
    assert ok is False
    assert err and ("timed out" in err.lower() or "timeout" in err.lower())
    assert usage["total_tokens"] is None
    assert usage["usage_json"] is None
    assert lat and lat >= 2000  # latency фиксируется всегда
    omni_server.set_completions(delay=0)


def test_probe_http_error_usage_null(core, omni_server):
    omni_server.set_completions(status=500)
    ok, lat, err, usage = core.probe.check_model("agy/test-model-b")
    assert ok is False and err
    assert usage["input_tokens"] is None


def test_probe_run_end_to_end_both_models(core, omni_server):
    """Полный прогон run(): разные модели — разный исход по usage."""
    omni_server.set_completions(usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15})
    omni_server.set_per_model("agy/test-model-b", {"no_usage": True})
    omni_server.set_per_model("zai/test-model-a",
                              {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}})

    core.probe.run()

    rows = {r["model_id"]: r for r in _rows(core.db_path)}
    assert set(rows) == {"zai/test-model-a", "agy/test-model-b"}
    a = rows["zai/test-model-a"]
    assert a["ok"] == 1 and (a["input_tokens"], a["output_tokens"], a["total_tokens"]) == (10, 5, 15)
    assert json.loads(a["usage_json"])["total_tokens"] == 15
    b = rows["agy/test-model-b"]
    assert b["ok"] == 1  # без usage проверка всё равно успешна
    assert b["input_tokens"] is None and b["total_tokens"] is None and b["usage_json"] is None


def test_probe_run_migrates_old_schema_standalone(tmp_path, omni_server):
    """Cron стартует раньше uvicorn: probe должен сам уметь мигрировать старую БД."""
    from conftest import _set_env, restore_env, MODULES
    import sys as _sys

    db = str(tmp_path / "probe-only.db")
    con = sqlite3.connect(db)
    con.execute("CREATE TABLE checks (id INTEGER PRIMARY KEY AUTOINCREMENT, model_id TEXT NOT NULL, "
                "checked_at TEXT NOT NULL, ok INTEGER NOT NULL, latency_ms INTEGER, error TEXT)")
    con.execute("CREATE TABLE models (id TEXT PRIMARY KEY, provider TEXT NOT NULL, name TEXT NOT NULL, "
                "raw_base TEXT NOT NULL, norm_base TEXT NOT NULL, clean_base TEXT NOT NULL, "
                "text_model_match TEXT, agent_model_match TEXT, updated_at TEXT)")
    con.execute("INSERT INTO models (id, provider, name, raw_base, norm_base, clean_base) "
                "VALUES ('zai/test-model-a', 'zai', 'zai/test-model-a', 'test-model-a', 'testmodela', 'testmodela')")
    con.commit()
    con.close()

    _set_env(db, omni_server.url)
    try:
        for m in MODULES:
            _sys.modules.pop(m, None)
        import probe as standalone_probe
        omni_server.set_per_model("zai/test-model-a",
                                  {"usage": {"input_tokens": 6, "output_tokens": 2}})
        standalone_probe.run()
    finally:
        restore_env()

    rows = _rows(db)
    assert len(rows) == 1
    assert rows[0]["input_tokens"] == 6
    assert rows[0]["output_tokens"] == 2
    assert rows[0]["total_tokens"] == 8  # вычислен из отчётных значений