ReMe/tests/test_message.py

"""Test cases for message schema and serialization."""

import unittest

from mcp.types import Tool

from reme_ai.core.enumeration import Role
from reme_ai.core.schema import ToolAttr, ToolCall, ContentBlock, Message


class TestModelDefinitions(unittest.TestCase):
    """Test suite for validating message models and their serialization methods."""

    def test_tool_attr_serialization(self):
        """Test if ToolAttr correctly dumps to JSON schema format."""
        # Test simple string attribute with enum
        attr = ToolAttr(
            type="string",
            description="The city name",
            enum=["Beijing", "London"],
        )
        dump = attr.simple_input_dump()

        print("\n=== ToolAttr.simple_input_dump() (string with enum) ===")
        print(dump)

        self.assertEqual(dump["type"], "string")
        self.assertEqual(dump["enum"], ["Beijing", "London"])
        self.assertIn("description", dump)

        # Test object attribute with required child properties
        obj_attr = ToolAttr(
            type="object",
            description="User information",
            properties={
                "name": ToolAttr(type="string", description="User name"),
                "age": ToolAttr(type="number", description="User age"),
            },
            required=["name"],  # 'name' is required, 'age' is optional
        )
        obj_dump = obj_attr.simple_input_dump()

        print("\n=== ToolAttr.simple_input_dump() (object with required) ===")
        print(obj_dump)

        self.assertEqual(obj_dump["type"], "object")
        self.assertIn("properties", obj_dump)
        self.assertEqual(obj_dump["required"], ["name"])
        self.assertIn("name", obj_dump["properties"])
        self.assertIn("age", obj_dump["properties"])

    def test_tool_call_initialization(self):
        """Test if ToolCall correctly parses raw OpenAI-style tool definitions."""
        raw_input = {
            "type": "function",
            "function": {
                "name": "get_weather",
                "description": "Check weather info",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {"type": "string", "description": "City name"},
                        "unit": {"type": "string", "description": "Temperature unit"},
                    },
                    "required": ["location"],
                },
            },
        }

        tc = ToolCall(**raw_input)

        print("\n=== ToolCall.simple_input_dump() ===")
        print(tc.simple_input_dump())
        print("\n=== ToolCall.simple_output_dump() ===")
        print(tc.simple_output_dump())

        self.assertEqual(tc.name, "get_weather")
        self.assertIn("location", tc.parameters.properties)
        self.assertIn("unit", tc.parameters.properties)
        # Check that 'location' is in the required list at ToolCall level
        self.assertIn("location", tc.parameters.required)
        self.assertNotIn("unit", tc.parameters.required)

    def test_tool_call_argument_parsing(self):
        """Test JSON argument parsing and validation."""
        tc = ToolCall(name="test", arguments='{"key": "value"}')

        self.assertTrue(tc.check_argument())
        self.assertEqual(tc.argument_dict["key"], "value")

        # Test invalid JSON
        tc.arguments = "{invalid_json}"
        self.assertFalse(tc.check_argument())

    def test_content_block_dynamic_mapping(self):
        """Test if ContentBlock correctly identifies content based on type key."""
        # Test Image Block
        img_data = {"type": "image_url", "image_url": {"url": "http://test.com/a.jpg"}}
        block = ContentBlock(**img_data)
        self.assertEqual(block.type, "image_url")
        self.assertEqual(block.content["url"], "http://test.com/a.jpg")

        # Test Text Block
        text_data = {"type": "text", "text": "Hello World"}
        block = ContentBlock(**text_data)
        self.assertEqual(block.content, "Hello World")

    def test_message_simple_dump(self):
        """Test the transformation of Message to standard API dictionary."""
        msg = Message(
            role=Role.ASSISTANT,
            content="Thinking...",
            reasoning_content="I should check the weather first.",
            tool_calls=[ToolCall(name="get_weather", arguments='{"city": "NY"}', id="call_123")],
        )

        dump = msg.simple_dump(add_reasoning=True)

        print("\n=== Message.simple_dump(add_reasoning=True) ===")
        print(dump)

        dump_no_reasoning = msg.simple_dump(add_reasoning=False)
        print("\n=== Message.simple_dump(add_reasoning=False) ===")
        print(dump_no_reasoning)

        self.assertEqual(dump["role"], "assistant")
        self.assertEqual(dump["reasoning_content"], "I should check the weather first.")
        self.assertEqual(len(dump["tool_calls"]), 1)
        self.assertEqual(dump["tool_calls"][0]["id"], "call_123")

    def test_message_format_human_readable(self):
        """Test the string representation of messages for logging/UI."""
        msg = Message(
            role=Role.USER,
            content=[
                ContentBlock(type="text", text="Look at this:"),
                ContentBlock(type="image_url", image_url={"url": "img.png"}),
            ],
        )

        formatted = msg.format_message(index=1, use_name=False)

        self.assertIn("round1", formatted)
        self.assertIn("user:", formatted)
        self.assertIn("Look at this:", formatted)
        self.assertIn("img.png", formatted)

    def test_mcp_conversion(self):
        """Test the interoperability with MCP Tool format."""
        # Create a mock MCP Tool
        mcp_tool = Tool(
            name="calculator",
            description="adds numbers",
            inputSchema={
                "type": "object",
                "properties": {
                    "a": {"type": "number", "description": "First number"},
                    "b": {"type": "number", "description": "Second number"},
                },
                "required": ["a"],
            },
        )

        # From MCP to ToolCall
        tc = ToolCall.from_mcp_tool(mcp_tool)

        print("\n=== ToolCall from MCP - simple_input_dump() ===")
        print(tc.simple_input_dump())
        print("\n=== ToolCall from MCP - simple_output_dump() ===")
        print(tc.simple_output_dump())

        self.assertEqual(tc.name, "calculator")
        self.assertIn("a", tc.parameters.properties)
        self.assertIn("b", tc.parameters.properties)
        # Check that 'a' is in the required list
        self.assertIn("a", tc.parameters.required)
        self.assertNotIn("b", tc.parameters.required)

        # From ToolCall back to MCP structure (via to_mcp_tool)
        # Note: This checks the logic of constructing the dict for Tool(...)
        mcp_compatible = tc.to_mcp_tool()

        print("\n=== MCP Tool converted back ===")
        print(f"Name: {mcp_compatible.name}")
        print(f"Description: {mcp_compatible.description}")
        print(f"InputSchema: {mcp_compatible.inputSchema}")

        self.assertEqual(mcp_compatible.name, "calculator")
        self.assertIn("a", mcp_compatible.inputSchema["properties"])
        self.assertIn("b", mcp_compatible.inputSchema["properties"])
        self.assertEqual(mcp_compatible.inputSchema["required"], ["a"])


if __name__ == "__main__":
    unittest.main()