phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix test_call_with_end_user_over_budget

Browse source
This commit is contained in:
Ishaan Jaff 2024-06-20 19:54:34 -07:00
parent 76145a94fb
commit 1db25d5b6a
1 changed files with 24 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

32
litellm/tests/test_key_generate_prisma.py
View file

@ -72,6 +72,7 @@ from litellm.proxy.proxy_server import (
image_generation,
model_list,
moderations,
new_end_user,
user_api_key_auth,
)
from litellm.proxy.spend_reporting_endpoints.spend_management_endpoints import (
@ -91,6 +92,7 @@ from litellm.proxy._types import (
GenerateKeyRequest,
KeyRequest,
LiteLLM_UpperboundKeyGenerateParams,
NewCustomerRequest,
NewTeamRequest,
NewUserRequest,
UpdateKeyRequest,
@ -511,15 +513,22 @@ def test_call_with_end_user_over_budget(prisma_client):
async def test():
await litellm.proxy.proxy_server.prisma_client.connect()
request = GenerateKeyRequest() # create a key with no budget
key = await new_user(request)
print(key)
generated_key = key.key
bearer_token = "Bearer " + generated_key
user = f"ishaan {random.randint(0, 10000)}"
request = NewCustomerRequest(
user_id=user, max_budget=0.000001
) # create a key with no budget
await new_end_user(
request,
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-1234",
user_id="1234",
),
)
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
bearer_token = "Bearer sk-1234"
result = await user_api_key_auth(request=request, api_key=bearer_token)
@ -556,7 +565,7 @@ def test_call_with_end_user_over_budget(prisma_client):
"stream": False,
"litellm_params": {
"metadata": {
"user_api_key": generated_key,
"user_api_key": "sk-1234",
"user_api_key_user_id": user,
},
"proxy_server_request": {
@ -571,7 +580,14 @@ def test_call_with_end_user_over_budget(prisma_client):
start_time=datetime.now(),
end_time=datetime.now(),
)
await asyncio.sleep(5)
await asyncio.sleep(10)
await update_spend(
prisma_client=prisma_client,
db_writer_client=None,
proxy_logging_obj=proxy_logging_obj,
)
# use generated key to auth in
result = await user_api_key_auth(request=request, api_key=bearer_token)
print("result from user auth with new key", result)