scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 99

test: take keys out of the legacy proxy, enterprise and mcp unit tests before moving them (#42901)

Browse source 
* ci: fix the litellm-tests unit job with sysmon coverage, an env allowlist and coverage upload on failure

* test: replace key-dependent proxy, enterprise and mcp unit tests with synthetic values and integration and e2e coverage

* test: drop key reads at the legacy proxy, enterprise and mcp paths and wire the gemini pass-through split

* ci: fail the unit shard when circleci tests split errors

* test: drop restating comments from the gemini pass-through split

* ci: exit the unit shard cleanly when circleci tests split assigns it no files

---------

Co-authored-by: yuneng <yuneng@berri.ai>
This commit is contained in:
devin-ai-integration[bot] 2026-09-24 14:56:52 -07:00 • committed by GitHub
parent 1628978db7
commit b72a030501
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 652 additions and 693 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.github/workflows/test-unit-proxy-db.yml vendored
View file

@ -106,6 +106,7 @@ jobs:
- test-group: proxy-server-core
test-path: >-
tests/proxy_unit_tests/test_proxy_server.py
tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py
tests/proxy_unit_tests/test_aproxy_startup.py
workers: 4
dist: loadscope
@ -115,7 +116,6 @@ jobs:
tests/proxy_unit_tests/test_proxy_config_unit_test.py
tests/proxy_unit_tests/test_proxy_routes.py
tests/proxy_unit_tests/test_server_root_path.py
tests/proxy_unit_tests/test_proxy_pass_user_config.py
tests/proxy_unit_tests/test_proxy_token_counter.py
tests/proxy_unit_tests/test_request_size_limit_middleware.py
tests/proxy_unit_tests/test_multipart_bypass_repro.py

78
tests/e2e/llm_translation/test_token_counter_gemini_contents_e2e.py Normal file
View file

@ -0,0 +1,78 @@
"""Live e2e: `/utils/token_counter?call_endpoint=true` counts Gemini `contents` upstream.
Google's countTokens API is the only tokenizer that knows Gemini's real token
boundaries, so the proxy must forward `contents` to it for both the AI Studio and
Vertex deployments and hand back the provider's `promptTokensDetails`. Claude on
Vertex is covered by `/v1/messages/count_tokens`; this is the Gemini `contents`
route the claude_code rows never reach
"""
from __future__ import annotations
import pytest
from e2e_config import unique_marker
from e2e_http import require_successful_call
from proxy_client import ProxyClient
from pydantic import BaseModel
pytestmark = pytest.mark.e2e
GEMINI_DEPLOYMENTS = ("gemini-2.5-flash", "gemini-2.5-flash-vertex")
class _Part(BaseModel):
text: str
class _Content(BaseModel):
parts: tuple[_Part, ...]
class _TokenCountBody(BaseModel):
model: str
contents: tuple[_Content, ...]
class _CallEndpoint(BaseModel):
call_endpoint: bool = True
class _ModalityTokens(BaseModel):
modality: str
tokenCount: int
class _CountTokensUpstream(BaseModel):
totalTokens: int
promptTokensDetails: tuple[_ModalityTokens, ...]
class _TokenCountResponse(BaseModel):
total_tokens: int
request_model: str
model_used: str
tokenizer_type: str
original_response: _CountTokensUpstream
class TestGeminiContentsTokenCounting:
@pytest.mark.parametrize("model", GEMINI_DEPLOYMENTS)
def test_contents_are_counted_by_the_provider_endpoint(
self, proxy: ProxyClient, scoped_key: str, model: str
) -> None:
text = f"Hello world, how are you doing today? {unique_marker()}"
body = _TokenCountBody(model=model, contents=(_Content(parts=(_Part(text=text),)),))
result = proxy.transport.send(
"/utils/token_counter",
headers=proxy.transport.bearer(scoped_key),
json=body,
params=_CallEndpoint(),
)
require_successful_call(result)
counted = _TokenCountResponse.model_validate_json(result.body)
assert counted.request_model == model, counted
assert counted.original_response.totalTokens == counted.total_tokens > 0, counted
assert counted.original_response.promptTokensDetails, counted
assert all(detail.tokenCount > 0 for detail in counted.original_response.promptTokensDetails), counted

10
tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py
View file

@ -13,7 +13,6 @@ import asyncio
from dotenv import load_dotenv
load_dotenv()
import os
from unittest.mock import MagicMock
@ -165,9 +164,9 @@ async def test_prometheus_metric_tracking():
"model_name": "gpt-5-mini", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": os.getenv("AZURE_AI_API_VERSION"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_key": "sk-azure-unit-test",
"api_version": "2025-01-01-preview",
"api_base": "https://unit-test.openai.azure.com",
},
"model_info": {"id": "azure-model-id"},
},
@ -180,9 +179,6 @@ async def test_prometheus_metric_tracking():
},
],
provider_budget_config=provider_budget_config,
redis_host=os.getenv("REDIS_HOST"),
redis_port=int(os.getenv("REDIS_PORT", 6379)),
redis_password=os.getenv("REDIS_PASSWORD"),
)
try:

103
tests/integration/routing/test_user_config_routing.py Normal file
View file

