test_retry.py

#!/usr/bin/env python3
"""Test script to verify retry logic in LLMClient."""

import time
import logging
from unittest.mock import Mock, patch
from src.services.llm_client import LLMClient
from src.config.settings import settings

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def test_retry_logic():
    """Test the retry logic with simulated failures."""
    
    # Create LLMClient instance
    client = LLMClient(settings)
    
    # Create a mock exception that simulates a 503 error
    class Mock503Error(Exception):
        def __init__(self):
            self.status_code = 503
            super().__init__("Service Unavailable")
    
    # Test with a mock that fails twice then succeeds
    with patch('openai.responses.create') as mock_create:
        # First two calls fail with 503, third succeeds
        mock_create.side_effect = [
            Mock503Error(),
            Mock503Error(),
            Mock(
                output=[Mock(content=[Mock(text="Success!")])],
                usage=Mock(input_tokens=10, output_tokens=5)
            )
        ]
        
        start_time = time.time()
        try:
            result = client.responses("Test prompt", max_retries=2, base_delay=0.1)
            end_time = time.time()
            
            logger.info(f"Test completed successfully!")
            logger.info(f"Result: {result}")
            logger.info(f"Time taken: {end_time - start_time:.2f} seconds")
            logger.info(f"Number of calls made: {mock_create.call_count}")
            
            assert result == "Success!"
            assert mock_create.call_count == 3  # 2 failures + 1 success
            
        except Exception as e:
            logger.error(f"Test failed: {e}")
            raise

def test_non_retryable_error():
    """Test that non-retryable errors are not retried."""
    
    client = LLMClient(settings)
    
    class Mock400Error(Exception):
        def __init__(self):
            self.status_code = 400
            super().__init__("Bad Request")
    
    with patch('openai.responses.create') as mock_create:
        # Should not retry 400 errors
        mock_create.side_effect = Mock400Error()
        
        try:
            client.responses("Test prompt", max_retries=3, base_delay=0.1)
            assert False, "Should have raised an exception"
        except Mock400Error:
            logger.info("Correctly did not retry 400 error")
            assert mock_create.call_count == 1  # Only one call, no retries

if __name__ == "__main__":
    logger.info("Testing retry logic...")
    test_retry_logic()
    
    logger.info("Testing non-retryable error...")
    test_non_retryable_error()
    
    logger.info("All tests passed!")