diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 09bf7cd7dc1..a876d387517 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -1521,3 +1521,34 @@ async def test_end_user_transactions_reset(): assert ( mock_client.end_user_list_transactons == {} ), "Transactions list should be empty after error" + + +@pytest.mark.asyncio +async def test_spend_logs_cleanup_after_error(): + # Setup test data + mock_client = MagicMock() + mock_client.spend_log_transactions = [ + {"id": 1, "amount": 10.0}, + {"id": 2, "amount": 20.0}, + {"id": 3, "amount": 30.0}, + ] + # Make the DB operation fail + mock_client.db.litellm_spendlogs.create_many = AsyncMock( + side_effect=Exception("DB Error") + ) + + original_logs = mock_client.spend_log_transactions.copy() + + # Call function - should raise error + with pytest.raises(Exception): + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=0, + prisma_client=mock_client, + db_writer_client=None, # Test DB write path + proxy_logging_obj=MagicMock(), + ) + + # Verify the first batch was removed from spend_log_transactions + assert ( + mock_client.spend_log_transactions == original_logs[100:] + ), "Should remove processed logs even after error"