phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

test(test_keys.py): add more retries for checking if budget updated

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-31 10:44:03 -08:00
parent cf5f6ab49c
commit eea4325e60
2 changed files with 15 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

5
litellm/proxy/proxy_server.py
View file

@ -1743,7 +1743,10 @@ async def startup_event():
### START BUDGET SCHEDULER ###
scheduler = AsyncIOScheduler()
scheduler.add_job(reset_budget, "interval", seconds=10, args=[prisma_client])
interval = random.randint(
7, 14
) # random interval, so multiple workers avoid resetting budget at the same time
scheduler.add_job(reset_budget, "interval", seconds=interval, args=[prisma_client])
scheduler.start()

14
tests/test_keys.py
View file

@ -355,9 +355,17 @@ async def test_key_with_budgets():
print(f"hashed_token: {hashed_token}")
key_info = await get_key_info(session=session, get_key=key, call_key=key)
reset_at_init_value = key_info["info"]["budget_reset_at"]
await asyncio.sleep(30)
key_info = await get_key_info(session=session, get_key=key, call_key=key)
reset_at_new_value = key_info["info"]["budget_reset_at"]
reset_at_new_value = None
i = 0
while i < 3:
await asyncio.sleep(30)
key_info = await get_key_info(session=session, get_key=key, call_key=key)
reset_at_new_value = key_info["info"]["budget_reset_at"]
try:
assert reset_at_init_value != reset_at_new_value
break
except:
i + 1
assert reset_at_init_value != reset_at_new_value