From 8b5e397abe643ebdfc0a3149d323e17ffa985192 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 26 Jan 2024 14:32:34 -0800 Subject: [PATCH] (fix) dynamo - use hashed tokens budget tracking --- litellm/tests/test_key_generate_dynamodb.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py index be55595fa..5729c0e87 100644 --- a/litellm/tests/test_key_generate_dynamodb.py +++ b/litellm/tests/test_key_generate_dynamodb.py @@ -33,7 +33,7 @@ from litellm.proxy.proxy_server import ( ) from litellm.proxy._types import NewUserRequest, DynamoDBArgs, GenerateKeyRequest -from litellm.proxy.utils import DBClient +from litellm.proxy.utils import DBClient, hash_token from starlette.datastructures import URL @@ -232,7 +232,7 @@ def test_call_with_user_over_budget(custom_db_client): "stream": False, "litellm_params": { "metadata": { - "user_api_key": generated_key, + "user_api_key": hash_token(generated_key), "user_api_key_user_id": user_id, } }, @@ -305,7 +305,7 @@ def test_call_with_user_over_budget_stream(custom_db_client): "complete_streaming_response": resp, "litellm_params": { "metadata": { - "user_api_key": generated_key, + "user_api_key": hash_token(generated_key), "user_api_key_user_id": user_id, } }, @@ -376,7 +376,7 @@ def test_call_with_user_key_budget(custom_db_client): "stream": False, "litellm_params": { "metadata": { - "user_api_key": generated_key, + "user_api_key": hash_token(generated_key), "user_api_key_user_id": user_id, } }, @@ -449,7 +449,7 @@ def test_call_with_key_over_budget_stream(custom_db_client): "complete_streaming_response": resp, "litellm_params": { "metadata": { - "user_api_key": generated_key, + "user_api_key": hash_token(generated_key), "user_api_key_user_id": user_id, } },