agent-framework/python/packages/ag-ui/tests/ag_ui/test_utils.py

# Copyright (c) Microsoft. All rights reserved.

"""Tests for utilities."""

from dataclasses import dataclass
from datetime import date, datetime

from agent_framework_ag_ui._utils import (
    generate_event_id,
    make_json_safe,
    merge_state,
)


def test_generate_event_id():
    """Test event ID generation."""
    id1 = generate_event_id()
    id2 = generate_event_id()

    assert id1 != id2
    assert isinstance(id1, str)
    assert len(id1) > 0


def test_merge_state():
    """Test state merging."""
    current: dict[str, int] = {"a": 1, "b": 2}
    update: dict[str, int] = {"b": 3, "c": 4}

    result = merge_state(current, update)

    assert result["a"] == 1
    assert result["b"] == 3
    assert result["c"] == 4


def test_merge_state_empty_update():
    """Test merging with empty update."""
    current: dict[str, int] = {"x": 10, "y": 20}
    update: dict[str, int] = {}

    result = merge_state(current, update)

    assert result == current
    assert result is not current


def test_merge_state_empty_current():
    """Test merging with empty current state."""
    current: dict[str, int] = {}
    update: dict[str, int] = {"a": 1, "b": 2}

    result = merge_state(current, update)

    assert result == update


def test_merge_state_deep_copy():
    """Test that merge_state creates a deep copy preventing mutation of original."""
    current: dict[str, dict[str, object]] = {"recipe": {"name": "Cake", "ingredients": ["flour", "sugar"]}}
    update: dict[str, str] = {"other": "value"}

    result = merge_state(current, update)

    result["recipe"]["ingredients"].append("eggs")

    assert "eggs" not in current["recipe"]["ingredients"]
    assert current["recipe"]["ingredients"] == ["flour", "sugar"]
    assert result["recipe"]["ingredients"] == ["flour", "sugar", "eggs"]


def test_make_json_safe_basic():
    """Test JSON serialization of basic types."""
    assert make_json_safe("text") == "text"
    assert make_json_safe(123) == 123
    assert make_json_safe(None) is None
    assert make_json_safe(3.14) == 3.14
    assert make_json_safe(True) is True
    assert make_json_safe(False) is False


def test_make_json_safe_datetime():
    """Test datetime serialization."""
    dt = datetime(2025, 10, 30, 12, 30, 45)
    result = make_json_safe(dt)
    assert result == "2025-10-30T12:30:45"


def test_make_json_safe_date():
    """Test date serialization."""
    d = date(2025, 10, 30)
    result = make_json_safe(d)
    assert result == "2025-10-30"


@dataclass
class SampleDataclass:
    """Sample dataclass for testing."""

    name: str
    value: int


def test_make_json_safe_dataclass():
    """Test dataclass serialization."""
    obj = SampleDataclass(name="test", value=42)
    result = make_json_safe(obj)
    assert result == {"name": "test", "value": 42}


class ModelDumpObject:
    """Object with model_dump method."""

    def model_dump(self):
        return {"type": "model", "data": "dump"}


def test_make_json_safe_model_dump():
    """Test object with model_dump method."""
    obj = ModelDumpObject()
    result = make_json_safe(obj)
    assert result == {"type": "model", "data": "dump"}


class ToDictObject:
    """Object with to_dict method (like SerializationMixin)."""

    def to_dict(self):
        return {"type": "serialization_mixin", "method": "to_dict"}


def test_make_json_safe_to_dict():
    """Test object with to_dict method (SerializationMixin pattern)."""
    obj = ToDictObject()
    result = make_json_safe(obj)
    assert result == {"type": "serialization_mixin", "method": "to_dict"}


class DictObject:
    """Object with dict method."""

    def dict(self):
        return {"type": "dict", "method": "call"}


def test_make_json_safe_dict_method():
    """Test object with dict method."""
    obj = DictObject()
    result = make_json_safe(obj)
    assert result == {"type": "dict", "method": "call"}


class CustomObject:
    """Custom object with __dict__."""

    def __init__(self):
        self.field1 = "value1"
        self.field2 = 123


def test_make_json_safe_dict_attribute():
    """Test object with __dict__ attribute."""
    obj = CustomObject()
    result = make_json_safe(obj)
    assert result == {"field1": "value1", "field2": 123}


def test_make_json_safe_list():
    """Test list serialization."""
    lst = [1, "text", None, {"key": "value"}]
    result = make_json_safe(lst)
    assert result == [1, "text", None, {"key": "value"}]


def test_make_json_safe_tuple():
    """Test tuple serialization."""
    tpl = (1, 2, 3)
    result = make_json_safe(tpl)
    assert result == [1, 2, 3]


def test_make_json_safe_dict():
    """Test dict serialization."""
    d = {"a": 1, "b": {"c": 2}}
    result = make_json_safe(d)
    assert result == {"a": 1, "b": {"c": 2}}


