Add server ping/pong test coverage

Author: atdrendel

Tests/WebSocketTests/Server/WebSocketServer.swift

@@ -7,6 +7,7 @@ import WebSocketKit
 
 enum WebSocketServerOutput {
     case message(WebSocketMessage)
+    case ping(Data)
     case remoteClose
     case remoteCloseWithReason(WebSocketErrorCode, Data)
 }
@@ -23,6 +24,7 @@ final class WebSocketServer {
 
     // Publisher that repeats everything sent to it by clients.
     private let inputSubject = PassthroughSubject<WebSocketMessage, Never>()
+    private let pongSubject = PassthroughSubject<Data, Never>()
 
     private let eventLoopGroup: EventLoopGroup
     private var channel: Channel?
@@ -52,6 +54,14 @@ final class WebSocketServer {
                     ) else { return }
                     self?.inputSubject.send(.data(data))
                 }
+                ws.onPong { [weak self] _, pong in
+                    var pong = pong
+                    guard let data = pong.readData(
+                        length: pong.readableBytes,
+                        byteTransferStrategy: .copy
+                    ) else { return }
+                    self?.pongSubject.send(data)
+                }
             }.bind(host: "127.0.0.1", port: 0).wait()
     }
 
@@ -69,6 +79,9 @@ final class WebSocketServer {
                 },
                 receiveValue: { output in
                     switch output {
+                    case let .ping(data):
+                        ws.sendPing(data)
+
                     case .remoteClose:
                         do { try ws.close(code: .goingAway).wait() }
                         catch {}
@@ -103,4 +116,8 @@ final class WebSocketServer {
     var inputPublisher: AnyPublisher<WebSocketMessage, Never> {
         inputSubject.eraseToAnyPublisher()
     }
+
+    var pongPublisher: AnyPublisher<Data, Never> {
+        pongSubject.eraseToAnyPublisher()
+    }
 }

Tests/WebSocketTests/SystemWebSocketTests.swift

@@ -160,7 +160,7 @@ class SystemWebSocketTests: XCTestCase {
             else { return XCTFail("Received wrong error: \(error)") }
         }
 
-        await fulfillment(of: [secondCloseEx], timeout: 0.1)
+        await fulfillment(of: [secondCloseEx], timeout: 0.05)
     }
 
     func testDelegateDoesNotReorderOpenAndCloseCallbacks() async throws {
@@ -292,6 +292,81 @@ class SystemWebSocketTests: XCTestCase {
         await fulfillment(of: [receivedEx], timeout: 2)
     }
 
+    func testServerPingReceivesPongAndDoesNotCloseClient() async throws {
+        let closeEx = expectation(description: "Should not close after ping")
+        closeEx.isInverted = true
+        let shouldFailOnClose = Locked(true)
+
+        let (server, client) = try await makeServerAndClient(
+            onClose: { _ in
+                guard shouldFailOnClose.access({ $0 }) else { return }
+                closeEx.fulfill()
+            }
+        )
+        defer { server.shutDown() }
+
+        let pingPayload = Data("server ping".utf8)
+        let pongEx = expectation(description: "Server should receive pong")
+        let pongSub = server.pongPublisher
+            .sink { pong in
+                XCTAssertEqual(pingPayload, pong)
+                pongEx.fulfill()
+            }
+        defer { pongSub.cancel() }
+
+        let readyEx = expectation(description: "Should receive initial app message")
+        let receivedEx = expectation(description: "Should still receive app messages")
+        let receivedSub = client.sink { message in
+            guard case let .text(text) = message else {
+                return XCTFail("Should have received text")
+            }
+
+            switch text {
+            case "ready":
+                readyEx.fulfill()
+
+            case "still open":
+                receivedEx.fulfill()
+
+            default:
+                XCTFail("Received unexpected text: \(text)")
+            }
+        }
+        defer { receivedSub.cancel() }
+
+        let sentEx = expectation(description: "Server should receive client message")
+        let sentSub = server.inputPublisher
+            .sink { message in
+                guard case let .text(text) = message else {
+                    return XCTFail("Should have received text")
+                }
+                XCTAssertEqual("client ready", text)
+                sentEx.fulfill()
+            }
+        defer { sentSub.cancel() }
+
+        try await client.open()
+
+        try await client.send(.text("client ready"))
+        await fulfillment(of: [sentEx], timeout: 2)
+
+        subject.send(.message(.text("ready")))
+        await fulfillment(of: [readyEx], timeout: 2)
+
+        subject.send(.ping(pingPayload))
+        await fulfillment(of: [pongEx], timeout: 2)
+
+        let isOpen = await client.isOpen
+        XCTAssertTrue(isOpen)
+
+        subject.send(.message(.text("still open")))
+        await fulfillment(of: [receivedEx], timeout: 2)
+        await fulfillment(of: [closeEx], timeout: 0.05)
+
+        shouldFailOnClose.access { $0 = false }
+        try await client.close()
+    }
+
     @available(iOS 15.0, macOS 12.0, *)
     func testPushAndReceiveDataWithAsyncPublisher() async throws {
         let (server, client) = try await makeServerAndClient()