Fix multiple handlers.

Source/Request.swift

@@ -1786,16 +1786,6 @@ public final class WebSocketRequest: Request {
         }
     }
 
-//    override func didCancel() {
-//        dispatchPrecondition(condition: .onQueue(underlyingQueue))
-//
-//        mutableState.write { mutableState in
-//            mutableState.error = mutableState.error ?? AFError.explicitlyCancelled
-//        }
-//
-//        eventMonitor?.requestDidCancel(self)
-//    }
-
     func didClose() {
         dispatchPrecondition(condition: .onQueue(underlyingQueue))
 
@@ -1868,13 +1858,11 @@ public final class WebSocketRequest: Request {
             return
         }
 
-        NSLog("Starting ping.")
         let start = Date()
         let startTimestamp = ProcessInfo.processInfo.systemUptime
         socket?.sendPing { error in
             // Calls back on delegate queue / rootQueue / underlyingQueue
             if let error = error {
-                NSLog("Ping error: \(error)")
                 queue.async {
                     onResponse(.error(error))
                 }
@@ -1883,7 +1871,7 @@ public final class WebSocketRequest: Request {
                 let end = Date()
                 let endTimestamp = ProcessInfo.processInfo.systemUptime
                 let pong = PingResponse.Pong(start: start, end: end, latency: endTimestamp - startTimestamp)
-                NSLog("Pong received: \(pong)")
+
                 queue.async {
                     onResponse(.pong(pong))
                 }
@@ -1940,7 +1928,9 @@ public final class WebSocketRequest: Request {
             case let .success(message):
                 self.socketMutableState.read { state in
                     state.handlers.forEach { handler in
-                        handler.queue.async { handler.handler(.receivedMessage(message)) }
+                        handler.queue.async {
+                            handler.handler(.receivedMessage(message))
+                        }
                     }
                 }
 
@@ -2044,15 +2034,11 @@ public final class WebSocketRequest: Request {
 
         appendResponseSerializer {
             self.responseSerializerDidComplete {
-                self.socketMutableState.read { state in
-                    state.handlers.forEach { storedHandler in
-                        storedHandler.queue.async {
-                            storedHandler.handler(.completed(.init(request: self.request,
-                                                                   response: self.response,
-                                                                   metrics: self.metrics,
-                                                                   error: self.error)))
-                        }
-                    }
+                queue.async {
+                    handler(.completed(.init(request: self.request,
+                                             response: self.response,
+                                             metrics: self.metrics,
+                                             error: self.error)))
                 }
             }
         }

Tests/WebSocketTests.swift

@@ -64,6 +64,7 @@ final class WebSocketTests: BaseTestCase {
     func testThatWebSocketsCanReceiveAMessage() {
         // Given
         let didReceiveMessage = expectation(description: "didReceiveMessage")
+        let didComplete = expectation(description: "didComplete")
 
         let session = stored(Session())
 
@@ -74,6 +75,9 @@ final class WebSocketTests: BaseTestCase {
             receivedMessage = message
             didReceiveMessage.fulfill()
         }
+        .onCompletion {
+            didComplete.fulfill()
+        }
 
         waitForExpectations(timeout: timeout)
 
@@ -132,6 +136,7 @@ final class WebSocketTests: BaseTestCase {
     func testThatWebSocketsCanReceiveADecodableValue() {
         // Given
         let didReceiveValue = expectation(description: "didReceiveMessage")
+        let didComplete = expectation(description: "didComplete")
 
         let session = stored(Session())
 
@@ -142,6 +147,9 @@ final class WebSocketTests: BaseTestCase {
             receivedValue = value
             didReceiveValue.fulfill()
         }
+        .onCompletion {
+            didComplete.fulfill()
+        }
 
         waitForExpectations(timeout: timeout)
 
@@ -392,7 +400,7 @@ final class WebSocketTests: BaseTestCase {
         XCTAssertEqual(closeCode, .normalClosure)
         XCTAssertNil(closeReason)
         XCTAssertNotNil(receivedCompletion)
-//        XCTAssertNil(receivedCompletion?.error)
+        XCTAssertNil(receivedCompletion?.error)
         XCTAssertNotNil(receivedPong)
     }
 
@@ -511,6 +519,94 @@ final class WebSocketTests: BaseTestCase {
         XCTAssertNil(closeReason)
         XCTAssertNil(receivedCompletion?.error)
     }
+
+    func testThatWebSocketsCanHaveMultipleHandlers() {
+        // Given
+        let didConnect = expectation(description: "didConnect")
+        didConnect.expectedFulfillmentCount = 2
+        let didReceiveMessage = expectation(description: "didReceiveMessage")
+        didReceiveMessage.expectedFulfillmentCount = 2
+        let didDisconnect = expectation(description: "didDisconnect")
+        didDisconnect.expectedFulfillmentCount = 2
+        let didComplete = expectation(description: "didComplete")
+        didComplete.expectedFulfillmentCount = 2
+        let session = stored(Session())
+
+        var firstConnectedProtocol: String?
+        var firstMessage: URLSessionWebSocketTask.Message?
+        var firstCloseCode: URLSessionWebSocketTask.CloseCode?
+        var firstCloseReason: Data?
+        var firstReceivedCompletion: WebSocketRequest.Completion?
+        var secondConnectedProtocol: String?
+        var secondMessage: URLSessionWebSocketTask.Message?
+        var secondCloseCode: URLSessionWebSocketTask.CloseCode?
+        var secondCloseReason: Data?
+        var secondReceivedCompletion: WebSocketRequest.Completion?
+
+        // When
+        session.websocketRequest(.websocket(closeCode: .goingAway)).streamMessageEvents { event in
+            switch event.kind {
+            case let .connected(`protocol`):
+                firstConnectedProtocol = `protocol`
+                didConnect.fulfill()
+            case let .receivedMessage(receivedMessage):
+                firstMessage = receivedMessage
+                didReceiveMessage.fulfill()
+            case let .disconnected(code, reason):
+                firstCloseCode = code
+                firstCloseReason = reason
+                didDisconnect.fulfill()
+            case let .completed(completion):
+                firstReceivedCompletion = completion
+                didComplete.fulfill()
+            }
+        }
+        .streamMessageEvents { event in
+            switch event.kind {
+            case let .connected(`protocol`):
+                secondConnectedProtocol = `protocol`
+                didConnect.fulfill()
+            case let .receivedMessage(receivedMessage):
+                secondMessage = receivedMessage
+                didReceiveMessage.fulfill()
+            case let .disconnected(code, reason):
+                secondCloseCode = code
+                secondCloseReason = reason
+                didDisconnect.fulfill()
+            case let .completed(completion):
+                secondReceivedCompletion = completion
+                didComplete.fulfill()
+            }
+        }
+
+        wait(for: [didConnect, didReceiveMessage, didDisconnect, didComplete],
+             timeout: timeout,
+             enforceOrder: false)
+
+        // Then
+        XCTAssertNil(firstConnectedProtocol)
+        XCTAssertEqual(firstConnectedProtocol, secondConnectedProtocol)
+        XCTAssertNotNil(firstMessage)
+        XCTAssertEqual(firstMessage, secondMessage)
+        XCTAssertEqual(firstCloseCode, .goingAway)
+        XCTAssertEqual(firstCloseCode, secondCloseCode)
+        XCTAssertNil(firstCloseReason)
+        XCTAssertEqual(firstCloseReason, secondCloseReason)
+        XCTAssertNil(firstReceivedCompletion?.error)
+        XCTAssertNil(secondReceivedCompletion?.error)
+    }
+}
+
+@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
+extension WebSocketRequest {
+    @discardableResult
+    func onCompletion(_ handler: @escaping () -> Void) -> Self {
+        streamMessageEvents { event in
+            guard case .completed = event.kind else { return }
+
+            handler()
+        }
+    }
 }
 
 @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)