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

test: add test

Browse Source
This commit is contained in:
Yuta Saito
2026-01-21 14:51:56 +09:00
parent 872e5b9897
commit caf5f7f8ae
2 changed files with 283 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+135
View File
@@ -1817,3 +1817,138 @@ class TestMCPServerManagerReload:
mock_get_all.assert_awaited_once()
mock_build.assert_awaited_once_with(db_row)
assert manager.registry["server-1"] is rebuilt_server
@pytest.mark.asyncio
async def test_call_mcp_tool_logs_failure_via_post_call_failure_hook():
"""
Regression test for 6267f168...:
Ensure proxy-side `call_mcp_tool` logs failures via `proxy_logging_obj.post_call_failure_hook`.
"""
try:
from litellm.proxy._experimental.mcp_server.server import (
call_mcp_tool,
global_mcp_server_manager,
)
from litellm.proxy._types import MCPTransport, UserAPIKeyAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
except ImportError:
pytest.skip("MCP server not available")
mock_server = MCPServer(
server_id="server-123",
name="test_server",
alias="test_server",
server_name="test_server",
url="https://test-server.com/mcp",
transport=MCPTransport.http,
mcp_info={"server_name": "test_server"},
)
proxy_logging_mock = MagicMock()
proxy_logging_mock.post_call_failure_hook = AsyncMock()
user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user")
with patch.object(
global_mcp_server_manager,
"get_allowed_mcp_servers",
new_callable=AsyncMock,
return_value=[mock_server.server_id],
), patch.object(
global_mcp_server_manager,
"get_mcp_server_by_id",
return_value=mock_server,
), patch(
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers_from_mcp_server_names",
new_callable=AsyncMock,
return_value=[mock_server],
), patch(
"litellm.proxy._experimental.mcp_server.server.execute_mcp_tool",
new_callable=AsyncMock,
side_effect=Exception("boom"),
), patch(
"litellm.proxy.proxy_server.proxy_logging_obj",
proxy_logging_mock,
):
with pytest.raises(Exception):
await call_mcp_tool(
name="test_server-any_tool",
arguments={"x": 1},
user_api_key_auth=user_auth,
litellm_call_id="cid",
)
proxy_logging_mock.post_call_failure_hook.assert_awaited_once()
assert (
proxy_logging_mock.post_call_failure_hook.await_args.kwargs.get("route")
== "/mcp/call_tool"
)
@pytest.mark.asyncio
async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enabled():
"""
Regression test for 872e5b98...:
Ensure list-tools logging path calls `async_success_handler` when enabled.
"""
try:
from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers
from litellm.proxy._types import UserAPIKeyAuth
except ImportError:
pytest.skip("MCP server not available")
user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user")
server_a = MagicMock(name="server_a_obj")
server_a.name = "server_a"
server_a.alias = "server_a"
server_a.server_name = "server_a"
server_a.server_id = "a"
server_a.auth_type = None
server_a.extra_headers = None
tool_1 = MagicMock()
tool_1.name = "server_a-tool_1"
dummy_logging_obj = MagicMock()
dummy_logging_obj.model_call_details = {"metadata": {"spend_logs_metadata": {}}}
dummy_logging_obj.async_success_handler = AsyncMock()
with patch(
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server_a]),
), patch(
"litellm.proxy._experimental.mcp_server.server._prepare_mcp_server_headers",
return_value=(None, None),
), patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
) as mock_manager, patch(
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools",
side_effect=lambda tools, _server: tools,
), patch(
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions",
new=AsyncMock(side_effect=lambda tools, **_: tools),
), patch(
"litellm.proxy._experimental.mcp_server.server.function_setup",
return_value=(dummy_logging_obj, None),
):
mock_manager._get_tools_from_server = AsyncMock(return_value=[tool_1])
tools = await _get_tools_from_mcp_servers(
user_api_key_auth=user_auth,
mcp_auth_header=None,
mcp_servers=["server_a"],
mcp_server_auth_headers=None,
log_list_tools_to_spendlogs=True,
list_tools_log_source="mcp_protocol",
)
assert tools == [tool_1]
dummy_logging_obj.async_success_handler.assert_awaited_once()
assert dummy_logging_obj.async_success_handler.await_args.kwargs["result"] == [tool_1]
spend_meta = dummy_logging_obj.model_call_details["metadata"]["spend_logs_metadata"]
assert spend_meta["tool_count_total"] == 1
assert spend_meta["allowed_server_count"] == 1
assert spend_meta["per_server_tool_counts"]["server_a"] == 1
tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
+148 -3
View File
@@ -1,12 +1,15 @@
import sys
import types
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import HTTPException
import importlib
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)
from typing import Any, cast
from litellm.types.utils import ModelResponse
from litellm.types.responses.main import OutputFunctionToolCall
@@ -22,7 +25,9 @@ def _setup_mcp_call_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock:
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module)
fake_manager = types.SimpleNamespace(
call_tool=AsyncMock(return_value=_DummyMCPResult())
call_tool=AsyncMock(return_value=_DummyMCPResult()),
# Newer logging path calls this to enrich spend logs metadata
_get_mcp_server_from_tool_name=MagicMock(return_value=None),
)
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
@@ -31,6 +36,15 @@ def _setup_mcp_call_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock:
return fake_manager.call_tool
def _setup_proxy_logging(monkeypatch: pytest.MonkeyPatch) -> AsyncMock:
"""Patch proxy_logging_obj so failure hook can be asserted."""
proxy_logging_obj = MagicMock()
proxy_logging_obj.post_call_failure_hook = AsyncMock()
proxy_module = types.SimpleNamespace(proxy_logging_obj=proxy_logging_obj)
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module)
return proxy_logging_obj.post_call_failure_hook
def test_deduplicate_mcp_tools_single_allowed_server():
tools = [{"name": "search"}, {"name": "search"}] # duplicate on purpose
@@ -184,7 +198,7 @@ def test_create_follow_up_input_handles_response_function_tool_call():
)
follow_up = LiteLLM_Proxy_MCP_Handler._create_follow_up_input(
response=response,
response=cast(Any, response),
tool_results=[],
original_input=None,
)
@@ -216,6 +230,8 @@ async def test_execute_tool_calls_strips_server_prefix(monkeypatch):
user_api_key_auth=None,
)
assert call_tool_mock.await_count == 1
assert call_tool_mock.await_args is not None
assert call_tool_mock.await_args.kwargs["name"] == "read_wiki_structure"
@@ -236,6 +252,8 @@ async def test_execute_tool_calls_keeps_tool_name_without_prefix(monkeypatch):
user_api_key_auth=None,
)
assert call_tool_mock.await_count == 1
assert call_tool_mock.await_args is not None
assert call_tool_mock.await_args.kwargs["name"] == tool_name
@@ -256,4 +274,131 @@ async def test_execute_tool_calls_keeps_tool_name_when_equal_to_server(monkeypat
user_api_key_auth=None,
)
assert call_tool_mock.await_count == 1
assert call_tool_mock.await_args is not None
assert call_tool_mock.await_args.kwargs["name"] == tool_name
@pytest.mark.asyncio
async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkeypatch):
"""
Regression test for ae4d92ad...:
Ensure responses-side MCP tool execution logs failures via proxy_logging_obj.post_call_failure_hook.
"""
post_call_failure_hook = _setup_proxy_logging(monkeypatch)
fake_manager = types.SimpleNamespace(
call_tool=AsyncMock(
side_effect=HTTPException(status_code=500, detail="boom")
)
)
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
fake_manager,
)
tool_name = "deepwiki-read_wiki_structure"
tool_calls = [
{"id": "call-err", "function": {"name": tool_name, "arguments": "{}"}}
]
user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user")
results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(
tool_server_map={tool_name: "deepwiki"},
tool_calls=tool_calls,
user_api_key_auth=user_auth,
litellm_call_id="cid",
litellm_trace_id="tid",
)
assert len(results) == 1
assert results[0]["tool_call_id"] == "call-err"
assert results[0]["name"] == tool_name
post_call_failure_hook.assert_awaited_once()
assert post_call_failure_hook.await_args is not None
assert (
post_call_failure_hook.await_args.kwargs.get("route")
== "/responses/mcp/call_tool"
)
@pytest.mark.asyncio
async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_function_setup(
monkeypatch,
):
"""
Regression test for ae4d92ad...:
Ensure litellm_call_id / litellm_trace_id are forwarded into function_setup kwargs.
"""
_setup_proxy_logging(monkeypatch)
call_tool_mock = _setup_mcp_call_environment(monkeypatch)
captured = {}
def fake_function_setup(*_args, **kwargs):
captured.update(kwargs)
return None, None
# NOTE: Don't patch via dotted string path here because `litellm.responses`
# is a function attribute on the `litellm` package (shadowing the submodule),
# which breaks monkeypatch's importpath resolution.
handler_module = importlib.import_module(
"litellm.responses.mcp.litellm_proxy_mcp_handler"
)
monkeypatch.setattr(handler_module, "function_setup", fake_function_setup)
tool_name = "deepwiki-read_wiki_structure"
tool_calls = [
{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}
]
await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(
tool_server_map={tool_name: "deepwiki"},
tool_calls=tool_calls,
user_api_key_auth=None,
litellm_call_id="cid",
litellm_trace_id="tid",
)
# Ensure the tool call was attempted (sanity)
assert call_tool_mock.await_count == 1
assert captured.get("litellm_call_id") == "cid"
assert captured.get("litellm_trace_id") == "tid"
@pytest.mark.asyncio
async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch):
"""
Regression test for 872e5b98...:
Ensure responses-side tool discovery enables list-tools SpendLogs logging flags.
"""
mock_get_tools = AsyncMock(return_value=[])
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.server._get_tools_from_mcp_servers",
mock_get_tools,
)
# Patch manager methods used by _get_mcp_tools_from_manager to avoid needing full UserAPIKeyAuth fields.
fake_manager = types.SimpleNamespace(
get_allowed_mcp_servers=AsyncMock(return_value=[]),
get_mcp_servers_from_ids=MagicMock(return_value=[]),
)
monkeypatch.setattr(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
fake_manager,
)
user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user")
tools, _server_names = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager(
user_api_key_auth=user_auth,
mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}],
)
assert tools == []
assert mock_get_tools.await_count == 1
assert mock_get_tools.await_args is not None
assert mock_get_tools.await_args.kwargs["log_list_tools_to_spendlogs"] is True
assert mock_get_tools.await_args.kwargs["list_tools_log_source"] == "responses"