Publish response from the WebSocketTransport to the ApolloStore (#1889)

* publish response from the WebsocketTransport to the Apollo store

* fix build by making methods public

* update tests

* update StarWars tests

* update Playground

* make the `store` optional to init `WebSocketTransport`

Author: tgyhlsb

Playgrounds/ApolloMacPlayground.playground/Pages/Subscriptions.xcplaygroundpage/Contents.swift

@@ -31,7 +31,10 @@ let normalTransport = RequestChainNetworkTransport(interceptorProvider: DefaultI
 //: Next, set up the `WebSocketTransport` to talk to the websocket endpoint. Note that this may take a different URL, sometimes with a `ws` prefix, than your normal http endpoint:
 
 let webSocketURL = URL(string: "ws://localhost:8080/websocket")!
-let webSocketTransport = WebSocketTransport(request: URLRequest(url: webSocketURL))
+let webSocketTransport = WebSocketTransport(
+    request: URLRequest(url: webSocketURL),
+    store: store
+)
 
 //: Then, set up the split transport with the two transports you've just created:
 

Sources/Apollo/ApolloStore.swift

@@ -64,7 +64,7 @@ public final class ApolloStore {
     }
   }
 
-  func publish(records: RecordSet, identifier: UUID? = nil, callbackQueue: DispatchQueue = .main, completion: ((Result<Void, Error>) -> Void)? = nil) {
+  public func publish(records: RecordSet, identifier: UUID? = nil, callbackQueue: DispatchQueue = .main, completion: ((Result<Void, Error>) -> Void)? = nil) {
     queue.async(flags: .barrier) {
       do {
         let changedKeys = try self.cache.merge(records: records)

Sources/Apollo/GraphQLResponse.swift

@@ -19,7 +19,7 @@ public final class GraphQLResponse<Data: GraphQLSelectionSet> {
     self.variables = operation.variables
   }
 
-  func parseResult(cacheKeyForObject: CacheKeyForObject? = nil) throws -> (GraphQLResult<Data>, RecordSet?) {
+  public func parseResult(cacheKeyForObject: CacheKeyForObject? = nil) throws -> (GraphQLResult<Data>, RecordSet?) {
     let errors: [GraphQLError]?
 
     if let errorsEntry = body["errors"] as? [JSONObject] {

Sources/ApolloWebSocket/WebSocketTransport.swift

@@ -29,6 +29,7 @@ public class WebSocketTransport {
   let connectOnInit: Bool
   let reconnect: Atomic<Bool>
   let websocket: WebSocketClient
+  let store: ApolloStore?
   let error: Atomic<Error?> = Atomic(nil)
   let serializationFormat = JSONSerializationFormat.self
   private let requestBodyCreator: RequestBodyCreator
@@ -88,6 +89,7 @@ public class WebSocketTransport {
   ///   - connectingPayload: [optional] The payload to send on connection. Defaults to an empty `GraphQLMap`.
   ///   - requestBodyCreator: The `RequestBodyCreator` to use when serializing requests. Defaults to an `ApolloRequestBodyCreator`.
   public init(websocket: WebSocketClient,
+              store: ApolloStore? = nil,
               clientName: String = WebSocketTransport.defaultClientName,
               clientVersion: String = WebSocketTransport.defaultClientVersion,
               sendOperationIdentifiers: Bool = false,
@@ -98,6 +100,7 @@ public class WebSocketTransport {
               connectingPayload: GraphQLMap? = [:],
               requestBodyCreator: RequestBodyCreator = ApolloRequestBodyCreator()) {
     self.websocket = websocket
+    self.store = store
     self.connectingPayload = connectingPayload
     self.sendOperationIdentifiers = sendOperationIdentifiers
     self.reconnect = Atomic(reconnect)
@@ -359,10 +362,21 @@ extension WebSocketTransport: NetworkTransport {
       return EmptyCancellable()
     }
 
-    return WebSocketTask(self, operation) { result in
+    return WebSocketTask(self, operation) { [store] result in
       switch result {
       case .success(let jsonBody):
         let response = GraphQLResponse(operation: operation, body: jsonBody)
+        if let store = store {
+          do {
+            let (_, records) = try response.parseResult(cacheKeyForObject: store.cacheKeyForObject)
+            if let records = records {
+              store.publish(records: records, identifier: nil)
+            }
+          } catch {
+            callCompletion(with: .failure(error))
+          }
+        }
+        
         do {
           let graphQLResult = try response.parseResultFast()
           callCompletion(with: .success(graphQLResult))

Tests/ApolloServerIntegrationTests/StarWarsSubscriptionTests.swift

@@ -24,7 +24,8 @@ class StarWarsSubscriptionTests: XCTestCase {
     webSocketTransport = WebSocketTransport(
       websocket: DefaultWebSocket(
         request: URLRequest(url: TestServerURL.starWarsWebSocket.url)
-      )
+      ),
+      store: ApolloStore()
     )
     webSocketTransport.delegate = self
     client = ApolloClient(networkTransport: webSocketTransport, store: ApolloStore())
@@ -411,7 +412,8 @@ class StarWarsSubscriptionTests: XCTestCase {
     let webSocketTransport = WebSocketTransport(
       websocket: MockWebSocket(
         request: URLRequest(url: TestServerURL.starWarsWebSocket.url)
-      )
+      ),
+      store: ApolloStore()
     )
 
     let expectation = self.expectation(description: "Connection closed")

Tests/ApolloServerIntegrationTests/StarWarsWebSocketTests.swift

@@ -24,7 +24,8 @@ class StarWarsWebSocketTests: XCTestCase, CacheDependentTesting {
     let networkTransport = WebSocketTransport(
       websocket: DefaultWebSocket(
         request: URLRequest(url: TestServerURL.starWarsWebSocket.url)
-      )
+      ),
+      store: store
     )
 
     client = ApolloClient(networkTransport: networkTransport, store: store)

Tests/ApolloTests/WebSocket/WebSocketTests.swift

@@ -21,9 +21,10 @@ class WebSocketTests: XCTestCase {
   override func setUp() {
     super.setUp()
 
+    let store = ApolloStore()
     let websocket = MockWebSocket(request:URLRequest(url: TestURL.mockServer.url))
-    networkTransport = WebSocketTransport(websocket: websocket)
-    client = ApolloClient(networkTransport: networkTransport!, store: ApolloStore())
+    networkTransport = WebSocketTransport(websocket: websocket, store: store)
+    client = ApolloClient(networkTransport: networkTransport!, store: store)
   }
 
   override func tearDown() {

Tests/ApolloTests/WebSocket/WebSocketTransportTests.swift

@@ -17,7 +17,8 @@ class WebSocketTransportTests: XCTestCase {
     var request = URLRequest(url: TestURL.mockServer.url)
     request.addValue("OldToken", forHTTPHeaderField: "Authorization")
 
-    self.webSocketTransport = WebSocketTransport(websocket: MockWebSocket(request: request))
+    self.webSocketTransport = WebSocketTransport(websocket: MockWebSocket(request: request),
+                                                 store: ApolloStore())
 
     self.webSocketTransport.updateHeaderValues(["Authorization": "UpdatedToken"])
 
@@ -28,6 +29,7 @@ class WebSocketTransportTests: XCTestCase {
     let request = URLRequest(url: TestURL.mockServer.url)
 
     self.webSocketTransport = WebSocketTransport(websocket: MockWebSocket(request: request),
+                                                 store: ApolloStore(),
                                                  connectingPayload: ["Authorization": "OldToken"])
 
     let mockWebSocketDelegate = MockWebSocketDelegate()
@@ -58,6 +60,7 @@ class WebSocketTransportTests: XCTestCase {
     let request = URLRequest(url: TestURL.mockServer.url)
 
     self.webSocketTransport = WebSocketTransport(websocket: MockWebSocket(request: request),
+                                                 store: ApolloStore(),
                                                  connectingPayload: ["Authorization": "OldToken"])
     self.webSocketTransport.closeConnection()
     self.webSocketTransport.updateConnectingPayload(["Authorization": "UpdatedToken"])