@ -0,0 +1,103 @@
import json
import uuid
from pathlib import Path
from typing import Final
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
USER_KEY: Final = "sk-user-supplied-" + uuid.uuid4().hex
def _completion(request: Request) -> Reply:
if request.target != "/v1/chat/completions":
return Reply(status=404, body=b"{}")
body: Final = json.loads(request.body)
return Reply(
body=json.dumps(
{
"id": "chatcmpl-user-config",
"object": "chat.completion",
"created": 0,
"model": body["model"],
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "routed"}, "finish_reason": "stop"}
],
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4},
}
).encode()
)
def _user_config(upstream_url: str) -> dict[str, object]:
return {
"model_list": [
{
"model_name": "user-config-deployment",
"litellm_params": {
"model": "openai/gpt-4.1-mini",
"api_base": upstream_url + "/v1",
"api_key": USER_KEY,
},
}
],
"num_retries": 0,
}
def _opt_in_config(directory: Path, upstream_url: str) -> Path:
config: Final = directory / "allow_client_side_credentials_config.yaml"
config.write_text(
json.dumps(
{
"model_list": [
{
"model_name": "admin-deployment",
"litellm_params": {
"model": "openai/gpt-4.1-mini",
"api_base": upstream_url + "/v1",
"api_key": "sk-admin-configured",
},
}
],
"general_settings": {
"master_key": "os.environ/LITELLM_MASTER_KEY",
"database_url": "os.environ/DATABASE_URL",
"store_model_in_db": True,
"allow_client_side_credentials": True,
},
}
)
)
return config
def _request_body(upstream_url: str) -> dict[str, object]:
return {
"model": "user-config-deployment",
"messages": [{"role": "user", "content": "user config control"}],
"user_config": _user_config(upstream_url),
}
def test_user_config_routes_to_the_user_supplied_deployment_when_opted_in(gateway: Gateway, tmp_path: Path) -> None:
with wire_server(_completion) as upstream:
config: Final = _opt_in_config(tmp_path, upstream.url)
with owned_proxy(gateway, tmp_path, {}, config=config) as candidate:
response: Final = candidate.request("POST", "/v1/chat/completions", _request_body(upstream.url))
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == "routed"
outbound: Final = tuple(upstream.received.get_nowait() for _ in range(upstream.received.qsize()))
completions: Final = tuple(request for request in outbound if request.target == "/v1/chat/completions")
assert len(completions) == 1, outbound
assert completions[0].headers["authorization"] == f"Bearer {USER_KEY}"
assert json.loads(completions[0].body)["model"] == "gpt-4.1-mini"
def test_user_config_is_rejected_without_the_opt_in(gateway: Gateway) -> None:
with wire_server(_completion) as upstream:
response: Final = gateway.request("POST", "/v1/chat/completions", _request_body(upstream.url))
assert response.status_code == 401, response.text
assert "user_config is not allowed in request body" in response.text
assert upstream.received.empty()

377
tests/mcp_tests/test_aresponses_api_with_mcp.py
View file

