BIT-267 make sure root navigation state doesn't get blown away after … (#32)

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

@@ -34,5 +34,5 @@ fun NavGraphBuilder.authDestinations(navController: NavHostController) {
 fun NavController.navigateToAuth(
     navOptions: NavOptions? = null,
 ) {
-    navigate(LANDING_ROUTE, navOptions)
+    navigate(AUTH_ROUTE, navOptions)
 }

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

@@ -5,15 +5,18 @@ import androidx.compose.runtime.getValue
 import androidx.hilt.navigation.compose.hiltViewModel
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import androidx.navigation.NavController
+import androidx.navigation.NavDestination
 import androidx.navigation.NavGraphBuilder
 import androidx.navigation.NavOptions
 import androidx.navigation.compose.NavHost
 import androidx.navigation.compose.composable
 import androidx.navigation.compose.rememberNavController
 import androidx.navigation.navOptions
+import com.x8bit.bitwarden.ui.auth.feature.auth.AUTH_ROUTE
 import com.x8bit.bitwarden.ui.auth.feature.auth.authDestinations
 import com.x8bit.bitwarden.ui.auth.feature.auth.navigateToAuth
 import com.x8bit.bitwarden.ui.platform.components.PlaceholderComposable
+import com.x8bit.bitwarden.ui.platform.feature.vaultunlocked.VAULT_UNLOCKED_ROUTE
 import com.x8bit.bitwarden.ui.platform.feature.vaultunlocked.navigateToVaultUnlocked
 import com.x8bit.bitwarden.ui.platform.feature.vaultunlocked.vaultUnlockedDestinations
 
@@ -29,7 +32,7 @@ fun RootNavScreen(
 
     NavHost(
         navController = navController,
-        startDestination = SplashRoute,
+        startDestination = SPLASH_ROUTE,
     ) {
         splashDestinations()
         authDestinations(navController)
@@ -40,9 +43,32 @@ fun RootNavScreen(
     val rootNavOptions = navOptions {
         // When changing root navigation state, pop everything else off the back stack:
         popUpTo(navController.graph.id) {
-            inclusive = true
+            inclusive = false
+            saveState = true
         }
+        launchSingleTop = true
+        restoreState = true
     }
+
+    // This workaround is for an issue where "launchSingleTop" flag does not work correctly
+    // and we "re-navigate" after rotation or process death. To prevent this, we are currently
+    // checking the root level route and no-opping the navigation, because we are already there.
+    // When upgrading to the latest compose nav version (which we currently aren't doing for
+    // other reasons), we can test that this workaround is no longer needed and remove it.
+    // To test, remove the (currentRoute == targetRoute) and test that state is saved
+    // on process death and rotation (BIT-201).
+    val targetRoute = when (state) {
+        RootNavState.Auth -> AUTH_ROUTE
+        RootNavState.Splash -> SPLASH_ROUTE
+        RootNavState.VaultUnlocked -> VAULT_UNLOCKED_ROUTE
+    }
+    val currentRoute = navController.currentDestination?.routeLevelRoute()
+
+    // Don't navigate if we are already at the correct root:
+    if (currentRoute == targetRoute) {
+        return
+    }
+
     when (state) {
         RootNavState.Auth -> navController.navigateToAuth(rootNavOptions)
         RootNavState.Splash -> navController.navigateToSplash(rootNavOptions)
@@ -50,6 +76,23 @@ fun RootNavScreen(
     }
 }
 
+/**
+ * Helper method that returns the highest level route for the given [NavDestination].
+ *
+ * As noted above, this can be removed after upgrading to latest compose navigation, since
+ * the nav args can prevent us from having to do this check.
+ */
+@Suppress("ReturnCount")
+private fun NavDestination?.routeLevelRoute(): String? {
+    if (this == null) {
+        return null
+    }
+    if (parent?.route == null) {
+        return route
+    }
+    return parent.routeLevelRoute()
+}
+
 /**
  * The functions below should be moved to their respective feature packages once they exist.
  *
@@ -59,8 +102,7 @@ fun RootNavScreen(
 /**
  * TODO: move to splash package (BIT-147)
  */
-@Suppress("TopLevelPropertyNaming")
-private const val SplashRoute = "splash"
+private const val SPLASH_ROUTE = "splash"
 
 /**
  * Add splash destinations to the nav graph.
@@ -68,7 +110,7 @@ private const val SplashRoute = "splash"
  * TODO: move to splash package (BIT-147)
  */
 private fun NavGraphBuilder.splashDestinations() {
-    composable(SplashRoute) {
+    composable(SPLASH_ROUTE) {
         PlaceholderComposable(text = "Splash")
     }
 }
@@ -83,5 +125,5 @@ private fun NavGraphBuilder.splashDestinations() {
 private fun NavController.navigateToSplash(
     navOptions: NavOptions? = null,
 ) {
-    navigate(SplashRoute, navOptions)
+    navigate(SPLASH_ROUTE, navOptions)
 }

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

@@ -1,26 +1,48 @@
 package com.x8bit.bitwarden.ui.platform.feature.rootnav
 
+import android.os.Parcelable
+import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.viewModelScope
 import com.x8bit.bitwarden.ui.platform.base.BaseViewModel
 import dagger.hilt.android.lifecycle.HiltViewModel
 import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
+import kotlinx.parcelize.Parcelize
 import javax.inject.Inject
 
+private const val KEY_NAV_DESTINATION = "nav_state"
+
 /**
  * Manages root level navigation state of the application.
  */
 @HiltViewModel
-class RootNavViewModel @Inject constructor() :
-    BaseViewModel<RootNavState, Unit, Unit>(
-        initialState = RootNavState.Splash,
-    ) {
+class RootNavViewModel @Inject constructor(
+    private val savedStateHandle: SavedStateHandle,
+) : BaseViewModel<RootNavState, Unit, Unit>(
+    initialState = RootNavState.Splash,
+) {
+
+    private var savedRootNavState: RootNavState?
+        get() = savedStateHandle[KEY_NAV_DESTINATION]
+        set(value) {
+            savedStateHandle[KEY_NAV_DESTINATION] = value
+        }
 
     init {
+        savedRootNavState?.let { savedState: RootNavState ->
+            mutableStateFlow.update { savedState }
+        }
+        // Every time the nav state changes, update saved state handle:
+        stateFlow
+            .onEach { savedRootNavState = it }
+            .launchIn(viewModelScope)
         viewModelScope.launch {
             @Suppress("MagicNumber")
             delay(1000)
-            mutableStateFlow.value = RootNavState.Auth
+            mutableStateFlow.update { RootNavState.Auth }
         }
     }
 
@@ -28,21 +50,24 @@ class RootNavViewModel @Inject constructor() :
 }
 
 /**
- * Models state of the root level navigation of the app.
+ * Models root level destinations for the app.
  */
-sealed class RootNavState {
+sealed class RootNavState : Parcelable {
     /**
-     * Show the vault unlocked screen.
-     */
-    data object VaultUnlocked : RootNavState()
-
-    /**
-     * Show the auth screens.
+     * App should show auth nav graph.
      */
+    @Parcelize
     data object Auth : RootNavState()
 
     /**
-     * Show the splash screen.
+     * App should show splash nav graph.
      */
+    @Parcelize
     data object Splash : RootNavState()
+
+    /**
+     * App should show vault unlocked nav graph.
+     */
+    @Parcelize
+    data object VaultUnlocked : RootNavState()
 }

app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/vaultunlocked/VaultUnlockedNavigation.kt

@@ -7,7 +7,7 @@ import androidx.navigation.navigation
 import com.x8bit.bitwarden.ui.platform.feature.vaultunlockednavbar.VAULT_UNLOCKED_NAV_BAR_ROUTE
 import com.x8bit.bitwarden.ui.platform.feature.vaultunlockednavbar.vaultUnlockedNavBarDestination
 
-private const val VAULT_UNLOCKED_ROUTE = "VaultUnlocked"
+const val VAULT_UNLOCKED_ROUTE: String = "VaultUnlocked"
 
 /**
  * Navigate to the vault unlocked screen. Note this will only work if vault unlocked destinations were added

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

@@ -1,24 +1,33 @@
 package com.x8bit.bitwarden.ui.platform.feature.rootnav
 
+import androidx.lifecycle.SavedStateHandle
 import app.cash.turbine.test
 import com.x8bit.bitwarden.ui.platform.base.BaseViewModelTest
 import kotlinx.coroutines.test.runTest
+import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Test
 
 class RootNavViewModelTest : BaseViewModelTest() {
 
     @Test
     fun `initial state should be splash`() {
-        val viewModel = RootNavViewModel()
-        assert(viewModel.stateFlow.value is RootNavState.Splash)
+        val viewModel = RootNavViewModel(SavedStateHandle())
+        assertEquals(viewModel.stateFlow.value, RootNavState.Splash)
+    }
+
+    @Test
+    fun `initial state should be the state in savedStateHandle`() {
+        val handle = SavedStateHandle(mapOf(("nav_state" to RootNavState.VaultUnlocked)))
+        val viewModel = RootNavViewModel(handle)
+        assertEquals(viewModel.stateFlow.value, RootNavState.VaultUnlocked)
     }
 
     @Test
     fun `state should move from splash to auth`() = runTest {
-        val viewModel = RootNavViewModel()
+        val viewModel = RootNavViewModel(SavedStateHandle())
         viewModel.stateFlow.test {
-            assert(awaitItem() is RootNavState.Splash)
-            assert(awaitItem() is RootNavState.Auth)
+            assertEquals(awaitItem(), RootNavState.Splash)
+            assertEquals(awaitItem(), RootNavState.Auth)
         }
     }
 }