tests/contrib/test_bulk_tool_caller.py

from typing import Any

import pytest
from mcp.types import TextContent

from fastmcp import FastMCP
from fastmcp.contrib.bulk_tool_caller.bulk_tool_caller import (
    BulkToolCaller,
    CallToolRequest,
    CallToolRequestResult,
)
from fastmcp.tools.tool import Tool


class ToolException(Exception):
    """Custom exception for tool errors."""

    pass


async def error_tool(arg1: str) -> dict[str, Any]:
    """A tool that raises an error for testing purposes."""
    raise ToolException(f"Error in tool with arg1: {arg1}")


def error_tool_result_factory(arg1: str) -> CallToolRequestResult:
    """Generates the expected error result for error_tool."""
    # Mimic the error message format generated by BulkToolCaller when catching ToolException
    formatted_error_text = (
        "Error calling tool 'error_tool': Error in tool with arg1: " + arg1
    )
    return CallToolRequestResult(
        isError=True,
        content=[TextContent(text=formatted_error_text, type="text")],
        tool="error_tool",
        arguments={"arg1": arg1},
    )


async def echo_tool(arg1: str) -> str:
    """A simple tool that echoes arguments or raises an error."""
    return arg1


def echo_tool_result_factory(arg1: str) -> CallToolRequestResult:
    """A tool that returns a result based on the input arguments."""
    return CallToolRequestResult(
        isError=False,
        content=[TextContent(text=f"{arg1}", type="text")],
        tool="echo_tool",
        arguments={"arg1": arg1},
    )


async def no_return_tool(arg1: str) -> None:
    """A simple tool that echoes arguments or raises an error."""


def no_return_tool_result_factory(arg1: str) -> CallToolRequestResult:
    """A tool that returns a result based on the input arguments."""
    return CallToolRequestResult(
        isError=False,
        content=[],
        tool="no_return_tool",
        arguments={"arg1": arg1},
    )


@pytest.fixture(scope="module")
def live_server_with_tool() -> FastMCP:
    """Fixture to create a FastMCP server instance with the echo_tool registered."""
    server = FastMCP()
    server.add_tool(Tool.from_function(echo_tool))
    server.add_tool(Tool.from_function(error_tool))
    server.add_tool(Tool.from_function(no_return_tool))
    return server


@pytest.fixture
def bulk_caller_live(live_server_with_tool: FastMCP) -> BulkToolCaller:
    """Fixture to create a BulkToolCaller instance connected to the live server."""
    bulk_tool_caller = BulkToolCaller()
    bulk_tool_caller.register_tools(live_server_with_tool)
    return bulk_tool_caller


ECHO_TOOL_NAME = "echo_tool"
ERROR_TOOL_NAME = "error_tool"
NO_RETURN_TOOL_NAME = "no_return_tool"


async def test_call_tool_bulk_single_success(bulk_caller_live: BulkToolCaller):
    """Test single successful call via call_tool_bulk using echo_tool."""
    tool_arguments = [{"arg1": "value1"}]
    expected_result = echo_tool_result_factory(**tool_arguments[0])

    results = await bulk_caller_live.call_tool_bulk(ECHO_TOOL_NAME, tool_arguments)

    assert len(results) == 1
    result = results[0]
    assert result == expected_result


async def test_call_tool_bulk_multiple_success(bulk_caller_live: BulkToolCaller):
    """Test multiple successful calls via call_tool_bulk using echo_tool."""
    tool_arguments = [{"arg1": "value1"}, {"arg1": "value2"}]
    expected_results = [echo_tool_result_factory(**args) for args in tool_arguments]

    results = await bulk_caller_live.call_tool_bulk(ECHO_TOOL_NAME, tool_arguments)

    assert len(results) == 2
    assert results == expected_results


