342 lines
12 KiB
Python
342 lines
12 KiB
Python
"""Integration tests for the admin usage dashboard routes (U7).
|
|
|
|
Uses FastAPI TestClient with a test app that mounts only the
|
|
``admin_router`` from ``routes.admin``. The ``_require_admin``
|
|
dependency is overridden via ``app.dependency_overrides`` so the tests
|
|
don't need real JWTs — they can simulate admin and non-admin callers
|
|
directly.
|
|
|
|
The LLM gateway is replaced with a stub that exposes a
|
|
``_usage_tracker.store`` attribute pointing at an
|
|
:class:`InMemoryUsageStore` pre-populated with test records.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import csv
|
|
import io
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.testclient import TestClient
|
|
|
|
from agentkit.llm.protocol import TokenUsage
|
|
from agentkit.llm.providers.usage_store import InMemoryUsageStore
|
|
from agentkit.server.admin.usage_service import set_usage_service
|
|
from agentkit.server.auth.models import init_auth_db
|
|
from agentkit.server.routes import admin as admin_routes_module
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _StubTracker:
|
|
"""Minimal stub matching the UsageTracker surface used by routes."""
|
|
|
|
def __init__(self, store: InMemoryUsageStore) -> None:
|
|
self.store = store
|
|
|
|
|
|
class _StubGateway:
|
|
"""Minimal stub matching the LLMGateway surface used by routes."""
|
|
|
|
def __init__(self, store: InMemoryUsageStore) -> None:
|
|
self._usage_tracker = _StubTracker(store)
|
|
|
|
|
|
@pytest.fixture
|
|
def store() -> InMemoryUsageStore:
|
|
return InMemoryUsageStore()
|
|
|
|
|
|
@pytest.fixture
|
|
def populated_store() -> InMemoryUsageStore:
|
|
"""Pre-populated store with a mix of records across users/depts/models."""
|
|
s = InMemoryUsageStore()
|
|
s.record(
|
|
"agent1",
|
|
"gpt-4o",
|
|
TokenUsage(prompt_tokens=60, completion_tokens=40),
|
|
cost=0.05,
|
|
latency_ms=200,
|
|
user_id="u1",
|
|
department_id="d1",
|
|
)
|
|
s.record(
|
|
"agent1",
|
|
"claude",
|
|
TokenUsage(prompt_tokens=120, completion_tokens=80),
|
|
cost=0.10,
|
|
latency_ms=300,
|
|
user_id="u1",
|
|
department_id="d1",
|
|
)
|
|
s.record(
|
|
"agent2",
|
|
"gpt-4o",
|
|
TokenUsage(prompt_tokens=30, completion_tokens=20),
|
|
cost=0.02,
|
|
latency_ms=100,
|
|
user_id="u2",
|
|
department_id="d2",
|
|
)
|
|
s.record(
|
|
"agent3",
|
|
"gpt-4o",
|
|
TokenUsage(prompt_tokens=300, completion_tokens=200),
|
|
cost=0.50,
|
|
latency_ms=400,
|
|
user_id="u3",
|
|
department_id="d1",
|
|
)
|
|
return s
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_singletons():
|
|
set_usage_service(None)
|
|
yield
|
|
set_usage_service(None)
|
|
|
|
|
|
@pytest.fixture
|
|
async def tmp_auth_db(tmp_path: Path) -> Path:
|
|
db_path = tmp_path / "usage_routes.db"
|
|
await init_auth_db(db_path)
|
|
return db_path
|
|
|
|
|
|
def _make_admin_app(store: InMemoryUsageStore, tmp_auth_db: Path) -> FastAPI:
|
|
"""Build a FastAPI app with admin router + stub gateway."""
|
|
app = FastAPI()
|
|
app.state.auth_db_path = str(tmp_auth_db)
|
|
app.state.llm_gateway = _StubGateway(store)
|
|
app.include_router(admin_routes_module.admin_router, prefix="/api/v1")
|
|
# Default: allow admin access.
|
|
app.dependency_overrides[admin_routes_module._require_admin] = lambda: _make_admin_user()
|
|
return app
|
|
|
|
|
|
def _make_admin_user() -> dict[str, Any]:
|
|
return {"user_id": "admin-1", "username": "admin", "role": "admin"}
|
|
|
|
|
|
def _raise_forbidden() -> dict[str, Any]:
|
|
raise HTTPException(status_code=403, detail="Admin permission required")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# /admin/usage/summary
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestUsageSummaryRoute:
|
|
def test_returns_200_with_data(self, populated_store: InMemoryUsageStore, tmp_auth_db: Path):
|
|
app = _make_admin_app(populated_store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/summary")
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert body["total_tokens"] == 850
|
|
assert abs(body["total_cost"] - 0.67) < 1e-6
|
|
assert body["total_requests"] == 4
|
|
assert "gpt-4o" in body["by_model"]
|
|
assert "u1" in body["by_user"]
|
|
assert "d1" in body["by_department"]
|
|
|
|
def test_with_department_filter(self, populated_store: InMemoryUsageStore, tmp_auth_db: Path):
|
|
app = _make_admin_app(populated_store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/summary", params={"department_id": "d2"})
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert body["total_tokens"] == 50
|
|
assert body["total_requests"] == 1
|
|
|
|
def test_empty_store_returns_200_with_zeros(self, store: InMemoryUsageStore, tmp_auth_db: Path):
|
|
app = _make_admin_app(store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/summary")
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert body["total_tokens"] == 0
|
|
assert body["total_cost"] == 0.0
|
|
assert body["total_requests"] == 0
|
|
|
|
def test_non_admin_returns_403(self, populated_store: InMemoryUsageStore, tmp_auth_db: Path):
|
|
app = _make_admin_app(populated_store, tmp_auth_db)
|
|
app.dependency_overrides[admin_routes_module._require_admin] = _raise_forbidden
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/summary")
|
|
assert resp.status_code == 403
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# /admin/usage/timeseries
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestUsageTimeseriesRoute:
|
|
def test_returns_200_with_data(self, populated_store: InMemoryUsageStore, tmp_auth_db: Path):
|
|
app = _make_admin_app(populated_store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/timeseries")
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert isinstance(body, list)
|
|
assert len(body) >= 1
|
|
assert "timestamp" in body[0]
|
|
assert "tokens" in body[0]
|
|
assert body[0]["tokens"] == 850
|
|
|
|
def test_invalid_interval_returns_400(
|
|
self, populated_store: InMemoryUsageStore, tmp_auth_db: Path
|
|
):
|
|
app = _make_admin_app(populated_store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/timeseries", params={"interval": "week"})
|
|
assert resp.status_code == 400
|
|
|
|
def test_empty_store_returns_200_empty_list(self, store: InMemoryUsageStore, tmp_auth_db: Path):
|
|
app = _make_admin_app(store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/timeseries")
|
|
assert resp.status_code == 200
|
|
assert resp.json() == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# /admin/usage/by-model
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestUsageByModelRoute:
|
|
def test_returns_200_with_breakdown(
|
|
self, populated_store: InMemoryUsageStore, tmp_auth_db: Path
|
|
):
|
|
app = _make_admin_app(populated_store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/by-model")
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert isinstance(body, list)
|
|
models = {row["model"] for row in body}
|
|
assert models == {"gpt-4o", "claude"}
|
|
|
|
def test_empty_store_returns_200_empty_list(self, store: InMemoryUsageStore, tmp_auth_db: Path):
|
|
app = _make_admin_app(store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/by-model")
|
|
assert resp.status_code == 200
|
|
assert resp.json() == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# /admin/usage/top-users
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestTopUsersRoute:
|
|
def test_returns_200_sorted_by_tokens(
|
|
self, populated_store: InMemoryUsageStore, tmp_auth_db: Path
|
|
):
|
|
app = _make_admin_app(populated_store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/top-users")
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert len(body) == 3
|
|
# u3 (500), u1 (300), u2 (50)
|
|
assert body[0]["user_id"] == "u3"
|
|
assert body[0]["tokens"] == 500
|
|
assert body[1]["user_id"] == "u1"
|
|
assert body[2]["user_id"] == "u2"
|
|
|
|
def test_limit_param_respected(self, populated_store: InMemoryUsageStore, tmp_auth_db: Path):
|
|
app = _make_admin_app(populated_store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/top-users", params={"limit": 2})
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert len(body) == 2
|
|
|
|
def test_empty_store_returns_200_empty_list(self, store: InMemoryUsageStore, tmp_auth_db: Path):
|
|
app = _make_admin_app(store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/top-users")
|
|
assert resp.status_code == 200
|
|
assert resp.json() == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# /admin/usage/export
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestUsageExportRoute:
|
|
def test_csv_export_returns_200(self, populated_store: InMemoryUsageStore, tmp_auth_db: Path):
|
|
app = _make_admin_app(populated_store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/export", params={"format": "csv"})
|
|
assert resp.status_code == 200
|
|
assert resp.headers["content-type"].startswith("text/csv")
|
|
reader = csv.DictReader(io.StringIO(resp.text))
|
|
rows = list(reader)
|
|
assert len(rows) == 4
|
|
assert "timestamp" in rows[0]
|
|
assert "user_id" in rows[0]
|
|
assert "department_id" in rows[0]
|
|
|
|
def test_json_export_returns_200(self, populated_store: InMemoryUsageStore, tmp_auth_db: Path):
|
|
app = _make_admin_app(populated_store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/export", params={"format": "json"})
|
|
assert resp.status_code == 200
|
|
# PlainTextResponse returns text/plain or application/json depending on media_type.
|
|
body = json.loads(resp.text)
|
|
assert isinstance(body, list)
|
|
assert len(body) == 4
|
|
|
|
def test_invalid_format_returns_400(
|
|
self, populated_store: InMemoryUsageStore, tmp_auth_db: Path
|
|
):
|
|
app = _make_admin_app(populated_store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/export", params={"format": "xml"})
|
|
assert resp.status_code == 400
|
|
|
|
def test_empty_store_csv_returns_header_only(
|
|
self, store: InMemoryUsageStore, tmp_auth_db: Path
|
|
):
|
|
app = _make_admin_app(store, tmp_auth_db)
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/export", params={"format": "csv"})
|
|
assert resp.status_code == 200
|
|
reader = csv.DictReader(io.StringIO(resp.text))
|
|
rows = list(reader)
|
|
assert rows == []
|
|
# Header should still be present.
|
|
assert "timestamp" in resp.text
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Missing gateway
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestMissingGateway:
|
|
def test_summary_returns_500_without_gateway(self, tmp_auth_db: Path):
|
|
app = FastAPI()
|
|
app.state.auth_db_path = str(tmp_auth_db)
|
|
# No llm_gateway on app.state.
|
|
app.include_router(admin_routes_module.admin_router, prefix="/api/v1")
|
|
app.dependency_overrides[admin_routes_module._require_admin] = lambda: _make_admin_user()
|
|
client = TestClient(app)
|
|
resp = client.get("/api/v1/admin/usage/summary")
|
|
assert resp.status_code == 500
|