@ -1,5 +1,3 @@
import logging
import os
import pytest
from mcp.types import Tool as MCPTool
from typing import List, Any, cast
@ -846,161 +844,6 @@ async def test_streaming_mcp_events_validation():
assert mock_get_tools.called, "MCP tools should have been fetched"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model",
[
pytest.param("gpt-4o-mini", id="openai"),
pytest.param("claude-haiku-4-5", id="anthropic"),
],
)
async def test_streaming_responses_api_with_mcp_tools(
model: str, caplog: pytest.LogCaptureFixture
):
"""
Test the streaming responses API with MCP tools when using server_url="litellm_proxy"
Under the hood the follow occurs
- MCP: responses called litellm MCP manager.list_tools (MOCKED)
- Request 1: Made to model under test with fetched tools (REAL LLM CALL)
- MCP: Execute tool call from request 1 and returns result (MOCKED)
- Request 2: Made to model under test with fetched tools and tool results (REAL LLM CALL)
Return the user the result of request 2
"""
# Skip test if API keys are not set for the respective models
if ("claude" in model.lower() or "anthropic" in model.lower()) and not os.getenv(
"ANTHROPIC_API_KEY"
):
pytest.skip("ANTHROPIC_API_KEY not set, skipping anthropic model test")
if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv(
"OPENAI_API_KEY"
):
pytest.skip("OPENAI_API_KEY not set, skipping openai model test")
from unittest.mock import AsyncMock, patch
print("🧪 Testing basic streaming with MCP tools...")
# Mock MCP tools that would be returned from the manager
mock_mcp_tools = [
MCPTool.model_validate({
"name": "search_repo",
"description": "Search BerriAI/litellm repository for information",
"inputSchema": {
"type": "object",
"properties": {
"query": {"type": "string", "description": "Search query"}
},
"required": ["query"],
},
}, by_name=False)
]
# Only mock the MCP-specific operations, let LLM responses be real
with caplog.at_level(logging.ERROR):
with (
patch.object(
LiteLLM_Proxy_MCP_Handler,
"_get_mcp_tools_from_manager",
new_callable=AsyncMock,
) as mock_get_tools,
patch.object(
LiteLLM_Proxy_MCP_Handler,
"_execute_tool_calls",
new_callable=AsyncMock,
) as mock_execute_tools,
):
# Setup MCP mocks only
mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"])
# Create a dynamic mock that will match the actual tool call ID from the LLM response
def mock_execute_tool_calls_side_effect(
tool_calls, user_api_key_auth, **kwargs
):
"""Mock function that returns results matching the actual tool call IDs from the LLM"""
results = []
for tool_call in tool_calls:
# Extract call_id from the tool call
call_id = None
if isinstance(tool_call, dict):
call_id = tool_call.get("call_id") or tool_call.get("id")
elif hasattr(tool_call, "call_id"):
call_id = tool_call.call_id
elif hasattr(tool_call, "id"):
call_id = tool_call.id
if call_id:
results.append(
{
"tool_call_id": call_id,
"result": "LiteLLM is a unified interface for 100+ LLMs that translates inputs to provider-specific completion endpoints and provides consistent OpenAI-format output.",
}
)
return results
mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect
# Make the actual call - LLM responses will be real
mcp_tool_config = cast(
Any,
{
"type": "mcp",
"server_url": "litellm_proxy",
"require_approval": "never",
},
)
response = await litellm.aresponses(
model=model,
tools=[mcp_tool_config],
tool_choice="required",
input=[
{
"role": "user",
"type": "message",
"content": "give me a TLDR of what BerriAI/litellm is about",
}
],
stream=True,
)
print(f"📋 Response type: {type(response)}")
assert hasattr(
response, "__aiter__"
), "Response should be an async streaming response"
# Collect streaming chunks
chunks = []
async for chunk in response:
chunks.append(chunk)
print(f"📦 Chunk type: {getattr(chunk, 'type', 'unknown')}")
print(f"📊 Total chunks received: {len(chunks)}")
# Verify MCP mocks were called (may be called multiple times in streaming)
assert (
mock_get_tools.call_count >= 1
), f"Expected MCP tools to be fetched at least once, got {mock_get_tools.call_count}"
print(f"MCP tools fetched: {len(mock_mcp_tools)}")
# Verify we got a response
assert response is not None
assert len(chunks) > 0, "Should have received streaming chunks"
print("Basic streaming responses API with MCP tools test passed!")
lite_errors = [
record
for record in caplog.records
if record.levelno >= logging.ERROR
and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage())
]
assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join(
record.getMessage() for record in lite_errors
)
@pytest.mark.asyncio
async def test_mcp_parameter_preparation_helpers():
"""
@ -1215,7 +1058,7 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e():
The test mocks the MCP manager response but validates the actual tools
sent to the LLM to ensure no duplication occurs.
"""
from unittest.mock import AsyncMock, patch, call
from unittest.mock import AsyncMock, patch
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)
@ -1432,221 +1275,3 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e():
"tools_per_call": [len(tools) for tools in llm_call_tools],
"duplicate_tools_found": False,
}
@pytest.mark.asyncio
@pytest.mark.parametrize("model", ["gpt-4o-mini"])
async def test_streaming_mcp_event_order_and_response_id_consistency(
model: str, caplog: pytest.LogCaptureFixture
):
"""
Test that:
1. Streaming events are emitted in correct order (response.created, response.in_progress, response.output_item.added before MCP events)
2. All response lifecycle events share the same response ID within a cycle
"""
if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv(
"OPENAI_API_KEY"
):
pytest.skip("OPENAI_API_KEY not set, skipping openai model test")
from unittest.mock import AsyncMock, patch
mock_mcp_tools = [
MCPTool.model_validate({
"name": "get_weather",
"description": "Get weather for a city",
"inputSchema": {
"type": "object",
"properties": {
"city": {"type": "string", "description": "City name"}
},
"required": ["city"],
},
}, by_name=False)
]
with caplog.at_level(logging.ERROR):
with (
patch.object(
LiteLLM_Proxy_MCP_Handler,
"_get_mcp_tools_from_manager",
new_callable=AsyncMock,
) as mock_get_tools,
patch.object(
LiteLLM_Proxy_MCP_Handler,
"_execute_tool_calls",
new_callable=AsyncMock,
) as mock_execute_tools,
):
mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"])
def mock_execute_side_effect(tool_calls, user_api_key_auth, **kwargs):
results = []
for tool_call in tool_calls:
call_id = None
if isinstance(tool_call, dict):
call_id = tool_call.get("call_id") or tool_call.get("id")
elif hasattr(tool_call, "call_id"):
call_id = tool_call.call_id
elif hasattr(tool_call, "id"):
call_id = tool_call.id
if call_id:
results.append(
{
"tool_call_id": call_id,
"result": "Sunny, 72°F",
}
)
return results
mock_execute_tools.side_effect = mock_execute_side_effect
mcp_tool_config = cast(
Any,
{
"type": "mcp",
"server_url": "litellm_proxy",
"require_approval": "never",
},
)
response = await litellm.aresponses(
model=model,
tools=[mcp_tool_config],
input=[
{
"role": "user",
"type": "message",
"content": "What's the weather in San Francisco?",
}
],
stream=True,
)
events = []
async for chunk in response:
events.append(chunk)
assert len(events) > 0, "Should receive streaming events"
created_idx = next(
(
i
for i, e in enumerate(events)
if getattr(e, "type", None) == "response.created"
),
None,
)
in_progress_idx = next(
(
i
for i, e in enumerate(events)
if getattr(e, "type", None) == "response.in_progress"
),
None,
)
output_item_added_idx = next(
(
i
for i, e in enumerate(events)
if getattr(e, "type", None) == "response.output_item.added"
),
None,
)
mcp_in_progress_idx = next(
(
i
for i, e in enumerate(events)
if "mcp_list_tools.in_progress" in str(getattr(e, "type", ""))
),
None,
)
completed_idx = next(
(
i
for i, e in enumerate(events)
if getattr(e, "type", None) == "response.completed"
),
None,
)
assert created_idx is not None, "response.created event should be present"
assert (
in_progress_idx is not None
), "response.in_progress event should be present"
assert (
output_item_added_idx is not None
), "response.output_item.added event should be present"
assert (
created_idx < in_progress_idx
), "response.created should come before response.in_progress"
assert (
in_progress_idx < output_item_added_idx
), "response.in_progress should come before response.output_item.added"
if mcp_in_progress_idx is not None:
assert (
output_item_added_idx < mcp_in_progress_idx
), "response.output_item.added should come before response.mcp_list_tools.in_progress"
response_ids = []
for i, event in enumerate(events):
event_type = getattr(event, "type", None)
if hasattr(event, "response"):
response_obj = getattr(event, "response", None)
if response_obj and hasattr(response_obj, "id"):
event_type_value = (
event_type.value
if hasattr(event_type, "value")
else str(event_type)
)
if any(
x in event_type_value
for x in [
"response.created",
"response.in_progress",
"response.completed",
]
):
response_ids.append((i, event_type_value, response_obj.id))
assert (
len(response_ids) >= 2
), f"Should have at least 2 response lifecycle events. Found {len(response_ids)}"
cycles = []
current_cycle = []
current_id = None
for idx, event_type, resp_id in response_ids:
if current_id is None or resp_id == current_id:
current_cycle.append((idx, event_type, resp_id))
current_id = resp_id
else:
if current_cycle:
cycles.append(current_cycle)
current_cycle = [(idx, event_type, resp_id)]
current_id = resp_id
if current_cycle:
cycles.append(current_cycle)
for cycle_num, cycle in enumerate(cycles):
cycle_ids = set(resp_id for _, _, resp_id in cycle)
assert (
len(cycle_ids) == 1
), f"Cycle {cycle_num + 1} should have consistent response ID. Found {len(cycle_ids)} unique IDs"
assert (
completed_idx is not None
), "response.completed event should be present"
lite_errors = [
record
for record in caplog.records
if record.levelno >= logging.ERROR
and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage())
]
assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join(
record.getMessage() for record in lite_errors
)

