test(e2e): pin team model block shape and prove override non-accrual

Adversarial-review hardening: blocks now assert 429 plus the per-model phrase
so a team-wide budget block cannot satisfy them, and the override test is
restructured to prove both precedence halves deterministically; the override
key serves past the team cap, its spend leaves the shared counter untouched
for a plain key, and the plain key's own use then exhausts the cap.
This commit is contained in:
ryan-crabbe-berri 2026-08-04 19:38:15 -07:00
parent 5ca606aa4e
commit 93d6c65f46

View file

@ -15,6 +15,7 @@ pytestmark = pytest.mark.e2e
CAPPED_MODEL = "claude-haiku-4-5"
FREE_MODEL = "gemini-2.5-flash"
BLOCK_DEADLINE_SECONDS = 60
SPEND_FLUSH_SECONDS = 2
def _call(client: BudgetClient, key: str, model: str) -> StreamingResponse:
@ -24,12 +25,18 @@ def _call(client: BudgetClient, key: str, model: str) -> StreamingResponse:
return result
def _assert_team_model_block(result: StreamingResponse, team_id: str) -> None:
assert result.status_code == 429, f"HTTP {result.status_code} {result.body}"
assert team_id in result.body, result.body
assert f"exceeded budget for model={CAPPED_MODEL}" in result.body, result.body
def _drive_to_team_block(client: BudgetClient, key: str, team_id: str) -> None:
deadline = time.monotonic() + BLOCK_DEADLINE_SECONDS
while time.monotonic() < deadline:
result = _call(client, key, CAPPED_MODEL)
if is_budget_block(result):
assert team_id in result.body, result.body
_assert_team_model_block(result, team_id)
return
time.sleep(1)
pytest.fail(f"team model_max_budget on {CAPPED_MODEL} never enforced")
@ -62,16 +69,19 @@ def test_team_model_cap_blocks_sibling_key(client: BudgetClient, resources: Reso
f"sibling team key served {CAPPED_MODEL} after the team cap was exhausted: "
f"HTTP {first.status_code} {first.body}"
)
assert team_id in first.body, first.body
_assert_team_model_block(first, team_id)
require_successful_call(_call(client, bystander, FREE_MODEL))
other = _call(client, bystander, FREE_MODEL)
assert not is_budget_block(other), f"{FREE_MODEL} was blocked by {CAPPED_MODEL}'s team cap: {other.body}"
require_successful_call(other)
# User flow (TLDR^2)
# 1. Admin caps claude-haiku tiny on the team
# 2. Admin issues Carol a team key with its own big claude budget
# 3. Dave's plain team key exhausts the team cap and is refused
# 4. Carol keeps getting claude answers on the same team
# 3. Carol keeps getting claude answers past the team cap
# 4. Dave's plain key still starts fresh: Carol spent none of the team cap
# 5. Dave's own claude use then exhausts the team cap and is refused
@pytest.mark.covers("quota_management.budget.team_model_max.key_override_wins")
def test_key_override_exempts_from_team_cap(client: BudgetClient, resources: ResourceManager) -> None:
team_id = client.create_team(
@ -79,11 +89,20 @@ def test_key_override_exempts_from_team_cap(client: BudgetClient, resources: Res
model_max_budget=model_budget(CAPPED_MODEL, 1e-6),
)
resources.defer(lambda: client.delete_team(team_id))
inherit = client.generate_key(team_id=team_id)
resources.defer(lambda: client.delete_key(inherit))
override = client.generate_key(team_id=team_id, model_max_budget=model_budget(CAPPED_MODEL, 1000.0))
resources.defer(lambda: client.delete_key(override))
_drive_to_team_block(client, inherit, team_id)
plain = client.generate_key(team_id=team_id)
resources.defer(lambda: client.delete_key(plain))
require_successful_call(_call(client, override, CAPPED_MODEL))
time.sleep(SPEND_FLUSH_SECONDS)
second = _call(client, override, CAPPED_MODEL)
assert not is_budget_block(second), f"override key hit the team cap it should be exempt from: {second.body}"
require_successful_call(second)
time.sleep(SPEND_FLUSH_SECONDS)
fresh = _call(client, plain, CAPPED_MODEL)
assert not is_budget_block(fresh), f"override key's spend drove the shared team counter: {fresh.body}"
require_successful_call(fresh)
_drive_to_team_block(client, plain, team_id)