Addressed PR comments

aws-amplify · Oct 4, 2023 · fd6c25e · fd6c25e
1 parent 9e3a5e9
commit fd6c25e
Show file tree

Hide file tree

Showing 2 changed files with 112 additions and 2 deletions.
diff --git a/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/RealAWSCognitoAuthPlugin.kt b/aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/RealAWSCognitoAuthPlugin.kt
@@ -2228,6 +2228,8 @@ internal class RealAWSCognitoAuthPlugin(
             onError.accept(InvalidParameterException("No mfa settings given"))
             return
         }
+        // If either of the params have preferred setting set then ignore fetched preference preferred property
+        val overridePreferredSetting: Boolean = !(sms?.mfaPreferred == true || totp?.mfaPreferred == true)
         fetchMFAPreference({ userPreference ->
             authStateMachine.getCurrentState { authState ->
                 when (authState.authNState) {
@@ -2243,15 +2245,23 @@ internal class RealAWSCognitoAuthPlugin(
                                             this.accessToken = token
                                             this.smsMfaSettings = sms?.let { it ->
                                                 val preferredMFASetting = it.mfaPreferred
-                                                    ?: (userPreference.preferred == MFAType.SMS && it.mfaEnabled)
+                                                    ?: (
+                                                        overridePreferredSetting &&
+                                                            userPreference.preferred == MFAType.SMS &&
+                                                            it.mfaEnabled
+                                                        )
                                                 SmsMfaSettingsType.invoke {
                                                     enabled = it.mfaEnabled
                                                     preferredMfa = preferredMFASetting
                                                 }
                                             }
                                             this.softwareTokenMfaSettings = totp?.let { it ->
                                                 val preferredMFASetting = it.mfaPreferred
-                                                    ?: (userPreference.preferred == MFAType.TOTP && it.mfaEnabled)
+                                                    ?: (
+                                                        overridePreferredSetting &&
+                                                            userPreference.preferred == MFAType.TOTP &&
+                                                            it.mfaEnabled
+                                                        )
                                                 SoftwareTokenMfaSettingsType.invoke {
                                                     enabled = it.mfaEnabled
                                                     preferredMfa = preferredMFASetting

diff --git a/...h-cognito/src/test/java/com/amplifyframework/auth/cognito/RealAWSCognitoAuthPluginTest.kt b/...h-cognito/src/test/java/com/amplifyframework/auth/cognito/RealAWSCognitoAuthPluginTest.kt
@@ -2516,4 +2516,104 @@ class RealAWSCognitoAuthPluginTest {
             setUserMFAPreferenceRequest.captured.softwareTokenMfaSettings
         )
     }
+
+    @Test
+    fun `updatepref  when currentpref is TOTP preferred and params include SMS preferred and TOTP enabled`() {
+        val currentAuthState = mockk<AuthState> {
+            every { authNState } returns AuthenticationState.SignedIn(mockk(), mockk())
+            every { authZState } returns AuthorizationState.SessionEstablished(credentials)
+        }
+        every { authStateMachine.getCurrentState(captureLambda()) } answers {
+            lambda<(AuthState) -> Unit>().invoke(currentAuthState)
+        }
+        val listenLatch = CountDownLatch(1)
+        val onSuccess = mockk<Action>()
+        every { onSuccess.call() }.answers { listenLatch.countDown() }
+        val onError = mockk<Consumer<AuthException>>()
+        val setUserMFAPreferenceRequest = slot<SetUserMfaPreferenceRequest>()
+
+        coEvery {
+            mockCognitoIPClient.getUser {
+                accessToken = credentials.signedInData.cognitoUserPoolTokens.accessToken
+            }
+        }.answers {
+            GetUserResponse.invoke {
+                userMfaSettingList = listOf("SOFTWARE_TOKEN_MFA")
+                preferredMfaSetting = "SOFTWARE_TOKEN_MFA"
+            }
+        }
+
+        coEvery { mockCognitoIPClient.setUserMfaPreference(capture(setUserMFAPreferenceRequest)) }.answers {
+            SetUserMfaPreferenceResponse.invoke {
+            }
+        }
+        plugin.updateMFAPreference(MFAPreference.PREFERRED, MFAPreference.ENABLED, onSuccess, onError)
+
+        assertTrue { listenLatch.await(5, TimeUnit.SECONDS) }
+        assertTrue(setUserMFAPreferenceRequest.isCaptured)
+        assertEquals(
+            SmsMfaSettingsType.invoke {
+                enabled = true
+                preferredMfa = true
+            },
+            setUserMFAPreferenceRequest.captured.smsMfaSettings
+        )
+        assertEquals(
+            SoftwareTokenMfaSettingsType.invoke {
+                enabled = true
+                preferredMfa = false
+            },
+            setUserMFAPreferenceRequest.captured.softwareTokenMfaSettings
+        )
+    }
+
+    @Test
+    fun `updatepref  when currentpref is SMS preferred and params include SMS enabled and TOTP preferred`() {
+        val currentAuthState = mockk<AuthState> {
+            every { authNState } returns AuthenticationState.SignedIn(mockk(), mockk())
+            every { authZState } returns AuthorizationState.SessionEstablished(credentials)
+        }
+        every { authStateMachine.getCurrentState(captureLambda()) } answers {
+            lambda<(AuthState) -> Unit>().invoke(currentAuthState)
+        }
+        val listenLatch = CountDownLatch(1)
+        val onSuccess = mockk<Action>()
+        every { onSuccess.call() }.answers { listenLatch.countDown() }
+        val onError = mockk<Consumer<AuthException>>()
+        val setUserMFAPreferenceRequest = slot<SetUserMfaPreferenceRequest>()
+
+        coEvery {
+            mockCognitoIPClient.getUser {
+                accessToken = credentials.signedInData.cognitoUserPoolTokens.accessToken
+            }
+        }.answers {
+            GetUserResponse.invoke {
+                userMfaSettingList = listOf("SMS_MFA")
+                preferredMfaSetting = "SMS_MFA"
+            }
+        }
+
+        coEvery { mockCognitoIPClient.setUserMfaPreference(capture(setUserMFAPreferenceRequest)) }.answers {
+            SetUserMfaPreferenceResponse.invoke {
+            }
+        }
+        plugin.updateMFAPreference(MFAPreference.ENABLED, MFAPreference.PREFERRED, onSuccess, onError)
+
+        assertTrue { listenLatch.await(5, TimeUnit.SECONDS) }
+        assertTrue(setUserMFAPreferenceRequest.isCaptured)
+        assertEquals(
+            SmsMfaSettingsType.invoke {
+                enabled = true
+                preferredMfa = false
+            },
+            setUserMFAPreferenceRequest.captured.smsMfaSettings
+        )
+        assertEquals(
+            SoftwareTokenMfaSettingsType.invoke {
+                enabled = true
+                preferredMfa = true
+            },
+            setUserMFAPreferenceRequest.captured.softwareTokenMfaSettings
+        )
+    }
 }