389
tests/mcp_tests/test_aresponses_api_with_mcp_providers.py Normal file
View file

@ -0,0 +1,389 @@
import logging
import os
import pytest
from mcp.types import Tool as MCPTool
from typing import Any, cast
import litellm
from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler
class MockUserAPIKeyAuth:
"""Mock UserAPIKeyAuth for testing"""
def __init__(self):
self.api_key = "test_key"
self.user_id = "test_user"
self.team_id = "test_team"
self.user_email = "test@example.com"
self.max_budget = 100.0
self.spend = 0.0
self.models = []
self.aliases = {}
self.config = {}
self.permissions = {}
self.metadata = {}
self.object_permission_id = "test_permission_id"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model",
[
pytest.param("gpt-4o-mini", id="openai"),
pytest.param("claude-haiku-4-5", id="anthropic"),
],
)
async def test_streaming_responses_api_with_mcp_tools(
model: str, caplog: pytest.LogCaptureFixture
):
"""
Test the streaming responses API with MCP tools when using server_url="litellm_proxy"
Under the hood the follow occurs
- MCP: responses called litellm MCP manager.list_tools (MOCKED)
- Request 1: Made to model under test with fetched tools (REAL LLM CALL)
- MCP: Execute tool call from request 1 and returns result (MOCKED)
- Request 2: Made to model under test with fetched tools and tool results (REAL LLM CALL)
Return the user the result of request 2
"""
if ("claude" in model.lower() or "anthropic" in model.lower()) and not os.getenv(
"ANTHROPIC_API_KEY"
):
pytest.skip("ANTHROPIC_API_KEY not set, skipping anthropic model test")
if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv(
"OPENAI_API_KEY"
):
pytest.skip("OPENAI_API_KEY not set, skipping openai model test")
from unittest.mock import AsyncMock, patch
print("🧪 Testing basic streaming with MCP tools...")
mock_mcp_tools = [
MCPTool.model_validate({
"name": "search_repo",
"description": "Search BerriAI/litellm repository for information",
"inputSchema": {
"type": "object",
"properties": {
"query": {"type": "string", "description": "Search query"}
},
"required": ["query"],
},
}, by_name=False)
]
with caplog.at_level(logging.ERROR):
with (
patch.object(
LiteLLM_Proxy_MCP_Handler,
"_get_mcp_tools_from_manager",
new_callable=AsyncMock,
) as mock_get_tools,
patch.object(
LiteLLM_Proxy_MCP_Handler,
"_execute_tool_calls",
new_callable=AsyncMock,
) as mock_execute_tools,
):
mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"])
def mock_execute_tool_calls_side_effect(
tool_calls, user_api_key_auth, **kwargs
):
"""Mock function that returns results matching the actual tool call IDs from the LLM"""
results = []
for tool_call in tool_calls:
call_id = None
if isinstance(tool_call, dict):
call_id = tool_call.get("call_id") or tool_call.get("id")
elif hasattr(tool_call, "call_id"):
call_id = tool_call.call_id
elif hasattr(tool_call, "id"):
call_id = tool_call.id
if call_id:
results.append(
{
"tool_call_id": call_id,
"result": "LiteLLM is a unified interface for 100+ LLMs that translates inputs to provider-specific completion endpoints and provides consistent OpenAI-format output.",
}
)
return results
mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect
mcp_tool_config = cast(
Any,
{
"type": "mcp",
"server_url": "litellm_proxy",
"require_approval": "never",
},
)
response = await litellm.aresponses(
model=model,
tools=[mcp_tool_config],
tool_choice="required",
input=[
{
"role": "user",
"type": "message",
"content": "give me a TLDR of what BerriAI/litellm is about",
}
],
stream=True,
)
print(f"📋 Response type: {type(response)}")
assert hasattr(
response, "__aiter__"
), "Response should be an async streaming response"
chunks = []
async for chunk in response:
chunks.append(chunk)
print(f"📦 Chunk type: {getattr(chunk, 'type', 'unknown')}")
print(f"📊 Total chunks received: {len(chunks)}")
assert (
mock_get_tools.call_count >= 1
), f"Expected MCP tools to be fetched at least once, got {mock_get_tools.call_count}"
print(f"MCP tools fetched: {len(mock_mcp_tools)}")
assert response is not None
assert len(chunks) > 0, "Should have received streaming chunks"
print("Basic streaming responses API with MCP tools test passed!")
lite_errors = [
record
for record in caplog.records
if record.levelno >= logging.ERROR
and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage())
]
assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join(
record.getMessage() for record in lite_errors
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model", ["gpt-4o-mini"])
async def test_streaming_mcp_event_order_and_response_id_consistency(
model: str, caplog: pytest.LogCaptureFixture
):
"""
Test that:
1. Streaming events are emitted in correct order (response.created, response.in_progress, response.output_item.added before MCP events)
2. All response lifecycle events share the same response ID within a cycle
"""
if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv(
"OPENAI_API_KEY"
):
pytest.skip("OPENAI_API_KEY not set, skipping openai model test")
from unittest.mock import AsyncMock, patch
mock_mcp_tools = [
MCPTool.model_validate({
"name": "get_weather",
"description": "Get weather for a city",
"inputSchema": {
"type": "object",
"properties": {
"city": {"type": "string", "description": "City name"}
},
"required": ["city"],
},
}, by_name=False)
]
with caplog.at_level(logging.ERROR):
with (
patch.object(
LiteLLM_Proxy_MCP_Handler,
"_get_mcp_tools_from_manager",
new_callable=AsyncMock,
) as mock_get_tools,
patch.object(
LiteLLM_Proxy_MCP_Handler,
"_execute_tool_calls",
new_callable=AsyncMock,
) as mock_execute_tools,
):
mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"])
def mock_execute_side_effect(tool_calls, user_api_key_auth, **kwargs):
results = []
for tool_call in tool_calls:
call_id = None
if isinstance(tool_call, dict):
call_id = tool_call.get("call_id") or tool_call.get("id")
elif hasattr(tool_call, "call_id"):
call_id = tool_call.call_id
elif hasattr(tool_call, "id"):
call_id = tool_call.id
if call_id:
results.append(
{
"tool_call_id": call_id,
"result": "Sunny, 72°F",
}
)
return results
mock_execute_tools.side_effect = mock_execute_side_effect
mcp_tool_config = cast(
Any,
{
"type": "mcp",
"server_url": "litellm_proxy",
"require_approval": "never",
},
)
response = await litellm.aresponses(
model=model,
tools=[mcp_tool_config],
input=[
{
"role": "user",
"type": "message",
"content": "What's the weather in San Francisco?",
}
],
stream=True,
)
events = []
async for chunk in response:
events.append(chunk)
assert len(events) > 0, "Should receive streaming events"
created_idx = next(
(
i
for i, e in enumerate(events)
if getattr(e, "type", None) == "response.created"
),
None,
)
in_progress_idx = next(
(
i
for i, e in enumerate(events)
if getattr(e, "type", None) == "response.in_progress"
),
None,
)
output_item_added_idx = next(
(
i
for i, e in enumerate(events)
if getattr(e, "type", None) == "response.output_item.added"
),
None,
)
mcp_in_progress_idx = next(
(
i
for i, e in enumerate(events)
if "mcp_list_tools.in_progress" in str(getattr(e, "type", ""))
),
None,
)
completed_idx = next(
(
i
for i, e in enumerate(events)
if getattr(e, "type", None) == "response.completed"
),
None,
)
assert created_idx is not None, "response.created event should be present"
assert (
in_progress_idx is not None
), "response.in_progress event should be present"
assert (
output_item_added_idx is not None
), "response.output_item.added event should be present"
assert (
created_idx < in_progress_idx
), "response.created should come before response.in_progress"
assert (
in_progress_idx < output_item_added_idx
), "response.in_progress should come before response.output_item.added"
if mcp_in_progress_idx is not None:
assert (
output_item_added_idx < mcp_in_progress_idx
), "response.output_item.added should come before response.mcp_list_tools.in_progress"
response_ids = []
for i, event in enumerate(events):
event_type = getattr(event, "type", None)
if hasattr(event, "response"):
response_obj = getattr(event, "response", None)
if response_obj and hasattr(response_obj, "id"):
event_type_value = (
event_type.value
if hasattr(event_type, "value")
else str(event_type)
)
if any(
x in event_type_value
for x in [
"response.created",
"response.in_progress",
"response.completed",
]
):
response_ids.append((i, event_type_value, response_obj.id))
assert (
len(response_ids) >= 2
), f"Should have at least 2 response lifecycle events. Found {len(response_ids)}"
cycles = []
current_cycle = []
current_id = None
for idx, event_type, resp_id in response_ids:
if current_id is None or resp_id == current_id:
current_cycle.append((idx, event_type, resp_id))
current_id = resp_id
else:
if current_cycle:
cycles.append(current_cycle)
current_cycle = [(idx, event_type, resp_id)]
current_id = resp_id
if current_cycle:
cycles.append(current_cycle)
for cycle_num, cycle in enumerate(cycles):
cycle_ids = set(resp_id for _, _, resp_id in cycle)
assert (
len(cycle_ids) == 1
), f"Cycle {cycle_num + 1} should have consistent response ID. Found {len(cycle_ids)} unique IDs"
assert (
completed_idx is not None
), "response.completed event should be present"
lite_errors = [
record
for record in caplog.records
if record.levelno >= logging.ERROR
and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage())
]
assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join(
record.getMessage() for record in lite_errors
)

