Color Picker Section Refactor (1/3)

Refactoring color picker section to use the same clean architecture as
color picker fragment. Also adjusted color picker section for large
screen to hide overflow option and show `more colors` button. This feature is gated by revamped UI flag

Test: unit tests, screen recording: https://drive.google.com/file/d/1HTfZtpEtSOayNn79N66q8GgynhNo3vDn/view?usp=sharing&resourcekey=0-j3W7aCh7_Sc_MhTiSQ6Kiw
Bug: 262924584
Change-Id: I1457da64bb01f504e41ba90565a83ebf6e2adfb8
diff --git a/res/drawable/ic_nav_color.xml b/res/drawable/ic_nav_color.xml
new file mode 100644
index 0000000..cfa64e2
--- /dev/null
+++ b/res/drawable/ic_nav_color.xml
@@ -0,0 +1,23 @@
+<!--
+     Copyright (C) 2019 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.
+-->
+<!-- Represents the color icon (a palette) -->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="24"
+    android:viewportHeight="24">
+    <path android:fillColor="@android:color/white" android:pathData="M12,22Q9.95,22 8.125,21.212Q6.3,20.425 4.938,19.062Q3.575,17.7 2.788,15.875Q2,14.05 2,12Q2,9.925 2.812,8.1Q3.625,6.275 5.013,4.925Q6.4,3.575 8.25,2.787Q10.1,2 12.2,2Q14.2,2 15.975,2.688Q17.75,3.375 19.087,4.588Q20.425,5.8 21.212,7.463Q22,9.125 22,11.05Q22,13.925 20.25,15.462Q18.5,17 16,17H14.15Q13.925,17 13.838,17.125Q13.75,17.25 13.75,17.4Q13.75,17.7 14.125,18.262Q14.5,18.825 14.5,19.55Q14.5,20.8 13.812,21.4Q13.125,22 12,22ZM12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12ZM6.5,13Q7.15,13 7.575,12.575Q8,12.15 8,11.5Q8,10.85 7.575,10.425Q7.15,10 6.5,10Q5.85,10 5.425,10.425Q5,10.85 5,11.5Q5,12.15 5.425,12.575Q5.85,13 6.5,13ZM9.5,9Q10.15,9 10.575,8.575Q11,8.15 11,7.5Q11,6.85 10.575,6.425Q10.15,6 9.5,6Q8.85,6 8.425,6.425Q8,6.85 8,7.5Q8,8.15 8.425,8.575Q8.85,9 9.5,9ZM14.5,9Q15.15,9 15.575,8.575Q16,8.15 16,7.5Q16,6.85 15.575,6.425Q15.15,6 14.5,6Q13.85,6 13.425,6.425Q13,6.85 13,7.5Q13,8.15 13.425,8.575Q13.85,9 14.5,9ZM17.5,13Q18.15,13 18.575,12.575Q19,12.15 19,11.5Q19,10.85 18.575,10.425Q18.15,10 17.5,10Q16.85,10 16.425,10.425Q16,10.85 16,11.5Q16,12.15 16.425,12.575Q16.85,13 17.5,13ZM12,20Q12.225,20 12.363,19.875Q12.5,19.75 12.5,19.55Q12.5,19.2 12.125,18.725Q11.75,18.25 11.75,17.3Q11.75,16.25 12.475,15.625Q13.2,15 14.25,15H16Q17.65,15 18.825,14.037Q20,13.075 20,11.05Q20,8.025 17.688,6.012Q15.375,4 12.2,4Q8.8,4 6.4,6.325Q4,8.65 4,12Q4,15.325 6.338,17.663Q8.675,20 12,20Z"/>
+</vector>
diff --git a/res/layout/color_section_view2.xml b/res/layout/color_section_view2.xml
index 0a3fc7f..687bcef 100644
--- a/res/layout/color_section_view2.xml
+++ b/res/layout/color_section_view2.xml
@@ -14,18 +14,14 @@
      See the License for the specific language governing permissions and
      limitations under the License.
 -->
-<com.android.customization.picker.color.ui.view.ColorSectionView2
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:id="@+id/color_section_option_container"
+<com.android.customization.picker.color.ui.view.ColorSectionView2 xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:tools="http://schemas.android.com/tools"
     android:layout_width="match_parent"
     android:layout_height="wrap_content"
     android:layout_marginBottom="@dimen/section_bottom_padding"
     android:layout_marginHorizontal="@dimen/section_horizontal_padding"
-    android:orientation="horizontal"
-    android:background="@drawable/top_connected_section_background"
-    android:paddingVertical="24dp"
-    android:paddingHorizontal="24dp"
-    android:weightSum="@integer/color_section_num_columns">
+    android:orientation="vertical"
+    android:background="@drawable/top_connected_section_background">
 
     <!--
         This is just an invisible placeholder put in place so that the parent keeps its height
@@ -36,10 +32,36 @@
 
         It's critical for any TextViews inside the included layout to have text.
         -->
-    <include
-        android:layout_width="0dp"
+    <LinearLayout
+        android:id="@+id/color_section_option_container"
+        android:layout_width="match_parent"
         android:layout_height="wrap_content"
