test(auth): Fixes HostedUI signIn flow and flaky test (#2456)

Author: royjit

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Operations/Helpers/HostedUISignInHelper.swift

@@ -37,6 +37,8 @@ struct HostedUISignInHelper {
                 continue
             }
             switch authenticationState {
+            case .signingIn:
+                continue
             case .signedIn:
                 throw AuthError.invalidState(
                     "There is already a user in signedIn state. SignOut the user first before calling signIn",

AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/Support/DefaultConfig.swift

@@ -256,14 +256,15 @@ struct MockCredentialStoreOperationClient: CredentialStoreStateBehavior {
 
 class MockAmplifyStore: AmplifyAuthCredentialStoreBehavior {
     let credentialsKey = "amplifyCredentials"
-    static var dict: [String: Data] = [:]
+    static var dict = AtomicDictionary<String, Data>()
 
     func saveCredential(_ credential: AmplifyCredentials) throws {
-        Self.dict[credentialsKey] = try? JSONEncoder().encode(credential)
+        let value = (try? JSONEncoder().encode(credential)) ?? Data()
+        Self.dict.set(value: value, forKey: credentialsKey)
     }
 
     func retrieveCredential() throws -> AmplifyCredentials {
-        guard let data = Self.dict[credentialsKey],
+        guard let data = Self.dict.getValue(forKey: credentialsKey),
               let cred = (try? JSONDecoder().decode(AmplifyCredentials.self, from: data)) else {
             return AmplifyCredentials.noCredentials
         }
@@ -279,11 +280,12 @@ class MockAmplifyStore: AmplifyAuthCredentialStoreBehavior {
     }
 
     func saveDevice(_ deviceMetadata: DeviceMetadata, for username: String) throws {
-        Self.dict[username] = try? JSONEncoder().encode(deviceMetadata)
+        let value = (try? JSONEncoder().encode(deviceMetadata)) ?? Data()
+        Self.dict.set(value: value, forKey: username)
     }
 
     func retrieveDevice(for username: String) throws -> DeviceMetadata {
-        guard let data = Self.dict[username],
+        guard let data = Self.dict.getValue(forKey: username),
               let device = (try? JSONDecoder().decode(DeviceMetadata.self, from: data)) else {
             return .noData
         }
@@ -295,12 +297,13 @@ class MockAmplifyStore: AmplifyAuthCredentialStoreBehavior {
     }
 
     func saveASFDevice(_ deviceId: String, for username: String) throws {
-        Self.dict[username] = try? JSONEncoder().encode(deviceId)
+        let value = (try? JSONEncoder().encode(deviceId)) ?? Data()
+        Self.dict.set(value: value, forKey: username)
 
     }
 
     func retrieveASFDevice(for username: String) throws -> String {
-        guard let data = Self.dict[username],
+        guard let data = Self.dict.getValue(forKey: username),
               let device = (try? JSONDecoder().decode(String.self, from: data)) else {
             throw KeychainStoreError.itemNotFound
         }

AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/TaskTests/HostedUITests/AWSAuthHostedUISignInTests.swift

@@ -99,7 +99,7 @@ class AWSAuthHostedUISignInTests: XCTestCase {
     }
 
     @MainActor
-    func testRestartAfterError() async {
+    func testRestartAfterError() async throws {
         mockHostedUIResult = .failure(.cancelled)
         let errorExpectation  = expectation(description: "SignIn operation should complete")
         do {
@@ -114,7 +114,7 @@ class AWSAuthHostedUISignInTests: XCTestCase {
             errorExpectation.fulfill()
         }
 
-        wait(for: [errorExpectation], timeout: networkTimeout)
+        waitForExpectations(timeout: networkTimeout)
         mockHostedUIResult = .success([
             .init(name: "state", value: mockState),
             .init(name: "code", value: mockProof)
@@ -127,7 +127,7 @@ class AWSAuthHostedUISignInTests: XCTestCase {
         } catch {
             XCTFail("Should not fail with error = \(error)")
         }
-        wait(for: [signInExpectation], timeout: networkTimeout)
+        waitForExpectations(timeout: networkTimeout)
     }
 
     @MainActor

AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/hierarchical-state-machine-swiftTests/StateMachineListenerTests.swift

@@ -123,22 +123,28 @@ class StateMachineListenerTests: XCTestCase {
         let task = Task {
             for try await _ in seq {
                 if Task.isCancelled {
-                    notified.fulfill()
-                    break
+                    print("Cancelled")
                 }
             }
-
+            notified.fulfill()
         }
 
-        Task {
-            for _ in 1...100 {
-                let event = Counter.Event(id: "test", eventType: .adjustBy(0))
+        let task2 = Task {
+
+            for index in 1...100 {
+
+                let event = Counter.Event(id: "test", eventType: .adjustBy(index))
                 await stateMachine.send(event)
+
+                if (index == 30) {
+                    task.cancel()
+                    await Task.yield()
+                }
             }
         }
 
-        task.cancel()
-        await waitForExpectations(timeout: 0.1)
+        await waitForExpectations(timeout: 1)
+        task2.cancel()
     }
 
 }

AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/hierarchical-state-machine-swiftTests/StateMachineTests.swift

@@ -37,15 +37,19 @@ class StateMachineTests: XCTestCase {
         let increment = Counter.Event(id: "increment", eventType: .increment)
         let decrement = Counter.Event(id: "decrement", eventType: .decrement)
 
+        let taskCompletion = expectation(description: "Completed sending event")
+        taskCompletion.expectedFulfillmentCount = 1_000
         Task {
             for i in 1...1_000 {
                 if i.isMultiple(of: 2) {
                     await testMachine.send(increment)
                 } else {
                     await testMachine.send(decrement)
                 }
+                taskCompletion.fulfill()
             }
         }
+        await waitForExpectations(timeout: 1)
         let state = await testMachine.currentState
         XCTAssertEqual(state.value, 0)
     }