5
tests/proxy_unit_tests/test_proxy_custom_auth.py
View file

@ -53,8 +53,7 @@ def test_custom_auth(client):
"max_tokens": 10,
}
# Your bearer token
token = os.getenv("PROXY_MASTER_KEY")
print(f"token: {token}")
token = "sk-unit-test-master"
headers = {"Authorization": f"Bearer {token}"}
with pytest.raises(Exception, match="Authentication Error, Failed custom auth") as exc_info:
client.post("/chat/completions", json=test_data, headers=headers)
@ -71,7 +70,7 @@ def test_custom_auth_bearer(client):
"max_tokens": 10,
}
# Your bearer token
token = os.getenv("PROXY_MASTER_KEY")
token = "sk-unit-test-master"
headers = {"Authorization": f"WITHOUT BEAR Er {token}"}
with pytest.raises(Exception, match="CustomAuth - Malformed API Key passed in") as exc_info:

114
tests/proxy_unit_tests/test_proxy_pass_user_config.py
View file

@ -1,114 +0,0 @@
import sys, os
import traceback
from dotenv import load_dotenv
load_dotenv()
import io
# this file is to test litellm/proxy
import pytest, logging, asyncio
import litellm
from litellm import embedding, completion, completion_cost, Timeout
from litellm import RateLimitError
# Configure logging
logging.basicConfig(
level=logging.DEBUG, # Set the desired logging level
format="%(asctime)s - %(levelname)s - %(message)s",
)
# test /chat/completion request to the proxy
from fastapi.testclient import TestClient
from fastapi import FastAPI
from litellm.proxy.proxy_server import (
router,
save_worker_config,
initialize,
) # Replace with the actual module where your FastAPI router is defined
# Your bearer token
token = "sk-1234"
headers = {"Authorization": f"Bearer {token}"}
@pytest.fixture(scope="function")
def client_no_auth():
# Assuming litellm.proxy.proxy_server is an object
from litellm.proxy.proxy_server import cleanup_router_config_variables
cleanup_router_config_variables()
filepath = os.path.dirname(os.path.abspath(__file__))
config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
# initialize can get run in parallel, it sets specific variables for the fast api app, sinc eit gets run in parallel different tests use the wrong variables
asyncio.run(initialize(config=config_fp, debug=True))
app = FastAPI()
app.include_router(router) # Include your router in the test app
return TestClient(app)
@pytest.mark.skipif(
os.environ.get("AZURE_AI_API_KEY") is None
or os.environ.get("OPENAI_API_KEY") is None,
reason="AZURE_AI_API_KEY or OPENAI_API_KEY not set - skipping integration test",
)
def test_chat_completion(client_no_auth):
global headers
from litellm.types.router import RouterConfig, ModelConfig
from litellm.types.completion import CompletionRequest
user_config = RouterConfig(
model_list=[
ModelConfig(
model_name="user-azure-instance",
litellm_params=CompletionRequest(
model="azure/gpt-4.1-mini",
api_key=os.getenv("AZURE_AI_API_KEY"),
api_version=os.getenv("AZURE_API_VERSION"),
api_base=os.getenv("AZURE_AI_API_BASE"),
timeout=10,
),
tpm=240000,
rpm=1800,
),
ModelConfig(
model_name="user-openai-instance",
litellm_params=CompletionRequest(
model="gpt-3.5-turbo",
api_key=os.getenv("OPENAI_API_KEY"),
timeout=10,
),
tpm=240000,
rpm=1800,
),
],
num_retries=2,
allowed_fails=3,
fallbacks=[{"user-azure-instance": ["user-openai-instance"]}],
).dict()
try:
# Your test data
test_data = {
"model": "user-azure-instance",
"messages": [
{"role": "user", "content": "hi"},
],
"max_tokens": 10,
"user_config": user_config,
}
print("testing proxy server with chat completions")
response = client_no_auth.post("/v1/chat/completions", json=test_data)
print(f"response - {response.text}")
assert response.status_code == 200
result = response.json()
print(f"Received response: {result}")
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
# Run the test

