mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-11 03:12:55 +00:00
Alex
6826313b60
docsgpt/quotas resolves a user's effective token and cost limits from the policy rows that apply to them: their own override, then the most generous allowance among their teams, then the instance default, then any defaults a registered provider supplies. Each budget resolves on its own, and a team counts once however many memberships the user holds in it. QuotaService compares those limits with the user's token_usage totals over the current QUOTA_PERIOD window (calendar-aligned, UTC, computed at read time). It fails open, and skips the usage query for unlimited users.
238 lines
9.6 KiB
Python
238 lines
9.6 KiB
Python
"""Tests for QuotaService against a real Postgres instance."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
import pytest
|
|
from sqlalchemy import text
|
|
|
|
from docsgpt.quotas import providers
|
|
from docsgpt.quotas.providers import QuotaDefaultsProvider, register_defaults_provider
|
|
from docsgpt.quotas.service import QuotaService
|
|
from docsgpt.storage.db.repositories.quota_policies import QuotaPoliciesRepository
|
|
from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository
|
|
|
|
NOW = datetime(2026, 9, 23, 12, tzinfo=timezone.utc)
|
|
THIS_MONTH = NOW - timedelta(days=2)
|
|
LAST_MONTH = NOW - timedelta(days=40)
|
|
|
|
|
|
@pytest.fixture
|
|
def conn(pg_conn, monkeypatch):
|
|
@contextmanager
|
|
def _readonly():
|
|
yield pg_conn
|
|
|
|
monkeypatch.setattr("docsgpt.quotas.service.db_readonly", _readonly)
|
|
monkeypatch.setattr("docsgpt.quotas.service.settings.QUOTA_PERIOD", "month")
|
|
return pg_conn
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _restore_provider():
|
|
original = providers.get_defaults_provider()
|
|
yield
|
|
register_defaults_provider(original)
|
|
|
|
|
|
def _use(conn, user_id="u1", tokens=0, cost=0.0, api_key=None, when=THIS_MONTH, source="agent_stream"):
|
|
TokenUsageRepository(conn).insert(
|
|
user_id=user_id, api_key=api_key, prompt_tokens=tokens, cost=cost, timestamp=when, source=source
|
|
)
|
|
|
|
|
|
def _policy(conn, scope, subject_id=None, **fields):
|
|
return QuotaPoliciesRepository(conn).upsert(scope=scope, subject_id=subject_id, **fields)
|
|
|
|
|
|
def _team_with_member(conn, slug, user_id="u1"):
|
|
team_id = str(
|
|
conn.execute(
|
|
text("INSERT INTO teams (name, slug, owner_id) VALUES (:n, :s, 'o') RETURNING id"),
|
|
{"n": slug, "s": slug},
|
|
).scalar()
|
|
)
|
|
conn.execute(
|
|
text("INSERT INTO team_members (team_id, user_id, role) VALUES (CAST(:t AS uuid), :u, 'team_member')"),
|
|
{"t": team_id, "u": user_id},
|
|
)
|
|
return team_id
|
|
|
|
|
|
class TestCheck:
|
|
def test_no_policies_allows(self, conn):
|
|
_use(conn, tokens=10**9)
|
|
assert QuotaService.check("u1", now=NOW) is None
|
|
|
|
def test_no_user_allows(self, conn):
|
|
assert QuotaService.check(None, now=NOW) is None
|
|
|
|
def test_under_the_limit_allows(self, conn):
|
|
_policy(conn, "instance", token_limit=100)
|
|
_use(conn, tokens=99)
|
|
assert QuotaService.check("u1", now=NOW) is None
|
|
|
|
def test_reaching_the_token_limit_blocks(self, conn):
|
|
_policy(conn, "instance", token_limit=100)
|
|
_use(conn, tokens=100)
|
|
exceeded = QuotaService.check("u1", now=NOW)
|
|
assert (exceeded.budget, exceeded.usage, exceeded.limit, exceeded.source) == ("tokens", 100, 100.0, "instance")
|
|
assert exceeded.resets_at == datetime(2026, 10, 1, tzinfo=timezone.utc)
|
|
|
|
def test_cost_limit_blocks(self, conn):
|
|
_policy(conn, "user", "u1", cost_limit_usd=1.5)
|
|
_use(conn, cost=1.0)
|
|
_use(conn, cost=0.5)
|
|
exceeded = QuotaService.check("u1", now=NOW)
|
|
assert (exceeded.budget, exceeded.usage, exceeded.source) == ("cost", 1.5, "user")
|
|
|
|
def test_zero_limit_blocks_without_usage(self, conn):
|
|
_policy(conn, "user", "u1", token_limit=0)
|
|
assert QuotaService.check("u1", now=NOW).limit == 0
|
|
|
|
def test_last_periods_usage_does_not_count(self, conn):
|
|
_policy(conn, "instance", token_limit=100)
|
|
_use(conn, tokens=500, when=LAST_MONTH)
|
|
assert QuotaService.check("u1", now=NOW) is None
|
|
|
|
def test_other_users_usage_does_not_count(self, conn):
|
|
_policy(conn, "instance", token_limit=100)
|
|
_use(conn, user_id="u2", tokens=500)
|
|
assert QuotaService.check("u1", now=NOW) is None
|
|
|
|
def test_scheduler_rollups_do_not_count(self, conn):
|
|
_policy(conn, "instance", token_limit=100)
|
|
_use(conn, tokens=60)
|
|
_use(conn, tokens=60, source="schedule")
|
|
assert QuotaService.check("u1", now=NOW) is None
|
|
|
|
def test_user_override_lifts_the_instance_limit(self, conn):
|
|
_policy(conn, "instance", token_limit=100)
|
|
_policy(conn, "user", "u1", token_unlimited=True)
|
|
_use(conn, tokens=10**6)
|
|
assert QuotaService.check("u1", now=NOW) is None
|
|
|
|
def test_period_setting_moves_the_window(self, conn, monkeypatch):
|
|
monkeypatch.setattr("docsgpt.quotas.service.settings.QUOTA_PERIOD", "day")
|
|
_policy(conn, "instance", token_limit=100)
|
|
_use(conn, tokens=500, when=NOW - timedelta(days=1))
|
|
assert QuotaService.check("u1", now=NOW) is None
|
|
_use(conn, tokens=500, when=NOW - timedelta(hours=1))
|
|
assert QuotaService.check("u1", now=NOW).resets_at == datetime(2026, 9, 24, tzinfo=timezone.utc)
|
|
|
|
def test_unknown_bucket_rejected(self, conn):
|
|
with pytest.raises(ValueError):
|
|
QuotaService.check("u1", bucket="all", now=NOW)
|
|
|
|
def test_failure_allows_the_request(self, monkeypatch):
|
|
def boom():
|
|
raise RuntimeError("db down")
|
|
|
|
monkeypatch.setattr("docsgpt.quotas.service.db_readonly", boom)
|
|
assert QuotaService.check("u1", now=NOW) is None
|
|
|
|
|
|
class TestTeams:
|
|
def test_the_most_generous_team_sets_the_allowance(self, conn):
|
|
_policy(conn, "team", _team_with_member(conn, "svc-small"), token_limit=100)
|
|
big = _team_with_member(conn, "svc-big")
|
|
_policy(conn, "team", big, token_limit=1000)
|
|
_use(conn, tokens=500)
|
|
assert QuotaService.check("u1", now=NOW) is None
|
|
_use(conn, tokens=500)
|
|
exceeded = QuotaService.check("u1", now=NOW)
|
|
assert (exceeded.limit, exceeded.source, exceeded.source_id) == (1000.0, "team", big)
|
|
|
|
def test_usage_is_one_total_not_one_per_team(self, conn):
|
|
for slug in ("svc-a", "svc-b", "svc-c"):
|
|
_policy(conn, "team", _team_with_member(conn, slug), token_limit=100)
|
|
_use(conn, tokens=100)
|
|
assert QuotaService.check("u1", now=NOW).limit == 100.0
|
|
|
|
def test_a_team_only_covers_its_members(self, conn):
|
|
_policy(conn, "instance", token_limit=100)
|
|
_policy(conn, "team", _team_with_member(conn, "svc-vip", user_id="u2"), token_unlimited=True)
|
|
_use(conn, tokens=100)
|
|
assert QuotaService.check("u1", now=NOW).source == "instance"
|
|
|
|
|
|
class TestBuckets:
|
|
def test_bucket_policies_only_see_their_traffic(self, conn):
|
|
_policy(conn, "instance", bucket="agent", token_limit=100)
|
|
_use(conn, tokens=500)
|
|
_use(conn, tokens=90, api_key="k")
|
|
assert QuotaService.check("u1", "direct", now=NOW) is None
|
|
assert QuotaService.check("u1", "agent", now=NOW) is None
|
|
_use(conn, tokens=10, api_key="k")
|
|
exceeded = QuotaService.check("u1", "agent", now=NOW)
|
|
assert (exceeded.bucket, exceeded.usage) == ("agent", 100)
|
|
assert QuotaService.check("u1", "direct", now=NOW) is None
|
|
|
|
def test_the_all_bucket_applies_to_both_kinds_of_traffic(self, conn):
|
|
_policy(conn, "instance", token_limit=100)
|
|
_use(conn, tokens=60)
|
|
_use(conn, tokens=60, api_key="k")
|
|
assert QuotaService.check("u1", "direct", now=NOW).bucket == "all"
|
|
assert QuotaService.check("u1", "agent", now=NOW).bucket == "all"
|
|
|
|
|
|
class TestProviderDefaults:
|
|
class _Plan(QuotaDefaultsProvider):
|
|
def default_policies(self, user_id):
|
|
return [{"cost_limit_usd": 5.0}] if user_id == "u1" else []
|
|
|
|
def error_payload(self, payload, user_id):
|
|
return {**payload, "error_code": "free-limit-reached"}
|
|
|
|
def test_defaults_apply_without_stored_rows(self, conn):
|
|
register_defaults_provider(self._Plan())
|
|
_use(conn, cost=5.0)
|
|
exceeded = QuotaService.check("u1", now=NOW)
|
|
assert exceeded.source == "default"
|
|
assert exceeded.to_payload()["error_code"] == "free-limit-reached"
|
|
|
|
def test_a_broken_provider_payload_falls_back(self, conn):
|
|
class _Broken(self._Plan):
|
|
def error_payload(self, payload, user_id):
|
|
raise RuntimeError("nope")
|
|
|
|
register_defaults_provider(_Broken())
|
|
_use(conn, cost=5.0)
|
|
assert QuotaService.check("u1", now=NOW).to_payload()["error_code"] == "quota-exceeded"
|
|
|
|
|
|
class TestStatusAndPayload:
|
|
def test_status_reports_limits_and_usage(self, conn):
|
|
_policy(conn, "instance", token_limit=100, cost_limit_usd=2)
|
|
_use(conn, tokens=40, cost=0.5)
|
|
(status,) = QuotaService.status("u1", now=NOW)
|
|
assert status.to_dict() == {
|
|
"bucket": "all",
|
|
"tokens": {"limit": 100.0, "used": 40, "source": "instance", "source_id": None},
|
|
"cost": {"limit": 2.0, "used": 0.5, "source": "instance", "source_id": None},
|
|
"resets_at": "2026-10-01T00:00:00+00:00",
|
|
}
|
|
|
|
def test_unlimited_users_skip_the_usage_query(self, conn, monkeypatch):
|
|
def fail(*args, **kwargs):
|
|
raise AssertionError("usage must not be summed for an unlimited user")
|
|
|
|
monkeypatch.setattr(TokenUsageRepository, "usage_totals", fail)
|
|
(status,) = QuotaService.status("u1", now=NOW)
|
|
assert status.limits.unlimited and status.tokens_used == 0
|
|
|
|
def test_payload_shape(self, conn):
|
|
_policy(conn, "user", "u1", cost_limit_usd=1)
|
|
_use(conn, cost=1.25)
|
|
exceeded = QuotaService.check("u1", now=NOW)
|
|
payload = exceeded.to_payload()
|
|
assert payload["success"] is False
|
|
assert payload["error_code"] == "quota-exceeded"
|
|
assert (payload["dimension"], payload["unit"]) == ("cost", "USD")
|
|
assert (payload["usage"], payload["limit"]) == (1.25, 1.0)
|
|
assert payload["resets_at"] == "2026-10-01T00:00:00+00:00"
|
|
assert "$1.25 of $1.00" in payload["message"]
|
|
assert exceeded.retry_after_seconds >= 1
|