diff --git a/Tests/WebSocketTests/Server/WebSocketServer.swift b/Tests/WebSocketTests/Server/WebSocketServer.swift index 03dbf16..1271d2b 100644 --- a/Tests/WebSocketTests/Server/WebSocketServer.swift +++ b/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() + private let pongSubject = PassthroughSubject() 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 { inputSubject.eraseToAnyPublisher() } + + var pongPublisher: AnyPublisher { + pongSubject.eraseToAnyPublisher() + } } diff --git a/Tests/WebSocketTests/SystemWebSocketTests.swift b/Tests/WebSocketTests/SystemWebSocketTests.swift index c02d61f..f0317c1 100644 --- a/Tests/WebSocketTests/SystemWebSocketTests.swift +++ b/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()