451 lines
16 KiB
Python
451 lines
16 KiB
Python
"""Unit tests for UsageStore (U4 — UsageStore Persistence)."""
|
|
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
from agentkit.llm.protocol import TokenUsage
|
|
from agentkit.llm.providers.usage_store import (
|
|
InMemoryUsageStore,
|
|
RedisUsageStore,
|
|
UsageRecord,
|
|
UsageBucket,
|
|
UsageSummary,
|
|
UsageStore,
|
|
create_usage_store,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# InMemoryUsageStore
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestInMemoryUsageStore:
|
|
"""Tests for InMemoryUsageStore."""
|
|
|
|
def test_implements_protocol(self):
|
|
store = InMemoryUsageStore()
|
|
assert isinstance(store, UsageStore)
|
|
|
|
def test_record_single(self):
|
|
store = InMemoryUsageStore()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store.record("agent1", "gpt-4", usage, cost=0.05, latency_ms=200)
|
|
summary = store.get_usage()
|
|
assert summary.total_tokens == 150
|
|
assert summary.total_cost == 0.05
|
|
assert len(summary.records) == 1
|
|
|
|
def test_record_multiple(self):
|
|
store = InMemoryUsageStore()
|
|
usage1 = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
usage2 = TokenUsage(prompt_tokens=200, completion_tokens=100)
|
|
store.record("agent1", "gpt-4", usage1, cost=0.05, latency_ms=200)
|
|
store.record("agent1", "gpt-4", usage2, cost=0.10, latency_ms=300)
|
|
summary = store.get_usage()
|
|
assert summary.total_tokens == 450
|
|
assert abs(summary.total_cost - 0.15) < 1e-6
|
|
assert len(summary.records) == 2
|
|
|
|
def test_get_usage_by_agent(self):
|
|
store = InMemoryUsageStore()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store.record("agent1", "gpt-4", usage, cost=0.05, latency_ms=200)
|
|
store.record("agent2", "gpt-4", usage, cost=0.05, latency_ms=200)
|
|
summary = store.get_usage(agent_name="agent1")
|
|
assert len(summary.records) == 1
|
|
assert summary.records[0].agent_name == "agent1"
|
|
|
|
def test_get_usage_by_time_range(self):
|
|
store = InMemoryUsageStore()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store.record("agent1", "gpt-4", usage, cost=0.05, latency_ms=200)
|
|
|
|
# Query with start_time in the future — should return empty
|
|
future = datetime.now(timezone.utc) + timedelta(hours=1)
|
|
summary = store.get_usage(start_time=future)
|
|
assert len(summary.records) == 0
|
|
|
|
# Query with end_time in the past — should return empty
|
|
past = datetime.now(timezone.utc) - timedelta(hours=1)
|
|
summary = store.get_usage(end_time=past)
|
|
assert len(summary.records) == 0
|
|
|
|
# Query with wide range — should return the record
|
|
start = datetime.now(timezone.utc) - timedelta(hours=1)
|
|
end = datetime.now(timezone.utc) + timedelta(hours=1)
|
|
summary = store.get_usage(start_time=start, end_time=end)
|
|
assert len(summary.records) == 1
|
|
|
|
def test_get_usage_by_model(self):
|
|
store = InMemoryUsageStore()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store.record("agent1", "gpt-4", usage, cost=0.05, latency_ms=200)
|
|
store.record("agent1", "gpt-3.5", usage, cost=0.01, latency_ms=100)
|
|
summary = store.get_usage()
|
|
assert "gpt-4" in summary.by_model
|
|
assert "gpt-3.5" in summary.by_model
|
|
assert summary.by_model["gpt-4"]["count"] == 1
|
|
assert summary.by_model["gpt-3.5"]["count"] == 1
|
|
|
|
def test_get_usage_empty(self):
|
|
store = InMemoryUsageStore()
|
|
summary = store.get_usage()
|
|
assert summary.total_tokens == 0
|
|
assert summary.total_cost == 0.0
|
|
assert len(summary.records) == 0
|
|
|
|
def test_max_records_trimming(self):
|
|
store = InMemoryUsageStore()
|
|
store.MAX_RECORDS = 5
|
|
usage = TokenUsage(prompt_tokens=1, completion_tokens=1)
|
|
for i in range(10):
|
|
store.record(f"agent{i}", "gpt-4", usage, cost=0.01, latency_ms=100)
|
|
assert len(store._records) == 5
|
|
# Should keep the last 5 records
|
|
assert store._records[0].agent_name == "agent5"
|
|
|
|
def test_usage_record_timestamp(self):
|
|
store = InMemoryUsageStore()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store.record("agent1", "gpt-4", usage, cost=0.05, latency_ms=200)
|
|
rec = store.get_usage().records[0]
|
|
assert rec.timestamp != ""
|
|
# Should be parseable as ISO 8601
|
|
datetime.fromisoformat(rec.timestamp)
|
|
|
|
def test_record_with_user_and_department(self):
|
|
store = InMemoryUsageStore()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store.record(
|
|
"agent1",
|
|
"gpt-4",
|
|
usage,
|
|
cost=0.05,
|
|
latency_ms=200,
|
|
user_id="u1",
|
|
department_id="d1",
|
|
)
|
|
rec = store.get_usage().records[0]
|
|
assert rec.user_id == "u1"
|
|
assert rec.department_id == "d1"
|
|
|
|
def test_record_defaults_user_department_to_none(self):
|
|
store = InMemoryUsageStore()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store.record("agent1", "gpt-4", usage, cost=0.05, latency_ms=200)
|
|
rec = store.get_usage().records[0]
|
|
assert rec.user_id is None
|
|
assert rec.department_id is None
|
|
|
|
def test_get_usage_filters_by_user(self):
|
|
store = InMemoryUsageStore()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store.record("agent1", "gpt-4", usage, cost=0.05, latency_ms=200, user_id="u1")
|
|
store.record("agent1", "gpt-4", usage, cost=0.05, latency_ms=200, user_id="u2")
|
|
summary = store.get_usage(user_id="u1")
|
|
assert len(summary.records) == 1
|
|
assert summary.records[0].user_id == "u1"
|
|
|
|
def test_get_usage_filters_by_department(self):
|
|
store = InMemoryUsageStore()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store.record("agent1", "gpt-4", usage, cost=0.05, latency_ms=200, department_id="d1")
|
|
store.record("agent1", "gpt-4", usage, cost=0.05, latency_ms=200, department_id="d2")
|
|
summary = store.get_usage(department_id="d1")
|
|
assert len(summary.records) == 1
|
|
assert summary.records[0].department_id == "d1"
|
|
|
|
def test_get_usage_by_user(self):
|
|
store = InMemoryUsageStore()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store.record(
|
|
"agent1",
|
|
"gpt-4",
|
|
usage,
|
|
cost=0.05,
|
|
latency_ms=200,
|
|
user_id="u1",
|
|
department_id="d1",
|
|
)
|
|
store.record(
|
|
"agent1",
|
|
"gpt-4",
|
|
usage,
|
|
cost=0.05,
|
|
latency_ms=200,
|
|
user_id="u2",
|
|
department_id="d2",
|
|
)
|
|
summary = store.get_usage_by_user("u1")
|
|
assert len(summary.records) == 1
|
|
assert summary.records[0].user_id == "u1"
|
|
assert summary.total_tokens == 150
|
|
|
|
def test_get_usage_by_department(self):
|
|
store = InMemoryUsageStore()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store.record(
|
|
"agent1",
|
|
"gpt-4",
|
|
usage,
|
|
cost=0.05,
|
|
latency_ms=200,
|
|
user_id="u1",
|
|
department_id="d1",
|
|
)
|
|
store.record(
|
|
"agent1",
|
|
"gpt-4",
|
|
usage,
|
|
cost=0.05,
|
|
latency_ms=200,
|
|
user_id="u2",
|
|
department_id="d2",
|
|
)
|
|
summary = store.get_usage_by_department("d1")
|
|
assert len(summary.records) == 1
|
|
assert summary.records[0].department_id == "d1"
|
|
assert summary.total_tokens == 150
|
|
|
|
def test_summary_includes_by_user_and_by_department(self):
|
|
store = InMemoryUsageStore()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store.record(
|
|
"agent1",
|
|
"gpt-4",
|
|
usage,
|
|
cost=0.05,
|
|
latency_ms=200,
|
|
user_id="u1",
|
|
department_id="d1",
|
|
)
|
|
store.record(
|
|
"agent1",
|
|
"gpt-4",
|
|
usage,
|
|
cost=0.05,
|
|
latency_ms=200,
|
|
user_id="u1",
|
|
department_id="d1",
|
|
)
|
|
summary = store.get_usage()
|
|
assert "u1" in summary.by_user
|
|
assert summary.by_user["u1"]["count"] == 2
|
|
assert summary.by_user["u1"]["total_tokens"] == 300
|
|
assert "d1" in summary.by_department
|
|
assert summary.by_department["d1"]["count"] == 2
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# UsageRecord / UsageBucket / UsageSummary dataclasses
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestDataclasses:
|
|
def test_usage_record_auto_timestamp(self):
|
|
rec = UsageRecord(
|
|
agent_name="a",
|
|
model="m",
|
|
prompt_tokens=1,
|
|
completion_tokens=1,
|
|
total_tokens=2,
|
|
cost=0.01,
|
|
latency_ms=100,
|
|
)
|
|
assert rec.timestamp != ""
|
|
|
|
def test_usage_record_explicit_timestamp(self):
|
|
rec = UsageRecord(
|
|
agent_name="a",
|
|
model="m",
|
|
prompt_tokens=1,
|
|
completion_tokens=1,
|
|
total_tokens=2,
|
|
cost=0.01,
|
|
latency_ms=100,
|
|
timestamp="2026-01-01T00:00:00+00:00",
|
|
)
|
|
assert rec.timestamp == "2026-01-01T00:00:00+00:00"
|
|
|
|
def test_usage_bucket_defaults(self):
|
|
bucket = UsageBucket()
|
|
assert bucket.prompt_tokens == 0
|
|
assert bucket.completion_tokens == 0
|
|
assert bucket.total_tokens == 0
|
|
assert bucket.cost == 0.0
|
|
assert bucket.count == 0
|
|
|
|
def test_usage_summary_defaults(self):
|
|
summary = UsageSummary()
|
|
assert summary.total_tokens == 0
|
|
assert summary.total_cost == 0.0
|
|
assert summary.by_model == {}
|
|
assert summary.records == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# RedisUsageStore (mocked)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestRedisUsageStoreMocked:
|
|
"""Tests for RedisUsageStore with mocked Redis."""
|
|
|
|
def _make_store(self):
|
|
store = RedisUsageStore(redis_url="redis://localhost:6379")
|
|
return store
|
|
|
|
def test_implements_protocol(self):
|
|
store = self._make_store()
|
|
assert isinstance(store, UsageStore)
|
|
|
|
def test_degrade_to_fallback(self):
|
|
store = self._make_store()
|
|
assert not store._degraded
|
|
store._degrade_to_fallback()
|
|
assert store._degraded
|
|
assert store._fallback is not None
|
|
|
|
def test_record_degraded_uses_fallback(self):
|
|
store = self._make_store()
|
|
store._degrade_to_fallback()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store.record("agent1", "gpt-4", usage, cost=0.05, latency_ms=200)
|
|
# Should be in fallback
|
|
summary = store._fallback.get_usage()
|
|
assert len(summary.records) == 1
|
|
|
|
def test_get_usage_degraded_uses_fallback(self):
|
|
store = self._make_store()
|
|
store._degrade_to_fallback()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store._fallback.record("agent1", "gpt-4", usage, cost=0.05, latency_ms=200)
|
|
summary = store.get_usage()
|
|
assert len(summary.records) == 1
|
|
|
|
def test_get_usage_degraded_no_fallback_returns_empty(self):
|
|
store = self._make_store()
|
|
store._degraded = True
|
|
# No fallback set — should return empty
|
|
summary = store.get_usage()
|
|
assert summary.total_tokens == 0
|
|
|
|
def test_today_key_format(self):
|
|
store = self._make_store()
|
|
key = store._today_key()
|
|
# Should be YYYY-MM-DD
|
|
assert len(key) == 10
|
|
assert key[4] == "-"
|
|
|
|
def test_v2_keys_with_user_and_department(self):
|
|
store = self._make_store()
|
|
hash_key, list_key = store._v2_keys("2026-06-21", "u1", "d1")
|
|
assert hash_key == "agentkit:usage:v2:2026-06-21:u1:d1"
|
|
assert list_key == "agentkit:usage_records:v2:2026-06-21:u1:d1"
|
|
|
|
def test_v2_keys_with_none_user_and_department(self):
|
|
store = self._make_store()
|
|
hash_key, list_key = store._v2_keys("2026-06-21", None, None)
|
|
# None values are normalized to "none" in the key.
|
|
assert hash_key == "agentkit:usage:v2:2026-06-21:none:none"
|
|
assert list_key == "agentkit:usage_records:v2:2026-06-21:none:none"
|
|
|
|
def test_record_degraded_with_user_and_department(self):
|
|
store = self._make_store()
|
|
store._degrade_to_fallback()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store.record(
|
|
"agent1",
|
|
"gpt-4",
|
|
usage,
|
|
cost=0.05,
|
|
latency_ms=200,
|
|
user_id="u1",
|
|
department_id="d1",
|
|
)
|
|
# Should be in fallback with user/department attached.
|
|
summary = store._fallback.get_usage()
|
|
assert len(summary.records) == 1
|
|
assert summary.records[0].user_id == "u1"
|
|
assert summary.records[0].department_id == "d1"
|
|
|
|
def test_get_usage_degraded_with_user_filter(self):
|
|
store = self._make_store()
|
|
store._degrade_to_fallback()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store._fallback.record(
|
|
"agent1",
|
|
"gpt-4",
|
|
usage,
|
|
cost=0.05,
|
|
latency_ms=200,
|
|
user_id="u1",
|
|
department_id="d1",
|
|
)
|
|
store._fallback.record(
|
|
"agent1",
|
|
"gpt-4",
|
|
usage,
|
|
cost=0.05,
|
|
latency_ms=200,
|
|
user_id="u2",
|
|
department_id="d2",
|
|
)
|
|
summary = store.get_usage(user_id="u1")
|
|
assert len(summary.records) == 1
|
|
assert summary.records[0].user_id == "u1"
|
|
|
|
def test_get_usage_by_user_degraded(self):
|
|
store = self._make_store()
|
|
store._degrade_to_fallback()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store._fallback.record(
|
|
"agent1",
|
|
"gpt-4",
|
|
usage,
|
|
cost=0.05,
|
|
latency_ms=200,
|
|
user_id="u1",
|
|
department_id="d1",
|
|
)
|
|
summary = store.get_usage_by_user("u1")
|
|
assert len(summary.records) == 1
|
|
assert summary.records[0].user_id == "u1"
|
|
|
|
def test_get_usage_by_department_degraded(self):
|
|
store = self._make_store()
|
|
store._degrade_to_fallback()
|
|
usage = TokenUsage(prompt_tokens=100, completion_tokens=50)
|
|
store._fallback.record(
|
|
"agent1",
|
|
"gpt-4",
|
|
usage,
|
|
cost=0.05,
|
|
latency_ms=200,
|
|
user_id="u1",
|
|
department_id="d1",
|
|
)
|
|
summary = store.get_usage_by_department("d1")
|
|
assert len(summary.records) == 1
|
|
assert summary.records[0].department_id == "d1"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Factory
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCreateUsageStore:
|
|
def test_memory_backend(self):
|
|
store = create_usage_store(backend="memory")
|
|
assert isinstance(store, InMemoryUsageStore)
|
|
|
|
def test_auto_backend_returns_store(self):
|
|
store = create_usage_store(backend="auto")
|
|
assert isinstance(store, (InMemoryUsageStore, RedisUsageStore))
|
|
|
|
def test_redis_backend_returns_store(self):
|
|
store = create_usage_store(backend="redis")
|
|
# May be InMemory if redis package unavailable
|
|
assert isinstance(store, (InMemoryUsageStore, RedisUsageStore))
|