/tests/mcp/test_e2e.py

"""
End-to-end integration tests for MCP with aisuite.

These tests verify the complete flow of using MCP tools with aisuite's
chat.completions.create() API, including:
- Config dict format
- Mixing MCP tools with Python functions
- Multiple MCP servers with prefixing
- Automatic cleanup

Requirements:
    - Node.js and npx must be installed
    - Tests are marked with @pytest.mark.integration
    - Run with: pytest tests/mcp/test_e2e.py -v -m integration
"""

import pytest
from unittest.mock import patch, MagicMock, Mock
from aisuite import Client


def create_mock_response(content="Test response", tool_calls=None):
    """Helper to create a mock chat completion response."""
    # Create a simple mock object that mimics the response structure
    response = MagicMock()
    response.choices = [MagicMock()]
    response.choices[0].message = MagicMock()
    response.choices[0].message.content = content
    response.choices[0].message.tool_calls = tool_calls
    response.choices[0].intermediate_messages = [response.choices[0].message]
    response.intermediate_responses = []

    return response


@pytest.mark.integration
class TestMCPConfigDictFormat:
    """Test using MCP config dict format in chat.completions.create()."""

    def test_basic_config_dict(self, temp_test_dir, skip_if_no_npx):
        """Test basic MCP config dict usage."""
        client = Client()

        # Mock the provider to avoid actual LLM API calls
        with patch.object(client.chat.completions, "_tool_runner") as mock_runner:
            mock_runner.return_value = create_mock_response("Files listed successfully")

            response = client.chat.completions.create(
                model="openai:gpt-4o",
                messages=[{"role": "user", "content": "List all files"}],
                tools=[
                    {
                        "type": "mcp",
                        "name": "filesystem",
                        "command": "npx",
                        "args": [
                            "-y",
                            "@modelcontextprotocol/server-filesystem",
                            temp_test_dir,
                        ],
                    }
                ],
                max_turns=2,
            )

            # Verify response
            assert response.choices[0].message.content == "Files listed successfully"

            # Verify tool_runner was called with processed tools
            assert mock_runner.called
            call_args = mock_runner.call_args
            tools_arg = call_args[0][3]  # tools is 4th positional arg

            # Verify tools were converted to callables
            assert isinstance(tools_arg, list)
            assert all(callable(t) for t in tools_arg)

    def test_config_dict_with_allowed_tools(self, temp_test_dir, skip_if_no_npx):
        """Test MCP config dict with allowed_tools filtering."""
        client = Client()

        with patch.object(client.chat.completions, "_tool_runner") as mock_runner:
            mock_runner.return_value = create_mock_response("File read successfully")

            response = client.chat.completions.create(
                model="openai:gpt-4o",
                messages=[{"role": "user", "content": "Read test.txt"}],
                tools=[
                    {
                        "type": "mcp",
                        "name": "filesystem",
                        "command": "npx",
                        "args": [
                            "-y",
                            "@modelcontextprotocol/server-filesystem",
                            temp_test_dir,
                        ],
                        "allowed_tools": ["read_file"],  # Only allow reading
                    }
                ],
                max_turns=2,
            )

            assert response.choices[0].message.content == "File read successfully"

            # Verify only read_file tool was passed
            call_args = mock_runner.call_args
            tools_arg = call_args[0][3]

            assert len(tools_arg) == 1
            assert tools_arg[0].__name__ == "read_file"

    def test_config_dict_with_prefixing(self, temp_test_dir, skip_if_no_npx):
        """Test MCP config dict with tool name prefixing."""
        client = Client()

        with patch.object(client.chat.completions, "_tool_runner") as mock_runner:
            mock_runner.return_value = create_mock_response("Success")

            response = client.chat.completions.create(
                model="openai:gpt-4o",
                messages=[{"role": "user", "content": "Test"}],
                tools=[
                    {
                        "type": "mcp",
                        "name": "docs",
                        "command": "npx",
                        "args": [
                            "-y",
                            "@modelcontextprotocol/server-filesystem",
                            temp_test_dir,
                        ],
                        "use_tool_prefix": True,
                    }
                ],
                max_turns=2,
            )

            # Verify tools have prefixes
            call_args = mock_runner.call_args
            tools_arg = call_args[0][3]

            tool_names = [t.__name__ for t in tools_arg]
            assert any(name.startswith("docs__") for name in tool_names)


