mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 04:15:35 +00:00
feat(quotas): limit resolution and the quota service
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.
This commit is contained in:
1 parent
6c42139224
commit
6826313b60
12 files changed
+1175
No files matched your search
@@ -0,0 +1,143 @@
|
||||
"""Tests for QuotaPoliciesRepository against a real Postgres instance."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
from docsgpt.storage.db.repositories.quota_policies import QuotaPoliciesRepository
|
||||
|
||||
|
||||
def _repo(conn) -> QuotaPoliciesRepository:
|
||||
return QuotaPoliciesRepository(conn)
|
||||
|
||||
|
||||
def _team(conn, slug: str) -> str:
|
||||
return str(
|
||||
conn.execute(
|
||||
text("INSERT INTO teams (name, slug, owner_id) VALUES (:n, :s, 'owner') RETURNING id"),
|
||||
{"n": slug, "s": slug},
|
||||
).scalar()
|
||||
)
|
||||
|
||||
|
||||
def _member(conn, team_id: str, user_id: str, role: str = "team_member", source: str = "manual") -> None:
|
||||
conn.execute(
|
||||
text(
|
||||
"INSERT INTO team_members (team_id, user_id, role, source) "
|
||||
"VALUES (CAST(:t AS uuid), :u, :r, :s)"
|
||||
),
|
||||
{"t": team_id, "u": user_id, "r": role, "s": source},
|
||||
)
|
||||
|
||||
|
||||
class TestUpsert:
|
||||
def test_creates_then_replaces(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
created = repo.upsert(scope="user", subject_id="u1", token_limit=100, note="trial", actor="admin1")
|
||||
assert created["token_limit"] == 100
|
||||
assert created["created_by"] == created["updated_by"] == "admin1"
|
||||
|
||||
replaced = repo.upsert(scope="user", subject_id="u1", cost_limit_usd=2.5, actor="admin2")
|
||||
assert replaced["id"] == created["id"]
|
||||
assert replaced["token_limit"] is None
|
||||
assert float(replaced["cost_limit_usd"]) == 2.5
|
||||
assert replaced["note"] is None
|
||||
assert (replaced["created_by"], replaced["updated_by"]) == ("admin1", "admin2")
|
||||
|
||||
def test_instance_row_is_a_singleton_per_bucket(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
repo.upsert(scope="instance", subject_id=None, token_limit=1)
|
||||
repo.upsert(scope="instance", subject_id=None, token_limit=2)
|
||||
repo.upsert(scope="instance", subject_id=None, bucket="agent", token_limit=3)
|
||||
rows = repo.list_by_scope("instance")
|
||||
assert [(r["bucket"], r["token_limit"]) for r in rows] == [("all", 2), ("agent", 3)]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kwargs",
|
||||
[
|
||||
{"scope": "org", "subject_id": "x"},
|
||||
{"scope": "user", "subject_id": None},
|
||||
{"scope": "instance", "subject_id": "x"},
|
||||
{"scope": "user", "subject_id": "u", "bucket": "nope"},
|
||||
{"scope": "user", "subject_id": "u", "token_limit": 1, "token_unlimited": True},
|
||||
{"scope": "user", "subject_id": "u", "cost_limit_usd": 1, "cost_unlimited": True},
|
||||
{"scope": "user", "subject_id": "u", "token_limit": -1},
|
||||
{"scope": "user", "subject_id": "u", "cost_limit_usd": -0.5},
|
||||
],
|
||||
)
|
||||
def test_rejects_invalid_policies(self, pg_conn, kwargs):
|
||||
with pytest.raises(ValueError):
|
||||
_repo(pg_conn).upsert(**kwargs)
|
||||
|
||||
|
||||
class TestReads:
|
||||
def test_get_and_list_for_subject(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
repo.upsert(scope="user", subject_id="u1", bucket="agent", token_limit=5)
|
||||
repo.upsert(scope="user", subject_id="u1", token_limit=9)
|
||||
assert repo.get("user", "u1")["token_limit"] == 9
|
||||
assert repo.get("user", "u1", "direct") is None
|
||||
assert [r["bucket"] for r in repo.list_for_subject("user", "u1")] == ["all", "agent"]
|
||||
assert repo.list_for_subject("user", "nobody") == []
|
||||
|
||||
def test_list_by_scope_rejects_unknown_scope(self, pg_conn):
|
||||
with pytest.raises(ValueError):
|
||||
_repo(pg_conn).list_by_scope("org")
|
||||
|
||||
|
||||
class TestPoliciesForUser:
|
||||
def test_collects_instance_user_and_team_rows(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
mine, other = _team(pg_conn, "qp-mine"), _team(pg_conn, "qp-other")
|
||||
_member(pg_conn, mine, "u1")
|
||||
repo.upsert(scope="instance", subject_id=None, token_limit=1)
|
||||
repo.upsert(scope="team", subject_id=mine, token_limit=2)
|
||||
repo.upsert(scope="team", subject_id=other, token_limit=3)
|
||||
repo.upsert(scope="user", subject_id="u1", token_limit=4)
|
||||
repo.upsert(scope="user", subject_id="u2", token_limit=5)
|
||||
|
||||
limits = sorted(r["token_limit"] for r in repo.policies_for_user("u1"))
|
||||
assert limits == [1, 2, 4]
|
||||
|
||||
def test_a_team_counts_once_however_many_memberships(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
team = _team(pg_conn, "qp-multi")
|
||||
_member(pg_conn, team, "u1", "team_member", "manual")
|
||||
_member(pg_conn, team, "u1", "team_admin", "manual")
|
||||
_member(pg_conn, team, "u1", "team_member", "oidc_group")
|
||||
repo.upsert(scope="team", subject_id=team, token_limit=7)
|
||||
assert [r["token_limit"] for r in repo.policies_for_user("u1")] == [7]
|
||||
|
||||
def test_every_team_of_the_user_is_included(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
for slug, limit in (("qp-a", 10), ("qp-b", 20), ("qp-c", 30)):
|
||||
team = _team(pg_conn, slug)
|
||||
_member(pg_conn, team, "u1")
|
||||
repo.upsert(scope="team", subject_id=team, token_limit=limit)
|
||||
assert sorted(r["token_limit"] for r in repo.policies_for_user("u1")) == [10, 20, 30]
|
||||
|
||||
def test_disabled_rows_are_left_out(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
repo.upsert(scope="user", subject_id="u1", token_limit=4, enabled=False)
|
||||
assert repo.policies_for_user("u1") == []
|
||||
|
||||
def test_leaving_a_team_drops_its_allowance(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
team = _team(pg_conn, "qp-leave")
|
||||
_member(pg_conn, team, "u1")
|
||||
repo.upsert(scope="team", subject_id=team, token_limit=7)
|
||||
pg_conn.execute(text("DELETE FROM team_members WHERE user_id = 'u1'"))
|
||||
assert repo.policies_for_user("u1") == []
|
||||
|
||||
|
||||
class TestDelete:
|
||||
def test_delete_one_bucket_or_all(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
repo.upsert(scope="user", subject_id="u1", token_limit=1)
|
||||
repo.upsert(scope="user", subject_id="u1", bucket="agent", token_limit=2)
|
||||
repo.upsert(scope="user", subject_id="u2", token_limit=3)
|
||||
assert repo.delete("user", "u1", "agent") == 1
|
||||
assert repo.delete("user", "u1", "agent") == 0
|
||||
assert repo.delete("user", "u1") == 1
|
||||
assert repo.get("user", "u2") is not None
|
||||
Reference in new issue
Block a user