phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

test(test_update_spend.py): allow db_client to be none

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-28 13:44:40 -07:00
parent a09818e72e
commit e8d80509b1
2 changed files with 12 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

10
litellm/proxy/utils.py
View file

@ -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. Batch write updates to db.
@ -1996,7 +1998,11 @@ async def update_spend(prisma_client: PrismaClient, db_writer_client: HTTPHandle
### UPDATE SPEND LOGS ### ### UPDATE SPEND LOGS ###
base_url = os.getenv("SPEND_LOGS_URL", None) 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("/"): if not base_url.endswith("/"):
base_url += "/" base_url += "/"
response = await db_writer_client.post( response = await db_writer_client.post(

5
litellm/tests/test_update_spend.py
View file

@ -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, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect() 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,
)