@pytest.mark.integration
class TestMCPWithPythonFunctions:
    """Test mixing MCP tools with regular Python functions."""

    def test_mix_mcp_and_python_functions(self, temp_test_dir, skip_if_no_npx):
        """Test using MCP config dict alongside Python functions."""
        client = Client()

        # Define a Python function
        def get_current_time() -> str:
            """Get the current time."""
            return "2025-01-01 12:00:00"

        with patch.object(client.chat.completions, "_tool_runner") as mock_runner:
            mock_runner.return_value = create_mock_response("Mixed tools work!")

            response = client.chat.completions.create(
                model="openai:gpt-4o",
                messages=[{"role": "user", "content": "What time is it?"}],
                tools=[
                    get_current_time,  # Python function
                    {
                        "type": "mcp",
                        "name": "filesystem",
                        "command": "npx",
                        "args": [
                            "-y",
                            "@modelcontextprotocol/server-filesystem",
                            temp_test_dir,
                        ],
                    },  # MCP config
                ],
                max_turns=2,
            )

            assert response.choices[0].message.content == "Mixed tools work!"

            # Verify both types of tools were passed
            call_args = mock_runner.call_args
            tools_arg = call_args[0][3]

            # Should have Python function + MCP tools
            assert len(tools_arg) > 1

            # Verify Python function is in there
            assert any(t.__name__ == "get_current_time" for t in tools_arg)

            # Verify MCP tools are in there
            assert any(t.__name__ == "read_file" for t in tools_arg)


@pytest.mark.integration
class TestMultipleMCPServers:
    """Test using multiple MCP servers simultaneously."""

    def test_multiple_servers_with_prefixing(self, temp_test_dir, skip_if_no_npx):
        """Test multiple MCP servers with tool name prefixing to avoid collisions."""
        import tempfile

        client = Client()

        # Create a second temp directory
        with tempfile.TemporaryDirectory() as temp_dir_2:
            with patch.object(client.chat.completions, "_tool_runner") as mock_runner:
                mock_runner.return_value = create_mock_response(
                    "Multiple servers work!"
                )

                response = client.chat.completions.create(
                    model="openai:gpt-4o",
                    messages=[{"role": "user", "content": "Compare directories"}],
                    tools=[
                        {
                            "type": "mcp",
                            "name": "dir1",
                            "command": "npx",
                            "args": [
                                "-y",
                                "@modelcontextprotocol/server-filesystem",
                                temp_test_dir,
                            ],
                            "use_tool_prefix": True,
                        },
                        {
                            "type": "mcp",
                            "name": "dir2",
                            "command": "npx",
                            "args": [
                                "-y",
                                "@modelcontextprotocol/server-filesystem",
                                temp_dir_2,
                            ],
                            "use_tool_prefix": True,
                        },
                    ],
                    max_turns=2,
                )

                assert response.choices[0].message.content == "Multiple servers work!"

                # Verify tools from both servers with prefixes
                call_args = mock_runner.call_args
                tools_arg = call_args[0][3]

                tool_names = [t.__name__ for t in tools_arg]

                # Should have tools from both servers
                assert any(name.startswith("dir1__") for name in tool_names)
                assert any(name.startswith("dir2__") for name in tool_names)

                # Should have both read_file tools with different prefixes
                assert "dir1__read_file" in tool_names
                assert "dir2__read_file" in tool_names


