diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 06b1194b1..4e6b88b1c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -384,7 +384,9 @@ class PrismaClient: ) if response is not None: # for prisma we need to cast the expires time to str - if isinstance(response.expires, datetime): + if response.expires is not None and isinstance( + response.expires, datetime + ): response.expires = response.expires.isoformat() elif query_type == "find_all" and user_id is not None: response = await self.db.litellm_verificationtoken.find_many( diff --git a/tests/test_keys.py b/tests/test_keys.py index fe8c6c2f3..e9ec39256 100644 --- a/tests/test_keys.py +++ b/tests/test_keys.py @@ -20,7 +20,7 @@ async def generate_key(session, i, budget=None, budget_duration=None): "models": ["azure-models", "gpt-4"], "aliases": {"mistral-7b": "gpt-3.5-turbo"}, "duration": None, - "budget": budget, + "max_budget": budget, "budget_duration": budget_duration, } @@ -303,14 +303,18 @@ async def test_key_with_budgets(): - wait 5s - Check if value updated """ + from litellm.proxy.utils import hash_token + async with aiohttp.ClientSession() as session: key_gen = await generate_key( session=session, i=0, budget=10, budget_duration="5s" ) key = key_gen["key"] + hashed_token = hash_token(token=key) + 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(5) + await asyncio.sleep(15) key_info = await get_key_info(session=session, get_key=key, call_key=key) reset_at_new_value = key_info["info"]["budget_reset_at"] assert reset_at_init_value != reset_at_new_value diff --git a/tests/test_users.py b/tests/test_users.py index 81b6166dc..e1d1e45e3 100644 --- a/tests/test_users.py +++ b/tests/test_users.py @@ -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