chore(auth): Migrate TOTP setup to use case classes

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

@@ -414,7 +414,7 @@ class AWSCognitoAuthPlugin : AuthPlugin<AWSCognitoAuthService>() {
     }
 
     override fun setUpTOTP(onSuccess: Consumer<TOTPSetupDetails>, onError: Consumer<AuthException>) =
-        enqueue(onSuccess, onError) { queueFacade.setUpTOTP() }
+        enqueue(onSuccess, onError) { useCaseFactory.setupTotp().execute() }
 
     override fun verifyTOTPSetup(code: String, onSuccess: Action, onError: Consumer<AuthException>) {
         verifyTOTPSetup(code, AWSCognitoAuthVerifyTOTPSetupOptions.CognitoBuilder().build(), onSuccess, onError)
@@ -425,7 +425,7 @@ class AWSCognitoAuthPlugin : AuthPlugin<AWSCognitoAuthService>() {
         options: AuthVerifyTOTPSetupOptions,
         onSuccess: Action,
         onError: Consumer<AuthException>
-    ) = enqueue(onSuccess, onError) { queueFacade.verifyTOTPSetup(code, options) }
+    ) = enqueue(onSuccess, onError) { useCaseFactory.verifyTotpSetup().execute(code, options) }
 
     override fun associateWebAuthnCredential(
         callingActivity: Activity,

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

@@ -20,7 +20,6 @@ import android.content.Intent
 import com.amplifyframework.auth.AuthCodeDeliveryDetails
 import com.amplifyframework.auth.AuthProvider
 import com.amplifyframework.auth.AuthSession
-import com.amplifyframework.auth.TOTPSetupDetails
 import com.amplifyframework.auth.cognito.options.FederateToIdentityPoolOptions
 import com.amplifyframework.auth.cognito.result.FederateToIdentityPoolResult
 import com.amplifyframework.auth.options.AuthConfirmResetPasswordOptions
@@ -32,7 +31,6 @@ import com.amplifyframework.auth.options.AuthResetPasswordOptions
 import com.amplifyframework.auth.options.AuthSignInOptions
 import com.amplifyframework.auth.options.AuthSignOutOptions
 import com.amplifyframework.auth.options.AuthSignUpOptions
-import com.amplifyframework.auth.options.AuthVerifyTOTPSetupOptions
 import com.amplifyframework.auth.options.AuthWebUISignInOptions
 import com.amplifyframework.auth.result.AuthResetPasswordResult
 import com.amplifyframework.auth.result.AuthSignInResult
@@ -286,21 +284,6 @@ internal class KotlinAuthFacadeInternal(private val delegate: RealAWSCognitoAuth
         )
     }
 
