(test) test_upperbound_key_params

litellm/tests/test_key_generate_prisma.py

@@ -1279,6 +1279,40 @@ async def test_default_key_params(prisma_client):
         pytest.fail(f"Got exception {e}")
 
 
+@pytest.mark.asyncio()
+async def test_upperbound_key_params(prisma_client):
+    """
+    - create key
+    - get key info
+    - assert key_name is not null
+    """
+    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",
+    }
+    await litellm.proxy.proxy_server.prisma_client.connect()
+    try:
+        request = GenerateKeyRequest(
+            max_budget=200000,
+            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)
+    except Exception as e:
+        print("Got Exception", e)
+        pytest.fail(f"Got exception {e}")
+
+
 def test_get_bearer_token():
     from litellm.proxy.proxy_server import _get_bearer_token