-        layout="@layout/color_option_overflow_no_background"
-        android:visibility="invisible"
-        android:layout_weight="1"/>
+        android:orientation="horizontal"
+        android:paddingVertical="24dp"
+        android:paddingHorizontal="24dp"
+        android:weightSum="@integer/color_section_num_columns">
+        <include
+            android:layout_width="0dp"
+            android:layout_height="wrap_content"
+            layout="@layout/color_option_overflow_no_background"
+            android:visibility="invisible"
+            android:layout_weight="1"/>
+    </LinearLayout>
+
+    <TextView
+        android:id="@+id/more_colors"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:layout_gravity="center_horizontal"
+        android:layout_marginTop="10dp"
+        android:minHeight="48dp"
+        android:gravity="center"
+        android:drawablePadding="12dp"
+        android:drawableStart="@drawable/ic_nav_color"
+        android:drawableTint="@color/text_color_primary"
+        android:text="@string/more_colors"
+        android:textAppearance="@android:style/TextAppearance.DeviceDefault.Small"
+        android:textColor="@color/text_color_primary"
+        android:visibility="gone"
+        tools:ignore="UseCompatTextViewDrawableXml" />
 </com.android.customization.picker.color.ui.view.ColorSectionView2>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index f24ea28..f4d525e 100755
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -412,4 +412,12 @@
     [CHAR LIMIT=64].
     -->
     <string name="more_settings_section_description">Text on lock screen, Now Playing, and more</string>
+
+    <!--
+    Label for button that lets the user navigate to a full-screen experience of selecting
+    system colors.
+
+    [CHAR LIMIT=128].
+    -->
+    <string name="more_colors">More Colors</string>
 </resources>
