465 lines
17 KiB
Python
465 lines
17 KiB
Python
"""Integration tests for the LLM config admin routes (U5).
|
|
|
|
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 :class:`LlmConfigService` is injected via
|
|
:func:`set_llm_config_service` so the tests can point it at a temp
|
|
YAML file.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import yaml
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.testclient import TestClient
|
|
|
|
from agentkit.server.admin.llm_config_service import (
|
|
LlmConfigService,
|
|
set_llm_config_service,
|
|
)
|
|
from agentkit.server.auth.models import init_auth_db
|
|
from agentkit.server.routes import admin as admin_routes_module
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _sample_config() -> dict[str, Any]:
|
|
"""A minimal agentkit.yaml-style config for testing."""
|
|
return {
|
|
"server": {"host": "0.0.0.0", "port": 8001},
|
|
"llm": {
|
|
"providers": {
|
|
"openai": {
|
|
"type": "openai",
|
|
"api_key": "sk-test-12345678",
|
|
"base_url": "https://api.openai.com/v1",
|
|
"models": {"gpt-4o": {}},
|
|
"max_tokens": 4096,
|
|
"timeout": 120.0,
|
|
},
|
|
},
|
|
"model_aliases": {"gpt4": "openai/gpt-4o"},
|
|
"fallbacks": {},
|
|
},
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def config_path(tmp_path: Path) -> Path:
|
|
"""Create a temporary agentkit.yaml config file."""
|
|
path = tmp_path / "agentkit.yaml"
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
yaml.dump(_sample_config(), f, default_flow_style=False, allow_unicode=True)
|
|
return path
|
|
|
|
|
|
@pytest.fixture
|
|
async def tmp_auth_db(tmp_path: Path) -> Path:
|
|
db_path = tmp_path / "admin_llm.db"
|
|
await init_auth_db(db_path)
|
|
return db_path
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_llm_config_singleton():
|
|
"""Reset the LlmConfigService singleton before and after each test."""
|
|
set_llm_config_service(None)
|
|
yield
|
|
set_llm_config_service(None)
|
|
|
|
|
|
@pytest.fixture
|
|
def admin_app(config_path: Path, tmp_auth_db: Path) -> FastAPI:
|
|
"""A minimal FastAPI app with only the admin router mounted.
|
|
|
|
The ``_require_admin`` dependency is overridden to return a fake
|
|
admin user. The :class:`LlmConfigService` singleton is pre-populated
|
|
so the routes use the temp YAML file.
|
|
"""
|
|
# Inject the LlmConfigService singleton pointing at the temp YAML.
|
|
set_llm_config_service(LlmConfigService(config_path))
|
|
|
|
app = FastAPI()
|
|
app.state.auth_db_path = str(tmp_auth_db)
|
|
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
|
|
|
|
|
|
@pytest.fixture
|
|
def admin_client(admin_app: FastAPI) -> TestClient:
|
|
return TestClient(admin_app)
|
|
|
|
|
|
def _make_admin_user() -> dict[str, Any]:
|
|
return {"user_id": "admin-1", "username": "admin", "role": "admin"}
|
|
|
|
|
|
def _raise_forbidden() -> dict[str, Any]:
|
|
"""Dependency override that simulates a non-admin (403) response."""
|
|
raise HTTPException(status_code=403, detail="Admin permission required")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Provider CRUD
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestListProviders:
|
|
def test_list_returns_existing_providers(self, admin_client: TestClient):
|
|
resp = admin_client.get("/api/v1/admin/llm/providers")
|
|
assert resp.status_code == 200
|
|
providers = resp.json()
|
|
assert len(providers) == 1
|
|
assert providers[0]["name"] == "openai"
|
|
|
|
def test_list_masks_api_keys(self, admin_client: TestClient):
|
|
resp = admin_client.get("/api/v1/admin/llm/providers")
|
|
providers = resp.json()
|
|
assert providers[0]["api_key"].startswith("****")
|
|
# Should not contain the plaintext key.
|
|
assert "sk-test-12345678" not in providers[0]["api_key"]
|
|
|
|
|
|
class TestCreateProvider:
|
|
def test_create_returns_201(self, admin_client: TestClient):
|
|
resp = admin_client.post(
|
|
"/api/v1/admin/llm/providers",
|
|
json={
|
|
"name": "anthropic",
|
|
"type": "anthropic",
|
|
"api_key": "sk-ant-test-abcd1234",
|
|
"base_url": "https://api.anthropic.com",
|
|
"models": {"claude-sonnet-4-20250514": {}},
|
|
"max_tokens": 8192,
|
|
"timeout": 90.0,
|
|
},
|
|
)
|
|
assert resp.status_code == 201
|
|
body = resp.json()
|
|
assert body["name"] == "anthropic"
|
|
assert body["type"] == "anthropic"
|
|
# API key should be masked.
|
|
assert body["api_key"].startswith("****")
|
|
|
|
def test_create_duplicate_returns_409(self, admin_client: TestClient):
|
|
resp = admin_client.post(
|
|
"/api/v1/admin/llm/providers",
|
|
json={
|
|
"name": "openai",
|
|
"type": "openai",
|
|
"api_key": "sk-duplicate",
|
|
},
|
|
)
|
|
assert resp.status_code == 409
|
|
|
|
def test_non_admin_returns_403(self, admin_app: FastAPI):
|
|
admin_app.dependency_overrides[admin_routes_module._require_admin] = _raise_forbidden
|
|
client = TestClient(admin_app)
|
|
resp = client.post(
|
|
"/api/v1/admin/llm/providers",
|
|
json={"name": "x", "type": "openai", "api_key": "sk-x"},
|
|
)
|
|
assert resp.status_code == 403
|
|
|
|
|
|
class TestUpdateProvider:
|
|
def test_update_partial_returns_200(self, admin_client: TestClient):
|
|
resp = admin_client.patch(
|
|
"/api/v1/admin/llm/providers/openai",
|
|
json={"max_tokens": 2048},
|
|
)
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert body["max_tokens"] == 2048
|
|
# Other fields should be preserved.
|
|
assert body["base_url"] == "https://api.openai.com/v1"
|
|
|
|
def test_update_unknown_returns_404(self, admin_client: TestClient):
|
|
resp = admin_client.patch(
|
|
"/api/v1/admin/llm/providers/nonexistent",
|
|
json={"max_tokens": 2048},
|
|
)
|
|
assert resp.status_code == 404
|
|
|
|
def test_update_with_masked_api_key_preserves_existing(
|
|
self, admin_client: TestClient, config_path: Path
|
|
):
|
|
resp = admin_client.patch(
|
|
"/api/v1/admin/llm/providers/openai",
|
|
json={"api_key": "****5678"},
|
|
)
|
|
assert resp.status_code == 200
|
|
# Verify the YAML file still has the original key.
|
|
with open(config_path, encoding="utf-8") as f:
|
|
saved = yaml.safe_load(f)
|
|
assert saved["llm"]["providers"]["openai"]["api_key"] == "sk-test-12345678"
|
|
|
|
|
|
class TestDeleteProvider:
|
|
def test_delete_returns_200(self, admin_client: TestClient):
|
|
resp = admin_client.delete("/api/v1/admin/llm/providers/openai")
|
|
assert resp.status_code == 200
|
|
assert resp.json() == {"deleted": True}
|
|
|
|
def test_delete_unknown_returns_404(self, admin_client: TestClient):
|
|
resp = admin_client.delete("/api/v1/admin/llm/providers/nonexistent")
|
|
assert resp.status_code == 404
|
|
|
|
def test_delete_provider_used_in_fallback_returns_400(self, admin_client: TestClient):
|
|
# Add a second provider and a fallback chain that references it.
|
|
admin_client.post(
|
|
"/api/v1/admin/llm/providers",
|
|
json={
|
|
"name": "anthropic",
|
|
"type": "anthropic",
|
|
"api_key": "sk-ant-test",
|
|
},
|
|
)
|
|
admin_client.put(
|
|
"/api/v1/admin/llm/fallbacks/gpt-4o",
|
|
json={"chain": ["anthropic"]},
|
|
)
|
|
resp = admin_client.delete("/api/v1/admin/llm/providers/anthropic")
|
|
assert resp.status_code == 400
|
|
assert "fallback" in resp.json()["detail"].lower()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# API key management
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSetApiKey:
|
|
def test_set_api_key_returns_200(self, admin_client: TestClient):
|
|
resp = admin_client.post(
|
|
"/api/v1/admin/llm/providers/openai/api-key",
|
|
json={"api_key": "sk-brand-new-key-1234"},
|
|
)
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert body["api_key"].startswith("****")
|
|
|
|
def test_set_api_key_writes_to_env_file(self, admin_client: TestClient, config_path: Path):
|
|
admin_client.post(
|
|
"/api/v1/admin/llm/providers/openai/api-key",
|
|
json={"api_key": "sk-brand-new-key-1234"},
|
|
)
|
|
env_path = config_path.parent / ".env"
|
|
assert env_path.exists()
|
|
with open(env_path, encoding="utf-8") as f:
|
|
content = f.read()
|
|
assert "OPENAI_API_KEY=sk-brand-new-key-1234" in content
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fallback chain management
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFallbackRoutes:
|
|
def test_get_fallbacks_returns_empty_when_none(self, admin_client: TestClient):
|
|
resp = admin_client.get("/api/v1/admin/llm/fallbacks")
|
|
assert resp.status_code == 200
|
|
assert resp.json() == {}
|
|
|
|
def test_set_fallback_returns_200(self, admin_client: TestClient):
|
|
resp = admin_client.put(
|
|
"/api/v1/admin/llm/fallbacks/gpt-4o",
|
|
json={"chain": ["openai", "anthropic"]},
|
|
)
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert body["model"] == "gpt-4o"
|
|
assert body["chain"] == ["openai", "anthropic"]
|
|
|
|
def test_get_fallbacks_returns_set_chain(self, admin_client: TestClient):
|
|
admin_client.put(
|
|
"/api/v1/admin/llm/fallbacks/gpt-4o",
|
|
json={"chain": ["openai", "anthropic"]},
|
|
)
|
|
resp = admin_client.get("/api/v1/admin/llm/fallbacks")
|
|
assert resp.status_code == 200
|
|
assert resp.json() == {"gpt-4o": ["openai", "anthropic"]}
|
|
|
|
def test_delete_fallback_returns_200(self, admin_client: TestClient):
|
|
admin_client.put(
|
|
"/api/v1/admin/llm/fallbacks/gpt-4o",
|
|
json={"chain": ["openai"]},
|
|
)
|
|
resp = admin_client.delete("/api/v1/admin/llm/fallbacks/gpt-4o")
|
|
assert resp.status_code == 200
|
|
assert resp.json() == {"deleted": True}
|
|
|
|
def test_delete_unknown_fallback_returns_404(self, admin_client: TestClient):
|
|
resp = admin_client.delete("/api/v1/admin/llm/fallbacks/nonexistent")
|
|
assert resp.status_code == 404
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Department quota endpoints
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestDepartmentQuotaRoutes:
|
|
def test_set_quota_returns_200(self, admin_client: TestClient):
|
|
dept_id = str(uuid.uuid4())
|
|
resp = admin_client.put(
|
|
f"/api/v1/admin/departments/{dept_id}/quotas",
|
|
json={
|
|
"quota_type": "token_limit",
|
|
"limit_value": 1000,
|
|
"period": "daily",
|
|
},
|
|
)
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert body["quota_type"] == "token_limit"
|
|
assert body["limit_value"] == 1000
|
|
assert body["period"] == "daily"
|
|
|
|
def test_set_quota_invalid_type_returns_400(self, admin_client: TestClient):
|
|
dept_id = str(uuid.uuid4())
|
|
resp = admin_client.put(
|
|
f"/api/v1/admin/departments/{dept_id}/quotas",
|
|
json={"quota_type": "invalid", "limit_value": 1000},
|
|
)
|
|
assert resp.status_code == 400
|
|
|
|
def test_set_quota_model_whitelist_accepts_list(self, admin_client: TestClient):
|
|
dept_id = str(uuid.uuid4())
|
|
resp = admin_client.put(
|
|
f"/api/v1/admin/departments/{dept_id}/quotas",
|
|
json={
|
|
"quota_type": "model_whitelist",
|
|
"limit_value": ["gpt-4o", "claude"],
|
|
"period": "daily",
|
|
},
|
|
)
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert body["limit_value"] == ["gpt-4o", "claude"]
|
|
|
|
def test_list_quotas_returns_empty_for_new_department(self, admin_client: TestClient):
|
|
dept_id = str(uuid.uuid4())
|
|
resp = admin_client.get(f"/api/v1/admin/departments/{dept_id}/quotas")
|
|
assert resp.status_code == 200
|
|
assert resp.json() == []
|
|
|
|
def test_list_quotas_returns_set_quotas(self, admin_client: TestClient):
|
|
dept_id = str(uuid.uuid4())
|
|
admin_client.put(
|
|
f"/api/v1/admin/departments/{dept_id}/quotas",
|
|
json={"quota_type": "token_limit", "limit_value": 1000},
|
|
)
|
|
admin_client.put(
|
|
f"/api/v1/admin/departments/{dept_id}/quotas",
|
|
json={
|
|
"quota_type": "cost_limit",
|
|
"limit_value": 5000,
|
|
"period": "monthly",
|
|
},
|
|
)
|
|
resp = admin_client.get(f"/api/v1/admin/departments/{dept_id}/quotas")
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert len(body) == 2
|
|
types = {q["quota_type"] for q in body}
|
|
assert types == {"token_limit", "cost_limit"}
|
|
|
|
def test_delete_quota_returns_200(self, admin_client: TestClient):
|
|
dept_id = str(uuid.uuid4())
|
|
admin_client.put(
|
|
f"/api/v1/admin/departments/{dept_id}/quotas",
|
|
json={"quota_type": "token_limit", "limit_value": 1000},
|
|
)
|
|
resp = admin_client.delete(
|
|
f"/api/v1/admin/departments/{dept_id}/quotas",
|
|
params={"quota_type": "token_limit", "period": "daily"},
|
|
)
|
|
assert resp.status_code == 200
|
|
assert resp.json() == {"deleted": True}
|
|
|
|
def test_delete_unknown_quota_returns_404(self, admin_client: TestClient):
|
|
dept_id = str(uuid.uuid4())
|
|
resp = admin_client.delete(
|
|
f"/api/v1/admin/departments/{dept_id}/quotas",
|
|
params={"quota_type": "token_limit", "period": "daily"},
|
|
)
|
|
assert resp.status_code == 404
|
|
|
|
def test_delete_quota_invalid_period_returns_400(self, admin_client: TestClient):
|
|
dept_id = str(uuid.uuid4())
|
|
resp = admin_client.delete(
|
|
f"/api/v1/admin/departments/{dept_id}/quotas",
|
|
params={"quota_type": "token_limit", "period": "weekly"},
|
|
)
|
|
assert resp.status_code == 400
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Non-admin access
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestNonAdminAccess:
|
|
"""All LLM config endpoints must return 403 for non-admin users."""
|
|
|
|
def test_non_admin_cannot_list_providers(self, admin_app: FastAPI):
|
|
admin_app.dependency_overrides[admin_routes_module._require_admin] = _raise_forbidden
|
|
client = TestClient(admin_app)
|
|
resp = client.get("/api/v1/admin/llm/providers")
|
|
assert resp.status_code == 403
|
|
|
|
def test_non_admin_cannot_create_provider(self, admin_app: FastAPI):
|
|
admin_app.dependency_overrides[admin_routes_module._require_admin] = _raise_forbidden
|
|
client = TestClient(admin_app)
|
|
resp = client.post(
|
|
"/api/v1/admin/llm/providers",
|
|
json={"name": "x", "type": "openai", "api_key": "sk-x"},
|
|
)
|
|
assert resp.status_code == 403
|
|
|
|
def test_non_admin_cannot_delete_provider(self, admin_app: FastAPI):
|
|
admin_app.dependency_overrides[admin_routes_module._require_admin] = _raise_forbidden
|
|
client = TestClient(admin_app)
|
|
resp = client.delete("/api/v1/admin/llm/providers/openai")
|
|
assert resp.status_code == 403
|
|
|
|
def test_non_admin_cannot_set_api_key(self, admin_app: FastAPI):
|
|
admin_app.dependency_overrides[admin_routes_module._require_admin] = _raise_forbidden
|
|
client = TestClient(admin_app)
|
|
resp = client.post(
|
|
"/api/v1/admin/llm/providers/openai/api-key",
|
|
json={"api_key": "sk-x"},
|
|
)
|
|
assert resp.status_code == 403
|
|
|
|
def test_non_admin_cannot_get_fallbacks(self, admin_app: FastAPI):
|
|
admin_app.dependency_overrides[admin_routes_module._require_admin] = _raise_forbidden
|
|
client = TestClient(admin_app)
|
|
resp = client.get("/api/v1/admin/llm/fallbacks")
|
|
assert resp.status_code == 403
|
|
|
|
def test_non_admin_cannot_set_quota(self, admin_app: FastAPI):
|
|
admin_app.dependency_overrides[admin_routes_module._require_admin] = _raise_forbidden
|
|
client = TestClient(admin_app)
|
|
resp = client.put(
|
|
f"/api/v1/admin/departments/{uuid.uuid4()}/quotas",
|
|
json={"quota_type": "token_limit", "limit_value": 1000},
|
|
)
|
|
assert resp.status_code == 403
|