async def test_call_tool_bulk_error_stops(bulk_caller_live: BulkToolCaller):
    """Test call_tool_bulk stops on first error using error_tool."""
    tool_arguments = [{"arg1": "error_value"}, {"arg1": "value2"}]
    expected_result = error_tool_result_factory(**tool_arguments[0])

    results = await bulk_caller_live.call_tool_bulk(
        ERROR_TOOL_NAME, tool_arguments, continue_on_error=False
    )

    assert len(results) == 1
    result = results[0]
    assert result == expected_result


async def test_call_tool_bulk_error_continues(bulk_caller_live: BulkToolCaller):
    """Test call_tool_bulk continues on error using error_tool and echo_tool."""
    tool_arguments = [{"arg1": "error_value"}, {"arg1": "success_value"}]
    expected_error_result = error_tool_result_factory(**tool_arguments[0])
    expected_success_result = echo_tool_result_factory(**tool_arguments[1])

    tool_calls = [
        CallToolRequest(tool=ERROR_TOOL_NAME, arguments=tool_arguments[0]),
        CallToolRequest(tool=ECHO_TOOL_NAME, arguments=tool_arguments[1]),
    ]

    results = await bulk_caller_live.call_tools_bulk(tool_calls, continue_on_error=True)

    assert len(results) == 2

    error_result = results[0]
    assert error_result == expected_error_result

    success_result = results[1]
    assert success_result == expected_success_result


async def test_call_tools_bulk_single_success(bulk_caller_live: BulkToolCaller):
    """Test single successful call via call_tools_bulk using echo_tool."""
    tool_calls = [CallToolRequest(tool=ECHO_TOOL_NAME, arguments={"arg1": "value1"})]
    expected_result = echo_tool_result_factory(**tool_calls[0].arguments)

    results = await bulk_caller_live.call_tools_bulk(tool_calls)

    assert len(results) == 1
    result = results[0]
    assert result == expected_result


async def test_call_tools_bulk_multiple_success(bulk_caller_live: BulkToolCaller):
    """Test multiple successful calls via call_tools_bulk with different tools."""
    tool_calls = [
        CallToolRequest(tool=ECHO_TOOL_NAME, arguments={"arg1": "echo_value"}),
        CallToolRequest(
            tool=NO_RETURN_TOOL_NAME, arguments={"arg1": "no_return_value"}
        ),
    ]
    expected_results = [
        echo_tool_result_factory(**tool_calls[0].arguments),
        no_return_tool_result_factory(**tool_calls[1].arguments),
    ]

    results = await bulk_caller_live.call_tools_bulk(tool_calls)

    assert len(results) == 2
    assert results == expected_results


async def test_call_tools_bulk_error_stops(bulk_caller_live: BulkToolCaller):
    """Test call_tools_bulk stops on first error using error_tool."""
    tool_calls = [
        CallToolRequest(tool=ERROR_TOOL_NAME, arguments={"arg1": "error_value"}),
        CallToolRequest(tool=ECHO_TOOL_NAME, arguments={"arg1": "skipped_value"}),
    ]
    expected_result = error_tool_result_factory(**tool_calls[0].arguments)

    results = await bulk_caller_live.call_tools_bulk(
        tool_calls, continue_on_error=False
    )

    assert len(results) == 1
    result = results[0]
    assert result == expected_result


async def test_call_tools_bulk_error_continues(bulk_caller_live: BulkToolCaller):
    """Test call_tools_bulk continues on error using error_tool and echo_tool."""
    tool_calls = [
        CallToolRequest(tool=ERROR_TOOL_NAME, arguments={"arg1": "error_value"}),
        CallToolRequest(tool=ECHO_TOOL_NAME, arguments={"arg1": "success_value"}),
    ]
    expected_error_result = error_tool_result_factory(**tool_calls[0].arguments)
    expected_success_result = echo_tool_result_factory(**tool_calls[1].arguments)

    results = await bulk_caller_live.call_tools_bulk(tool_calls, continue_on_error=True)

    assert len(results) == 2

    error_result = results[0]
    assert error_result == expected_error_result

    success_result = results[1]
    assert success_result == expected_success_result