From 7a2442d6c052dd1f82a56d70087f8400ea58d2e2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 07:12:29 -0700 Subject: [PATCH] test_batch_update_spend --- tests/local_testing/test_update_spend.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/local_testing/test_update_spend.py b/tests/local_testing/test_update_spend.py index fffa3062d7..cc2c94af27 100644 --- a/tests/local_testing/test_update_spend.py +++ b/tests/local_testing/test_update_spend.py @@ -62,6 +62,8 @@ from litellm.proxy._types import ( KeyRequest, NewUserRequest, UpdateKeyRequest, + SpendUpdateQueueItem, + Litellm_EntityType, ) proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) @@ -93,7 +95,13 @@ def prisma_client(): @pytest.mark.asyncio async def test_batch_update_spend(prisma_client): - prisma_client.user_list_transactions["test-litellm-user-5"] = 23 + await proxy_logging_obj.db_spend_update_writer.spend_update_queue.add_update( + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.USER, + entity_id="test-litellm-user-5", + response_cost=23, + ) + ) setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") await litellm.proxy.proxy_server.prisma_client.connect()