331 lines
12 KiB
Python
331 lines
12 KiB
Python
"""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
|