arcade-mcp/libs/tests/core/test_output.py

from typing import Any

import pytest
from arcade_core.output import ToolOutputFactory
from pydantic import BaseModel


@pytest.fixture
def output_factory():
    return ToolOutputFactory()


class SampleOutputModel(BaseModel):
    result: Any


@pytest.mark.parametrize(
    "data, expected_value",
    [
        (None, ""),
        ("success", "success"),
        ("", ""),
        (None, ""),
        (123, 123),
        (0, 0),
        (123.45, 123.45),
        (True, True),
        (False, False),
    ],
)
def test_success(output_factory, data, expected_value):
    data_obj = SampleOutputModel(result=data) if data is not None else None
    output = output_factory.success(data=data_obj)
    assert output.value == expected_value
    assert output.error is None


@pytest.mark.parametrize(
    "data, expected_value",
    [
        # Dict types (simulating TypedDict at runtime)
        ({"name": "test", "value": 123}, {"name": "test", "value": 123}),
        ({}, {}),
        ({"nested": {"key": "value"}}, {"nested": {"key": "value"}}),
        # List types
        (["a", "b", "c"], ["a", "b", "c"]),
        ([1, 2, 3], [1, 2, 3]),
        ([], []),
        # List of dicts (simulating list[TypedDict])
        (
            [{"id": 1, "name": "a"}, {"id": 2, "name": "b"}],
            [{"id": 1, "name": "a"}, {"id": 2, "name": "b"}],
        ),
        ([{}], [{}]),
        # Mixed lists
        ([1, "two", 3.0, True], [1, "two", 3.0, True]),
    ],
)
def test_success_complex_types(output_factory, data, expected_value):
    """Test that dict and list types are properly handled by ToolOutputFactory."""
    data_obj = SampleOutputModel(result=data)
    output = output_factory.success(data=data_obj)
    assert output.value == expected_value
    assert output.error is None


def test_success_with_basemodel_direct(output_factory):
    """Test that BaseModel instances are converted to dict via model_dump()."""

    class TestModel(BaseModel):
        name: str
        value: int

    model = TestModel(name="test", value=42)
    output = output_factory.success(data=model)
    assert output.value == {"name": "test", "value": 42}
    assert output.error is None


def test_success_raw_dict(output_factory):
    """Test that raw dict values (not wrapped in model) are handled correctly."""
    raw_dict = {"key": "value", "number": 123}
    output = output_factory.success(data=raw_dict)
    assert output.value == raw_dict
    assert output.error is None


def test_success_raw_list(output_factory):
    """Test that raw list values (not wrapped in model) are handled correctly."""
    raw_list = [{"id": 1}, {"id": 2}, {"id": 3}]
    output = output_factory.success(data=raw_list)
    assert output.value == raw_list
    assert output.error is None


@pytest.mark.parametrize(
    "message, developer_message",
    [
        ("Error occurred", None),
        ("Error occurred", "Detailed error message"),
    ],
)
def test_fail(output_factory, message, developer_message):
    output = output_factory.fail(message=message, developer_message=developer_message)
    assert output.error is not None
    assert output.error.message == message
    assert output.error.developer_message == developer_message
    assert output.error.can_retry is False


def test_fail_empty_message_gets_default(output_factory):
    output = output_factory.fail(message="")
    assert output.error is not None
    assert output.error.message == "Unspecified error during tool execution"


def test_fail_whitespace_message_gets_default(output_factory):
    output = output_factory.fail(message="  ")
    assert output.error is not None
    assert output.error.message == "Unspecified error during tool execution"


def test_fail_nonempty_message_unchanged(output_factory):
    output = output_factory.fail(message="real error")
    assert output.error is not None
    assert output.error.message == "real error"


def test_fail_retry_empty_message_gets_default(output_factory):
    output = output_factory.fail_retry(message="")
    assert output.error is not None
    assert output.error.message == "Unspecified error during tool execution"


def test_fail_retry_whitespace_message_gets_default(output_factory):
    output = output_factory.fail_retry(message="   ")
    assert output.error is not None
    assert output.error.message == "Unspecified error during tool execution"


@pytest.mark.parametrize(
    "message, developer_message, additional_prompt_content, retry_after_ms",
    [
        ("Retry error", None, None, None),
        ("Retry error", "Retrying", "Please try again with this additional data: foobar", 1000),
    ],
)
def test_fail_retry(
    output_factory, message, developer_message, additional_prompt_content, retry_after_ms
):
    output = output_factory.fail_retry(
        message=message,
        developer_message=developer_message,
        additional_prompt_content=additional_prompt_content,
        retry_after_ms=retry_after_ms,
    )
    assert output.error is not None
    assert output.error.message == message
    assert output.error.developer_message == developer_message
    assert output.error.can_retry is True
    assert output.error.additional_prompt_content == additional_prompt_content
    assert output.error.retry_after_ms == retry_after_ms