Add tests for ChatAgent

Author: RussellLuo

coagent/core/types.py

@@ -293,6 +293,7 @@ async def run(
         msg: RawMessage,
         stream: bool = False,
         session_id: str = "",
+        request: bool = True,
         timeout: float = 0.5,
     ) -> AsyncIterator[RawMessage] | RawMessage | None:
         """Create an agent and run it with the given message."""
@@ -303,7 +304,7 @@ async def run(
         addr = Address(name=self.name, id=session_id)
 
         return await self.__runtime.channel.publish(
-            addr, msg, stream=stream, request=True, timeout=timeout
+            addr, msg, stream=stream, request=request, timeout=timeout
         )
 
 

tests/agents/test_chat_agent.py

@@ -1,9 +1,21 @@
+import sys
 from typing import AsyncIterator
 
 from pydantic import Field
 import pytest
 
-from coagent.agents.chat_agent import wrap_error
+from coagent.agents import ChatAgent
+from coagent.agents.chat_agent import (
+    CallTool,
+    CallToolResult,
+    ListToolsResult,
+    MCPTextContent,
+    MCPTool,
+    wrap_error,
+)
+from coagent.core import Address, RawMessage
+from coagent.core.runtime import NopChannel
+import jsonschema
 
 
 @pytest.mark.asyncio
@@ -37,3 +49,91 @@ async def func(
 
     result = await func(a=1, b=0)
     assert await anext(result) == "Error: division by zero"
+
+
+class MCPServerTestChannel(NopChannel):
+    async def publish(
+        self,
+        addr: Address,
+        msg: RawMessage,
+        stream: bool = False,
+        request: bool = False,
+        reply: str = "",
+        timeout: float = 0.5,
+        probe: bool = True,
+    ) -> AsyncIterator[RawMessage] | RawMessage | None:
+        match msg.header.type:
+            case "ListTools":
+                return ListToolsResult(
+                    tools=[
+                        MCPTool(
+                            name="query_weather",
+                            description="Query the weather in the given city.",
+                            inputSchema={
+                                "title": "query_weatherArguments",
+                                "type": "object",
+                                "properties": {
+                                    "city": {
+                                        "title": "City",
+                                        "type": "string",
+                                    },
+                                },
+                                "required": ["city"],
+                            },
+                        )
+                    ]
+                ).encode()
+
+            case "CallTool":
+                call_tool = CallTool.decode(msg)
+                city = call_tool.arguments["city"]
+                return CallToolResult(
+                    content=[
+                        MCPTextContent(
+                            type="text", text=f"The weather in {city} is sunny."
+                        )
+                    ],
+                ).encode()
+
+
+class TestChatAgent:
+    @pytest.mark.skipif(sys.platform == "win32", reason="Does not run on Windows.")
+    @pytest.mark.asyncio
+    async def test_get_mcp_tools(self):
+        agent = ChatAgent()
+        addr = Address(name="test", id="0")
+        agent.init(MCPServerTestChannel(), addr)
+
+        tools = await agent._get_mcp_tools(["server1"])
+        assert len(tools) == 1
+        # Tool query_weather
+        tool = tools[0]
+
+        # Validate the tool
+        assert tool.__name__ == "query_weather"
+        assert tool.__doc__ == "Query the weather in the given city."
+        assert tool.__mcp_tool_schema__ == {
+            "description": "Query the weather in the given city.",
+            "name": "query_weather",
+            "parameters": {
+                "properties": {
+                    "city": {
+                        "title": "City",
+                        "type": "string",
+                    }
+                },
+                "required": ["city"],
+                "title": "query_weatherArguments",
+                "type": "object",
+            },
+        }
+        assert tool.__mcp_tool_args__ == ("city",)
+
+        # Call the tool with no arguments
+        with pytest.raises(jsonschema.exceptions.ValidationError) as exc:
+            await tool()
+        assert str(exc.value).startswith("'city' is a required property")
+
+        # Call the tool with required arguments
+        result = await tool(city="Beijing")
+        assert result == "The weather in Beijing is sunny."