Merge pull request #22327 from BerriAI/fix/mcp-test-mock-filter-method-name

fix(mcp): update test mocks for renamed filter_server_ids_by_ip_with_info

Author: jquinter

tests/mcp_tests/test_mcp_server.py

@@ -659,9 +659,9 @@ async def test_list_tools_rest_api_server_not_found():
         mock_manager.get_allowed_mcp_servers = AsyncMock(
             return_value=["non_existent_server_id"]
         )
-        # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test)
-        mock_manager.filter_server_ids_by_ip = MagicMock(
-            side_effect=lambda server_ids, client_ip: server_ids
+        # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test)
+        mock_manager.filter_server_ids_by_ip_with_info = MagicMock(
+            side_effect=lambda server_ids, client_ip: (server_ids, 0)
         )
         # Return None when trying to get the server (server doesn't exist)
         mock_manager.get_mcp_server_by_id = MagicMock(return_value=None)
@@ -732,9 +732,9 @@ async def test_list_tools_rest_api_success():
             return_value=["test-server-123"]
         )
         mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server)
-        # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test)
-        mock_manager.filter_server_ids_by_ip = MagicMock(
-            side_effect=lambda server_ids, client_ip: server_ids
+        # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test)
+        mock_manager.filter_server_ids_by_ip_with_info = MagicMock(
+            side_effect=lambda server_ids, client_ip: (server_ids, 0)
         )
 
         # Mock the _get_tools_for_single_server function
@@ -814,9 +814,9 @@ def mock_get_server_by_id(server_id):
         )
         mock_manager.get_mcp_server_by_id = lambda server_id: mock_server_1 if server_id == "server1_id" else mock_server_2
         mock_manager._get_tools_from_server = AsyncMock(return_value=[mock_tool_1])
-        # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test)
-        mock_manager.filter_server_ids_by_ip = MagicMock(
-            side_effect=lambda server_ids, client_ip: server_ids
+        # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test)
+        mock_manager.filter_server_ids_by_ip_with_info = MagicMock(
+            side_effect=lambda server_ids, client_ip: (server_ids, 0)
         )
 
         with patch(
@@ -853,9 +853,9 @@ async def mock_get_tools_side_effect(
             mock_manager_2._get_tools_from_server = AsyncMock(
                 side_effect=mock_get_tools_side_effect
             )
-            # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test)
-            mock_manager_2.filter_server_ids_by_ip = MagicMock(
-                side_effect=lambda server_ids, client_ip: server_ids
+            # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test)
+            mock_manager_2.filter_server_ids_by_ip_with_info = MagicMock(
+                side_effect=lambda server_ids, client_ip: (server_ids, 0)
             )
 
         with patch(
@@ -881,9 +881,9 @@ async def mock_get_tools_side_effect(
         )
         mock_manager.get_mcp_server_by_id = lambda server_id: mock_server_1 if server_id == "server1_id" else (mock_server_2 if server_id == "server2_id" else mock_server_3)
         mock_manager._get_tools_from_server = AsyncMock(return_value=[mock_tool_1])
-        # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test)
-        mock_manager.filter_server_ids_by_ip = MagicMock(
-            side_effect=lambda server_ids, client_ip: server_ids
+        # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test)
+        mock_manager.filter_server_ids_by_ip_with_info = MagicMock(
+            side_effect=lambda server_ids, client_ip: (server_ids, 0)
         )
 
         with patch(
@@ -1817,9 +1817,9 @@ async def test_list_tool_rest_api_with_server_specific_auth():
                 mock_server.mcp_info = {"server_name": "zapier"}
 
                 mock_manager.get_mcp_server_by_id.return_value = mock_server
-                # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test)
-                mock_manager.filter_server_ids_by_ip = MagicMock(
-                    side_effect=lambda server_ids, client_ip: server_ids
+                # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test)
+                mock_manager.filter_server_ids_by_ip_with_info = MagicMock(
+                    side_effect=lambda server_ids, client_ip: (server_ids, 0)
                 )
 
                 mock_user_api_key_dict = UserAPIKeyAuth(
@@ -1911,9 +1911,9 @@ async def test_list_tool_rest_api_with_default_auth():
                 mock_server.mcp_info = {"server_name": "unknown_server"}
 
                 mock_manager.get_mcp_server_by_id.return_value = mock_server
-                # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test)
-                mock_manager.filter_server_ids_by_ip = MagicMock(
-                    side_effect=lambda server_ids, client_ip: server_ids
+                # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test)
+                mock_manager.filter_server_ids_by_ip_with_info = MagicMock(
+                    side_effect=lambda server_ids, client_ip: (server_ids, 0)
                 )
 
                 mock_user_api_key_dict = UserAPIKeyAuth(
@@ -2021,9 +2021,9 @@ async def test_list_tool_rest_api_all_servers_with_auth():
                         server_id
                     )
                 )
-                # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test)
-                mock_manager.filter_server_ids_by_ip = MagicMock(
-                    side_effect=lambda server_ids, client_ip: server_ids
+                # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test)
+                mock_manager.filter_server_ids_by_ip_with_info = MagicMock(
+                    side_effect=lambda server_ids, client_ip: (server_ids, 0)
                 )
 
                 mock_user_api_key_dict = UserAPIKeyAuth(
@@ -2154,9 +2154,9 @@ def mock_client_constructor(*args, **kwargs):
             return_value=["test-server-123"]
         )
         mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server)
-        # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test)
-        mock_manager.filter_server_ids_by_ip = MagicMock(
-            side_effect=lambda server_ids, client_ip: server_ids
+        # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test)
+        mock_manager.filter_server_ids_by_ip_with_info = MagicMock(
+            side_effect=lambda server_ids, client_ip: (server_ids, 0)
         )
 
         # Mock the _get_tools_from_server method to return all tools
@@ -2268,9 +2268,9 @@ def mock_client_constructor(*args, **kwargs):
             return_value=["test-server-456"]
         )
         mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server)
-        # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test)
-        mock_manager.filter_server_ids_by_ip = MagicMock(
-            side_effect=lambda server_ids, client_ip: server_ids
+        # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test)
+        mock_manager.filter_server_ids_by_ip_with_info = MagicMock(
+            side_effect=lambda server_ids, client_ip: (server_ids, 0)
         )
         # Mock the _get_tools_from_server method to return all tools
         mock_manager._get_tools_from_server = AsyncMock(return_value=mock_tools)
@@ -2368,9 +2368,9 @@ def mock_client_constructor(*args, **kwargs):
             return_value=["test-server-000"]
         )
         mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server)
-        # Mock filter_server_ids_by_ip to return input unchanged (no IP filtering in test)
-        mock_manager.filter_server_ids_by_ip = MagicMock(
-            side_effect=lambda server_ids, client_ip: server_ids
+        # Mock filter_server_ids_by_ip_with_info to return input unchanged (no IP filtering in test)
+        mock_manager.filter_server_ids_by_ip_with_info = MagicMock(
+            side_effect=lambda server_ids, client_ip: (server_ids, 0)
         )
 
         # Mock the _get_tools_from_server method to return all tools