def test_make_json_safe_nested():
    """Test nested structure serialization."""
    obj = {
        "datetime": datetime(2025, 10, 30),
        "list": [1, 2, CustomObject()],
        "nested": {"value": SampleDataclass(name="nested", value=99)},
    }
    result = make_json_safe(obj)

    assert result["datetime"] == "2025-10-30T00:00:00"
    assert result["list"][0] == 1
    assert result["list"][2] == {"field1": "value1", "field2": 123}
    assert result["nested"]["value"] == {"name": "nested", "value": 99}


class UnserializableObject:
    """Object that can't be serialized by standard methods."""

    def __init__(self):
        # Add attribute to trigger __dict__ fallback path
        pass


def test_make_json_safe_fallback():
    """Test fallback to dict for objects with __dict__."""
    obj = UnserializableObject()
    result = make_json_safe(obj)
    # Objects with __dict__ return their __dict__ dict
    assert isinstance(result, dict)


def test_make_json_safe_dataclass_with_nested_to_dict_object():
    """Test dataclass containing a to_dict object (like HandoffAgentUserRequest with AgentResponse).

    This test verifies the fix for the AG-UI JSON serialization error when
    HandoffAgentUserRequest (a dataclass) contains an AgentResponse (SerializationMixin).
    """

    class NestedToDictObject:
        """Simulates SerializationMixin objects like AgentResponse."""

        def __init__(self, contents: list[str]):
            self.contents = contents

        def to_dict(self):
            return {"type": "response", "contents": self.contents}

    @dataclass
    class ContainerDataclass:
        """Simulates HandoffAgentUserRequest dataclass."""

        response: NestedToDictObject

    obj = ContainerDataclass(response=NestedToDictObject(contents=["hello", "world"]))
    result = make_json_safe(obj)

    # Verify the nested to_dict object was properly serialized
    assert result == {"response": {"type": "response", "contents": ["hello", "world"]}}

    # Verify the result is actually JSON serializable
    import json

    json_str = json.dumps(result)
    assert json_str is not None


def test_convert_tools_to_agui_format_with_tool():
    """Test converting FunctionTool to AG-UI format."""
    from agent_framework import tool

    from agent_framework_ag_ui._utils import convert_tools_to_agui_format

    @tool
    def test_func(param: str, count: int = 5) -> str:
        """Test function."""
        return f"{param} {count}"

    result = convert_tools_to_agui_format([test_func])

    assert result is not None
    assert len(result) == 1
    assert result[0]["name"] == "test_func"
    assert result[0]["description"] == "Test function."
    assert "parameters" in result[0]
    assert "properties" in result[0]["parameters"]


def test_convert_tools_to_agui_format_with_callable():
    """Test converting plain callable to AG-UI format."""
    from agent_framework_ag_ui._utils import convert_tools_to_agui_format

    def plain_func(x: int) -> int:
        """A plain function."""
        return x * 2

    result = convert_tools_to_agui_format([plain_func])

    assert result is not None
    assert len(result) == 1
    assert result[0]["name"] == "plain_func"
    assert result[0]["description"] == "A plain function."
    assert "parameters" in result[0]


def test_convert_tools_to_agui_format_with_dict():
    """Test converting dict tool to AG-UI format."""
    from agent_framework_ag_ui._utils import convert_tools_to_agui_format

    tool_dict = {
        "name": "custom_tool",
        "description": "Custom tool",
        "parameters": {"type": "object"},
    }

    result = convert_tools_to_agui_format([tool_dict])

    assert result is not None
    assert len(result) == 1
    assert result[0] == tool_dict


def test_convert_tools_to_agui_format_with_none():
    """Test converting None tools."""
    from agent_framework_ag_ui._utils import convert_tools_to_agui_format

    result = convert_tools_to_agui_format(None)

    assert result is None


def test_convert_tools_to_agui_format_with_single_tool():
    """Test converting single tool (not in list)."""
    from agent_framework import tool

    from agent_framework_ag_ui._utils import convert_tools_to_agui_format

    @tool
    def single_tool(arg: str) -> str:
        """Single tool."""
        return arg

    result = convert_tools_to_agui_format(single_tool)

    assert result is not None
    assert len(result) == 1
    assert result[0]["name"] == "single_tool"


def test_convert_tools_to_agui_format_with_multiple_tools():
    """Test converting multiple tools."""
    from agent_framework import tool

    from agent_framework_ag_ui._utils import convert_tools_to_agui_format

    @tool
    def tool1(x: int) -> int:
        """Tool 1."""
        return x

    @tool
    def tool2(y: str) -> str:
        """Tool 2."""
        return y

    result = convert_tools_to_agui_format([tool1, tool2])

    assert result is not None
    assert len(result) == 2
    assert result[0]["name"] == "tool1"
    assert result[1]["name"] == "tool2"


