test(test_users.py): test budgets with resets

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(

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

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