(test) /key/gen fails when crossing budget

litellm/tests/test_key_generate_dynamodb.py

@@ -27,7 +27,7 @@ sys.path.insert(
 )  # Adds the parent directory to the system path
 import pytest, logging, asyncio
 import litellm, asyncio
-from litellm.proxy.proxy_server import new_user, user_api_key_auth
+from litellm.proxy.proxy_server import new_user, user_api_key_auth, user_update
 
 from litellm.proxy._types import NewUserRequest, DynamoDBArgs
 from litellm.proxy.utils import DBClient
@@ -92,7 +92,7 @@ def test_call_with_invalid_key():
 
             # use generated key to auth in
             result = await user_api_key_auth(request=request, api_key=bearer_token)
-            pytest.fail(f"This should have failed!")
+            pytest.fail(f"This should have failed!. IT's an invalid key")
 
         asyncio.run(test())
     except Exception as e:
@@ -129,11 +129,68 @@ def test_call_with_invalid_key():
 #     assert result is True
 
 
-# def test_call_with_key_under_budget():
-#     # 7. Make a call with a key under budget, expect to pass
-#     key = new_user(KeyUnderBudgetRequest())
-#     result = user_auth(ValidRequest(key))
-#     assert result is True
+def test_call_with_key_under_budget():
+    # 7. Make a call with a key under budget, expect to pass
+    setattr(litellm.proxy.proxy_server, "custom_db_client", custom_db_client)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    try:
+
+        async def test():
+            request = NewUserRequest(max_budget=0.00001)
+            key = await new_user(request)
+            print(key)
+
+            generated_key = key.key
+            user_id = key.user_id
+            bearer_token = "Bearer " + generated_key
+
+            request = Request(scope={"type": "http"})
+            request._url = URL(url="/chat/completions")
+
+            # use generated key to auth in
+            result = await user_api_key_auth(request=request, api_key=bearer_token)
+            print("result from user auth with new key", result)
+
+            # update spend using track_cost callback, make 2nd request, it should fail
+            from litellm.proxy.proxy_server import track_cost_callback
+            from litellm import ModelResponse, Choices, Message, Usage
+
+            resp = ModelResponse(
+                id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac",
+                choices=[
+                    Choices(
+                        finish_reason=None,
+                        index=0,
+                        message=Message(
+                            content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a",
+                            role="assistant",
+                        ),
+                    )
+                ],
+                model="gpt-35-turbo",  # azure always has model written like this
+                usage=Usage(prompt_tokens=21, completion_tokens=17, total_tokens=38),
+            )
+            await track_cost_callback(
+                kwargs={
+                    "stream": False,
+                    "litellm_params": {
+                        "metadata": {
+                            "user_api_key": generated_key,
+                            "user_api_key_user_id": user_id,
+                        }
+                    },
+                },
+                completion_response=resp,
+            )
+
+            # use generated key to auth in
+            result = await user_api_key_auth(request=request, api_key=bearer_token)
+            print("result from user auth with new key", result)
+
+        asyncio.run(test())
+    except Exception as e:
+        assert "ExceededBudget:" in str(e)
+        print("Got Exception", e)
 
 
 # def test_call_with_key_over_budget():