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.
56 lines
2.0 KiB
Python
56 lines
2.0 KiB
Python
"""Tests for docsgpt/quotas/windows.py."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
import pytest
|
|
|
|
from docsgpt.quotas.windows import window_bounds
|
|
|
|
UTC = timezone.utc
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestWindowBounds:
|
|
@pytest.mark.parametrize(
|
|
"period, start, end",
|
|
[
|
|
("day", datetime(2026, 9, 23, tzinfo=UTC), datetime(2026, 9, 24, tzinfo=UTC)),
|
|
("week", datetime(2026, 9, 21, tzinfo=UTC), datetime(2026, 9, 28, tzinfo=UTC)),
|
|
("month", datetime(2026, 9, 1, tzinfo=UTC), datetime(2026, 10, 1, tzinfo=UTC)),
|
|
],
|
|
)
|
|
def test_midweek(self, period, start, end):
|
|
now = datetime(2026, 9, 23, 15, 30, 12, 99, tzinfo=UTC) # a Wednesday
|
|
assert window_bounds(period, now) == (start, end)
|
|
|
|
def test_week_starts_on_monday_itself(self):
|
|
monday = datetime(2026, 9, 21, 0, 0, tzinfo=UTC)
|
|
assert window_bounds("week", monday)[0] == monday
|
|
|
|
def test_month_rolls_over_the_year(self):
|
|
start, end = window_bounds("month", datetime(2026, 12, 31, 23, 59, tzinfo=UTC))
|
|
assert (start, end) == (datetime(2026, 12, 1, tzinfo=UTC), datetime(2027, 1, 1, tzinfo=UTC))
|
|
|
|
def test_leap_february(self):
|
|
start, end = window_bounds("month", datetime(2028, 2, 29, 12, tzinfo=UTC))
|
|
assert (end - start).days == 29
|
|
|
|
def test_other_timezones_are_read_in_utc(self):
|
|
tokyo = timezone(timedelta(hours=9))
|
|
# 2026-10-01 08:00 in Tokyo is still 2026-09-30 in UTC.
|
|
start, _ = window_bounds("month", datetime(2026, 10, 1, 8, tzinfo=tokyo))
|
|
assert start == datetime(2026, 9, 1, tzinfo=UTC)
|
|
|
|
def test_naive_datetimes_are_utc(self):
|
|
assert window_bounds("day", datetime(2026, 9, 23, 5))[0] == datetime(2026, 9, 23, tzinfo=UTC)
|
|
|
|
def test_defaults_to_now(self):
|
|
start, end = window_bounds("day")
|
|
assert start <= datetime.now(UTC) < end
|
|
|
|
def test_unknown_period(self):
|
|
with pytest.raises(ValueError):
|
|
window_bounds("year")
|