Add tests/test_storage.py

This commit is contained in:
2026-09-02 16:47:03 +03:00
parent ff7420e153
commit 58b640c088
+238
View File
@@ -0,0 +1,238 @@
"""Unit-тесты SQLite-кэша и истории."""
from __future__ import annotations
import datetime as dt
import json
import sqlite3
from decimal import Decimal
import pytest
from selectel_ml_tui.models import (
Consumption,
HistoryPoint,
Metrics,
ModelMetrics,
TokenUsage,
)
from selectel_ml_tui.storage import Storage, default_db_path
def _consumption(period: str, amount: str) -> Consumption:
return Consumption.model_validate(
{
"period": period,
"amount_rub": amount,
"input_tokens": 100,
"output_tokens": 200,
"cache_read_tokens": 30,
"cache_write_tokens": 40,
}
)
def _metrics() -> Metrics:
return Metrics.model_validate(
{
"availability_pct": 99.9,
"avg_ttft_ms": 0.4,
"input_tokens": 1000,
"output_tokens": 900,
"cache_read_tokens": 800,
"cache_write_tokens": 700,
}
)
def _models() -> list[ModelMetrics]:
return [
ModelMetrics.model_validate(
{
"model_id": "gpt-4o",
"model_name": "GPT-4o",
"availability_pct": 99.9,
"avg_ttft_ms": 0.3,
}
),
ModelMetrics.model_validate({"model_id": "llama-3", "model_name": "Llama 3"}),
]
def test_default_db_path(monkeypatch, tmp_path) -> None:
monkeypatch.setattr("selectel_ml_tui.storage.data_dir", lambda: tmp_path)
assert default_db_path() == tmp_path / "cache.db"
assert default_db_path().parent == tmp_path
def test_save_and_get_latest_consumption(tmp_path) -> None:
storage = Storage(tmp_path / "cache.db")
item = _consumption("2026-08-01T00:00:00", "854.83")
storage.save_consumption(item)
latest = storage.get_latest_consumption()
assert latest is not None
assert latest.amount_rub == Decimal("854.83")
assert latest.period == dt.datetime(2026, 8, 1)
assert latest.input_tokens == 100
storage.close()
def test_get_latest_empty(tmp_path) -> None:
storage = Storage(tmp_path / "cache.db")
assert storage.get_latest_consumption() is None
assert storage.get_latest_metrics() is None
assert storage.get_latest_model_metrics() == []
storage.close()
def test_save_and_get_latest_metrics(tmp_path) -> None:
storage = Storage(tmp_path / "cache.db")
storage.save_metrics(_metrics())
latest = storage.get_latest_metrics()
assert latest is not None
assert latest.availability_pct == 99.9
assert latest.input_tokens == 1000
storage.close()
def test_save_and_get_latest_model_metrics(tmp_path) -> None:
storage = Storage(tmp_path / "cache.db")
storage.save_model_metrics(_models())
latest = storage.get_latest_model_metrics()
assert len(latest) == 2
assert {m.model_id for m in latest} == {"gpt-4o", "llama-3"}
gpt = next(m for m in latest if m.model_id == "gpt-4o")
assert gpt.availability_pct == 99.9
storage.close()
def test_history_upsert_same_day(tmp_path) -> None:
storage = Storage(tmp_path / "cache.db")
point = HistoryPoint(
timestamp=dt.datetime(2026, 8, 1, 12, 0),
amount_rub=Decimal("10.00"),
tokens=TokenUsage(input_tokens=5),
)
storage.save_history(point)
point2 = HistoryPoint(
timestamp=dt.datetime(2026, 8, 1, 12, 0),
amount_rub=Decimal("15.00"),
tokens=TokenUsage(input_tokens=8),
)
storage.save_history(point2)
history = storage.get_history(days=30)
assert len(history) == 1 # upsert, без дублей
assert history[0].amount_rub == Decimal("15.00")
storage.close()
def test_history_days_filter(tmp_path) -> None:
storage = Storage(tmp_path / "cache.db")
now = dt.datetime.now(dt.timezone.utc)
storage.save_history(
HistoryPoint(timestamp=now, amount_rub=Decimal("1"), tokens=TokenUsage())
)
storage.save_history(
HistoryPoint(
timestamp=now - dt.timedelta(days=30),
amount_rub=Decimal("2"),
tokens=TokenUsage(),
)
)
assert len(storage.get_history(days=7)) == 1
assert len(storage.get_history(days=60)) == 2
storage.close()
def test_save_snapshot_combined(tmp_path) -> None:
storage = Storage(tmp_path / "cache.db")
storage.save_snapshot(
consumption=_consumption("2026-08-01T00:00:00", "854.83"),
metrics=_metrics(),
models=_models(),
history_point=HistoryPoint(
timestamp=dt.datetime(2026, 8, 1, 12, 0),
amount_rub=Decimal("854.83"),
tokens=TokenUsage(input_tokens=100),
),
)
assert storage.get_latest_consumption() is not None
assert storage.get_latest_metrics() is not None
assert len(storage.get_latest_model_metrics()) == 2
assert len(storage.get_history(days=30)) == 1
storage.close()
def test_cleanup_removes_old(tmp_path) -> None:
storage = Storage(tmp_path / "cache.db")
old = dt.datetime(2020, 1, 1, 0, 0)
storage.save_history(
HistoryPoint(timestamp=old, amount_rub=Decimal("1"), tokens=TokenUsage())
)
storage.save_consumption(_consumption("2020-01-01T00:00:00", "1"))
deleted = storage.cleanup(max_age_days=30)
# history удаляется по timestamp; consumption_snapshots — по created_at (сейчас)
assert deleted == 1
assert storage.get_history(days=3650) == []
assert storage.get_latest_consumption() is not None
storage.close()
def test_export_json(tmp_path) -> None:
storage = Storage(tmp_path / "cache.db")
storage.save_history(
HistoryPoint(
timestamp=dt.datetime(2026, 8, 1, 12, 0),
amount_rub=Decimal("854.83"),
tokens=TokenUsage(input_tokens=100, output_tokens=200),
)
)
out = tmp_path / "history.json"
storage.export_json(out)
data = json.loads(out.read_text(encoding="utf-8"))
assert len(data) == 1
assert data[0]["amount_rub"] == "854.83"
assert data[0]["input_tokens"] == 100
storage.close()
def test_export_csv(tmp_path) -> None:
storage = Storage(tmp_path / "cache.db")
storage.save_history(
HistoryPoint(
timestamp=dt.datetime(2026, 8, 1, 12, 0),
amount_rub=Decimal("854.83"),
tokens=TokenUsage(input_tokens=100, output_tokens=200),
)
)
out = tmp_path / "history.csv"
storage.export_csv(out)
lines = out.read_text(encoding="utf-8").strip().splitlines()
assert lines[0].startswith("timestamp,amount_rub")
assert "854.83" in lines[1]
assert "100" in lines[1]
storage.close()
def test_close_idempotent(tmp_path) -> None:
storage = Storage(tmp_path / "cache.db")
storage.close()
with pytest.raises(sqlite3.ProgrammingError):
storage.get_history()
def test_clear_removes_all_rows(tmp_path) -> None:
storage = Storage(tmp_path / "cache.db")
storage.save_history(
HistoryPoint(
timestamp=dt.datetime(2026, 8, 18),
amount_rub=Decimal("1"),
tokens=TokenUsage(input_tokens=1),
)
)
storage.save_consumption(Consumption(period="2026-08-01T00:00:00"))
assert storage.get_history() != []
deleted = storage.clear()
assert deleted >= 1
assert storage.get_history() == []
assert storage.get_latest_consumption() is None