chore: improved unit tests a little bit

Author: yehorsobko-mac

Sources/MCP/Base/Transports/NetworkTransport.swift

@@ -349,16 +349,20 @@ import Logging
         /// - Throws: Error if the connection fails
         private func attemptConnection() async throws {
             // Create stream for state changes with proper cleanup
-            let (stateStream, continuation) = AsyncStream.makeStream(of: NWConnection.State.self)
-
-            connection.stateUpdateHandler = { state in
-                continuation.yield(state)
-            }
+            let stateStream = AsyncStream<NWConnection.State> { continuation in
+                continuation.onTermination = { [weak self] _ in
+                    self?.connection.stateUpdateHandler = nil
+                }
 
-            // Ensure cleanup happens when we exit
-            defer {
-                continuation.finish()
-                connection.stateUpdateHandler = nil
+                connection.stateUpdateHandler = { state in
+                    continuation.yield(state)
+                    switch state {
+                    case .ready, .failed, .cancelled:
+                        continuation.finish()
+                    default:
+                        break
+                    }
+                }
             }
 
             connection.start(queue: .main)

Tests/MCPTests/ClientTests.swift

@@ -3,7 +3,7 @@ import Testing
 
 @testable import MCP
 
-@Suite("Client Tests")
+@Suite("Client Tests", .timeLimit(.minutes(1)))
 struct ClientTests {
     @Test("Client connect and disconnect")
     func testClientConnectAndDisconnect() async throws {

Tests/MCPTests/NetworkTransportTests.swift

@@ -268,7 +268,7 @@ import Testing
 
             // Simulate failure before connecting
             mockConnection.simulateFailure(error: NWError.posix(POSIXErrorCode.ECONNRESET))
-
+            
             do {
                 try await transport.connect()
                 Issue.record("Expected connect to throw an error")
@@ -496,33 +496,6 @@ import Testing
             _ = await receiveTask.result
         }
 
-        @Test("Connection State Transitions")
-        func testConnectionStateTransitions() async throws {
-            let mockConnection = MockNetworkConnection()
-            let transport = NetworkTransport(
-                mockConnection,
-                heartbeatConfig: .disabled
-            )
-
-            // Test setup -> preparing -> ready transition
-            mockConnection.simulatePreparing()
-            try await Task.sleep(for: .milliseconds(100))
-            mockConnection.simulateReady()
-            try await transport.connect()
-            #expect(mockConnection.state == .ready)
-
-            // Test ready -> failed transition
-            mockConnection.simulateFailure(error: NWError.posix(POSIXErrorCode.ECONNRESET))
-            try await Task.sleep(for: .milliseconds(100))
-            if case .failed = mockConnection.state {
-                // expected
-            } else {
-                Issue.record("Expected state to be failed")
-            }
-
-            await transport.disconnect()
-        }
-
         @Test("Partial Message Reception")
         func testPartialMessageReception() async throws {
             let mockConnection = MockNetworkConnection()

Tests/MCPTests/ServerTests.swift

@@ -39,7 +39,7 @@ struct ServerTests {
         try await server.start(transport: transport)
 
         // Wait for message processing and response
-        try await Task.sleep(for: .milliseconds(100))
+        try await Task.sleep(for: .milliseconds(200))
 
         #expect(await transport.sentMessages.count == 1)
 
@@ -143,49 +143,66 @@ struct ServerTests {
 
     @Test("JSON-RPC batch processing")
     func testJSONRPCBatchProcessing() async throws {
-        let transport = MockTransport()
+        let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair()
         let server = Server(name: "TestServer", version: "1.0")
 
+        // Connect transports
+        try await clientTransport.connect()
+        try await serverTransport.connect()
+
+        // Start receiving messages on client side
+        let receiveTask = Task {
+            var responses: [String] = []
+            for try await data in await clientTransport.receive() {
+                if let response = String(data: data, encoding: .utf8) {
+                    responses.append(response)
+                }
+                // Stop after receiving 2 responses (initialize + batch)
+                if responses.count == 2 {
+                    break
+                }
+            }
+            return responses
+        }
+
         // Start the server
-        try await server.start(transport: transport)
+        try await server.start(transport: serverTransport)
 
         // Initialize the server first
-        try await transport.queue(
-            request: Initialize.request(
-                .init(
-                    protocolVersion: Version.latest,
-                    capabilities: .init(),
-                    clientInfo: .init(name: "TestClient", version: "1.0")
-                )
+        let initRequest = Initialize.request(
+            .init(
+                protocolVersion: Version.latest,
+                capabilities: .init(),
+                clientInfo: .init(name: "TestClient", version: "1.0")
             )
         )
+        let initData = try JSONEncoder().encode(AnyRequest(initRequest))
+        try await clientTransport.send(initData)
 
-        // Wait for server to initialize and respond
+        // Wait for initialization
         try await Task.sleep(for: .milliseconds(100))
 
-        // Clear sent messages
-        await transport.clearMessages()
-
         // Create a batch with multiple requests
         let batchJSON = """
             [
                 {"jsonrpc":"2.0","id":1,"method":"ping","params":{}},
                 {"jsonrpc":"2.0","id":2,"method":"ping","params":{}}
             ]
             """
-        let batch = try JSONDecoder().decode([AnyRequest].self, from: batchJSON.data(using: .utf8)!)
-
-        // Send the batch request
-        try await transport.queue(batch: batch)
+        let batchData = batchJSON.data(using: .utf8)!
+        try await clientTransport.send(batchData)
 
         // Wait for batch processing
-        try await Task.sleep(for: .milliseconds(100))
+        try await Task.sleep(for: .milliseconds(200))
 
-        // Verify response
-        let sentMessages = await transport.sentMessages
-        #expect(sentMessages.count == 1)
+        // Get responses
+        let responses = try await receiveTask.value
+        #expect(responses.count == 2)
+
+        // Verify the batch response (second response)
+        if responses.count >= 2 {
+            let batchResponse = responses[1]
 
-        if let batchResponse = sentMessages.first {
             // Should be an array
             #expect(batchResponse.hasPrefix("["))
             #expect(batchResponse.hasSuffix("]"))
@@ -196,6 +213,7 @@ struct ServerTests {
         }
 
         await server.stop()
-        await transport.disconnect()
+        await clientTransport.disconnect()
+        await serverTransport.disconnect()
     }
 }