wangweiandCopilot
5d132981ad
- Add seed() method to ModelUsageTracker for bulk-loading persisted entries at startup - Create PostgresModelUsageStore for persistence of model usage counters to Postgres - Store only current cumulative snapshots (no historical time-series) - Use standard CREATE TABLE IF NOT EXISTS idiom matching other Postgres stores - Add comprehensive mocked unit tests for both components Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
110 lines
4.6 KiB
Python
110 lines
4.6 KiB
Python
"""Unit tests for ModelUsageTracker — no mocking needed, pure in-memory state."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from app.shared.model_usage_tracker import ModelUsageEntry, ModelUsageTracker, get_model_usage_tracker
|
|
|
|
|
|
def test_never_called_model_has_no_entry():
|
|
"""A tracker that has never recorded a call returns None from get()."""
|
|
tracker = ModelUsageTracker()
|
|
assert tracker.get("deepseek", "deepseek-v4-flash") is None
|
|
|
|
|
|
def test_record_success_accumulates_tokens_and_calls():
|
|
"""Two successful calls accumulate tokens and call_count_ok."""
|
|
tracker = ModelUsageTracker()
|
|
tracker.record(
|
|
provider="deepseek", model="deepseek-v4-flash", success=True,
|
|
usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, latency_ms=100,
|
|
)
|
|
tracker.record(
|
|
provider="deepseek", model="deepseek-v4-flash", success=True,
|
|
usage={"prompt_tokens": 20, "completion_tokens": 8, "total_tokens": 28}, latency_ms=200,
|
|
)
|
|
entry = tracker.get("deepseek", "deepseek-v4-flash")
|
|
assert entry is not None
|
|
assert entry.total_tokens == 43
|
|
assert entry.prompt_tokens == 30
|
|
assert entry.completion_tokens == 13
|
|
assert entry.call_count_ok == 2
|
|
assert entry.call_count_error == 0
|
|
assert entry.status == "ok"
|
|
assert entry.last_latency_ms == 200
|
|
|
|
|
|
def test_record_error_sets_error_status_without_losing_prior_tokens():
|
|
"""A failed call after successful ones flips status to 'error' but keeps accumulated tokens."""
|
|
tracker = ModelUsageTracker()
|
|
tracker.record(provider="qwen", model="qwen3.5-flash", success=True, usage={"total_tokens": 50}, latency_ms=50)
|
|
tracker.record(provider="qwen", model="qwen3.5-flash", success=False, error="HTTP 500", latency_ms=30)
|
|
entry = tracker.get("qwen", "qwen3.5-flash")
|
|
assert entry.total_tokens == 50
|
|
assert entry.call_count_ok == 1
|
|
assert entry.call_count_error == 1
|
|
assert entry.status == "error"
|
|
assert entry.last_error == "HTTP 500"
|
|
|
|
|
|
def test_record_success_after_error_clears_last_error():
|
|
"""A later successful call clears last_error and status returns to 'ok'."""
|
|
tracker = ModelUsageTracker()
|
|
tracker.record(provider="qwen", model="qwen3.5-flash", success=False, error="timeout", latency_ms=30)
|
|
tracker.record(provider="qwen", model="qwen3.5-flash", success=True, usage={"total_tokens": 5}, latency_ms=40)
|
|
entry = tracker.get("qwen", "qwen3.5-flash")
|
|
assert entry.status == "ok"
|
|
assert entry.last_error is None
|
|
|
|
|
|
def test_record_never_raises_on_bad_usage_dict():
|
|
"""A malformed usage value (wrong type) is swallowed, not raised, and does not corrupt other entries."""
|
|
tracker = ModelUsageTracker()
|
|
tracker.record(provider="embedding", model="text-embedding-v3", success=True, usage="not-a-dict", latency_ms=10) # type: ignore[arg-type]
|
|
# Must not raise, and must not have created a corrupted entry that breaks snapshot().
|
|
snapshot = tracker.snapshot()
|
|
assert isinstance(snapshot, dict)
|
|
|
|
|
|
def test_snapshot_returns_independent_copy():
|
|
"""snapshot() returns a dict that can be safely mutated without affecting the tracker."""
|
|
tracker = ModelUsageTracker()
|
|
tracker.record(provider="deepseek", model="deepseek-v4-flash", success=True, usage={"total_tokens": 1}, latency_ms=1)
|
|
snap = tracker.snapshot()
|
|
snap.clear()
|
|
assert tracker.get("deepseek", "deepseek-v4-flash") is not None
|
|
|
|
|
|
def test_get_model_usage_tracker_returns_singleton():
|
|
"""get_model_usage_tracker() always returns the same process-wide instance."""
|
|
assert get_model_usage_tracker() is get_model_usage_tracker()
|
|
|
|
|
|
def test_seed_populates_registry_from_persisted_entries():
|
|
"""seed() must bulk-load entries (e.g. from Postgres at startup) into the registry."""
|
|
tracker = ModelUsageTracker()
|
|
persisted = {
|
|
"deepseek:deepseek-v4-flash": ModelUsageEntry(
|
|
provider="deepseek", model="deepseek-v4-flash", total_tokens=500, call_count_ok=20,
|
|
),
|
|
}
|
|
|
|
tracker.seed(persisted)
|
|
|
|
entry = tracker.get("deepseek", "deepseek-v4-flash")
|
|
assert entry.total_tokens == 500
|
|
assert entry.call_count_ok == 20
|
|
|
|
|
|
def test_seed_then_record_accumulates_on_top_of_seeded_value():
|
|
"""A call recorded after seeding must add to the seeded total, not replace it."""
|
|
tracker = ModelUsageTracker()
|
|
tracker.seed({
|
|
"deepseek:deepseek-v4-flash": ModelUsageEntry(
|
|
provider="deepseek", model="deepseek-v4-flash", total_tokens=500,
|
|
),
|
|
})
|
|
|
|
tracker.record(provider="deepseek", model="deepseek-v4-flash", success=True, usage={"total_tokens": 10})
|
|
|
|
assert tracker.get("deepseek", "deepseek-v4-flash").total_tokens == 510
|