Fix dismiss dialog shows on dark mode apply (2/2)

Flag: com.android.systemui.shared.new_customization_picker_ui
Bug: 388730120
Test: manually verified by applying dark mode and color
Change-Id: Iefa318507fec4a5998c37ae235e44a5a310b6f36
diff --git a/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModel.kt b/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModel.kt
index aacf001..4c79405 100644
--- a/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModel.kt
+++ b/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModel.kt
@@ -18,17 +18,26 @@
 
 import com.android.customization.module.logging.ThemesUserEventLogger
 import com.android.customization.picker.mode.domain.interactor.DarkModeInteractor
+import com.android.wallpaper.picker.customization.ui.viewmodel.ColorUpdateViewModel
 import dagger.hilt.android.scopes.ViewModelScoped
 import javax.inject.Inject
+import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.drop
+import kotlinx.coroutines.flow.take
+import kotlinx.coroutines.launch
 
 @ViewModelScoped
 class DarkModeViewModel
 @Inject
-constructor(private val interactor: DarkModeInteractor, private val logger: ThemesUserEventLogger) {
+constructor(
+    private val colorUpdateViewModel: ColorUpdateViewModel,
+    private val interactor: DarkModeInteractor,
+    private val logger: ThemesUserEventLogger,
+) {
     private val isDarkMode = interactor.isDarkMode
     val isEnabled = interactor.isEnabled
 
@@ -54,8 +63,21 @@
         combine(overridingIsDarkMode, isDarkMode, isEnabled) { override, current, isEnabled ->
             if (override != null && override != current && isEnabled) {
                 {
-                    interactor.setIsDarkMode(override)
-                    logger.logDarkThemeApplied(override)
+                    coroutineScope {
+                        launch { interactor.setIsDarkMode(override) }
+                        // Dark mode change also invokes a color update. Suspend until both dark
+                        // mode and color are updated.
+                        combine(
+                                // Omit the first value which is emitted on subscribe.
+                                isDarkMode.drop(1).take(1),
+                                colorUpdateViewModel.systemColorsUpdatedNoReplay.take(1),
+                                ::Pair,
+                            )
+                            .collect { (_, _) ->
+                                return@collect
+                            }
+                        logger.logDarkThemeApplied(override)
+                    }
                 }
             } else null
         }
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
index 1a68f4e..7b5c9e6 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
@@ -35,6 +35,7 @@
 import dagger.hilt.android.qualifiers.ApplicationContext
 import dagger.hilt.android.scopes.ViewModelScoped
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
@@ -169,16 +170,18 @@
                     null
                 } else {
                     {
-                        interactor.select(it)
-                        // Suspend until first color update
-                        colorUpdateViewModel.systemColorsUpdatedNoReplay.take(1).collect {
-                            return@collect
+                        coroutineScope {
+                            launch { interactor.select(it) }
+                            // Suspend until first color update
+                            colorUpdateViewModel.systemColorsUpdatedNoReplay.take(1).collect {
+                                return@collect
+                            }
+                            logger.logThemeColorApplied(
+                                it.sourceForLogging,
+                                it.styleForLogging,
+                                it.seedColor,
+                            )
                         }
-                        logger.logThemeColorApplied(
-                            it.sourceForLogging,
-                            it.styleForLogging,
-                            it.seedColor,
-                        )
                     }
                 }
             }
diff --git a/tests/robotests/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModelTest.kt
index 3bdbb9d..61de0ce 100644
--- a/tests/robotests/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModelTest.kt
@@ -16,20 +16,27 @@
 
 package com.android.customization.picker.mode.ui.viewmodel
 
+import android.content.Context
+import androidx.test.platform.app.InstrumentationRegistry
 import com.android.customization.module.logging.TestThemesUserEventLogger
 import com.android.customization.picker.mode.data.repository.DarkModeRepository
+import com.android.customization.picker.mode.data.repository.DarkModeStateRepository
 import com.android.customization.picker.mode.domain.interactor.DarkModeInteractor
+import com.android.wallpaper.picker.customization.ui.viewmodel.ColorUpdateViewModel
 import com.android.wallpaper.testing.FakePowerManager
 import com.android.wallpaper.testing.FakeUiModeManager
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
+import dagger.hilt.android.internal.lifecycle.RetainedLifecycleImpl
 import dagger.hilt.android.testing.HiltAndroidRule
 import dagger.hilt.android.testing.HiltAndroidTest
 import javax.inject.Inject
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.TestDispatcher
 import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.test.setMain
 import org.junit.Before
@@ -46,20 +53,26 @@
 
     @Inject lateinit var uiModeManager: FakeUiModeManager
     @Inject lateinit var powerManager: FakePowerManager
+    @Inject lateinit var darkModeStateRepository: DarkModeStateRepository
     @Inject lateinit var darkModeRepository: DarkModeRepository
     @Inject lateinit var darkModeInteractor: DarkModeInteractor
     @Inject lateinit var logger: TestThemesUserEventLogger
-    private lateinit var darkModeViewModel: DarkModeViewModel
-
     @Inject lateinit var testDispatcher: TestDispatcher
     @Inject lateinit var testScope: TestScope
 
+    private lateinit var context: Context
+    private lateinit var colorUpdateViewModel: ColorUpdateViewModel
+    private lateinit var darkModeViewModel: DarkModeViewModel
+
     @Before
     fun setUp() {
         hiltRule.inject()
         Dispatchers.setMain(testDispatcher)
 
-        darkModeViewModel = DarkModeViewModel(darkModeInteractor, logger)
+        context = InstrumentationRegistry.getInstrumentation().targetContext
+        colorUpdateViewModel =
+            ColorUpdateViewModel(context, RetainedLifecycleImpl(), darkModeStateRepository)
+        darkModeViewModel = DarkModeViewModel(colorUpdateViewModel, darkModeInteractor, logger)
     }
 
     @Test
@@ -141,10 +154,9 @@
             uiModeManager.setNightModeActivated(false)
             darkModeRepository.refreshIsDarkMode()
             val getToggleDarkMode = collectLastValue(darkModeViewModel.toggleDarkMode)
-            val onApply = collectLastValue(darkModeViewModel.onApply)
 
             getToggleDarkMode()?.invoke()
-            onApply()?.invoke()
+            applyDarkMode()
 
             assertThat(logger.useDarkTheme).isTrue()
         }
@@ -156,12 +168,23 @@
             uiModeManager.setNightModeActivated(false)
             darkModeRepository.refreshIsDarkMode()
             val getToggleDarkMode = collectLastValue(darkModeViewModel.toggleDarkMode)
-            val onApply = collectLastValue(darkModeViewModel.onApply)
 
             getToggleDarkMode()?.invoke()
-            onApply()?.invoke()
+            applyDarkMode()
 
             assertThat(uiModeManager.getIsNightModeActivated()).isTrue()
         }
     }
+
+    /** Simulates a user applying the previewing dark mode, and the apply completes. */
+    private fun TestScope.applyDarkMode() {
+        val onApply = collectLastValue(darkModeViewModel.onApply)()
+        testScope.launch { onApply?.invoke() }
+        // Run coroutine launched in DarkModeViewModel#onApply
+        runCurrent()
+        // Simulate dark mode and color update config change
+        colorUpdateViewModel.updateDarkModeAndColors()
+        // Run coroutine launched in colorUpdateViewModel#updateColors
+        runCurrent()
+    }
 }