chore(auth): Cleanup RealAWSCognitoAuthPluginTest (#2690)

aws-amplify · Feb 7, 2024 · 932c4f8 · 932c4f8
1 parent 2ee5eb2
commit 932c4f8
Show file tree

Hide file tree

Showing 3 changed files with 433 additions and 822 deletions.
diff --git a/.idea/codeStyles/Project.xml b/.idea/codeStyles/Project.xml
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
@@ -27,6 +27,7 @@ import aws.sdk.kotlin.services.cognitoidentityprovider.model.AttributeType
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.ChallengeNameType
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.ChangePasswordRequest
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.DeviceRememberedStatusType
+import aws.sdk.kotlin.services.cognitoidentityprovider.model.ForgetDeviceRequest
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.GetUserAttributeVerificationCodeRequest
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.GetUserRequest
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.ListDevicesRequest
@@ -1291,16 +1292,21 @@ internal class RealAWSCognitoAuthPlugin(
         onError: Consumer<AuthException>
     ) {
         authStateMachine.getCurrentState { authState ->
-            when (val authState = authState.authNState) {
+            when (val authNState = authState.authNState) {
                 is AuthenticationState.SignedIn -> {
-                    if (device.id.isEmpty()) {
-                        GlobalScope.launch {
-                            val deviceKey = authEnvironment.getDeviceMetadata(authState.signedInData.username)
-                                ?.deviceKey
-                            updateDevice(deviceKey, DeviceRememberedStatusType.NotRemembered, onSuccess, onError)
+                    GlobalScope.launch {
+                        try {
+                            if (device.id.isEmpty()) {
+                                val deviceKey = authEnvironment.getDeviceMetadata(authNState.signedInData.username)
+                                    ?.deviceKey
+                                forgetDevice(deviceKey)
+                            } else {
+                                forgetDevice(device.id)
+                            }
+                            onSuccess.call()
+                        } catch (e: Exception) {
+                            onError.accept(CognitoAuthExceptionConverter.lookup(e, "Failed to forget device."))
                         }
-                    } else {
-                        updateDevice(device.id, DeviceRememberedStatusType.NotRemembered, onSuccess, onError)
                     }
                 }
                 is AuthenticationState.SignedOut -> {
@@ -1313,6 +1319,16 @@ internal class RealAWSCognitoAuthPlugin(
         }
     }
 
+    private suspend fun forgetDevice(alternateDeviceId: String?) {
+        val tokens = getSession().userPoolTokensResult
+        authEnvironment.cognitoAuthService.cognitoIdentityProviderClient?.forgetDevice(
+            ForgetDeviceRequest.invoke {
+                accessToken = tokens.value?.accessToken
+                deviceKey = alternateDeviceId
+            }
+        )
+    }
+
     override fun fetchDevices(
         onSuccess: Consumer<List<AuthDevice>>,
         onError: Consumer<AuthException>
@@ -1342,12 +1358,14 @@ internal class RealAWSCognitoAuthPlugin(
                             accessToken = tokens.value?.accessToken
                         }
                     )
-                val _devices = response?.devices
-                val authdeviceList = mutableListOf<AuthDevice>()
-                _devices?.forEach {
-                    authdeviceList.add(AuthDevice.fromId(it.deviceKey ?: ""))
-                }
-                onSuccess.accept(authdeviceList)
+
+                val devices = response?.devices?.map { device ->
+                    val id = device.deviceKey ?: ""
+                    val name = device.deviceAttributes?.find { it.name == "device_name" }?.value
+                    AuthDevice.fromId(id, name)
+                } ?: emptyList()
+
+                onSuccess.accept(devices)
             } catch (e: Exception) {
                 onError.accept(CognitoAuthExceptionConverter.lookup(e, "Fetch devices failed."))
             }