diff --git a/src/com/android/customization/model/color/ColorSectionController2.java b/src/com/android/customization/model/color/ColorSectionController2.java
deleted file mode 100644
index 791a9a0..0000000
--- a/src/com/android/customization/model/color/ColorSectionController2.java
+++ /dev/null
@@ -1,243 +0,0 @@
-/*
- * Copyright (C) 2023 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.color;
-
-import static com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_HOME;
-import static com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_LOCK;
-import static com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_PRESET;
-
-import android.app.Activity;
-import android.app.WallpaperColors;
-import android.content.Context;
-import android.os.Handler;
-import android.os.SystemClock;
-import android.stats.style.StyleEnums;
-import android.text.TextUtils;
-import android.util.Log;
-import android.view.LayoutInflater;
-
-import androidx.annotation.Nullable;
-import androidx.lifecycle.LifecycleOwner;
-
-import com.android.customization.model.CustomizationManager;
-import com.android.customization.model.theme.OverlayManagerCompat;
-import com.android.customization.module.CustomizationInjector;
-import com.android.customization.module.ThemesUserEventLogger;
-import com.android.customization.picker.color.ui.fragment.ColorPickerFragment;
-import com.android.customization.picker.color.ui.view.ColorSectionView2;
-import com.android.wallpaper.R;
-import com.android.wallpaper.model.CustomizationSectionController;
-import com.android.wallpaper.model.WallpaperColorsViewModel;
-import com.android.wallpaper.module.InjectorProvider;
-
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Lists;
-
-import java.util.ArrayList;
-import java.util.List;
-
-/**
- * Color section view's controller for the logic of color customization.
- *
- * TODO (b/262924584): Convert ColorSectionController2 into Kotlin & use new architecture
- */
-public class ColorSectionController2 implements CustomizationSectionController<ColorSectionView2> {
-
-    private static final String TAG = "ColorSectionController";
-    private static final long MIN_COLOR_APPLY_PERIOD = 500L;
-
-    private final ThemesUserEventLogger mEventLogger;
-    private final ColorCustomizationManager mColorManager;
-    private final WallpaperColorsViewModel mWallpaperColorsViewModel;
-    private final LifecycleOwner mLifecycleOwner;
-    private final CustomizationSectionNavigationController mSectionNavigationController;
-
-    private List<ColorOption> mWallpaperColorOptions = new ArrayList<>();
-    private List<ColorOption> mPresetColorOptions = new ArrayList<>();
-    private ColorOption mSelectedColor;
-    @Nullable private WallpaperColors mHomeWallpaperColors;
-    @Nullable private WallpaperColors mLockWallpaperColors;
-    // Uses a boolean value to indicate whether wallpaper color is ready because WallpaperColors
-    // maybe be null when it's ready.
-    private boolean mHomeWallpaperColorsReady;
-    private boolean mLockWallpaperColorsReady;
-    private long mLastColorApplyingTime = 0L;
-    private ColorSectionView2 mColorSectionView;
-
-    public ColorSectionController2(Activity activity, WallpaperColorsViewModel viewModel,
-            LifecycleOwner lifecycleOwner,
-            CustomizationSectionNavigationController sectionNavigationController) {
-        CustomizationInjector injector = (CustomizationInjector) InjectorProvider.getInjector();
-        mEventLogger = (ThemesUserEventLogger) injector.getUserEventLogger(activity);
-        mColorManager = ColorCustomizationManager.getInstance(activity,
-                new OverlayManagerCompat(activity));
-        mWallpaperColorsViewModel = viewModel;
-        mLifecycleOwner = lifecycleOwner;
-        mSectionNavigationController = sectionNavigationController;
-    }
-
-    @Override
-    public boolean isAvailable(@Nullable Context context) {
-        return context != null && ColorUtils.isMonetEnabled(context) && mColorManager.isAvailable();
-    }
-
-    @Override
-    public ColorSectionView2 createView(Context context) {
-        mColorSectionView = (ColorSectionView2) LayoutInflater.from(context).inflate(
-                R.layout.color_section_view2, /* root= */ null);
-
-        mWallpaperColorsViewModel.getHomeWallpaperColorsLiveData().observe(mLifecycleOwner,
-                homeColors -> {
-                    mHomeWallpaperColors = homeColors;
-                    mHomeWallpaperColorsReady = true;
-                    maybeLoadColors();
-                });
-        mWallpaperColorsViewModel.getLockWallpaperColorsLiveData().observe(mLifecycleOwner,
-                lockColors -> {
-                    mLockWallpaperColors = lockColors;
-                    mLockWallpaperColorsReady = true;
-                    maybeLoadColors();
-                });
-        return mColorSectionView;
-    }
-
-    private void maybeLoadColors() {
-        if (mHomeWallpaperColorsReady && mLockWallpaperColorsReady) {
-            mColorManager.setWallpaperColors(mHomeWallpaperColors, mLockWallpaperColors);
-            loadColorOptions(/* reload= */ false);
-        }
-    }
-
-    private void loadColorOptions(boolean reload) {
-        mColorManager.fetchOptions(new CustomizationManager.OptionsFetchedListener<ColorOption>() {
-            @Override
-            public void onOptionsLoaded(List<ColorOption> options) {
-                List<ColorOption> wallpaperColorOptions = new ArrayList<>();
-                List<ColorOption> presetColorOptions = new ArrayList<>();
-                for (ColorOption option : options) {
-                    if (option instanceof ColorSeedOption) {
-                        wallpaperColorOptions.add(option);
-                    } else if (option instanceof ColorBundle) {
-                        presetColorOptions.add(option);
-                    }
-                }
-                mWallpaperColorOptions = wallpaperColorOptions;
-                mPresetColorOptions = presetColorOptions;
-                mSelectedColor = findActiveColorOption(mWallpaperColorOptions,
-                        mPresetColorOptions);
-
-                mColorSectionView.post(() -> setUpColorSectionView(mWallpaperColorOptions,
-                        mPresetColorOptions));
-            }
-
-            @Override
-            public void onError(@Nullable Throwable throwable) {
-                if (throwable != null) {
-                    Log.e(TAG, "Error loading theme bundles", throwable);
-                }
-            }
-        }, reload);
-    }
-
-    private void setUpColorSectionView(List<ColorOption> wallpaperColorOptions,
-            List<ColorOption> presetColorOptions) {
-        int wallpaperOptionSize = wallpaperColorOptions.size();
-
-        List<ColorOption> subOptions = wallpaperColorOptions.subList(0,
-                Math.min(5, wallpaperOptionSize));
-        // add additional options based on preset colors if there are less than 5 wallpaper colors
-        List<ColorOption> additionalSubOptions = presetColorOptions.subList(0,
-                Math.min(Math.max(0, 5 - wallpaperOptionSize), presetColorOptions.size()));
-        subOptions.addAll(additionalSubOptions);
-
-        mColorSectionView.setOverflowOnClick(() -> {
-            mSectionNavigationController.navigateTo(new ColorPickerFragment());
-            return null;
-        });
-        mColorSectionView.setColorOptionOnClick(selectedOption -> {
-            if (mSelectedColor.equals(selectedOption)) {
-                return null;
-            }
-            mSelectedColor = (ColorOption) selectedOption;
-            // Post with delay for color option to run ripple.
-            new Handler().postDelayed(()-> applyColor(mSelectedColor), /* delayMillis= */ 100);
-            return null;
-        });
-        mColorSectionView.setItems(subOptions, mColorManager);
-    }
-
-    private ColorOption findActiveColorOption(List<ColorOption> wallpaperColorOptions,
-            List<ColorOption> presetColorOptions) {
-        ColorOption activeColorOption = null;
-        for (ColorOption colorOption : Lists.newArrayList(
-                Iterables.concat(wallpaperColorOptions, presetColorOptions))) {
-            if (colorOption.isActive(mColorManager)) {
-                activeColorOption = colorOption;
-                break;
-            }
-        }
-        // Use the first one option by default. This should not happen as above should have an
-        // active option found.
-        if (activeColorOption == null) {
-            activeColorOption = wallpaperColorOptions.isEmpty()
-                    ? presetColorOptions.get(0)
-                    : wallpaperColorOptions.get(0);
-        }
-        return activeColorOption;
-    }
-
-    private void applyColor(ColorOption colorOption) {
-        if (SystemClock.elapsedRealtime() - mLastColorApplyingTime < MIN_COLOR_APPLY_PERIOD) {
-            return;
-        }
-        mLastColorApplyingTime = SystemClock.elapsedRealtime();
-        mColorManager.apply(colorOption, new CustomizationManager.Callback() {
-            @Override
-            public void onSuccess() {
-                mColorSectionView.announceForAccessibility(
-                        mColorSectionView.getContext().getString(R.string.color_changed));
-                mEventLogger.logColorApplied(getColorAction(colorOption), colorOption);
-            }
-
-            @Override
-            public void onError(@Nullable Throwable throwable) {
-                Log.w(TAG, "Apply theme with error: " + throwable);
-            }
-        });
-    }
-
-    private int getColorAction(ColorOption colorOption) {
-        int action = StyleEnums.DEFAULT_ACTION;
-        boolean isForBoth = mLockWallpaperColors == null || mLockWallpaperColors.equals(
-                mHomeWallpaperColors);
-
-        if (TextUtils.equals(colorOption.getSource(), COLOR_SOURCE_PRESET)) {
-            action = StyleEnums.COLOR_PRESET_APPLIED;
-        } else if (isForBoth) {
-            action = StyleEnums.COLOR_WALLPAPER_HOME_LOCK_APPLIED;
-        } else {
-            switch (colorOption.getSource()) {
-                case COLOR_SOURCE_HOME:
-                    action = StyleEnums.COLOR_WALLPAPER_HOME_APPLIED;
-                    break;
-                case COLOR_SOURCE_LOCK:
-                    action = StyleEnums.COLOR_WALLPAPER_LOCK_APPLIED;
-                    break;
-            }
-        }
-        return action;
-    }
-}
diff --git a/src/com/android/customization/module/DefaultCustomizationSections.java b/src/com/android/customization/module/DefaultCustomizationSections.java
index ebe1b17..487bca8 100644
--- a/src/com/android/customization/module/DefaultCustomizationSections.java
+++ b/src/com/android/customization/module/DefaultCustomizationSections.java
@@ -8,7 +8,6 @@
 import androidx.lifecycle.ViewModelProvider;
 
 import com.android.customization.model.color.ColorSectionController;