53
tests/proxy_unit_tests/test_proxy_server.py
View file

@ -2065,59 +2065,6 @@ async def test_add_callback_via_key_litellm_pre_call_utils_langsmith(
assert new_data["failure_callback"] == expected_failure_callbacks
@pytest.mark.skipif(
not os.getenv("GEMINI_API_KEY") and not os.getenv("GOOGLE_API_KEY"),
reason="Requires GEMINI_API_KEY or GOOGLE_API_KEY.",
)
@pytest.mark.asyncio
async def test_gemini_pass_through_endpoint():
from starlette.datastructures import URL
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
Request,
Response,
gemini_proxy_route,
)
body = b"""
{
"contents": [{
"parts":[{
"text": "The quick brown fox jumps over the lazy dog."
}]
}]
}
"""
# Construct the scope dictionary
scope = {
"type": "http",
"method": "POST",
"path": "/gemini/v1beta/models/gemini-2.5-flash:countTokens",
"query_string": b"key=sk-1234",
"headers": [
(b"content-type", b"application/json"),
],
}
# Create a new Request object
async def async_receive():
return {"type": "http.request", "body": body, "more_body": False}
request = Request(
scope=scope,
receive=async_receive,
)
resp = await gemini_proxy_route(
endpoint="v1beta/models/gemini-2.5-flash:countTokens?key=sk-1234",
request=request,
fastapi_response=Response(),
)
print(resp.body)
@pytest.mark.parametrize("hidden", [True, False])
@pytest.mark.asyncio
async def test_model_info_alias_without_prisma(hidden):

