tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-11 19:48:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

test: add unit tests

Browse Source
This commit is contained in:
Krrish Dholakia
2025-10-06 13:53:36 -07:00
parent 6177b23ebf
commit d7aab1d5de
1 changed files with 164 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_litellm/passthrough/test_passthrough_main.py
+164 -64
View File
@@ -53,27 +53,38 @@ def test_llm_passthrough_route():
def test_bedrock_application_inference_profile_url_encoding():
client = HTTPHandler()
mock_provider_config = MagicMock()
mock_provider_config.get_complete_url.return_value = (
httpx.URL("https://bedrock-runtime.us-east-1.amazonaws.com/model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse"),
"https://bedrock-runtime.us-east-1.amazonaws.com"
httpx.URL(
"https://bedrock-runtime.us-east-1.amazonaws.com/model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse"
),
"https://bedrock-runtime.us-east-1.amazonaws.com",
)
mock_provider_config.get_api_key.return_value = "test-key"
mock_provider_config.validate_environment.return_value = {}
mock_provider_config.sign_request.return_value = ({}, None)
mock_provider_config.is_streaming_request.return_value = False
with patch("litellm.utils.ProviderConfigManager.get_provider_passthrough_config", return_value=mock_provider_config), \
patch("litellm.litellm_core_utils.get_litellm_params.get_litellm_params", return_value={}), \
patch("litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("test-model", "bedrock", "test-key", "test-base")), \
patch.object(client.client, "send", return_value=MagicMock(status_code=200)) as mock_send, \
patch.object(client.client, "build_request") as mock_build_request:
with patch(
"litellm.utils.ProviderConfigManager.get_provider_passthrough_config",
return_value=mock_provider_config,
), patch(
"litellm.litellm_core_utils.get_litellm_params.get_litellm_params",
return_value={},
), patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
return_value=("test-model", "bedrock", "test-key", "test-base"),
), patch.object(
client.client, "send", return_value=MagicMock(status_code=200)
) as mock_send, patch.object(
client.client, "build_request"
) as mock_build_request:
# Mock logging object
mock_logging_obj = MagicMock()
mock_logging_obj.update_environment_variables = MagicMock()
response = llm_passthrough_route(
model="arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd",
endpoint="model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse",
@@ -86,7 +97,7 @@ def test_bedrock_application_inference_profile_url_encoding():
# Verify that build_request was called with the encoded URL
mock_build_request.assert_called_once()
call_args = mock_build_request.call_args
# The URL should have the application-inference-profile ID encoded
actual_url = str(call_args.kwargs["url"])
assert "application-inference-profile%2Fr742sbn2zckd" in actual_url
@@ -95,28 +106,39 @@ def test_bedrock_application_inference_profile_url_encoding():
def test_bedrock_non_application_inference_profile_no_encoding():
client = HTTPHandler()
# Mock the provider config and its methods
mock_provider_config = MagicMock()
mock_provider_config.get_complete_url.return_value = (
httpx.URL("https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-sonnet-20240229-v1:0/converse"),
"https://bedrock-runtime.us-east-1.amazonaws.com"
httpx.URL(
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-sonnet-20240229-v1:0/converse"
),
"https://bedrock-runtime.us-east-1.amazonaws.com",
)
mock_provider_config.get_api_key.return_value = "test-key"
mock_provider_config.validate_environment.return_value = {}
mock_provider_config.sign_request.return_value = ({}, None)
mock_provider_config.is_streaming_request.return_value = False
with patch("litellm.utils.ProviderConfigManager.get_provider_passthrough_config", return_value=mock_provider_config), \
patch("litellm.litellm_core_utils.get_litellm_params.get_litellm_params", return_value={}), \
patch("litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("test-model", "bedrock", "test-key", "test-base")), \
patch.object(client.client, "send", return_value=MagicMock(status_code=200)) as mock_send, \
patch.object(client.client, "build_request") as mock_build_request:
with patch(
"litellm.utils.ProviderConfigManager.get_provider_passthrough_config",
return_value=mock_provider_config,
), patch(
"litellm.litellm_core_utils.get_litellm_params.get_litellm_params",
return_value={},
), patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
return_value=("test-model", "bedrock", "test-key", "test-base"),
), patch.object(
client.client, "send", return_value=MagicMock(status_code=200)
) as mock_send, patch.object(
client.client, "build_request"
) as mock_build_request:
# Mock logging object
mock_logging_obj = MagicMock()
mock_logging_obj.update_environment_variables = MagicMock()
response = llm_passthrough_route(
model="anthropic.claude-3-sonnet-20240229-v1:0",
endpoint="model/anthropic.claude-3-sonnet-20240229-v1:0/converse",
@@ -129,7 +151,7 @@ def test_bedrock_non_application_inference_profile_no_encoding():
# Verify that build_request was called with the original URL (no encoding)
mock_build_request.assert_called_once()
call_args = mock_build_request.call_args
# The URL should NOT have application-inference-profile encoding
actual_url = str(call_args.kwargs["url"])
assert "application-inference-profile%2F" not in actual_url
@@ -151,21 +173,21 @@ def test_update_stream_param_based_on_request_body():
parsed_body=parsed_body, stream=False
)
assert result is True
# Test 2: no stream in request body should return original stream param
parsed_body = {"model": "test-model"}
result = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body(
parsed_body=parsed_body, stream=False
)
assert result is False
# Test 3: stream=False in request body should return False
parsed_body = {"stream": False, "model": "test-model"}
result = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body(
parsed_body=parsed_body, stream=True
)
assert result is False
# Test 4: no stream param provided, no stream in body
parsed_body = {"model": "test-model"}
result = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body(
@@ -178,14 +200,14 @@ def test_update_stream_param_based_on_request_body():
def mock_request():
"""Create a mock request with headers"""
from typing import Optional
class QueryParams:
def __init__(self):
self._dict = {}
def __iter__(self):
return iter(self._dict)
def items(self):
return self._dict.items()
@@ -210,6 +232,7 @@ def mock_request():
def mock_user_api_key_dict():
"""Create a mock user API key dictionary"""
from litellm.proxy._types import UserAPIKeyAuth
return UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
@@ -223,8 +246,8 @@ async def test_pass_through_request_stream_param_override(
mock_request, mock_user_api_key_dict
):
"""
Test that when stream=None is passed as parameter but stream=True
is in request body, the request body value takes precedence and
Test that when stream=None is passed as parameter but stream=True
is in request body, the request body value takes precedence and
the eventual POST request uses streaming.
"""
from unittest.mock import AsyncMock, Mock, patch
@@ -238,29 +261,29 @@ async def test_pass_through_request_stream_param_override(
"model": "claude-3-5-sonnet-20241022",
"max_tokens": 256,
"messages": [{"role": "user", "content": "Hello, world"}],
"stream": True # This should override the function parameter
"stream": True, # This should override the function parameter
}
# Create a mock streaming response
mock_response = AsyncMock()
mock_response.status_code = 200
mock_response.headers = {"content-type": "text/event-stream"}
# Mock the streaming response behavior
async def mock_aiter_bytes():
yield b'data: {"content": "Hello"}\n\n'
yield b'data: {"content": "World"}\n\n'
yield b'data: [DONE]\n\n'
yield b"data: [DONE]\n\n"
mock_response.aiter_bytes = mock_aiter_bytes
# Create mocks for the async client
mock_async_client = AsyncMock()
mock_request_obj = AsyncMock()
# Mock build_request to return a request object (it's a sync method)
mock_async_client.build_request = Mock(return_value=mock_request_obj)
# Mock send to return the streaming response
mock_async_client.send.return_value = mock_response
@@ -269,9 +292,7 @@ async def test_pass_through_request_stream_param_override(
mock_client_obj.client = mock_async_client
# Create the request
request = mock_request(
headers={}, method="POST", request_body=request_body
)
request = mock_request(headers={}, method="POST", request_body=request_body)
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client",
@@ -298,33 +319,32 @@ async def test_pass_through_request_stream_param_override(
httpx.URL("https://api.anthropic.com/v1/messages"),
json=request_body,
params={},
headers={
"Authorization": "Bearer test-key"
},
headers={"Authorization": "Bearer test-key"},
)
# Verify that send was called with stream=True
mock_async_client.send.assert_called_once_with(
mock_request_obj,
stream=True # This proves that stream=True from request body was used
mock_request_obj,
stream=True, # This proves that stream=True from request body was used
)
# Verify that the non-streaming request method was NOT called
mock_async_client.request.assert_not_called()
# Verify response is a StreamingResponse
from fastapi.responses import StreamingResponse
assert isinstance(response, StreamingResponse)
assert response.status_code == 200
@pytest.mark.asyncio
@pytest.mark.asyncio
async def test_pass_through_request_stream_param_no_override(
mock_request, mock_user_api_key_dict
):
"""
Test that when stream=False is passed as parameter and no stream
is in request body, the function parameter is used and
Test that when stream=False is passed as parameter and no stream
is in request body, the function parameter is used and
the eventual request uses non-streaming.
"""
from unittest.mock import AsyncMock, Mock, patch
@@ -335,7 +355,7 @@ async def test_pass_through_request_stream_param_no_override(
# Create request body without stream parameter
request_body = {
"model": "claude-3-5-sonnet-20241022",
"model": "claude-3-5-sonnet-20241022",
"max_tokens": 256,
"messages": [{"role": "user", "content": "Hello, world"}],
# No stream parameter - should use function parameter stream=False
@@ -346,15 +366,15 @@ async def test_pass_through_request_stream_param_no_override(
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response._content = b'{"response": "Hello world"}'
async def mock_aread():
return mock_response._content
mock_response.aread = mock_aread
# Create mocks for the async client
mock_async_client = AsyncMock()
# Mock request to return the non-streaming response
mock_async_client.request.return_value = mock_response
@@ -363,9 +383,7 @@ async def test_pass_through_request_stream_param_no_override(
mock_client_obj.client = mock_async_client
# Create the request
request = mock_request(
headers={}, method="POST", request_body=request_body
)
request = mock_request(headers={}, method="POST", request_body=request_body)
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client",
@@ -388,23 +406,105 @@ async def test_pass_through_request_stream_param_no_override(
# Verify that build_request was NOT called (no streaming path)
mock_async_client.build_request.assert_not_called()
# Verify that send was NOT called (no streaming path)
mock_async_client.send.assert_not_called()
# Verify that the non-streaming request method WAS called
mock_async_client.request.assert_called_once_with(
method="POST",
url=httpx.URL("https://api.anthropic.com/v1/messages"),
headers={
"Authorization": "Bearer test-key"
},
headers={"Authorization": "Bearer test-key"},
params={},
json=request_body,
)
# Verify response is a regular Response (not StreamingResponse)
from fastapi.responses import Response, StreamingResponse
assert not isinstance(response, StreamingResponse)
assert isinstance(response, Response)
assert response.status_code == 200
assert response.status_code == 200
def test_azure_with_custom_api_base_and_key():
"""
Test that llm_passthrough_route correctly handles Azure OpenAI
with custom api_base and api_key.
"""
client = HTTPHandler()
# Mock the provider config and its methods
mock_provider_config = MagicMock()
mock_provider_config.get_complete_url.return_value = (
httpx.URL(
"https://my-custom-base/openai/deployments/gpt-4.1/chat/completions?api-version=2024-02-01"
),
"https://my-custom-base",
)
mock_provider_config.get_api_key.return_value = "my-custom-key"
mock_provider_config.validate_environment.return_value = {
"api-key": "my-custom-key"
}
mock_provider_config.sign_request.return_value = (
{"api-key": "my-custom-key"},
None,
)
mock_provider_config.is_streaming_request.return_value = False
with patch(
"litellm.utils.ProviderConfigManager.get_provider_passthrough_config",
return_value=mock_provider_config,
), patch(
"litellm.litellm_core_utils.get_litellm_params.get_litellm_params",
return_value={},
), patch(
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
return_value=("gpt-4.1", "azure", "my-custom-key", "https://my-custom-base"),
), patch.object(
client.client,
"send",
return_value=MagicMock(
status_code=200, json=lambda: {"id": "chatcmpl-123", "choices": []}
),
) as mock_send, patch.object(
client.client, "build_request"
) as mock_build_request:
# Mock logging object
mock_logging_obj = MagicMock()
mock_logging_obj.update_environment_variables = MagicMock()
response = llm_passthrough_route(
model="azure/gpt-4.1",
endpoint="openai/deployments/gpt-4.1/chat/completions",
method="POST",
custom_llm_provider="azure",
api_base="https://my-custom-base",
api_key="my-custom-key",
json={
"model": "gpt-4.1",
"messages": [{"role": "user", "content": "Hello!"}],
},
client=client,
litellm_logging_obj=mock_logging_obj,
)
# Verify that build_request was called with the correct parameters
mock_build_request.assert_called_once()
call_args = mock_build_request.call_args
# Verify the URL contains the custom base
actual_url = str(call_args.kwargs["url"])
assert "my-custom-base" in actual_url
assert "gpt-4.1" in actual_url
# Verify the headers contain the custom API key
headers = call_args.kwargs["headers"]
assert headers["api-key"] == "my-custom-key"
# Verify the model in JSON body is updated
json_body = call_args.kwargs["json"]
assert json_body["model"] == "gpt-4.1"
assert response.status_code == 200