litellm/tests/litellm_utils_tests/test_aiohttp_handler.py

import asyncio
import copy
import sys
import time
from datetime import datetime
from unittest import mock

from dotenv import load_dotenv

from litellm.types.utils import StandardCallbackDynamicParams

load_dotenv()
import os

sys.path.insert(
    0, os.path.abspath("../..")
)  # Adds the parent directory to the system-path
import pytest

import litellm
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler

@pytest.mark.asyncio
async def test_client_session_helper():
    """Test that the client session helper handles event loop changes correctly"""
    try:
        # Create a transport with the new helper
        transport = AsyncHTTPHandler._create_aiohttp_transport()
        if transport is not None:
            print('✅ Successfully created aiohttp transport with helper')
            
            # Test the helper function directly if it's a LiteLLMAiohttpTransport
            if hasattr(transport, '_get_valid_client_session'):
                session1 = transport._get_valid_client_session()  # type: ignore
                print(f'✅ First session created: {type(session1).__name__}')
                
                # Call it again to test reuse
                session2 = transport._get_valid_client_session()  # type: ignore
                print(f'✅ Second session call: {type(session2).__name__}')
                
                # In the same event loop, should be the same session
                print(f'✅ Same session reused: {session1 is session2}')
                
            return True
        else:
            print('ℹ️  No aiohttp transport available (probably missing httpx-aiohttp)')
            return True
    except Exception as e:
        print(f'❌ Error: {e}')
        import traceback
        traceback.print_exc()
        return False

async def test_event_loop_robustness():
    """Test behavior when event loops change (simulating CI/CD scenario)"""
    try:
        # Test session creation in multiple scenarios
        transport = AsyncHTTPHandler._create_aiohttp_transport()
        
        if transport and hasattr(transport, '_get_valid_client_session'):
            # Test 1: Normal usage
            session = transport._get_valid_client_session()  # type: ignore
            print(f'✅ Normal session creation works: {session is not None}')
            
            # Test 2: Force recreation by setting client to a callable
            from aiohttp import ClientSession
            transport.client = lambda: ClientSession()  # type: ignore
            session2 = transport._get_valid_client_session()  # type: ignore
            print(f'✅ Session recreation after callable works: {session2 is not None}')
            
            return True
        else:
            print('ℹ️  Transport not available or no helper method')
            return True
            
    except Exception as e:
        print(f'❌ Error in event loop robustness test: {e}')
        import traceback
        traceback.print_exc()
        return False

async def test_httpx_request_simulation():
    """Test that the transport can handle a simulated HTTP request"""
    try:
        transport = AsyncHTTPHandler._create_aiohttp_transport()
        
        if transport is not None:
            print('✅ Transport created for request simulation')
            
            # Create a simple httpx request to test with
            import httpx
            request = httpx.Request('GET', 'https://httpbin.org/headers')
            
            # Just test that we can get a valid session for this request context
            if hasattr(transport, '_get_valid_client_session'):
                session = transport._get_valid_client_session()  # type: ignore
                print(f'✅ Got valid session for request: {session is not None}')
                
                # Test that session has required aiohttp methods
                has_request_method = hasattr(session, 'request')
                print(f'✅ Session has request method: {has_request_method}')
                
                return has_request_method
            
            return True
        else:
            print('ℹ️  No transport available for request simulation')
            return True
            
    except Exception as e:
        print(f'❌ Error in request simulation: {e}')
        return False

if __name__ == "__main__":
    print("Testing client session helper and event loop handling fix...")
    
    result1 = asyncio.run(test_client_session_helper())
    result2 = asyncio.run(test_event_loop_robustness()) 
    result3 = asyncio.run(test_httpx_request_simulation())
    
    if result1 and result2 and result3:
        print("🎉 All tests passed! The helper function approach should fix the CI/CD event loop issues.")
    else:
        print("💥 Some tests failed")