diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index de26168591..b4c86afb25 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/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