Spaces:
Configuration error
Configuration error
File size: 6,512 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 |
import io
import os
import pathlib
import ssl
import sys
from unittest.mock import MagicMock, patch
import httpx
import pytest
from aiohttp import ClientSession, TCPConnector
sys.path.insert(
0, os.path.abspath("../../../..")
) # Adds the parent directory to the system path
import litellm
from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
@pytest.mark.asyncio
async def test_ssl_security_level(monkeypatch):
# Set environment variable for SSL security level
monkeypatch.setenv("SSL_SECURITY_LEVEL", "DEFAULT@SECLEVEL=1")
# Create async client with SSL verification disabled to isolate SSL context testing
client = AsyncHTTPHandler(ssl_verify=False)
# Get the transport (should be LiteLLMAiohttpTransport)
transport = client.client._transport
# Get the aiohttp ClientSession
client_session = transport._get_valid_client_session()
# Get the connector from the session
connector = client_session.connector
# Get the SSL context from the connector
ssl_context = connector._ssl
print("ssl_context", ssl_context)
# Verify that the SSL context exists and has the correct cipher string
assert isinstance(ssl_context, ssl.SSLContext)
# Optionally, check the ciphers string if needed
# assert "DEFAULT@SECLEVEL=1" in ssl_context.get_ciphers()
@pytest.mark.asyncio
async def test_force_ipv4_transport():
"""Test transport creation with force_ipv4 enabled"""
litellm.force_ipv4 = True
litellm.disable_aiohttp_transport = True
transport = AsyncHTTPHandler._create_async_transport()
# Should get an AsyncHTTPTransport
assert isinstance(transport, httpx.AsyncHTTPTransport)
# Verify IPv4 configuration through a request
client = httpx.AsyncClient(transport=transport)
try:
response = await client.get("http://example.com")
assert response.status_code == 200
finally:
await client.aclose()
@pytest.mark.asyncio
async def test_ssl_context_transport():
"""Test transport creation with SSL context"""
# Create a test SSL context
ssl_context = ssl.create_default_context()
transport = AsyncHTTPHandler._create_async_transport(ssl_context=ssl_context)
assert transport is not None
if isinstance(transport, LiteLLMAiohttpTransport):
# Get the client session and verify SSL context is passed through
client_session = transport._get_valid_client_session()
assert isinstance(client_session, ClientSession)
assert isinstance(client_session.connector, TCPConnector)
# Verify the connector has SSL context set by checking if it's using SSL
assert client_session.connector._ssl is not None
@pytest.mark.asyncio
async def test_aiohttp_disabled_transport():
"""Test transport creation with aiohttp disabled"""
litellm.disable_aiohttp_transport = True
litellm.force_ipv4 = False
transport = AsyncHTTPHandler._create_async_transport()
# Should get None when both aiohttp is disabled and force_ipv4 is False
assert transport is None
@pytest.mark.asyncio
async def test_ssl_verification_with_aiohttp_transport():
"""
Test aiohttp respects ssl_verify=False
We validate that the ssl settings for a litellm transport match what a ssl verify=False aiohttp client would have.
"""
import aiohttp
# Create a test SSL context
litellm_async_client = AsyncHTTPHandler(ssl_verify=False)
transport_connector = (
litellm_async_client.client._transport._get_valid_client_session().connector
)
print("transport_connector", transport_connector)
print("transport_connector._ssl", transport_connector._ssl)
aiohttp_session = aiohttp.ClientSession(
connector=aiohttp.TCPConnector(verify_ssl=False)
)
print("aiohttp_session", aiohttp_session)
print("aiohttp_session._ssl", aiohttp_session.connector._ssl)
# assert both litellm transport and aiohttp session have ssl_verify=False
assert transport_connector._ssl == aiohttp_session.connector._ssl
@pytest.mark.asyncio
async def test_disable_aiohttp_trust_env_with_env_variable(monkeypatch):
"""Test aiohttp transport respects DISABLE_AIOHTTP_TRUST_ENV environment variable"""
# Set environment variable to disable trust env
monkeypatch.setenv("DISABLE_AIOHTTP_TRUST_ENV", "True")
# Ensure aiohttp transport is enabled
litellm.disable_aiohttp_transport = False
# Create async client
client = AsyncHTTPHandler()
# Get the transport (should be LiteLLMAiohttpTransport)
transport = client.client._transport
assert isinstance(transport, LiteLLMAiohttpTransport)
# Get the aiohttp ClientSession
client_session = transport._get_valid_client_session()
# Verify that trust_env is False when DISABLE_AIOHTTP_TRUST_ENV is True
assert client_session._trust_env is False
@pytest.mark.asyncio
async def test_disable_aiohttp_trust_env_with_litellm_setting():
"""Test aiohttp transport respects litellm.disable_aiohttp_trust_env setting"""
# Set litellm setting to disable trust env
litellm.disable_aiohttp_trust_env = True
# Ensure aiohttp transport is enabled
litellm.disable_aiohttp_transport = False
# Create async client
client = AsyncHTTPHandler()
# Get the transport (should be LiteLLMAiohttpTransport)
transport = client.client._transport
assert isinstance(transport, LiteLLMAiohttpTransport)
# Get the aiohttp ClientSession
client_session = transport._get_valid_client_session()
# Verify that trust_env is False when litellm.disable_aiohttp_trust_env is True
assert client_session._trust_env is False
@pytest.mark.asyncio
async def test_enable_aiohttp_trust_env_default():
"""Test aiohttp transport enables trust_env by default"""
# Ensure both settings are disabled/default
litellm.disable_aiohttp_trust_env = False
litellm.disable_aiohttp_transport = False
# Create async client
client = AsyncHTTPHandler()
# Get the transport (should be LiteLLMAiohttpTransport)
transport = client.client._transport
assert isinstance(transport, LiteLLMAiohttpTransport)
# Get the aiohttp ClientSession
client_session = transport._get_valid_client_session()
# Verify that trust_env is True by default
assert client_session._trust_env is True
|