test(test_key_generate_prisma.py): fix check for raising error if key generate max budget > key upperbound

litellm/proxy/_types.py

@@ -44,6 +44,7 @@ class LiteLLM_UpperboundKeyGenerateParams(LiteLLMBase):
     """
 
     max_budget: Optional[float] = None
+    budget_duration: Optional[str] = None
     max_parallel_requests: Optional[int] = None
     tpm_limit: Optional[int] = None
     rpm_limit: Optional[int] = None

litellm/tests/test_key_generate_prisma.py

@@ -66,6 +66,7 @@ from litellm.proxy._types import (
     GenerateKeyRequest,
     NewTeamRequest,
     UserAPIKeyAuth,
+    LiteLLM_UpperboundKeyGenerateParams,
 )
 from litellm.proxy.utils import DBClient
 from starlette.datastructures import URL
@@ -1627,10 +1628,9 @@ async def test_upperbound_key_params(prisma_client):
     """
     setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
     setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
-    litellm.upperbound_key_generate_params = {
-        "max_budget": 0.001,
-        "budget_duration": "1m",
-    }
+    litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(
+        max_budget=0.001, budget_duration="1m"
+    )
     await litellm.proxy.proxy_server.prisma_client.connect()
     try:
         request = GenerateKeyRequest(
@@ -1638,18 +1638,9 @@ async def test_upperbound_key_params(prisma_client):
             budget_duration="30d",
         )
         key = await generate_key_fn(request)
-        generated_key = key.key
-
-        result = await info_key_fn(key=generated_key)
-        key_info = result["info"]
-        # assert it used the upper bound for max_budget, and budget_duration
-        assert key_info["max_budget"] == 0.001
-        assert key_info["budget_duration"] == "1m"
-
-        print(result)
+        # print(result)
     except Exception as e:
-        print("Got Exception", e)
-        pytest.fail(f"Got exception {e}")
+        assert e.code == 400
 
 
 def test_get_bearer_token():