From cbf35087c7e1d6cbdaeee8d26d5a93cca45cd046 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 29 Mar 2024 20:30:43 -0700 Subject: [PATCH] test(test_key_generate_prisma.py): fix test --- litellm/tests/test_key_generate_prisma.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 98126f9a16..8430aa7fd2 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -576,9 +576,17 @@ def test_call_with_proxy_over_budget_stream(prisma_client): litellm_proxy_budget_name = f"litellm-proxy-budget-{time.time()}" setattr( litellm.proxy.proxy_server, - "litellm_proxy_budget_name", + "litellm_proxy_admin_name", litellm_proxy_budget_name, ) + setattr(litellm, "max_budget", 0.00001) + from litellm.proxy.proxy_server import user_api_key_cache + + user_api_key_cache.set_cache( + key="{}:spend".format(litellm_proxy_budget_name), value=0 + ) + setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) + from litellm._logging import verbose_proxy_logger import logging @@ -589,10 +597,10 @@ def test_call_with_proxy_over_budget_stream(prisma_client): async def test(): await litellm.proxy.proxy_server.prisma_client.connect() ## CREATE PROXY + USER BUDGET ## - request = NewUserRequest( - max_budget=0.00001, user_id=litellm_proxy_budget_name - ) - await new_user(request) + # request = NewUserRequest( + # max_budget=0.00001, user_id=litellm_proxy_budget_name + # ) + # await new_user(request) request = NewUserRequest() key = await new_user(request) print(key)