arcade-mcp/libs/tests/arcade_mcp_server/conftest.py

"""Shared fixtures and utilities for arcade-mcp-server tests."""

import asyncio
from collections.abc import AsyncGenerator
from typing import Annotated, Any
from unittest.mock import AsyncMock, Mock

import pytest
import pytest_asyncio
from arcade_core.catalog import MaterializedTool, ToolCatalog, ToolMeta, create_func_models
from arcade_core.schema import (
    InputParameter,
    OAuth2Requirement,
    ToolAuthRequirement,
    ToolDefinition,
    ToolInput,
    ToolkitDefinition,
    ToolOutput,
    ToolRequirements,
    ValueSchema,
)
from arcade_mcp_server import tool
from arcade_mcp_server.context import Context
from arcade_mcp_server.server import MCPServer
from arcade_mcp_server.session import ServerSession
from arcade_mcp_server.settings import MCPSettings
from arcade_tdk.auth import OAuth2


@pytest.fixture
def event_loop():
    """Create an instance of the default event loop for the test session."""
    loop = asyncio.new_event_loop()
    yield loop
    loop.close()


@pytest.fixture
def sample_tool_def() -> ToolDefinition:
    """Create a sample tool definition."""
    return ToolDefinition(
        name="test_tool",
        fully_qualified_name="TestToolkit.test_tool",
        description="A test tool",
        toolkit=ToolkitDefinition(name="TestToolkit", description="Test toolkit", version="1.0.0"),
        input=ToolInput(
            parameters=[
                InputParameter(
                    name="text",
                    required=True,
                    description="Input text",
                    value_schema=ValueSchema(val_type="string"),
                )
            ]
        ),
        output=ToolOutput(description="Tool output", value_schema=ValueSchema(val_type="string")),
        requirements=ToolRequirements(),
    )


@pytest.fixture
def sample_tool_def_with_auth() -> ToolDefinition:
    """Create a sample tool definition."""
    return ToolDefinition(
        name="sample_tool_with_auth",
        fully_qualified_name="TestToolkit.sample_tool_with_auth",
        description="A test tool",
        toolkit=ToolkitDefinition(name="TestToolkit", description="Test toolkit", version="1.0.0"),
        input=ToolInput(
            parameters=[
                InputParameter(
                    name="text",
                    required=True,
                    description="Input text",
                    value_schema=ValueSchema(val_type="string"),
                )
            ]
        ),
        output=ToolOutput(description="Tool output", value_schema=ValueSchema(val_type="string")),
        requirements=ToolRequirements(
            authorization=ToolAuthRequirement(
                provider_type="oauth2",
                provider_id="test-provider",
                id="test-provider",
                oauth2=OAuth2Requirement(
                    scopes=["test.scope", "another.scope"],
                ),
            ),
        ),
    )


@pytest.fixture
def sample_tool_func():
    """Create a sample tool function."""

    @tool
    def sample_tool(
        text: Annotated[str, "Input text to echo"],
    ) -> Annotated[str, "Echoed text result"]:
        """Echo input text back to the caller."""
        return f"Echo: {text}"

    return sample_tool


@pytest.fixture
def sample_tool_func_with_auth():
    """Create a sample tool function."""

    @tool(
        requires_auth=OAuth2(
            id="test-provider",
            scopes=["test.scope", "another.scope"],
        ),
    )
    def sample_tool_with_auth(
        text: Annotated[str, "Input text to echo"],
    ) -> Annotated[str, "Echoed text result"]:
        """Echo input text back to the caller."""
        return f"Echo: {text}"

    return sample_tool_with_auth


@pytest.fixture
def materialized_tool(sample_tool_func, sample_tool_def) -> MaterializedTool:
    """Create a materialized tool with required models and metadata."""
    input_model, output_model = create_func_models(sample_tool_func)
    meta = ToolMeta(module=sample_tool_func.__module__, toolkit=sample_tool_def.toolkit.name)
    return MaterializedTool(
        tool=sample_tool_func,
        definition=sample_tool_def,
        meta=meta,
        input_model=input_model,
        output_model=output_model,
    )


