Don't reuse web socket task after it's been closed

Author: cbaker6

Sources/ParseSwift/LiveQuery/ParseLiveQuery.swift

@@ -53,7 +53,13 @@ public final class ParseLiveQuery: NSObject {
     let notificationQueue: DispatchQueue
 
     //Task
-    var task: URLSessionWebSocketTask!
+    var task: URLSessionWebSocketTask! {
+        willSet {
+            if newValue == nil && isSocketEstablished == true {
+                isSocketEstablished = false
+            }
+        }
+    }
     var url: URL!
     var clientId: String!
     var attempts: Int = 1 {
@@ -205,14 +211,19 @@ public final class ParseLiveQuery: NSObject {
         close(useDedicatedQueue: false)
         authenticationDelegate = nil
         receiveDelegate = nil
-        URLSession.liveQuery.delegates.removeValue(forKey: task)
+        if task != nil {
+            URLSession.liveQuery.delegates.removeValue(forKey: task)
+        } else {
+            task = nil
+        }
     }
 }
 
 // MARK: Helpers
 @available(macOS 10.15, iOS 13.0, macCatalyst 13.0, watchOS 6.0, tvOS 13.0, *)
 extension ParseLiveQuery {
 
+    /// Current LiveQuery client.
     public private(set) static var client = try? ParseLiveQuery()
 
     var reconnectInterval: Int {
@@ -533,7 +544,10 @@ extension ParseLiveQuery {
                 self.task.cancel()
                 self.isDisconnectedByUser = true
             }
-            URLSession.liveQuery.delegates.removeValue(forKey: self.task)
+            if task != nil {
+                URLSession.liveQuery.delegates.removeValue(forKey: self.task)
+            }
+            self.task = nil
         }
     }
 
@@ -568,7 +582,9 @@ extension ParseLiveQuery {
             if self.isConnected {
                 self.task.cancel()
             }
-            URLSession.liveQuery.delegates.removeValue(forKey: self.task)
+            if self.task != nil {
+                URLSession.liveQuery.delegates.removeValue(forKey: self.task)
+            }
         }
     }
 
@@ -577,6 +593,8 @@ extension ParseLiveQuery {
             self.pendingSubscriptions.append((requestId, record))
             if self.isConnected {
                 URLSession.liveQuery.send(record.messageData, task: self.task, completion: completion)
+            } else {
+                self.open(completion: completion)
             }
         }
     }

Tests/ParseSwiftTests/ParseLiveQueryTests.swift

@@ -50,7 +50,7 @@ class ParseLiveQueryTests: XCTestCase {
                               masterKey: "masterKey",
                               serverURL: url,
                               testing: true)