@pytest.mark.integration
class TestMCPCleanup:
    """Test that MCP clients are properly cleaned up."""

    def test_cleanup_after_success(self, temp_test_dir, skip_if_no_npx):
        """Test MCP clients are cleaned up after successful request."""
        client = Client()

        with patch.object(client.chat.completions, "_tool_runner") as mock_runner:
            mock_runner.return_value = create_mock_response("Success")

            # Patch MCPClient to track close() calls
            with patch("aisuite.client.MCPClient") as mock_mcp_class:
                mock_mcp_instance = MagicMock()
                mock_mcp_class.from_config.return_value = mock_mcp_instance
                mock_mcp_instance.get_callable_tools.return_value = []

                response = client.chat.completions.create(
                    model="openai:gpt-4o",
                    messages=[{"role": "user", "content": "Test"}],
                    tools=[
                        {
                            "type": "mcp",
                            "name": "filesystem",
                            "command": "npx",
                            "args": [
                                "-y",
                                "@modelcontextprotocol/server-filesystem",
                                temp_test_dir,
                            ],
                        }
                    ],
                    max_turns=2,
                )

                # Verify MCP client was used as context manager (cleanup called)
                mock_mcp_instance.__enter__.assert_called_once()
                mock_mcp_instance.__exit__.assert_called_once()

    def test_cleanup_after_error(self, temp_test_dir, skip_if_no_npx):
        """Test MCP clients are cleaned up even after error."""
        client = Client()

        with patch.object(client.chat.completions, "_tool_runner") as mock_runner:
            # Make tool_runner raise an error
            mock_runner.side_effect = ValueError("Test error")

            # Patch MCPClient to track close() calls
            with patch("aisuite.client.MCPClient") as mock_mcp_class:
                mock_mcp_instance = MagicMock()
                mock_mcp_class.from_config.return_value = mock_mcp_instance
                mock_mcp_instance.get_callable_tools.return_value = []

                with pytest.raises(ValueError, match="Test error"):
                    client.chat.completions.create(
                        model="openai:gpt-4o",
                        messages=[{"role": "user", "content": "Test"}],
                        tools=[
                            {
                                "type": "mcp",
                                "name": "filesystem",
                                "command": "npx",
                                "args": [
                                    "-y",
                                    "@modelcontextprotocol/server-filesystem",
                                    temp_test_dir,
                                ],
                            }
                        ],
                        max_turns=2,
                    )

                # Even after error, MCP client context manager exit should be called
                mock_mcp_instance.__enter__.assert_called_once()
                mock_mcp_instance.__exit__.assert_called_once()


@pytest.mark.integration
class TestMCPErrorHandling:
    """Test error handling for MCP integration."""

    def test_invalid_mcp_config_raises_error(self):
        """Test that invalid MCP config raises clear error."""
        client = Client()

        with pytest.raises(ValueError, match="must have 'name'"):
            client.chat.completions.create(
                model="openai:gpt-4o",
                messages=[{"role": "user", "content": "Test"}],
                tools=[
                    {
                        "type": "mcp",
                        # Missing 'name' field
                        "command": "npx",
                        "args": ["server"],
                    }
                ],
                max_turns=2,
            )

    def test_mcp_not_installed_raises_error(self, temp_test_dir, skip_if_no_npx):
        """Test that helpful error is raised if MCP package not installed."""
        client = Client()

        # Simulate MCP not being installed
        with patch("aisuite.client.MCP_AVAILABLE", False):
            with pytest.raises(ImportError, match="mcp.*package"):
                client.chat.completions.create(
                    model="openai:gpt-4o",
                    messages=[{"role": "user", "content": "Test"}],
                    tools=[
                        {
                            "type": "mcp",
                            "name": "filesystem",
                            "command": "npx",
                            "args": [
                                "-y",
                                "@modelcontextprotocol/server-filesystem",
                                temp_test_dir,
                            ],
                        }
                    ],
                    max_turns=2,
                )