PM-13067 Navigate to setup unlock screen from action card in security settings (#4023)

app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/accountsetup/SetupUnlockNavigation.kt

@@ -1,29 +1,88 @@
 package com.x8bit.bitwarden.ui.auth.feature.accountsetup
 
+import androidx.lifecycle.SavedStateHandle
 import androidx.navigation.NavController
 import androidx.navigation.NavGraphBuilder
 import androidx.navigation.NavOptions
+import androidx.navigation.NavType
+import androidx.navigation.navArgument
+import com.x8bit.bitwarden.data.platform.annotation.OmitFromCoverage
 import com.x8bit.bitwarden.ui.platform.base.util.composableWithPushTransitions
+import com.x8bit.bitwarden.ui.platform.base.util.composableWithSlideTransitions
 
 /**
- * Route for [SetupUnlockScreen]
+ * Route constants for [SetupUnlockScreen]
  */
-const val SETUP_UNLOCK_ROUTE = "setup_unlock"
+private const val SETUP_UNLOCK_PREFIX = "setup_unlock"
+private const val SETUP_UNLOCK_AS_ROOT_PREFIX = "${SETUP_UNLOCK_PREFIX}_as_root"
+private const val SETUP_UNLOCK_INITIAL_SETUP_ARG = "isInitialSetup"
+const val SETUP_UNLOCK_AS_ROOT_ROUTE = "$SETUP_UNLOCK_AS_ROOT_PREFIX/" +
+    "{$SETUP_UNLOCK_INITIAL_SETUP_ARG}"
+private const val SETUP_UNLOCK_ROUTE = "$SETUP_UNLOCK_PREFIX/{$SETUP_UNLOCK_INITIAL_SETUP_ARG}"
+
+/**
+ * Class to retrieve setup unlock arguments from the [SavedStateHandle].
+ */
+@OmitFromCoverage
+data class SetupUnlockArgs(
+    val isInitialSetup: Boolean,
+) {
+    constructor(savedStateHandle: SavedStateHandle) : this(
+        isInitialSetup = requireNotNull(savedStateHandle[SETUP_UNLOCK_INITIAL_SETUP_ARG]),
+    )
+}
 
 /**
  * Navigate to the setup unlock screen.
  */
 fun NavController.navigateToSetupUnlockScreen(navOptions: NavOptions? = null) {
-    this.navigate(SETUP_UNLOCK_ROUTE, navOptions)
+    this.navigate("$SETUP_UNLOCK_PREFIX/false", navOptions)
 }
 
 /**
- * Add the setup unlock screen to the nav graph.
+ * Navigate to the setup unlock screen as root.
  */
-fun NavGraphBuilder.setupUnlockDestination() {
+fun NavController.navigateToSetupUnlockScreenAsRoot(navOptions: NavOptions? = null) {
-    composableWithPushTransitions(
+    this.navigate("$SETUP_UNLOCK_AS_ROOT_PREFIX/true", navOptions)
+}
+
+/**
+ * Add the setup unlock screen to a nav graph.
+ */
+fun NavGraphBuilder.setupUnlockDestination(
+    onNavigateBack: () -> Unit,
+) {
+    composableWithSlideTransitions(
         route = SETUP_UNLOCK_ROUTE,
+        arguments = setupUnlockArguments,
     ) {
-        SetupUnlockScreen()
+        SetupUnlockScreen(
+            onNavigateBack = onNavigateBack,
+        )
     }
 }
+
+/**
+ * Add the setup unlock screen to the root nav graph.
+ */
+fun NavGraphBuilder.setupUnlockDestinationAsRoot() {
+    composableWithPushTransitions(
+        route = SETUP_UNLOCK_AS_ROOT_ROUTE,
+        arguments = setupUnlockArguments,
+    ) {
+        SetupUnlockScreen(
+            onNavigateBack = {
+                // No-Op
+            },
+        )
+    }
+}
+
+private val setupUnlockArguments = listOf(
+    navArgument(
+        name = SETUP_UNLOCK_INITIAL_SETUP_ARG,
+        builder = {
+            type = NavType.BoolType
+        },
+    ),
+)

app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/accountsetup/SetupUnlockScreen.kt

@@ -40,6 +40,7 @@ import com.x8bit.bitwarden.ui.auth.feature.accountsetup.handlers.SetupUnlockHand
 import com.x8bit.bitwarden.ui.platform.base.util.EventsEffect
 import com.x8bit.bitwarden.ui.platform.base.util.standardHorizontalMargin
 import com.x8bit.bitwarden.ui.platform.components.appbar.BitwardenTopAppBar
+import com.x8bit.bitwarden.ui.platform.components.appbar.NavigationIcon
 import com.x8bit.bitwarden.ui.platform.components.button.BitwardenFilledButton
 import com.x8bit.bitwarden.ui.platform.components.button.BitwardenTextButton
 import com.x8bit.bitwarden.ui.platform.components.dialog.BasicDialogState
@@ -60,10 +61,12 @@ import com.x8bit.bitwarden.ui.platform.util.isPortrait
  * Top level composable for the setup unlock screen.
  */
 @OptIn(ExperimentalMaterial3Api::class)
+@Suppress("LongMethod")
 @Composable
 fun SetupUnlockScreen(
     viewModel: SetupUnlockViewModel = hiltViewModel(),
     biometricsManager: BiometricsManager = LocalBiometricsManager.current,
+    onNavigateBack: () -> Unit,
 ) {
     val state by viewModel.stateFlow.collectAsStateWithLifecycle()
     val handler = remember(viewModel) { SetupUnlockHandler.create(viewModel = viewModel) }
@@ -83,6 +86,8 @@ fun SetupUnlockScreen(
                     cipher = event.cipher,
                 )
             }
+
+            SetupUnlockEvent.NavigateBack -> onNavigateBack()
         }
     }
 
@@ -100,9 +105,27 @@ fun SetupUnlockScreen(
             .nestedScroll(scrollBehavior.nestedScrollConnection),
         topBar = {
             BitwardenTopAppBar(
-                title = stringResource(id = R.string.account_setup),
+                title = stringResource(
+                    id = if (state.isInitialSetup) {
+                        R.string.account_setup
+                    } else {
+                        R.string.set_up_unlock
+                    },
+                ),
                 scrollBehavior = scrollBehavior,
-                navigationIcon = null,
+                navigationIcon = if (state.isInitialSetup) {
+                    null
+                } else {
+                    NavigationIcon(
+                        navigationIcon = rememberVectorPainter(id = R.drawable.ic_close),
+                        navigationIconContentDescription = stringResource(id = R.string.close),
+                        onNavigationIconClick = remember(viewModel) {
+                            {
+                                viewModel.trySendAction(SetupUnlockAction.CloseClick)
+                            }
+                        },
+                    )
+                },
             )
         },
     ) { innerPadding ->
@@ -169,14 +192,16 @@ private fun SetupUnlockScreenContent(
         )
 
         Spacer(modifier = Modifier.height(height = 12.dp))
-        SetUpLaterButton(
+        if (state.isInitialSetup) {
-            onConfirmClick = handler.onSetUpLaterClick,
+            SetUpLaterButton(
-            modifier = Modifier
+                onConfirmClick = handler.onSetUpLaterClick,
-                .fillMaxWidth()
+                modifier = Modifier
-                .standardHorizontalMargin(),
+                    .fillMaxWidth()
-        )
+                    .standardHorizontalMargin(),
+            )
 
-        Spacer(modifier = Modifier.height(height = 12.dp))
+            Spacer(modifier = Modifier.height(height = 12.dp))
+        }
         Spacer(modifier = Modifier.navigationBarsPadding())
     }
 }

app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/accountsetup/SetupUnlockViewModel.kt

@@ -25,6 +25,7 @@ private const val KEY_STATE = "state"
 /**
  * Models logic for the setup unlock screen.
  */
+@Suppress("TooManyFunctions")
 @HiltViewModel
 class SetupUnlockViewModel @Inject constructor(
     savedStateHandle: SavedStateHandle,
@@ -38,6 +39,8 @@ class SetupUnlockViewModel @Inject constructor(
             userId = userId,
             cipher = biometricsEncryptionManager.getOrCreateCipher(userId = userId),
         )
+        // whether or not the user has completed the initial setup prior to this.
+        val isInitialSetup = SetupUnlockArgs(savedStateHandle).isInitialSetup
         SetupUnlockState(
             userId = userId,
             isUnlockWithPasswordEnabled = authRepository
@@ -49,6 +52,7 @@ class SetupUnlockViewModel @Inject constructor(
             isUnlockWithBiometricsEnabled = settingsRepository.isUnlockWithBiometricsEnabled &&
                 isBiometricsValid,
             dialogState = null,
+            isInitialSetup = isInitialSetup,
         )
     },
 ) {
@@ -64,11 +68,20 @@ class SetupUnlockViewModel @Inject constructor(
 
             is SetupUnlockAction.UnlockWithPinToggle -> handleUnlockWithPinToggle(action)
             is SetupUnlockAction.Internal -> handleInternalActions(action)
+            SetupUnlockAction.CloseClick -> handleCloseClick()
         }
     }
 
