xsl
93b88db2f5
Covers routing, failover, the circuit breaker, structured-output retries, the cache-prefix validator, cost computation including cached-token pricing, the daily budget breaker, the quota protocol, the task state machine and log redaction. HTTP is mocked -- no unit test spends money. The integration test runs the real slice against PostgreSQL and Redis with a faked provider.
112 lines
4.0 KiB
Python
Executable File
112 lines
4.0 KiB
Python
Executable File
"""Quota three-phase arithmetic (架构文档 §7.4, FR-BILL-01/05/06)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
|
|
import pytest
|
|
|
|
from packages.domain.billing import quota
|
|
from packages.domain.errors import QuotaExceededError
|
|
|
|
|
|
def test_period_key_day_uses_utc_plus_8() -> None:
|
|
# 17:00Z is already 01:00 the next day in UTC+8.
|
|
assert quota.period_key(datetime(2026, 8, 13, 17, 0, tzinfo=UTC), "day") == "d:2026-08-14"
|
|
assert quota.period_key(datetime(2026, 8, 13, 15, 0, tzinfo=UTC), "day") == "d:2026-08-13"
|
|
|
|
|
|
def test_period_key_month_and_lifetime() -> None:
|
|
now = datetime(2026, 8, 13, 1, 0, tzinfo=UTC)
|
|
assert quota.period_key(now, "month") == "m:2026-08"
|
|
assert quota.period_key(now, "lifetime") == "life"
|
|
|
|
|
|
def test_unknown_period_type_is_rejected() -> None:
|
|
with pytest.raises(ValueError, match="unknown period_type"):
|
|
quota.period_key(datetime.now(UTC), "weekly") # type: ignore[arg-type]
|
|
|
|
|
|
def test_period_reset_at_rolls_to_the_next_local_midnight() -> None:
|
|
reset = quota.period_reset_at(datetime(2026, 8, 13, 1, 0, tzinfo=UTC), "day")
|
|
assert reset is not None
|
|
assert reset == datetime(2026, 8, 13, 16, 0, tzinfo=UTC) # 2026-08-14 00:00 +08
|
|
|
|
|
|
def test_period_reset_at_handles_december_rollover() -> None:
|
|
reset = quota.period_reset_at(datetime(2026, 12, 20, 1, 0, tzinfo=UTC), "month")
|
|
assert reset is not None
|
|
assert reset.year == 2026 # 2027-01-01 00:00 +08 == 2026-12-31 16:00 UTC
|
|
assert reset.month == 12
|
|
|
|
|
|
def test_lifetime_quota_never_resets() -> None:
|
|
assert quota.period_reset_at(datetime.now(UTC), "lifetime") is None
|
|
|
|
|
|
def test_task_bonus_raises_the_effective_limit() -> None:
|
|
state = quota.QuotaState(used=20, base_limit=20, bonus=3)
|
|
assert state.effective_limit == 23
|
|
assert state.remaining == 3
|
|
assert quota.evaluate(state, 1).allowed
|
|
|
|
|
|
def test_reserve_freezes_the_amount() -> None:
|
|
state = quota.reserve(quota.QuotaState(used=0, base_limit=20), 1)
|
|
assert state.used == 1
|
|
assert state.remaining == 19
|
|
|
|
|
|
def test_reserve_beyond_the_limit_is_rejected_server_side() -> None:
|
|
state = quota.QuotaState(used=20, base_limit=20)
|
|
with pytest.raises(QuotaExceededError) as excinfo:
|
|
quota.reserve(state, 1)
|
|
assert excinfo.value.remaining == 0
|
|
assert excinfo.value.required == 1
|
|
|
|
|
|
def test_reserve_rejects_a_batch_that_would_overshoot() -> None:
|
|
with pytest.raises(QuotaExceededError):
|
|
quota.reserve(quota.QuotaState(used=18, base_limit=20), 3)
|
|
|
|
|
|
def test_non_positive_amounts_are_programming_errors() -> None:
|
|
with pytest.raises(ValueError, match="positive"):
|
|
quota.evaluate(quota.QuotaState(used=0, base_limit=1), 0)
|
|
with pytest.raises(ValueError, match="positive"):
|
|
quota.refund(quota.QuotaState(used=1, base_limit=1), -1)
|
|
|
|
|
|
def test_refund_releases_the_reservation() -> None:
|
|
reserved = quota.reserve(quota.QuotaState(used=5, base_limit=20), 1)
|
|
assert quota.refund(reserved, 1).used == 5
|
|
|
|
|
|
def test_refund_never_drives_usage_negative() -> None:
|
|
assert quota.refund(quota.QuotaState(used=0, base_limit=20), 5).used == 0
|
|
|
|
|
|
def test_balance_after_records_remaining_quota() -> None:
|
|
assert quota.balance_after(quota.QuotaState(used=3, base_limit=20)) == 17
|
|
|
|
|
|
def test_refund_key_is_derived_so_a_retry_cannot_double_refund() -> None:
|
|
reserve_key = quota.reserve_idempotency_key(
|
|
user_id=7, quota_key="agent_chat", request_key="req-1"
|
|
)
|
|
assert reserve_key == "rsv:7:agent_chat:req-1"
|
|
assert quota.refund_idempotency_key(reserve_key) == "rfd:rsv:7:agent_chat:req-1"
|
|
assert quota.refund_idempotency_key(reserve_key) == quota.refund_idempotency_key(reserve_key)
|
|
|
|
|
|
def test_ledger_can_recompute_the_balance() -> None:
|
|
"""流水是真相，余额是缓存 (架构文档 §6.1 原则 4)."""
|
|
state = quota.QuotaState(used=0, base_limit=20)
|
|
deltas = []
|
|
for _ in range(3):
|
|
state = quota.reserve(state, 1)
|
|
deltas.append(-1)
|
|
state = quota.refund(state, 1)
|
|
deltas.append(1)
|
|
assert state.used == -sum(deltas)
|