fixup! [CIS-224] Implement token refresh

GetStream · May 11, 2021 · 9b186ef · 9b186ef
1 parent d154056
commit 9b186ef
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 35 deletions.
diff --git a/Sources/StreamChat/ChatClient.swift b/Sources/StreamChat/ChatClient.swift
@@ -411,38 +411,51 @@ extension _ChatClient: ConnectionStateDelegate {
     func webSocketClient(_ client: WebSocketClient, didUpdateConectionState state: WebSocketConnectionState) {
         connectionStatus = .init(webSocketConnectionState: state)
 
-        // Reacquire a new token if current one has expired
-        if case let .disconnected(error) = state,
-           let error = error,
-           error.isTokenExpiredError {
-            clientUpdater.reloadUserIfNeeded()
+        // We should notify waiters if connectionId was obtained (i.e. state is .connected)
+        // or for .disconnected state except for disconnect caused by an expired token
+        let shouldNotifyConnectionIdWaiters: Bool
+        let connectionId: String?
+        switch state {
+        case let .connected(connectionId: id):
+            shouldNotifyConnectionIdWaiters = true
+            connectionId = id
+        case let .disconnected(error: error):
+            if let error = error,
+               error.isTokenExpiredError {
+                clientUpdater.reloadUserIfNeeded()
+                shouldNotifyConnectionIdWaiters = false
+            } else {
+                shouldNotifyConnectionIdWaiters = true
+            }
+            connectionId = nil
+        case .connecting,
+             .disconnecting,
+             .waitingForConnectionId,
+             .waitingForReconnect:
+            shouldNotifyConnectionIdWaiters = false
+            connectionId = nil
         }
-        norifyConnectionIdWaitersIfNeeded(for: state)
+
+        updateConnectionId(
+            connectionId: connectionId,
+            shouldNotifyWaiters: shouldNotifyConnectionIdWaiters
+        )
     }
 
-    private func norifyConnectionIdWaitersIfNeeded(for state: WebSocketConnectionState) {
+    /// Update connectionId and notify waiters if needed
+    /// - Parameters:
+    ///   - connectionId: new connectionId (if present)
+    ///   - shouldFailWaiters: Whether it's necessary to notify waiters or not
+    private func updateConnectionId(
+        connectionId: String?,
+        shouldNotifyWaiters: Bool
+    ) {
         _connectionId.mutate { mutableConnectionId in
             _connectionIdWaiters.mutate { connectionIdWaiters in
-
-                if case let .connected(connectionId) = state {
-                    mutableConnectionId = connectionId
+                mutableConnectionId = connectionId
+                if shouldNotifyWaiters {
                     connectionIdWaiters.forEach { $0(connectionId) }
                     connectionIdWaiters.removeAll()
-
-                } else {
-                    mutableConnectionId = nil
-
-                    if case let .disconnected(error) = state {
-                        // Do not notify waiters in case of token expired error
-                        if let error = error,
-                           error.isTokenExpiredError {
-                            return
-                        }
-
-                        // No reconnection attempt schedule, we should fail all existing connectionId waiters.
-                        connectionIdWaiters.forEach { $0(nil) }
-                        connectionIdWaiters.removeAll()
-                    }
                 }
             }
         }

diff --git a/Sources/StreamChat/ChatClient_Tests.swift b/Sources/StreamChat/ChatClient_Tests.swift
@@ -392,7 +392,7 @@ class ChatClient_Tests: StressTestCase {
         XCTAssertNil(providedConnectionId)
     }
 
-    func test_client_webSocketIsDisconnected_becauseTokenExpired_callsReloadUserIfNeeded() {
+    func test_client_webSocketIsDisconnected_becauseTokenExpired_callsReloadUserIfNeeded() throws {
         // Create a new chat client
         let client = ChatClient(
             config: inMemoryStorageConfig,
@@ -405,12 +405,12 @@ class ChatClient_Tests: StressTestCase {
         // Simulate access to `webSocketClient` so it is initialized
         _ = client.webSocketClient
 
-        // Set a connection Id waiter and set `providedConnectionId` to a non-nil value
-        var providedConnectionId: ConnectionId? = .unique
-        client.provideConnectionId {
-            providedConnectionId = $0
-        }
-        XCTAssertNotNil(providedConnectionId)
+        // Simulate .connected state to obtain connection id
+        let connectionId: ConnectionId = .unique
+        testEnv.webSocketClient?.connectionStateDelegate?
+            .webSocketClient(testEnv.webSocketClient!, didUpdateConectionState: .connected(connectionId: connectionId))
+
+        XCTAssertEqual(client.connectionId, connectionId)
 
         // Was called on ChatClient init
         XCTAssertEqual(testEnv.clientUpdater!.reloadUserIfNeeded_callsCount, 1)
@@ -427,9 +427,8 @@ class ChatClient_Tests: StressTestCase {
         // Was called one more time on receiving token expired error
         XCTAssertEqual(testEnv.clientUpdater!.reloadUserIfNeeded_callsCount, 2)
 
-        // Assert the provided connection id is not `nil`, because we don't reset connectionId
-        // on receiving token expiration error
-        XCTAssertNotNil(providedConnectionId)
+        // We set connectionId to nil after token expiration disconnect
+        XCTAssertNil(client.connectionId)
     }
 
     // MARK: - APIClient tests