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

2024-11-21 17:05:44 +03:00 · 2024-10-04 14:29:47 -04:00 · 2024-10-04 14:29:47 -04:00 · 8ae6433906
commit 8ae6433906
parent 83652c9699
14 changed files with 299 additions and 76 deletions
--- a/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/accountsetup/SetupUnlockNavigation.kt
+++ b/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() {
-    composableWithPushTransitions(
+fun NavController.navigateToSetupUnlockScreenAsRoot(navOptions: NavOptions? = null) {
+    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
+        },
+    ),
+)
--- a/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/accountsetup/SetupUnlockScreen.kt
+++ b/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,6 +192,7 @@ private fun SetupUnlockScreenContent(
        )

        Spacer(modifier = Modifier.height(height = 12.dp))
+        if (state.isInitialSetup) {
            SetUpLaterButton(
                onConfirmClick = handler.onSetUpLaterClick,
                modifier = Modifier
@ -177,6 +201,7 @@ private fun SetupUnlockScreenContent(
            )

            Spacer(modifier = Modifier.height(height = 12.dp))
+        }
        Spacer(modifier = Modifier.navigationBarsPadding())
    }
 }
--- a/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/accountsetup/SetupUnlockViewModel.kt
+++ b/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() {
+        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.
     */
--- a/app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/rootnav/RootNavScreen.kt
+++ b/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 -> {
--- a/app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/settings/SettingsNavigation.kt
+++ b/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(
--- a/app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/settings/accountsecurity/AccountSecurityNavigation.kt
+++ b/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,
        )
    }
 }
--- a/app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/settings/accountsecurity/AccountSecurityScreen.kt
+++ b/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()
        }
    }

--- a/app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/settings/accountsecurity/AccountSecurityViewModel.kt
+++ b/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()
 }

 /**
--- a/app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/vaultunlockednavbar/VaultUnlockedNavBarScreen.kt
+++ b/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()
+                },
+            )
        }
    }
 }
--- a/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/accountsetup/SetupUnlockScreenTest.kt
+++ b/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>()
--- a/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/accountsetup/SetupUnlockViewModelTest.kt
+++ b/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`() =
-        runTest {
+    fun `ContinueClick should call setOnboardingStatus and set to AUTOFILL_SETUP if AutoFill is not enabled`() {
        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,8 +112,7 @@ class SetupUnlockViewModelTest : BaseViewModelTest() {

    @Suppress("MaxLineLength")
    @Test
-    fun `ContinueClick should call setOnboardingStatus and set to FINAL_STEP if AutoFill is already enabled`() =
-        runTest {
+    fun `ContinueClick should call setOnboardingStatus and set to FINAL_STEP if AutoFill is already enabled`() {
        mutableAutofillEnabledStateFlow.update { true }
        val viewModel = createViewModel()
        viewModel.trySendAction(SetupUnlockAction.ContinueClick)
@ -116,8 +140,7 @@ class SetupUnlockViewModelTest : BaseViewModelTest() {
        }

    @Test
-    fun `on UnlockWithBiometricToggle false should call clearBiometricsKey and update the state`() =
-        runTest {
+    fun `on UnlockWithBiometricToggle false should call clearBiometricsKey and update the state`() {
        val initialState = DEFAULT_STATE.copy(isUnlockWithBiometricsEnabled = true)
        every { settingsRepository.isUnlockWithBiometricsEnabled } returns true
        every { settingsRepository.clearBiometricsKey() } just runs
@ -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,13 +368,10 @@ private val DEFAULT_STATE: SetupUnlockState = SetupUnlockState(
    isUnlockWithPasswordEnabled = true,
    isUnlockWithBiometricsEnabled = false,
    dialogState = null,
+    isInitialSetup = true,
 )

-private val CIPHER = mockk<Cipher>()
-private val DEFAULT_USER_STATE: UserState = UserState(
-    activeUserId = DEFAULT_USER_ID,
-    accounts = listOf(
-        UserState.Account(
+private val DEFAULT_USER_ACCOUNT = UserState.Account(
    userId = DEFAULT_USER_ID,
    name = "Active User",
    email = "active@bitwarden.com",
@ -350,7 +387,13 @@ private val DEFAULT_USER_STATE: UserState = UserState(
    trustedDevice = null,
    hasMasterPassword = true,
    isUsingKeyConnector = false,
-            onboardingStatus = OnboardingStatus.COMPLETE,
-        ),
+    onboardingStatus = OnboardingStatus.ACCOUNT_LOCK_SETUP,
+)
+
+private val CIPHER = mockk<Cipher>()
+private val DEFAULT_USER_STATE: UserState = UserState(
+    activeUserId = DEFAULT_USER_ID,
+    accounts = listOf(
+        DEFAULT_USER_ACCOUNT,
    ),
 )
--- a/app/src/test/java/com/x8bit/bitwarden/ui/platform/feature/rootnav/RootNavScreenTest.kt
+++ b/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,
            )
        }
--- a/app/src/test/java/com/x8bit/bitwarden/ui/platform/feature/settings/accountsecurity/AccountSecurityScreenTest.kt
+++ b/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>()
--- a/app/src/test/java/com/x8bit/bitwarden/ui/platform/feature/settings/accountsecurity/AccountSecurityViewModelTest.kt
+++ b/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.eventFlow.test {
                viewModel.trySendAction(AccountSecurityAction.UnlockActionCardCtaClick)
+                assertEquals(
+                    AccountSecurityEvent.NavigateToSetupUnlockScreen,
+                    awaitItem(),
+                )
+            }
        verify {
            settingsRepository.storeShowUnlockSettingBadge(DEFAULT_STATE.userId, false)
        }