chore(auth): Move MFA Preferences to use cases

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/AWSCognitoAuthPlugin.kt

@@ -526,23 +526,35 @@ class AWSCognitoAuthPlugin : AuthPlugin<AWSCognitoAuthService>() {
         enqueue(onSuccess, onError) { queueFacade.clearFederationToIdentityPool() }
 
     fun fetchMFAPreference(onSuccess: Consumer<UserMFAPreference>, onError: Consumer<AuthException>) =
-        enqueue(onSuccess, onError) { queueFacade.fetchMFAPreference() }
+        enqueue(onSuccess, onError) { useCaseFactory.fetchMfaPreference().execute() }
 
     @Deprecated("Use updateMFAPreference(sms, totp, email, onSuccess, onError) instead")
     fun updateMFAPreference(
         sms: MFAPreference?,
         totp: MFAPreference?,
         onSuccess: Action,
         onError: Consumer<AuthException>
-    ) = enqueue(onSuccess, onError) { queueFacade.updateMFAPreference(sms, totp, null) }
+    ) = enqueue(onSuccess, onError) {
+        useCaseFactory.updateMfaPreference().execute(
+            sms = sms,
+            totp = totp,
+            email = null
+        )
+    }
 
     fun updateMFAPreference(
         sms: MFAPreference? = null,
         totp: MFAPreference? = null,
         email: MFAPreference? = null,
         onSuccess: Action,
         onError: Consumer<AuthException>
-    ) = enqueue(onSuccess, onError) { queueFacade.updateMFAPreference(sms, totp, email) }
+    ) = enqueue(onSuccess, onError) {
+        useCaseFactory.updateMfaPreference().execute(
+            sms = sms,
+            totp = totp,
+            email = email
+        )
+    }
 
     private fun enqueue(onSuccess: Action, onError: Consumer<AuthException>, block: suspend () -> Unit) =
         enqueue({ onSuccess.call() }, onError::accept, block)

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/KotlinAuthFacadeInternal.kt

@@ -169,24 +169,6 @@ internal class KotlinAuthFacadeInternal(private val delegate: RealAWSCognitoAuth
         )
     }
 