# Additional tests for utils coverage


def test_safe_json_parse_with_dict():
    """Test safe_json_parse with dict input."""
    from agent_framework_ag_ui._utils import safe_json_parse

    input_dict = {"key": "value"}
    result = safe_json_parse(input_dict)
    assert result == input_dict


def test_safe_json_parse_with_json_string():
    """Test safe_json_parse with JSON string."""
    from agent_framework_ag_ui._utils import safe_json_parse

    result = safe_json_parse('{"key": "value"}')
    assert result == {"key": "value"}


def test_safe_json_parse_with_invalid_json():
    """Test safe_json_parse with invalid JSON."""
    from agent_framework_ag_ui._utils import safe_json_parse

    result = safe_json_parse("not json")
    assert result is None


def test_safe_json_parse_with_non_dict_json():
    """Test safe_json_parse with JSON that parses to non-dict."""
    from agent_framework_ag_ui._utils import safe_json_parse

    result = safe_json_parse("[1, 2, 3]")
    assert result is None


def test_safe_json_parse_with_none():
    """Test safe_json_parse with None input."""
    from agent_framework_ag_ui._utils import safe_json_parse

    result = safe_json_parse(None)
    assert result is None


def test_get_role_value_with_enum():
    """Test get_role_value with enum role."""
    from agent_framework import ChatMessage, Content

    from agent_framework_ag_ui._utils import get_role_value

    message = ChatMessage(role="user", contents=[Content.from_text("test")])
    result = get_role_value(message)
    assert result == "user"


def test_get_role_value_with_string():
    """Test get_role_value with string role."""
    from agent_framework_ag_ui._utils import get_role_value

    class MockMessage:
        role = "assistant"

    result = get_role_value(MockMessage())
    assert result == "assistant"


def test_get_role_value_with_none():
    """Test get_role_value with no role."""
    from agent_framework_ag_ui._utils import get_role_value

    class MockMessage:
        pass

    result = get_role_value(MockMessage())
    assert result == ""


def test_normalize_agui_role_developer():
    """Test normalize_agui_role maps developer to system."""
    from agent_framework_ag_ui._utils import normalize_agui_role

    assert normalize_agui_role("developer") == "system"


def test_normalize_agui_role_valid():
    """Test normalize_agui_role with valid roles."""
    from agent_framework_ag_ui._utils import normalize_agui_role

    assert normalize_agui_role("user") == "user"
    assert normalize_agui_role("assistant") == "assistant"
    assert normalize_agui_role("system") == "system"
    assert normalize_agui_role("tool") == "tool"


def test_normalize_agui_role_invalid():
    """Test normalize_agui_role with invalid role defaults to user."""
    from agent_framework_ag_ui._utils import normalize_agui_role

    assert normalize_agui_role("invalid") == "user"
    assert normalize_agui_role(123) == "user"


def test_extract_state_from_tool_args():
    """Test extract_state_from_tool_args."""
    from agent_framework_ag_ui._utils import extract_state_from_tool_args

    # Specific key
    assert extract_state_from_tool_args({"key": "value"}, "key") == "value"

    # Wildcard
    args = {"a": 1, "b": 2}
    assert extract_state_from_tool_args(args, "*") == args

    # Missing key
    assert extract_state_from_tool_args({"other": "value"}, "key") is None

    # None args
    assert extract_state_from_tool_args(None, "key") is None


def test_convert_agui_tools_to_agent_framework():
    """Test convert_agui_tools_to_agent_framework."""
    from agent_framework_ag_ui._utils import convert_agui_tools_to_agent_framework

    agui_tools = [
        {
            "name": "test_tool",
            "description": "A test tool",
            "parameters": {"type": "object", "properties": {"arg": {"type": "string"}}},
        }
    ]

    result = convert_agui_tools_to_agent_framework(agui_tools)

    assert result is not None
    assert len(result) == 1
    assert result[0].name == "test_tool"
    assert result[0].description == "A test tool"
    assert result[0].declaration_only is True


def test_convert_agui_tools_to_agent_framework_none():
    """Test convert_agui_tools_to_agent_framework with None."""
    from agent_framework_ag_ui._utils import convert_agui_tools_to_agent_framework

    result = convert_agui_tools_to_agent_framework(None)
    assert result is None


def test_convert_agui_tools_to_agent_framework_empty():
    """Test convert_agui_tools_to_agent_framework with empty list."""
    from agent_framework_ag_ui._utils import convert_agui_tools_to_agent_framework

    result = convert_agui_tools_to_agent_framework([])
    assert result is None


def test_make_json_safe_unconvertible():
    """Test make_json_safe with object that has no standard conversion."""

    class NoConversion:
        __slots__ = ()  # No __dict__

    from agent_framework_ag_ui._utils import make_json_safe

    result = make_json_safe(NoConversion())
    # Falls back to str()
    assert isinstance(result, str)