"""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)