Expand roundtrip tests

Author: mattt

Tests/MCPTests/RoundtripTests.swift

@@ -8,10 +8,9 @@ import Testing
 @Suite("Roundtrip Tests")
 struct RoundtripTests {
     @Test(
-        "Initialize roundtrip",
         .timeLimit(.minutes(1))
     )
-    func testInitializeRoundtrip() async throws {
+    func testRoundtrip() async throws {
         let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe()
         let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe()
 
@@ -36,6 +35,32 @@ struct RoundtripTests {
             version: "1.0.0",
             capabilities: .init(prompts: .init(), tools: .init())
         )
+        await server.withMethodHandler(ListTools.self) { _ in
+            return ListTools.Result(tools: [
+                Tool(
+                    name: "add",
+                    description: "Adds two numbers together",
+                    inputSchema: [
+                        "a": ["type": "integer", "description": "The first number"],
+                        "a": ["type": "integer", "description": "The second number"],
+                    ])
+            ])
+        }
+        await server.withMethodHandler(CallTool.self) { request in
+            guard request.name == "add" else {
+                return CallTool.Result(content: [.text("Invalid tool name")], isError: true)
+            }
+
+            guard let a = request.arguments?["a"]?.intValue,
+                  let b = request.arguments?["b"]?.intValue
+            else {
+                return CallTool.Result(
+                    content: [.text("Did not receive valid arguments")], isError: true)
+            }
+
+            return CallTool.Result(content: [.text("\(a + b)")], isError: false)
+        }
+
         let client = Client(name: "TestClient", version: "1.0")
 
         try await server.start(transport: serverTransport)
@@ -63,6 +88,31 @@ struct RoundtripTests {
             group.cancelAll()
         }
 
+        let listToolsTask = Task {
+            let result = try await client.listTools()
+            #expect(result.count == 1)
+            #expect(result[0].name == "add")
+        }
+
+        let callToolTask = Task {
+            let result = try await client.callTool(name: "add", arguments: ["a": 1, "b": 2])
+            #expect(result.isError == false)
+            #expect(result.content == [.text("3")])
+        }
+
+        try await withThrowingTaskGroup(of: Void.self) { group in
+            group.addTask {
+                try await Task.sleep(for: .seconds(1))
+                listToolsTask.cancel()
+                throw CancellationError()
+            }
+            group.addTask {
+                try await callToolTask.value
+            }
+            try await group.next()
+            group.cancelAll()
+        }
+
         await server.stop()
         await client.disconnect()
         try? clientToServerRead.close()