scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 133

test(e2e): cover key budget_duration resets on personal, team, and team-member keys (#33896)

Browse source
This commit is contained in:
ryan-crabbe-berri 2026-07-23 07:18:35 -07:00 • committed by GitHub
parent 3bba3633c7
commit 2bfd50ed37
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 41 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
tests/e2e/quota_management/budgets/test_budget_reset_e2e.py
View file

@ -12,6 +12,7 @@ from lifecycle import ResourceManager
pytestmark = pytest.mark.e2e
TINY_CAP = 3e-6
ROOMY_CAP = 100.0
WINDOW = "30s"
RESET_DEADLINE_SECONDS = 150
@ -46,7 +47,7 @@ def _poll_until_serves_again(client: BudgetClient, key: str) -> None:
pytest.fail(f"budget never reset within {RESET_DEADLINE_SECONDS}s")
class TestBudgetResetDiagonal:
class TestBudgetResetPerLevel:
@pytest.mark.covers("quota_management.budget.key.resets_after_window")
def test_bare_key_budget_resets_after_window(self, client: BudgetClient, resources: ResourceManager) -> None:
key = client.generate_key(max_budget=TINY_CAP, budget_duration=WINDOW)
@ -115,3 +116,42 @@ class TestBudgetResetDiagonal:
_drive_to_block(client, key)
_poll_until_serves_again(client, key)
class TestKeyBudgetResetAcrossKeyKinds:
"""The tiny max_budget and its 30s window sit on the key itself while the user,
team, and membership around it are roomy (100.0), so the key's own budget is
the only thing that can block and the only thing that has to reset."""
@pytest.mark.covers("quota_management.budget.key.resets_after_window")
def test_personal_key_resets_after_window(self, client: BudgetClient, resources: ResourceManager) -> None:
user_id = client.create_user(max_budget=ROOMY_CAP)
resources.defer(lambda: client.delete_user(user_id))
key = client.generate_key(user_id=user_id, max_budget=TINY_CAP, budget_duration=WINDOW)
resources.defer(lambda: client.delete_key(key))
_drive_to_block(client, key)
_poll_until_serves_again(client, key)
@pytest.mark.covers("quota_management.budget.key.resets_after_window")
def test_team_key_resets_after_window(self, client: BudgetClient, resources: ResourceManager) -> None:
team_id = client.create_team(alias=f"e2e-key-reset-team-{unique_marker()}", max_budget=ROOMY_CAP)
resources.defer(lambda: client.delete_team(team_id))
key = client.generate_key(team_id=team_id, max_budget=TINY_CAP, budget_duration=WINDOW)
resources.defer(lambda: client.delete_key(key))
_drive_to_block(client, key)
_poll_until_serves_again(client, key)
@pytest.mark.covers("quota_management.budget.key.resets_after_window")
def test_team_member_key_resets_after_window(self, client: BudgetClient, resources: ResourceManager) -> None:
team_id = client.create_team(alias=f"e2e-key-reset-team-{unique_marker()}", max_budget=ROOMY_CAP)
resources.defer(lambda: client.delete_team(team_id))
member_id = client.create_user(max_budget=ROOMY_CAP)
resources.defer(lambda: client.delete_user(member_id))
client.add_team_member(team_id, member_id, max_budget_in_team=ROOMY_CAP)
key = client.generate_key(team_id=team_id, user_id=member_id, max_budget=TINY_CAP, budget_duration=WINDOW)
resources.defer(lambda: client.delete_key(key))
_drive_to_block(client, key)
_poll_until_serves_again(client, key)