-        ParseLiveQuery.client = try? ParseLiveQuery(isDefault: true)
+        ParseLiveQuery.setDefault(try ParseLiveQuery(isDefault: true))
     }
 
     override func tearDownWithError() throws {
@@ -314,24 +314,28 @@ class ParseLiveQueryTests: XCTestCase {
             XCTFail("Should be able to get client")
             return
         }
-        client.isSocketEstablished = true // Socket neets to be true
+        client.isSocketEstablished = true // Socket needs to be true
         client.isConnecting = true
         client.attempts = 50
         client.isConnected = true
         client.clientId = "yolo"
-
-        XCTAssertEqual(client.isSocketEstablished, true)
-        XCTAssertEqual(client.isConnecting, false)
-        XCTAssertEqual(client.clientId, "yolo")
-        XCTAssertEqual(client.attempts, 1)
+        let expectation1 = XCTestExpectation(description: "Synch")
+        client.synchronizationQueue.sync {
+            XCTAssertEqual(client.isSocketEstablished, true)
+            XCTAssertEqual(client.isConnecting, false)
+            XCTAssertEqual(client.clientId, "yolo")
+            XCTAssertEqual(client.attempts, 1)
+            expectation1.fulfill()
+        }
+        wait(for: [expectation1], timeout: 20.0)
     }
 
     func testDisconnectedState() throws {
         guard let client = ParseLiveQuery.getDefault() else {
             XCTFail("Should be able to get client")
             return
         }
-        client.isSocketEstablished = true // Socket neets to be true
+        client.isSocketEstablished = true // Socket needs to be true
         client.isConnecting = true
         client.isConnected = true
         client.clientId = "yolo"
@@ -352,7 +356,7 @@ class ParseLiveQueryTests: XCTestCase {
             XCTFail("Should be able to get client")
             return
         }
-        client.isSocketEstablished = true // Socket neets to be true
+        client.isSocketEstablished = true // Socket needs to be true
         client.isConnecting = true
         client.isConnected = true
         client.clientId = "yolo"
@@ -384,11 +388,13 @@ class ParseLiveQueryTests: XCTestCase {
         XCTAssertEqual(client.clientId, "yolo")
         client.close()
 
-        XCTAssertEqual(client.isSocketEstablished, true)
-        XCTAssertEqual(client.isConnected, false)
-        XCTAssertEqual(client.isConnecting, false)
-        XCTAssertNil(client.clientId)
-        XCTAssertEqual(client.isDisconnectedByUser, true)
+        DispatchQueue.main.asyncAfter(deadline: .now()) {
+            XCTAssertEqual(client.isSocketEstablished, true)
+            XCTAssertEqual(client.isConnected, false)
+            XCTAssertEqual(client.isConnecting, false)
+            XCTAssertNil(client.clientId)
+            XCTAssertEqual(client.isDisconnectedByUser, true)
+        }
     }
 
     func testReconnectInterval() throws {
@@ -584,6 +590,87 @@ class ParseLiveQueryTests: XCTestCase {
         wait(for: [expectation1, expectation2], timeout: 20.0)
     }
 
+    func testSubscribeCloseSubscribe() throws {
+        let query = GameScore.query("score" > 9)
+        let handler = SubscriptionCallback(query: query)
+        var subscription = try Query<GameScore>.subscribe(handler)
+
+        guard let client = ParseLiveQuery.getDefault() else {
+            XCTFail("Should be able to get client")
+            return
+        }
+        XCTAssertEqual(subscription.query, query)
+
+        let expectation1 = XCTestExpectation(description: "Subscribe Handler")
+        let expectation2 = XCTestExpectation(description: "Resubscribe Handler")
+        var count = 0
+        var originalTask: URLSessionWebSocketTask?
+        subscription.handleSubscribe { subscribedQuery, isNew in
+            XCTAssertEqual(query, subscribedQuery)
+            if count == 0 {
+                XCTAssertTrue(isNew)
+                XCTAssertEqual(client.pendingSubscriptions.count, 0)
+                XCTAssertEqual(client.subscriptions.count, 1)
+                XCTAssertNotNil(ParseLiveQuery.client?.task)
+                originalTask = ParseLiveQuery.client?.task
+                expectation1.fulfill()
+            } else if count == 2 {
+                XCTAssertNotNil(ParseLiveQuery.client?.task)
+                XCTAssertFalse(originalTask == ParseLiveQuery.client?.task)
+                expectation2.fulfill()
+                return
+            }
+
+            ParseLiveQuery.client?.close()
+            ParseLiveQuery.client?.synchronizationQueue.sync {
+            XCTAssertNil(ParseLiveQuery.client?.task)
+            if let socketEstablished = ParseLiveQuery.client?.isSocketEstablished {
+                XCTAssertFalse(socketEstablished)
+            } else {
+                XCTFail("Should have socket that isn't established")
+            }
+
+            //Resubscribe
+            do {
+                count += 1
+                subscription = try Query<GameScore>.subscribe(handler)
+            } catch {
+                XCTFail("\(error)")
+            }
+
+            try? self.pretendToBeConnected()
+
+            let response2 = PreliminaryMessageResponse(op: .subscribed,
+                                                       requestId: 2,
+                                                       clientId: "yolo",
+                                                       installationId: "naw")
+            guard let encoded2 = try? ParseCoding.jsonEncoder().encode(response2) else {
+                expectation2.fulfill()
+                return
+            }
+            client.received(encoded2)
+            }
+        }
+
+        XCTAssertFalse(try client.isSubscribed(query))
+        XCTAssertTrue(try client.isPendingSubscription(query))
+        XCTAssertEqual(client.subscriptions.count, 0)
+        XCTAssertEqual(client.pendingSubscriptions.count, 1)
+        try pretendToBeConnected()
+        let response = PreliminaryMessageResponse(op: .subscribed,
+                                                           requestId: 1,
+                                                           clientId: "yolo",
+                                                           installationId: "naw")
+        let encoded = try ParseCoding.jsonEncoder().encode(response)
+        client.received(encoded)
+        XCTAssertTrue(try client.isSubscribed(query))
+        XCTAssertFalse(try client.isPendingSubscription(query))
+        XCTAssertEqual(client.subscriptions.count, 1)
+        XCTAssertEqual(client.pendingSubscriptions.count, 0)
+
+        wait(for: [expectation1, expectation2], timeout: 20.0)
+    }
+
     func testServerRedirectResponse() throws {
         guard let client = ParseLiveQuery.getDefault() else {
             XCTFail("Should be able to get client")