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.
85 lines
2.9 KiB
Python
Executable File
85 lines
2.9 KiB
Python
Executable File
"""Global daily cost circuit breaker (架构文档 §13 三道成本闸门).
|
|
|
|
The requirement is enforcement, not alerting: once the platform-wide daily
|
|
spend crosses the threshold, non-critical AI calls are rejected.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
from decimal import Decimal
|
|
|
|
import pytest
|
|
|
|
from packages.llm.budget import DailyBudgetGuard, InMemoryCostCounter
|
|
from packages.llm.errors import DailyBudgetExceededError
|
|
|
|
|
|
def make_guard(
|
|
budget: str = "100", *, enabled: bool = True, now: datetime | None = None
|
|
) -> tuple[DailyBudgetGuard, InMemoryCostCounter]:
|
|
counter = InMemoryCostCounter()
|
|
guard = DailyBudgetGuard(
|
|
counter,
|
|
daily_budget_cents=Decimal(budget),
|
|
enabled=enabled,
|
|
clock=(lambda: now) if now else None,
|
|
)
|
|
return guard, counter
|
|
|
|
|
|
async def test_allows_calls_below_the_threshold() -> None:
|
|
guard, _ = make_guard("100")
|
|
await guard.add_cost(Decimal("99.9"))
|
|
await guard.ensure_within_budget(scene="agent_chat")
|
|
|
|
|
|
async def test_rejects_non_critical_scenes_once_the_budget_is_reached() -> None:
|
|
guard, _ = make_guard("100")
|
|
await guard.add_cost(Decimal("100"))
|
|
with pytest.raises(DailyBudgetExceededError) as excinfo:
|
|
await guard.ensure_within_budget(scene="agent_chat")
|
|
assert "agent_chat" in str(excinfo.value)
|
|
|
|
|
|
async def test_critical_scenes_bypass_the_breaker() -> None:
|
|
guard, _ = make_guard("100")
|
|
await guard.add_cost(Decimal("500"))
|
|
# Degrade the product, do not break it: 意图识别/越界拒答 must keep working.
|
|
await guard.ensure_within_budget(scene="intent_classify", critical=True)
|
|
|
|
|
|
async def test_disabled_guard_never_blocks_and_never_counts() -> None:
|
|
guard, counter = make_guard("1", enabled=False)
|
|
await guard.add_cost(Decimal("999"))
|
|
await guard.ensure_within_budget(scene="agent_chat")
|
|
assert await counter.get(guard.current_key()) == Decimal("0")
|
|
|
|
|
|
async def test_cost_accumulates_across_calls() -> None:
|
|
guard, _ = make_guard("100")
|
|
await guard.add_cost(Decimal("1.5"))
|
|
await guard.add_cost(Decimal("2.25"))
|
|
assert await guard.spent_today() == Decimal("3.75")
|
|
|
|
|
|
async def test_zero_cost_is_not_counted() -> None:
|
|
guard, _ = make_guard("100")
|
|
await guard.add_cost(Decimal("0"))
|
|
assert await guard.spent_today() == Decimal("0")
|
|
|
|
|
|
async def test_budget_day_boundary_follows_utc_plus_8() -> None:
|
|
"""FR-BILL-01: 跨时区按 UTC+8 统一。17:00Z already belongs to the next day."""
|
|
guard, _ = make_guard(now=datetime(2026, 8, 13, 17, 0, tzinfo=UTC))
|
|
assert guard.current_key().endswith("2026-08-14")
|
|
|
|
guard_before, _ = make_guard(now=datetime(2026, 8, 13, 15, 0, tzinfo=UTC))
|
|
assert guard_before.current_key().endswith("2026-08-13")
|
|
|
|
|
|
async def test_budget_properties_are_exposed_for_dashboards() -> None:
|
|
guard, _ = make_guard("42")
|
|
assert guard.enabled is True
|
|
assert guard.budget_cents == Decimal("42")
|