scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 227

test(spend): update daily tag spend retry test for tx-wrapped batch

Browse source 
_update_daily_spend now runs its batch inside db.tx(...).batch_(); wire the retry test's prisma mock through db.tx so the batch context manager resolves.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
shivam 2026-07-23 01:41:02 +00:00
parent 250de625af
commit f845c31f50
1 changed files with 12 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
tests/proxy_unit_tests/test_update_daily_tag_spend.py
View file

@ -96,7 +96,8 @@ async def test_daily_tag_spend_retries_then_succeeds():
mock_batcher.litellm_dailytagspend = mock_table
# Fail entering batch context 3 times with retryable DB errors, then succeed.
prisma_client.db.batch_.return_value.__aenter__ = AsyncMock(
mock_batch_context = MagicMock()
mock_batch_context.__aenter__ = AsyncMock(
side_effect=[
httpx.ConnectError("x"),
httpx.ConnectError("x"),
@ -104,6 +105,14 @@ async def test_daily_tag_spend_retries_then_succeeds():
mock_batcher,
]
)
mock_batch_context.__aexit__ = AsyncMock(return_value=False)
mock_transaction = MagicMock()
mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction)
mock_transaction.__aexit__ = AsyncMock(return_value=False)
mock_transaction.batch_ = MagicMock(return_value=mock_batch_context)
prisma_client.db.tx = MagicMock(return_value=mock_transaction)
daily_spend_transactions: Dict[str, DailyTagSpendTransaction] = {
"k": {
@ -138,6 +147,7 @@ async def test_daily_tag_spend_retries_then_succeeds():
daily_spend_transactions=daily_spend_transactions,
)
assert prisma_client.db.batch_.return_value.__aenter__.await_count == 4
assert mock_batch_context.__aenter__.await_count == 4
assert prisma_client.db.tx.call_count == 4
assert sleep_mock.await_count == 3
mock_table.upsert.assert_called_once()