Delay color picker back navigation until apply is complete

Suspend the apply function until apply is complete and color
configuration is updated to delay proceeding actions, including back
navigation. Prevent double apply in the process by checking if an apply
is already in progress.

There is currently no loading indicator, and corner cases
where config change update does not happen is not currently handled.

Flag: com.android.systemui.shared.new_customization_picker_ui
Test: manually verified & unit tests
Bug: 350718581
Change-Id: Ie6ea517dfb634b6a801f077e5ff655da4e05f663
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
index e7efebd..4029dbe 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
@@ -33,6 +33,8 @@
 import dagger.assisted.AssistedInject
 import dagger.hilt.android.qualifiers.ApplicationContext
 import dagger.hilt.android.scopes.ViewModelScoped
+import kotlin.coroutines.resume
+import kotlinx.coroutines.CancellableContinuation
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -42,6 +44,7 @@
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.suspendCancellableCoroutine
 
 /** Models UI state for a color picker experience. */
 class ColorPickerViewModel2
@@ -57,6 +60,7 @@
     val previewingColorOption = overridingColorOption.asStateFlow()
 
     private val selectedColorTypeTabId = MutableStateFlow<ColorType?>(null)
+    private var onApplyContinuation: CancellableContinuation<Unit>? = null
 
     /** View-models for each color tab. */
     val colorTypeTabs: Flow<List<FloatingToolbarTabViewModel>> =
@@ -165,6 +169,10 @@
                 .toMap()
         }
 
+    /**
+     * This function suspends until onApplyComplete is called to accommodate for configuration
+     * change updates, which are applied with a latency.
+     */
     val onApply: Flow<(suspend () -> Unit)?> =
         previewingColorOption.map { previewingColorOption ->
             previewingColorOption?.let {
@@ -173,6 +181,12 @@
                 } else {
                     {
                         interactor.select(it)
+                        // Suspend until onApplyComplete is called, e.g. on configuration change
+                        suspendCancellableCoroutine { continuation: CancellableContinuation<Unit> ->
+                            onApplyContinuation?.cancel()
+                            onApplyContinuation = continuation
+                            continuation.invokeOnCancellation { onApplyContinuation = null }
+                        }
                         logger.logThemeColorApplied(
                             previewingColorOption.colorOption.sourceForLogging,
                             previewingColorOption.colorOption.styleForLogging,
@@ -187,6 +201,12 @@
         overridingColorOption.value = null
     }
 
+    /** Resumes the onApply function if apply is in progress, otherwise no-op */
+    fun onApplyComplete() {
+        onApplyContinuation?.resume(Unit)
+        onApplyContinuation = null
+    }
+
     /** The list of all available color options for the selected Color Type. */
     val colorOptions: Flow<List<OptionItemViewModel2<ColorOptionIconViewModel>>> =
         combine(allColorOptions, selectedColorTypeTabId) {
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
index 62a96fe..6bc6180 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
@@ -27,6 +27,7 @@
 import dagger.hilt.android.scopes.ViewModelScoped
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.combine
@@ -59,6 +60,8 @@
     val shapeGridPickerViewModel =
         shapeGridPickerViewModelFactory.create(viewModelScope = viewModelScope)
 
+    private var onApplyJob: Job? = null
+
     override val selectedOption = defaultCustomizationOptionsViewModel.selectedOption
 
     override fun handleBackPressed(): Boolean {
@@ -167,9 +170,14 @@
             .map { onApply ->
                 if (onApply != null) {
                     fun(onComplete: () -> Unit) {
-                        viewModelScope.launch {
-                            onApply()
-                            onComplete()
+                        // Prevent double apply
+                        if (onApplyJob?.isActive != true) {
+                            onApplyJob =
+                                viewModelScope.launch {
+                                    onApply()
+                                    onComplete()
+                                    onApplyJob = null
+                                }
                         }
                     }
                 } else {
diff --git a/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt b/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
index 6952d91..4b79739 100644
--- a/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
+++ b/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
@@ -159,13 +159,18 @@
                         }
 
                         launch {
+                            colorUpdateViewModel.systemColorsUpdated.collect {
+                                viewModel.colorPickerViewModel2.onApplyComplete()
+                            }
+                        }
+
+                        launch {
                             combine(
                                     viewModel.colorPickerViewModel2.previewingColorOption,
                                     viewModel.darkModeViewModel.overridingIsDarkMode,
-                                    colorUpdateViewModel.systemColorsUpdated,
-                                    ::Triple,
+                                    ::Pair,
                                 )
-                                .collect { (colorModel, darkMode, _) ->
+                                .collect { (colorModel, darkMode) ->
                                     val bundle =
                                         Bundle().apply {
                                             if (colorModel != null) {
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
index 421d8a0..e4a4884 100644
--- a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
@@ -36,6 +36,7 @@
 import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
@@ -96,6 +97,27 @@
     }
 
     @Test
+    fun onApply_suspendsUntilOnApplyCompleteIsCalled() =
+        testScope.runTest {
+            val colorTypes = collectLastValue(underTest.colorTypeTabs)
+            val colorOptions = collectLastValue(underTest.colorOptions)
+            val onApply = collectLastValue(underTest.onApply)
+
+            // Select "Wallpaper colors" tab
+            colorTypes()?.get(0)?.onClick?.invoke()
+            // Select a color option to preview
+            selectColorOption(colorOptions, 1)
+            // Apply the selected color option
+            val job = testScope.launch { onApply()?.invoke() }
+
+            assertThat(job.isActive).isTrue()
+
+            underTest.onApplyComplete()
+
+            assertThat(job.isActive).isFalse()
+        }
+
+    @Test
     fun onApply_wallpaperColor_shouldLogColor() =
         testScope.runTest {
             repository.setOptions(
@@ -203,7 +225,7 @@
             )
         }
 
-    /** Simulates a user selecting the affordance at the given index, if that is clickable. */
+    /** Simulates a user selecting the color option at the given index. */
     private fun TestScope.selectColorOption(
         colorOptions: () -> List<OptionItemViewModel2<ColorOptionIconViewModel>>?,
         index: Int,
@@ -217,10 +239,11 @@
         }
     }
 
-    /** Simulates a user selecting the affordance at the given index, if that is clickable. */
+    /** Simulates a user applying the color option at the given index, and the apply completes. */
     private suspend fun TestScope.applySelectedColorOption() {
         val onApply = collectLastValue(underTest.onApply)()
-        onApply?.invoke()
+        testScope.launch { onApply?.invoke() }
+        underTest.onApplyComplete()
     }
 
     /**