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

test(test_users.py): test budgets with resets

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-24 15:30:30 -08:00
parent e471157d57
commit bb7705b494
3 changed files with 39 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

32
tests/test_users.py
View file

@ -6,13 +6,15 @@ import aiohttp
import time
async def new_user(session, i, user_id=None):
async def new_user(session, i, user_id=None, budget=None, budget_duration=None):
url = "http://0.0.0.0:4000/user/new"
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
data = {
"models": ["azure-models"],
"aliases": {"mistral-7b": "gpt-3.5-turbo"},
"duration": None,
"max_budget": budget,
"budget_duration": budget_duration,
}
if user_id is not None:
@ -46,7 +48,7 @@ async def get_user_info(session, get_user, call_user):
"""
Make sure only models user has access to are returned
"""
url = f"http://0.0.0.0:4000/user/info?key={get_user}"
url = f"http://0.0.0.0:4000/user/info?user_id={get_user}"
headers = {
"Authorization": f"Bearer {call_user}",
"Content-Type": "application/json",
@ -100,3 +102,29 @@ async def test_user_update():
Make chat completion call
"""
pass
@pytest.mark.asyncio
async def test_users_with_budgets():
"""
- Create key with budget and 5s duration
- Get 'reset_at' value
- wait 5s
- Check if value updated
"""
get_user = f"krrish_{time.time()}@berri.ai"
async with aiohttp.ClientSession() as session:
key_gen = await new_user(
session, 0, user_id=get_user, budget=10, budget_duration="5s"
)
key = key_gen["key"]
user_info = await get_user_info(
session=session, get_user=get_user, call_user=key
)
reset_at_init_value = user_info["user_info"]["budget_reset_at"]
await asyncio.sleep(15)
user_info = await get_user_info(
session=session, get_user=get_user, call_user=key
)
reset_at_new_value = user_info["user_info"]["budget_reset_at"]
assert reset_at_init_value != reset_at_new_value