From f845c31f502f9fa387f244e0e8ca2bb193c7703f Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 23 Jul 2026 01:41:02 +0000 Subject: [PATCH] test(spend): update daily tag spend retry test for tx-wrapped batch _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> --- .../test_update_daily_tag_spend.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/tests/proxy_unit_tests/test_update_daily_tag_spend.py b/tests/proxy_unit_tests/test_update_daily_tag_spend.py index 80616ade5ef..204f3061a67 100644 --- a/tests/proxy_unit_tests/test_update_daily_tag_spend.py +++ b/tests/proxy_unit_tests/test_update_daily_tag_spend.py @@ -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()