@pytest.fixture
def materialized_tool_with_auth(
    sample_tool_func_with_auth, sample_tool_def_with_auth
) -> MaterializedTool:
    """Create a materialized tool with required models and metadata."""
    input_model, output_model = create_func_models(sample_tool_func_with_auth)
    meta = ToolMeta(
        module=sample_tool_func_with_auth.__module__, toolkit=sample_tool_def_with_auth.toolkit.name
    )
    return MaterializedTool(
        tool=sample_tool_func_with_auth,
        definition=sample_tool_def_with_auth,
        meta=meta,
        input_model=input_model,
        output_model=output_model,
    )


@pytest.fixture
def tool_catalog(
    materialized_tool: MaterializedTool, materialized_tool_with_auth: MaterializedTool
) -> ToolCatalog:
    """Create a tool catalog with sample tools."""
    catalog = ToolCatalog()
    catalog._tools[materialized_tool.definition.get_fully_qualified_name()] = materialized_tool
    catalog._tools[materialized_tool_with_auth.definition.get_fully_qualified_name()] = (
        materialized_tool_with_auth
    )
    return catalog


@pytest.fixture
def mcp_settings() -> MCPSettings:
    """Create test MCP settings."""
    settings = MCPSettings()
    settings.debug = True
    settings.middleware.enable_logging = True
    settings.middleware.mask_error_details = False
    return settings


@pytest_asyncio.fixture
async def mcp_server(tool_catalog, mcp_settings) -> AsyncGenerator[MCPServer, None]:
    """Create and start an MCP server."""
    server = MCPServer(
        catalog=tool_catalog,
        name="Test Server",
        version="1.0.0",
        settings=mcp_settings,
    )
    await server.start()
    yield server
    await server.stop()


@pytest.fixture
def mock_read_stream() -> AsyncMock:
    """Create a mock read stream."""
    stream = AsyncMock()
    stream.read = AsyncMock()
    return stream


@pytest.fixture
def mock_write_stream() -> AsyncMock:
    """Create a mock write stream."""
    stream = AsyncMock()
    stream.write = AsyncMock()
    stream.send = AsyncMock()
    return stream


@pytest_asyncio.fixture
async def server_session(mcp_server, mock_read_stream, mock_write_stream) -> ServerSession:
    """Create a server session."""
    session = ServerSession(
        server=mcp_server,
        read_stream=mock_read_stream,
        write_stream=mock_write_stream,
    )
    return session


@pytest_asyncio.fixture
async def initialized_server_session(server_session) -> ServerSession:
    """Create an initialized server session."""
    server_session.mark_initialized()
    return server_session


@pytest.fixture
def sample_messages() -> dict[str, Any]:
    """Sample MCP protocol messages for testing."""
    return {
        "initialize": {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "capabilities": {"tools": {}, "sampling": {}},
                "clientInfo": {"name": "test-client", "version": "1.0.0"},
            },
        },
        "initialized": {"jsonrpc": "2.0", "method": "notifications/initialized"},
        "list_tools": {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}},
        "call_tool": {
            "jsonrpc": "2.0",
            "id": 3,
            "method": "tools/call",
            "params": {"name": "TestToolkit.test_tool", "arguments": {"text": "Hello, world!"}},
        },
        "ping": {"jsonrpc": "2.0", "id": 4, "method": "ping"},
    }


@pytest.fixture
def mock_context() -> Context:
    """Create a mock context."""
    context = Mock(spec=Context)
    context.server = Mock()
    context.request_id = "test-request-123"
    context.session_id = "test-session-456"
    context.state = {}

    # Mock async methods
    context.log = AsyncMock()
    context.debug = AsyncMock()
    context.info = AsyncMock()
    context.warning = AsyncMock()
    context.error = AsyncMock()
    context.report_progress = AsyncMock()
    context.read_resource = AsyncMock(return_value=[])
    context.list_roots = AsyncMock(return_value=[])
    context.sample = AsyncMock()
    context.elicit = AsyncMock()
    context.send_tool_list_changed = AsyncMock()
    context.send_resource_list_changed = AsyncMock()
    context.send_prompt_list_changed = AsyncMock()

    return context


# Async test helpers
async def wait_for(condition, timeout=1.0):
    """Wait for a condition to become true."""
    start = asyncio.get_event_loop().time()
    while not condition():
        if asyncio.get_event_loop().time() - start > timeout:
            raise TimeoutError("Condition not met within timeout")
        await asyncio.sleep(0.01)


async def collect_messages(stream, count):
    """Collect a specific number of messages from a stream."""
    messages = []
    for _ in range(count):
        msg = await stream.read()
        messages.append(msg)
    return messages