From 93053cc9870aa337095c119546d3e42ebeb6e35d Mon Sep 17 00:00:00 2001 From: Alejandro Ponce Date: Thu, 23 Oct 2025 13:31:46 +0300 Subject: [PATCH 1/2] Remove URL manipulation in MCPServerClient MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Removed _prepare_url_for_proxy_mode method that appended /mcp or /sse paths - Workload URLs now passed directly to underlying MCP clients without modification - Eliminated temporary URL override logic in _execute_with_session - Added 7 unit tests verifying URL remains unchanged during: - MCPServerClient initialization - list_tools operations (both streamable-http and SSE) - call_tool operations - Multiple consecutive operations - All tests verify both the workload URL property and the URL passed to MCP clients 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- src/mcp_optimizer/mcp_client.py | 33 ----- tests/test_mcp_client.py | 246 ++++++++++++++++++++++++++++++++ 2 files changed, 246 insertions(+), 33 deletions(-) diff --git a/src/mcp_optimizer/mcp_client.py b/src/mcp_optimizer/mcp_client.py index 6dfbd12..af8d3f9 100644 --- a/src/mcp_optimizer/mcp_client.py +++ b/src/mcp_optimizer/mcp_client.py @@ -106,29 +106,6 @@ def _determine_proxy_mode(self) -> ToolHiveProxyMode: ) return url_to_toolhive_proxy_mode(self.workload.url) - def _prepare_url_for_proxy_mode(self, proxy_mode: ToolHiveProxyMode) -> str: - """ - Prepare the URL with the correct path for the transport type. - - Args: - proxy_mode: The proxy mode being used - - Returns: - The URL with the correct path appended - """ - url = self.workload.url - if proxy_mode == ToolHiveProxyMode.STREAMABLE: - # Streamable HTTP expects /mcp path - if not url.endswith("/mcp"): - url = url.rstrip("/") + "/mcp" - logger.debug(f"Appended /mcp path to URL: {url}", workload=self.workload.name) - elif proxy_mode == ToolHiveProxyMode.SSE: - # SSE expects /sse path - if not url.endswith("/sse"): - url = url.rstrip("/") + "/sse" - logger.debug(f"Appended /sse path to URL: {url}", workload=self.workload.name) - return url - async def _execute_with_session(self, operation: Callable[[ClientSession], Awaitable]) -> Any: """ Execute an operation with an MCP session. @@ -153,13 +130,6 @@ async def _execute_with_session(self, operation: Callable[[ClientSession], Await proxy_mode_field=self.workload.proxy_mode, ) - # Prepare URL with correct path for the transport type - url = self._prepare_url_for_proxy_mode(proxy_mode) - - # Temporarily override the URL for this connection - original_url = self.workload.url - self.workload.url = url - try: if proxy_mode == ToolHiveProxyMode.STREAMABLE: return await asyncio.wait_for( @@ -197,9 +167,6 @@ async def _execute_with_session(self, operation: Callable[[ClientSession], Await error=str(e), ) raise WorkloadConnectionError(f"MCP protocol error: {e}") from e - finally: - # Restore original URL - self.workload.url = original_url async def _execute_streamable_session( self, operation: Callable[[ClientSession], Awaitable] diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index 062701d..53dc3a0 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -236,3 +236,249 @@ def test_extract_error_from_exception_group(): result = client._extract_error_from_exception_group(eg) assert "ValueError" in result assert "Some error" in result + + +def test_workload_url_unchanged_after_init(): + """Test that workload URL is not modified during MCPServerClient initialization.""" + original_url = "http://localhost:8080/sse/test-server" + workload = Workload( + name="test-server", + url=original_url, + status="running", + tool_type="mcp", + ) + + # Create client + _client = MCPServerClient(workload, timeout=10) + + # Verify URL is unchanged + assert workload.url == original_url + + +def test_workload_url_unchanged_with_proxy_mode(): + """Test that workload URL is not modified when proxy_mode is set.""" + original_url = "http://localhost:8080/mcp/test-server" + workload = Workload( + name="test-server", + url=original_url, + proxy_mode="streamable-http", + status="running", + tool_type="mcp", + ) + + # Create client + _client = MCPServerClient(workload, timeout=10) + + # Verify URL is unchanged + assert workload.url == original_url + + +@pytest.mark.asyncio +async def test_workload_url_unchanged_during_list_tools_streamable(): + """Test that workload URL remains unchanged during list_tools with streamable-http.""" + original_url = "http://localhost:8080/mcp/test-server" + workload = Workload( + name="test-server", + url=original_url, + status="running", + tool_type="mcp", + ) + + client = MCPServerClient(workload, timeout=10) + + # Mock the MCP client session + mock_result = AsyncMock() + mock_result.tools = [] + + mock_session = AsyncMock() + mock_session.list_tools.return_value = mock_result + + with ( + patch("mcp_optimizer.mcp_client.streamablehttp_client") as mock_client, + patch( + "mcp_optimizer.mcp_client.ClientSession", return_value=mock_session + ) as mock_session_class, + ): + # Mock the context manager + mock_client.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock(), AsyncMock()) + mock_session_class.return_value.__aenter__.return_value = mock_session + + # Call list_tools + await client.list_tools() + + # Verify URL is unchanged in workload + assert workload.url == original_url + + # Verify the client was called with the original URL + mock_client.assert_called_once_with(original_url) + + +@pytest.mark.asyncio +async def test_workload_url_unchanged_during_list_tools_sse(): + """Test that workload URL remains unchanged during list_tools with SSE.""" + original_url = "http://localhost:8080/sse/test-server" + workload = Workload( + name="test-server", + url=original_url, + status="running", + tool_type="mcp", + ) + + client = MCPServerClient(workload, timeout=10) + + # Mock the MCP client session + mock_result = AsyncMock() + mock_result.tools = [] + + mock_session = AsyncMock() + mock_session.list_tools.return_value = mock_result + + with ( + patch("mcp_optimizer.mcp_client.sse_client") as mock_client, + patch( + "mcp_optimizer.mcp_client.ClientSession", return_value=mock_session + ) as mock_session_class, + ): + # Mock the context manager + mock_client.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock()) + mock_session_class.return_value.__aenter__.return_value = mock_session + + # Call list_tools + await client.list_tools() + + # Verify URL is unchanged in workload + assert workload.url == original_url + + # Verify the client was called with the original URL + mock_client.assert_called_once_with(original_url) + + +@pytest.mark.asyncio +async def test_workload_url_unchanged_during_call_tool(): + """Test that workload URL remains unchanged during call_tool.""" + original_url = "http://localhost:8080/mcp/test-server" + workload = Workload( + name="test-server", + url=original_url, + status="running", + tool_type="mcp", + ) + + client = MCPServerClient(workload, timeout=10) + + # Mock the MCP client session and result + mock_result = AsyncMock() + mock_result.content = [AsyncMock(text="Tool result")] + + mock_session = AsyncMock() + mock_session.call_tool.return_value = mock_result + + with ( + patch("mcp_optimizer.mcp_client.streamablehttp_client") as mock_client, + patch( + "mcp_optimizer.mcp_client.ClientSession", return_value=mock_session + ) as mock_session_class, + ): + # Mock the context manager + mock_client.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock(), AsyncMock()) + mock_session_class.return_value.__aenter__.return_value = mock_session + + # Call tool + await client.call_tool("test_tool", {"param": "value"}) + + # Verify URL is unchanged in workload + assert workload.url == original_url + + # Verify the client was called with the original URL + mock_client.assert_called_once_with(original_url) + + +@pytest.mark.asyncio +async def test_workload_url_unchanged_with_proxy_mode_explicit(): + """Test that workload URL is not modified when explicit proxy_mode is provided.""" + original_url = "http://localhost:8080/custom/endpoint" + workload = Workload( + name="test-server", + url=original_url, + proxy_mode="streamable-http", + status="running", + tool_type="mcp", + ) + + client = MCPServerClient(workload, timeout=10) + + # Mock the MCP client session + mock_result = AsyncMock() + mock_result.tools = [] + + mock_session = AsyncMock() + mock_session.list_tools.return_value = mock_result + + with ( + patch("mcp_optimizer.mcp_client.streamablehttp_client") as mock_client, + patch( + "mcp_optimizer.mcp_client.ClientSession", return_value=mock_session + ) as mock_session_class, + ): + # Mock the context manager + mock_client.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock(), AsyncMock()) + mock_session_class.return_value.__aenter__.return_value = mock_session + + # Call list_tools + await client.list_tools() + + # Verify URL is unchanged in workload + assert workload.url == original_url + + # Verify the client was called with the original URL (no modification) + mock_client.assert_called_once_with(original_url) + + +@pytest.mark.asyncio +async def test_workload_url_unchanged_multiple_operations(): + """Test that workload URL remains unchanged across multiple operations.""" + original_url = "http://localhost:8080/mcp/test-server" + workload = Workload( + name="test-server", + url=original_url, + status="running", + tool_type="mcp", + ) + + client = MCPServerClient(workload, timeout=10) + + # Mock the MCP client session + mock_list_result = AsyncMock() + mock_list_result.tools = [] + + mock_call_result = AsyncMock() + mock_call_result.content = [AsyncMock(text="Tool result")] + + mock_session = AsyncMock() + mock_session.list_tools.return_value = mock_list_result + mock_session.call_tool.return_value = mock_call_result + + with ( + patch("mcp_optimizer.mcp_client.streamablehttp_client") as mock_client, + patch( + "mcp_optimizer.mcp_client.ClientSession", return_value=mock_session + ) as mock_session_class, + ): + # Mock the context manager + mock_client.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock(), AsyncMock()) + mock_session_class.return_value.__aenter__.return_value = mock_session + + # Perform multiple operations + await client.list_tools() + assert workload.url == original_url + + await client.call_tool("test_tool", {"param": "value"}) + assert workload.url == original_url + + await client.list_tools() + assert workload.url == original_url + + # Verify the client was always called with the original URL + assert mock_client.call_count == 3 + for call in mock_client.call_args_list: + assert call[0][0] == original_url From c5139eeebb08f6922a2caae94f9645c59b46a7d0 Mon Sep 17 00:00:00 2001 From: Alejandro Ponce Date: Thu, 23 Oct 2025 13:40:45 +0300 Subject: [PATCH 2/2] Refactor URL preservation tests to reduce duplication MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Added mock_mcp_session fixture to eliminate repeated mock setup - Consolidated init tests using pytest.mark.parametrize (3 test cases in one) - Parametrized list_tools tests to cover both transport types (streamable, SSE) - Parametrized call_tool tests with different URL and proxy_mode combinations - Updated multiple operations test to use shared fixture - Reduced test code from ~245 lines to ~175 lines while maintaining coverage - All 16 tests pass (3 init variations + 2 list_tools + 2 call_tool + 1 multi-op + 8 existing) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- src/mcp_optimizer/mcp_client.py | 1 + tests/test_mcp_client.py | 224 +++++++++++--------------------- 2 files changed, 78 insertions(+), 147 deletions(-) diff --git a/src/mcp_optimizer/mcp_client.py b/src/mcp_optimizer/mcp_client.py index af8d3f9..300d23e 100644 --- a/src/mcp_optimizer/mcp_client.py +++ b/src/mcp_optimizer/mcp_client.py @@ -128,6 +128,7 @@ async def _execute_with_session(self, operation: Callable[[ClientSession], Await f"Using {proxy_mode} client for workload", workload=self.workload.name, proxy_mode_field=self.workload.proxy_mode, + url=self.workload.url, ) try: diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index 53dc3a0..2cced53 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -238,30 +238,35 @@ def test_extract_error_from_exception_group(): assert "Some error" in result -def test_workload_url_unchanged_after_init(): - """Test that workload URL is not modified during MCPServerClient initialization.""" - original_url = "http://localhost:8080/sse/test-server" - workload = Workload( - name="test-server", - url=original_url, - status="running", - tool_type="mcp", - ) +@pytest.fixture +def mock_mcp_session(): + """Create a mock MCP session for testing.""" + mock_session = AsyncMock() + mock_list_result = AsyncMock() + mock_list_result.tools = [] + mock_session.list_tools.return_value = mock_list_result - # Create client - _client = MCPServerClient(workload, timeout=10) + mock_call_result = AsyncMock() + mock_call_result.content = [AsyncMock(text="Tool result")] + mock_session.call_tool.return_value = mock_call_result - # Verify URL is unchanged - assert workload.url == original_url + return mock_session -def test_workload_url_unchanged_with_proxy_mode(): - """Test that workload URL is not modified when proxy_mode is set.""" - original_url = "http://localhost:8080/mcp/test-server" +@pytest.mark.parametrize( + "url,proxy_mode", + [ + ("http://localhost:8080/sse/test-server", None), + ("http://localhost:8080/mcp/test-server", "streamable-http"), + ("http://localhost:8080/custom/endpoint", "sse"), + ], +) +def test_workload_url_unchanged_after_init(url, proxy_mode): + """Test that workload URL is not modified during MCPServerClient initialization.""" workload = Workload( name="test-server", - url=original_url, - proxy_mode="streamable-http", + url=url, + proxy_mode=proxy_mode, status="running", tool_type="mcp", ) @@ -270,172 +275,108 @@ def test_workload_url_unchanged_with_proxy_mode(): _client = MCPServerClient(workload, timeout=10) # Verify URL is unchanged - assert workload.url == original_url + assert workload.url == url @pytest.mark.asyncio -async def test_workload_url_unchanged_during_list_tools_streamable(): - """Test that workload URL remains unchanged during list_tools with streamable-http.""" - original_url = "http://localhost:8080/mcp/test-server" +@pytest.mark.parametrize( + "url,client_mock_name,context_return", + [ + ( + "http://localhost:8080/mcp/test-server", + "streamablehttp_client", + (AsyncMock(), AsyncMock(), AsyncMock()), + ), + ("http://localhost:8080/sse/test-server", "sse_client", (AsyncMock(), AsyncMock())), + ], +) +async def test_workload_url_unchanged_during_list_tools( + url, client_mock_name, context_return, mock_mcp_session +): + """Test that workload URL remains unchanged during list_tools for both transport types.""" workload = Workload( name="test-server", - url=original_url, + url=url, status="running", tool_type="mcp", ) client = MCPServerClient(workload, timeout=10) - # Mock the MCP client session - mock_result = AsyncMock() - mock_result.tools = [] - - mock_session = AsyncMock() - mock_session.list_tools.return_value = mock_result - with ( - patch("mcp_optimizer.mcp_client.streamablehttp_client") as mock_client, + patch(f"mcp_optimizer.mcp_client.{client_mock_name}") as mock_client, patch( - "mcp_optimizer.mcp_client.ClientSession", return_value=mock_session + "mcp_optimizer.mcp_client.ClientSession", return_value=mock_mcp_session ) as mock_session_class, ): # Mock the context manager - mock_client.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock(), AsyncMock()) - mock_session_class.return_value.__aenter__.return_value = mock_session + mock_client.return_value.__aenter__.return_value = context_return + mock_session_class.return_value.__aenter__.return_value = mock_mcp_session # Call list_tools await client.list_tools() # Verify URL is unchanged in workload - assert workload.url == original_url + assert workload.url == url # Verify the client was called with the original URL - mock_client.assert_called_once_with(original_url) + mock_client.assert_called_once_with(url) @pytest.mark.asyncio -async def test_workload_url_unchanged_during_list_tools_sse(): - """Test that workload URL remains unchanged during list_tools with SSE.""" - original_url = "http://localhost:8080/sse/test-server" - workload = Workload( - name="test-server", - url=original_url, - status="running", - tool_type="mcp", - ) - - client = MCPServerClient(workload, timeout=10) - - # Mock the MCP client session - mock_result = AsyncMock() - mock_result.tools = [] - - mock_session = AsyncMock() - mock_session.list_tools.return_value = mock_result - - with ( - patch("mcp_optimizer.mcp_client.sse_client") as mock_client, - patch( - "mcp_optimizer.mcp_client.ClientSession", return_value=mock_session - ) as mock_session_class, - ): - # Mock the context manager - mock_client.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock()) - mock_session_class.return_value.__aenter__.return_value = mock_session - - # Call list_tools - await client.list_tools() - - # Verify URL is unchanged in workload - assert workload.url == original_url - - # Verify the client was called with the original URL - mock_client.assert_called_once_with(original_url) - - -@pytest.mark.asyncio -async def test_workload_url_unchanged_during_call_tool(): +@pytest.mark.parametrize( + "url,proxy_mode,client_mock_name,context_return", + [ + ( + "http://localhost:8080/mcp/test-server", + None, + "streamablehttp_client", + (AsyncMock(), AsyncMock(), AsyncMock()), + ), + ( + "http://localhost:8080/custom/endpoint", + "streamable-http", + "streamablehttp_client", + (AsyncMock(), AsyncMock(), AsyncMock()), + ), + ], +) +async def test_workload_url_unchanged_during_call_tool( + url, proxy_mode, client_mock_name, context_return, mock_mcp_session +): """Test that workload URL remains unchanged during call_tool.""" - original_url = "http://localhost:8080/mcp/test-server" workload = Workload( name="test-server", - url=original_url, + url=url, + proxy_mode=proxy_mode, status="running", tool_type="mcp", ) client = MCPServerClient(workload, timeout=10) - # Mock the MCP client session and result - mock_result = AsyncMock() - mock_result.content = [AsyncMock(text="Tool result")] - - mock_session = AsyncMock() - mock_session.call_tool.return_value = mock_result - with ( - patch("mcp_optimizer.mcp_client.streamablehttp_client") as mock_client, + patch(f"mcp_optimizer.mcp_client.{client_mock_name}") as mock_client, patch( - "mcp_optimizer.mcp_client.ClientSession", return_value=mock_session + "mcp_optimizer.mcp_client.ClientSession", return_value=mock_mcp_session ) as mock_session_class, ): # Mock the context manager - mock_client.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock(), AsyncMock()) - mock_session_class.return_value.__aenter__.return_value = mock_session + mock_client.return_value.__aenter__.return_value = context_return + mock_session_class.return_value.__aenter__.return_value = mock_mcp_session # Call tool await client.call_tool("test_tool", {"param": "value"}) # Verify URL is unchanged in workload - assert workload.url == original_url + assert workload.url == url # Verify the client was called with the original URL - mock_client.assert_called_once_with(original_url) + mock_client.assert_called_once_with(url) @pytest.mark.asyncio -async def test_workload_url_unchanged_with_proxy_mode_explicit(): - """Test that workload URL is not modified when explicit proxy_mode is provided.""" - original_url = "http://localhost:8080/custom/endpoint" - workload = Workload( - name="test-server", - url=original_url, - proxy_mode="streamable-http", - status="running", - tool_type="mcp", - ) - - client = MCPServerClient(workload, timeout=10) - - # Mock the MCP client session - mock_result = AsyncMock() - mock_result.tools = [] - - mock_session = AsyncMock() - mock_session.list_tools.return_value = mock_result - - with ( - patch("mcp_optimizer.mcp_client.streamablehttp_client") as mock_client, - patch( - "mcp_optimizer.mcp_client.ClientSession", return_value=mock_session - ) as mock_session_class, - ): - # Mock the context manager - mock_client.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock(), AsyncMock()) - mock_session_class.return_value.__aenter__.return_value = mock_session - - # Call list_tools - await client.list_tools() - - # Verify URL is unchanged in workload - assert workload.url == original_url - - # Verify the client was called with the original URL (no modification) - mock_client.assert_called_once_with(original_url) - - -@pytest.mark.asyncio -async def test_workload_url_unchanged_multiple_operations(): +async def test_workload_url_unchanged_multiple_operations(mock_mcp_session): """Test that workload URL remains unchanged across multiple operations.""" original_url = "http://localhost:8080/mcp/test-server" workload = Workload( @@ -447,26 +388,15 @@ async def test_workload_url_unchanged_multiple_operations(): client = MCPServerClient(workload, timeout=10) - # Mock the MCP client session - mock_list_result = AsyncMock() - mock_list_result.tools = [] - - mock_call_result = AsyncMock() - mock_call_result.content = [AsyncMock(text="Tool result")] - - mock_session = AsyncMock() - mock_session.list_tools.return_value = mock_list_result - mock_session.call_tool.return_value = mock_call_result - with ( patch("mcp_optimizer.mcp_client.streamablehttp_client") as mock_client, patch( - "mcp_optimizer.mcp_client.ClientSession", return_value=mock_session + "mcp_optimizer.mcp_client.ClientSession", return_value=mock_mcp_session ) as mock_session_class, ): # Mock the context manager mock_client.return_value.__aenter__.return_value = (AsyncMock(), AsyncMock(), AsyncMock()) - mock_session_class.return_value.__aenter__.return_value = mock_session + mock_session_class.return_value.__aenter__.return_value = mock_mcp_session # Perform multiple operations await client.list_tools()