"""Unit tests for UsageStore (U4 — UsageStore Persistence).""" import pytest 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) # --------------------------------------------------------------------------- # 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] == "-" # --------------------------------------------------------------------------- # 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))