PopWithTransition on dismiss

DialogNavigator should be using popWithTransition when calling dismiss
on a dialog as well. That way we don't dismiss too early when the dialog
is being destroyed.

RelNote: "Fixed a race condition when using a `ViewModel` within `dialog` destination that would cause an `IllegalStateException` when dismissing the dialog by tapping outside the dialog."
Test: existing tests pass
Bug: 226552301

Change-Id: Id7376c0b8db5be869d8ff53185e15b0603bf8582
diff --git a/navigation/navigation-compose/src/androidTest/java/androidx/navigation/compose/DialogNavigatorTest.kt b/navigation/navigation-compose/src/androidTest/java/androidx/navigation/compose/DialogNavigatorTest.kt
index 98b57cb..adab446 100644
--- a/navigation/navigation-compose/src/androidTest/java/androidx/navigation/compose/DialogNavigatorTest.kt
+++ b/navigation/navigation-compose/src/androidTest/java/androidx/navigation/compose/DialogNavigatorTest.kt
@@ -20,9 +20,12 @@
 import androidx.compose.ui.test.assertIsDisplayed
 import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.compose.ui.test.onNodeWithText
+import androidx.lifecycle.viewmodel.compose.viewModel
+import androidx.navigation.NavHostController
 import androidx.navigation.testing.TestNavigatorState
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.MediumTest
+import com.google.common.truth.Truth.assertThat
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -77,4 +80,33 @@
 
         rule.onNodeWithText(defaultText).assertDoesNotExist()
     }
+
+    @Test
+    fun testNestedNavHostInDialogDismissed() {
+        lateinit var navController: NavHostController
+
+        rule.setContent {
+            navController = rememberNavController()
+            NavHost(navController, "first") {
+                composable("first") { }
+                dialog("second") {
+                    viewModel<TestViewModel>(it)
+                }
+            }
+        }
+
+        rule.runOnIdle {
+            navController.navigate("second")
+        }
+
+        // Now trigger the back button
+        rule.runOnIdle {
+            navController.navigatorProvider.getNavigator(DialogNavigator::class.java).dismiss(
+                navController.getBackStackEntry("second")
+            )
+        }
+
+        rule.waitForIdle()
+        assertThat(navController.currentDestination?.route).isEqualTo("first")
+    }
 }
diff --git a/navigation/navigation-compose/src/main/java/androidx/navigation/compose/DialogNavigator.kt b/navigation/navigation-compose/src/main/java/androidx/navigation/compose/DialogNavigator.kt
index bfac256..9de4e8f 100644
--- a/navigation/navigation-compose/src/main/java/androidx/navigation/compose/DialogNavigator.kt
+++ b/navigation/navigation-compose/src/main/java/androidx/navigation/compose/DialogNavigator.kt
@@ -43,7 +43,7 @@
      * Dismiss the dialog destination associated with the given [backStackEntry].
      */
     internal fun dismiss(backStackEntry: NavBackStackEntry) {
-        state.pop(backStackEntry, false)
+        state.popWithTransition(backStackEntry, false)
     }
 
     override fun navigate(