diff --git a/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py b/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py index 9c7a65444bc..80f2b6bd67d 100644 --- a/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py +++ b/tests/e2e/quota_management/budgets/test_multi_window_budget_e2e.py @@ -93,14 +93,6 @@ def _assert_short_window_blocks_then_resets(client: BudgetClient, key: str) -> N pytest.fail(f"{WINDOW_SECONDS}s window never reset within 150s") -@pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") -def test_short_window_blocks_then_resets(client: BudgetClient, resources: ResourceManager) -> None: - key = client.generate_key(models=[MODEL], budget_limits=_short_roomy_limits()) - resources.defer(lambda: client.delete_key(key)) - - _assert_short_window_blocks_then_resets(client, key) - - def _mint_key_of_kind(client: BudgetClient, resources: ResourceManager, kind: str) -> str: """A key carrying the tight+roomy budget_limits pair, minted to roomy (100.0) surroundings so only the key's own windows can block.""" @@ -128,68 +120,79 @@ def _mint_key_of_kind(client: BudgetClient, resources: ResourceManager, kind: st return key -@pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") -@pytest.mark.parametrize("kind", ["personal", "team", "team_member"]) -def test_short_window_blocks_then_resets_across_key_kinds( - client: BudgetClient, resources: ResourceManager, kind: str -) -> None: - key = _mint_key_of_kind(client, resources, kind) +class TestKeyMultiWindowBudget: + @pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") + def test_short_window_blocks_then_resets(self, client: BudgetClient, resources: ResourceManager) -> None: + key = client.generate_key(models=[MODEL], budget_limits=_short_roomy_limits()) + resources.defer(lambda: client.delete_key(key)) - _assert_short_window_blocks_then_resets(client, key) + _assert_short_window_blocks_then_resets(client, key) -@pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") -def test_long_window_blocks_after_short_window_resets(client: BudgetClient, resources: ResourceManager) -> None: - key = client.generate_key( - models=[MODEL], - budget_limits=[ - BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP), - BudgetWindow(budget_duration=LONG_WINDOW, max_budget=LONG_CAP), - ], - ) - resources.defer(lambda: client.delete_key(key)) + @pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") + @pytest.mark.parametrize("kind", ["personal", "team", "team_member"]) + def test_short_window_blocks_then_resets_across_key_kinds( + self, client: BudgetClient, resources: ResourceManager, kind: str + ) -> None: + key = _mint_key_of_kind(client, resources, kind) - # 1. drive the key to get blocked by SHORT_WINDOW, assert it's budget error - blocked = _drive_to_block(client, key) - assert blocked.status_code == 429, f"budget block was not a 429: {blocked.status_code} {blocked.body[:200]}" + _assert_short_window_blocks_then_resets(client, key) - # 2. check the reset times of both budget windows after we drove to being blocked - blocked_reset_at = window_reset_at(client.key_budget_windows(key), SHORT_WINDOW) - assert blocked_reset_at is not None, "short window missing from /key/info budget_limits" - blocked_long_reset_at = window_reset_at(client.key_budget_windows(key), LONG_WINDOW) - assert blocked_long_reset_at is not None, "long window missing from /key/info budget_limits" - # 3. poll every 5s for the SHORT_WINDOW reset time until it is past it, fails if it doesnt reset - deadline = time.monotonic() + RESET_DEADLINE_SECONDS - while time.monotonic() < deadline: - time.sleep(5) - current = window_reset_at(client.key_budget_windows(key), SHORT_WINDOW) - if current is not None and current > blocked_reset_at: - break - else: - pytest.fail( - f"{SHORT_WINDOW} window's reset_at never advanced past {blocked_reset_at} within {RESET_DEADLINE_SECONDS}s" + @pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets") + def test_long_window_blocks_after_short_window_resets( + self, client: BudgetClient, resources: ResourceManager + ) -> None: + key = client.generate_key( + models=[MODEL], + budget_limits=[ + BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP), + BudgetWindow(budget_duration=LONG_WINDOW, max_budget=LONG_CAP), + ], ) + resources.defer(lambda: client.delete_key(key)) - # 4. short window just reset in 3, so now make a call, check that its blocked (should be blocked by LONG_WINDOW because short window reset), also make sure its budget error - deadline = time.monotonic() + RESET_DEADLINE_SECONDS - last_body = "" - while time.monotonic() < deadline: - result = _call(client, key) - if result.ok: - rolled = window_reset_at(client.key_budget_windows(key), LONG_WINDOW) != blocked_long_reset_at + # 1. drive the key to get blocked by SHORT_WINDOW, assert it's budget error + blocked = _drive_to_block(client, key) + assert blocked.status_code == 429, f"budget block was not a 429: {blocked.status_code} {blocked.body[:200]}" + + # 2. check the reset times of both budget windows after we drove to being blocked + blocked_reset_at = window_reset_at(client.key_budget_windows(key), SHORT_WINDOW) + assert blocked_reset_at is not None, "short window missing from /key/info budget_limits" + blocked_long_reset_at = window_reset_at(client.key_budget_windows(key), LONG_WINDOW) + assert blocked_long_reset_at is not None, "long window missing from /key/info budget_limits" + + # 3. poll every 5s for the SHORT_WINDOW reset time until it is past it, fails if it doesnt reset + deadline = time.monotonic() + RESET_DEADLINE_SECONDS + while time.monotonic() < deadline: + time.sleep(5) + current = window_reset_at(client.key_budget_windows(key), SHORT_WINDOW) + if current is not None and current > blocked_reset_at: + break + else: pytest.fail( - f"{LONG_WINDOW} window failed to block after the {SHORT_WINDOW} window reset" - + (f" (the {LONG_WINDOW} window itself rolled mid-test - boundary crossed; rerun)" if rolled else "") + f"{SHORT_WINDOW} window's reset_at never advanced past {blocked_reset_at} within {RESET_DEADLINE_SECONDS}s" ) - assert is_budget_block(result), ( - f"non-budget error while waiting for {LONG_WINDOW} attribution: " - f"status={result.status_code} body={result.body[:200]}" + + # 4. short window just reset in 3, so now make a call, check that its blocked (should be blocked by LONG_WINDOW because short window reset), also make sure its budget error + deadline = time.monotonic() + RESET_DEADLINE_SECONDS + last_body = "" + while time.monotonic() < deadline: + result = _call(client, key) + if result.ok: + rolled = window_reset_at(client.key_budget_windows(key), LONG_WINDOW) != blocked_long_reset_at + pytest.fail( + f"{LONG_WINDOW} window failed to block after the {SHORT_WINDOW} window reset" + + (f" (the {LONG_WINDOW} window itself rolled mid-test - boundary crossed; rerun)" if rolled else "") + ) + assert is_budget_block(result), ( + f"non-budget error while waiting for {LONG_WINDOW} attribution: " + f"status={result.status_code} body={result.body[:200]}" + ) + if f"over {LONG_WINDOW} budget" in result.body: + return + last_body = result.body + time.sleep(5) + pytest.fail( + f"block never attributed to the {LONG_WINDOW} window within {RESET_DEADLINE_SECONDS}s: {last_body[:200]}" ) - if f"over {LONG_WINDOW} budget" in result.body: - return - last_body = result.body - time.sleep(5) - pytest.fail( - f"block never attributed to the {LONG_WINDOW} window within {RESET_DEADLINE_SECONDS}s: {last_body[:200]}" - ) diff --git a/tests/e2e/quota_management/budgets/test_team_multi_window_budget_e2e.py b/tests/e2e/quota_management/budgets/test_team_multi_window_budget_e2e.py index 1db68e6afe9..4334c036b3a 100644 --- a/tests/e2e/quota_management/budgets/test_team_multi_window_budget_e2e.py +++ b/tests/e2e/quota_management/budgets/test_team_multi_window_budget_e2e.py @@ -51,96 +51,99 @@ def _drive_to_block(client: BudgetClient, key: str) -> StreamingResponse: pytest.fail("team budget never enforced before block") -@pytest.mark.covers("quota_management.budget.team_multi_window.blocks_then_resets") -def test_team_short_window_blocks_then_resets(client: BudgetClient, resources: ResourceManager) -> None: - team_id = client.create_team( - alias=f"e2e-team-window-{unique_marker()}", - budget_limits=[ - BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=3e-6), - BudgetWindow(budget_duration="1m", max_budget=1.0), # roomy: never blocks - ], - ) - resources.defer(lambda: client.delete_team(team_id)) - key = client.generate_key(team_id=team_id, models=["claude-haiku-4-5"]) - resources.defer(lambda: client.delete_key(key)) +class TestTeamMultiWindowBudget: + @pytest.mark.covers("quota_management.budget.team_multi_window.blocks_then_resets") + def test_team_short_window_blocks_then_resets(self, client: BudgetClient, resources: ResourceManager) -> None: + team_id = client.create_team( + alias=f"e2e-team-window-{unique_marker()}", + budget_limits=[ + BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=3e-6), + BudgetWindow(budget_duration="1m", max_budget=1.0), # roomy: never blocks + ], + ) + resources.defer(lambda: client.delete_team(team_id)) + key = client.generate_key(team_id=team_id, models=["claude-haiku-4-5"]) + resources.defer(lambda: client.delete_key(key)) - # 1. exhaust the tight window -> litellm returns budget_exceeded - start = time.monotonic() - _drive_to_block(client, key) + # 1. exhaust the tight window -> litellm returns budget_exceeded + start = time.monotonic() + _drive_to_block(client, key) - # 2. the window resets at the next wall-clock-aligned boundary (up to a window - # after start), then the reset job (~15-20s rescheduler) zeroes the spend. - # Allow generous headroom for that alignment + rescheduler latency; a stuck - # rescheduler is caught by the wait-loop timeout, not this elapsed bound. - deadline = time.monotonic() + 150 - while time.monotonic() < deadline: - time.sleep(5) - result = _call(client, key) - if result.ok: - elapsed = time.monotonic() - start - assert elapsed < WINDOW_SECONDS + 90, f"reset took {elapsed:.0f}s - too long for a {WINDOW_SECONDS}s window" - return - assert is_budget_block(result), f"non-budget error during reset wait: {result.body[:200]}" - pytest.fail(f"team {WINDOW_SECONDS}s window never reset within 150s") + # 2. the window resets at the next wall-clock-aligned boundary (up to a window + # after start), then the reset job (~15-20s rescheduler) zeroes the spend. + # Allow generous headroom for that alignment + rescheduler latency; a stuck + # rescheduler is caught by the wait-loop timeout, not this elapsed bound. + deadline = time.monotonic() + 150 + while time.monotonic() < deadline: + time.sleep(5) + result = _call(client, key) + if result.ok: + elapsed = time.monotonic() - start + assert elapsed < WINDOW_SECONDS + 90, f"reset took {elapsed:.0f}s - too long for a {WINDOW_SECONDS}s window" + return + assert is_budget_block(result), f"non-budget error during reset wait: {result.body[:200]}" + pytest.fail(f"team {WINDOW_SECONDS}s window never reset within 150s") -@pytest.mark.covers("quota_management.budget.team_multi_window.blocks_then_resets") -def test_team_long_window_blocks_after_short_window_resets(client: BudgetClient, resources: ResourceManager) -> None: + @pytest.mark.covers("quota_management.budget.team_multi_window.blocks_then_resets") + def test_team_long_window_blocks_after_short_window_resets( + self, client: BudgetClient, resources: ResourceManager + ) -> None: - # 0. key with a short budget window and a long budget window - team_id = client.create_team( - alias=f"e2e-team-long-window-{unique_marker()}", - budget_limits=[ - BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP), - BudgetWindow(budget_duration=LONG_WINDOW, max_budget=LONG_CAP), - ], - ) - resources.defer(lambda: client.delete_team(team_id)) - key = client.generate_key(team_id=team_id, models=["claude-haiku-4-5"]) - resources.defer(lambda: client.delete_key(key)) + # 0. key with a short budget window and a long budget window + team_id = client.create_team( + alias=f"e2e-team-long-window-{unique_marker()}", + budget_limits=[ + BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP), + BudgetWindow(budget_duration=LONG_WINDOW, max_budget=LONG_CAP), + ], + ) + resources.defer(lambda: client.delete_team(team_id)) + key = client.generate_key(team_id=team_id, models=["claude-haiku-4-5"]) + resources.defer(lambda: client.delete_key(key)) - # 1. drive the key to being blocked, assert its blocked by budget budget_exceeded - blocked = _drive_to_block(client, key) - assert blocked.status_code == 429, f"budget block was not a 429: {blocked.status_code} {blocked.body[:200]}" + # 1. drive the key to being blocked, assert its blocked by budget budget_exceeded + blocked = _drive_to_block(client, key) + assert blocked.status_code == 429, f"budget block was not a 429: {blocked.status_code} {blocked.body[:200]}" - # 2. check the the teams budget windows - blocked_reset_at = window_reset_at(client.team_budget_windows(team_id), SHORT_WINDOW) - assert blocked_reset_at is not None, "short window missing from /team/info budget_limits" - blocked_long_reset_at = window_reset_at(client.team_budget_windows(team_id), LONG_WINDOW) - assert blocked_long_reset_at is not None, "long window missing from /team/info budget_limits" + # 2. check the the teams budget windows + blocked_reset_at = window_reset_at(client.team_budget_windows(team_id), SHORT_WINDOW) + assert blocked_reset_at is not None, "short window missing from /team/info budget_limits" + blocked_long_reset_at = window_reset_at(client.team_budget_windows(team_id), LONG_WINDOW) + assert blocked_long_reset_at is not None, "long window missing from /team/info budget_limits" - # 3. keep checking that the short budget window reset, if it doesnt within the deadline then fail - deadline = time.monotonic() + RESET_DEADLINE_SECONDS - while time.monotonic() < deadline: - time.sleep(5) - current = window_reset_at(client.team_budget_windows(team_id), SHORT_WINDOW) - if current is not None and current > blocked_reset_at: - break - else: - pytest.fail( - f"team {SHORT_WINDOW} window's reset_at never advanced past " - f"{blocked_reset_at} within {RESET_DEADLINE_SECONDS}s" - ) - - # 4. short window just reset in 3, so now make another call, assert that the long window blocks the next call with budget_exceeded - deadline = time.monotonic() + RESET_DEADLINE_SECONDS - last_body = "" - while time.monotonic() < deadline: - result = _call(client, key) - if result.ok: - rolled = window_reset_at(client.team_budget_windows(team_id), LONG_WINDOW) != blocked_long_reset_at + # 3. keep checking that the short budget window reset, if it doesnt within the deadline then fail + deadline = time.monotonic() + RESET_DEADLINE_SECONDS + while time.monotonic() < deadline: + time.sleep(5) + current = window_reset_at(client.team_budget_windows(team_id), SHORT_WINDOW) + if current is not None and current > blocked_reset_at: + break + else: pytest.fail( - f"team {LONG_WINDOW} window failed to block after the {SHORT_WINDOW} window reset" - + (f" (the {LONG_WINDOW} window itself rolled mid-test - boundary crossed; rerun)" if rolled else "") + f"team {SHORT_WINDOW} window's reset_at never advanced past " + f"{blocked_reset_at} within {RESET_DEADLINE_SECONDS}s" ) - assert is_budget_block(result), ( - f"non-budget error while waiting for {LONG_WINDOW} attribution: " - f"status={result.status_code} body={result.body[:200]}" + + # 4. short window just reset in 3, so now make another call, assert that the long window blocks the next call with budget_exceeded + deadline = time.monotonic() + RESET_DEADLINE_SECONDS + last_body = "" + while time.monotonic() < deadline: + result = _call(client, key) + if result.ok: + rolled = window_reset_at(client.team_budget_windows(team_id), LONG_WINDOW) != blocked_long_reset_at + pytest.fail( + f"team {LONG_WINDOW} window failed to block after the {SHORT_WINDOW} window reset" + + (f" (the {LONG_WINDOW} window itself rolled mid-test - boundary crossed; rerun)" if rolled else "") + ) + assert is_budget_block(result), ( + f"non-budget error while waiting for {LONG_WINDOW} attribution: " + f"status={result.status_code} body={result.body[:200]}" + ) + if f"over {LONG_WINDOW} budget" in result.body: + return + last_body = result.body + time.sleep(5) + pytest.fail( + f"team block never attributed to the {LONG_WINDOW} window within {RESET_DEADLINE_SECONDS}s: {last_body[:200]}" ) - if f"over {LONG_WINDOW} budget" in result.body: - return - last_body = result.body - time.sleep(5) - pytest.fail( - f"team block never attributed to the {LONG_WINDOW} window within {RESET_DEADLINE_SECONDS}s: {last_body[:200]}" - )