diff --git a/Sources/ApolloWebSocket/WebSocketTransport.swift b/Sources/ApolloWebSocket/WebSocketTransport.swift index 74101f2e68..9f8b98adce 100644 --- a/Sources/ApolloWebSocket/WebSocketTransport.swift +++ b/Sources/ApolloWebSocket/WebSocketTransport.swift @@ -323,6 +323,23 @@ public class WebSocketTransport { reconnect.value = oldReconnectValue } + + /// Disconnects the websocket while setting the auto-reconnect value to false, + /// allowing purposeful disconnects that do not dump existing subscriptions. + /// NOTE: You will receive an error on the subscription (should be a `Starscream.WSError` with code 1000) when the socket disconnects. + /// ALSO NOTE: To reconnect after calling this, you will need to call `resumeWebSocketConnection`. + public func pauseWebSocketConnection() { + self.reconnect.value = false + self.websocket.disconnect() + } + + /// Reconnects a paused web socket. + /// + /// - Parameter autoReconnect: `true` if you want the websocket to automatically reconnect if the connection drops. Defaults to true. + public func resumeWebSocketConnection(autoReconnect: Bool = true) { + self.reconnect.value = autoReconnect + self.websocket.connect() + } } // MARK: - HTTPNetworkTransport conformance diff --git a/Tests/ApolloWebsocketTests/StarWarsSubscriptionTests.swift b/Tests/ApolloWebsocketTests/StarWarsSubscriptionTests.swift index 7cb62c67cd..ec3c40cc5d 100644 --- a/Tests/ApolloWebsocketTests/StarWarsSubscriptionTests.swift +++ b/Tests/ApolloWebsocketTests/StarWarsSubscriptionTests.swift @@ -3,15 +3,18 @@ import Apollo import ApolloTestSupport @testable import ApolloWebSocket import StarWarsAPI +import Starscream class StarWarsSubscriptionTests: XCTestCase { - let SERVER: String = "ws://localhost:8080/websocket" + let SERVER = "ws://localhost:8080/websocket" let concurrentQueue = DispatchQueue(label: "com.apollographql.testing", attributes: .concurrent) var client: ApolloClient! var webSocketTransport: WebSocketTransport! var connectionStartedExpectation: XCTestExpectation? + var disconnectedExpectation: XCTestExpectation? + var reconnectedExpectation: XCTestExpectation? override func setUp() { super.setUp() @@ -408,6 +411,72 @@ class StarWarsSubscriptionTests: XCTestCase { waitForExpectations(timeout: 10, handler: nil) } + + func testPausingAndResumingWebSocketConnection() { + let subscription = ReviewAddedSubscription() + let reviewMutation = CreateAwesomeReviewMutation() + + // Send the mutations via a separate transport so they can still be sent when the websocket is disconnected + let httpTransport = HTTPNetworkTransport(url: URL(string: "http://localhost:8080/graphql")!) + let httpClient = ApolloClient(networkTransport: httpTransport) + + func sendReview() { + let reviewSentExpectation = self.expectation(description: "review sent") + httpClient.perform(mutation: reviewMutation) { mutationResult in + switch mutationResult { + case .success: + break + case .failure(let error): + XCTFail("Unexpected error sending review: \(error)") + } + + reviewSentExpectation.fulfill() + } + self.wait(for: [reviewSentExpectation], timeout: 10) + } + + let subscriptionExpectation = self.expectation(description: "Received review") + // This should get hit twice - once before we pause the web socket and once after. + subscriptionExpectation.expectedFulfillmentCount = 2 + let reviewAddedSubscription = self.client.subscribe(subscription: subscription) { subscriptionResult in + switch subscriptionResult { + case .success(let graphQLResult): + XCTAssertEqual(graphQLResult.data?.reviewAdded?.episode, .jedi) + subscriptionExpectation.fulfill() + case .failure(let error): + if let wsError = error as? Starscream.WSError { + // This is an expected error on disconnection, ignore it. + XCTAssertEqual(wsError.code, 1000) + } else { + XCTFail("Unexpected error receiving subscription: \(error)") + subscriptionExpectation.fulfill() + } + } + } + + self.waitForSubscriptionsToStart() + sendReview() + + self.disconnectedExpectation = self.expectation(description: "Web socket disconnected") + webSocketTransport.pauseWebSocketConnection() + self.wait(for: [self.disconnectedExpectation!], timeout: 10) + + // This should not go through since the socket is paused + sendReview() + + self.reconnectedExpectation = self.expectation(description: "Web socket reconnected") + webSocketTransport.resumeWebSocketConnection() + self.wait(for: [self.reconnectedExpectation!], timeout: 10) + self.waitForSubscriptionsToStart() + + // Now that we've reconnected, this should go through to the same subscription. + sendReview() + + self.wait(for: [subscriptionExpectation], timeout: 10) + + // Cancel subscription so it doesn't keep receiving from other tests. + reviewAddedSubscription.cancel() + } } extension StarWarsSubscriptionTests: WebSocketTransportDelegate { @@ -415,4 +484,12 @@ extension StarWarsSubscriptionTests: WebSocketTransportDelegate { func webSocketTransportDidConnect(_ webSocketTransport: WebSocketTransport) { self.connectionStartedExpectation?.fulfill() } + + func webSocketTransportDidReconnect(_ webSocketTransport: WebSocketTransport) { + self.reconnectedExpectation?.fulfill() + } + + func webSocketTransport(_ webSocketTransport: WebSocketTransport, didDisconnectWithError error: Error?) { + self.disconnectedExpectation?.fulfill() + } }