From 5132deb2fe19d8ed39e7d9180dbbe134532cf68b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 22 Jan 2025 16:11:32 -0800 Subject: [PATCH] test(test_proxy_utils.py): add test to ensure bad log is dropped --- tests/proxy_unit_tests/test_proxy_utils.py | 26 ++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index a3de35a2abe..09bf7cd7dc1 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -1495,3 +1495,29 @@ def test_custom_openapi(mock_get_openapi_schema): openapi_schema = custom_openapi() assert openapi_schema is not None + + +import pytest +from unittest.mock import MagicMock, AsyncMock +import asyncio +from datetime import timedelta +from litellm.proxy.utils import ProxyUpdateSpend + + +@pytest.mark.asyncio +async def test_end_user_transactions_reset(): + # Setup + mock_client = MagicMock() + mock_client.end_user_list_transactons = {"1": 10.0} # Bad log + mock_client.db.tx = AsyncMock(side_effect=Exception("DB Error")) + + # Call function - should raise error + with pytest.raises(Exception): + await ProxyUpdateSpend.update_end_user_spend( + n_retry_times=0, prisma_client=mock_client, proxy_logging_obj=MagicMock() + ) + + # Verify cleanup happened + assert ( + mock_client.end_user_list_transactons == {} + ), "Transactions list should be empty after error"