Merge "Reset preview only after transition completed (2/2)" into main
diff --git a/res/layout/floating_sheet_shape_grid.xml b/res/layout/floating_sheet_shape_grid.xml
index 4e2409b..cd0a709 100644
--- a/res/layout/floating_sheet_shape_grid.xml
+++ b/res/layout/floating_sheet_shape_grid.xml
@@ -47,7 +47,7 @@
             It's critical for any TextViews inside the included layout to have text.
             -->
             <include
-                layout="@layout/shape_option"
+                layout="@layout/shape_option2"
                 android:layout_width="64dp"
                 android:layout_height="64dp"
                 android:visibility="invisible" />
@@ -78,7 +78,7 @@
             It's critical for any TextViews inside the included layout to have text.
             -->
             <include
-                layout="@layout/grid_option"
+                layout="@layout/grid_option2"
                 android:id="@+id/invisible_grid_option"
                 android:layout_width="wrap_content"
                 android:layout_height="wrap_content"
diff --git a/res/layout/shape_option2.xml b/res/layout/shape_option2.xml
new file mode 100644
index 0000000..88d5437
--- /dev/null
+++ b/res/layout/shape_option2.xml
@@ -0,0 +1,36 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ 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.
+  ~
+  -->
+<FrameLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="64dp"
+    android:layout_height="64dp"
+    android:clipChildren="false">
+
+    <com.android.wallpaper.picker.option.ui.view.OptionItemBackground
+        android:id="@id/background"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        android:background="@drawable/option_item_background"
+        android:importantForAccessibility="no" />
+
+    <ImageView
+        android:id="@id/foreground"
+        android:layout_width="36dp"
+        android:layout_height="36dp"
+        android:layout_gravity="center" />
+</FrameLayout>
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/ColorPickerRepositoryImpl.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
index f393880..43d510c 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -24,7 +24,6 @@
 import com.android.customization.picker.color.shared.model.ColorOptionModel
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.systemui.monet.Style
-import com.android.wallpaper.config.BaseFlags
 import com.android.wallpaper.picker.customization.data.repository.WallpaperColorsRepository
 import com.android.wallpaper.picker.customization.shared.model.WallpaperColorsModel
 import javax.inject.Inject
@@ -47,8 +46,6 @@
     private val colorManager: ColorCustomizationManager,
 ) : ColorPickerRepository {
 
-    private val isNewPickerUi = BaseFlags.get().isNewPickerUi()
-
     private val homeWallpaperColors: StateFlow<WallpaperColorsModel?> =
         wallpaperColorsRepository.homeWallpaperColors
     private val lockWallpaperColors: StateFlow<WallpaperColorsModel?> =
@@ -59,7 +56,7 @@
     private val _isApplyingSystemColor = MutableStateFlow(false)
     override val isApplyingSystemColor = _isApplyingSystemColor.asStateFlow()
 
-    private val generatedColorOptions: Flow<Map<ColorType, List<ColorOptionImpl>>> =
+    override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
         combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
                 homeColors to lockColors
             }
