diff --git a/tests/e2e/quota_management/budgets/test_team_model_max_budget_e2e.py b/tests/e2e/quota_management/budgets/test_team_model_max_budget_e2e.py index 21514fee753..d6b8bd31e60 100644 --- a/tests/e2e/quota_management/budgets/test_team_model_max_budget_e2e.py +++ b/tests/e2e/quota_management/budgets/test_team_model_max_budget_e2e.py @@ -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)