test(test_proxy_utils.py): ensure in-memory spend logs reset after bad log error

This commit is contained in:
Krrish Dholakia 2025-01-22 16:13:55 -08:00
parent 5132deb2fe
commit b3b06e9387

View file

@ -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"