-import com.android.customization.model.color.ColorSectionController2;
 import com.android.customization.model.grid.GridOptionsManager;
 import com.android.customization.model.grid.GridSectionController;
 import com.android.customization.model.mode.DarkModeSectionController;
@@ -18,6 +17,8 @@
 import com.android.customization.model.themedicon.domain.interactor.ThemedIconInteractor;
 import com.android.customization.model.themedicon.domain.interactor.ThemedIconSnapshotRestorer;
 import com.android.customization.picker.clock.data.repository.ClockRegistryProvider;
+import com.android.customization.picker.color.ui.section.ColorSectionController2;
+import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel;
 import com.android.customization.picker.notifications.ui.section.NotificationSectionController;
 import com.android.customization.picker.notifications.ui.viewmodel.NotificationSectionViewModel;
 import com.android.customization.picker.preview.ui.section.PreviewWithClockCarouselSectionController;
@@ -48,6 +49,7 @@
 /** {@link CustomizationSections} for the customization picker. */
 public final class DefaultCustomizationSections implements CustomizationSections {
 
+    private final ColorPickerViewModel.Factory mColorPickerViewModelFactory;
     private final KeyguardQuickAffordancePickerInteractor mKeyguardQuickAffordancePickerInteractor;
     private final KeyguardQuickAffordancePickerViewModel.Factory
             mKeyguardQuickAffordancePickerViewModelFactory;
@@ -63,6 +65,7 @@
     private final ThemedIconInteractor mThemedIconInteractor;
 
     public DefaultCustomizationSections(
+            ColorPickerViewModel.Factory colorPickerViewModelFactory,
             KeyguardQuickAffordancePickerInteractor keyguardQuickAffordancePickerInteractor,
             KeyguardQuickAffordancePickerViewModel.Factory
                     keyguardQuickAffordancePickerViewModelFactory,
@@ -74,6 +77,7 @@
             DarkModeSnapshotRestorer darkModeSnapshotRestorer,
             ThemedIconSnapshotRestorer themedIconSnapshotRestorer,
             ThemedIconInteractor themedIconInteractor) {
+        mColorPickerViewModelFactory = colorPickerViewModelFactory;
         mKeyguardQuickAffordancePickerInteractor = keyguardQuickAffordancePickerInteractor;
         mKeyguardQuickAffordancePickerViewModelFactory =
                 keyguardQuickAffordancePickerViewModelFactory;
@@ -127,10 +131,12 @@
                 new ConnectedSectionController(
                         // Theme color section.
                         new ColorSectionController2(
-                                activity,
-                                wallpaperColorsViewModel,
-                                lifecycleOwner,
-                                sectionNavigationController),
+                                sectionNavigationController,
+                                new ViewModelProvider(
+                                        activity,
+                                        mColorPickerViewModelFactory)
+                                        .get(ColorPickerViewModel.class),
+                                lifecycleOwner),
                         // Wallpaper quick switch section.
                         new WallpaperQuickSwitchSectionController(
                                 screen,
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index 1a8fe6f..56ef2e6 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -23,6 +23,7 @@
 import androidx.activity.ComponentActivity
 import androidx.fragment.app.Fragment
 import androidx.fragment.app.FragmentActivity
+import androidx.lifecycle.ViewModelProvider
 import com.android.customization.model.mode.DarkModeSnapshotRestorer
 import com.android.customization.model.theme.OverlayManagerCompat
 import com.android.customization.model.theme.ThemeBundleProvider
@@ -99,6 +100,11 @@
     override fun getCustomizationSections(activity: ComponentActivity): CustomizationSections {
         return customizationSections
             ?: DefaultCustomizationSections(
+                    getColorPickerViewModelFactory(
+                        context = activity,
+                        wallpaperColorsViewModel =
+                            ViewModelProvider(activity)[WallpaperColorsViewModel::class.java],
+                    ),
                     getKeyguardQuickAffordancePickerInteractor(activity),
                     getKeyguardQuickAffordancePickerViewModelFactory(activity),
                     NotificationSectionViewModel.Factory(
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 e163c15..4ce5ed9 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -31,8 +31,11 @@
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.suspendCancellableCoroutine
 
 // TODO (b/262924623): refactor to remove dependency on ColorCustomizationManager & ColorOption
+// TODO (b/268203200): Create test for ColorPickerRepositoryImpl
 class ColorPickerRepositoryImpl(
     context: Context,
     wallpaperColorsViewModel: WallpaperColorsViewModel,
@@ -48,31 +51,48 @@
     /** List of wallpaper and preset color options on the device, categorized by Color Type */
     override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
         combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
-            colorManager.setWallpaperColors(homeColors, lockColors)
-            val wallpaperColorOptions: MutableList<ColorOptionModel> = mutableListOf()
-            val presetColorOptions: MutableList<ColorOptionModel> = mutableListOf()
-            colorManager.fetchOptions(
-                object : CustomizationManager.OptionsFetchedListener<ColorOption?> {
-                    override fun onOptionsLoaded(options: MutableList<ColorOption?>?) {
-                        options?.forEach { option ->
-                            when (option) {
-                                is ColorSeedOption -> wallpaperColorOptions.add(option.toModel())
-                                is ColorBundle -> presetColorOptions.add(option.toModel())
+                homeColors to lockColors
+            }
+            .map { (homeColors, lockColors) ->
+                suspendCancellableCoroutine { continuation ->
+                    colorManager.setWallpaperColors(homeColors, lockColors)
+                    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) {
+                                        is ColorSeedOption ->
+                                            wallpaperColorOptions.add(option.toModel())
+                                        is ColorBundle -> presetColorOptions.add(option.toModel())
+                                    }
+                                }
+                                continuation.resumeWith(
+                                    Result.success(
+                                        mapOf(
+                                            ColorType.WALLPAPER_COLOR to wallpaperColorOptions,
+                                            ColorType.BASIC_COLOR to presetColorOptions
+                                        )
+                                    )
+                                )
                             }
-                        }
-                    }
 
-                    override fun onError(throwable: Throwable?) {
-                        Log.e("ColorPickerRepository", "Error loading theme bundles", throwable)
-                    }
-                },
-                /* reload= */ false
-            )
-            mapOf(
-                ColorType.WALLPAPER_COLOR to wallpaperColorOptions,
-                ColorType.BASIC_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 fun select(colorOptionModel: ColorOptionModel) {
         val colorOption: ColorOption = colorOptionModel.colorOption
@@ -82,7 +102,7 @@
                 override fun onSuccess() = Unit
 
                 override fun onError(throwable: Throwable?) {
-                    Log.w("ColorPickerRepository", "Apply theme with error", throwable)
+                    Log.w(TAG, "Apply theme with error", throwable)
                 }
             }
         )
@@ -94,4 +114,8 @@
             isSelected = isActive(colorManager),
         )
     }
+
+    companion object {
+        private const val TAG = "ColorPickerRepositoryImpl"
+    }
 }
diff --git a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
index 6d8b7dc..331d635 100644
--- a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
@@ -136,13 +136,33 @@
 
     override fun select(colorOptionModel: ColorOptionModel) {
         val colorOptions = _colorOptions.value
-        colorOptions[ColorType.WALLPAPER_COLOR]?.forEach {
-            it.isSelected = (it.testEquals(colorOptionModel))
+        val wallpaperColorOptions = colorOptions[ColorType.WALLPAPER_COLOR]!!
+        val newWallpaperColorOptions = buildList {
+            wallpaperColorOptions.forEach { option ->
+                add(
+                    ColorOptionModel(
+                        colorOption = option.colorOption,
+                        isSelected = option.testEquals(colorOptionModel),
+                    )
+                )
+            }
         }
-        colorOptions[ColorType.BASIC_COLOR]?.forEach {
-            it.isSelected = (it.testEquals(colorOptionModel))
+        val basicColorOptions = colorOptions[ColorType.BASIC_COLOR]!!
+        val newBasicColorOptions = buildList {
+            basicColorOptions.forEach { option ->
+                add(
+                    ColorOptionModel(
+                        colorOption = option.colorOption,
+                        isSelected = option.testEquals(colorOptionModel),
+                    )
+                )
+            }
         }
-        _colorOptions.value = colorOptions
+        _colorOptions.value =
+            mapOf(
+                ColorType.WALLPAPER_COLOR to newWallpaperColorOptions,
+                ColorType.BASIC_COLOR to newBasicColorOptions
+            )
     }
 
     private fun ColorOptionModel.testEquals(other: Any?): Boolean {
diff --git a/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt b/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt
new file mode 100644
index 0000000..0842870
--- /dev/null
+++ b/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt
@@ -0,0 +1,126 @@
+/*
+ * Copyright (C) 2023 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.ui.binder
+
+import android.graphics.BlendMode
+import android.graphics.BlendModeColorFilter
+import android.view.LayoutInflater
+import android.view.View
+import android.widget.ImageView
+import android.widget.LinearLayout
+import androidx.core.view.isVisible
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.lifecycleScope
+import androidx.lifecycle.repeatOnLifecycle
+import com.android.customization.picker.color.ui.viewmodel.ColorOptionViewModel
+import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
+import com.android.wallpaper.R
+import kotlinx.coroutines.launch
+
+object ColorSectionViewBinder {
+
+    /**
+     * Binds view with view-model for color picker section. The view should include a linear layout
+     * with id [R.id.color_section_option_container]
+     */
+    @JvmStatic
+    fun bind(
+        view: View,
+        viewModel: ColorPickerViewModel,
+        lifecycleOwner: LifecycleOwner,
+        navigationOnClick: (View) -> Unit,
+        isConnectedHorizontallyToOtherSections: Boolean = false,
+    ) {
+        val optionContainer: LinearLayout =
+            view.requireViewById(R.id.color_section_option_container)
+        val moreColorsButton: View = view.requireViewById(R.id.more_colors)
+        if (isConnectedHorizontallyToOtherSections) {
+            moreColorsButton.isVisible = true
+            moreColorsButton.setOnClickListener(navigationOnClick)
+        } else {
+            moreColorsButton.isVisible = false
+        }
+        lifecycleOwner.lifecycleScope.launch {
+            lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                launch {
+                    viewModel.colorSectionOptions.collect { colorOptions ->
+                        setOptions(
+                            options = colorOptions,
+                            view = optionContainer,
+                            addOverflowOption = !isConnectedHorizontallyToOtherSections,
+                            overflowOnClick = navigationOnClick
+                        )
+                    }
+                }
+            }
+        }
+    }
+
+    fun setOptions(
+        options: List<ColorOptionViewModel>,
+        view: LinearLayout,
+        addOverflowOption: Boolean = false,
+        overflowOnClick: (View) -> Unit = {}
+    ) {
+        view.removeAllViews()
+        // Color option slot size is the minimum between the color option size and the view column
+        // count. When having an overflow option, a slot is reserved for the overflow option.
+        val colorOptionSlotSize =
+            (if (addOverflowOption) {
+                    minOf(view.weightSum.toInt() - 1, options.size)
+                } else {
+                    minOf(view.weightSum.toInt(), options.size)
+                })
+                .let { if (it < 0) 0 else it }
+        options.subList(0, colorOptionSlotSize).forEach { item ->
+            val itemView =
+                LayoutInflater.from(view.context)
+                    .inflate(R.layout.color_option_no_background, view, false)
+
+            val color0View: ImageView = itemView.requireViewById(R.id.color_preview_0)
+            val color1View: ImageView = itemView.requireViewById(R.id.color_preview_1)
+            val color2View: ImageView = itemView.requireViewById(R.id.color_preview_2)
+            val color3View: ImageView = itemView.requireViewById(R.id.color_preview_3)
+            color0View.drawable.colorFilter = BlendModeColorFilter(item.color0, BlendMode.SRC)
+            color1View.drawable.colorFilter = BlendModeColorFilter(item.color1, BlendMode.SRC)
+            color2View.drawable.colorFilter = BlendModeColorFilter(item.color2, BlendMode.SRC)
+            color3View.drawable.colorFilter = BlendModeColorFilter(item.color3, BlendMode.SRC)
+
+            val optionSelectedView = itemView.findViewById<ImageView>(R.id.option_selected)
+            optionSelectedView.isVisible = item.isSelected
+
+            itemView.setOnClickListener(
+                if (item.onClick != null) {
+                    View.OnClickListener { item.onClick.invoke() }
+                } else {
+                    null
+                }
+            )
+            view.addView(itemView)
+        }
+        // add overflow option
+        if (addOverflowOption) {
+            val itemView =
+                LayoutInflater.from(view.context)
+                    .inflate(R.layout.color_option_overflow_no_background, view, false)
+            itemView.setOnClickListener(overflowOnClick)
+            view.addView(itemView)
+        }
+    }
+}
diff --git a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
index fad7def..0bc22f8 100644
--- a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
+++ b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
@@ -37,6 +37,12 @@
 
 @OptIn(ExperimentalCoroutinesApi::class)
 class ColorPickerFragment : AppbarFragment() {
+    companion object {
+        @JvmStatic
+        fun newInstance(): ColorPickerFragment {
+            return ColorPickerFragment()
+        }
+    }
 
     override fun onCreateView(
         inflater: LayoutInflater,
diff --git a/src/com/android/customization/picker/color/ui/section/ColorSectionController2.kt b/src/com/android/customization/picker/color/ui/section/ColorSectionController2.kt
new file mode 100644
index 0000000..f1c982b
--- /dev/null
+++ b/src/com/android/customization/picker/color/ui/section/ColorSectionController2.kt
@@ -0,0 +1,67 @@
+/*
+ * Copyright (C) 2023 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.ui.section
+
+import android.content.Context
+import android.view.LayoutInflater
+import androidx.lifecycle.LifecycleOwner
+import com.android.customization.picker.color.ui.binder.ColorSectionViewBinder
+import com.android.customization.picker.color.ui.fragment.ColorPickerFragment
+import com.android.customization.picker.color.ui.view.ColorSectionView2
+import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
+import com.android.wallpaper.R
+import com.android.wallpaper.model.CustomizationSectionController
+import com.android.wallpaper.model.CustomizationSectionController.CustomizationSectionNavigationController as NavigationController
+
+class ColorSectionController2(
+    private val navigationController: NavigationController,
+    private val viewModel: ColorPickerViewModel,
+    private val lifecycleOwner: LifecycleOwner
+) : CustomizationSectionController<ColorSectionView2> {
+
+    override fun isAvailable(context: Context): Boolean {
+        return true
+    }
+
+    override fun createView(context: Context): ColorSectionView2 {
+        return createView(context, CustomizationSectionController.ViewCreationParams())
+    }
+
+    override fun createView(
+        context: Context,
+        params: CustomizationSectionController.ViewCreationParams
+    ): ColorSectionView2 {
+        @SuppressWarnings("It is fine to inflate with null parent for our need.")
+        val view =
+            LayoutInflater.from(context)
+                .inflate(
+                    R.layout.color_section_view2,
+                    null,
+                ) as ColorSectionView2
+        ColorSectionViewBinder.bind(
+            view = view,
+            viewModel = viewModel,
+            lifecycleOwner = lifecycleOwner,
+            navigationOnClick = {
+                navigationController.navigateTo(ColorPickerFragment.newInstance())
+            },
+            isConnectedHorizontallyToOtherSections = params.isConnectedHorizontallyToOtherSections,
+        )
+        return view
+    }
+}
diff --git a/src/com/android/customization/picker/color/ui/view/ColorSectionView2.kt b/src/com/android/customization/picker/color/ui/view/ColorSectionView2.kt
index 358514e..7a8f21a 100644
--- a/src/com/android/customization/picker/color/ui/view/ColorSectionView2.kt
+++ b/src/com/android/customization/picker/color/ui/view/ColorSectionView2.kt
@@ -17,77 +17,10 @@
 
 import android.content.Context
 import android.util.AttributeSet
-import android.view.LayoutInflater
-import android.widget.ImageView
-import android.widget.LinearLayout
-import com.android.customization.model.color.ColorCustomizationManager
-import com.android.customization.model.color.ColorOption
-import com.android.wallpaper.R
 import com.android.wallpaper.picker.SectionView
 
 /**
  * The class inherits from {@link SectionView} as the view representing the color section of the
  * customization picker. It displays a list of color options and an overflow option.
  */
-class ColorSectionView2(context: Context, attrs: AttributeSet?) : SectionView(context, attrs) {
-
-    private val items = mutableListOf<ColorOption>()
-    private var onClick: ((ColorOption) -> Unit)? = null
-    private var overflowOnClick: (() -> Unit)? = null
-
-    // TODO (b/262924623): make adjustments for large screen
-    fun setItems(items: List<ColorOption>, manager: ColorCustomizationManager) {
-        this.items.clear()
-        this.items.addAll(items)
-        val optionContainer = findViewById<LinearLayout>(R.id.color_section_option_container)
-        optionContainer.removeAllViews()
-        // Last color option is either the last index of the items list, or the second last index
-        // of column count. Save the last column of the option container for the overflow option
-        val lastOptionIndex = minOf(optionContainer.weightSum.toInt() - 2, items.size - 1)
-        if (items.isNotEmpty()) {
-            for (position in 0..lastOptionIndex) {
-                val item = items[position]
-                val itemView =
-                    LayoutInflater.from(context)
-                        .inflate(R.layout.color_option_no_background, optionContainer, false)
-                item.bindThumbnailTile(itemView.findViewById(R.id.option_tile))
-                if (item.isActive(manager)) {
-                    val optionSelectedView = itemView.findViewById<ImageView>(R.id.option_selected)
-                    optionSelectedView.visibility = VISIBLE
-                }
-                itemView.setOnClickListener { onClick?.invoke(item) }
-                optionContainer.addView(itemView)
-            }
-        }
-        // add overflow option
-        val itemView =
-            LayoutInflater.from(context)
-                .inflate(R.layout.color_option_overflow_no_background, optionContainer, false)
-        itemView.setOnClickListener { overflowOnClick?.invoke() }
-        optionContainer.addView(itemView)
-    }
-
-    /** Sets the on click callback for a color option. */
-    fun setColorOptionOnClick(onClick: (ColorOption) -> Unit) {
-        this.onClick = onClick
-        if (items.isNotEmpty()) {
-            val optionContainer = findViewById<LinearLayout>(R.id.color_section_option_container)
-            val lastOptionIndex = minOf(optionContainer.childCount - 2, items.size - 1)
-            for (position in 0..lastOptionIndex) {
-                val item = items[position]
-                val itemView = optionContainer.getChildAt(position)
-                itemView.setOnClickListener { onClick.invoke(item) }
-            }
-        }
-    }
-
-    /** Sets the on click callback for the overflow option. */
-    fun setOverflowOnClick(onClick: () -> Unit) {
-        this.overflowOnClick = onClick
-        val optionContainer = findViewById<LinearLayout>(R.id.color_section_option_container)
-        if (optionContainer.childCount > 0) {
-            val itemView = optionContainer.getChildAt(optionContainer.childCount - 1)
-            itemView.setOnClickListener { onClick.invoke() }
-        }
-    }
-}
+class ColorSectionView2(context: Context, attrs: AttributeSet?) : SectionView(context, attrs)
diff --git a/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
index 15445fa..7eb5488 100644
--- a/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
+++ b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
@@ -22,12 +22,14 @@
 import com.android.customization.model.color.ColorBundle
 import com.android.customization.model.color.ColorSeedOption
 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.wallpaper.R
+import kotlin.math.max
+import kotlin.math.min
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.map
 
 /** Models UI state for a color picker experience. */
 class ColorPickerViewModel
@@ -70,62 +72,85 @@
                 .toMap()
         }
 
+    /** The list of all available wallpaper colors */
+    private val wallpaperColorOptions: Flow<List<ColorOptionViewModel>> =
+        interactor.colorOptions.map { colorOptions ->
+            colorOptions[ColorType.WALLPAPER_COLOR]!!.map { colorOptionModel ->
+                val colorSeedOption: ColorSeedOption =
+                    colorOptionModel.colorOption as ColorSeedOption
+                val colors = colorSeedOption.previewInfo.resolveColors(context.resources)
+                ColorOptionViewModel(
+                    color0 = colors[0],
+                    color1 = colors[1],
+                    color2 = colors[2],
+                    color3 = colors[3],
+                    contentDescription = colorSeedOption.getContentDescription(context).toString(),
+                    isSelected = colorOptionModel.isSelected,
+                    onClick =
+                        if (colorOptionModel.isSelected) {
+                            null
+                        } else {
+                            { interactor.select(colorOptionModel) }
+                        }
+                )
+            }
+        }
+
+    /** The list of all available preset colors */
+    private val presetColorOptions: Flow<List<ColorOptionViewModel>> =
+        interactor.colorOptions.map { colorOptions ->
+            colorOptions[ColorType.BASIC_COLOR]!!.map { colorOptionModel ->
+                val colorBundle: ColorBundle = colorOptionModel.colorOption as ColorBundle
+                val primaryColor = colorBundle.previewInfo.resolvePrimaryColor(context.resources)
+                val secondaryColor =
+                    colorBundle.previewInfo.resolveSecondaryColor(context.resources)
+                ColorOptionViewModel(
+                    color0 = primaryColor,
+                    color1 = secondaryColor,
+                    color2 = primaryColor,
+                    color3 = secondaryColor,
+                    contentDescription = colorBundle.getContentDescription(context).toString(),
+                    isSelected = colorOptionModel.isSelected,
+                    onClick =
+                        if (colorOptionModel.isSelected) {
+                            null
+                        } else {
+                            { interactor.select(colorOptionModel) }
+                        },
+                )
+            }
+        }
+
     /** The list of all available color options for the selected Color Type. */
     val colorOptions: Flow<List<ColorOptionViewModel>> =
-        combine(interactor.colorOptions, selectedColorTypeId) {
-            colorOptions,
+        combine(wallpaperColorOptions, presetColorOptions, selectedColorTypeId) {
+            wallpaperOptions,
+            presetOptions,
             selectedColorTypeIdOrNull ->
-            val selectedColorType: ColorType =
-                selectedColorTypeIdOrNull ?: ColorType.WALLPAPER_COLOR
-            val selectedColorOptions: List<ColorOptionModel> = colorOptions[selectedColorType]!!
-            selectedColorOptions.map { colorOptionModel ->
-                when (selectedColorType) {
-                    ColorType.BASIC_COLOR -> {
-                        val colorBundle: ColorBundle = colorOptionModel.colorOption as ColorBundle
-                        val primaryColor =
-                            colorBundle.previewInfo.resolvePrimaryColor(context.resources)
-                        val secondaryColor =
-                            colorBundle.previewInfo.resolveSecondaryColor(context.resources)
-                        ColorOptionViewModel(
-                            color0 = primaryColor,
-                            color1 = secondaryColor,
-                            color2 = primaryColor,
-                            color3 = secondaryColor,
-                            contentDescription =
-                                colorBundle.getContentDescription(context).toString(),
-                            isSelected = colorOptionModel.isSelected,
-                            onClick =
-                                if (colorOptionModel.isSelected) {
-                                    null
-                                } else {
-                                    { interactor.select(colorOptionModel) }
-                                },
-                        )
-                    }
-                    ColorType.WALLPAPER_COLOR -> {
-                        val colorSeedOption: ColorSeedOption =
-                            colorOptionModel.colorOption as ColorSeedOption
-                        val colors = colorSeedOption.previewInfo.resolveColors(context.resources)
-                        ColorOptionViewModel(
-                            color0 = colors[0],
-                            color1 = colors[1],
-                            color2 = colors[2],
-                            color3 = colors[3],
-                            contentDescription =
-                                colorSeedOption.getContentDescription(context).toString(),
-                            isSelected = colorOptionModel.isSelected,
-                            onClick =
-                                if (colorOptionModel.isSelected) {
-                                    null
-                                } else {
-                                    { interactor.select(colorOptionModel) }
-                                },
-                        )
-                    }
-                }
+            when (selectedColorTypeIdOrNull ?: ColorType.WALLPAPER_COLOR) {
+                ColorType.WALLPAPER_COLOR -> wallpaperOptions
+                ColorType.BASIC_COLOR -> presetOptions
             }
         }
 
+    /** The list of color options for the color section */
+    val colorSectionOptions: Flow<List<ColorOptionViewModel>> =
+        combine(wallpaperColorOptions, presetColorOptions) { wallpaperOptions, presetOptions ->
+            val subOptions =
+                wallpaperOptions.subList(0, min(COLOR_SECTION_OPTION_SIZE, wallpaperOptions.size))
+            // Add additional options based on preset colors if size of wallpaper color options is
+            // less than COLOR_SECTION_OPTION_SIZE
+            val additionalSubOptions =
+                presetOptions.subList(
+                    0,
+                    min(
+                        max(0, COLOR_SECTION_OPTION_SIZE - wallpaperOptions.size),
+                        presetOptions.size,
+                    )
+                )
+            subOptions + additionalSubOptions
+        }
+
     class Factory(
         private val context: Context,
         private val interactor: ColorPickerInteractor,
@@ -139,4 +164,8 @@
                 as T
         }
     }
+
+    companion object {
+        private const val COLOR_SECTION_OPTION_SIZE = 5
+    }
 }
diff --git a/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt b/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
index 19d5dd1..6e5f776 100644
--- a/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
+++ b/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
@@ -59,6 +59,19 @@
     }
 
     @Test
+    fun `Select a color section color`() = runTest {
+        val colorSectionOptions = collectLastValue(underTest.colorSectionOptions)
+
+        assertColorOptionUiState(colorOptions = colorSectionOptions(), selectedColorOptionIndex = 0)
+
+        colorSectionOptions()?.get(2)?.onClick?.invoke()
+        assertColorOptionUiState(colorOptions = colorSectionOptions(), selectedColorOptionIndex = 2)
+
+        colorSectionOptions()?.get(4)?.onClick?.invoke()
+        assertColorOptionUiState(colorOptions = colorSectionOptions(), selectedColorOptionIndex = 4)
+    }
+
+    @Test
     fun `Select a preset color`() = runTest {
         val colorTypes = collectLastValue(underTest.colorTypes)
         val colorOptions = collectLastValue(underTest.colorOptions)
@@ -128,7 +141,20 @@
             colorTypeId = ColorType.BASIC_COLOR,
             isSelected = "Basic colors" == selectedColorTypeText,
         )
+        assertColorOptionUiState(colorOptions, selectedColorOptionIndex)
+    }
 
+    /**
+     * Asserts the picker section UI state is what is expected.
+     *
+     * @param colorOptions The observed color options
+     * @param selectedColorOptionIndex The index of the color option that's expected to be selected,
+     * -1 stands for no color option should be selected
+     */
+    private fun assertColorOptionUiState(
+        colorOptions: List<ColorOptionViewModel>?,
+        selectedColorOptionIndex: Int,
+    ) {
         var foundSelectedColorOption = false
         assertThat(colorOptions).isNotNull()
         if (colorOptions != null) {