mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
refactor(e2e): group multi-window budget tests into spec classes
Pure mechanical wrap: the tests move into TestKeyMultiWindowBudget and TestTeamMultiWindowBudget per the suite's lay-the-pattern-down-in-a-class convention; helpers stay module-level and bodies are unchanged besides indentation and the two signatures that would cross 120 chars
This commit is contained in:
ryan-crabbe-berri
parent
734e633c64
commit
2ad30c04c7
2 changed files with 150 additions and 144 deletions
|
|
@ -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]}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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]}"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue