diff --git a/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/RealAWSCognitoAuthPluginTest.kt b/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/RealAWSCognitoAuthPluginTest.kt
index 3099bca7f..4c8ae9736 100644
--- a/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/RealAWSCognitoAuthPluginTest.kt
+++ b/aws-auth-cognito/src/test/java/com/amplifyframework/auth/cognito/RealAWSCognitoAuthPluginTest.kt
@@ -98,13 +98,14 @@ import com.amplifyframework.statemachine.codegen.errors.SessionError
 import com.amplifyframework.statemachine.codegen.states.AuthState
 import com.amplifyframework.statemachine.codegen.states.AuthenticationState
 import com.amplifyframework.statemachine.codegen.states.AuthorizationState
+import com.amplifyframework.testutils.await
 import featureTest.utilities.APICaptorFactory.Companion.onError
+import io.kotest.assertions.throwables.shouldThrowWithMessage
+import io.kotest.matchers.booleans.shouldBeTrue
 import io.mockk.coEvery
-import io.mockk.coJustRun
 import io.mockk.coVerify
 import io.mockk.every
 import io.mockk.invoke
-import io.mockk.justRun
 import io.mockk.mockk
 import io.mockk.mockkConstructor
 import io.mockk.mockkObject
@@ -120,11 +121,15 @@ import kotlin.test.assertFalse
 import kotlin.test.assertNotNull
 import kotlin.test.assertNull
 import kotlin.test.assertTrue
+import kotlin.time.Duration
+import kotlin.time.Duration.Companion.seconds
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import org.json.JSONObject
 import org.junit.After
 import org.junit.Before
 import org.junit.Test
 
+@OptIn(ExperimentalCoroutinesApi::class)
 class RealAWSCognitoAuthPluginTest {
 
     private val dummyToken = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VybmFtZSI6IkpvaG4gRG9lIiwiZXhwIjoxNTE2Mj" +
@@ -249,7 +254,7 @@ class RealAWSCognitoAuthPluginTest {
     @Test
     fun testGetCurrentUserSucceedsIfSignedIn() {
         // GIVEN
-        val onSuccess = mockk<Consumer<AuthUser>>()
+        val onSuccess = ConsumerWithLatch<AuthUser>()
         val onError = mockk<Consumer<AuthException>>()
         mockkObject(SessionHelper)
         every { SessionHelper.getUsername(any()) } returns "username"
@@ -258,7 +263,7 @@ class RealAWSCognitoAuthPluginTest {
         plugin.getCurrentUser(onSuccess, onError)
 
         // THEN
-        verify { onSuccess.accept(any()) }
+        onSuccess.shouldBeCalled()
         verify(exactly = 0) { onError.accept(any()) }
     }
 
@@ -266,7 +271,7 @@ class RealAWSCognitoAuthPluginTest {
     fun testGetCurrentUserFailsWithInvalidStateException() {
         // GIVEN
         val onSuccess = mockk<Consumer<AuthUser>>()
-        val onError = mockk<Consumer<AuthException>>(relaxed = true)
+        val onError = ConsumerWithLatch<AuthException>(expect = InvalidStateException())
 
         setupCurrentAuthState(authNState = AuthenticationState.NotConfigured())
 
@@ -274,15 +279,15 @@ class RealAWSCognitoAuthPluginTest {
         plugin.getCurrentUser(onSuccess, onError)
 
         // THEN
+        onError.shouldBeCalled()
         verify(exactly = 0) { onSuccess.accept(any()) }
-        verify { onError.accept(InvalidStateException()) }
     }
 
     @Test
     fun testGetCurrentUserFailsWithSignedOutException() {
         // GIVEN
         val onSuccess = mockk<Consumer<AuthUser>>()
-        val onError = mockk<Consumer<AuthException>>(relaxed = true)
+        val onError = ConsumerWithLatch<AuthException>(expect = SignedOutException())
 
         setupCurrentAuthState(
             authNState = AuthenticationState.SignedOut(mockk()),
@@ -292,23 +297,18 @@ class RealAWSCognitoAuthPluginTest {
         plugin.getCurrentUser(onSuccess, onError)
 
         // THEN
+        onError.shouldBeCalled()
         verify(exactly = 0) { onSuccess.accept(any()) }
-        verify { onError.accept(SignedOutException()) }
     }
 
     @Test
     fun testGetCurrentUserFailsWithExpiredSessionException() {
         // GIVEN
         val onGetCurrentUserSuccess = mockk<Consumer<AuthUser>>()
-        val onGetCurrentUserError = mockk<Consumer<AuthException>>(relaxed = true)
+        val onGetCurrentUserError = ConsumerWithLatch<AuthException>(expect = SessionExpiredException())
         val sessionExpiredException = SessionExpiredException()
         val sessionError = SessionError(sessionExpiredException, credentials)
-        val authNState = AuthenticationState.SignedIn(
-            mockk {
-                every { username } returns "username"
-            },
-            mockk()
-        )
+        val authNState = AuthenticationState.SignedIn(mockk { every { username } returns "username" }, mockk())
         val authZState = AuthorizationState.Error(sessionError)
 
         setupCurrentAuthState(
@@ -318,9 +318,7 @@ class RealAWSCognitoAuthPluginTest {
 
         val sessionErrorState = mockk<AuthState> {
             every { this@mockk.authNState } returns AuthenticationState.SignedIn(
-                mockk {
-                    every { username } returns "username"
-                },
+                mockk { every { username } returns "username" },
                 mockk()
             )
             every { this@mockk.authZState } returns AuthorizationState.Error(sessionError)
@@ -336,8 +334,8 @@ class RealAWSCognitoAuthPluginTest {
         plugin.getCurrentUser(onGetCurrentUserSuccess, onGetCurrentUserError)
 
         // THEN
+        onGetCurrentUserError.shouldBeCalled()
         verify(exactly = 0) { onGetCurrentUserSuccess.accept(any()) }
-        verify(timeout = 1000L) { onGetCurrentUserError.accept(sessionExpiredException) }
     }
 
     @Test
@@ -364,9 +362,9 @@ class RealAWSCognitoAuthPluginTest {
     @Test
     fun testConfirmSignUpFailsIfNotConfigured() {
         // GIVEN
-        val onSuccess = mockk<Consumer<AuthSignUpResult>>()
-        val onError = mockk<Consumer<AuthException>>(relaxed = true)
         val expectedAuthError = InvalidUserPoolConfigurationException()
+        val onSuccess = mockk<Consumer<AuthSignUpResult>>()
+        val onError = ConsumerWithLatch<AuthException>(expect = expectedAuthError)
 
         setupCurrentAuthState(authNState = AuthenticationState.NotConfigured())
 
@@ -374,16 +372,16 @@ class RealAWSCognitoAuthPluginTest {
         plugin.confirmSignUp("user", "pass", AuthConfirmSignUpOptions.defaults(), onSuccess, onError)
 
         // THEN
+        onError.shouldBeCalled()
         verify(exactly = 0) { onSuccess.accept(any()) }
-        verify { onError.accept(expectedAuthError) }
     }
 
     @Test
     fun testSignInFailsIfNotConfigured() {
         // GIVEN
-        val onSuccess = mockk<Consumer<AuthSignInResult>>()
-        val onError = mockk<Consumer<AuthException>>(relaxed = true)
         val expectedAuthError = InvalidUserPoolConfigurationException()
+        val onSuccess = mockk<Consumer<AuthSignInResult>>()
+        val onError = ConsumerWithLatch<AuthException>(expect = expectedAuthError)
 
         setupCurrentAuthState(authNState = AuthenticationState.NotConfigured())
 
@@ -392,14 +390,14 @@ class RealAWSCognitoAuthPluginTest {
         plugin.signIn("user", "password", AuthSignInOptions.defaults(), onSuccess, onError)
 
         // THEN
+        onError.shouldBeCalled()
         verify(exactly = 0) { onSuccess.accept(any()) }
-        verify { onError.accept(expectedAuthError) }
     }
 
     @Test
     fun testSignInFailsIfAlreadySignedIn() {
         // GIVEN
-        val onError = mockk<Consumer<AuthException>>(relaxed = true)
+        val onError = ConsumerWithLatch<AuthException>()
         coEvery { authConfiguration.authFlowType } returns AuthFlowType.USER_SRP_AUTH
 
         setupCurrentAuthState(
@@ -419,14 +417,13 @@ class RealAWSCognitoAuthPluginTest {
         plugin.signIn("user", "password", AuthSignInOptions.defaults(), mockk(), onError)
 
         // THEN
-        verify { onError.accept(any()) }
+        onError.shouldBeCalled()
     }
 
     @Test
     fun testResendSignUpCodeFailsIfNotConfigured() {
         // GIVEN
-        val onError = mockk<Consumer<AuthException>>(relaxed = true)
-        val expectedAuthError = InvalidUserPoolConfigurationException()
+        val onError = ConsumerWithLatch<AuthException>(expect = InvalidUserPoolConfigurationException())
 
         setupCurrentAuthState(authNState = AuthenticationState.NotConfigured())
 
@@ -434,7 +431,7 @@ class RealAWSCognitoAuthPluginTest {
         plugin.resendSignUpCode("user", AuthResendSignUpCodeOptions.defaults(), mockk(), onError)
 
         // THEN
-        verify { onError.accept(expectedAuthError) }
+        onError.shouldBeCalled()
     }
 
     @Test
@@ -449,7 +446,7 @@ class RealAWSCognitoAuthPluginTest {
         // WHEN
         plugin.updatePassword("old", "new", onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
     }
 
     @Test
@@ -461,7 +458,7 @@ class RealAWSCognitoAuthPluginTest {
 
         // WHEN
         plugin.updatePassword("old", "new", mockk(), onError)
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
     }
 
     @Test
@@ -469,92 +466,74 @@ class RealAWSCognitoAuthPluginTest {
         val onError = ConsumerWithLatch<AuthException>()
 
         plugin.updatePassword("old", "new", mockk(), onError)
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
     }
 
     @Test
     fun `reset password fails if cognitoIdentityProviderClient is not set`() {
         // GIVEN
         val onSuccess = mockk<Consumer<AuthResetPasswordResult>>()
-        val onError = mockk<Consumer<AuthException>>(relaxed = true)
-        val expectedAuthError = InvalidUserPoolConfigurationException()
+        val onError = ConsumerWithLatch<AuthException>(expect = InvalidUserPoolConfigurationException())
 
         val userPool = UserPoolConfiguration.invoke { appClientId = "app Client Id" }
         every { authService.cognitoIdentityProviderClient } returns null
         every { authConfiguration.userPool } returns userPool
 
-        val errorCaptor = slot<InvalidUserPoolConfigurationException>()
-        justRun { onError.accept(capture(errorCaptor)) }
-
         // WHEN
         plugin.resetPassword("user", AuthResetPasswordOptions.defaults(), onSuccess, onError)
 
         // THEN
+        onError.shouldBeCalled()
         verify(exactly = 0) { onSuccess.accept(any()) }
-        assertEquals(expectedAuthError.toString(), errorCaptor.captured.toString())
     }
 
     @Test
     fun `reset password fails if appClientId is not set`() {
         // GIVEN
         val onSuccess = mockk<Consumer<AuthResetPasswordResult>>()
-        val onError = mockk<Consumer<AuthException>>(relaxed = true)
-        val expectedAuthError = InvalidUserPoolConfigurationException()
+        val onError = ConsumerWithLatch<AuthException>(expect = InvalidUserPoolConfigurationException())
 
         val userPool = UserPoolConfiguration.invoke { appClientId = null }
         every { authService.cognitoIdentityProviderClient } returns mockk()
         every { authConfiguration.userPool } returns userPool
 
-        val errorCaptor = slot<InvalidUserPoolConfigurationException>()
-        justRun { onError.accept(capture(errorCaptor)) }
-
         // WHEN
         plugin.resetPassword("user", AuthResetPasswordOptions.defaults(), onSuccess, onError)
 
         // THEN
+        onError.shouldBeCalled()
         verify(exactly = 0) { onSuccess.accept(any()) }
-        assertEquals(expectedAuthError.toString(), errorCaptor.captured.toString())
     }
+
     @Test
     fun `reset password executes ResetPasswordUseCase if required params are set`() {
         // GIVEN
-        val onSuccess = mockk<Consumer<AuthResetPasswordResult>>()
+        val onSuccess = ConsumerWithLatch<AuthResetPasswordResult>()
         val onError = mockk<Consumer<AuthException>>()
         val options = mockk<AuthResetPasswordOptions>()
         val username = "user"
         val pinpointAppId = "abc"
         val encodedData = "encodedData"
 
-        coEvery {
-            authEnvironment.getUserContextData(username)
-        } returns encodedData
-
-        every {
-            authEnvironment.getPinpointEndpointId()
-        } returns pinpointAppId
+        coEvery { authEnvironment.getUserContextData(username) } returns encodedData
+        every { authEnvironment.getPinpointEndpointId() } returns pinpointAppId
 
         mockkConstructor(ResetPasswordUseCase::class)
 
         every { authService.cognitoIdentityProviderClient } returns mockk()
         every { authConfiguration.userPool } returns UserPoolConfiguration.invoke { appClientId = "app Client Id" }
-        coJustRun {
-            anyConstructed<ResetPasswordUseCase>().execute(
-                username,
-                options,
-                encodedData,
-                pinpointAppId,
-                onSuccess,
-                onError
-            )
+
+        coEvery {
+            anyConstructed<ResetPasswordUseCase>().execute(any(), any(), any(), any(), any(), any())
+        } answers {
+            arg<Consumer<AuthResetPasswordResult>>(4).accept(mockk())
         }
 
         // WHEN
         plugin.resetPassword(username, options, onSuccess, onError)
 
         // THEN
-        coVerify {
-            anyConstructed<ResetPasswordUseCase>().execute(username, options, any(), pinpointAppId, onSuccess, onError)
-        }
+        onSuccess.shouldBeCalled()
     }
 
     @Test
@@ -592,7 +571,7 @@ class RealAWSCognitoAuthPluginTest {
         // WHEN
         plugin.fetchUserAttributes(onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertEquals(expectedResult, onSuccess.captured)
     }
 
@@ -606,7 +585,7 @@ class RealAWSCognitoAuthPluginTest {
         // WHEN
         plugin.fetchUserAttributes(mockk(), onError)
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
     }
 
     @Test
@@ -623,7 +602,7 @@ class RealAWSCognitoAuthPluginTest {
         plugin.confirmResetPassword("user", "pass", "code", mockk(), mockk(), onError)
 
         // Then
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
     }
 
     @Test
@@ -688,7 +667,7 @@ class RealAWSCognitoAuthPluginTest {
         plugin.confirmResetPassword(user, pass, code, AuthConfirmResetPasswordOptions.defaults(), onSuccess, mockk())
 
         // THEN
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
     }
 
     @Test
@@ -709,7 +688,7 @@ class RealAWSCognitoAuthPluginTest {
         plugin.confirmResetPassword(user, pass, code, AuthConfirmResetPasswordOptions.defaults(), mockk(), onError)
 
         // THEN
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
         assertEquals(expectedException, onError.captured.cause)
     }
 
@@ -807,9 +786,9 @@ class RealAWSCognitoAuthPluginTest {
 
         // WHEN
         plugin.signUp(username, password, options, onSuccess, mockk())
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
 
         // THEN
+        onSuccess.shouldBeCalled()
         assertEquals(expectedAuthSignUpResult, onSuccess.captured)
     }
 
@@ -967,7 +946,7 @@ class RealAWSCognitoAuthPluginTest {
 
         // WHEN
         plugin.resendSignUpCode(username, AuthResendSignUpCodeOptions.defaults(), onSuccess, onError)
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
 
         // THEN
         assertEquals(expectedCodeDeliveryDetails, onSuccess.captured)
@@ -987,7 +966,7 @@ class RealAWSCognitoAuthPluginTest {
             onError
         )
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
     }
 
     @Test
@@ -1008,7 +987,7 @@ class RealAWSCognitoAuthPluginTest {
             onError
         )
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
     }
 
     @Test
@@ -1027,7 +1006,7 @@ class RealAWSCognitoAuthPluginTest {
             onError
         )
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(onSuccess.captured.isUpdated, "attribute should be successfully updated")
         assertNotNull(onSuccess.captured.nextStep, "next step should not be null")
         assertNull(onSuccess.captured.nextStep.codeDeliveryDetails, "code delivery details should be null")
@@ -1058,7 +1037,7 @@ class RealAWSCognitoAuthPluginTest {
             mockk()
         )
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(onSuccess.captured.isUpdated, "attribute should be successfully updated")
         assertNotNull(onSuccess.captured.nextStep, "next step should not be null")
         assertNull(onSuccess.captured.nextStep.codeDeliveryDetails, "code delivery details should be null")
@@ -1100,7 +1079,7 @@ class RealAWSCognitoAuthPluginTest {
             mockk()
         )
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertEquals(onSuccess.captured.size, 2)
         // nickname
         assertNotNull(onSuccess.captured[AuthUserAttributeKey.nickname()], "nick name should be in result")
@@ -1163,7 +1142,7 @@ class RealAWSCognitoAuthPluginTest {
             onError
         )
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
     }
 
     @Test
@@ -1194,7 +1173,7 @@ class RealAWSCognitoAuthPluginTest {
             onError
         )
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
     }
 
     @Test
@@ -1214,7 +1193,7 @@ class RealAWSCognitoAuthPluginTest {
             onError
         )
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
         assertEquals(expectedException, onError.captured.cause)
     }
 
@@ -1234,7 +1213,7 @@ class RealAWSCognitoAuthPluginTest {
             mockk()
         )
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
     }
 
     @Test
@@ -1251,7 +1230,7 @@ class RealAWSCognitoAuthPluginTest {
             onError
         )
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
     }
 
     @Test
@@ -1281,7 +1260,7 @@ class RealAWSCognitoAuthPluginTest {
             onError
         )
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
     }
 
     @Test
@@ -1304,7 +1283,7 @@ class RealAWSCognitoAuthPluginTest {
             onError
         )
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
         assertEquals(expectedException, onError.captured.cause)
     }
 
@@ -1337,7 +1316,7 @@ class RealAWSCognitoAuthPluginTest {
             mockk()
         )
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         // nickname
         assertEquals(
             onSuccess.captured.attributeName,
@@ -1372,12 +1351,10 @@ class RealAWSCognitoAuthPluginTest {
         configJsonObject.put("Endpoint", invalidEndpoint)
         val expectedErrorMessage = "Invalid endpoint value $invalidEndpoint. Expected fully qualified hostname with " +
             "no scheme, no path and no query"
-        val message = try {
+
+        shouldThrowWithMessage<Exception>(expectedErrorMessage) {
             UserPoolConfiguration.fromJson(configJsonObject).build()
-        } catch (ex: Exception) {
-            ex.message
         }
-        assertEquals(message, expectedErrorMessage, "Error message do not match expected one")
     }
 
     @Test
@@ -1390,12 +1367,10 @@ class RealAWSCognitoAuthPluginTest {
         configJsonObject.put("Endpoint", invalidEndpoint)
         val expectedErrorMessage = "Invalid endpoint value $invalidEndpoint. Expected fully qualified hostname with " +
             "no scheme, no path and no query"
-        val message = try {
+
+        shouldThrowWithMessage<Exception>(expectedErrorMessage) {
             UserPoolConfiguration.fromJson(configJsonObject).build()
-        } catch (ex: Exception) {
-            ex.message
         }
-        assertEquals(message, expectedErrorMessage, "Error message do not match expected one")
     }
 
     @Test
@@ -1409,12 +1384,10 @@ class RealAWSCognitoAuthPluginTest {
         configJsonObject.put("Endpoint", invalidEndpoint)
         val expectedErrorMessage = "Invalid endpoint value $invalidEndpoint. Expected fully qualified hostname with " +
             "no scheme, no path and no query"
-        val message = try {
+
+        shouldThrowWithMessage<Exception>(expectedErrorMessage) {
             UserPoolConfiguration.fromJson(configJsonObject).build()
-        } catch (ex: Exception) {
-            ex.message
         }
-        assertEquals(message, expectedErrorMessage, "Error message do not match expected one")
     }
 
     @Test
@@ -1479,7 +1452,7 @@ class RealAWSCognitoAuthPluginTest {
 
         plugin.setUpTOTP(onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertEquals(secretCode, onSuccess.captured.sharedSecret)
     }
 
@@ -1496,7 +1469,7 @@ class RealAWSCognitoAuthPluginTest {
 
         plugin.setUpTOTP(mockk(), onError)
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
         assertEquals(expectedErrorMessage, onError.captured.cause?.message)
     }
 
@@ -1526,7 +1499,7 @@ class RealAWSCognitoAuthPluginTest {
             mockk()
         )
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
     }
 
     @Test
@@ -1557,7 +1530,7 @@ class RealAWSCognitoAuthPluginTest {
             onError
         )
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
         assertEquals(errorMessage, onError.captured.cause?.message)
     }
 
@@ -1579,7 +1552,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.fetchMFAPreference(onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertEquals(setOf(MFAType.SMS, MFAType.TOTP), onSuccess.captured.enabled)
         assertEquals(MFAType.TOTP, onSuccess.captured.preferred)
     }
@@ -1605,7 +1578,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.ENABLED, MFAPreference.PREFERRED, onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -1646,7 +1619,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.ENABLED, MFAPreference.ENABLED, onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -1687,7 +1660,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.ENABLED, MFAPreference.ENABLED, onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -1711,7 +1684,7 @@ class RealAWSCognitoAuthPluginTest {
         coEvery { mockCognitoIPClient.setUserMfaPreference(any()) } throws Exception()
         plugin.updateMFAPreference(null, null, mockk(), onError)
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
     }
 
     @Test
@@ -1725,7 +1698,7 @@ class RealAWSCognitoAuthPluginTest {
         } throws Exception()
         plugin.updateMFAPreference(null, null, mockk(), onError)
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
     }
 
     @Test
@@ -1751,7 +1724,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.ENABLED, MFAPreference.ENABLED, onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -1792,7 +1765,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.ENABLED, MFAPreference.DISABLED, onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -1833,7 +1806,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.DISABLED, MFAPreference.ENABLED, onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -1874,7 +1847,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.PREFERRED, MFAPreference.ENABLED, onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -1915,7 +1888,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.ENABLED, MFAPreference.PREFERRED, onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -1956,7 +1929,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.DISABLED, MFAPreference.PREFERRED, onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -1997,7 +1970,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.PREFERRED, MFAPreference.DISABLED, onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -2038,7 +2011,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.ENABLED, MFAPreference.DISABLED, onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -2079,7 +2052,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.DISABLED, MFAPreference.ENABLED, onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -2121,7 +2094,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.DISABLED, MFAPreference.ENABLED, onSuccess, onError)
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -2162,7 +2135,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.PREFERRED, MFAPreference.ENABLED, onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -2203,7 +2176,7 @@ class RealAWSCognitoAuthPluginTest {
         }
         plugin.updateMFAPreference(MFAPreference.ENABLED, MFAPreference.PREFERRED, onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertTrue(setUserMFAPreferenceRequest.isCaptured)
         assertEquals(
             SmsMfaSettingsType.invoke {
@@ -2233,7 +2206,7 @@ class RealAWSCognitoAuthPluginTest {
 
         plugin.forgetDevice(onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         coVerify { mockCognitoIPClient.forgetDevice(match { it.deviceKey == "test" }) }
     }
 
@@ -2249,7 +2222,7 @@ class RealAWSCognitoAuthPluginTest {
 
         plugin.forgetDevice(mockk(), onError)
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
         assertEquals("failed", onError.captured.cause?.message)
     }
 
@@ -2261,7 +2234,7 @@ class RealAWSCognitoAuthPluginTest {
 
         plugin.forgetDevice(AuthDevice.fromId("test"), onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         coVerify { mockCognitoIPClient.forgetDevice(match { it.deviceKey == "test" }) }
     }
 
@@ -2273,7 +2246,7 @@ class RealAWSCognitoAuthPluginTest {
 
         plugin.forgetDevice(AuthDevice.fromId("test"), mockk(), onError)
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
         assertEquals("failed", onError.captured.cause?.message)
     }
 
@@ -2286,7 +2259,10 @@ class RealAWSCognitoAuthPluginTest {
                 DeviceType.invoke {
                     deviceKey = "id1"
                     deviceAttributes = listOf(
-                        AttributeType.invoke { name = "device_name"; value = "name1" }
+                        AttributeType.invoke {
+                            name = "device_name"
+                            value = "name1"
+                        }
                     )
                 }
             )
@@ -2294,7 +2270,7 @@ class RealAWSCognitoAuthPluginTest {
 
         plugin.fetchDevices(onSuccess, mockk())
 
-        assertTrue { onSuccess.latch.await(5, TimeUnit.SECONDS) }
+        onSuccess.shouldBeCalled()
         assertEquals("id1", onSuccess.captured.first().id)
         assertEquals("name1", onSuccess.captured.first().name)
     }
@@ -2306,7 +2282,7 @@ class RealAWSCognitoAuthPluginTest {
 
         plugin.fetchDevices(mockk(), onError)
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
     }
 
     @Test
@@ -2316,7 +2292,7 @@ class RealAWSCognitoAuthPluginTest {
 
         plugin.fetchDevices(mockk(), onError)
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
         assertEquals(SignedOutException(), onError.captured)
     }
 
@@ -2327,14 +2303,11 @@ class RealAWSCognitoAuthPluginTest {
 
         plugin.fetchDevices(mockk(), onError)
 
-        assertTrue { onError.latch.await(5, TimeUnit.SECONDS) }
+        onError.shouldBeCalled()
         assertEquals(InvalidStateException(), onError.captured)
     }
 
-    private fun setupCurrentAuthState(
-        authNState: AuthenticationState? = null,
-        authZState: AuthorizationState? = null
-    ) {
+    private fun setupCurrentAuthState(authNState: AuthenticationState? = null, authZState: AuthorizationState? = null) {
         val currentAuthState = mockk<AuthState> {
             every { this@mockk.authNState } returns authNState
             every { this@mockk.authZState } returns authZState
@@ -2350,6 +2323,8 @@ class RealAWSCognitoAuthPluginTest {
             assertTrue(latch.count > 0)
             latch.countDown()
         }
+
+        fun shouldBeCalled(timeout: Duration = 5.seconds) = latch.await(timeout).shouldBeTrue()
     }
 
     private class ConsumerWithLatch<T : Any>(private val expect: T? = null, count: Int = 1) : Consumer<T> {
@@ -2363,5 +2338,7 @@ class RealAWSCognitoAuthPluginTest {
                 latch.countDown()
             }
         }
+
+        fun shouldBeCalled(timeout: Duration = 5.seconds) = latch.await(timeout).shouldBeTrue()
     }
 }
diff --git a/testutils/src/main/java/com/amplifyframework/testutils/CountdownLatchExtensions.kt b/testutils/src/main/java/com/amplifyframework/testutils/CountdownLatchExtensions.kt
new file mode 100644
index 000000000..a8786b559
--- /dev/null
+++ b/testutils/src/main/java/com/amplifyframework/testutils/CountdownLatchExtensions.kt
@@ -0,0 +1,25 @@
+/*
+ * Copyright 2024 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.testutils
+
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
+import kotlin.time.Duration
+
+/**
+ * Await using a [Duration]
+ */
+fun CountDownLatch.await(duration: Duration) = await(duration.inWholeMilliseconds, TimeUnit.MILLISECONDS)