-    suspend fun setUpTOTP(): TOTPSetupDetails = suspendCoroutine { continuation ->
-        delegate.setUpTOTP(
-            { continuation.resume(it) },
-            { continuation.resumeWithException(it) }
-        )
-    }
-    suspend fun verifyTOTPSetup(code: String, options: AuthVerifyTOTPSetupOptions) = suspendCoroutine { continuation ->
-        delegate.verifyTOTPSetup(
-            code,
-            options,
-            { continuation.resume(Unit) },
-            { continuation.resumeWithException(it) }
-        )
-    }
-
     suspend fun fetchMFAPreference(): UserMFAPreference = suspendCoroutine { continuation ->
         delegate.fetchMFAPreference(
             { continuation.resume(it) },

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

@@ -18,7 +18,6 @@ package com.amplifyframework.auth.cognito
 import android.app.Activity
 import android.content.Intent
 import androidx.annotation.WorkerThread
-import aws.sdk.kotlin.services.cognitoidentityprovider.associateSoftwareToken
 import aws.sdk.kotlin.services.cognitoidentityprovider.confirmForgotPassword
 import aws.sdk.kotlin.services.cognitoidentityprovider.getUser
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.AnalyticsMetadataType
@@ -27,10 +26,8 @@ import aws.sdk.kotlin.services.cognitoidentityprovider.model.ChangePasswordReque
 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.model.VerifySoftwareTokenResponseType
 import aws.sdk.kotlin.services.cognitoidentityprovider.resendConfirmationCode
 import aws.sdk.kotlin.services.cognitoidentityprovider.setUserMfaPreference
-import aws.sdk.kotlin.services.cognitoidentityprovider.verifySoftwareToken
 import com.amplifyframework.AmplifyException
 import com.amplifyframework.annotations.InternalAmplifyApi
 import com.amplifyframework.auth.AWSCognitoAuthMetadataType
@@ -43,7 +40,6 @@ import com.amplifyframework.auth.AuthFactorType
 import com.amplifyframework.auth.AuthProvider
 import com.amplifyframework.auth.AuthSession
 import com.amplifyframework.auth.MFAType
-import com.amplifyframework.auth.TOTPSetupDetails
 import com.amplifyframework.auth.cognito.exceptions.configuration.InvalidOauthConfigurationException
 import com.amplifyframework.auth.cognito.exceptions.configuration.InvalidUserPoolConfigurationException
 import com.amplifyframework.auth.cognito.exceptions.invalidstate.SignedInException
@@ -53,7 +49,6 @@ import com.amplifyframework.auth.cognito.exceptions.service.InvalidParameterExce
 import com.amplifyframework.auth.cognito.exceptions.service.UserCancelledException
 import com.amplifyframework.auth.cognito.helpers.AuthHelper
 import com.amplifyframework.auth.cognito.helpers.HostedUIHelper
-import com.amplifyframework.auth.cognito.helpers.SessionHelper
 import com.amplifyframework.auth.cognito.helpers.SignInChallengeHelper
 import com.amplifyframework.auth.cognito.helpers.collectWhile
 import com.amplifyframework.auth.cognito.helpers.getAllowedMFATypesFromChallengeParameters
@@ -70,7 +65,6 @@ import com.amplifyframework.auth.cognito.options.AWSCognitoAuthResendSignUpCodeO
 import com.amplifyframework.auth.cognito.options.AWSCognitoAuthSignInOptions
 import com.amplifyframework.auth.cognito.options.AWSCognitoAuthSignOutOptions
 import com.amplifyframework.auth.cognito.options.AWSCognitoAuthSignUpOptions
-import com.amplifyframework.auth.cognito.options.AWSCognitoAuthVerifyTOTPSetupOptions
 import com.amplifyframework.auth.cognito.options.AWSCognitoAuthWebUISignInOptions
 import com.amplifyframework.auth.cognito.options.AuthFlowType
 import com.amplifyframework.auth.cognito.options.FederateToIdentityPoolOptions
@@ -96,7 +90,6 @@ import com.amplifyframework.auth.options.AuthResetPasswordOptions
 import com.amplifyframework.auth.options.AuthSignInOptions
 import com.amplifyframework.auth.options.AuthSignOutOptions
 import com.amplifyframework.auth.options.AuthSignUpOptions
-import com.amplifyframework.auth.options.AuthVerifyTOTPSetupOptions
 import com.amplifyframework.auth.options.AuthWebUISignInOptions
 import com.amplifyframework.auth.result.AuthResetPasswordResult
 import com.amplifyframework.auth.result.AuthSignInResult
@@ -1948,56 +1941,6 @@ internal class RealAWSCognitoAuthPlugin(
         }
     }
 
-    fun setUpTOTP(onSuccess: Consumer<TOTPSetupDetails>, onError: Consumer<AuthException>) {
-        authStateMachine.getCurrentState { authState ->
-            when (authState.authNState) {
-                is AuthenticationState.SignedIn -> {
-                    GlobalScope.launch {
-                        try {
-                            val accessToken = getSession().userPoolTokensResult.value?.accessToken
-                            accessToken?.let { token ->
-                                SessionHelper.getUsername(token)?.let { username ->
-                                    authEnvironment.cognitoAuthService
-                                        .cognitoIdentityProviderClient?.associateSoftwareToken {
-                                            this.accessToken = token
-                                        }?.also { response ->
-                                            response.secretCode?.let { secret ->
-                                                onSuccess.accept(
-                                                    TOTPSetupDetails(
-                                                        secret,
-                                                        username
-                                                    )
-                                                )
-                                            }
-                                        }
-                                }
-                            } ?: onError.accept(SignedOutException())
-                        } catch (error: Exception) {
-                            onError.accept(
-                                CognitoAuthExceptionConverter.lookup(
-                                    error,
-                                    "Cannot find a multi-factor authentication (MFA) method."
-                                )
-                            )
-                        }
-                    }
-                }
-
-                else -> onError.accept(InvalidStateException())
-            }
-        }
-    }
-
-    fun verifyTOTPSetup(
-        code: String,
-        options: AuthVerifyTOTPSetupOptions,
-        onSuccess: Action,
-        onError: Consumer<AuthException>
-    ) {
-        val cognitoOptions = options as? AWSCognitoAuthVerifyTOTPSetupOptions
-        verifyTotp(code, cognitoOptions?.friendlyDeviceName, onSuccess, onError)
-    }
-
     fun fetchMFAPreference(onSuccess: Consumer<UserMFAPreference>, onError: Consumer<AuthException>) {
         authStateMachine.getCurrentState { authState ->
             when (authState.authNState) {
@@ -2131,50 +2074,6 @@ internal class RealAWSCognitoAuthPlugin(
         })
     }
 
-    private fun verifyTotp(
-        code: String,
-        friendlyDeviceName: String?,
-        onSuccess: Action,
-        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?.verifySoftwareToken {
-                                        this.userCode = code
-                                        this.friendlyDeviceName = friendlyDeviceName
-                                        this.accessToken = token
-                                    }?.also {
-                                        when (it.status) {
-                                            is VerifySoftwareTokenResponseType.Success -> onSuccess.call()
-                                            else -> throw ServiceException(
-                                                message = "An unknown service error has occurred",
-                                                recoverySuggestion = AmplifyException.TODO_RECOVERY_SUGGESTION
-                                            )
-                                        }
-                                    }
-                            } ?: onError.accept(SignedOutException())
-                        } catch (error: Exception) {
-                            onError.accept(
-                                CognitoAuthExceptionConverter.lookup(
-                                    error,
-                                    "Amazon Cognito cannot find a multi-factor authentication (MFA) method."
-                                )
-                            )
-                        }
-                    }
-                }
-
-                else -> onError.accept(InvalidStateException())
-            }
-        }
-    }
-
     private fun _clearFederationToIdentityPool(onSuccess: Action, onError: Consumer<AuthException>) {
         _signOut(sendHubEvent = false) {
             when (it) {

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

@@ -96,4 +96,16 @@ internal class AuthUseCaseFactory(
         fetchAuthSession = fetchAuthSession(),
         stateMachine = stateMachine
     )
+
+    fun setupTotp() = SetupTotpUseCase(
+        client = authEnvironment.requireIdentityProviderClient(),
+        fetchAuthSession = fetchAuthSession(),
+        stateMachine = stateMachine
+    )
+
+    fun verifyTotpSetup() = VerifyTotpSetupUseCase(
+        client = authEnvironment.requireIdentityProviderClient(),
+        fetchAuthSession = fetchAuthSession(),
+        stateMachine = stateMachine
+    )
 }

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

@@ -0,0 +1,52 @@
+/*
+ * 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.associateSoftwareToken
+import com.amplifyframework.AmplifyException
+import com.amplifyframework.auth.AuthException
+import com.amplifyframework.auth.TOTPSetupDetails
+import com.amplifyframework.auth.cognito.AuthStateMachine
+import com.amplifyframework.auth.cognito.requireAccessToken
+import com.amplifyframework.auth.cognito.requireSignedInState
+
+internal class SetupTotpUseCase(
+    private val fetchAuthSession: FetchAuthSessionUseCase,
+    private val client: CognitoIdentityProviderClient,
+    private val stateMachine: AuthStateMachine
+) {
+
+    suspend fun execute(): TOTPSetupDetails {
+        val state = stateMachine.requireSignedInState()
+
+        val token = fetchAuthSession.execute().requireAccessToken()
+
+        val response = client.associateSoftwareToken {
+            accessToken = token
+        }
+
+        val sharedSecret = response.secretCode ?: throw AuthException(
+            "Shared secret missing from response",
+            AmplifyException.REPORT_BUG_TO_AWS_SUGGESTION
+        )
+
+        return TOTPSetupDetails(
+            sharedSecret = sharedSecret,
+            username = state.signedInData.username
+        )
+    }
+}

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

@@ -0,0 +1,56 @@
+/*
+ * 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.VerifySoftwareTokenResponseType
+import aws.sdk.kotlin.services.cognitoidentityprovider.verifySoftwareToken
+import com.amplifyframework.AmplifyException
+import com.amplifyframework.auth.cognito.AuthStateMachine
+import com.amplifyframework.auth.cognito.options.AWSCognitoAuthVerifyTOTPSetupOptions
+import com.amplifyframework.auth.cognito.requireAccessToken
+import com.amplifyframework.auth.cognito.requireSignedInState
+import com.amplifyframework.auth.exceptions.ServiceException
+import com.amplifyframework.auth.options.AuthVerifyTOTPSetupOptions
+
+internal class VerifyTotpSetupUseCase(
+    private val fetchAuthSession: FetchAuthSessionUseCase,
+    private val client: CognitoIdentityProviderClient,
+    private val stateMachine: AuthStateMachine
+) {
+
+    suspend fun execute(code: String, options: AuthVerifyTOTPSetupOptions) {
+        val cognitoOptions = options as? AWSCognitoAuthVerifyTOTPSetupOptions
+        val deviceName = cognitoOptions?.friendlyDeviceName
+
+        stateMachine.requireSignedInState()
+
+        val token = fetchAuthSession.execute().requireAccessToken()
+
+        val response = client.verifySoftwareToken {
+            userCode = code
+            friendlyDeviceName = deviceName
+            accessToken = token
+        }
+
+        if (response.status != VerifySoftwareTokenResponseType.Success) {
+            throw ServiceException(
+                message = "An unknown service error has occurred",
+                recoverySuggestion = AmplifyException.TODO_RECOVERY_SUGGESTION
+            )
+        }
+    }
+}

aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/AWSCognitoAuthPluginTest.kt

@@ -692,17 +692,23 @@ class AWSCognitoAuthPluginTest {
     fun setUpTOTP() {
         val expectedOnSuccess = Consumer<TOTPSetupDetails> { }
         val expectedOnError = Consumer<AuthException> { }
+
+        val useCase = authPlugin.useCaseFactory.setupTotp()
+
         authPlugin.setUpTOTP(expectedOnSuccess, expectedOnError)
-        verify(timeout = CHANNEL_TIMEOUT) { realPlugin.setUpTOTP(any(), any()) }
+        coVerify(timeout = CHANNEL_TIMEOUT) { useCase.execute() }
     }
 
     @Test
     fun verifyTOTPSetup() {
         val code = "123456"
         val expectedOnSuccess = Action { }
         val expectedOnError = Consumer<AuthException> { }
+
+        val useCase = authPlugin.useCaseFactory.verifyTotpSetup()
+
         authPlugin.verifyTOTPSetup(code, expectedOnSuccess, expectedOnError)
-        verify(timeout = CHANNEL_TIMEOUT) { realPlugin.verifyTOTPSetup(code, any(), any(), any()) }
+        coVerify(timeout = CHANNEL_TIMEOUT) { useCase.execute(code, any()) }
     }
 
     @Test
@@ -711,8 +717,11 @@ class AWSCognitoAuthPluginTest {
         val options = AWSCognitoAuthVerifyTOTPSetupOptions.CognitoBuilder().friendlyDeviceName("DEVICE_NAME").build()
         val expectedOnSuccess = Action { }
         val expectedOnError = Consumer<AuthException> { }
+
+        val useCase = authPlugin.useCaseFactory.verifyTotpSetup()
+
         authPlugin.verifyTOTPSetup(code, options, expectedOnSuccess, expectedOnError)
-        verify(timeout = CHANNEL_TIMEOUT) { realPlugin.verifyTOTPSetup(code, options, any(), any()) }
+        coVerify(timeout = CHANNEL_TIMEOUT) { useCase.execute(code, options) }
     }
 
     @Test