51
tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py Normal file
View file

@ -0,0 +1,51 @@
import os
import pytest
@pytest.mark.skipif(
not os.getenv("GEMINI_API_KEY") and not os.getenv("GOOGLE_API_KEY"),
reason="Requires GEMINI_API_KEY or GOOGLE_API_KEY.",
)
@pytest.mark.asyncio
async def test_gemini_pass_through_endpoint():
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
Request,
Response,
gemini_proxy_route,
)
body = b"""
{
"contents": [{
"parts":[{
"text": "The quick brown fox jumps over the lazy dog."
}]
}]
}
"""
scope = {
"type": "http",
"method": "POST",
"path": "/gemini/v1beta/models/gemini-2.5-flash:countTokens",
"query_string": b"key=sk-1234",
"headers": [
(b"content-type", b"application/json"),
],
}
async def async_receive():
return {"type": "http.request", "body": body, "more_body": False}
request = Request(
scope=scope,
receive=async_receive,
)
await gemini_proxy_route(
endpoint="v1beta/models/gemini-2.5-flash:countTokens?key=sk-1234",
request=request,
fastapi_response=Response(),
)

159
tests/proxy_unit_tests/test_proxy_token_counter.py
View file

@ -2,10 +2,7 @@
# 1. Generate a Key, and use it to make a call
import json
import logging
import os
import tempfile
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -35,79 +32,6 @@ from litellm.types.utils import TokenCountResponse
verbose_proxy_logger.setLevel(level=logging.DEBUG)
def get_vertex_ai_creds_json() -> dict:
# Define the path to the vertex_key.json file
print("loading vertex ai credentials")
filepath = os.path.dirname(os.path.abspath(__file__))
vertex_key_path = filepath + "/vertex_key.json"
# Read the existing content of the file or create an empty dictionary
try:
with open(vertex_key_path, "r") as file:
# Read the file content
print("Read vertexai file path")
content = file.read()
# If the file is empty or not valid JSON, create an empty dictionary
if not content or not content.strip():
service_account_key_data = {}
else:
# Attempt to load the existing JSON content
file.seek(0)
service_account_key_data = json.load(file)
except FileNotFoundError:
# If the file doesn't exist, create an empty dictionary
service_account_key_data = {}
# Update the service_account_key_data with environment variables
private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "")
private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "")
private_key = private_key.replace("\\n", "\n")
service_account_key_data["private_key_id"] = private_key_id
service_account_key_data["private_key"] = private_key
return service_account_key_data
def load_vertex_ai_credentials():
# Define the path to the vertex_key.json file
print("loading vertex ai credentials")
filepath = os.path.dirname(os.path.abspath(__file__))
vertex_key_path = filepath + "/vertex_key.json"
# Read the existing content of the file or create an empty dictionary
try:
with open(vertex_key_path, "r") as file:
# Read the file content
print("Read vertexai file path")
content = file.read()
# If the file is empty or not valid JSON, create an empty dictionary
if not content or not content.strip():
service_account_key_data = {}
else:
# Attempt to load the existing JSON content
file.seek(0)
service_account_key_data = json.load(file)
except FileNotFoundError:
# If the file doesn't exist, create an empty dictionary
service_account_key_data = {}
# Update the service_account_key_data with environment variables
private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "")
private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "")
private_key = private_key.replace("\\n", "\n")
service_account_key_data["private_key_id"] = private_key_id
service_account_key_data["private_key"] = private_key
# Create a temporary file
with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file:
# Write the updated content to the temporary files
json.dump(service_account_key_data, temp_file, indent=2)
# Export the temporary file as GOOGLE_APPLICATION_CREDENTIALS
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name)
@pytest.mark.asyncio
async def test_vLLM_token_counting():
"""
@ -223,10 +147,12 @@ async def test_anthropic_messages_count_tokens_endpoint():
- Should return response in Anthropic format: {"input_tokens": <count>}
- Should work as wrapper around internal token_counter function
"""
from litellm.proxy.anthropic_endpoints.endpoints import count_tokens
from fastapi import Request
from unittest.mock import MagicMock
from fastapi import Request
from litellm.proxy.anthropic_endpoints.endpoints import count_tokens
# Mock request object
mock_request = MagicMock(spec=Request)
mock_request_data = {
@ -295,10 +221,12 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model():
- Should still work and return Anthropic format
- Should call internal token_counter with from_anthropic_endpoint=True
"""
from litellm.proxy.anthropic_endpoints.endpoints import count_tokens
from fastapi import Request
from unittest.mock import MagicMock
from fastapi import Request
from litellm.proxy.anthropic_endpoints.endpoints import count_tokens
# Mock request object
mock_request = MagicMock(spec=Request)
mock_request_data = {
@ -435,10 +363,12 @@ async def test_anthropic_endpoint_error_handling():
"""
Test error handling in the /v1/messages/count_tokens endpoint
"""
from litellm.proxy.anthropic_endpoints.endpoints import count_tokens
from fastapi import Request, HTTPException
from unittest.mock import MagicMock
from fastapi import HTTPException, Request
from litellm.proxy.anthropic_endpoints.endpoints import count_tokens
# Mock request object
mock_request = MagicMock(spec=Request)
mock_user_api_key_dict = MagicMock()
@ -474,8 +404,10 @@ async def test_anthropic_endpoint_error_handling():
@pytest.mark.asyncio
async def test_factory_anthropic_endpoint_calls_anthropic_counter():
"""Test that /v1/messages/count_tokens with Anthropic model uses Anthropic counter."""
from unittest.mock import patch, AsyncMock, MagicMock
from unittest.mock import AsyncMock, MagicMock, patch
from fastapi.testclient import TestClient
from litellm.proxy.proxy_server import app
# Mock the global handler instance in token_counter module
@ -531,8 +463,10 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter():
@pytest.mark.asyncio
async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter():
"""Test that /v1/messages/count_tokens with GPT-4 does NOT use Anthropic counter."""
from unittest.mock import patch, AsyncMock, MagicMock
from unittest.mock import AsyncMock, MagicMock, patch
from fastapi.testclient import TestClient
from litellm.proxy.proxy_server import app
# Mock the global handler instance in token_counter module
@ -590,8 +524,10 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter():
@pytest.mark.asyncio
async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic():
"""Test that /utils/token_counter does NOT use Anthropic counter even with Anthropic model."""
from unittest.mock import patch, AsyncMock, MagicMock
from unittest.mock import AsyncMock, MagicMock, patch
from fastapi.testclient import TestClient
from litellm.proxy.proxy_server import app
# Mock the global handler instance in token_counter module
@ -678,57 +614,6 @@ async def test_factory_registration():
assert not counter.should_use_token_counting_api(custom_llm_provider=None)
@pytest.mark.skip(
reason="Requires Google/Vertex AI credentials (GEMINI_API_KEY or VERTEX_AI_PRIVATE_KEY)."
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", ["gemini-2.5-pro", "vertex-ai-gemini-2.5-pro"])
async def test_vertex_ai_gemini_token_counting_with_contents(model_name):
"""
Test token counting for Vertex AI Gemini model using contents format with call_endpoint=True
"""
load_vertex_ai_credentials()
llm_router = Router(
model_list=[
{
"model_name": "gemini-2.5-pro",
"litellm_params": {
"model": "gemini/gemini-2.5-pro",
},
},
{
"model_name": "vertex-ai-gemini-2.5-pro",
"litellm_params": {
"model": "vertex_ai/gemini-2.5-pro",
},
},
]
)
setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
# Test with contents format and call_endpoint=True
response = await token_counter(
request=TokenCountRequest(
model=model_name,
contents=[
{"parts": [{"text": "Hello world, how are you doing today? i am ij"}]}
],
),
call_endpoint=True,
)
print("Vertex AI Gemini token counting response:", response)
# validate we have original response
assert response.original_response is not None
assert response.original_response.get("totalTokens") is not None
assert response.original_response.get("promptTokensDetails") is not None
prompt_tokens_details = response.original_response.get("promptTokensDetails")
assert prompt_tokens_details is not None
@pytest.mark.asyncio
async def test_bedrock_count_tokens_endpoint():
"""
@ -779,7 +664,7 @@ async def test_vertex_ai_anthropic_token_counting():
This tests the token counting implementation for Vertex AI partner models
without making actual API calls. Mocks at the handler level to test the full flow.
"""
from unittest.mock import AsyncMock, patch, MagicMock
from unittest.mock import patch
# Mock the Vertex AI partner models token counter response
mock_token_response = {

4
tests/proxy_unit_tests/test_proxy_utils.py
View file

@ -438,7 +438,7 @@ def test_is_request_body_safe_global_enabled(
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
"api_key": "sk-openai-unit-test",
},
}
]
@ -475,7 +475,7 @@ def test_is_request_body_safe_model_enabled(
"model_name": "fireworks_ai/*",
"litellm_params": {
"model": "fireworks_ai/*",
"api_key": os.getenv("FIREWORKS_API_KEY"),
"api_key": "sk-fireworks-unit-test",
"configurable_clientside_auth_params": (
["api_base"] if allow_client_side_credentials else []
),