test_batch_update_spend

This commit is contained in:
Ishaan Jaff
2025-04-01 07:12:29 -07:00
parent d0a7e44a6e
commit 7a2442d6c0
+9 -1
View File
@@ -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()