+    private fun handleCloseClick() {
+        sendEvent(SetupUnlockEvent.NavigateBack)
+    }
+
     private fun handleContinueClick() {
-        updateOnboardingStatusToNextStep()
+        if (state.isInitialSetup) {
+            updateOnboardingStatusToNextStep()
+        } else {
+            sendEvent(SetupUnlockEvent.NavigateBack)
+        }
     }
 
     private fun handleEnableBiometricsClick() {
@@ -196,6 +209,7 @@ data class SetupUnlockState(
     val isUnlockWithPinEnabled: Boolean,
     val isUnlockWithBiometricsEnabled: Boolean,
     val dialogState: DialogState?,
+    val isInitialSetup: Boolean,
 ) : Parcelable {
     /**
      * Indicates whether the continue button should be enabled or disabled.
@@ -237,6 +251,11 @@ sealed class SetupUnlockEvent {
     data class ShowBiometricsPrompt(
         val cipher: Cipher,
     ) : SetupUnlockEvent()
+
+    /**
+     * Navigates back to the previous screen.
+     */
+    data object NavigateBack : SetupUnlockEvent()
 }
 
 /**
@@ -277,6 +296,11 @@ sealed class SetupUnlockAction {
      */
     data object DismissDialog : SetupUnlockAction()
 
+    /**
+     * The user has clicked the close button.
+     */
+    data object CloseClick : SetupUnlockAction()
+
     /**
      * Models actions that can be sent by the view model itself.
      */

app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/rootnav/RootNavScreen.kt

@@ -17,13 +17,13 @@ import androidx.navigation.compose.rememberNavController
 import androidx.navigation.navOptions
 import com.x8bit.bitwarden.ui.auth.feature.accountsetup.SETUP_AUTO_FILL_ROUTE
 import com.x8bit.bitwarden.ui.auth.feature.accountsetup.SETUP_COMPLETE_ROUTE
-import com.x8bit.bitwarden.ui.auth.feature.accountsetup.SETUP_UNLOCK_ROUTE
+import com.x8bit.bitwarden.ui.auth.feature.accountsetup.SETUP_UNLOCK_AS_ROOT_ROUTE
 import com.x8bit.bitwarden.ui.auth.feature.accountsetup.navigateToSetupAutoFillScreen
 import com.x8bit.bitwarden.ui.auth.feature.accountsetup.navigateToSetupCompleteScreen
-import com.x8bit.bitwarden.ui.auth.feature.accountsetup.navigateToSetupUnlockScreen
+import com.x8bit.bitwarden.ui.auth.feature.accountsetup.navigateToSetupUnlockScreenAsRoot
 import com.x8bit.bitwarden.ui.auth.feature.accountsetup.setupAutoFillDestination
 import com.x8bit.bitwarden.ui.auth.feature.accountsetup.setupCompleteDestination
-import com.x8bit.bitwarden.ui.auth.feature.accountsetup.setupUnlockDestination
+import com.x8bit.bitwarden.ui.auth.feature.accountsetup.setupUnlockDestinationAsRoot
 import com.x8bit.bitwarden.ui.auth.feature.auth.AUTH_GRAPH_ROUTE
 import com.x8bit.bitwarden.ui.auth.feature.auth.authGraph
 import com.x8bit.bitwarden.ui.auth.feature.auth.navigateToAuthGraph
@@ -99,7 +99,7 @@ fun RootNavScreen(
         vaultUnlockDestination()
         vaultUnlockedGraph(navController)
         setupDebugMenuDestination(onNavigateBack = { navController.popBackStack() })
-        setupUnlockDestination()
+        setupUnlockDestinationAsRoot()
         setupAutoFillDestination()
         setupCompleteDestination()
     }
@@ -127,7 +127,7 @@ fun RootNavScreen(
         is RootNavState.VaultUnlockedForFido2GetCredentials,
         -> VAULT_UNLOCKED_GRAPH_ROUTE
 
-        RootNavState.OnboardingAccountLockSetup -> SETUP_UNLOCK_ROUTE
+        RootNavState.OnboardingAccountLockSetup -> SETUP_UNLOCK_AS_ROOT_ROUTE
         RootNavState.OnboardingAutoFillSetup -> SETUP_AUTO_FILL_ROUTE
         RootNavState.OnboardingStepsComplete -> SETUP_COMPLETE_ROUTE
     }
@@ -235,7 +235,7 @@ fun RootNavScreen(
             }
 
             RootNavState.OnboardingAccountLockSetup -> {
-                navController.navigateToSetupUnlockScreen(rootNavOptions)
+                navController.navigateToSetupUnlockScreenAsRoot(rootNavOptions)
             }
 
             RootNavState.OnboardingAutoFillSetup -> {

app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/settings/SettingsNavigation.kt

@@ -4,6 +4,7 @@ import androidx.navigation.NavController
 import androidx.navigation.NavGraphBuilder
 import androidx.navigation.NavOptions
 import androidx.navigation.navigation
+import com.x8bit.bitwarden.ui.auth.feature.accountsetup.navigateToSetupUnlockScreen
 import com.x8bit.bitwarden.ui.platform.base.util.composableWithRootPushTransitions
 import com.x8bit.bitwarden.ui.platform.feature.settings.about.aboutDestination
 import com.x8bit.bitwarden.ui.platform.feature.settings.about.navigateToAbout
@@ -26,6 +27,7 @@ private const val SETTINGS_ROUTE: String = "settings"
 /**
  * Add settings destinations to the nav graph.
  */
+@Suppress("LongParameterList")
 fun NavGraphBuilder.settingsGraph(
     navController: NavController,
     onNavigateToDeleteAccount: () -> Unit,
@@ -54,6 +56,7 @@ fun NavGraphBuilder.settingsGraph(
             onNavigateBack = { navController.popBackStack() },
             onNavigateToDeleteAccount = onNavigateToDeleteAccount,
             onNavigateToPendingRequests = onNavigateToPendingRequests,
+            onNavigateToSetupUnlockScreen = { navController.navigateToSetupUnlockScreen() },
         )
         appearanceDestination(onNavigateBack = { navController.popBackStack() })
         autoFillDestination(

app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/settings/accountsecurity/AccountSecurityNavigation.kt

@@ -14,6 +14,7 @@ fun NavGraphBuilder.accountSecurityDestination(
     onNavigateBack: () -> Unit,
     onNavigateToDeleteAccount: () -> Unit,
     onNavigateToPendingRequests: () -> Unit,
+    onNavigateToSetupUnlockScreen: () -> Unit,
 ) {
     composableWithPushTransitions(
         route = ACCOUNT_SECURITY_ROUTE,
@@ -22,6 +23,7 @@ fun NavGraphBuilder.accountSecurityDestination(
             onNavigateBack = onNavigateBack,
             onNavigateToDeleteAccount = onNavigateToDeleteAccount,
             onNavigateToPendingRequests = onNavigateToPendingRequests,
+            onNavigateToSetupUnlockScreen = onNavigateToSetupUnlockScreen,
         )
     }
 }

app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/settings/accountsecurity/AccountSecurityScreen.kt

@@ -82,6 +82,7 @@ fun AccountSecurityScreen(
     onNavigateBack: () -> Unit,
     onNavigateToDeleteAccount: () -> Unit,
     onNavigateToPendingRequests: () -> Unit,
+    onNavigateToSetupUnlockScreen: () -> Unit,
     viewModel: AccountSecurityViewModel = hiltViewModel(),
     biometricsManager: BiometricsManager = LocalBiometricsManager.current,
     intentManager: IntentManager = LocalIntentManager.current,
@@ -140,6 +141,8 @@ fun AccountSecurityScreen(
             is AccountSecurityEvent.ShowToast -> {
                 Toast.makeText(context, event.text(resources), Toast.LENGTH_SHORT).show()
             }
+
+            AccountSecurityEvent.NavigateToSetupUnlockScreen -> onNavigateToSetupUnlockScreen()
         }
     }
 

app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/settings/accountsecurity/AccountSecurityViewModel.kt

@@ -165,7 +165,7 @@ class AccountSecurityViewModel @Inject constructor(
 
     private fun handleUnlockCardCtaClick() {
         dismissUnlockNotificationBadge()
-        // TODO: Navigate to unlock set up screen PM-13067
+        sendEvent(AccountSecurityEvent.NavigateToSetupUnlockScreen)
     }
 
     private fun handleAccountFingerprintPhraseClick() {
@@ -564,6 +564,11 @@ sealed class AccountSecurityEvent {
     data class ShowToast(
         val text: Text,
     ) : AccountSecurityEvent()
+
+    /**
+     * Navigate to the setup unlock screen.
+     */
+    data object NavigateToSetupUnlockScreen : AccountSecurityEvent()
 }
 
 /**

app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/vaultunlockednavbar/VaultUnlockedNavBarScreen.kt

@@ -34,6 +34,7 @@ import androidx.navigation.compose.NavHost
 import androidx.navigation.compose.currentBackStackEntryAsState
 import androidx.navigation.compose.rememberNavController
 import androidx.navigation.navOptions
+import com.x8bit.bitwarden.ui.auth.feature.accountsetup.setupUnlockDestination
 import com.x8bit.bitwarden.ui.platform.base.util.EventsEffect
 import com.x8bit.bitwarden.ui.platform.base.util.max
 import com.x8bit.bitwarden.ui.platform.base.util.toDp
@@ -235,6 +236,11 @@ private fun VaultUnlockedNavBarScaffold(
                 onNavigateToFolders = navigateToFolders,
                 onNavigateToPendingRequests = navigateToPendingRequests,
             )
+            setupUnlockDestination(
+                onNavigateBack = {
+                    navController.popBackStack()
+                },
+            )
         }
     }
 }

app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/accountsetup/SetupUnlockScreenTest.kt

@@ -8,6 +8,7 @@ import androidx.compose.ui.test.filterToOne
 import androidx.compose.ui.test.hasAnyAncestor
 import androidx.compose.ui.test.isDialog
 import androidx.compose.ui.test.onAllNodesWithText
+import androidx.compose.ui.test.onNodeWithContentDescription
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
 import androidx.compose.ui.test.performScrollTo
@@ -24,6 +25,7 @@ import io.mockk.mockk
 import io.mockk.runs
 import io.mockk.slot
 import io.mockk.verify
+import junit.framework.TestCase.assertTrue
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.update
 import org.junit.Before
@@ -32,7 +34,7 @@ import org.robolectric.annotation.Config
 import javax.crypto.Cipher
 
 class SetupUnlockScreenTest : BaseComposeTest() {
-
+    private var onNavigateBackCalled = false
     private val captureBiometricsSuccess = slot<(cipher: Cipher?) -> Unit>()
     private val captureBiometricsCancel = slot<() -> Unit>()
     private val captureBiometricsLockOut = slot<() -> Unit>()
@@ -64,6 +66,7 @@ class SetupUnlockScreenTest : BaseComposeTest() {
             SetupUnlockScreen(
                 viewModel = viewModel,
                 biometricsManager = biometricsManager,
+                onNavigateBack = { onNavigateBackCalled = true },
             )
         }
     }
@@ -509,6 +512,15 @@ class SetupUnlockScreenTest : BaseComposeTest() {
         }
     }
 
+    @Test
+    fun `on Set up later component should not be displayed when not in initial setup`() {
+        mutableStateFlow.update { it.copy(isInitialSetup = false) }
+        composeTestRule.assertNoDialogExists()
+        composeTestRule
+            .onNodeWithText(text = "Set up later")
+            .assertDoesNotExist()
+    }
+
     @Test
     fun `on Set up later click should display confirmation dialog`() {
         composeTestRule.assertNoDialogExists()
@@ -610,6 +622,30 @@ class SetupUnlockScreenTest : BaseComposeTest() {
         mutableStateFlow.update { it.copy(dialogState = null) }
         composeTestRule.assertNoDialogExists()
     }
+
+    @Test
+    fun `on NavigateBack event should invoke onNavigateBack`() {
+        mutableEventFlow.tryEmit(SetupUnlockEvent.NavigateBack)
+        assertTrue(onNavigateBackCalled)
+    }
+
+    @Test
+    fun `close icon should not show when in initial setup`() {
+        composeTestRule
+            .onNodeWithContentDescription(label = "Close")
+            .assertDoesNotExist()
+    }
+
+    @Test
+    fun `close icon should show when not initial setup and send action when clicked`() {
+        mutableStateFlow.update { it.copy(isInitialSetup = false) }
+        composeTestRule
+            .onNodeWithContentDescription(label = "Close")
+            .assertIsDisplayed()
+            .performClick()
+
+        verify { viewModel.trySendAction(SetupUnlockAction.CloseClick) }
+    }
 }
 
 private const val DEFAULT_USER_ID: String = "user_id"
@@ -619,6 +655,7 @@ private val DEFAULT_STATE: SetupUnlockState = SetupUnlockState(
     isUnlockWithPasswordEnabled = true,
     isUnlockWithBiometricsEnabled = false,
     dialogState = null,
+    isInitialSetup = true,
 )
 
 private val CIPHER = mockk<Cipher>()

app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/accountsetup/SetupUnlockViewModelTest.kt

@@ -56,10 +56,18 @@ class SetupUnlockViewModelTest : BaseViewModelTest() {
         assertEquals(DEFAULT_STATE, viewModel.stateFlow.value)
     }
 
+    @Test
+    fun `initial state should be correct when not initial setup`() {
+        val viewModel = createViewModel(DEFAULT_STATE.copy(isInitialSetup = false))
+        assertEquals(
+            DEFAULT_STATE.copy(isInitialSetup = false),
+            viewModel.stateFlow.value,
+        )
+    }
+
     @Suppress("MaxLineLength")
     @Test
-    fun `ContinueClick should call setOnboardingStatus and set to AUTOFILL_SETUP if AutoFill is not enabled`() =
+    fun `ContinueClick should call setOnboardingStatus and set to AUTOFILL_SETUP if AutoFill is not enabled`() {
-        runTest {
         val viewModel = createViewModel()
         viewModel.trySendAction(SetupUnlockAction.ContinueClick)
         verify {
@@ -72,8 +80,25 @@ class SetupUnlockViewModelTest : BaseViewModelTest() {
 
     @Suppress("MaxLineLength")
     @Test
-    fun `SetUpLaterClick should call setOnboardingStatus and set to AUTOFILL_SETUP if AutoFill is not enabled`() =
+    fun `ContinueClick should send NavigateBack event if this is not the initial setup`() =
         runTest {
+            val viewModel = createViewModel(DEFAULT_STATE.copy(isInitialSetup = false))
+            viewModel.eventFlow.test {
+                viewModel.trySendAction(SetupUnlockAction.ContinueClick)
+                assertEquals(SetupUnlockEvent.NavigateBack, awaitItem())
+            }
+
+            verify(exactly = 0) {
+                authRepository.setOnboardingStatus(
+                    userId = DEFAULT_USER_ID,
+                    status = OnboardingStatus.AUTOFILL_SETUP,
+                )
+            }
+        }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `SetUpLaterClick should call setOnboardingStatus and set to AUTOFILL_SETUP if AutoFill is not enabled`() {
         val viewModel = createViewModel()
         viewModel.trySendAction(SetupUnlockAction.SetUpLaterClick)
         verify {
@@ -87,18 +112,17 @@ class SetupUnlockViewModelTest : BaseViewModelTest() {
 
     @Suppress("MaxLineLength")
     @Test
-    fun `ContinueClick should call setOnboardingStatus and set to FINAL_STEP if AutoFill is already enabled`() =
+    fun `ContinueClick should call setOnboardingStatus and set to FINAL_STEP if AutoFill is already enabled`() {
-        runTest {
+        mutableAutofillEnabledStateFlow.update { true }
-            mutableAutofillEnabledStateFlow.update { true }
+        val viewModel = createViewModel()
-            val viewModel = createViewModel()
+        viewModel.trySendAction(SetupUnlockAction.ContinueClick)
-            viewModel.trySendAction(SetupUnlockAction.ContinueClick)
+        verify {
-            verify {
+            authRepository.setOnboardingStatus(
-                authRepository.setOnboardingStatus(
+                userId = DEFAULT_USER_ID,
-                    userId = DEFAULT_USER_ID,
+                status = OnboardingStatus.FINAL_STEP,
-                    status = OnboardingStatus.FINAL_STEP,
+            )
-                )
-            }
         }
+    }
 
     @Suppress("MaxLineLength")
     @Test
@@ -116,24 +140,23 @@ class SetupUnlockViewModelTest : BaseViewModelTest() {
         }
 
     @Test
-    fun `on UnlockWithBiometricToggle false should call clearBiometricsKey and update the state`() =
+    fun `on UnlockWithBiometricToggle false should call clearBiometricsKey and update the state`() {
-        runTest {
+        val initialState = DEFAULT_STATE.copy(isUnlockWithBiometricsEnabled = true)
-            val initialState = DEFAULT_STATE.copy(isUnlockWithBiometricsEnabled = true)
+        every { settingsRepository.isUnlockWithBiometricsEnabled } returns true
-            every { settingsRepository.isUnlockWithBiometricsEnabled } returns true
+        every { settingsRepository.clearBiometricsKey() } just runs
-            every { settingsRepository.clearBiometricsKey() } just runs
+        val viewModel = createViewModel(initialState)
-            val viewModel = createViewModel(initialState)
+        assertEquals(initialState, viewModel.stateFlow.value)
-            assertEquals(initialState, viewModel.stateFlow.value)
 
-            viewModel.trySendAction(SetupUnlockAction.UnlockWithBiometricToggle(isEnabled = false))
+        viewModel.trySendAction(SetupUnlockAction.UnlockWithBiometricToggle(isEnabled = false))
 
-            assertEquals(
+        assertEquals(
-                initialState.copy(isUnlockWithBiometricsEnabled = false),
+            initialState.copy(isUnlockWithBiometricsEnabled = false),
-                viewModel.stateFlow.value,
+            viewModel.stateFlow.value,
-            )
+        )
-            verify(exactly = 1) {
+        verify(exactly = 1) {
-                settingsRepository.clearBiometricsKey()
+            settingsRepository.clearBiometricsKey()
-            }
         }
+    }
 
     @Suppress("MaxLineLength")
     @Test
@@ -310,11 +333,28 @@ class SetupUnlockViewModelTest : BaseViewModelTest() {
         )
     }
 
+    @Test
+    fun `CloseClick action should send NavigateBack event`() = runTest {
+        val viewModel = createViewModel()
+        viewModel.eventFlow.test {
+            viewModel.trySendAction(SetupUnlockAction.CloseClick)
+            assertEquals(
+                SetupUnlockEvent.NavigateBack,
+                awaitItem(),
+            )
+        }
+    }
+
     private fun createViewModel(
         state: SetupUnlockState? = null,
     ): SetupUnlockViewModel =
         SetupUnlockViewModel(
-            savedStateHandle = SavedStateHandle(mapOf("state" to state)),
+            savedStateHandle = SavedStateHandle(
+                mapOf(
+                    "state" to state,
+                    "isInitialSetup" to true,
+                ),
+            ),
             authRepository = authRepository,
             settingsRepository = settingsRepository,
             biometricsEncryptionManager = biometricsEncryptionManager,
@@ -328,29 +368,32 @@ private val DEFAULT_STATE: SetupUnlockState = SetupUnlockState(
     isUnlockWithPasswordEnabled = true,
     isUnlockWithBiometricsEnabled = false,
     dialogState = null,
+    isInitialSetup = true,
+)
+
+private val DEFAULT_USER_ACCOUNT = UserState.Account(
+    userId = DEFAULT_USER_ID,
+    name = "Active User",
+    email = "active@bitwarden.com",
+    avatarColorHex = "#aa00aa",
+    environment = Environment.Us,
+    isPremium = true,
+    isLoggedIn = true,
+    isVaultUnlocked = true,
+    needsPasswordReset = false,
+    isBiometricsEnabled = false,
+    organizations = emptyList(),
+    needsMasterPassword = false,
+    trustedDevice = null,
+    hasMasterPassword = true,
+    isUsingKeyConnector = false,
+    onboardingStatus = OnboardingStatus.ACCOUNT_LOCK_SETUP,
 )
 
 private val CIPHER = mockk<Cipher>()
 private val DEFAULT_USER_STATE: UserState = UserState(
     activeUserId = DEFAULT_USER_ID,
     accounts = listOf(
-        UserState.Account(
+        DEFAULT_USER_ACCOUNT,
-            userId = DEFAULT_USER_ID,
-            name = "Active User",
-            email = "active@bitwarden.com",
-            avatarColorHex = "#aa00aa",
-            environment = Environment.Us,
-            isPremium = true,
-            isLoggedIn = true,
-            isVaultUnlocked = true,
-            needsPasswordReset = false,
-            isBiometricsEnabled = false,
-            organizations = emptyList(),
-            needsMasterPassword = false,
-            trustedDevice = null,
-            hasMasterPassword = true,
-            isUsingKeyConnector = false,
-            onboardingStatus = OnboardingStatus.COMPLETE,
-        ),
     ),
 )

app/src/test/java/com/x8bit/bitwarden/ui/platform/feature/rootnav/RootNavScreenTest.kt

@@ -231,7 +231,7 @@ class RootNavScreenTest : BaseComposeTest() {
             RootNavState.OnboardingAccountLockSetup
         composeTestRule.runOnIdle {
             fakeNavHostController.assertLastNavigation(
-                route = "setup_unlock",
+                route = "setup_unlock_as_root/true",
                 navOptions = expectedNavOptions,
             )
         }

app/src/test/java/com/x8bit/bitwarden/ui/platform/feature/settings/accountsecurity/AccountSecurityScreenTest.kt

@@ -48,6 +48,7 @@ class AccountSecurityScreenTest : BaseComposeTest() {
     private var onNavigateBackCalled = false
     private var onNavigateToDeleteAccountCalled = false
     private var onNavigateToPendingRequestsCalled = false
+    private var onNavigateToUnlockSetupScreenCalled = false
 
     private val intentManager = mockk<IntentManager> {
         every { launchUri(any()) } just runs
@@ -85,6 +86,7 @@ class AccountSecurityScreenTest : BaseComposeTest() {
                 onNavigateBack = { onNavigateBackCalled = true },
                 onNavigateToDeleteAccount = { onNavigateToDeleteAccountCalled = true },
                 onNavigateToPendingRequests = { onNavigateToPendingRequestsCalled = true },
+                onNavigateToSetupUnlockScreen = { onNavigateToUnlockSetupScreenCalled = true },
                 viewModel = viewModel,
                 biometricsManager = biometricsManager,
                 intentManager = intentManager,
@@ -1524,6 +1526,12 @@ class AccountSecurityScreenTest : BaseComposeTest() {
             .performClick()
         verify { viewModel.trySendAction(AccountSecurityAction.UnlockActionCardDismiss) }
     }
+
+    @Test
+    fun `on NavigateToSetupUnlockScreen event invokes the correct lambda`() {
+        mutableEventFlow.tryEmit(AccountSecurityEvent.NavigateToSetupUnlockScreen)
+        assertTrue(onNavigateToUnlockSetupScreenCalled)
+    }
 }
 
 private val CIPHER = mockk<Cipher>()

app/src/test/java/com/x8bit/bitwarden/ui/platform/feature/settings/accountsecurity/AccountSecurityViewModelTest.kt

@@ -727,11 +727,19 @@ class AccountSecurityViewModelTest : BaseViewModelTest() {
         }
     }
 
+    @Suppress("MaxLineLength")
     @Test
-    fun `when UnlockActionCardCtaClick action received, should dismiss unlock action card`() {
+    fun `when UnlockActionCardCtaClick action received, should dismiss unlock action card and send NavigateToSetupUnlockScreen event`() =
+        runTest {
         mutableShowUnlockBadgeFlow.update { true }
         val viewModel = createViewModel()
-        viewModel.trySendAction(AccountSecurityAction.UnlockActionCardCtaClick)
+            viewModel.eventFlow.test {
+                viewModel.trySendAction(AccountSecurityAction.UnlockActionCardCtaClick)
+                assertEquals(
+                    AccountSecurityEvent.NavigateToSetupUnlockScreen,
+                    awaitItem(),
+                )
+            }
         verify {
             settingsRepository.storeShowUnlockSettingBadge(DEFAULT_STATE.userId, false)
         }