-    suspend fun fetchMFAPreference(): UserMFAPreference = suspendCoroutine { continuation ->
-        delegate.fetchMFAPreference(
-            { continuation.resume(it) },
-            { continuation.resumeWithException(it) }
-        )
-    }
-
-    suspend fun updateMFAPreference(sms: MFAPreference?, totp: MFAPreference?, email: MFAPreference?) =
-        suspendCoroutine { continuation ->
-            delegate.updateMFAPreference(
-                sms,
-                totp,
-                email,
-                { continuation.resume(Unit) },
-                { continuation.resumeWithException(it) }
-            )
-        }
-
     suspend fun autoSignIn(): AuthSignInResult = suspendCoroutine { continuation ->
         delegate.autoSignIn(
             { continuation.resume(it) },

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/RealAWSCognitoAuthPlugin.kt

@@ -18,12 +18,7 @@ package com.amplifyframework.auth.cognito
 import android.app.Activity
 import android.content.Intent
 import androidx.annotation.WorkerThread
-import aws.sdk.kotlin.services.cognitoidentityprovider.getUser
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.ChallengeNameType
-import aws.sdk.kotlin.services.cognitoidentityprovider.model.EmailMfaSettingsType
-import aws.sdk.kotlin.services.cognitoidentityprovider.model.SmsMfaSettingsType
-import aws.sdk.kotlin.services.cognitoidentityprovider.model.SoftwareTokenMfaSettingsType
-import aws.sdk.kotlin.services.cognitoidentityprovider.setUserMfaPreference
 import com.amplifyframework.AmplifyException
 import com.amplifyframework.annotations.InternalAmplifyApi
 import com.amplifyframework.auth.AWSCognitoAuthMetadataType
@@ -47,7 +42,6 @@ import com.amplifyframework.auth.cognito.helpers.HostedUIHelper
 import com.amplifyframework.auth.cognito.helpers.SignInChallengeHelper
 import com.amplifyframework.auth.cognito.helpers.getAllowedMFATypesFromChallengeParameters
 import com.amplifyframework.auth.cognito.helpers.getMFASetupTypeOrNull
-import com.amplifyframework.auth.cognito.helpers.getMFAType
 import com.amplifyframework.auth.cognito.helpers.getMFATypeOrNull
 import com.amplifyframework.auth.cognito.helpers.identityProviderName
 import com.amplifyframework.auth.cognito.helpers.isMfaSetupSelectionChallenge
@@ -1566,139 +1560,6 @@ internal class RealAWSCognitoAuthPlugin(
         }
     }
 
-    fun fetchMFAPreference(onSuccess: Consumer<UserMFAPreference>, onError: Consumer<AuthException>) {
-        authStateMachine.getCurrentState { authState ->
-            when (authState.authNState) {
-                is AuthenticationState.SignedIn -> {
-                    GlobalScope.launch {
-                        try {
-                            val accessToken = getSession().userPoolTokensResult.value?.accessToken
-                            accessToken?.let { token ->
-                                authEnvironment.cognitoAuthService
-                                    .cognitoIdentityProviderClient?.getUser {
-                                        this.accessToken = token
-                                    }?.also { response ->
-                                        var enabledSet: MutableSet<MFAType>? = null
-                                        var preferred: MFAType? = null
-                                        if (!response.userMfaSettingList.isNullOrEmpty()) {
-                                            enabledSet = mutableSetOf()
-                                            response.userMfaSettingList?.forEach { mfaType ->
-                                                enabledSet.add(getMFAType(mfaType))
-                                            }
-                                        }
-                                        response.preferredMfaSetting?.let { preferredMFA ->
-                                            preferred = getMFAType(preferredMFA)
-                                        }
-                                        onSuccess.accept(UserMFAPreference(enabledSet, preferred))
-                                    }
-                            } ?: onError.accept(SignedOutException())
-                        } catch (error: Exception) {
-                            onError.accept(
-                                CognitoAuthExceptionConverter.lookup(
-                                    error,
-                                    "Cannot update the MFA preferences"
-                                )
-                            )
-                        }
-                    }
-                }
-
-                else -> onError.accept(InvalidStateException())
-            }
-        }
-    }
-
-    fun updateMFAPreference(
-        sms: MFAPreference?,
-        totp: MFAPreference?,
-        email: MFAPreference?,
-        onSuccess: Action,
-        onError: Consumer<AuthException>
-    ) {
-        if (sms == null && totp == null && email == null) {
-            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) {
-                    is AuthenticationState.SignedIn -> {
-                        GlobalScope.launch {
-                            try {
-                                val accessToken = getSession().userPoolTokensResult.value?.accessToken
-                                accessToken?.let { token ->
-                                    authEnvironment
-                                        .cognitoAuthService
-                                        .cognitoIdentityProviderClient
-                                        ?.setUserMfaPreference {
-                                            this.accessToken = token
-                                            this.smsMfaSettings = sms?.let {
-                                                val preferredMFASetting = it.mfaPreferred
-                                                    ?: (
-                                                        overridePreferredSetting &&
-                                                            userPreference.preferred == MFAType.SMS &&
-                                                            it.mfaEnabled
-                                                        )
-                                                SmsMfaSettingsType.invoke {
-                                                    enabled = it.mfaEnabled
-                                                    preferredMfa = preferredMFASetting
-                                                }
-                                            }
-                                            this.softwareTokenMfaSettings = totp?.let {
-                                                val preferredMFASetting = it.mfaPreferred
-                                                    ?: (
-                                                        overridePreferredSetting &&
-                                                            userPreference.preferred == MFAType.TOTP &&
-                                                            it.mfaEnabled
-                                                        )
-                                                SoftwareTokenMfaSettingsType.invoke {
-                                                    enabled = it.mfaEnabled
-                                                    preferredMfa = preferredMFASetting
-                                                }
-                                            }
-                                            this.emailMfaSettings = email?.let {
-                                                val preferredMFASetting = it.mfaPreferred
-                                                    ?: (
-                                                        overridePreferredSetting &&
-                                                            userPreference.preferred == MFAType.EMAIL &&
-                                                            it.mfaEnabled
-                                                        )
-                                                EmailMfaSettingsType.invoke {
-                                                    enabled = it.mfaEnabled
-                                                    preferredMfa = preferredMFASetting
-                                                }
-                                            }
-                                        }?.also {
-                                            onSuccess.call()
-                                        }
-                                } ?: onError.accept(SignedOutException())
-                            } catch (error: Exception) {
-                                onError.accept(
-                                    CognitoAuthExceptionConverter.lookup(
-                                        error,
-                                        "Amazon Cognito cannot update the MFA preferences"
-                                    )
-                                )
-                            }
-                        }
-                    }
-                    else -> onError.accept(InvalidStateException())
-                }
-            }
-        }, {
-            onError.accept(
-                AuthException(
-                    message = "Failed to fetch current MFA preferences " +
-                        "which is a pre-requisite to update MFA preferences",
-                    recoverySuggestion = AmplifyException.TODO_RECOVERY_SUGGESTION,
-                    cause = it
-                )
-            )
-        })
-    }
-
     private fun _clearFederationToIdentityPool(onSuccess: Action, onError: Consumer<AuthException>) {
         _signOut(sendHubEvent = false) {
             when (it) {

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/UserMFAPreference.kt

@@ -22,18 +22,12 @@ import com.amplifyframework.auth.MFAType
  * @param enabled MFA types
  * @param preferred MFA type. null if not set
  */
-data class UserMFAPreference(
-    val enabled: Set<MFAType>?,
-    val preferred: MFAType?
-)
+data class UserMFAPreference(val enabled: Set<MFAType>?, val preferred: MFAType?)
 
 /**
  * Input for updating the MFA preference for a MFA Type
  */
-enum class MFAPreference(
-    internal val mfaEnabled: Boolean,
-    internal val mfaPreferred: Boolean? = null
-) {
+enum class MFAPreference(internal val mfaEnabled: Boolean, internal val mfaPreferred: Boolean? = null) {
     /**
      * MFA not enabled
      */

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/usecases/AuthUseCaseFactory.kt

@@ -139,4 +139,17 @@ internal class AuthUseCaseFactory(
         environment = authEnvironment,
         stateMachine = stateMachine
     )
+
+    fun fetchMfaPreference() = FetchMfaPreferenceUseCase(
+        client = authEnvironment.requireIdentityProviderClient(),
+        fetchAuthSession = fetchAuthSession(),
+        stateMachine = stateMachine
+    )
+
+    fun updateMfaPreference() = UpdateMfaPreferenceUseCase(
+        client = authEnvironment.requireIdentityProviderClient(),
+        fetchAuthSession = fetchAuthSession(),
+        fetchMfaPreference = fetchMfaPreference(),
+        stateMachine = stateMachine
+    )
 }

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/usecases/FetchMfaPreferenceUseCase.kt

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.amplifyframework.auth.cognito.usecases
+
+import aws.sdk.kotlin.services.cognitoidentityprovider.CognitoIdentityProviderClient
+import aws.sdk.kotlin.services.cognitoidentityprovider.getUser
+import com.amplifyframework.auth.cognito.AuthStateMachine
+import com.amplifyframework.auth.cognito.UserMFAPreference
+import com.amplifyframework.auth.cognito.helpers.getMFAType
+import com.amplifyframework.auth.cognito.requireAccessToken
+import com.amplifyframework.auth.cognito.requireSignedInState
+
+internal class FetchMfaPreferenceUseCase(
+    private val client: CognitoIdentityProviderClient,
+    private val fetchAuthSession: FetchAuthSessionUseCase,
+    private val stateMachine: AuthStateMachine
+) {
+    suspend fun execute(): UserMFAPreference {
+        stateMachine.requireSignedInState()
+        val token = fetchAuthSession.execute().requireAccessToken()
+
+        val response = client.getUser { accessToken = token }
+
+        val enabled = response.userMfaSettingList?.map { getMFAType(it) }?.toSet()
+        val preferred = response.preferredMfaSetting?.let { getMFAType(it) }
+
+        return UserMFAPreference(enabled = enabled, preferred = preferred)
+    }
+}

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/usecases/UpdateMfaPreferenceUseCase.kt

@@ -0,0 +1,91 @@
+/*
+ * Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.amplifyframework.auth.cognito.usecases
+
+import aws.sdk.kotlin.services.cognitoidentityprovider.CognitoIdentityProviderClient
+import aws.sdk.kotlin.services.cognitoidentityprovider.model.EmailMfaSettingsType
+import aws.sdk.kotlin.services.cognitoidentityprovider.model.SmsMfaSettingsType
+import aws.sdk.kotlin.services.cognitoidentityprovider.model.SoftwareTokenMfaSettingsType
+import aws.sdk.kotlin.services.cognitoidentityprovider.setUserMfaPreference
+import com.amplifyframework.auth.MFAType
+import com.amplifyframework.auth.cognito.AuthStateMachine
+import com.amplifyframework.auth.cognito.MFAPreference
+import com.amplifyframework.auth.cognito.UserMFAPreference
+import com.amplifyframework.auth.cognito.exceptions.service.InvalidParameterException
+import com.amplifyframework.auth.cognito.requireAccessToken
+import com.amplifyframework.auth.cognito.requireSignedInState
+
+internal class UpdateMfaPreferenceUseCase(
+    private val client: CognitoIdentityProviderClient,
+    private val fetchAuthSession: FetchAuthSessionUseCase,
+    private val fetchMfaPreference: FetchMfaPreferenceUseCase,
+    private val stateMachine: AuthStateMachine
+) {
+    suspend fun execute(sms: MFAPreference?, totp: MFAPreference?, email: MFAPreference?) {
+        if (sms == null && totp == null && email == null) {
+            throw InvalidParameterException("No mfa settings given")
+        }
+
+        stateMachine.requireSignedInState()
+        val token = fetchAuthSession.execute().requireAccessToken()
+
+        // If none of the params are marked as preferred then keep the existing preferred property
+        val keepExistingPreference = !(
+            sms?.mfaPreferred == true ||
+                totp?.mfaPreferred == true ||
+                email?.mfaPreferred == true
+            )
+        val existingPreference = fetchMfaPreference.execute()
+
+        client.setUserMfaPreference {
+            accessToken = token
+            smsMfaSettings = sms?.let {
+                val preferred = isPreferred(MFAType.SMS, sms, keepExistingPreference, existingPreference)
+                SmsMfaSettingsType {
+                    enabled = it.mfaEnabled
+                    preferredMfa = preferred
+                }
+            }
+            softwareTokenMfaSettings = totp?.let {
+                val preferred = isPreferred(MFAType.TOTP, totp, keepExistingPreference, existingPreference)
+                SoftwareTokenMfaSettingsType {
+                    enabled = it.mfaEnabled
+                    preferredMfa = preferred
+                }
+            }
+            emailMfaSettings = email?.let {
+                val preferred = isPreferred(MFAType.EMAIL, email, keepExistingPreference, existingPreference)
+                EmailMfaSettingsType {
+                    enabled = it.mfaEnabled
+                    preferredMfa = preferred
+                }
+            }
+        }
+    }
+
+    // If preference.mfaPreferred is set then we use that. Otherwise, we will keep the existing preference if
+    // keepExistingPreference is true and the preference is enabled
+    private fun isPreferred(
+        type: MFAType,
+        preference: MFAPreference,
+        keepExistingPreference: Boolean,
+        existingPreference: UserMFAPreference
+    ): Boolean = preference.mfaPreferred ?: (
+        keepExistingPreference &&
+            existingPreference.preferred == type &&
+            preference.mfaEnabled
+        )
+}