fix(auth): fetch device metadata from credential store for remember and forgetDevice (#2601)

Co-authored-by: Thomas Leing <[email protected]>

Author: tjleing

aws-auth-cognito/src/androidTest/java/com/amplifyframework/auth/cognito/AWSCognitoAuthPluginInstrumentationTests.kt

@@ -179,6 +179,70 @@ class AWSCognitoAuthPluginInstrumentationTests {
         }
     }
 
+    @Test
+    fun rememberDevice_succeeds_after_signIn_and_signOut() {
+        signInWithCognito()
+
+        val rememberLatch = CountDownLatch(1)
+
+        auth.rememberDevice(
+            {
+                rememberLatch.countDown()
+            },
+            {
+                rememberLatch.countDown()
+                assertTrue(false)
+            }
+        )
+
+        rememberLatch.await(10, TimeUnit.SECONDS)
+
+        val forgetLatch = CountDownLatch(1)
+
+        auth.forgetDevice(
+            {
+                forgetLatch.countDown()
+            },
+            {
+                forgetLatch.countDown()
+                assertTrue(false)
+            }
+        )
+
+        forgetLatch.await(10, TimeUnit.SECONDS)
+
+        signOut()
+        signInWithCognito()
+
+        val rememberLatch2 = CountDownLatch(1)
+
+        auth.rememberDevice(
+            {
+                rememberLatch2.countDown()
+            },
+            {
+                assertTrue(false)
+                rememberLatch2.countDown()
+            }
+        )
+
+        rememberLatch2.await(10, TimeUnit.SECONDS)
+
+        val forgetLatch2 = CountDownLatch(1)
+
+        auth.forgetDevice(
+            {
+                forgetLatch2.countDown()
+            },
+            {
+                assertTrue(false)
+                forgetLatch2.countDown()
+            }
+        )
+
+        forgetLatch2.await(10, TimeUnit.SECONDS)
+    }
+
     private fun signInWithCognito(synchronous: Boolean = true) {
         val context = ApplicationProvider.getApplicationContext<Context>()
         val (username, password) = Credentials.load(context)

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

@@ -133,7 +133,6 @@ import com.amplifyframework.statemachine.StateChangeListenerToken
 import com.amplifyframework.statemachine.codegen.data.AmplifyCredential
 import com.amplifyframework.statemachine.codegen.data.AuthChallenge
 import com.amplifyframework.statemachine.codegen.data.AuthConfiguration
-import com.amplifyframework.statemachine.codegen.data.DeviceMetadata
 import com.amplifyframework.statemachine.codegen.data.FederatedToken
 import com.amplifyframework.statemachine.codegen.data.HostedUIErrorData
 import com.amplifyframework.statemachine.codegen.data.SignInData
@@ -1223,12 +1222,14 @@ internal class RealAWSCognitoAuthPlugin(
         authStateMachine.getCurrentState { authState ->
             when (val state = authState.authNState) {
                 is AuthenticationState.SignedIn -> {
-                    updateDevice(
-                        (state.deviceMetadata as? DeviceMetadata.Metadata)?.deviceKey,
-                        DeviceRememberedStatusType.Remembered,
-                        onSuccess,
-                        onError
-                    )
+                    GlobalScope.launch {
+                        updateDevice(
+                            authEnvironment.getDeviceMetadata(state.signedInData.username)?.deviceKey,
+                            DeviceRememberedStatusType.Remembered,
+                            onSuccess,
+                            onError
+                        )
+                    }
                 }
                 is AuthenticationState.SignedOut -> {
                     onError.accept(SignedOutException())
@@ -1276,8 +1277,11 @@ internal class RealAWSCognitoAuthPlugin(
             when (val authState = authState.authNState) {
                 is AuthenticationState.SignedIn -> {
                     if (device.id.isEmpty()) {
-                        val deviceKey = (authState.deviceMetadata as? DeviceMetadata.Metadata)?.deviceKey
-                        updateDevice(deviceKey, DeviceRememberedStatusType.NotRemembered, onSuccess, onError)
+                        GlobalScope.launch {
+                            val deviceKey = authEnvironment.getDeviceMetadata(authState.signedInData.username)
+                                ?.deviceKey
+                            updateDevice(deviceKey, DeviceRememberedStatusType.NotRemembered, onSuccess, onError)
+                        }
                     } else {
                         updateDevice(device.id, DeviceRememberedStatusType.NotRemembered, onSuccess, onError)
                     }