Repository: jerryjliu/gpt_index
File: llama-index-integrations/tools/llama-index-tools-mcp/tests/test_tools_mcp.py (Python)

Code before:
    assert json_schema["properties"]["name"]["$ref"] == "#/$defs/TestName"
    assert json_schema["properties"]["method"]["$ref"] == "#/$defs/TestMethod"
    assert json_schema["properties"]["lst"]["$ref"] == "#/$defs/TestList"

    # Check model types
    assert json_schema["$defs"]["TestName"]["properties"]["name"]["type"] == "string"
    assert json_schema["$defs"]["TestList"]["properties"]["lst"]["type"] == "array"
    assert (
        json_schema["$defs"]["TestList"]["properties"]["lst"]["items"]["type"]
        == "integer"
    )


def test_schema_structure_exact_match(client: BasicMCPClient):
    """Test that the generated schema structure exactly matches the expected format."""
    json_schema = (
        McpToolSpec(client, allowed_tools=["test_pydantic"])
        .to_tool_list()[0]
        .metadata.fn_schema.model_json_schema()
    )

    assert json_schema["type"] == "object"
    assert set(json_schema["required"]) == {"name", "method", "lst"}


def test_union_of_literal_types_resolution(client: BasicMCPClient):
    """Test that union of Literal types (enum in anyOf) are resolved correctly.

    Regression test for https://github.com/run-llama/llama_index/issues/20109
    When Literal types appear inside a union (anyOf in JSON Schema),
    _resolve_union_option should handle the enum case just like
    _resolve_field_type already does.
    """
    from typing import get_args, Literal

    tool_spec = McpToolSpec(client)

    # Test 1: _resolve_union_option handles a single enum option
    option_with_enum = {"enum": ["red", "green", "blue"]}
    resolved = tool_spec._resolve_union_option(option_with_enum, {})
    assert resolved == Literal["red", "green", "blue"]

    # Test 2: _resolve_union_type handles anyOf with multiple enum options
    union_schema = {
        "anyOf": [
            {"enum": ["red", "green", "blue"]},
            {"enum": ["cyan", "magenta", "yellow"]},
        ]
    }
    resolved_union = tool_spec._resolve_union_type(union_schema, {})
    union_args = get_args(resolved_union)

Review comment (suggestion):
This test should always check that the expected literal types are within the resolved union options, which is done only for the first case

Code after:

    # Check property types
    assert json_schema["properties"]["name"]["$ref"] == "#/$defs/TestName"
    assert json_schema["properties"]["method"]["$ref"] == "#/$defs/TestMethod"
    assert json_schema["properties"]["lst"]["$ref"] == "#/$defs/TestList"

    # Check model types
    assert json_schema["$defs"]["TestName"]["properties"]["name"]["type"] == "string"
    assert json_schema["$defs"]["TestList"]["properties"]["lst"]["type"] == "array"
    assert (
        json_schema["$defs"]["TestList"]["properties"]["lst"]["items"]["type"]
        == "integer"
    )


def test_schema_structure_exact_match(client: BasicMCPClient):
    """Test that the generated schema structure exactly matches the expected format."""
    json_schema = (
        McpToolSpec(client, allowed_tools=["test_pydantic"])
        .to_tool_list()[0]
        .metadata.fn_schema.model_json_schema()
    )

    assert json_schema["type"] == "object"
    assert set(json_schema["required"]) == {"name", "method", "lst"}


def test_resolve_union_option_with_enum(client: BasicMCPClient):
    """
    Regression test for https://github.com/run-llama/llama_index/issues/20109

    _resolve_union_option should handle enum entries in anyOf schemas,
    resolving them to Literal types instead of falling through to str.
    """
    tool_spec = McpToolSpec(client)

    result = tool_spec._resolve_union_option({"enum": ["red", "green", "blue"]}, {})
    assert result == Literal["red", "green", "blue"]


def test_resolve_union_type_multiple_enums(client: BasicMCPClient):
    """
    Test that anyOf with multiple enum entries resolves to a Union of Literals.
    """
    tool_spec = McpToolSpec(client)

    schema = {
        "anyOf": [
            {"enum": ["red", "green", "blue"]},
            {"enum": ["cyan", "magenta", "yellow"]},
        ]