Fix download error if token needs refresh

Previously the token in the WealsimpleDownloader class was not updated when the token was refreshed.
Nef10 · Jun 4, 2021 · 9ed7e8f · 9ed7e8f
1 parent 48c7ef4
commit 9ed7e8f
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 24 deletions.
diff --git a/Sources/Wealthsimple/Token.swift b/Sources/Wealthsimple/Token.swift
@@ -78,14 +78,29 @@ struct Token {
         sendTokenRequest(parameters: json, request: request, credentialStorage: credentialStorage, completion: completion)
     }
 
-    static func getToken(from credentialStorage: CredentialStorage) -> Self? {
+    static func getToken(from credentialStorage: CredentialStorage, completion: @escaping (Self?) -> Void) {
         guard let accessToken = credentialStorage.read(credentialStorageKeyAccessToken),
               let refreshToken = credentialStorage.read(credentialStorageKeyRefreshToken),
               let expiryString = credentialStorage.read(credentialStorageKeyExpiry),
               let expiryDouble = Double(expiryString) else {
-            return nil
+            completion(nil)
+            return
+        }
+        let token = Self(accessToken: accessToken, refreshToken: refreshToken, expiry: Date(timeIntervalSince1970: expiryDouble), credentialStorage: credentialStorage)
+        token.refreshIfNeeded {
+            switch $0 {
+            case .failure:
+                completion(nil)
+            case let .success(newToken):
+                token.testIfValid {
+                    if $0 {
+                        completion(newToken)
+                    } else {
+                        completion(nil)
+                    }
+                }
+            }
         }
-        return Self(accessToken: accessToken, refreshToken: refreshToken, expiry: Date(timeIntervalSince1970: expiryDouble), credentialStorage: credentialStorage)
     }
 
     private static func sendTokenRequest(
@@ -150,7 +165,7 @@ struct Token {
         }
     }
 
-    func testIfValid(completion: @escaping (Bool) -> Void) {
+    private func testIfValid(completion: @escaping (Bool) -> Void) {
         var request = URLRequest(url: Self.testUrl)
         let session = URLSession.shared
         request.setValue("application/json", forHTTPHeaderField: "Content-Type")
@@ -177,18 +192,11 @@ struct Token {
 
     func authenticateRequest(_ request: URLRequest, completion: @escaping (Result<URLRequest, TokenError>) -> Void) {
         var requestCopy = request
-        refreshIfNeeded {
-            switch $0 {
-            case let .failure(error):
-                completion(.failure(error))
-            case let .success(token):
-                requestCopy.setValue("Bearer \(token.accessToken)", forHTTPHeaderField: "Authorization")
-                completion(.success(requestCopy))
-            }
-        }
+        requestCopy.setValue("Bearer \(self.accessToken)", forHTTPHeaderField: "Authorization")
+        completion(.success(requestCopy))
     }
 
-    private func refreshIfNeeded(completion: @escaping (Result<Self, TokenError>) -> Void) {
+    func refreshIfNeeded(completion: @escaping (Result<Self, TokenError>) -> Void) {
         if needsRefresh() {
             refresh(completion: completion)
         } else {

diff --git a/Sources/Wealthsimple/WealthsimpleDownloader.swift b/Sources/Wealthsimple/WealthsimpleDownloader.swift
@@ -51,16 +51,23 @@ public final class WealthsimpleDownloader {
     }
 
     /// Authneticates against the API. Call before calling any other method.
-    /// - Parameter completion: Get an error in case soemthing went wrong, otherwise nil
+    /// - Parameter completion: Gets an error in case something went wrong, otherwise nil
     public func authenticate(completion: @escaping (Error?) -> Void) {
-        guard token == nil else {
-            completion(nil)
-            return
-        }
-        token = Token.getToken(from: credentialStorage)
         if let token = token {
-            token.testIfValid {
-                if $0 {
+            token.refreshIfNeeded {
+                switch $0 {
+                case .failure:
+                    self.getNewToken(completion: completion)
+                case let .success(newToken):
+                    self.token = newToken
+                    completion(nil)
+                }
+            }
+            return
+        } else {
+            Token.getToken(from: credentialStorage) {
+                if let token = $0 {
+                    self.token = token
                     completion(nil)
                     return
                 } else {
@@ -69,8 +76,6 @@ public final class WealthsimpleDownloader {
                     return
                 }
             }
-        } else {
-            getNewToken(completion: completion)
         }
     }