Refactor color picker to better track selected option (1/2)

The color picker workflow used to involve creating a list of
ColorOptions, converting each option to a ColorOptionModel in the
repository, and then converting each option to an OptionItemViewModel in
the view model. The selected state of each color option used to be kept
individually in each ColorOptionModel. Create a new flagged color picker
workflow, and centralize the selected state by keeping a single
selectedColorOption variable in ColorPickerRepository2. Also remove the
usage of ColorOptionModel in the new flow completely since this
abstraction layer is no longer needed. This should simplify the new flow
and reduce bugs, and will simplify upcoming work.

Flag: com.android.systemui.shared.new_customization_picker_ui
Test: manually verified & unit tests
Bug: 350718581
Change-Id: Ibec699fc45e4ebd8b3fa37dc6c4cb23e099ff2f9
diff --git a/src/com/android/customization/model/color/ColorCustomizationManager.java b/src/com/android/customization/model/color/ColorCustomizationManager.java
index 61a7967..99916b5 100644
--- a/src/com/android/customization/model/color/ColorCustomizationManager.java
+++ b/src/com/android/customization/model/color/ColorCustomizationManager.java
@@ -88,6 +88,7 @@
     private String mCurrentStyle;
     private WallpaperColors mHomeWallpaperColors;
     private WallpaperColors mLockWallpaperColors;
+    private SettingsChangedListener mListener;
 
     /** Returns the {@link ColorCustomizationManager} instance. */
     public static ColorCustomizationManager getInstance(Context context,
@@ -116,6 +117,7 @@
         mProvider = provider;
         mContentResolver = contentResolver;
         mExecutorService = executorService;
+        mListener = null;
         ContentObserver observer = new ContentObserver(/* handler= */ null) {
             @Override
             public void onChange(boolean selfChange, Uri uri) {
@@ -127,6 +129,9 @@
                     mCurrentOverlays = null;
                     mCurrentStyle = null;
                     mCurrentSource = null;
+                    if (mListener != null) {
+                        mListener.onSettingsChanged();
+                    }
                 }
             }
         };
@@ -314,4 +319,19 @@
         }
         return overlayPackages;
     }
+
+    /**
+     * Sets a listener that is called when ColorCustomizationManager is updated.
+     */
+    public void setListener(SettingsChangedListener listener) {
+        mListener = listener;
+    }
+
+    /**
+     * A listener for listening to when ColorCustomizationManager is updated.
+     */
+    public interface SettingsChangedListener {
+        /** */
+        void onSettingsChanged();
+    }
 }
diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepository2.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepository2.kt
new file mode 100644
index 0000000..0f8a86e
--- /dev/null
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepository2.kt
@@ -0,0 +1,36 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+package com.android.customization.picker.color.data.repository
+
+import com.android.customization.model.color.ColorOption
+import com.android.customization.picker.color.shared.model.ColorType
+import kotlinx.coroutines.flow.Flow
+
+/**
+ * Abstracts access to application state related to functionality for selecting, picking, or setting
+ * system color.
+ */
+interface ColorPickerRepository2 {
+    /** List of wallpaper and preset color options on the device, categorized by Color Type */
+    val colorOptions: Flow<Map<ColorType, List<ColorOption>>>
+
+    /** The system selected color option from the generated list of color options */
+    val selectedColorOption: Flow<ColorOption?>
+
+    /** Selects a color option with optimistic update */
+    suspend fun select(colorOption: ColorOption)
+}
diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl2.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl2.kt
new file mode 100644
index 0000000..5e90b41
--- /dev/null
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl2.kt
@@ -0,0 +1,160 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+package com.android.customization.picker.color.data.repository
+
+import android.util.Log
+import com.android.customization.model.CustomizationManager
+import com.android.customization.model.color.ColorCustomizationManager
+import com.android.customization.model.color.ColorOption
+import com.android.customization.model.color.ColorOptionImpl
+import com.android.customization.picker.color.shared.model.ColorType
+import com.android.wallpaper.picker.customization.data.repository.WallpaperColorsRepository
+import com.android.wallpaper.picker.customization.shared.model.WallpaperColorsModel
+import com.android.wallpaper.picker.di.modules.BackgroundDispatcher
+import javax.inject.Inject
+import javax.inject.Singleton
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.callbackFlow
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.suspendCancellableCoroutine
+
+@Singleton
+class ColorPickerRepositoryImpl2
+@Inject
+constructor(
+    @BackgroundDispatcher private val scope: CoroutineScope,
+    wallpaperColorsRepository: WallpaperColorsRepository,
+    private val colorManager: ColorCustomizationManager,
+) : ColorPickerRepository2 {
+
+    private val homeWallpaperColors: StateFlow<WallpaperColorsModel?> =
+        wallpaperColorsRepository.homeWallpaperColors
+    private val lockWallpaperColors: StateFlow<WallpaperColorsModel?> =
+        wallpaperColorsRepository.lockWallpaperColors
+
+    override val colorOptions: Flow<Map<ColorType, List<ColorOption>>> =
+        combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
+                homeColors to lockColors
+            }
+            .map { (homeColors, lockColors) ->
+                suspendCancellableCoroutine { continuation ->
+                    if (
+                        homeColors is WallpaperColorsModel.Loading ||
+                            lockColors is WallpaperColorsModel.Loading
+                    ) {
+                        continuation.resumeWith(
+                            Result.success(
+                                mapOf(
+                                    ColorType.WALLPAPER_COLOR to listOf(),
+                                    ColorType.PRESET_COLOR to listOf(),
+                                )
+                            )
+                        )
+                        return@suspendCancellableCoroutine
+                    }
+                    val homeColorsLoaded = homeColors as WallpaperColorsModel.Loaded
+                    val lockColorsLoaded = lockColors as WallpaperColorsModel.Loaded
+                    colorManager.setWallpaperColors(
+                        homeColorsLoaded.colors,
+                        lockColorsLoaded.colors,
+                    )
+                    colorManager.fetchOptions(
+                        object : CustomizationManager.OptionsFetchedListener<ColorOption> {
+                            override fun onOptionsLoaded(options: MutableList<ColorOption>?) {
+                                val wallpaperColorOptions: MutableList<ColorOption> =
+                                    mutableListOf()
+                                val presetColorOptions: MutableList<ColorOption> = mutableListOf()
+                                options?.forEach { option ->
+                                    when ((option as ColorOptionImpl).type) {
+                                        ColorType.WALLPAPER_COLOR ->
+                                            wallpaperColorOptions.add(option)
+                                        ColorType.PRESET_COLOR -> presetColorOptions.add(option)
+                                    }
+                                }
+                                continuation.resumeWith(
+                                    Result.success(
+                                        mapOf(
+                                            ColorType.WALLPAPER_COLOR to wallpaperColorOptions,
+                                            ColorType.PRESET_COLOR to presetColorOptions,
+                                        )
+                                    )
+                                )
+                            }
+
+                            override fun onError(throwable: Throwable?) {
+                                Log.e(TAG, "Error loading theme bundles", throwable)
+                                continuation.resumeWith(
+                                    Result.failure(
+                                        throwable ?: Throwable("Error loading theme bundles")
+                                    )
+                                )
+                            }
+                        },
+                        /* reload= */ false,
+                    )
+                }
+            }
+
+    private val settingsChanged = callbackFlow {
+        trySend(Unit)
+        colorManager.setListener { trySend(Unit) }
+        awaitClose { colorManager.setListener(null) }
+    }
+
+    override val selectedColorOption =
+        combine(colorOptions, settingsChanged) { options, _ ->
+                options.forEach { (_, optionsByType) ->
+                    optionsByType.forEach {
+                        if (it.isActive(colorManager)) {
+                            return@combine it
+                        }
+                    }
+                }
+                return@combine null
+            }
+            .stateIn(scope = scope, started = SharingStarted.WhileSubscribed(), initialValue = null)
+
+    override suspend fun select(colorOption: ColorOption) {
+        suspendCancellableCoroutine { continuation ->
+            colorManager.apply(
+                colorOption,
+                object : CustomizationManager.Callback {
+                    override fun onSuccess() {
+                        continuation.resumeWith(Result.success(Unit))
+                    }
+
+                    override fun onError(throwable: Throwable?) {
+                        Log.w(TAG, "Apply theme with error", throwable)
+                        continuation.resumeWith(
+                            Result.failure(throwable ?: Throwable("Error loading theme bundles"))
+                        )
+                    }
+                },
+            )
+        }
+    }
+
+    companion object {
+        private const val TAG = "ColorPickerRepositoryImpl"
+    }
+}
diff --git a/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor2.kt b/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor2.kt
new file mode 100644
index 0000000..df69660
--- /dev/null
+++ b/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor2.kt
@@ -0,0 +1,35 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+package com.android.customization.picker.color.domain.interactor
+
+import com.android.customization.model.color.ColorOption
+import com.android.customization.picker.color.data.repository.ColorPickerRepository2
+import javax.inject.Inject
+import javax.inject.Singleton
+
+/** Single entry-point for all application state and business logic related to system color. */
+@Singleton
+class ColorPickerInteractor2 @Inject constructor(private val repository: ColorPickerRepository2) {
+    val selectedColorOption = repository.selectedColorOption
+
+    /** List of wallpaper and preset color options on the device, categorized by Color Type */
+    val colorOptions = repository.colorOptions
+
+    suspend fun select(colorOption: ColorOption) {
+        repository.select(colorOption)
+    }
+}
diff --git a/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt
index 7ddcb01..4845121 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt
@@ -93,11 +93,11 @@
                 }
 
                 launch {
-                    viewModel.previewingColorOption.collect { colorModel ->
-                        if (colorModel != null) {
+                    viewModel.previewingColorOption.collect { colorOption ->
+                        if (colorOption != null) {
                             colorUpdateViewModel.previewColors(
-                                colorModel.colorOption.seedColor,
-                                colorModel.colorOption.style,
+                                colorOption.seedColor,
+                                colorOption.style,
                             )
                         } else colorUpdateViewModel.resetPreview()
                     }
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
index 2a1a8c9..e055d48 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
@@ -19,15 +19,14 @@
 import android.content.res.Resources
 import android.graphics.drawable.Drawable
 import androidx.core.graphics.ColorUtils
+import com.android.customization.model.color.ColorOption
 import com.android.customization.model.color.ColorOptionImpl
 import com.android.customization.module.logging.ThemesUserEventLogger
 import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
 import com.android.customization.picker.clock.shared.ClockSize
 import com.android.customization.picker.clock.shared.model.ClockMetadataModel
 import com.android.customization.picker.clock.ui.viewmodel.ClockColorViewModel
-import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
-import com.android.customization.picker.color.shared.model.ColorOptionModel
-import com.android.customization.picker.color.shared.model.ColorType
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor2
 import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
 import com.android.systemui.plugins.clocks.ClockFontAxisSetting
 import com.android.themepicker.R
@@ -67,7 +66,7 @@
     @ApplicationContext context: Context,
     resources: Resources,
     private val clockPickerInteractor: ClockPickerInteractor,
-    colorPickerInteractor: ColorPickerInteractor,
+    colorPickerInteractor: ColorPickerInteractor2,
     private val logger: ThemesUserEventLogger,
     @BackgroundDispatcher private val backgroundDispatcher: CoroutineDispatcher,
     @Assisted private val viewModelScope: CoroutineScope,
@@ -292,19 +291,12 @@
         }
 
     val clockColorOptions: Flow<List<OptionItemViewModel<ColorOptionIconViewModel>>> =
-        colorPickerInteractor.colorOptions.map { colorOptions ->
+        colorPickerInteractor.selectedColorOption.map { selectedColorOption ->
             // Use mapLatest and delay(100) here to prevent too many selectedClockColor update
             // events from ClockRegistry upstream, caused by sliding the saturation level bar.
             delay(COLOR_OPTIONS_EVENT_UPDATE_DELAY_MILLIS)
             buildList {
-                val defaultThemeColorOptionViewModel =
-                    (colorOptions[ColorType.WALLPAPER_COLOR]?.find { it.isSelected })
-                        ?.toOptionItemViewModel(context)
-                        ?: (colorOptions[ColorType.PRESET_COLOR]?.find { it.isSelected })
-                            ?.toOptionItemViewModel(context)
-                if (defaultThemeColorOptionViewModel != null) {
-                    add(defaultThemeColorOptionViewModel)
-                }
+                selectedColorOption?.let { add(it.toOptionItemViewModel(context)) }
 
                 colorMap.values.forEachIndexed { index, colorModel ->
                     val isSelectedFlow =
@@ -352,23 +344,24 @@
             }
         }
 
-    private suspend fun ColorOptionModel.toOptionItemViewModel(
+    private suspend fun ColorOption.toOptionItemViewModel(
         context: Context
     ): OptionItemViewModel<ColorOptionIconViewModel> {
         val lightThemeColors =
-            (colorOption as ColorOptionImpl)
+            (this as ColorOptionImpl)
                 .previewInfo
                 .resolveColors(
                     /** darkTheme= */
                     false
                 )
         val darkThemeColors =
-            colorOption.previewInfo.resolveColors(
+            this.previewInfo.resolveColors(
                 /** darkTheme= */
                 true
             )
         val isSelectedFlow =
             previewingClockColorId.map { it == DEFAULT_CLOCK_COLOR_ID }.stateIn(viewModelScope)
+        val key = "${this.type}::${this.style}::${this.serializedPackages}"
         return OptionItemViewModel<ColorOptionIconViewModel>(
             key = MutableStateFlow(key) as StateFlow<String>,
             payload =
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
index 9e2353a..02af6a6 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
@@ -17,10 +17,11 @@
 package com.android.wallpaper.customization.ui.viewmodel
 
 import android.content.Context
+import androidx.lifecycle.viewModelScope
+import com.android.customization.model.color.ColorOption
 import com.android.customization.model.color.ColorOptionImpl
 import com.android.customization.module.logging.ThemesUserEventLogger
-import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
-import com.android.customization.picker.color.shared.model.ColorOptionModel
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor2
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
 import com.android.themepicker.R
@@ -51,12 +52,12 @@
 constructor(
     @ApplicationContext context: Context,
     private val colorUpdateViewModel: ColorUpdateViewModel,
-    private val interactor: ColorPickerInteractor,
+    private val interactor: ColorPickerInteractor2,
     private val logger: ThemesUserEventLogger,
     @Assisted private val viewModelScope: CoroutineScope,
 ) {
 
-    private val overridingColorOption = MutableStateFlow<ColorOptionModel?>(null)
+    private val overridingColorOption = MutableStateFlow<ColorOption?>(null)
     val previewingColorOption = overridingColorOption.asStateFlow()
 
     private val selectedColorTypeTabId = MutableStateFlow<ColorType?>(null)
@@ -117,22 +118,25 @@
             colorOptions
                 .map { colorOptionEntry ->
                     colorOptionEntry.key to
-                        colorOptionEntry.value.map { colorOptionModel ->
-                            val colorOption: ColorOptionImpl =
-                                colorOptionModel.colorOption as ColorOptionImpl
+                        colorOptionEntry.value.map { colorOption ->
+                            colorOption as ColorOptionImpl
                             val lightThemeColors =
                                 colorOption.previewInfo.resolveColors(/* darkTheme= */ false)
                             val darkThemeColors =
                                 colorOption.previewInfo.resolveColors(/* darkTheme= */ true)
                             val isSelectedFlow: StateFlow<Boolean> =
-                                previewingColorOption
-                                    .map {
-                                        it?.colorOption?.isEquivalent(colorOptionModel.colorOption)
-                                            ?: colorOptionModel.isSelected
+                                combine(previewingColorOption, interactor.selectedColorOption) {
+                                        previewing,
+                                        selected ->
+                                        previewing?.isEquivalent(colorOption)
+                                            ?: selected?.isEquivalent(colorOption)
+                                            ?: false
                                     }
                                     .stateIn(viewModelScope)
+                            val key =
+                                "${colorOption.type}::${colorOption.style}::${colorOption.serializedPackages}"
                             OptionItemViewModel2<ColorOptionIconViewModel>(
-                                key = MutableStateFlow(colorOptionModel.key) as StateFlow<String>,
+                                key = MutableStateFlow(key) as StateFlow<String>,
                                 payload =
                                     ColorOptionIconViewModel(
                                         lightThemeColor0 = lightThemeColors[0],
@@ -157,7 +161,7 @@
                                         } else {
                                             {
                                                 viewModelScope.launch {
-                                                    overridingColorOption.value = colorOptionModel
+                                                    overridingColorOption.value = colorOption
                                                 }
                                             }
                                         }
@@ -173,9 +177,9 @@
      * change updates, which are applied with a latency.
      */
     val onApply: Flow<(suspend () -> Unit)?> =
-        previewingColorOption.map { previewingColorOption ->
-            previewingColorOption?.let {
-                if (it.isSelected) {
+        combine(previewingColorOption, interactor.selectedColorOption) { previewing, selected ->
+            previewing?.let {
+                if (previewing.isEquivalent(selected)) {
                     null
                 } else {
                     {
@@ -185,9 +189,9 @@
                             return@collect
                         }
                         logger.logThemeColorApplied(
-                            previewingColorOption.colorOption.sourceForLogging,
-                            previewingColorOption.colorOption.styleForLogging,
-                            previewingColorOption.colorOption.seedColor,
+                            it.sourceForLogging,
+                            it.styleForLogging,
+                            it.seedColor,
                         )
                     }
                 }
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 5e5bc1f..d1c5695 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
@@ -216,14 +216,14 @@
                                     viewModel.darkModeViewModel.overridingIsDarkMode,
                                     ::Pair,
                                 )
-                                .collect { (colorModel, darkMode) ->
+                                .collect { (colorOption, darkMode) ->
                                     val bundle =
                                         Bundle().apply {
-                                            if (colorModel != null) {
+                                            if (colorOption != null) {
                                                 val (ids, colors) =
                                                     materialColorsGenerator.generate(
-                                                        colorModel.colorOption.seedColor,
-                                                        colorModel.colorOption.style,
+                                                        colorOption.seedColor,
+                                                        colorOption.style,
                                                     )
                                                 putIntArray(KEY_COLOR_RESOURCE_IDS, ids)
                                                 putIntArray(KEY_COLOR_VALUES, colors)
diff --git a/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt b/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt
index 1c4ecc9..d5c3005 100644
--- a/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt
+++ b/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt
@@ -27,7 +27,9 @@
 import com.android.customization.picker.clock.data.repository.ClockPickerRepositoryImpl
 import com.android.customization.picker.clock.data.repository.ClockRegistryProvider
 import com.android.customization.picker.color.data.repository.ColorPickerRepository
+import com.android.customization.picker.color.data.repository.ColorPickerRepository2
 import com.android.customization.picker.color.data.repository.ColorPickerRepositoryImpl
+import com.android.customization.picker.color.data.repository.ColorPickerRepositoryImpl2
 import com.android.systemui.shared.clocks.ClockRegistry
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClientImpl
@@ -89,6 +91,12 @@
 
     @Binds
     @Singleton
+    abstract fun bindColorPickerRepository2(
+        impl: ColorPickerRepositoryImpl2
+    ): ColorPickerRepository2
+
+    @Binds
+    @Singleton
     abstract fun bindCreativeCategoryInteractor(
         impl: CreativeCategoryInteractorImpl
     ): CreativeCategoryInteractor
diff --git a/tests/common/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository2.kt b/tests/common/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository2.kt
new file mode 100644
index 0000000..93f46cc
--- /dev/null
+++ b/tests/common/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository2.kt
@@ -0,0 +1,204 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+package com.android.customization.picker.color.data.repository
+
+import android.graphics.Color
+import com.android.customization.model.ResourceConstants
+import com.android.customization.model.color.ColorOption
+import com.android.customization.model.color.ColorOptionImpl
+import com.android.customization.model.color.ColorOptionsProvider
+import com.android.customization.model.color.ColorUtils.toColorString
+import com.android.customization.picker.color.shared.model.ColorType
+import com.android.systemui.monet.Style
+import javax.inject.Inject
+import javax.inject.Singleton
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+
+@Singleton
+class FakeColorPickerRepository2 @Inject constructor() : ColorPickerRepository2 {
+
+    private val _selectedColorOption = MutableStateFlow<ColorOption?>(null)
+    override val selectedColorOption = _selectedColorOption.asStateFlow()
+
+    private val _colorOptions =
+        MutableStateFlow(
+            mapOf<ColorType, List<ColorOption>>(
+                ColorType.WALLPAPER_COLOR to listOf(),
+                ColorType.PRESET_COLOR to listOf(),
+            )
+        )
+    override val colorOptions: StateFlow<Map<ColorType, List<ColorOption>>> =
+        _colorOptions.asStateFlow()
+
+    init {
+        setOptions(4, 4, ColorType.WALLPAPER_COLOR, 0)
+    }
+
+    fun setOptions(
+        wallpaperOptions: List<ColorOptionImpl>,
+        presetOptions: List<ColorOptionImpl>,
+        selectedColorOptionType: ColorType,
+        selectedColorOptionIndex: Int,
+    ) {
+        _colorOptions.value =
+            mapOf(
+                ColorType.WALLPAPER_COLOR to
+                    buildList {
+                        for ((index, colorOption) in wallpaperOptions.withIndex()) {
+                            val isSelected =
+                                selectedColorOptionType == ColorType.WALLPAPER_COLOR &&
+                                    selectedColorOptionIndex == index
+                            if (isSelected) {
+                                _selectedColorOption.value = colorOption
+                            }
+                            add(colorOption)
+                        }
+                    },
+                ColorType.PRESET_COLOR to
+                    buildList {
+                        for ((index, colorOption) in presetOptions.withIndex()) {
+                            val isSelected =
+                                selectedColorOptionType == ColorType.PRESET_COLOR &&
+                                    selectedColorOptionIndex == index
+                            if (isSelected) {
+                                _selectedColorOption.value = colorOption
+                            }
+                            add(colorOption)
+                        }
+                    },
+            )
+    }
+
+    fun setOptions(
+        numWallpaperOptions: Int,
+        numPresetOptions: Int,
+        selectedColorOptionType: ColorType,
+        selectedColorOptionIndex: Int,
+    ) {
+        _colorOptions.value =
+            mapOf(
+                ColorType.WALLPAPER_COLOR to
+                    buildList {
+                        repeat(times = numWallpaperOptions) { index ->
+                            val isSelected =
+                                selectedColorOptionType == ColorType.WALLPAPER_COLOR &&
+                                    selectedColorOptionIndex == index
+                            val colorOption = buildWallpaperOption(index)
+                            if (isSelected) {
+                                _selectedColorOption.value = colorOption
+                            }
+                            add(colorOption)
+                        }
+                    },
+                ColorType.PRESET_COLOR to
+                    buildList {
+                        repeat(times = numPresetOptions) { index ->
+                            val isSelected =
+                                selectedColorOptionType == ColorType.PRESET_COLOR &&
+                                    selectedColorOptionIndex == index
+                            val colorOption = buildPresetOption(index)
+                            if (isSelected) {
+                                _selectedColorOption.value = colorOption
+                            }
+                            add(colorOption)
+                        }
+                    },
+            )
+    }
+
+    private fun buildPresetOption(index: Int): ColorOptionImpl {
+        val builder = ColorOptionImpl.Builder()
+        builder.lightColors =
+            intArrayOf(Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT)
+        builder.darkColors =
+            intArrayOf(Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT)
+        builder.index = index
+        builder.type = ColorType.PRESET_COLOR
+        builder.source = ColorOptionsProvider.COLOR_SOURCE_PRESET
+        builder.title = "Preset"
+        builder
+            .addOverlayPackage("TEST_PACKAGE_TYPE", "preset_color")
+            .addOverlayPackage("TEST_PACKAGE_INDEX", "$index")
+        return builder.build()
+    }
+
+    fun buildPresetOption(@Style.Type style: Int, seedColor: Int): ColorOptionImpl {
+        val builder = ColorOptionImpl.Builder()
+        builder.lightColors =
+            intArrayOf(Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT)
+        builder.darkColors =
+            intArrayOf(Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT)
+        builder.type = ColorType.PRESET_COLOR
+        builder.source = ColorOptionsProvider.COLOR_SOURCE_PRESET
+        builder.style = style
+        builder.title = "Preset"
+        builder.seedColor = seedColor
+        builder
+            .addOverlayPackage("TEST_PACKAGE_TYPE", "preset_color")
+            .addOverlayPackage(
+                ResourceConstants.OVERLAY_CATEGORY_SYSTEM_PALETTE,
+                toColorString(seedColor),
+            )
+        return builder.build()
+    }
+
+    private fun buildWallpaperOption(index: Int): ColorOptionImpl {
+        val builder = ColorOptionImpl.Builder()
+        builder.lightColors =
+            intArrayOf(Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT)
+        builder.darkColors =
+            intArrayOf(Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT)
+        builder.index = index
+        builder.type = ColorType.WALLPAPER_COLOR
+        builder.source = ColorOptionsProvider.COLOR_SOURCE_HOME
+        builder.title = "Dynamic"
+        builder
+            .addOverlayPackage("TEST_PACKAGE_TYPE", "wallpaper_color")
+            .addOverlayPackage("TEST_PACKAGE_INDEX", "$index")
+        return builder.build()
+    }
+
+    fun buildWallpaperOption(
+        source: String,
+        @Style.Type style: Int,
+        seedColor: Int,
+    ): ColorOptionImpl {
+        val builder = ColorOptionImpl.Builder()
+        builder.lightColors =
+            intArrayOf(Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT)
+        builder.darkColors =
+            intArrayOf(Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT, Color.TRANSPARENT)
+        builder.type = ColorType.WALLPAPER_COLOR
+        builder.source = source
+        builder.style = style
+        builder.title = "Dynamic"
+        builder.seedColor = seedColor
+        builder
+            .addOverlayPackage("TEST_PACKAGE_TYPE", "wallpaper_color")
+            .addOverlayPackage(
+                ResourceConstants.OVERLAY_CATEGORY_SYSTEM_PALETTE,
+                toColorString(seedColor),
+            )
+        return builder.build()
+    }
+
+    override suspend fun select(colorOption: ColorOption) {
+        _selectedColorOption.value = colorOption
+    }
+}
diff --git a/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt b/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
index bc03f12..36b95da 100644
--- a/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
+++ b/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
@@ -27,7 +27,9 @@
 import com.android.customization.picker.clock.data.repository.ClockPickerRepositoryImpl
 import com.android.customization.picker.clock.data.repository.ClockRegistryProvider
 import com.android.customization.picker.color.data.repository.ColorPickerRepository
+import com.android.customization.picker.color.data.repository.ColorPickerRepository2
 import com.android.customization.picker.color.data.repository.ColorPickerRepositoryImpl
+import com.android.customization.picker.color.data.repository.FakeColorPickerRepository2
 import com.android.customization.testing.TestCustomizationInjector
 import com.android.customization.testing.TestDefaultCustomizationPreferences
 import com.android.systemui.shared.clocks.ClockRegistry
@@ -90,6 +92,12 @@
 
     @Binds
     @Singleton
+    abstract fun bindColorPickerRepository2(
+        impl: FakeColorPickerRepository2
+    ): ColorPickerRepository2
+
+    @Binds
+    @Singleton
     abstract fun bindCustomizationInjector(impl: TestCustomizationInjector): CustomizationInjector
 
     @Binds
diff --git a/tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractor2Test.kt b/tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractor2Test.kt
new file mode 100644
index 0000000..00152dd
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractor2Test.kt
@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+package com.android.customization.model.picker.color.domain.interactor
+
+import android.content.Context
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.customization.picker.color.data.repository.FakeColorPickerRepository2
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor2
+import com.android.customization.picker.color.shared.model.ColorType
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::class)
+class ColorPickerInteractor2Test {
+    private lateinit var underTest: ColorPickerInteractor2
+    private lateinit var repository: FakeColorPickerRepository2
+    private lateinit var store: FakeSnapshotStore
+
+    private lateinit var context: Context
+
+    @Before
+    fun setUp() {
+        context = InstrumentationRegistry.getInstrumentation().targetContext
+        repository = FakeColorPickerRepository2()
+        store = FakeSnapshotStore()
+        underTest = ColorPickerInteractor2(repository = repository)
+        repository.setOptions(4, 4, ColorType.WALLPAPER_COLOR, 0)
+    }
+
+    @Test
+    fun select() = runTest {
+        val colorOptions = collectLastValue(underTest.colorOptions)
+        val selectedColorOption = collectLastValue(underTest.selectedColorOption)
+
+        val wallpaperColorOption = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(2)
+        assertThat(selectedColorOption()).isNotEqualTo(wallpaperColorOption)
+
+        wallpaperColorOption?.let { underTest.select(colorOption = it) }
+        assertThat(selectedColorOption()).isEqualTo(wallpaperColorOption)
+
+        val presetColorOption = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(1)
+        assertThat(selectedColorOption()).isNotEqualTo(presetColorOption)
+
+        presetColorOption?.let { underTest.select(colorOption = it) }
+        assertThat(selectedColorOption()).isEqualTo(presetColorOption)
+    }
+}
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModelTest.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModelTest.kt
index 76df409..b035c84 100644
--- a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModelTest.kt
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModelTest.kt
@@ -26,9 +26,8 @@
 import com.android.customization.picker.clock.shared.model.ClockMetadataModel
 import com.android.customization.picker.clock.ui.viewmodel.ClockColorViewModel
 import com.android.customization.picker.clock.ui.viewmodel.ClockSettingsViewModel
-import com.android.customization.picker.color.data.repository.FakeColorPickerRepository
-import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
-import com.android.customization.picker.color.domain.interactor.ColorPickerSnapshotRestorer
+import com.android.customization.picker.color.data.repository.FakeColorPickerRepository2
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor2
 import com.android.wallpaper.customization.ui.viewmodel.ClockPickerViewModel.Tab
 import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
@@ -83,15 +82,8 @@
                         runBlocking { setUpSnapshotRestorer(store = FakeSnapshotStore()) }
                     },
             )
-        val colorPickerRepository = FakeColorPickerRepository(context = context)
-        val colorPickerInteractor =
-            ColorPickerInteractor(
-                repository = colorPickerRepository,
-                snapshotRestorer =
-                    ColorPickerSnapshotRestorer(repository = colorPickerRepository).apply {
-                        runBlocking { setUpSnapshotRestorer(store = FakeSnapshotStore()) }
-                    },
-            )
+        val colorPickerRepository = FakeColorPickerRepository2()
+        val colorPickerInteractor = ColorPickerInteractor2(repository = colorPickerRepository)
         colorMap = ClockColorViewModel.getPresetColorMap(context.resources)
         underTest =
             ClockPickerViewModel(
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 07c3a16..649c298 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
@@ -22,9 +22,8 @@
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.customization.model.color.ColorOptionsProvider
 import com.android.customization.module.logging.TestThemesUserEventLogger
-import com.android.customization.picker.color.data.repository.FakeColorPickerRepository
-import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
-import com.android.customization.picker.color.domain.interactor.ColorPickerSnapshotRestorer
+import com.android.customization.picker.color.data.repository.FakeColorPickerRepository2
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor2
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
 import com.android.systemui.monet.Style
@@ -39,7 +38,6 @@
 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
 import kotlinx.coroutines.test.resetMain
@@ -57,8 +55,8 @@
 class ColorPickerViewModel2Test {
     private val logger = TestThemesUserEventLogger()
     private lateinit var underTest: ColorPickerViewModel2
-    private lateinit var repository: FakeColorPickerRepository
-    private lateinit var interactor: ColorPickerInteractor
+    private lateinit var repository: FakeColorPickerRepository2
+    private lateinit var interactor: ColorPickerInteractor2
     private lateinit var store: FakeSnapshotStore
     private lateinit var colorUpdateViewModel: ColorUpdateViewModel
 
@@ -71,17 +69,10 @@
         val testDispatcher = UnconfinedTestDispatcher()
         Dispatchers.setMain(testDispatcher)
         testScope = TestScope(testDispatcher)
-        repository = FakeColorPickerRepository(context = context)
+        repository = FakeColorPickerRepository2()
         store = FakeSnapshotStore()
 
-        interactor =
-            ColorPickerInteractor(
-                repository = repository,
-                snapshotRestorer =
-                    ColorPickerSnapshotRestorer(repository = repository).apply {
-                        runBlocking { setUpSnapshotRestorer(store = store) }
-                    },
-            )
+        interactor = ColorPickerInteractor2(repository = repository)
 
         colorUpdateViewModel = ColorUpdateViewModel(context, RetainedLifecycleImpl())