"""Unit tests for UsageService (U7 — usage dashboard aggregations).""" from __future__ import annotations import csv import io import json import pytest from agentkit.llm.protocol import TokenUsage from agentkit.llm.providers.usage_store import InMemoryUsageStore from agentkit.server.admin.usage_service import ( UsageService, get_usage_service, set_usage_service, ) # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @pytest.fixture def store() -> InMemoryUsageStore: return InMemoryUsageStore() @pytest.fixture def service() -> UsageService: return UsageService() @pytest.fixture(autouse=True) def _reset_singleton(): """Reset the UsageService singleton before and after each test.""" set_usage_service(None) yield set_usage_service(None) def _populate_store(store: InMemoryUsageStore) -> None: """Populate ``store`` with a mix of records for testing.""" # User u1 in dept d1, gpt-4o, 100 tokens, $0.05 store.record( "agent1", "gpt-4o", TokenUsage(prompt_tokens=60, completion_tokens=40), cost=0.05, latency_ms=200, user_id="u1", department_id="d1", ) # User u1 in dept d1, claude, 200 tokens, $0.10 store.record( "agent1", "claude", TokenUsage(prompt_tokens=120, completion_tokens=80), cost=0.10, latency_ms=300, user_id="u1", department_id="d1", ) # User u2 in dept d2, gpt-4o, 50 tokens, $0.02 store.record( "agent2", "gpt-4o", TokenUsage(prompt_tokens=30, completion_tokens=20), cost=0.02, latency_ms=100, user_id="u2", department_id="d2", ) # User u3 in dept d1, gpt-4o, 500 tokens, $0.50 (top user) store.record( "agent3", "gpt-4o", TokenUsage(prompt_tokens=300, completion_tokens=200), cost=0.50, latency_ms=400, user_id="u3", department_id="d1", ) # --------------------------------------------------------------------------- # get_usage_summary # --------------------------------------------------------------------------- class TestGetUsageSummary: async def test_summary_aggregates_all(self, service: UsageService, store: InMemoryUsageStore): _populate_store(store) result = await service.get_usage_summary(store) assert result["total_tokens"] == 850 assert abs(result["total_cost"] - 0.67) < 1e-6 assert result["total_requests"] == 4 # by_model: gpt-4o (3 records, 650 tokens), claude (1, 200) assert "gpt-4o" in result["by_model"] assert "claude" in result["by_model"] assert result["by_model"]["gpt-4o"]["count"] == 3 assert result["by_model"]["gpt-4o"]["total_tokens"] == 650 # by_user: u1 (2 records, 300 tokens), u2 (1, 50), u3 (1, 500) assert result["by_user"]["u1"]["total_tokens"] == 300 assert result["by_user"]["u2"]["total_tokens"] == 50 assert result["by_user"]["u3"]["total_tokens"] == 500 # by_department: d1 (3 records, 800 tokens), d2 (1, 50) assert result["by_department"]["d1"]["total_tokens"] == 800 assert result["by_department"]["d2"]["total_tokens"] == 50 async def test_summary_with_department_filter( self, service: UsageService, store: InMemoryUsageStore ): _populate_store(store) result = await service.get_usage_summary(store, department_id="d1") assert result["total_tokens"] == 800 assert result["total_requests"] == 3 # Only u1 and u3 are in d1. assert "u1" in result["by_user"] assert "u3" in result["by_user"] assert "u2" not in result["by_user"] async def test_summary_with_user_filter(self, service: UsageService, store: InMemoryUsageStore): _populate_store(store) result = await service.get_usage_summary(store, user_id="u2") assert result["total_tokens"] == 50 assert result["total_requests"] == 1 assert "u2" in result["by_user"] async def test_summary_with_empty_store(self, service: UsageService, store: InMemoryUsageStore): result = await service.get_usage_summary(store) assert result["total_tokens"] == 0 assert result["total_cost"] == 0.0 assert result["total_requests"] == 0 assert result["by_model"] == {} assert result["by_user"] == {} assert result["by_department"] == {} # --------------------------------------------------------------------------- # get_usage_timeseries # --------------------------------------------------------------------------- class TestGetUsageTimeseries: async def test_timeseries_day_buckets(self, service: UsageService, store: InMemoryUsageStore): _populate_store(store) result = await service.get_usage_timeseries(store, interval="day") # All records are within the same day (today), so we expect one bucket. assert len(result) >= 1 bucket = result[0] assert "timestamp" in bucket assert bucket["tokens"] == 850 assert abs(bucket["cost"] - 0.67) < 1e-6 assert bucket["requests"] == 4 async def test_timeseries_hour_buckets(self, service: UsageService, store: InMemoryUsageStore): _populate_store(store) result = await service.get_usage_timeseries(store, interval="hour") # All records are within the same hour (now), so we expect one bucket. assert len(result) >= 1 assert result[0]["tokens"] == 850 async def test_timeseries_with_department_filter( self, service: UsageService, store: InMemoryUsageStore ): _populate_store(store) result = await service.get_usage_timeseries(store, department_id="d2", interval="day") assert len(result) >= 1 assert result[0]["tokens"] == 50 async def test_timeseries_empty_store(self, service: UsageService, store: InMemoryUsageStore): result = await service.get_usage_timeseries(store, interval="day") assert result == [] # --------------------------------------------------------------------------- # get_usage_by_model # --------------------------------------------------------------------------- class TestGetUsageByModel: async def test_by_model_breakdown(self, service: UsageService, store: InMemoryUsageStore): _populate_store(store) result = await service.get_usage_by_model(store) # Sorted by model name: claude, gpt-4o assert len(result) == 2 models = {row["model"] for row in result} assert models == {"gpt-4o", "claude"} gpt_row = next(r for r in result if r["model"] == "gpt-4o") assert gpt_row["tokens"] == 650 assert gpt_row["requests"] == 3 claude_row = next(r for r in result if r["model"] == "claude") assert claude_row["tokens"] == 200 assert claude_row["requests"] == 1 async def test_by_model_with_department_filter( self, service: UsageService, store: InMemoryUsageStore ): _populate_store(store) result = await service.get_usage_by_model(store, department_id="d2") # d2 only has gpt-4o (50 tokens, 1 request) assert len(result) == 1 assert result[0]["model"] == "gpt-4o" assert result[0]["tokens"] == 50 async def test_by_model_empty_store(self, service: UsageService, store: InMemoryUsageStore): result = await service.get_usage_by_model(store) assert result == [] # --------------------------------------------------------------------------- # get_top_users # --------------------------------------------------------------------------- class TestGetTopUsers: async def test_top_users_sorted_by_tokens( self, service: UsageService, store: InMemoryUsageStore ): _populate_store(store) result = await service.get_top_users(store, limit=10) # u3 (500), u1 (300), u2 (50) assert len(result) == 3 assert result[0]["user_id"] == "u3" assert result[0]["tokens"] == 500 assert result[1]["user_id"] == "u1" assert result[1]["tokens"] == 300 assert result[2]["user_id"] == "u2" assert result[2]["tokens"] == 50 async def test_top_users_respects_limit(self, service: UsageService, store: InMemoryUsageStore): _populate_store(store) result = await service.get_top_users(store, limit=2) assert len(result) == 2 assert result[0]["user_id"] == "u3" assert result[1]["user_id"] == "u1" async def test_top_users_with_department_filter( self, service: UsageService, store: InMemoryUsageStore ): _populate_store(store) result = await service.get_top_users(store, department_id="d1", limit=10) # d1 has u1 and u3 assert len(result) == 2 assert result[0]["user_id"] == "u3" assert result[1]["user_id"] == "u1" async def test_top_users_empty_store(self, service: UsageService, store: InMemoryUsageStore): result = await service.get_top_users(store, limit=10) assert result == [] # --------------------------------------------------------------------------- # export_usage # --------------------------------------------------------------------------- class TestExportUsage: async def test_export_csv_has_header_and_rows( self, service: UsageService, store: InMemoryUsageStore ): _populate_store(store) body = await service.export_usage(store, format="csv") reader = csv.DictReader(io.StringIO(body)) rows = list(reader) assert len(rows) == 4 # Verify headers assert "timestamp" in rows[0] assert "agent_name" in rows[0] assert "model" in rows[0] assert "user_id" in rows[0] assert "department_id" in rows[0] # Verify a known record gpt_rows = [r for r in rows if r["model"] == "gpt-4o"] assert len(gpt_rows) == 3 async def test_export_json_returns_valid_json( self, service: UsageService, store: InMemoryUsageStore ): _populate_store(store) body = await service.export_usage(store, format="json") data = json.loads(body) assert isinstance(data, list) assert len(data) == 4 assert "timestamp" in data[0] assert "user_id" in data[0] assert "department_id" in data[0] async def test_export_csv_empty_store_returns_header_only( self, service: UsageService, store: InMemoryUsageStore ): body = await service.export_usage(store, format="csv") reader = csv.DictReader(io.StringIO(body)) rows = list(reader) assert rows == [] # Header should still be present. assert "timestamp" in body async def test_export_with_department_filter( self, service: UsageService, store: InMemoryUsageStore ): _populate_store(store) body = await service.export_usage(store, department_id="d2", format="csv") reader = csv.DictReader(io.StringIO(body)) rows = list(reader) assert len(rows) == 1 assert rows[0]["department_id"] == "d2" # --------------------------------------------------------------------------- # Singleton helpers # --------------------------------------------------------------------------- class TestSingletonHelpers: def test_get_usage_service_returns_singleton(self): first = get_usage_service() second = get_usage_service() assert first is second def test_set_usage_service_overrides(self): custom = UsageService() set_usage_service(custom) assert get_usage_service() is custom # Clearing falls back to a new lazy instance. set_usage_service(None) new_one = get_usage_service() assert new_one is not custom