From 5303d52a5e3bee7e131dcabd098e94f0613a7bb9 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 25 Sep 2025 19:24:33 -0400 Subject: [PATCH] fix lint --- litellm/proxy/utils.py | 13 +- .../test_prisma_client_backoff_retry.py | 329 ------------------ 2 files changed, 7 insertions(+), 335 deletions(-) delete mode 100644 tests/proxy_unit_tests/test_prisma_client_backoff_retry.py diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a4a3a2a6e12..7b85037595f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2816,13 +2816,13 @@ class PrismaClient: ) ) raise e - + @backoff.on_exception( - backoff.expo, - Exception, - max_tries=3, - max_time=10, - on_backoff=on_backoff, + backoff.expo, + Exception, + max_tries=3, + max_time=10, + on_backoff=on_backoff, ) async def health_check(self): """ @@ -2858,6 +2858,7 @@ class PrismaClient: """ Get the row count from LiteLLM_SpendLogs table using PostgreSQL system statistics. """ + @backoff.on_exception( backoff.expo, Exception, diff --git a/tests/proxy_unit_tests/test_prisma_client_backoff_retry.py b/tests/proxy_unit_tests/test_prisma_client_backoff_retry.py deleted file mode 100644 index 7698fbb8287..00000000000 --- a/tests/proxy_unit_tests/test_prisma_client_backoff_retry.py +++ /dev/null @@ -1,329 +0,0 @@ -""" -Test backoff retry mechanisms for PrismaClient methods during _setup_prisma_client state. - -This test validates that intermittent database connection issues are handled correctly -with exponential backoff retries for critical startup operations. -""" - -import asyncio -import pytest -import time -from unittest.mock import AsyncMock, MagicMock, patch, call -from unittest.mock import Mock -import sys -import os - -# Add project root to path -sys.path.insert(0, os.path.abspath("../..")) - -from litellm.proxy.utils import PrismaClient, ProxyLogging -from prisma.errors import PrismaError, ClientNotConnectedError -import httpx -import backoff - - -@pytest.fixture -def mock_prisma_client(): - """Create a mock PrismaClient with necessary attributes""" - client = MagicMock(spec=PrismaClient) - - # Mock database connection - client.db = AsyncMock() - client.db.query_raw = AsyncMock() - - # Mock proxy logging object - client.proxy_logging_obj = AsyncMock() - client.proxy_logging_obj.failure_handler = AsyncMock() - - return client - - -@pytest.fixture -def mock_proxy_logging(): - """Create a mock ProxyLogging object""" - proxy_logging = AsyncMock(spec=ProxyLogging) - proxy_logging.failure_handler = AsyncMock() - return proxy_logging - - -@pytest.fixture -def connection_errors(): - """Common database connection errors that should trigger retries""" - return [ - httpx.ConnectError("Connection failed"), - httpx.TimeoutException("Request timeout"), - ClientNotConnectedError(), - PrismaError("Database connection lost"), - ConnectionError("Network connection failed"), - OSError("Connection refused"), - ] - - -class TestPrismaClientBackoffRetry: - """Test suite for PrismaClient backoff retry mechanisms""" - - @pytest.mark.asyncio - async def test_health_check_success_no_retry(self, mock_prisma_client, mock_proxy_logging): - """Test health_check succeeds immediately without retries""" - # Mock successful query response - mock_prisma_client.db.query_raw.return_value = [{"result": 1}] - - # Create real PrismaClient instance with mocked db - client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) - client.db = mock_prisma_client.db - client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj - - # Call health_check - result = await client.health_check() - - # Verify success - assert result == [{"result": 1}] - mock_prisma_client.db.query_raw.assert_called_once_with("SELECT 1") - - @pytest.mark.asyncio - async def test_health_check_retry_then_success(self, mock_prisma_client, mock_proxy_logging, connection_errors): - """Test health_check retries on connection errors and eventually succeeds""" - # Mock first two calls to fail, third to succeed - mock_prisma_client.db.query_raw.side_effect = [ - connection_errors[0], # First call fails - connection_errors[1], # Second call fails - [{"result": 1}] # Third call succeeds - ] - - # Create real PrismaClient instance - client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) - client.db = mock_prisma_client.db - client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj - - # Measure execution time to verify backoff delay - start_time = time.time() - result = await client.health_check() - end_time = time.time() - - # Verify eventual success - assert result == [{"result": 1}] - - # Verify retry attempts (3 calls total) - assert mock_prisma_client.db.query_raw.call_count == 3 - - # Verify backoff delay occurred (should be at least a few milliseconds) - assert end_time - start_time > 0.01 - - @pytest.mark.asyncio - async def test_health_check_max_retries_exceeded(self, mock_prisma_client, mock_proxy_logging, connection_errors): - """Test health_check fails after max retries (3) are exceeded""" - # Mock all calls to fail - mock_prisma_client.db.query_raw.side_effect = connection_errors[0] - - # Create real PrismaClient instance - client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) - client.db = mock_prisma_client.db - client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj - - # Expect final exception after retries - with pytest.raises(httpx.ConnectError): - await client.health_check() - - # Verify max retries attempted (3 attempts) - assert mock_prisma_client.db.query_raw.call_count == 3 - - @pytest.mark.asyncio - async def test_get_spend_logs_row_count_success_no_retry(self, mock_prisma_client, mock_proxy_logging): - """Test _get_spend_logs_row_count succeeds immediately""" - # Mock successful query response - mock_prisma_client.db.query_raw.return_value = [{"reltuples": 1000}] - - # Create real PrismaClient instance - client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) - client.db = mock_prisma_client.db - client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj - - result = await client._get_spend_logs_row_count() - - assert result == 1000 - mock_prisma_client.db.query_raw.assert_called_once() - - @pytest.mark.asyncio - async def test_get_spend_logs_row_count_retry_then_success(self, mock_prisma_client, mock_proxy_logging, connection_errors): - """Test _get_spend_logs_row_count retries and eventually succeeds""" - # Mock first call fails, second succeeds - mock_prisma_client.db.query_raw.side_effect = [ - connection_errors[2], # First call fails (ClientNotConnectedError) - [{"reltuples": 500}] # Second call succeeds - ] - - client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) - client.db = mock_prisma_client.db - client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj - - result = await client._get_spend_logs_row_count() - - assert mock_prisma_client.db.query_raw.call_count == 2 - - @pytest.mark.asyncio - async def test_get_spend_logs_row_count_handles_errors_gracefully(self, mock_prisma_client, mock_proxy_logging): - """Test _get_spend_logs_row_count returns 0 on persistent errors""" - # Mock all calls to fail - mock_prisma_client.db.query_raw.side_effect = PrismaError("Persistent DB error") - - client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) - client.db = mock_prisma_client.db - client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj - - result = await client._get_spend_logs_row_count() - - assert result == 0 - assert mock_prisma_client.db.query_raw.call_count == 3 # Max retries attempted - - @pytest.mark.asyncio - async def test_set_spend_logs_row_count_in_proxy_state_success(self, mock_prisma_client, mock_proxy_logging): - """Test _set_spend_logs_row_count_in_proxy_state succeeds""" - # Mock successful query response - mock_prisma_client.db.query_raw.return_value = [{"reltuples": 2000}] - - client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) - client.db = mock_prisma_client.db - client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj - - # Mock proxy_state - with patch('litellm.proxy.proxy_server.proxy_state') as mock_proxy_state: - mock_proxy_state.set_proxy_state_variable = Mock() - - await client._set_spend_logs_row_count_in_proxy_state() - - # Verify proxy state was updated - mock_proxy_state.set_proxy_state_variable.assert_called_once_with( - variable_name="spend_logs_row_count", - value=2000 - ) - - @pytest.mark.asyncio - async def test_set_spend_logs_row_count_retry_behavior(self, mock_prisma_client, mock_proxy_logging, connection_errors): - """Test _set_spend_logs_row_count_in_proxy_state retries on database errors""" - mock_prisma_client.db.query_raw.side_effect = [ - connection_errors[3], # First call fails (PrismaError) - [{"reltuples": 1500}] # Second call succeeds - ] - - client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) - client.db = mock_prisma_client.db - client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj - - with patch('litellm.proxy.proxy_server.proxy_state') as mock_proxy_state: - mock_proxy_state.set_proxy_state_variable = Mock() - - await client._set_spend_logs_row_count_in_proxy_state() - - assert mock_prisma_client.db.query_raw.call_count == 2 - mock_proxy_state.set_proxy_state_variable.assert_called_once_with( - variable_name="spend_logs_row_count", - value=1500 - ) - - @pytest.mark.asyncio - async def test_backoff_configuration_parameters(self, mock_proxy_logging): - """Test that backoff decorators are configured with correct parameters""" - client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) - - # Check that methods have backoff decorators - assert hasattr(client.health_check, '__wrapped__') - assert hasattr(client._set_spend_logs_row_count_in_proxy_state, '__wrapped__') - - # Verify backoff configuration exists (methods should have retry behavior) - # This is implicit verification - the decorators are applied in the source code - - @pytest.mark.asyncio - async def test_multiple_connection_error_types(self, mock_prisma_client, mock_proxy_logging): - """Test that different types of connection errors all trigger retries""" - error_types = [ - httpx.ConnectError("Connection error"), - httpx.TimeoutException("Timeout error"), - ClientNotConnectedError(), - PrismaError("Database error"), - ConnectionError("Network error"), - OSError("OS-level connection error") - ] - - for error_type in error_types: - # Reset mock for each test - mock_prisma_client.db.query_raw.reset_mock() - mock_prisma_client.db.query_raw.side_effect = [ - error_type, # First call fails - [{"result": 1}] # Second call succeeds - ] - - client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) - client.db = mock_prisma_client.db - client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj - - # Should succeed after retry - result = await client.health_check() - assert result == [{"result": 1}] - assert mock_prisma_client.db.query_raw.call_count == 2 - - @pytest.mark.asyncio - async def test_setup_prisma_client_integration(self, mock_prisma_client, mock_proxy_logging): - """Test simulated _setup_prisma_client flow with intermittent failures""" - # This simulates the actual flow that happens in proxy_server.py _setup_prisma_client - - client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) - client.db = mock_prisma_client.db - client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj - - # Simulate intermittent failures followed by success - call_count = 0 - def mock_query_side_effect(*args, **kwargs): - nonlocal call_count - call_count += 1 - if call_count <= 2: # First two calls fail - raise httpx.ConnectError("Intermittent connection issue") - else: - if "SELECT 1" in str(args): - return [{"result": 1}] - else: - return [{"reltuples": 1000}] - - mock_prisma_client.db.query_raw.side_effect = mock_query_side_effect - - with patch('litellm.proxy.proxy_server.proxy_state') as mock_proxy_state: - mock_proxy_state.set_proxy_state_variable = Mock() - - # Execute the two critical calls from _setup_prisma_client - health_result = await client.health_check() - await client._set_spend_logs_row_count_in_proxy_state() - - # Verify both operations eventually succeeded despite initial failures - assert health_result == [{"result": 1}] - mock_proxy_state.set_proxy_state_variable.assert_called_once() - - # Verify retries occurred (should have more than 2 total query calls) - assert mock_prisma_client.db.query_raw.call_count >= 4 - - @pytest.mark.asyncio - async def test_backoff_timing_constraints(self, mock_prisma_client, mock_proxy_logging, connection_errors): - """Test that backoff respects max_time constraint (10 seconds)""" - # Mock all calls to fail to test max_time - mock_prisma_client.db.query_raw.side_effect = connection_errors[0] - - client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) - client.db = mock_prisma_client.db - client.proxy_logging_obj = mock_prisma_client.proxy_logging_obj - - start_time = time.time() - - with pytest.raises(httpx.ConnectError): - await client.health_check() - - end_time = time.time() - duration = end_time - start_time - - # Should not exceed max_time of 10 seconds significantly - # Adding small buffer for test execution overhead - assert duration < 12.0 - - # Should have attempted max_tries (3) retries - assert mock_prisma_client.db.query_raw.call_count == 3 - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) \ No newline at end of file