@@ -88,15 +85,16 @@
                     colorManager.fetchOptions(
                         object : CustomizationManager.OptionsFetchedListener<ColorOption?> {
                             override fun onOptionsLoaded(options: MutableList<ColorOption?>?) {
-                                val wallpaperColorOptions: MutableList<ColorOptionImpl> =
+                                val wallpaperColorOptions: MutableList<ColorOptionModel> =
                                     mutableListOf()
-                                val presetColorOptions: MutableList<ColorOptionImpl> =
+                                val presetColorOptions: MutableList<ColorOptionModel> =
                                     mutableListOf()
                                 options?.forEach { option ->
                                     when ((option as ColorOptionImpl).type) {
                                         ColorType.WALLPAPER_COLOR ->
-                                            wallpaperColorOptions.add(option)
-                                        ColorType.PRESET_COLOR -> presetColorOptions.add(option)
+                                            wallpaperColorOptions.add(option.toModel())
+                                        ColorType.PRESET_COLOR ->
+                                            presetColorOptions.add(option.toModel())
                                     }
                                 }
                                 continuation.resumeWith(
@@ -123,83 +121,6 @@
                 }
             }
 
-    override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
-        if (isNewPickerUi) {
-            // Convert to ColorOptionModel. When the selected color option changes, update each
-            // ColorOptionModel's isSelected by calling toModel again.
-            combine(generatedColorOptions, selectedColorOption) { generatedColorOptions, _ ->
-                generatedColorOptions
-                    .map { entry ->
-                        entry.key to entry.value.map { colorOption -> colorOption.toModel() }
-                    }
-                    .toMap()
-            }
-        } else {
-            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<ColorOptionModel> =
-                                        mutableListOf()
-                                    val presetColorOptions: MutableList<ColorOptionModel> =
-                                        mutableListOf()
-                                    options?.forEach { option ->
-                                        when ((option as ColorOptionImpl).type) {
-                                            ColorType.WALLPAPER_COLOR ->
-                                                wallpaperColorOptions.add(option.toModel())
-                                            ColorType.PRESET_COLOR ->
-                                                presetColorOptions.add(option.toModel())
-                                        }
-                                    }
-                                    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,
-                        )
-                    }
-                }
-        }
-
     override suspend fun select(colorOptionModel: ColorOptionModel) {
         _isApplyingSystemColor.value = true
         suspendCancellableCoroutine { continuation ->
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/binder/ShapeGridFloatingSheetBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ShapeGridFloatingSheetBinder.kt
index 138a253..9cebd27 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ShapeGridFloatingSheetBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ShapeGridFloatingSheetBinder.kt
@@ -42,9 +42,7 @@
 import com.android.wallpaper.picker.customization.ui.view.FloatingToolbar
 import com.android.wallpaper.picker.customization.ui.view.adapter.FloatingToolbarTabAdapter
 import com.android.wallpaper.picker.customization.ui.viewmodel.ColorUpdateViewModel
-import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter
 import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter2
-import com.android.wallpaper.picker.option.ui.binder.OptionItemBinder
 import java.lang.ref.WeakReference
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.flow.Flow
@@ -90,7 +88,7 @@
 
         val shapeContent = view.requireViewById<View>(R.id.app_shape_container)
         val shapeOptionListAdapter =
-            createShapeOptionItemAdapter(view.context, lifecycleOwner, backgroundDispatcher)
+            createShapeOptionItemAdapter(lifecycleOwner, backgroundDispatcher)
         val shapeOptionList =
             view.requireViewById<RecyclerView>(R.id.shape_options).also {
                 it.initShapeOptionList(view.context, shapeOptionListAdapter)
@@ -209,30 +207,23 @@
     }
 
     private fun createShapeOptionItemAdapter(
-        context: Context,
         lifecycleOwner: LifecycleOwner,
         backgroundDispatcher: CoroutineDispatcher,
-    ): OptionItemAdapter<ShapeIconViewModel> =
-        OptionItemAdapter(
-            layoutResourceId = R.layout.shape_option,
+    ): OptionItemAdapter2<ShapeIconViewModel> =
+        OptionItemAdapter2(
+            layoutResourceId = R.layout.shape_option2,
             lifecycleOwner = lifecycleOwner,
             backgroundDispatcher = backgroundDispatcher,
-            foregroundTintSpec =
-                OptionItemBinder.TintSpec(
-                    selectedColor =
-                        context.getColor(com.android.wallpaper.R.color.system_on_surface),
-                    unselectedColor =
-                        context.getColor(com.android.wallpaper.R.color.system_on_surface),
-                ),
-            bindIcon = { foregroundView: View, shapeIcon: ShapeIconViewModel ->
-                val imageView = foregroundView as? ImageView
+            bindPayload = { view: View, shapeIcon: ShapeIconViewModel ->
+                val imageView = view.findViewById(R.id.foreground) as? ImageView
                 imageView?.let { ShapeIconViewBinder.bind(imageView, shapeIcon) }
+                return@OptionItemAdapter2 null
             },
         )
 
     private fun RecyclerView.initShapeOptionList(
         context: Context,
-        adapter: OptionItemAdapter<ShapeIconViewModel>,
+        adapter: OptionItemAdapter2<ShapeIconViewModel>,
     ) {
         apply {
             this.layoutManager = LinearLayoutManager(context, RecyclerView.HORIZONTAL, false)
diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
index c17775a..e35e473 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
@@ -28,6 +28,12 @@
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.lifecycleScope
 import androidx.lifecycle.repeatOnLifecycle
+import com.android.themepicker.R as ThemePickerR
+import com.android.wallpaper.R
+import com.android.wallpaper.customization.ui.util.ThemePickerCustomizationOptionUtil.ThemePickerHomeCustomizationOption.APP_SHAPE_GRID
+import com.android.wallpaper.customization.ui.util.ThemePickerCustomizationOptionUtil.ThemePickerHomeCustomizationOption.COLORS
+import com.android.wallpaper.customization.ui.util.ThemePickerCustomizationOptionUtil.ThemePickerLockCustomizationOption.CLOCK
+import com.android.wallpaper.customization.ui.util.ThemePickerCustomizationOptionUtil.ThemePickerLockCustomizationOption.SHORTCUTS
 import com.android.wallpaper.customization.ui.viewmodel.ThemePickerCustomizationOptionsViewModel
 import com.android.wallpaper.customization.ui.viewmodel.ToolbarHeightsViewModel
 import com.android.wallpaper.picker.customization.ui.binder.ColorUpdateBinder
@@ -160,6 +166,21 @@
                 }
 
                 launch {
+                    viewModel.selectedOption.collect {
+                        val stringResId =
+                            when (it) {
+                                COLORS -> ThemePickerR.string.color_picker_title
+                                APP_SHAPE_GRID -> ThemePickerR.string.shape_and_grid_title
+                                CLOCK -> ThemePickerR.string.custom_clocks_label
+                                SHORTCUTS ->
+                                    ThemePickerR.string.keyguard_quick_affordance_section_title
+                                else -> R.string.app_name
+                            }
+                        toolbar.title = toolbar.resources.getString(stringResId)
+                    }
+                }
+
+                launch {
                     combine(toolbarHeights, viewModel.isToolbarCollapsed, ::Pair).collect {
                         (toolbarHeights, isToolbarCollapsed) ->
                         val (navButtonHeight, toolbarHeight, applyButtonHeight) = toolbarHeights
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/customization/ui/viewmodel/ShapeGridPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
index 1e19e80..85422dd 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
@@ -28,7 +28,6 @@
 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.FloatingToolbarTabViewModel
-import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel2
 import dagger.assisted.Assisted
 import dagger.assisted.AssistedFactory
@@ -105,7 +104,7 @@
             overridingShapeOptionKey ?: selectedShapeKey
         }
 
-    val shapeOptions: Flow<List<OptionItemViewModel<ShapeIconViewModel>>> =
+    val shapeOptions: Flow<List<OptionItemViewModel2<ShapeIconViewModel>>> =
         interactor.shapeOptions
             .filterNotNull()
             .map { shapeOptions -> shapeOptions.map { toShapeOptionItemViewModel(it) } }
@@ -157,7 +156,7 @@
 
     private fun toShapeOptionItemViewModel(
         option: ShapeOptionModel
-    ): OptionItemViewModel<ShapeIconViewModel> {
+    ): OptionItemViewModel2<ShapeIconViewModel> {
         val isSelected =
             previewingShapeKey
                 .map { it == option.key }
@@ -167,7 +166,7 @@
                     initialValue = false,
                 )
 
-        return OptionItemViewModel(
+        return OptionItemViewModel2(
             key = MutableStateFlow(option.key),
             payload = ShapeIconViewModel(option.key, option.path),
             text = Text.Loaded(option.title),
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..a2989bb 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
@@ -46,10 +48,12 @@
 import com.android.wallpaper.picker.category.domain.interactor.CategoriesLoadingStatusInteractor
 import com.android.wallpaper.picker.category.domain.interactor.CategoryInteractor
 import com.android.wallpaper.picker.category.domain.interactor.CreativeCategoryInteractor
+import com.android.wallpaper.picker.category.domain.interactor.CuratedPhotosInteractor
 import com.android.wallpaper.picker.category.domain.interactor.ThirdPartyCategoryInteractor
 import com.android.wallpaper.picker.category.domain.interactor.implementations.CategoryInteractorImpl
 import com.android.wallpaper.picker.category.domain.interactor.implementations.CreativeCategoryInteractorImpl
 import com.android.wallpaper.picker.category.domain.interactor.implementations.DefaultCategoriesLoadingStatusInteractor
+import com.android.wallpaper.picker.category.domain.interactor.implementations.DefaultCuratedPhotosInteractorImpl
 import com.android.wallpaper.picker.category.domain.interactor.implementations.ThirdPartyCategoryInteractorImpl
 import com.android.wallpaper.picker.category.ui.view.providers.IndividualPickerFactory
 import com.android.wallpaper.picker.category.ui.view.providers.implementation.DefaultIndividualPickerFactory
@@ -89,6 +93,12 @@
 
     @Binds
     @Singleton
+    abstract fun bindColorPickerRepository2(
+        impl: ColorPickerRepositoryImpl2
+    ): ColorPickerRepository2
+
+    @Binds
+    @Singleton
     abstract fun bindCreativeCategoryInteractor(
         impl: CreativeCategoryInteractorImpl
     ): CreativeCategoryInteractor
@@ -101,6 +111,12 @@
 
     @Binds
     @Singleton
+    abstract fun bindCuratedPhotosInteractor(
+        impl: DefaultCuratedPhotosInteractorImpl
+    ): CuratedPhotosInteractor
+
+    @Binds
+    @Singleton
     abstract fun bindCustomizationInjector(impl: ThemePickerInjector): CustomizationInjector
 
     @Binds
diff --git a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt b/tests/common/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
similarity index 100%
rename from src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
rename to tests/common/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
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..2eb5c10 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
@@ -46,6 +48,7 @@
 import com.android.wallpaper.modules.ThemePickerAppModule
 import com.android.wallpaper.network.Requester
 import com.android.wallpaper.picker.category.domain.interactor.CategoryInteractor
+import com.android.wallpaper.picker.category.domain.interactor.CuratedPhotosInteractor
 import com.android.wallpaper.picker.category.domain.interactor.ThirdPartyCategoryInteractor
 import com.android.wallpaper.picker.category.ui.view.providers.IndividualPickerFactory
 import com.android.wallpaper.picker.category.ui.view.providers.implementation.DefaultIndividualPickerFactory
@@ -60,6 +63,7 @@
 import com.android.wallpaper.picker.preview.ui.util.DefaultImageEffectDialogUtil
 import com.android.wallpaper.picker.preview.ui.util.ImageEffectDialogUtil
 import com.android.wallpaper.testing.FakeCategoryInteractor
+import com.android.wallpaper.testing.FakeCuratedPhotosInteractorImpl
 import com.android.wallpaper.testing.FakeDefaultRequester
 import com.android.wallpaper.testing.FakeThirdPartyCategoryInteractor
 import com.android.wallpaper.testing.FakeWallpaperCategoryWrapper
@@ -90,6 +94,12 @@
 
     @Binds
     @Singleton
+    abstract fun bindColorPickerRepository2(
+        impl: FakeColorPickerRepository2
+    ): ColorPickerRepository2
+
+    @Binds
+    @Singleton
     abstract fun bindCustomizationInjector(impl: TestCustomizationInjector): CustomizationInjector
 
     @Binds
@@ -130,6 +140,12 @@
     @Singleton
     abstract fun bindCategoryInteractor(impl: FakeCategoryInteractor): CategoryInteractor
 
+    @Binds
+    @Singleton
+    abstract fun bindCuratedPhotosInteractor(
+        impl: FakeCuratedPhotosInteractorImpl
+    ): CuratedPhotosInteractor
+
     @Binds @Singleton abstract fun bindInjector(impl: TestCustomizationInjector): Injector
 
     @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())
 
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModelTest.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModelTest.kt
index 71ea0d9..2bca39c 100644
--- a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModelTest.kt
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModelTest.kt
@@ -26,7 +26,6 @@
 import com.android.customization.picker.grid.ui.viewmodel.GridIconViewModel
 import com.android.customization.picker.grid.ui.viewmodel.ShapeIconViewModel
 import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
-import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel2
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
@@ -251,7 +250,7 @@
         }
 
     private fun TestScope.assertShapeItem(
-        optionItem: OptionItemViewModel<ShapeIconViewModel>?,
+        optionItem: OptionItemViewModel2<ShapeIconViewModel>?,
         key: String,
         payload: ShapeIconViewModel?,
         text: Text,