fix(tests): Add missing mocks for MCP IP filtering and updated APIs#20652
fix(tests): Add missing mocks for MCP IP filtering and updated APIs#20652ishaan-jaff merged 1 commit intomainfrom
Conversation
Fixes 15 failing tests in the MCP test suite: 1. **OAuth discoverable endpoints** (test_discoverable_endpoints.py): - Added autouse fixture to mock IPAddressUtils.get_mcp_client_ip - This bypasses IP-based access control which was blocking server lookup - Fixes: test_authorize_*, test_token_*, test_oauth_*, test_register_* 2. **A2A endpoints** (test_a2a_endpoints.py): - Fixed mock path for add_litellm_data_to_request - Was patching litellm_pre_call_utils but function is called from common_request_processing 3. **MCP guardrail handler** (test_mcp_guardrail_handler.py): - Updated tests to match new handler behavior - Handler now passes tools (not texts) to guardrail - Handler checks for mcp_tool_name (not messages array) 4. **MCP path-based segregation** (test_user_api_key_auth_mcp.py): - Added client_ip to get_auth_context unpacking (7 values now) - get_auth_context was updated to include client_ip 5. **MCP registry** (test_mcp_management_endpoints.py): - Added mock for get_filtered_registry (not just get_registry) - Registry endpoint uses get_filtered_registry for IP filtering
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
Shin seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account. You have signed the CLA already but the status is still pending? Let us recheck it. |
Greptile OverviewGreptile SummaryThis PR updates the MCP-related test suite to match the recent IP-based access control change and a few API surface shifts. It:
Overall, these are test-only adjustments intended to restore stability after the MCP IP filtering feature landed, with the main concern being that the new IP mock currently patches a single call site rather than the source symbol used across modules. Confidence Score: 4/5
|
| Filename | Overview |
|---|---|
| tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py | Updates MCP path-based segregation test to unpack new client_ip value from get_auth_context(). |
| tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py | Reworks guardrail handler tests to assert tool-based inputs; drops coverage of guardrail text-modification behavior. |
| tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py | Adds autouse fixture patching get_mcp_client_ip at a single import path; may not bypass IP filtering if other modules call the source function directly. |
| tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py | Fixes mock patch target for add_litellm_data_to_request to match actual call site in common_request_processing. |
| tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py | Adds missing mock for get_filtered_registry to match management endpoint behavior that filters registry by client IP. |
Sequence Diagram
sequenceDiagram
participant T as Tests
participant DE as discoverable_endpoints
participant MEP as mcp_management_endpoints
participant IPU as IPAddressUtils
participant MSM as MCPServerManager
T->>DE: authorize()/token_endpoint()/register_client()
DE->>IPU: get_mcp_client_ip(request)
Note over T,IPU: PR adds autouse patch returning None (in DE module)
DE->>DE: build redirect/perform token exchange
T->>MEP: GET /v1/mcp/registry.json
MEP->>IPU: get_mcp_client_ip(request)
MEP->>MSM: get_filtered_registry(client_ip)
MSM-->>MEP: filtered servers
MEP-->>T: {servers: builtin + filtered entries}
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| @pytest.fixture(autouse=True) | ||
| def mock_mcp_client_ip(): | ||
| """Mock IPAddressUtils.get_mcp_client_ip to return None for all tests. | ||
|
|
||
| This bypasses IP-based access control in tests, since the MCP server's | ||
| available_on_public_internet defaults to False and mock requests don't | ||
| have proper client IP context. | ||
| """ | ||
| with patch( | ||
| "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.get_mcp_client_ip", | ||
| return_value=None, | ||
| ): | ||
| yield |
There was a problem hiding this comment.
Autouse patch targets wrong symbol
This fixture patches IPAddressUtils.get_mcp_client_ip only at litellm.proxy._experimental.mcp_server.discoverable_endpoints..., so any code under test that calls the source litellm.proxy.auth.ip_address_utils.IPAddressUtils.get_mcp_client_ip (or other modules’ imported references) won’t be affected. That makes the “bypass IP-based access control in tests” brittle across refactors and can re-break tests.
Consider patching the source symbol (litellm.proxy.auth.ip_address_utils.IPAddressUtils.get_mcp_client_ip) instead of a single call site.
| async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): | ||
| self.call_count += 1 | ||
| self.last_inputs = inputs | ||
|
|
||
| if self.return_texts is not None: | ||
| return {"texts": self.return_texts} | ||
|
|
||
| texts = inputs.get("texts", []) | ||
| return {"texts": [f"{text} [SAFE]" for text in texts]} | ||
| self.last_request_data = request_data | ||
| return None # Guardrail doesn't modify for MCP tools |
There was a problem hiding this comment.
Lost behavior coverage
MockGuardrail.apply_guardrail() now always returns None and the tests only assert the handler returns data unchanged, which removes coverage for the handler’s behavior when a guardrail returns a modification payload. If MCPGuardrailTranslationHandler.process_input_messages() still supports applying updates from guardrail output, it’d be good to keep at least one test asserting those updates are handled (adapted to the new tool-based shape).
3 participants
Summary
Fixes 15 failing tests in the MCP test suite after the IP-based access control feature was added (#20620, #20607).
Failing Tests Fixed
OAuth Discoverable Endpoints (
test_discoverable_endpoints.py)test_authorize_endpoint_forwards_pkce_parameterstest_authorize_endpoint_includes_response_typetest_authorize_endpoint_respects_x_forwarded_hosttest_authorize_endpoint_respects_x_forwarded_prototest_oauth_authorization_server_respects_x_forwarded_prototest_oauth_protected_resource_respects_x_forwarded_prototest_register_client_remote_registration_successtest_token_endpoint_forwards_code_verifiertest_token_endpoint_respects_x_forwarded_hosttest_token_endpoint_respects_x_forwarded_protoFix: Added
autousefixture to mockIPAddressUtils.get_mcp_client_ipto returnNone, bypassing IP-based access control in tests.A2A Endpoints (
test_a2a_endpoints.py)test_invoke_agent_a2a_adds_litellm_dataFix: Changed mock path from
litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_requesttolitellm.proxy.common_request_processing.add_litellm_data_to_request(where it's actually called).MCP Guardrail Handler (
test_mcp_guardrail_handler.py)test_process_input_messages_skips_when_no_messagestest_process_input_messages_updates_contentFix: Updated tests to match new handler behavior - handler now checks for
mcp_tool_name(notmessagesarray) and passestools(nottexts) to guardrail.MCP Path-Based Segregation (
test_user_api_key_auth_mcp.py)test_mcp_path_based_server_segregationFix: Added
client_iptoget_auth_context()unpacking (now returns 7 values instead of 6).MCP Registry (
test_mcp_management_endpoints.py)test_registry_returns_entries_when_enabledFix: Added mock for
get_filtered_registry()(not justget_registry()) since the registry endpoint usesget_filtered_registryfor IP filtering.Root Cause
The IP-based access control feature (#20620) added a new
client_ipparameter throughout the MCP codebase. Tests were not updated to:get_auth_context()get_filtered_registry()instead of justget_registry()Tests verified locally - all 15 previously failing tests now pass