334 lines
14 KiB
Python
334 lines
14 KiB
Python
"""Unit tests for QuotaService (U5 — per-department LLM quotas).
|
|
|
|
Covers:
|
|
- Happy path: set → get → list → delete quotas
|
|
- Quota check: token_limit (allowed / denied)
|
|
- Quota check: cost_limit (allowed / denied)
|
|
- Quota check: no quota set → always allowed
|
|
- model_whitelist serialization (list → JSON → list)
|
|
- Upsert: set same quota twice updates the value
|
|
- Validation: invalid quota_type / period raises ValueError
|
|
- is_model_allowed: whitelist enforcement
|
|
- Singleton helpers (get/set_quota_service)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from agentkit.server.admin.quota_service import (
|
|
QuotaService,
|
|
get_quota_service,
|
|
set_quota_service,
|
|
)
|
|
from agentkit.server.auth.models import init_auth_db
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
async def fresh_db(tmp_path: Path) -> Path:
|
|
"""A brand-new auth DB on a fresh path (no data)."""
|
|
db_path = tmp_path / "auth.db"
|
|
await init_auth_db(db_path)
|
|
return db_path
|
|
|
|
|
|
@pytest.fixture
|
|
def service() -> QuotaService:
|
|
return QuotaService()
|
|
|
|
|
|
def _random_dept_id() -> str:
|
|
return str(uuid.uuid4())
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Quota CRUD happy path
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestQuotaCrudHappyPath:
|
|
async def test_set_token_limit_returns_quota_dict(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
quota = await service.set_quota(fresh_db, dept_id, "token_limit", 1000, period="daily")
|
|
assert quota["department_id"] == dept_id
|
|
assert quota["quota_type"] == "token_limit"
|
|
assert quota["limit_value"] == 1000
|
|
assert quota["period"] == "daily"
|
|
assert "id" in quota
|
|
assert "updated_at" in quota
|
|
|
|
async def test_set_cost_limit_returns_quota_dict(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
quota = await service.set_quota(fresh_db, dept_id, "cost_limit", 5000, period="monthly")
|
|
assert quota["quota_type"] == "cost_limit"
|
|
assert quota["limit_value"] == 5000
|
|
assert quota["period"] == "monthly"
|
|
|
|
async def test_set_model_whitelist_serializes_list(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
whitelist = ["gpt-4o", "claude-sonnet-4-20250514", "gemini-pro"]
|
|
quota = await service.set_quota(
|
|
fresh_db, dept_id, "model_whitelist", whitelist, period="daily"
|
|
)
|
|
assert quota["quota_type"] == "model_whitelist"
|
|
# The deserialized value should match the input list.
|
|
assert quota["limit_value"] == whitelist
|
|
|
|
async def test_get_returns_set_quota(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
await service.set_quota(fresh_db, dept_id, "token_limit", 1000)
|
|
fetched = await service.get_quota(fresh_db, dept_id, "token_limit")
|
|
assert fetched is not None
|
|
assert fetched["limit_value"] == 1000
|
|
assert fetched["quota_type"] == "token_limit"
|
|
|
|
async def test_get_returns_none_for_unset_quota(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
fetched = await service.get_quota(fresh_db, dept_id, "token_limit")
|
|
assert fetched is None
|
|
|
|
async def test_list_returns_all_quotas_for_department(
|
|
self, service: QuotaService, fresh_db: Path
|
|
):
|
|
dept_id = _random_dept_id()
|
|
await service.set_quota(fresh_db, dept_id, "token_limit", 1000, period="daily")
|
|
await service.set_quota(fresh_db, dept_id, "cost_limit", 5000, period="monthly")
|
|
await service.set_quota(fresh_db, dept_id, "model_whitelist", ["gpt-4o"], period="daily")
|
|
quotas = await service.list_department_quotas(fresh_db, dept_id)
|
|
assert len(quotas) == 3
|
|
types = {q["quota_type"] for q in quotas}
|
|
assert types == {"token_limit", "cost_limit", "model_whitelist"}
|
|
|
|
async def test_list_returns_empty_for_department_with_no_quotas(
|
|
self, service: QuotaService, fresh_db: Path
|
|
):
|
|
dept_id = _random_dept_id()
|
|
quotas = await service.list_department_quotas(fresh_db, dept_id)
|
|
assert quotas == []
|
|
|
|
async def test_delete_removes_quota(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
await service.set_quota(fresh_db, dept_id, "token_limit", 1000)
|
|
deleted = await service.delete_quota(fresh_db, dept_id, "token_limit")
|
|
assert deleted is True
|
|
# Confirm it's gone.
|
|
assert await service.get_quota(fresh_db, dept_id, "token_limit") is None
|
|
|
|
async def test_delete_returns_false_for_unset_quota(
|
|
self, service: QuotaService, fresh_db: Path
|
|
):
|
|
dept_id = _random_dept_id()
|
|
deleted = await service.delete_quota(fresh_db, dept_id, "token_limit")
|
|
assert deleted is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Upsert behavior
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestQuotaUpsert:
|
|
async def test_set_same_quota_twice_updates_value(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
await service.set_quota(fresh_db, dept_id, "token_limit", 1000)
|
|
updated = await service.set_quota(fresh_db, dept_id, "token_limit", 2000)
|
|
assert updated["limit_value"] == 2000
|
|
# Only one row should exist.
|
|
quotas = await service.list_department_quotas(fresh_db, dept_id)
|
|
assert len(quotas) == 1
|
|
assert quotas[0]["limit_value"] == 2000
|
|
|
|
async def test_set_same_quota_with_different_period_creates_new_row(
|
|
self, service: QuotaService, fresh_db: Path
|
|
):
|
|
dept_id = _random_dept_id()
|
|
await service.set_quota(fresh_db, dept_id, "token_limit", 1000, period="daily")
|
|
await service.set_quota(fresh_db, dept_id, "token_limit", 30000, period="monthly")
|
|
quotas = await service.list_department_quotas(fresh_db, dept_id)
|
|
assert len(quotas) == 2
|
|
periods = {q["period"] for q in quotas}
|
|
assert periods == {"daily", "monthly"}
|
|
|
|
async def test_upsert_model_whitelist_replaces_list(
|
|
self, service: QuotaService, fresh_db: Path
|
|
):
|
|
dept_id = _random_dept_id()
|
|
await service.set_quota(fresh_db, dept_id, "model_whitelist", ["gpt-4o", "claude"])
|
|
await service.set_quota(fresh_db, dept_id, "model_whitelist", ["gpt-4o-mini"])
|
|
fetched = await service.get_quota(fresh_db, dept_id, "model_whitelist")
|
|
assert fetched is not None
|
|
assert fetched["limit_value"] == ["gpt-4o-mini"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Quota enforcement
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestQuotaCheck:
|
|
async def test_token_limit_allowed_when_under_limit(
|
|
self, service: QuotaService, fresh_db: Path
|
|
):
|
|
dept_id = _random_dept_id()
|
|
await service.set_quota(fresh_db, dept_id, "token_limit", 1000)
|
|
allowed, reason = await service.check_quota(
|
|
fresh_db, dept_id, "token_limit", "daily", current_usage=500
|
|
)
|
|
assert allowed is True
|
|
assert reason == "ok"
|
|
|
|
async def test_token_limit_denied_at_limit(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
await service.set_quota(fresh_db, dept_id, "token_limit", 1000)
|
|
allowed, reason = await service.check_quota(
|
|
fresh_db, dept_id, "token_limit", "daily", current_usage=1000
|
|
)
|
|
assert allowed is False
|
|
assert "exceeded" in reason
|
|
|
|
async def test_token_limit_denied_over_limit(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
await service.set_quota(fresh_db, dept_id, "token_limit", 1000)
|
|
allowed, reason = await service.check_quota(
|
|
fresh_db, dept_id, "token_limit", "daily", current_usage=1500
|
|
)
|
|
assert allowed is False
|
|
assert "exceeded" in reason
|
|
|
|
async def test_cost_limit_allowed_when_under_limit(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
await service.set_quota(fresh_db, dept_id, "cost_limit", 5000, period="monthly")
|
|
allowed, reason = await service.check_quota(
|
|
fresh_db, dept_id, "cost_limit", "monthly", current_usage=2500
|
|
)
|
|
assert allowed is True
|
|
assert reason == "ok"
|
|
|
|
async def test_cost_limit_denied_at_limit(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
await service.set_quota(fresh_db, dept_id, "cost_limit", 5000, period="monthly")
|
|
allowed, _ = await service.check_quota(
|
|
fresh_db, dept_id, "cost_limit", "monthly", current_usage=5000
|
|
)
|
|
assert allowed is False
|
|
|
|
async def test_no_quota_set_always_allowed(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
allowed, reason = await service.check_quota(
|
|
fresh_db, dept_id, "token_limit", "daily", current_usage=999999
|
|
)
|
|
assert allowed is True
|
|
assert reason == "ok"
|
|
|
|
async def test_model_whitelist_not_a_quota_check(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
await service.set_quota(fresh_db, dept_id, "model_whitelist", ["gpt-4o"], period="daily")
|
|
allowed, reason = await service.check_quota(
|
|
fresh_db, dept_id, "model_whitelist", "daily", current_usage=0
|
|
)
|
|
assert allowed is False
|
|
assert "is_model_allowed" in reason
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model whitelist enforcement
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestModelWhitelist:
|
|
async def test_is_model_allowed_when_in_whitelist(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
await service.set_quota(fresh_db, dept_id, "model_whitelist", ["gpt-4o", "claude"])
|
|
allowed, reason = await service.is_model_allowed(fresh_db, dept_id, "gpt-4o")
|
|
assert allowed is True
|
|
assert reason == "ok"
|
|
|
|
async def test_is_model_denied_when_not_in_whitelist(
|
|
self, service: QuotaService, fresh_db: Path
|
|
):
|
|
dept_id = _random_dept_id()
|
|
await service.set_quota(fresh_db, dept_id, "model_whitelist", ["gpt-4o", "claude"])
|
|
allowed, reason = await service.is_model_allowed(fresh_db, dept_id, "gemini-pro")
|
|
assert allowed is False
|
|
assert "not in" in reason
|
|
|
|
async def test_is_model_allowed_when_no_whitelist(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
# No whitelist set — all models allowed.
|
|
allowed, reason = await service.is_model_allowed(fresh_db, dept_id, "any-model")
|
|
assert allowed is True
|
|
assert reason == "ok"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Validation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestQuotaValidation:
|
|
async def test_set_invalid_quota_type_raises(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
with pytest.raises(ValueError, match="Invalid quota_type"):
|
|
await service.set_quota(fresh_db, dept_id, "invalid_type", 1000)
|
|
|
|
async def test_set_invalid_period_raises(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
with pytest.raises(ValueError, match="Invalid period"):
|
|
await service.set_quota(fresh_db, dept_id, "token_limit", 1000, period="weekly")
|
|
|
|
async def test_set_token_limit_with_list_raises(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
with pytest.raises(ValueError, match="must be an int"):
|
|
await service.set_quota(fresh_db, dept_id, "token_limit", ["gpt-4o"])
|
|
|
|
async def test_set_model_whitelist_with_int_raises(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
with pytest.raises(ValueError, match="must be a list"):
|
|
await service.set_quota(fresh_db, dept_id, "model_whitelist", 1000)
|
|
|
|
async def test_get_invalid_quota_type_raises(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
with pytest.raises(ValueError, match="Invalid quota_type"):
|
|
await service.get_quota(fresh_db, dept_id, "invalid_type")
|
|
|
|
async def test_delete_invalid_period_raises(self, service: QuotaService, fresh_db: Path):
|
|
dept_id = _random_dept_id()
|
|
with pytest.raises(ValueError, match="Invalid period"):
|
|
await service.delete_quota(fresh_db, dept_id, "token_limit", period="weekly")
|
|
|
|
async def test_check_quota_invalid_quota_type_raises(
|
|
self, service: QuotaService, fresh_db: Path
|
|
):
|
|
dept_id = _random_dept_id()
|
|
with pytest.raises(ValueError, match="Invalid quota_type"):
|
|
await service.check_quota(fresh_db, dept_id, "invalid_type", "daily", current_usage=0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Singleton helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSingletonHelpers:
|
|
def test_get_quota_service_returns_singleton(self):
|
|
# Save the original singleton so we don't disturb other tests.
|
|
original = get_quota_service()
|
|
try:
|
|
custom = QuotaService()
|
|
set_quota_service(custom)
|
|
assert get_quota_service() is custom
|
|
# Clearing falls back to a new lazy instance.
|
|
set_quota_service(None)
|
|
new_one = get_quota_service()
|
|
assert new_one is not custom
|
|
finally:
|
|
set_quota_service(original)
|