test_call_with_key_over_model_budget

BerriAI · Dec 19, 2024 · 1b6a1e1 · 1b6a1e1
1 parent 664aa53
commit 1b6a1e1
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py
@@ -1694,14 +1694,15 @@ async def test():
             await litellm.proxy.proxy_server.prisma_client.connect()
 
             # set budget for chatgpt-v-2 to 0.000001, expect the next request to fail
+            model_max_budget = {
+                "chatgpt-v-2": {
+                    "budget_limit": "0.000001",
+                    "time_period": "1d",
+                },
+            }
             request = GenerateKeyRequest(
                 max_budget=1000,
-                model_max_budget={
-                    "chatgpt-v-2": {
-                        "budget_limit": "0.000001",
-                        "time_period": "1d",
-                    },
-                },
+                model_max_budget=model_max_budget,
                 metadata={"user_api_key": 0.0001},
             )
             key = await generate_key_fn(request)
@@ -1759,6 +1760,7 @@ async def return_body():
                         "metadata": {
                             "user_api_key": hash_token(generated_key),
                             "user_api_key_user_id": user_id,
+                            "user_api_key_model_max_budget": model_max_budget,
                         }
                     },
                     "response_cost": 0.00002,