From c88078befbfa21dbfb811fd85659889b8f2538bf Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 2 Feb 2024 10:23:31 -0800 Subject: [PATCH] test: fix import --- litellm/tests/test_key_generate_dynamodb.py | 4 +++- litellm/tests/test_key_generate_prisma.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_key_generate_dynamodb.py b/litellm/tests/test_key_generate_dynamodb.py index 418098333..bacb886d8 100644 --- a/litellm/tests/test_key_generate_dynamodb.py +++ b/litellm/tests/test_key_generate_dynamodb.py @@ -213,7 +213,9 @@ def test_call_with_user_over_budget(custom_db_client): print("result from user auth with new key", result) # update spend using track_cost callback, make 2nd request, it should fail - from litellm.proxy.proxy_server import track_cost_callback + from litellm.proxy.proxy_server import ( + _PROXY_track_cost_callback as track_cost_callback, + ) from litellm import ModelResponse, Choices, Message, Usage resp = ModelResponse( diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index d8ffcf022..5e0cefd2d 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -242,7 +242,9 @@ def test_call_with_user_over_budget(prisma_client): print("result from user auth with new key", result) # update spend using track_cost callback, make 2nd request, it should fail - from litellm.proxy.proxy_server import track_cost_callback + from litellm.proxy.proxy_server import ( + _PROXY_track_cost_callback as track_cost_callback, + ) from litellm import ModelResponse, Choices, Message, Usage resp = ModelResponse(