From e8d80509b1ce560e57d20898f1c2b2bb681c0d6d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 28 Mar 2024 13:44:40 -0700 Subject: [PATCH] test(test_update_spend.py): allow db_client to be none --- litellm/proxy/utils.py | 10 ++++++++-- litellm/tests/test_update_spend.py | 5 ++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1102766f9f8..5fe1fee3579 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1867,7 +1867,9 @@ async def reset_budget(prisma_client: PrismaClient): ) -async def update_spend(prisma_client: PrismaClient, db_writer_client: HTTPHandler): +async def update_spend( + prisma_client: PrismaClient, db_writer_client: Optional[HTTPHandler] +): """ Batch write updates to db. @@ -1996,7 +1998,11 @@ async def update_spend(prisma_client: PrismaClient, db_writer_client: HTTPHandle ### UPDATE SPEND LOGS ### base_url = os.getenv("SPEND_LOGS_URL", None) - if len(prisma_client.spend_log_transactons) > 0 and base_url is not None: + if ( + len(prisma_client.spend_log_transactons) > 0 + and base_url is not None + and db_writer_client is not None + ): if not base_url.endswith("/"): base_url += "/" response = await db_writer_client.post( diff --git a/litellm/tests/test_update_spend.py b/litellm/tests/test_update_spend.py index 0fd5d9bcf9d..4bc7f9fe4fb 100644 --- a/litellm/tests/test_update_spend.py +++ b/litellm/tests/test_update_spend.py @@ -92,4 +92,7 @@ async def test_batch_update_spend(prisma_client): 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() - await update_spend(prisma_client=litellm.proxy.proxy_server.prisma_client) + await update_spend( + prisma_client=litellm.proxy.proxy_server.prisma_client, + db_writer_client=None, + )