chore: fix OTP integration tests

Author: harsh62

AmplifyPlugins/Auth/Tests/AuthHostApp/AuthIntegrationTests/AWSAuthBaseTest.swift

@@ -135,13 +135,31 @@ class AWSAuthBaseTest: XCTestCase {
     """
 
     /// Function to create a subscription and store OTP codes in a dictionary
-    func subscribeToOTPCreation() {
+    func subscribeToOTPCreation() async {
         subscription = Amplify.API.subscribe(request: .init(document: document, responseType: [String: JSONValue].self))
 
+        func waitForSubscriptionConnection(
+            subscription: AmplifyAsyncThrowingSequence<GraphQLSubscriptionEvent<[String: JSONValue]>>
+        ) async throws {
+            for try await subscriptionEvent in subscription {
+                if case .connection(let subscriptionConnectionState) = subscriptionEvent {
+                    print("Subscription connect state is \(subscriptionConnectionState)")
+                    if subscriptionConnectionState == .connected {
+                        return
+                    }
+                }
+            }
+        }
+
+        guard let subscription = subscription else { return }
+
+        await wait(name: "Subscription Connection Waiter", timeout: 5.0) {
+            try await waitForSubscriptionConnection(subscription: subscription)
+        }
+
         // Create the subscription and listen for OTP code events
         Task {
             do {
-                guard let subscription = subscription else { return }
                 for try await subscriptionEvent in subscription {
                     switch subscriptionEvent {
                     case .connection(let subscriptionConnectionState):
@@ -153,7 +171,7 @@ class AWSAuthBaseTest: XCTestCase {
                             if let eventUsername = otpResult["onCreateMfaInfo"]?.asObject?["username"]?.stringValue,
                                let code = otpResult["onCreateMfaInfo"]?.asObject?["code"]?.stringValue {
                                 // Store the code in the dictionary for the given username
-                                usernameOTPDictionary[eventUsername] = code
+                                usernameOTPDictionary[eventUsername.lowercased()] = code
                             }
                         case .failure(let error):
                             print("Got failed result with \(error.errorDescription)")
@@ -168,13 +186,14 @@ class AWSAuthBaseTest: XCTestCase {
 
     /// Test that waits for the OTP code using XCTestExpectation
     func otp(for username: String) async throws -> String? {
+        let lowerCasedUsername = username.lowercased()
         let expectation = XCTestExpectation(description: "Wait for OTP")
         expectation.expectedFulfillmentCount = 1
 
         let task = Task { () -> String? in
             var code: String?
             for _ in 0..<30 { // Poll for the code, max 30 times (once per second)
-                if let otp = usernameOTPDictionary[username] {
+                if let otp = usernameOTPDictionary[lowerCasedUsername] {
                     code = otp
                     expectation.fulfill() // Fulfill the expectation when the value is found
                     break
@@ -190,10 +209,9 @@ class AWSAuthBaseTest: XCTestCase {
         if result == .timedOut {
             // Task cancels if timed out
             task.cancel()
-            subscription?.cancel()
             return nil
         }
-        usernameOTPDictionary.removeValue(forKey: username)
+        usernameOTPDictionary.removeValue(forKey: lowerCasedUsername)
         return try await task.value
     }
 }

AmplifyPlugins/Auth/Tests/AuthHostApp/AuthIntegrationTests/MFATests/EmailMFATests/EmailMFAOnlyTests.swift

@@ -60,7 +60,7 @@ class EmailMFARequiredTests: AWSAuthBaseTest {
     func testSuccessfulEmailMFASetupStep() async {
         do {
             // Step 1: Set up a subscription to receive MFA codes
-            subscribeToOTPCreation()
+            await subscribeToOTPCreation()
 
             // Step 2: Sign up a new user
             let uniqueId = UUID().uuidString
@@ -144,7 +144,7 @@ class EmailMFARequiredTests: AWSAuthBaseTest {
     func testSuccessfulEmailMFAWithIncorrectCodeFirstAndThenValidOne() async {
         do {
             // Step 1: Set up a subscription to receive MFA codes
-            subscribeToOTPCreation()
+            await subscribeToOTPCreation()
 
             // Step 2: Sign up a new user
             let uniqueId = UUID().uuidString

AmplifyPlugins/Auth/Tests/AuthHostApp/AuthIntegrationTests/MFATests/EmailMFATests/EmailMFAWithAllMFATypesRequiredTests.swift

@@ -94,7 +94,7 @@ class EmailMFAWithAllMFATypesRequiredTests: AWSAuthBaseTest {
     func testSuccessfulEmailMFACodeStep() async {
         do {
             // Step 1: Set up a subscription to receive MFA codes
-            subscribeToOTPCreation()
+            await subscribeToOTPCreation()
             let uniqueId = UUID().uuidString
             let username = randomEmail
             let password = "Pp123@\(uniqueId)"
@@ -152,7 +152,7 @@ class EmailMFAWithAllMFATypesRequiredTests: AWSAuthBaseTest {
     func testConfirmSignInForEmailMFASetupSelectionStep() async {
         do {
             // Step 1: Set up a subscription to receive MFA codes
-            subscribeToOTPCreation()
+            await subscribeToOTPCreation()
             let uniqueId = UUID().uuidString
             let username = "\(uniqueId)"
             let password = "Pp123@\(uniqueId)"

AmplifyPlugins/Auth/Tests/AuthHostApp/AuthIntegrationTests/PasswordlessTests/PasswordlessAutoSignInTests.swift

@@ -65,7 +65,7 @@ class PasswordlessAutoSignInTests: AWSAuthBaseTest {
     ///
     func testSuccessfulPasswordlessSignUpAndAutoSignInEndtoEnd() async throws {
 
-        subscribeToOTPCreation()
+        await subscribeToOTPCreation()
 
         let username = "integTest\(UUID().uuidString)"
         let options = AuthSignUpRequest.Options(
@@ -114,7 +114,7 @@ class PasswordlessAutoSignInTests: AWSAuthBaseTest {
     ///
     func testFailureMultipleAutoSignInWithSameSession() async throws {
 
-        subscribeToOTPCreation()
+        await subscribeToOTPCreation()
 
         let username = "integTest\(UUID().uuidString)"
         let options = AuthSignUpRequest.Options(

AmplifyPlugins/Auth/Tests/AuthHostApp/AuthIntegrationTests/PasswordlessTests/PasswordlessConfirmSignUpTests.swift

@@ -127,7 +127,7 @@ class PasswordlessConfirmSignUpTests: AWSAuthBaseTest {
     ///
     func testSuccessfulPasswordlessSignUpAndConfirmSignUpEndtoEnd() async throws {
 
-        subscribeToOTPCreation()
+        await subscribeToOTPCreation()
 
         let username = "integTest\(UUID().uuidString)"
         let options = AuthSignUpRequest.Options(

AmplifyPlugins/Auth/Tests/AuthHostApp/AuthIntegrationTests/PasswordlessTests/PasswordlessSignInTests.swift

@@ -30,7 +30,7 @@ class PasswordlessSignInTests: AWSAuthBaseTest {
         try await super.setUp()
         AuthSessionHelper.clearSession()
 
-        subscribeToOTPCreation()
+        await subscribeToOTPCreation()
     }
 
     override func tearDown() async throws {