test_batch_update_spend

Author: ishaan-jaff

tests/local_testing/test_update_spend.py

@@ -62,6 +62,8 @@
     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()