Make ColorUpdateViewModel ActivityRetainedScoped (2/2)

Changed ColorUpdateViewModel from activity scoped to activity retained
scoped so it can be easily injected and used in other view model scoped
classes. Injected the class into ColorPickerViewModel2 so the view model
can directly listen to system color updates and determine when apply is
complete.

Flag: com.android.systemui.shared.new_customization_picker_ui
Test: manually verified that applying color still works
Bug: 299510645
Change-Id: I13a2bd06fbe27b5e052311e82abb7d2f40856795
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
index 4029dbe..9e2353a 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
@@ -26,6 +26,7 @@
 import com.android.themepicker.R
 import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
 import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
+import com.android.wallpaper.picker.customization.ui.viewmodel.ColorUpdateViewModel
 import com.android.wallpaper.picker.customization.ui.viewmodel.FloatingToolbarTabViewModel
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel2
 import dagger.assisted.Assisted
@@ -33,8 +34,6 @@
 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
@@ -43,14 +42,15 @@
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.flow.take
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.suspendCancellableCoroutine
 
 /** Models UI state for a color picker experience. */
 class ColorPickerViewModel2
 @AssistedInject
 constructor(
     @ApplicationContext context: Context,
+    private val colorUpdateViewModel: ColorUpdateViewModel,
     private val interactor: ColorPickerInteractor,
     private val logger: ThemesUserEventLogger,
     @Assisted private val viewModelScope: CoroutineScope,
@@ -60,7 +60,6 @@
     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>> =
@@ -181,11 +180,9 @@
                 } 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 }
+                        // Suspend until first color update
+                        colorUpdateViewModel.systemColorsUpdatedNoReplay.take(1).collect {
+                            return@collect
                         }
                         logger.logThemeColorApplied(
                             previewingColorOption.colorOption.sourceForLogging,
@@ -201,12 +198,6 @@
         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/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt b/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
index 4b79739..cb6057d 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,12 +159,6 @@
                         }
 
                         launch {
-                            colorUpdateViewModel.systemColorsUpdated.collect {
-                                viewModel.colorPickerViewModel2.onApplyComplete()
-                            }
-                        }
-
-                        launch {
                             combine(
                                     viewModel.colorPickerViewModel2.previewingColorOption,
                                     viewModel.darkModeViewModel.overridingIsDarkMode,
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 e4a4884..9827a56 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
@@ -28,12 +28,14 @@
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
 import com.android.systemui.monet.Style
+import com.android.wallpaper.picker.customization.ui.viewmodel.ColorUpdateViewModel
 import com.android.wallpaper.picker.customization.ui.viewmodel.FloatingToolbarTabViewModel
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel2
 import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
+import dagger.hilt.android.internal.lifecycle.RetainedLifecycleImpl
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.launch
@@ -58,6 +60,7 @@
     private lateinit var repository: FakeColorPickerRepository
     private lateinit var interactor: ColorPickerInteractor
     private lateinit var store: FakeSnapshotStore
+    private lateinit var colorUpdateViewModel: ColorUpdateViewModel
 
     private lateinit var context: Context
     private lateinit var testScope: TestScope
@@ -80,11 +83,14 @@
                     },
             )
 
+        colorUpdateViewModel = ColorUpdateViewModel(context, RetainedLifecycleImpl())
+
         underTest =
             ColorPickerViewModel2(
                 context = context,
                 interactor = interactor,
                 logger = logger,
+                colorUpdateViewModel = colorUpdateViewModel,
                 viewModelScope = testScope.backgroundScope,
             )
 
@@ -112,7 +118,7 @@
 
             assertThat(job.isActive).isTrue()
 
-            underTest.onApplyComplete()
+            colorUpdateViewModel.updateColors()
 
             assertThat(job.isActive).isFalse()
         }
@@ -243,7 +249,7 @@
     private suspend fun TestScope.applySelectedColorOption() {
         val onApply = collectLastValue(underTest.onApply)()
         testScope.launch { onApply?.invoke() }
-        underTest.onApplyComplete()
+        colorUpdateViewModel.updateColors()
     }
 
     /**