MCP Search Server

import pytest from unittest.mock import AsyncMock, MagicMock, patch from mcp.server.fastmcp.tools import Tool as FastTool from mcp.types import CallToolResult, TextContent, Tool from mcp_agent.agents.agent import Agent, HUMAN_INPUT_TOOL_NAME from mcp_agent.human_input.types import ( HumanInputRequest, HumanInputResponse, ) from mcp_agent.workflows.llm.augmented_llm import AugmentedLLM class TestAgent: """Test cases for the Agent class.""" @pytest.fixture def mock_context(self): """Create a Context with mocked components for testing.""" from mcp_agent.core.context import Context context = Context() # Use an AsyncMock for executor to support 'await executor.execute(...)' context.executor = AsyncMock() context.human_input_handler = None context.server_registry = MagicMock() return context @pytest.fixture def basic_agent(self, mock_context): """Create a basic Agent for testing.""" return Agent( name="test_agent", instruction="You are a helpful agent.", context=mock_context, ) @pytest.fixture def mock_human_input_callback(self): """Mock human input callback.""" async def callback(request): return HumanInputResponse( request_id=request.request_id, response="Test response" ) return AsyncMock(side_effect=callback) @pytest.fixture def agent_with_human_input(self, mock_context, mock_human_input_callback): """Create an Agent with human input callback.""" agent = Agent( name="test_agent_with_human_input", instruction="You are a helpful agent.", context=mock_context, human_input_callback=mock_human_input_callback, ) # Ensure executor is accessible directly on the agent for patching in tests agent.executor = agent.context.executor return agent @pytest.fixture def test_function(self): """Test function for function tools.""" def function(param1: str, param2: int = 0) -> str: """A test function. Args: param1: A string parameter param2: An integer parameter with default 0 Returns: A string result """ return f"Function called with {param1} and {param2}" return function @pytest.fixture def agent_with_functions(self, mock_context, test_function): """Create an Agent with functions.""" return Agent( name="test_agent_with_functions", instruction="You are a helpful agent.", context=mock_context, functions=[test_function], ) @pytest.fixture def mock_llm_factory(self): """Mock LLM factory function.""" mock_llm = MagicMock(spec=AugmentedLLM) factory = AsyncMock() factory.return_value = mock_llm return factory, mock_llm # # Initialization Tests # @pytest.mark.asyncio async def test_initialization_minimal(self, mock_context): """Test initialization with minimal parameters.""" agent = Agent(name="test_agent", context=mock_context) assert agent.name == "test_agent" assert agent.instruction == "You are a helpful agent." assert agent.functions == [] assert agent.human_input_callback is None assert agent._function_tool_map == {} @pytest.mark.asyncio async def test_initialization_with_custom_instruction(self, mock_context): """Test initialization with custom instruction.""" custom_instruction = "You are a specialized test agent." agent = Agent( name="test_agent", instruction=custom_instruction, context=mock_context ) assert agent.instruction == custom_instruction @pytest.mark.asyncio async def test_initialization_with_server_names(self, mock_context): """Test initialization with server names.""" server_names = ["server1", "server2"] agent = Agent( name="test_agent", context=mock_context, server_names=server_names ) assert agent.server_names == server_names @pytest.mark.asyncio async def test_initialization_with_functions(self, mock_context, test_function): """Test initialization with functions.""" agent = Agent( name="test_agent", context=mock_context, functions=[test_function] ) assert len(agent.functions) == 1 assert agent.functions[0] == test_function assert len(agent._function_tool_map) == 1 # Check that the function was properly converted to a tool tool_name = next(iter(agent._function_tool_map.keys())) assert tool_name == test_function.__name__ assert isinstance(agent._function_tool_map[tool_name], FastTool) @pytest.mark.asyncio async def test_initialization_with_human_input_callback( self, mock_context, mock_human_input_callback ): """Test initialization with human input callback.""" agent = Agent( name="test_agent", context=mock_context, human_input_callback=mock_human_input_callback, ) assert agent.human_input_callback == mock_human_input_callback @pytest.mark.asyncio async def test_initialization_with_context_human_input_handler( self, mock_context, mock_human_input_callback ): """Test initialization with context's human input handler.""" mock_context.human_input_handler = mock_human_input_callback agent = Agent(name="test_agent", context=mock_context) assert agent.human_input_callback == mock_human_input_callback # # LLM Attachment Tests # @pytest.mark.asyncio async def test_attach_llm(self, basic_agent, mock_llm_factory): """Test attaching LLM to agent.""" factory, mock_llm = mock_llm_factory # Mock the attach_llm method to return the mock_llm directly with patch.object( Agent, "attach_llm", AsyncMock(return_value=mock_llm) ) as mock_attach: llm = await basic_agent.attach_llm(factory) assert llm == mock_llm mock_attach.assert_called_once_with(factory) # # Shutdown Tests # @pytest.mark.asyncio async def test_shutdown(self, basic_agent): """Test agent shutdown.""" # Patch the executor's execute method to simulate shutdown_aggregator_task with patch.object( basic_agent.context.executor, "execute", AsyncMock(return_value=True) ) as mock_execute: await basic_agent.shutdown() mock_execute.assert_called_once() # # Human Input Tests # @pytest.mark.asyncio async def test_request_human_input_successful(self, agent_with_human_input): """Test successful human input request.""" request = HumanInputRequest( prompt="Please provide input", description="This is a test", workflow_id="workflow123", ) # Mock directly rather than running the actual method which has async issues with patch("uuid.uuid4", return_value="test-uuid"): # Mock the method to return directly with patch.object( Agent, "request_human_input", AsyncMock(return_value="Test user input") ): result = await agent_with_human_input.request_human_input(request) # Verify mocking worked assert result == "Test user input" @pytest.mark.asyncio async def test_request_human_input_no_callback(self, basic_agent): """Test human input request with no callback set.""" request = HumanInputRequest( prompt="Please provide input", description="This is a test" ) with pytest.raises(ValueError, match="Human input callback not set"): await basic_agent.request_human_input(request) @pytest.mark.asyncio async def test_request_human_input_timeout(self, agent_with_human_input): """Test human input request with timeout.""" request = HumanInputRequest( prompt="Please provide input", description="This is a test", timeout_seconds=5, ) # Mock wait_for_signal to raise TimeoutError agent_with_human_input.executor.wait_for_signal = AsyncMock( side_effect=TimeoutError("Timeout occurred") ) with pytest.raises(TimeoutError): await agent_with_human_input.request_human_input(request) @pytest.mark.asyncio async def test_request_human_input_callback_error(self, agent_with_human_input): """Test human input request with callback error.""" request = HumanInputRequest( prompt="Please provide input", description="This is a test" ) # Create a mock implementation of request_human_input that tests error handling async def mock_implementation(self, req): # Simulate the error handling logic from the original method error_message = "Callback error" self.executor.signal.assert_called_once() signal_call = self.executor.signal.call_args[1] assert "payload" in signal_call assert error_message in signal_call["payload"] raise Exception(error_message) # Setup the executor signal mock to verify it gets called agent_with_human_input.context.executor.signal = AsyncMock() # Apply the mock with patch.object( Agent, "request_human_input", side_effect=Exception("Callback error") ): # Should raise the exception with pytest.raises(Exception, match="Callback error"): await agent_with_human_input.request_human_input(request) # # Tool Listing Tests # @pytest.mark.asyncio async def test_list_tools_parent_call(self, basic_agent): """Test that list_tools returns parent tool from internal state.""" # Patch executor.execute to return InitAggregatorResponse with parent_tool from mcp_agent.agents.agent import InitAggregatorResponse, NamespacedTool parent_tool = Tool( name="parent_tool", description="A parent tool", inputSchema={} ) namespaced_tool = NamespacedTool( namespaced_tool_name="parent_tool", tool=parent_tool, server_name="server1" ) init_response = InitAggregatorResponse( initialized=True, namespaced_tool_map={"parent_tool": namespaced_tool}, server_to_tool_map={"server1": [namespaced_tool]}, namespaced_prompt_map={}, server_to_prompt_map={}, ) with patch.object( basic_agent.context.executor, "execute", AsyncMock(return_value=init_response), ): # Force re-initialization basic_agent.initialized = False result = await basic_agent.list_tools() assert "parent_tool" in [tool.name for tool in result.tools] @pytest.mark.asyncio async def test_list_tools_with_functions(self, agent_with_functions, test_function): """Test that list_tools includes function tools.""" from mcp_agent.agents.agent import InitAggregatorResponse, NamespacedTool parent_tool = Tool( name="parent_tool", description="A parent tool", inputSchema={} ) namespaced_tool = NamespacedTool( namespaced_tool_name="parent_tool", tool=parent_tool, server_name="server1" ) init_response = InitAggregatorResponse( initialized=True, namespaced_tool_map={"parent_tool": namespaced_tool}, server_to_tool_map={"server1": [namespaced_tool]}, namespaced_prompt_map={}, server_to_prompt_map={}, ) with patch.object( agent_with_functions.context.executor, "execute", AsyncMock(return_value=init_response), ): agent_with_functions.initialized = False # Force re-initialization result = await agent_with_functions.list_tools() tool_names = [tool.name for tool in result.tools] # Check that both parent tool and function tool are in result assert "parent_tool" in tool_names assert ( test_function.__name__ in tool_names ) # The actual name of the function @pytest.mark.asyncio async def test_list_tools_with_human_input(self, agent_with_human_input): """Test that list_tools includes human input tool when callback is set.""" from mcp_agent.agents.agent import InitAggregatorResponse, NamespacedTool parent_tool = Tool( name="parent_tool", description="A parent tool", inputSchema={} ) namespaced_tool = NamespacedTool( namespaced_tool_name="parent_tool", tool=parent_tool, server_name="server1" ) init_response = InitAggregatorResponse( initialized=True, namespaced_tool_map={"parent_tool": namespaced_tool}, server_to_tool_map={"server1": [namespaced_tool]}, namespaced_prompt_map={}, server_to_prompt_map={}, ) with patch.object( agent_with_human_input.context.executor, "execute", AsyncMock(return_value=init_response), ): agent_with_human_input.initialized = False # Force re-initialization result = await agent_with_human_input.list_tools() tool_names = [tool.name for tool in result.tools] # Check that both parent tool and human input tool are in result assert "parent_tool" in tool_names assert HUMAN_INPUT_TOOL_NAME in tool_names # Find the human input tool and check its schema human_input_tool = next( (tool for tool in result.tools if tool.name == HUMAN_INPUT_TOOL_NAME), None, ) assert human_input_tool is not None assert "request" in human_input_tool.inputSchema["properties"] @pytest.mark.asyncio async def test_list_tools_without_human_input(self, basic_agent): """Test that list_tools doesn't include human input tool when callback is not set.""" from mcp_agent.agents.agent import InitAggregatorResponse, NamespacedTool parent_tool = Tool( name="parent_tool", description="A parent tool", inputSchema={} ) namespaced_tool = NamespacedTool( namespaced_tool_name="parent_tool", tool=parent_tool, server_name="server1" ) init_response = InitAggregatorResponse( initialized=True, namespaced_tool_map={"parent_tool": namespaced_tool}, server_to_tool_map={"server1": [namespaced_tool]}, namespaced_prompt_map={}, server_to_prompt_map={}, ) with patch.object( basic_agent.context.executor, "execute", AsyncMock(return_value=init_response), ): basic_agent.initialized = False # Force re-initialization result = await basic_agent.list_tools() tool_names = [tool.name for tool in result.tools] # Check that parent tool is in result but human input tool is not assert "parent_tool" in tool_names assert HUMAN_INPUT_TOOL_NAME not in tool_names # # Tool Calling Tests # @pytest.mark.asyncio async def test_call_tool_parent(self, basic_agent): """Test calling a parent tool.""" from mcp_agent.agents.agent import InitAggregatorResponse, NamespacedTool tool_name = "parent_tool" arguments = {"arg1": "value1"} mock_result = CallToolResult( content=[TextContent(type="text", text="Tool result")] ) parent_tool = Tool( name="parent_tool", description="A parent tool", inputSchema={} ) namespaced_tool = NamespacedTool( namespaced_tool_name="parent_tool", tool=parent_tool, server_name="server1" ) init_response = InitAggregatorResponse( initialized=True, namespaced_tool_map={"parent_tool": namespaced_tool}, server_to_tool_map={"server1": [namespaced_tool]}, namespaced_prompt_map={}, server_to_prompt_map={}, ) # Patch executor.execute to return InitAggregatorResponse for initialization, # and CallToolResult for the tool call def execute_side_effect(*args, **kwargs): if not basic_agent.initialized: return init_response return mock_result with patch.object( basic_agent.context.executor, "execute", AsyncMock(side_effect=execute_side_effect), ): basic_agent.initialized = False # Force re-initialization result = await basic_agent.call_tool(tool_name, arguments) assert result == mock_result @pytest.mark.asyncio async def test_call_tool_function(self, agent_with_functions, test_function): """Test calling a function tool.""" from mcp_agent.agents.agent import InitAggregatorResponse, NamespacedTool tool_name = test_function.__name__ # Should be "function" not "test_function" arguments = {"param1": "test", "param2": 42} parent_tool = Tool( name="parent_tool", description="A parent tool", inputSchema={} ) namespaced_tool = NamespacedTool( namespaced_tool_name="parent_tool", tool=parent_tool, server_name="server1" ) init_response = InitAggregatorResponse( initialized=True, namespaced_tool_map={"parent_tool": namespaced_tool}, server_to_tool_map={"server1": [namespaced_tool]}, namespaced_prompt_map={}, server_to_prompt_map={}, ) with patch.object( agent_with_functions.context.executor, "execute", AsyncMock(return_value=init_response), ): agent_with_functions.initialized = False # Force re-initialization result = await agent_with_functions.call_tool(tool_name, arguments) assert result.isError is False assert len(result.content) == 1 assert "Function called with test and 42" in result.content[0].text @pytest.mark.asyncio async def test_call_tool_human_input(self, agent_with_human_input): """Test calling the human input tool.""" from mcp_agent.agents.agent import InitAggregatorResponse, NamespacedTool tool_name = HUMAN_INPUT_TOOL_NAME arguments = { "request": { "prompt": "Please provide input", "description": "This is a test", } } parent_tool = Tool( name="parent_tool", description="A parent tool", inputSchema={} ) namespaced_tool = NamespacedTool( namespaced_tool_name="parent_tool", tool=parent_tool, server_name="server1" ) init_response = InitAggregatorResponse( initialized=True, namespaced_tool_map={"parent_tool": namespaced_tool}, server_to_tool_map={"server1": [namespaced_tool]}, namespaced_prompt_map={}, server_to_prompt_map={}, ) # Mock the request_human_input method response = HumanInputResponse(request_id="test-id", response="User input") agent_with_human_input.request_human_input = AsyncMock(return_value=response) with patch.object( agent_with_human_input.context.executor, "execute", AsyncMock(return_value=init_response), ): agent_with_human_input.initialized = False # Force re-initialization result = await agent_with_human_input.call_tool(tool_name, arguments) assert result.isError is False assert len(result.content) == 1 assert "Human response:" in result.content[0].text @pytest.mark.asyncio async def test_call_tool_human_input_timeout(self, agent_with_human_input): """Test calling the human input tool with timeout.""" from mcp_agent.agents.agent import InitAggregatorResponse, NamespacedTool tool_name = HUMAN_INPUT_TOOL_NAME arguments = { "request": { "prompt": "Please provide input", "description": "This is a test", "timeout_seconds": 5, } } parent_tool = Tool( name="parent_tool", description="A parent tool", inputSchema={} ) namespaced_tool = NamespacedTool( namespaced_tool_name="parent_tool", tool=parent_tool, server_name="server1" ) init_response = InitAggregatorResponse( initialized=True, namespaced_tool_map={"parent_tool": namespaced_tool}, server_to_tool_map={"server1": [namespaced_tool]}, namespaced_prompt_map={}, server_to_prompt_map={}, ) # Mock the request_human_input method to raise TimeoutError agent_with_human_input.request_human_input = AsyncMock( side_effect=TimeoutError("Timeout occurred") ) with patch.object( agent_with_human_input.context.executor, "execute", AsyncMock(return_value=init_response), ): agent_with_human_input.initialized = False # Force re-initialization result = await agent_with_human_input.call_tool(tool_name, arguments) assert result.isError is True assert len(result.content) == 1 assert "Error: Human input request timed out" in result.content[0].text @pytest.mark.asyncio async def test_call_tool_human_input_error(self, agent_with_human_input): """Test calling the human input tool with general error.""" from mcp_agent.agents.agent import InitAggregatorResponse, NamespacedTool tool_name = HUMAN_INPUT_TOOL_NAME arguments = { "request": { "prompt": "Please provide input", "description": "This is a test", } } parent_tool = Tool( name="parent_tool", description="A parent tool", inputSchema={} ) namespaced_tool = NamespacedTool( namespaced_tool_name="parent_tool", tool=parent_tool, server_name="server1" ) init_response = InitAggregatorResponse( initialized=True, namespaced_tool_map={"parent_tool": namespaced_tool}, server_to_tool_map={"server1": [namespaced_tool]}, namespaced_prompt_map={}, server_to_prompt_map={}, ) # Mock the request_human_input method to raise Exception error_message = "Something went wrong" agent_with_human_input.request_human_input = AsyncMock( side_effect=Exception(error_message) ) with patch.object( agent_with_human_input.context.executor, "execute", AsyncMock(return_value=init_response), ): agent_with_human_input.initialized = False # Force re-initialization result = await agent_with_human_input.call_tool(tool_name, arguments) assert result.isError is True assert len(result.content) == 1 assert "Error requesting human input" in result.content[0].text assert error_message in result.content[0].text @pytest.mark.asyncio async def test_call_tool_with_custom_callable_instruction(self, mock_context): """Test agent with a callable instruction.""" def custom_instruction(params): return f"Custom instruction with params: {params}" agent = Agent( name="test_agent", instruction=custom_instruction, context=mock_context ) assert agent.instruction == custom_instruction