mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
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:
parent
5ca606aa4e
commit
93d6c65f46
1 changed files with 28 additions and 9 deletions
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue