Merge "Make ThemePicker override any other AOSP picker" into main am: 09461d20e3 am: 98be76a984 am: cca5353b72 am: 0a5a1d4ca6

Original change: https://android-review.googlesource.com/c/platform/packages/apps/ThemePicker/+/2706335

Change-Id: I835d974791b5f869539c6f537277ff1f0ead1480
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..bff2d76
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+*.iml
diff --git a/Android.bp b/Android.bp
index b7c0d00..c53ec8d 100644
--- a/Android.bp
+++ b/Android.bp
@@ -82,6 +82,7 @@
         "androidx.lifecycle_lifecycle-viewmodel-ktx",
         "androidx.recyclerview_recyclerview",
         "SystemUICustomizationLib",
+        "hilt_android",
     ],
 
     jni_libs: [
diff --git a/res/drawable/carousel_item_card_background.xml b/res/drawable/carousel_item_card_background.xml
index 946d279..4197698 100644
--- a/res/drawable/carousel_item_card_background.xml
+++ b/res/drawable/carousel_item_card_background.xml
@@ -16,6 +16,6 @@
 -->
 <shape xmlns:android="http://schemas.android.com/apk/res/android"
     android:shape="rectangle">
-    <solid android:color="@color/color_surface" />
+    <solid android:color="@color/system_surface_container_highest" />
     <corners android:radius="16dp" />
 </shape>
\ No newline at end of file
diff --git a/res/drawable/clock_color_size_button_background.xml b/res/drawable/clock_color_size_button_background.xml
index 08d1c45..09b43c8 100644
--- a/res/drawable/clock_color_size_button_background.xml
+++ b/res/drawable/clock_color_size_button_background.xml
@@ -19,7 +19,7 @@
     <item android:id="@android:id/mask">
         <shape android:shape="rectangle">
             <corners android:radius="@dimen/clock_color_size_button_corner_radius" />
-            <solid android:color="?android:colorControlHighlight" />
+            <solid android:color="@color/clock_color_and_size_button_background" />
         </shape>
     </item>
 
diff --git a/res/drawable/picker_section_icon_background.xml b/res/drawable/picker_section_icon_background.xml
new file mode 100644
index 0000000..2a124e0
--- /dev/null
+++ b/res/drawable/picker_section_icon_background.xml
@@ -0,0 +1,30 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+     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.
+-->
+<ripple xmlns:android="http://schemas.android.com/apk/res/android"
+    android:color="?android:attr/colorControlHighlight">
+    <item android:id="@android:id/mask">
+        <shape android:shape="rectangle">
+            <solid android:color="@android:color/white"/>
+            <corners android:radius="@dimen/option_tile_radius" />
+        </shape>
+    </item>
+    <item>
+        <shape android:shape="rectangle">
+            <solid android:color="@color/picker_section_icon_background"/>
+            <corners android:radius="@dimen/option_tile_radius" />
+        </shape>
+    </item>
+</ripple>
diff --git a/res/drawable/top_connected_section_background.xml b/res/drawable/top_connected_section_background.xml
deleted file mode 100644
index f96824b..0000000
--- a/res/drawable/top_connected_section_background.xml
+++ /dev/null
@@ -1,29 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?><!--
-  ~ 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.
-  ~
-  -->
-
-<shape
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:shape="rectangle">
-
-    <solid android:color="@color/color_surface" />
-
-    <corners
-        android:topLeftRadius="32dp"
-        android:topRightRadius="32dp"
-        android:bottomLeftRadius="8dp"
-        android:bottomRightRadius="8dp" />
-</shape>
\ No newline at end of file
diff --git a/res/layout/clock_size_radio_button_group.xml b/res/layout/clock_size_radio_button_group.xml
index 2aa78d2..56d0ab7 100644
--- a/res/layout/clock_size_radio_button_group.xml
+++ b/res/layout/clock_size_radio_button_group.xml
@@ -68,11 +68,21 @@
             android:layout_marginEnd="8dp"
             android:clickable="false" />
 
-        <TextView
-            style="@style/SectionTitleTextStyle"
-            android:layout_width="wrap_content"
+        <LinearLayout
+            android:layout_width="match_parent"
             android:layout_height="wrap_content"
-            android:layout_gravity="center_vertical"
-            android:text="@string/clock_size_small" />
+            android:orientation="vertical">
+            <TextView
+                style="@style/SectionTitleTextStyle"
+                android:layout_width="wrap_content"
+                android:layout_height="wrap_content"
+                android:text="@string/clock_size_small" />
+
+            <TextView
+                style="@style/SectionSubtitleTextStyle"
+                android:layout_width="wrap_content"
+                android:layout_height="wrap_content"
+                android:text="@string/clock_size_small_description" />
+        </LinearLayout>
     </LinearLayout>
 </LinearLayout>
\ No newline at end of file
diff --git a/res/layout/color_section_view.xml b/res/layout/color_section_view.xml
deleted file mode 100644
index 6095c9b..0000000
--- a/res/layout/color_section_view.xml
+++ /dev/null
@@ -1,38 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-     Copyright (C) 2022 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.
--->
-<com.android.customization.picker.color.ColorSectionView
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="match_parent"
-    android:layout_height="wrap_content"
-    android:paddingBottom="@dimen/section_bottom_padding"
-    android:orientation="vertical">
-
-    <FrameLayout
-        android:layout_width="match_parent"
-        android:layout_height="wrap_content"
-        android:paddingHorizontal="@dimen/separated_tabs_horizontal_padding">
-        <include layout="@layout/separated_tabs" />
-    </FrameLayout>
-
-    <androidx.viewpager2.widget.ViewPager2
-        android:id="@+id/color_section_view_pager"
-        android:layout_width="match_parent"
-        android:layout_height="wrap_content"
-        android:layout_marginTop="@dimen/color_options_container_top_margin"
-        android:clipChildren="false"
-        android:clipToPadding="false"/>
-</com.android.customization.picker.color.ColorSectionView>
diff --git a/res/layout/fragment_clock_settings.xml b/res/layout/fragment_clock_settings.xml
index bb4a0c2..5736265 100644
--- a/res/layout/fragment_clock_settings.xml
+++ b/res/layout/fragment_clock_settings.xml
@@ -19,7 +19,8 @@
     android:layout_width="match_parent"
     android:layout_height="match_parent"
     android:orientation="vertical"
-    android:clipChildren="false">
+    android:clipChildren="false"
+    android:transitionGroup="true">
 
     <FrameLayout
         android:id="@+id/section_header_container"
diff --git a/res/layout/fragment_lock_screen_quick_affordances.xml b/res/layout/fragment_lock_screen_quick_affordances.xml
index 36d1697..ee259f1 100644
--- a/res/layout/fragment_lock_screen_quick_affordances.xml
+++ b/res/layout/fragment_lock_screen_quick_affordances.xml
@@ -19,7 +19,8 @@
     android:layout_width="match_parent"
     android:layout_height="match_parent"
     android:orientation="vertical"
-    android:clipChildren="false">
+    android:clipChildren="false"
+    android:transitionGroup="true">
 
     <FrameLayout
         android:id="@+id/section_header_container"
diff --git a/res/layout/grid_section_view.xml b/res/layout/grid_section_view.xml
index a90ee64..117691c 100644
--- a/res/layout/grid_section_view.xml
+++ b/res/layout/grid_section_view.xml
@@ -49,7 +49,7 @@
         android:layout_height="@dimen/option_tile_width"
         android:scaleType="center"
         android:src="@drawable/ic_grid_24px"
-        android:background="@drawable/option_border_color"
+        android:background="@drawable/picker_section_icon_background"
         android:contentDescription="@string/gird_picker_entry_content_description" />
 
 </com.android.customization.picker.grid.GridSectionView>
\ No newline at end of file
diff --git a/res/layout/keyguard_quick_affordance_section_view.xml b/res/layout/keyguard_quick_affordance_section_view.xml
index 7c5babb..d577060 100644
--- a/res/layout/keyguard_quick_affordance_section_view.xml
+++ b/res/layout/keyguard_quick_affordance_section_view.xml
@@ -48,7 +48,7 @@
         android:layout_width="@dimen/option_tile_width"
         android:layout_height="@dimen/option_tile_width"
         android:orientation="horizontal"
-        android:background="@drawable/option_border_color"
+        android:background="@drawable/picker_section_icon_background"
         android:importantForAccessibility="noHideDescendants"
         android:gravity="center"
         android:divider="@drawable/horizontal_divider_14dp"
diff --git a/res/layout/single_clock_view.xml b/res/layout/single_clock_view.xml
deleted file mode 100644
index ffc6039..0000000
--- a/res/layout/single_clock_view.xml
+++ /dev/null
@@ -1,27 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?><!--
-     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.
--->
-<com.android.wallpaper.picker.FixedWidthDisplayRatioFrameLayout
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="@dimen/screen_preview_width"
-    android:layout_height="wrap_content"
-    android:layout_gravity="center"
-    android:clipChildren="false">
-    <com.android.customization.picker.clock.ui.view.ClockHostView
-        android:id="@+id/single_clock_host_view"
-        android:layout_width="match_parent"
-        android:layout_height="match_parent"
-        android:layout_gravity="center" />
-</com.android.wallpaper.picker.FixedWidthDisplayRatioFrameLayout>
\ No newline at end of file
diff --git a/res/values-h800dp/dimens.xml b/res/values-h800dp/dimens.xml
deleted file mode 100644
index 36d0d62..0000000
--- a/res/values-h800dp/dimens.xml
+++ /dev/null
@@ -1,9 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-
-<resources>
-    <!--
-    Dimension for clock translationY when swiping clock carousel in picker,
-    copied from sysui resources
-    -->
-    <dimen name="keyguard_large_clock_top_margin">-112dp</dimen>
-</resources>
\ No newline at end of file
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index a3e0d4a..47171e5 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -170,14 +170,4 @@
     <dimen name="clock_color_size_button_corner_radius">16dp</dimen>
     <dimen name="clock_color_size_button_icon_size">20dp</dimen>
     <dimen name="clock_color_size_button_icon_margin_end">12dp</dimen>
-
-    <!--
-    Dimension for clock translationY when swiping clock carousel in picker,
-    copied from sysui resources
-    -->
-    <dimen name="keyguard_large_clock_top_margin">-60dp</dimen>
-    <!-- Dimension for the clock view, copied from sysui resources. -->
-    <dimen name="small_clock_height">114dp</dimen>
-    <dimen name="small_clock_padding_top">28dp</dimen>
-    <dimen name="clock_padding_start">28dp</dimen>
 </resources>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 680947d..eb68dce 100755
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -90,6 +90,9 @@
     <!-- Title of a radio button to apply clock size small. [CHAR LIMIT=15] -->
     <string name="clock_size_small">Small</string>
 
+    <!-- Description of a radio button to apply clock size small. [CHAR LIMIT=NONE] -->
+    <string name="clock_size_small_description">A small clock shows in the corner of your screen</string>
+
     <!-- Title of a section of the customization picker where the user can select a Grid size for
         the home screen. [CHAR LIMIT=15] -->
     <string name="grid_title">App grid</string>
diff --git a/res/values/styles.xml b/res/values/styles.xml
index c70b596..b2c22b2 100644
--- a/res/values/styles.xml
+++ b/res/values/styles.xml
@@ -48,6 +48,7 @@
         <item name="colorSurface">@color/design_default_color_surface</item>
         <item name="colorOnSurface">@color/design_default_color_on_surface</item>
         <item name="android:textAllCaps">false</item>
+        <item name="motionDurationLong1">500</item>
     </style>
 
     <!-- Snackbar margin -->
diff --git a/src/com/android/customization/model/CustomizationManager.java b/src/com/android/customization/model/CustomizationManager.java
index 104cc83..e6d3872 100644
--- a/src/com/android/customization/model/CustomizationManager.java
+++ b/src/com/android/customization/model/CustomizationManager.java
@@ -72,6 +72,9 @@
      */
     void apply(T option, Callback callback);
 
+    /** Preview the given option without committing the change. */
+    default void preview(T option) {}
+
     /**
      * Loads the available options for the type of Customization managed by this class, calling the
      * given callback when done.
diff --git a/src/com/android/customization/model/color/ColorProvider.kt b/src/com/android/customization/model/color/ColorProvider.kt
index 63d298d..74a4051 100644
--- a/src/com/android/customization/model/color/ColorProvider.kt
+++ b/src/com/android/customization/model/color/ColorProvider.kt
@@ -16,6 +16,7 @@
 package com.android.customization.model.color
 
 import android.app.WallpaperColors
+import android.app.WallpaperManager
 import android.content.Context
 import android.content.res.ColorStateList
 import android.content.res.Resources
@@ -38,7 +39,6 @@
 import com.android.systemui.monet.ColorScheme
 import com.android.systemui.monet.Style
 import com.android.wallpaper.R
-import com.android.wallpaper.compat.WallpaperManagerCompat
 import com.android.wallpaper.module.InjectorProvider
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -130,9 +130,9 @@
     private fun isLockScreenWallpaperLastApplied(): Boolean {
         // The WallpaperId increases every time a new wallpaper is set, so the larger wallpaper id
         // is the most recently set wallpaper
-        val manager = InjectorProvider.getInjector().getWallpaperManagerCompat(mContext)
-        return manager.getWallpaperId(WallpaperManagerCompat.FLAG_LOCK) >
-            manager.getWallpaperId(WallpaperManagerCompat.FLAG_SYSTEM)
+        val manager = WallpaperManager.getInstance(mContext)
+        return manager.getWallpaperId(WallpaperManager.FLAG_LOCK) >
+            manager.getWallpaperId(WallpaperManager.FLAG_SYSTEM)
     }
 
     private fun loadSeedColors(
@@ -432,7 +432,8 @@
             val extractor = ColorBundlePreviewExtractor(mContext)
             val bundles: MutableList<ColorOption> = ArrayList()
 
-            val bundleNames = getItemsFromStub(COLOR_BUNDLES_ARRAY_NAME)
+            val bundleNames =
+                if (isAvailable) getItemsFromStub(COLOR_BUNDLES_ARRAY_NAME) else emptyArray()
             // Color option index value starts from 1.
             var index = 1
             val maxPresetColors = if (themeStyleEnabled) bundleNames.size else MAX_PRESET_COLORS
diff --git a/src/com/android/customization/model/color/ColorSectionController.java b/src/com/android/customization/model/color/ColorSectionController.java
deleted file mode 100644
index be051ac..0000000
--- a/src/com/android/customization/model/color/ColorSectionController.java
+++ /dev/null
@@ -1,538 +0,0 @@
-/*
- * Copyright (C) 2022 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 android.view.View.IMPORTANT_FOR_ACCESSIBILITY_NO;
-import static android.view.View.VISIBLE;
-
-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 static com.android.customization.widget.OptionSelectorController.CheckmarkStyle.CENTER;
-
-import android.app.Activity;
-import android.app.WallpaperColors;
-import android.content.Context;
-import android.os.Bundle;
-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 android.view.View;
-import android.view.ViewGroup;
-import android.view.accessibility.AccessibilityNodeInfo;
-import android.widget.FrameLayout;
-
-import androidx.annotation.Nullable;
-import androidx.lifecycle.LifecycleOwner;
-import androidx.recyclerview.widget.RecyclerView;
-import androidx.viewpager2.widget.ViewPager2;
-
-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.ColorSectionView;
-import com.android.customization.widget.OptionSelectorController;
-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.android.wallpaper.widget.PageIndicator;
-import com.android.wallpaper.widget.SeparatedTabLayout;
-
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Lists;
-
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Locale;
-import java.util.Optional;
-
-/**
- * Color section view's controller for the logic of color customization.
- */
-public class ColorSectionController implements CustomizationSectionController<ColorSectionView> {
-
-    private static final String TAG = "ColorSectionController";
-    private static final String KEY_COLOR_TAB_POSITION = "COLOR_TAB_POSITION";
-    private static final String KEY_COLOR_PAGE_POSITION = "COLOR_PAGE_POSITION";
-    private static final String ID_VIEWPAGER = "ColorSectionController_colorSectionViewPager";
-    private static final String ID_ITEMVIEW = "ColorSectionController_itemView";
-    private static final String ID_CONTAINER = "ColorSectionController_container";
-    private static final long MIN_COLOR_APPLY_PERIOD = 500L;
-
-    private static final int WALLPAPER_TAB_INDEX = 0;
-    private static final int PRESET_TAB_INDEX = 1;
-
-    private final ThemesUserEventLogger mEventLogger;
-    private final ColorCustomizationManager mColorManager;
-    private final WallpaperColorsViewModel mWallpaperColorsViewModel;
-    private final LifecycleOwner mLifecycleOwner;
-    private final ColorSectionAdapter mColorSectionAdapter = new ColorSectionAdapter();
-
-    private List<ColorOption> mWallpaperColorOptions = new ArrayList<>();
-    private List<ColorOption> mPresetColorOptions = new ArrayList<>();
-    private ViewPager2 mColorSectionViewPager;
-    private ColorOption mSelectedColor;
-    private SeparatedTabLayout mTabLayout;
-    @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 Optional<Integer> mTabPositionToRestore = Optional.empty();
-    private Optional<Integer>[] mPagePositionToRestore =
-            new Optional[]{Optional.empty(), Optional.empty()};
-    private long mLastColorApplyingTime = 0L;
-    private ColorSectionView mColorSectionView;
-
-    private static int getNumPages(int optionsPerPage, int totalOptions) {
-        return (int) Math.ceil((float) totalOptions / optionsPerPage);
-    }
-
-    public ColorSectionController(Activity activity, WallpaperColorsViewModel viewModel,
-            LifecycleOwner lifecycleOwner, @Nullable Bundle savedInstanceState) {
-        CustomizationInjector injector = (CustomizationInjector) InjectorProvider.getInjector();
-        mEventLogger = (ThemesUserEventLogger) injector.getUserEventLogger(activity);
-        mColorManager = ColorCustomizationManager.getInstance(activity,
-                new OverlayManagerCompat(activity));
-        mWallpaperColorsViewModel = viewModel;
-        mLifecycleOwner = lifecycleOwner;
-
-        if (savedInstanceState != null) {
-            if (savedInstanceState.containsKey(KEY_COLOR_TAB_POSITION)) {
-                mTabPositionToRestore = Optional.of(
-                        savedInstanceState.getInt(KEY_COLOR_TAB_POSITION));
-            }
-
-            for (int i = 0; i < mPagePositionToRestore.length; i++) {
-                String keyColorPage = getPagePositionKey(i);
-                if (keyColorPage != null && savedInstanceState.containsKey(keyColorPage)) {
-                    setPagePositionToRestore(i, savedInstanceState.getInt(keyColorPage));
-                }
-            }
-        }
-    }
-
-    private String getPagePositionKey(int index) {
-        return String.format(Locale.US, "%s_%d", KEY_COLOR_PAGE_POSITION, index);
-    }
-
-    private void setPagePositionToRestore(int pagePositionKeyIndex, int pagePosition) {
-        if (pagePositionKeyIndex >= 0 && pagePositionKeyIndex < mPagePositionToRestore.length) {
-            mPagePositionToRestore[pagePositionKeyIndex] = Optional.of(pagePosition);
-        }
-    }
-
-    private int getPagePositionToRestore(int pagePositionKeyIndex, int defaultPagePosition) {
-        if (pagePositionKeyIndex >= 0 && pagePositionKeyIndex < mPagePositionToRestore.length) {
-            return mPagePositionToRestore[pagePositionKeyIndex].orElse(defaultPagePosition);
-        }
-        return 0;
-    }
-
-    @Override
-    public boolean isAvailable(@Nullable Context context) {
-        return context != null && ColorUtils.isMonetEnabled(context) && mColorManager.isAvailable();
-    }
-
-    @Override
-    public ColorSectionView createView(Context context) {
-        mColorSectionView = (ColorSectionView) LayoutInflater.from(context).inflate(
-                R.layout.color_section_view, /* root= */ null);
-        mColorSectionViewPager = mColorSectionView.findViewById(R.id.color_section_view_pager);
-        mColorSectionViewPager.setAccessibilityDelegate(createAccessibilityDelegate(ID_VIEWPAGER));
-        mColorSectionViewPager.setAdapter(mColorSectionAdapter);
-        mColorSectionViewPager.setUserInputEnabled(false);
-        if (ColorProvider.themeStyleEnabled) {
-            mColorSectionViewPager.setImportantForAccessibility(IMPORTANT_FOR_ACCESSIBILITY_NO);
-        }
-        mTabLayout = mColorSectionView.findViewById(R.id.separated_tabs);
-        mColorSectionAdapter.setNumColors(context.getResources().getInteger(
-                R.integer.options_grid_num_columns));
-        // TODO(b/202145216): Use just 2 views when tapping either button on top.
-        mTabLayout.setViewPager(mColorSectionViewPager);
-
-        mWallpaperColorsViewModel.getHomeWallpaperColorsLiveData().observe(mLifecycleOwner,
-                homeColors -> {
-                    mHomeWallpaperColors = homeColors;
-                    mHomeWallpaperColorsReady = true;
-                    maybeLoadColors();
-                });
-        mWallpaperColorsViewModel.getLockWallpaperColorsLiveData().observe(mLifecycleOwner,
-                lockColors -> {
-                    mLockWallpaperColors = lockColors;
-                    mLockWallpaperColorsReady = true;
-                    maybeLoadColors();
-                });
-        return mColorSectionView;
-    }
-
-    @Override
-    public void onSaveInstanceState(Bundle savedInstanceState) {
-        if (mColorSectionViewPager != null) {
-            savedInstanceState.putInt(KEY_COLOR_TAB_POSITION,
-                    mColorSectionViewPager.getCurrentItem());
-
-            for (int i = 0; i < mPagePositionToRestore.length; i++) {
-                savedInstanceState.putInt(getPagePositionKey(i), getPagePositionToRestore(i, 0));
-            }
-        }
-    }
-
-    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);
-                mTabLayout.post(()-> setUpColorViewPager());
-            }
-
-            @Override
-            public void onError(@Nullable Throwable throwable) {
-                if (throwable != null) {
-                    Log.e(TAG, "Error loading theme bundles", throwable);
-                }
-            }
-        }, reload);
-    }
-
-    private void setUpColorViewPager() {
-        mColorSectionAdapter.notifyDataSetChanged();
-
-        if (mTabLayout != null && mTabLayout.getTabCount() == 0) {
-            mTabLayout.addTab(mTabLayout.newTab().setText(R.string.wallpaper_color_tab),
-                    WALLPAPER_TAB_INDEX);
-            mTabLayout.addTab(mTabLayout.newTab().setText(R.string.preset_color_tab),
-                    PRESET_TAB_INDEX);
-        }
-
-        if (mWallpaperColorOptions.isEmpty()) {
-            // Select preset tab and disable wallpaper tab.
-            mTabLayout.getTabAt(WALLPAPER_TAB_INDEX).view.setEnabled(false);
-            mColorSectionViewPager.setCurrentItem(PRESET_TAB_INDEX, /* smoothScroll= */ false);
-            return;
-        }
-
-        mColorSectionViewPager.setCurrentItem(
-                mTabPositionToRestore.orElseGet(
-                        () -> COLOR_SOURCE_PRESET.equals(mColorManager.getCurrentColorSource())
-                                ? PRESET_TAB_INDEX
-                                : WALLPAPER_TAB_INDEX),
-                /* smoothScroll= */ false);
-
-        // Disable "wallpaper colors" and "basic colors" swiping for new color style.
-        mColorSectionViewPager.setUserInputEnabled(!ColorProvider.themeStyleEnabled);
-    }
-
-    private void setupColorPages(ViewPager2 container, int colorsPerPage, int sectionPosition,
-            List<ColorOption> options, PageIndicator pageIndicator) {
-        container.setAdapter(new ColorPageAdapter(options, /* pageEnabled= */ true,
-                colorsPerPage));
-        if (ColorProvider.themeStyleEnabled) {
-            // Update page index to show selected items.
-            int selectedIndex = options.indexOf(mSelectedColor);
-            if (colorsPerPage != 0) {
-                int pageIndex = selectedIndex / colorsPerPage;
-                int position = getPagePositionToRestore(sectionPosition, pageIndex);
-                container.setCurrentItem(position, /* smoothScroll= */ false);
-            }
-            pageIndicator.setNumPages(getNumPages(colorsPerPage, options.size()));
-            registerOnPageChangeCallback(sectionPosition, container, pageIndicator);
-        }
-    }
-
-    private void registerOnPageChangeCallback(int sectionPosition, ViewPager2 container,
-            PageIndicator pageIndicator) {
-        container.registerOnPageChangeCallback(new ViewPager2.OnPageChangeCallback() {
-            @Override
-            public void onPageSelected(int position) {
-                super.onPageSelected(position);
-                if (mColorSectionViewPager.getCurrentItem() == sectionPosition) {
-                    pageIndicator.setLocation(getPagePosition(pageIndicator, position));
-                    setPagePositionToRestore(sectionPosition, position);
-                }
-            }
-
-            @Override
-            public void onPageScrolled(int position, float positionOffset,
-                    int positionOffsetPixels) {
-                super.onPageScrolled(position, positionOffset, positionOffsetPixels);
-                if (mColorSectionViewPager.getCurrentItem() == sectionPosition) {
-                    pageIndicator.setLocation(getPagePosition(pageIndicator, position));
-                    setPagePositionToRestore(sectionPosition, position);
-                }
-            }
-
-            private int getPagePosition(PageIndicator pageIndicator, int position) {
-                return pageIndicator.isLayoutRtl() ? pageIndicator.getChildCount() - 1 - position
-                        : position;
-            }
-        });
-    }
-
-    private void setupColorOptions(RecyclerView container, List<ColorOption> colorOptions,
-            boolean pageEnabled, int index, int colorsPerPage) {
-        int totalSize = colorOptions.size();
-        if (totalSize == 0) {
-            return;
-        }
-
-        List<ColorOption> subOptions;
-        if (pageEnabled && ColorProvider.themeStyleEnabled) {
-            subOptions = colorOptions.subList(colorsPerPage * index,
-                    Math.min(colorsPerPage * (index + 1), totalSize));
-        } else {
-            subOptions = colorOptions;
-        }
-
-        OptionSelectorController<ColorOption> adaptiveController = new OptionSelectorController<>(
-                container, subOptions, /* useGrid= */ true, CENTER);
-        adaptiveController.initOptions(mColorManager);
-        setUpColorOptionsController(adaptiveController);
-    }
-
-    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 setUpColorOptionsController(
-            OptionSelectorController<ColorOption> optionSelectorController) {
-        if (mSelectedColor != null && optionSelectorController.containsOption(mSelectedColor)) {
-            optionSelectorController.setSelectedOption(mSelectedColor);
-        }
-
-        optionSelectorController.addListener(selectedOption -> {
-            ColorOption selectedColor = (ColorOption) selectedOption;
-            if (mSelectedColor.equals(selectedColor)) {
-                return;
-            }
-            mSelectedColor = (ColorOption) selectedOption;
-            // Post with delay for color option to run ripple.
-            new Handler().postDelayed(()-> applyColor(mSelectedColor), /* delayMillis= */ 100);
-        });
-    }
-
-    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;
-    }
-
-    private View.AccessibilityDelegate createAccessibilityDelegate(String id) {
-        return new View.AccessibilityDelegate() {
-            @Override
-            public void onInitializeAccessibilityNodeInfo(View host, AccessibilityNodeInfo info) {
-                super.onInitializeAccessibilityNodeInfo(host, info);
-                info.setUniqueId(id);
-            }
-        };
-    }
-
-    private class ColorSectionAdapter extends
-            RecyclerView.Adapter<ColorSectionAdapter.ColorPageViewHolder> {
-
-        private final int mItemCounts = new int[]{WALLPAPER_TAB_INDEX, PRESET_TAB_INDEX}.length;
-        private int mNumColors;
-
-        @Override
-        public int getItemCount() {
-            return mItemCounts;
-        }
-
-        @Override
-        public void onBindViewHolder(ColorPageViewHolder viewHolder, int position) {
-            switch (position) {
-                case WALLPAPER_TAB_INDEX:
-                    setupColorPages(viewHolder.mContainer, mNumColors, position,
-                            mWallpaperColorOptions, viewHolder.mPageIndicator);
-                    break;
-                case PRESET_TAB_INDEX:
-                    setupColorPages(viewHolder.mContainer, mNumColors, position,
-                            mPresetColorOptions, viewHolder.mPageIndicator);
-                    break;
-                default:
-                    break;
-            }
-        }
-
-        @Override
-        public ColorPageViewHolder onCreateViewHolder(ViewGroup viewGroup, int viewType) {
-            return new ColorPageViewHolder(LayoutInflater.from(viewGroup.getContext()).inflate(
-                    viewType, viewGroup, false));
-        }
-
-        @Override
-        public int getItemViewType(int position) {
-            return R.layout.color_pages_view;
-        }
-
-        public void setNumColors(int numColors) {
-            mNumColors = numColors;
-        }
-
-        private class ColorPageViewHolder extends RecyclerView.ViewHolder {
-            private ViewPager2 mContainer;
-            private PageIndicator mPageIndicator;
-
-            ColorPageViewHolder(View itemView) {
-                super(itemView);
-                mContainer = itemView.findViewById(R.id.color_page_container);
-                // Correct scrolling goes under collapsing toolbar while scrolling oclor options.
-                mContainer.getChildAt(0).setNestedScrollingEnabled(false);
-                mPageIndicator = itemView.findViewById(R.id.color_page_indicator);
-                if (ColorProvider.themeStyleEnabled) {
-                    mPageIndicator.setVisibility(VISIBLE);
-                }
-                itemView.setAccessibilityDelegate(createAccessibilityDelegate(ID_ITEMVIEW));
-                mContainer.setAccessibilityDelegate(createAccessibilityDelegate(ID_CONTAINER));
-            }
-        }
-    }
-
-    private class ColorPageAdapter extends
-            RecyclerView.Adapter<ColorPageAdapter.ColorOptionViewHolder> {
-
-        private final boolean mPageEnabled;
-        private final List<ColorOption> mColorOptions;
-        private final int mColorsPerPage;
-
-        private ColorPageAdapter(List<ColorOption> colorOptions, boolean pageEnabled,
-                int colorsPerPage) {
-            mPageEnabled = pageEnabled;
-            mColorOptions = colorOptions;
-            mColorsPerPage = colorsPerPage;
-        }
-
-        @Override
-        public int getItemCount() {
-            if (!mPageEnabled || !ColorProvider.themeStyleEnabled) {
-                return 1;
-            }
-            // Color page size.
-            return getNumPages(mColorsPerPage, mColorOptions.size());
-        }
-
-        @Override
-        public void onBindViewHolder(ColorOptionViewHolder viewHolder, int position) {
-            setupColorOptions(viewHolder.mContainer, mColorOptions, mPageEnabled, position,
-                    mColorsPerPage);
-        }
-
-        @Override
-        public ColorOptionViewHolder onCreateViewHolder(ViewGroup viewGroup, int viewType) {
-            return new ColorOptionViewHolder(
-                    LayoutInflater.from(viewGroup.getContext()).inflate(viewType, viewGroup,
-                            false));
-        }
-
-        @Override
-        public int getItemViewType(int position) {
-            return R.layout.color_options_view;
-        }
-
-        private class ColorOptionViewHolder extends RecyclerView.ViewHolder {
-            private RecyclerView mContainer;
-
-            ColorOptionViewHolder(View itemView) {
-                super(itemView);
-                mContainer = itemView.findViewById(R.id.color_option_container);
-                // Sets layout with margins to separate color options.
-                final FrameLayout.LayoutParams layoutParams = new FrameLayout.LayoutParams(
-                        mContainer.getLayoutParams());
-                final int margin = itemView.getContext().getResources().getDimensionPixelSize(
-                        R.dimen.section_horizontal_padding);
-                layoutParams.setMargins(margin, /* top= */ 0, margin, /* bottom= */ 0);
-                mContainer.setLayoutParams(layoutParams);
-            }
-        }
-    }
-}
diff --git a/src/com/android/customization/model/grid/GridOption.java b/src/com/android/customization/model/grid/GridOption.java
index 19c5d4f..17ba115 100644
--- a/src/com/android/customization/model/grid/GridOption.java
+++ b/src/com/android/customization/model/grid/GridOption.java
@@ -49,7 +49,6 @@
     };
 
     private final String mTitle;
-    private final boolean mIsCurrent;
     private final String mIconShapePath;
     private final GridTileDrawable mTileDrawable;
     public final String name;
@@ -57,6 +56,7 @@
     public final int cols;
     public final Uri previewImageUri;
     public final int previewPagesCount;
+    private boolean mIsCurrent;
 
     public GridOption(String title, String name, boolean isCurrent, int rows, int cols,
             Uri previewImageUri, int previewPagesCount, String iconShapePath) {
@@ -71,6 +71,10 @@
         this.previewPagesCount = previewPagesCount;
     }
 
+    public void setIsCurrent(boolean isCurrent) {
+        mIsCurrent = isCurrent;
+    }
+
     protected GridOption(Parcel in) {
         mTitle = in.readString();
         mIsCurrent = in.readByte() != 0;
diff --git a/src/com/android/customization/model/grid/GridOptionsManager.java b/src/com/android/customization/model/grid/GridOptionsManager.java
index b7ee37f..78dbb5b 100644
--- a/src/com/android/customization/model/grid/GridOptionsManager.java
+++ b/src/com/android/customization/model/grid/GridOptionsManager.java
@@ -99,6 +99,11 @@
     }
 
     @Override
+    public void preview(GridOption option) {
+        mProvider.updateView();
+    }
+
+    @Override
     public void fetchOptions(OptionsFetchedListener<GridOption> callback, boolean reload) {
         sExecutorService.submit(() -> {
             List<GridOption> gridOptions = mProvider.fetch(reload);
diff --git a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
index 4e775c6..8f1860e 100644
--- a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
+++ b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
@@ -117,6 +117,10 @@
         mPreviewUtils.renderPreview(bundle, callback);
     }
 
+    void updateView() {
+        mLiveData.postValue(new Object());
+    }
+
     int applyGrid(String name) {
         ContentValues values = new ContentValues();
         values.put("name", name);
diff --git a/src/com/android/customization/model/grid/data/repository/GridRepository.kt b/src/com/android/customization/model/grid/data/repository/GridRepository.kt
index 9a3be0c..e65c18e 100644
--- a/src/com/android/customization/model/grid/data/repository/GridRepository.kt
+++ b/src/com/android/customization/model/grid/data/repository/GridRepository.kt
@@ -38,12 +38,14 @@
     suspend fun isAvailable(): Boolean
     fun getOptionChanges(): Flow<Unit>
     suspend fun getOptions(): GridOptionItemsModel
+    fun getSelectedOption(): GridOption?
 }
 
 class GridRepositoryImpl(
     private val applicationScope: CoroutineScope,
     private val manager: GridOptionsManager,
     private val backgroundDispatcher: CoroutineDispatcher,
+    private val isGridApplyButtonEnabled: Boolean,
 ) : GridRepository {
 
     override suspend fun isAvailable(): Boolean {
@@ -55,6 +57,8 @@
 
     private val selectedOption = MutableStateFlow<GridOption?>(null)
 
+    override fun getSelectedOption() = selectedOption.value
+
     override suspend fun getOptions(): GridOptionItemsModel {
         return withContext(backgroundDispatcher) {
             suspendCancellableCoroutine { continuation ->
@@ -62,7 +66,11 @@
                     object : CustomizationManager.OptionsFetchedListener<GridOption> {
                         override fun onOptionsLoaded(options: MutableList<GridOption>?) {
                             val optionsOrEmpty = options ?: emptyList()
-                            selectedOption.value = optionsOrEmpty.find { it.isActive(manager) }
+                            // After Apply Button is added, we will rely on onSelected() method
+                            // to update selectedOption.
+                            if (!isGridApplyButtonEnabled || selectedOption.value == null) {
+                                selectedOption.value = optionsOrEmpty.find { it.isActive(manager) }
+                            }
                             continuation.resume(
                                 GridOptionItemsModel.Loaded(
                                     optionsOrEmpty.map { option -> toModel(option) }
@@ -105,18 +113,26 @@
     private suspend fun onSelected(option: GridOption) {
         withContext(backgroundDispatcher) {
             suspendCancellableCoroutine { continuation ->
-                manager.apply(
-                    option,
-                    object : CustomizationManager.Callback {
-                        override fun onSuccess() {
-                            continuation.resume(true)
-                        }
+                if (isGridApplyButtonEnabled) {
+                    selectedOption.value?.setIsCurrent(false)
+                    selectedOption.value = option
+                    selectedOption.value?.setIsCurrent(true)
+                    manager.preview(option)
+                    continuation.resume(true)
+                } else {
+                    manager.apply(
+                        option,
+                        object : CustomizationManager.Callback {
+                            override fun onSuccess() {
+                                continuation.resume(true)
+                            }
 
-                        override fun onError(throwable: Throwable?) {
-                            continuation.resume(false)
-                        }
-                    },
-                )
+                            override fun onError(throwable: Throwable?) {
+                                continuation.resume(false)
+                            }
+                        },
+                    )
+                }
             }
         }
     }
diff --git a/src/com/android/customization/model/grid/domain/interactor/GridInteractor.kt b/src/com/android/customization/model/grid/domain/interactor/GridInteractor.kt
index cdb679d..6cb00ff 100644
--- a/src/com/android/customization/model/grid/domain/interactor/GridInteractor.kt
+++ b/src/com/android/customization/model/grid/domain/interactor/GridInteractor.kt
@@ -95,4 +95,6 @@
             model
         }
     }
+
+    fun getSelectedOptionName(): String? = repository.getSelectedOption()?.name
 }
diff --git a/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt b/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt
index 6ab561f..14d1344 100644
--- a/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt
+++ b/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt
@@ -21,7 +21,12 @@
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
+import androidx.core.content.ContextCompat
+import androidx.core.view.isVisible
 import androidx.lifecycle.ViewModelProvider
+import androidx.transition.Transition
+import androidx.transition.doOnStart
+import com.android.customization.model.grid.domain.interactor.GridInteractor
 import com.android.customization.model.grid.ui.binder.GridScreenBinder
 import com.android.customization.model.grid.ui.viewmodel.GridScreenViewModel
 import com.android.customization.module.ThemePickerInjector
@@ -61,7 +66,8 @@
             bindScreenPreview(
                 view,
                 wallpaperInfoFactory,
-                injector.getWallpaperInteractor(requireContext())
+                injector.getWallpaperInteractor(requireContext()),
+                injector.getGridInteractor(requireContext())
             )
 
         val viewModelFactory = injector.getGridScreenViewModelFactory(requireContext())
@@ -80,11 +86,16 @@
                     bindScreenPreview(
                         view,
                         wallpaperInfoFactory,
-                        injector.getWallpaperInteractor(requireContext())
+                        injector.getWallpaperInteractor(requireContext()),
+                        injector.getGridInteractor(requireContext())
                     )
             }
         )
 
+        (returnTransition as? Transition)?.doOnStart {
+            view.requireViewById<View>(R.id.preview).isVisible = false
+        }
+
         return view
     }
 
@@ -92,10 +103,19 @@
         return getString(R.string.grid_title)
     }
 
+    override fun getToolbarColorId(): Int {
+        return android.R.color.transparent
+    }
+
+    override fun getToolbarTextColor(): Int {
+        return ContextCompat.getColor(requireContext(), R.color.system_on_surface)
+    }
+
     private fun bindScreenPreview(
         view: View,
         wallpaperInfoFactory: CurrentWallpaperInfoFactory,
         wallpaperInteractor: WallpaperInteractor,
+        gridInteractor: GridInteractor
     ): ScreenPreviewBinder.Binding {
         return ScreenPreviewBinder.bind(
             activity = requireActivity(),
@@ -111,6 +131,11 @@
                                         R.string.grid_control_metadata_name,
                                     ),
                         ),
+                    initialExtrasProvider = {
+                        val bundle = Bundle()
+                        bundle.putString("name", gridInteractor.getSelectedOptionName())
+                        bundle
+                    },
                     wallpaperInfoProvider = {
                         suspendCancellableCoroutine { continuation ->
                             wallpaperInfoFactory.createCurrentWallpaperInfos(
diff --git a/src/com/android/customization/module/CustomizationInjector.kt b/src/com/android/customization/module/CustomizationInjector.kt
index 5f8f9d3..8fd3768 100644
--- a/src/com/android/customization/module/CustomizationInjector.kt
+++ b/src/com/android/customization/module/CustomizationInjector.kt
@@ -16,6 +16,7 @@
 package com.android.customization.module
 
 import android.content.Context
+import android.content.res.Resources
 import androidx.activity.ComponentActivity
 import androidx.fragment.app.FragmentActivity
 import com.android.customization.model.theme.OverlayManagerCompat
@@ -78,5 +79,5 @@
         clockViewFactory: ClockViewFactory,
     ): ClockSettingsViewModel.Factory
 
-    fun getClockDescriptionUtils(): ClockDescriptionUtils
+    fun getClockDescriptionUtils(resources: Resources): ClockDescriptionUtils
 }
diff --git a/src/com/android/customization/module/DefaultCustomizationSections.java b/src/com/android/customization/module/DefaultCustomizationSections.java
index b408a89..bbe6bef 100644
--- a/src/com/android/customization/module/DefaultCustomizationSections.java
+++ b/src/com/android/customization/module/DefaultCustomizationSections.java
@@ -8,10 +8,8 @@
 import androidx.lifecycle.LifecycleOwner;
 import androidx.lifecycle.ViewModelProvider;
 
-import com.android.customization.model.color.ColorSectionController;
 import com.android.customization.model.grid.GridOptionsManager;
 import com.android.customization.model.grid.GridSectionController;
-import com.android.customization.model.mode.DarkModeSectionController;
 import com.android.customization.model.mode.DarkModeSnapshotRestorer;
 import com.android.customization.model.themedicon.ThemedIconSectionController;
 import com.android.customization.model.themedicon.ThemedIconSwitchProvider;
@@ -19,6 +17,7 @@
 import com.android.customization.model.themedicon.domain.interactor.ThemedIconSnapshotRestorer;
 import com.android.customization.picker.clock.ui.view.ClockViewFactory;
 import com.android.customization.picker.clock.ui.viewmodel.ClockCarouselViewModel;
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor;
 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;
@@ -35,7 +34,6 @@
 import com.android.wallpaper.model.PermissionRequester;
 import com.android.wallpaper.model.WallpaperColorsViewModel;
 import com.android.wallpaper.model.WallpaperPreviewNavigator;
-import com.android.wallpaper.model.WallpaperSectionController;
 import com.android.wallpaper.module.CurrentWallpaperInfoFactory;
 import com.android.wallpaper.module.CustomizationSections;
 import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor;
@@ -61,6 +59,7 @@
     private final DarkModeSnapshotRestorer mDarkModeSnapshotRestorer;
     private final ThemedIconSnapshotRestorer mThemedIconSnapshotRestorer;
     private final ThemedIconInteractor mThemedIconInteractor;
+    private final ColorPickerInteractor mColorPickerInteractor;
 
     public DefaultCustomizationSections(
             ColorPickerViewModel.Factory colorPickerViewModelFactory,
@@ -73,7 +72,8 @@
             ClockViewFactory clockViewFactory,
             DarkModeSnapshotRestorer darkModeSnapshotRestorer,
             ThemedIconSnapshotRestorer themedIconSnapshotRestorer,
-            ThemedIconInteractor themedIconInteractor) {
+            ThemedIconInteractor themedIconInteractor,
+            ColorPickerInteractor colorPickerInteractor) {
         mColorPickerViewModelFactory = colorPickerViewModelFactory;
         mKeyguardQuickAffordancePickerInteractor = keyguardQuickAffordancePickerInteractor;
         mKeyguardQuickAffordancePickerViewModelFactory =
@@ -85,10 +85,11 @@
         mDarkModeSnapshotRestorer = darkModeSnapshotRestorer;
         mThemedIconSnapshotRestorer = themedIconSnapshotRestorer;
         mThemedIconInteractor = themedIconInteractor;
+        mColorPickerInteractor = colorPickerInteractor;
     }
 
     @Override
-    public List<CustomizationSectionController<?>> getRevampedUISectionControllersForScreen(
+    public List<CustomizationSectionController<?>> getSectionControllersForScreen(
             Screen screen,
             FragmentActivity activity,
             LifecycleOwner lifecycleOwner,
@@ -121,8 +122,10 @@
                         sectionNavigationController,
                         wallpaperInteractor,
                         mThemedIconInteractor,
+                        mColorPickerInteractor,
                         wallpaperManager,
-                        isTwoPaneAndSmallWidth)
+                        isTwoPaneAndSmallWidth,
+                        customizationPickerViewModel)
                         : new PreviewWithThemeSectionController(
                                 activity,
                                 lifecycleOwner,
@@ -133,8 +136,10 @@
                                 wallpaperPreviewNavigator,
                                 wallpaperInteractor,
                                 mThemedIconInteractor,
+                                mColorPickerInteractor,
                                 wallpaperManager,
-                                isTwoPaneAndSmallWidth));
+                                isTwoPaneAndSmallWidth,
+                                customizationPickerViewModel));
 
         sectionControllers.add(
                 new ConnectedSectionController(
@@ -202,57 +207,4 @@
 
         return sectionControllers;
     }
-
-    @Override
-    public List<CustomizationSectionController<?>> getAllSectionControllers(
-            FragmentActivity activity,
-            LifecycleOwner lifecycleOwner,
-            WallpaperColorsViewModel wallpaperColorsViewModel,
-            PermissionRequester permissionRequester,
-            WallpaperPreviewNavigator wallpaperPreviewNavigator,
-            CustomizationSectionNavigationController sectionNavigationController,
-            @Nullable Bundle savedInstanceState,
-            DisplayUtils displayUtils) {
-        List<CustomizationSectionController<?>> sectionControllers = new ArrayList<>();
-
-        // Wallpaper section.
-        sectionControllers.add(
-                new WallpaperSectionController(
-                activity,
-                lifecycleOwner,
-                permissionRequester,
-                wallpaperColorsViewModel,
-                mThemedIconInteractor.isActivatedAsLiveData(),
-                sectionNavigationController,
-                wallpaperPreviewNavigator,
-                savedInstanceState,
-                displayUtils));
-
-        // Theme color section.
-        sectionControllers.add(new ColorSectionController(
-                activity, wallpaperColorsViewModel, lifecycleOwner, savedInstanceState));
-
-        // Dark/Light theme section.
-        sectionControllers.add(new DarkModeSectionController(
-                activity,
-                lifecycleOwner.getLifecycle(),
-                mDarkModeSnapshotRestorer));
-
-        // Themed app icon section.
-        sectionControllers.add(new ThemedIconSectionController(
-                ThemedIconSwitchProvider.getInstance(activity),
-                mThemedIconInteractor,
-                savedInstanceState,
-                mThemedIconSnapshotRestorer));
-
-        // App grid section.
-        sectionControllers.add(
-                new GridSectionController(
-                        GridOptionsManager.getInstance(activity),
-                        sectionNavigationController,
-                        lifecycleOwner,
-                        /* isRevampedUiEnabled= */ false));
-
-        return sectionControllers;
-    }
 }
diff --git a/src/com/android/customization/module/StatsLogUserEventLogger.java b/src/com/android/customization/module/StatsLogUserEventLogger.java
index 32e0599..647cdc9 100644
--- a/src/com/android/customization/module/StatsLogUserEventLogger.java
+++ b/src/com/android/customization/module/StatsLogUserEventLogger.java
@@ -37,7 +37,6 @@
 import static com.android.wallpaper.util.LaunchSourceUtils.WALLPAPER_LAUNCH_SOURCE;
 
 import android.app.WallpaperManager;
-import android.content.Context;
 import android.content.Intent;
 import android.stats.style.StyleEnums;
 import android.text.TextUtils;
@@ -47,8 +46,6 @@
 import com.android.customization.model.color.ColorOption;
 import com.android.customization.model.grid.GridOption;
 import com.android.customization.model.theme.ThemeBundle;
-import com.android.wallpaper.module.Injector;
-import com.android.wallpaper.module.InjectorProvider;
 import com.android.wallpaper.module.NoOpUserEventLogger;
 import com.android.wallpaper.module.WallpaperPreferences;
 import com.android.wallpaper.module.WallpaperStatusChecker;
@@ -61,86 +58,59 @@
  */
 public class StatsLogUserEventLogger extends NoOpUserEventLogger implements ThemesUserEventLogger {
 
-    private static final String TAG = "StatsLogUserEventLogger";
-    private final Context mContext;
     private final WallpaperPreferences mPreferences;
     private final WallpaperStatusChecker mWallpaperStatusChecker;
 
-    public StatsLogUserEventLogger(Context appContext) {
-        mContext = appContext;
-        Injector injector = InjectorProvider.getInjector();
-        mPreferences = injector.getPreferences(appContext);
-        mWallpaperStatusChecker = injector.getWallpaperStatusChecker();
+    public StatsLogUserEventLogger(
+            WallpaperPreferences preferences,
+            WallpaperStatusChecker wallpaperStatusChecker) {
+        mPreferences = preferences;
+        mWallpaperStatusChecker = wallpaperStatusChecker;
     }
 
     @Override
     public void logAppLaunched(Intent launchSource) {
-        new SysUiStatsLogger()
-                .setAction(STYLE_UICHANGED__ACTION__APP_LAUNCHED)
+        new SysUiStatsLogger(STYLE_UICHANGED__ACTION__APP_LAUNCHED)
                 .setLaunchedPreference(getAppLaunchSource(launchSource))
                 .log();
     }
 
     @Override
     public void logResumed(boolean provisioned, boolean wallpaper) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.ONRESUME)
+        new SysUiStatsLogger(StyleEnums.ONRESUME)
                 .log();
     }
 
     @Override
     public void logStopped() {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.ONSTOP)
+        new SysUiStatsLogger(StyleEnums.ONSTOP)
                 .log();
     }
 
     @Override
     public void logActionClicked(String collectionId, int actionLabelResId) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.WALLPAPER_EXPLORE)
+        new SysUiStatsLogger(StyleEnums.WALLPAPER_EXPLORE)
                 .setWallpaperCategoryHash(getIdHashCode(collectionId))
                 .log();
     }
 
     @Override
     public void logIndividualWallpaperSelected(String collectionId) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.WALLPAPER_SELECT)
+        new SysUiStatsLogger(StyleEnums.WALLPAPER_SELECT)
                 .setWallpaperCategoryHash(getIdHashCode(collectionId))
                 .log();
     }
 
     @Override
     public void logCategorySelected(String collectionId) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.WALLPAPER_OPEN_CATEGORY)
+        new SysUiStatsLogger(StyleEnums.WALLPAPER_OPEN_CATEGORY)
                 .setWallpaperCategoryHash(getIdHashCode(collectionId))
                 .log();
     }
 
     @Override
-    public void logLiveWallpaperInfoSelected(String collectionId, @Nullable String wallpaperId) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.LIVE_WALLPAPER_INFO_SELECT)
-                .setWallpaperCategoryHash(getIdHashCode(collectionId))
-                .setWallpaperIdHash(getIdHashCode(wallpaperId))
-                .log();
-    }
-
-    @Override
-    public void logLiveWallpaperCustomizeSelected(String collectionId,
-            @Nullable String wallpaperId) {
-        new SysUiStatsLogger().setAction(StyleEnums.LIVE_WALLPAPER_CUSTOMIZE_SELECT)
-                .setWallpaperCategoryHash(getIdHashCode(collectionId))
-                .setWallpaperIdHash(getIdHashCode(wallpaperId))
-                .log();
-
-    }
-
-    @Override
     public void logSnapshot() {
-        final boolean isLockWallpaperSet = mWallpaperStatusChecker.isLockWallpaperSet(mContext);
+        final boolean isLockWallpaperSet = mWallpaperStatusChecker.isLockWallpaperSet();
         final String homeCollectionId = mPreferences.getHomeWallpaperCollectionId();
         final String homeRemoteId = mPreferences.getHomeWallpaperRemoteId();
         final String effects = mPreferences.getHomeWallpaperEffects();
@@ -151,7 +121,7 @@
         String lockWallpaperId = isLockWallpaperSet ? mPreferences.getLockWallpaperRemoteId()
                 : homeWallpaperId;
 
-        new SysUiStatsLogger().setAction(StyleEnums.SNAPSHOT)
+        new SysUiStatsLogger(StyleEnums.SNAPSHOT)
                 .setWallpaperCategoryHash(getIdHashCode(homeCollectionId))
                 .setWallpaperIdHash(getIdHashCode(homeWallpaperId))
                 .setLockWallpaperCategoryHash(getIdHashCode(lockCollectionId))
@@ -167,8 +137,7 @@
     @Override
     public void logWallpaperSet(String collectionId, @Nullable String wallpaperId,
             @Nullable String effects) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.WALLPAPER_APPLIED)
+        new SysUiStatsLogger(StyleEnums.WALLPAPER_APPLIED)
                 .setWallpaperCategoryHash(getIdHashCode(collectionId))
                 .setWallpaperIdHash(getIdHashCode(wallpaperId))
                 .setEffectIdHash(getIdHashCode(effects))
@@ -178,8 +147,7 @@
     @Override
     public void logEffectApply(String effect, @EffectStatus int status, long timeElapsedMillis,
             int resultCode) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.WALLPAPER_EFFECT_APPLIED)
+        new SysUiStatsLogger(StyleEnums.WALLPAPER_EFFECT_APPLIED)
                 .setEffectPreference(status)
                 .setEffectIdHash(getIdHashCode(effect))
                 .setTimeElapsed(timeElapsedMillis)
@@ -189,13 +157,22 @@
 
     @Override
     public void logEffectProbe(String effect, @EffectStatus int status) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.WALLPAPER_EFFECT_PROBE)
+        new SysUiStatsLogger(StyleEnums.WALLPAPER_EFFECT_PROBE)
                 .setEffectPreference(status)
                 .setEffectIdHash(getIdHashCode(effect))
                 .log();
     }
 
+    @Override
+    public void logEffectForegroundDownload(String effect, @EffectStatus int status,
+            long timeElapsedMillis) {
+        new SysUiStatsLogger(StyleEnums.WALLPAPER_EFFECT_FG_DOWNLOAD)
+                .setEffectPreference(status)
+                .setEffectIdHash(getIdHashCode(effect))
+                .setTimeElapsed(timeElapsedMillis)
+                .log();
+    }
+
     @Nullable
     private String getThemePackage(ThemeBundle theme, String category) {
         Map<String, String> packages = theme.getPackagesByCategory();
@@ -204,8 +181,7 @@
 
     @Override
     public void logThemeSelected(ThemeBundle theme, boolean isCustomTheme) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.PICKER_SELECT)
+        new SysUiStatsLogger(StyleEnums.PICKER_SELECT)
                 .setColorPackageHash(
                         Objects.hashCode(getThemePackage(theme, OVERLAY_CATEGORY_COLOR)))
                 .setFontPackageHash(Objects.hashCode(getThemePackage(theme, OVERLAY_CATEGORY_FONT)))
@@ -216,8 +192,7 @@
 
     @Override
     public void logThemeApplied(ThemeBundle theme, boolean isCustomTheme) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.PICKER_APPLIED)
+        new SysUiStatsLogger(StyleEnums.PICKER_APPLIED)
                 .setColorPackageHash(
                         Objects.hashCode(getThemePackage(theme, OVERLAY_CATEGORY_COLOR)))
                 .setFontPackageHash(Objects.hashCode(getThemePackage(theme, OVERLAY_CATEGORY_FONT)))
@@ -228,8 +203,7 @@
 
     @Override
     public void logColorApplied(int action, ColorOption colorOption) {
-        new SysUiStatsLogger()
-                .setAction(action)
+        new SysUiStatsLogger(action)
                 .setColorPreference(colorOption.getIndex())
                 .setColorVariant(colorOption.getStyle().ordinal() + 1)
                 .log();
@@ -237,16 +211,14 @@
 
     @Override
     public void logGridSelected(GridOption grid) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.PICKER_SELECT)
+        new SysUiStatsLogger(StyleEnums.PICKER_SELECT)
                 .setLauncherGrid(grid.cols)
                 .log();
     }
 
     @Override
     public void logGridApplied(GridOption grid) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.PICKER_APPLIED)
+        new SysUiStatsLogger(StyleEnums.PICKER_APPLIED)
                 .setLauncherGrid(grid.cols)
                 .log();
     }
diff --git a/src/com/android/customization/module/SysUiStatsLogger.kt b/src/com/android/customization/module/SysUiStatsLogger.kt
index 318bf1f..8e97b0b 100644
--- a/src/com/android/customization/module/SysUiStatsLogger.kt
+++ b/src/com/android/customization/module/SysUiStatsLogger.kt
@@ -20,10 +20,8 @@
 import com.android.systemui.shared.system.SysUiStatsLog.STYLE_UI_CHANGED
 
 /** The builder for [SysUiStatsLog]. */
-class SysUiStatsLogger {
+class SysUiStatsLogger(val action: Int) {
 
-    private var atom = STYLE_UI_CHANGED
-    private var action = StyleEnums.DEFAULT_ACTION
     private var colorPackageHash = 0
     private var fontPackageHash = 0
     private var shapePackageHash = 0
@@ -46,85 +44,83 @@
     private var timeElapsedMillis = 0L
     private var effectResultCode = -1
 
-    fun setAction(action: Int) = apply { this.action = action }
-
-    fun setColorPackageHash(color_package_hash: Int) = apply {
-        this.colorPackageHash = color_package_hash
+    fun setColorPackageHash(colorPackageHash: Int) = apply {
+        this.colorPackageHash = colorPackageHash
     }
 
-    fun setFontPackageHash(font_package_hash: Int) = apply {
-        this.fontPackageHash = font_package_hash
+    fun setFontPackageHash(fontPackageHash: Int) = apply {
+        this.fontPackageHash = fontPackageHash
     }
 
-    fun setShapePackageHash(shape_package_hash: Int) = apply {
-        this.shapePackageHash = shape_package_hash
+    fun setShapePackageHash(shapePackageHash: Int) = apply {
+        this.shapePackageHash = shapePackageHash
     }
 
-    fun setClockPackageHash(clock_package_hash: Int) = apply {
-        this.clockPackageHash = clock_package_hash
+    fun setClockPackageHash(clockPackageHash: Int) = apply {
+        this.clockPackageHash = clockPackageHash
     }
 
-    fun setLauncherGrid(launcher_grid: Int) = apply { this.launcherGrid = launcher_grid }
+    fun setLauncherGrid(launcherGrid: Int) = apply { this.launcherGrid = launcherGrid }
 
-    fun setWallpaperCategoryHash(wallpaper_category_hash: Int) = apply {
-        this.wallpaperCategoryHash = wallpaper_category_hash
+    fun setWallpaperCategoryHash(wallpaperCategoryHash: Int) = apply {
+        this.wallpaperCategoryHash = wallpaperCategoryHash
     }
 
-    fun setWallpaperIdHash(wallpaper_id_hash: Int) = apply {
-        this.wallpaperIdHash = wallpaper_id_hash
+    fun setWallpaperIdHash(wallpaperIdHash: Int) = apply {
+        this.wallpaperIdHash = wallpaperIdHash
     }
 
-    fun setColorPreference(color_preference: Int) = apply {
-        this.colorPreference = color_preference
+    fun setColorPreference(colorPreference: Int) = apply {
+        this.colorPreference = colorPreference
     }
 
-    fun setLocationPreference(location_preference: Int) = apply {
-        this.locationPreference = location_preference
+    fun setLocationPreference(locationPreference: Int) = apply {
+        this.locationPreference = locationPreference
     }
 
-    fun setDatePreference(date_preference: Int) = apply { this.datePreference = date_preference }
+    fun setDatePreference(datePreference: Int) = apply { this.datePreference = datePreference }
 
-    fun setLaunchedPreference(launched_preference: Int) = apply {
-        this.launchedPreference = launched_preference
+    fun setLaunchedPreference(launchedPreference: Int) = apply {
+        this.launchedPreference = launchedPreference
     }
 
-    fun setEffectPreference(effect_preference: Int) = apply {
-        this.effectPreference = effect_preference
+    fun setEffectPreference(effectPreference: Int) = apply {
+        this.effectPreference = effectPreference
     }
 
-    fun setEffectIdHash(effect_id_hash: Int) = apply { this.effectIdHash = effect_id_hash }
+    fun setEffectIdHash(effectIdHash: Int) = apply { this.effectIdHash = effectIdHash }
 
-    fun setLockWallpaperCategoryHash(lock_wallpaper_category_hash: Int) = apply {
-        this.lockWallpaperCategoryHash = lock_wallpaper_category_hash
+    fun setLockWallpaperCategoryHash(lockWallpaperCategoryHash: Int) = apply {
+        this.lockWallpaperCategoryHash = lockWallpaperCategoryHash
     }
 
-    fun setLockWallpaperIdHash(lock_wallpaper_id_hash: Int) = apply {
-        this.lockWallpaperIdHash = lock_wallpaper_id_hash
+    fun setLockWallpaperIdHash(lockWallpaperIdHash: Int) = apply {
+        this.lockWallpaperIdHash = lockWallpaperIdHash
     }
 
-    fun setFirstLaunchDateSinceSetup(first_launch_date_since_setup: Int) = apply {
-        this.firstLaunchDateSinceSetup = first_launch_date_since_setup
+    fun setFirstLaunchDateSinceSetup(firstLaunchDateSinceSetup: Int) = apply {
+        this.firstLaunchDateSinceSetup = firstLaunchDateSinceSetup
     }
 
-    fun setFirstWallpaperApplyDateSinceSetup(first_wallpaper_apply_date_since_setup: Int) = apply {
-        this.firstWallpaperApplyDateSinceSetup = first_wallpaper_apply_date_since_setup
+    fun setFirstWallpaperApplyDateSinceSetup(firstWallpaperApplyDateSinceSetup: Int) = apply {
+        this.firstWallpaperApplyDateSinceSetup = firstWallpaperApplyDateSinceSetup
     }
 
-    fun setAppLaunchCount(app_launch_count: Int) = apply { this.appLaunchCount = app_launch_count }
+    fun setAppLaunchCount(appLaunchCount: Int) = apply { this.appLaunchCount = appLaunchCount }
 
-    fun setColorVariant(color_variant: Int) = apply { this.colorVariant = color_variant }
+    fun setColorVariant(colorVariant: Int) = apply { this.colorVariant = colorVariant }
 
-    fun setTimeElapsed(time_elapsed_millis: Long) = apply {
-      this.timeElapsedMillis = time_elapsed_millis
+    fun setTimeElapsed(timeElapsedMillis: Long) = apply {
+      this.timeElapsedMillis = timeElapsedMillis
     }
 
-    fun setEffectResultCode(effect_result_code: Int) = apply {
-        this.effectResultCode = effect_result_code
+    fun setEffectResultCode(effectResultCode: Int) = apply {
+        this.effectResultCode = effectResultCode
     }
 
     fun log() {
         SysUiStatsLog.write(
-            atom,
+            STYLE_UI_CHANGED,
             action,
             colorPackageHash,
             fontPackageHash,
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index cceb896..3197bf6 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -19,6 +19,7 @@
 import android.app.WallpaperManager
 import android.content.Context
 import android.content.Intent
+import android.content.res.Resources
 import android.net.Uri
 import android.os.Bundle
 import android.text.TextUtils
@@ -68,6 +69,9 @@
 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
+import com.android.wallpaper.config.BaseFlags
+import com.android.wallpaper.dispatchers.BackgroundDispatcher
+import com.android.wallpaper.dispatchers.MainDispatcher
 import com.android.wallpaper.model.LiveWallpaperInfo
 import com.android.wallpaper.model.WallpaperColorsViewModel
 import com.android.wallpaper.model.WallpaperInfo
@@ -85,9 +89,19 @@
 import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
 import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
 import com.android.wallpaper.util.ScreenSizeCalculator
-import kotlinx.coroutines.Dispatchers
+import javax.inject.Inject
+import javax.inject.Singleton
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
 
-open class ThemePickerInjector : WallpaperPicker2Injector(), CustomizationInjector {
+@Singleton
+open class ThemePickerInjector
+@Inject
+internal constructor(
+    @MainDispatcher private val mainScope: CoroutineScope,
+    @MainDispatcher private val mainDispatcher: CoroutineDispatcher,
+    @BackgroundDispatcher private val bgDispatcher: CoroutineDispatcher,
+) : WallpaperPicker2Injector(mainScope, bgDispatcher), CustomizationInjector {
     private var customizationSections: CustomizationSections? = null
     private var userEventLogger: UserEventLogger? = null
     private var prefs: WallpaperPreferences? = null
@@ -143,6 +157,7 @@
                     getDarkModeSnapshotRestorer(activity),
                     getThemedIconSnapshotRestorer(activity),
                     getThemedIconInteractor(),
+                    getColorPickerInteractor(activity, getWallpaperColorsViewModel()),
                 )
                 .also { customizationSections = it }
     }
@@ -183,8 +198,12 @@
 
     @Synchronized
     override fun getUserEventLogger(context: Context): ThemesUserEventLogger {
-        return if (userEventLogger != null) userEventLogger as ThemesUserEventLogger
-        else StatsLogUserEventLogger(context.applicationContext).also { userEventLogger = it }
+        return userEventLogger as? ThemesUserEventLogger
+            ?: StatsLogUserEventLogger(
+                    getPreferences(context.applicationContext),
+                    getWallpaperStatusChecker(context.applicationContext),
+                )
+                .also { userEventLogger = it }
     }
 
     @Synchronized
@@ -244,7 +263,7 @@
                                     wallpaperManager = WallpaperManager.getInstance(appContext)
                                 ),
                             wallpaperPreferences = getPreferences(context = appContext),
-                            backgroundDispatcher = Dispatchers.IO,
+                            backgroundDispatcher = bgDispatcher,
                         ),
                     shouldHandleReload = {
                         TextUtils.equals(
@@ -284,7 +303,7 @@
         val client = getKeyguardQuickAffordancePickerProviderClient(context)
         val appContext = context.applicationContext
         return KeyguardQuickAffordancePickerInteractor(
-            KeyguardQuickAffordancePickerRepository(client, Dispatchers.IO),
+            KeyguardQuickAffordancePickerRepository(client, bgDispatcher),
             client
         ) {
             getKeyguardQuickAffordanceSnapshotRestorer(appContext)
@@ -295,7 +314,7 @@
         context: Context
     ): CustomizationProviderClient {
         return customizationProviderClient
-            ?: CustomizationProviderClientImpl(context.applicationContext, Dispatchers.IO).also {
+            ?: CustomizationProviderClientImpl(context.applicationContext, bgDispatcher).also {
                 customizationProviderClient = it
             }
     }
@@ -330,7 +349,7 @@
                     repository =
                         NotificationsRepository(
                             scope = getApplicationCoroutineScope(),
-                            backgroundDispatcher = Dispatchers.IO,
+                            backgroundDispatcher = bgDispatcher,
                             secureSettingsRepository = getSecureSettingsRepository(context),
                         ),
                     snapshotRestorer = { getNotificationsSnapshotRestorer(appContext) },
@@ -354,8 +373,8 @@
                 ?: ClockRegistryProvider(
                         context = context.applicationContext,
                         coroutineScope = getApplicationCoroutineScope(),
-                        mainDispatcher = Dispatchers.Main,
-                        backgroundDispatcher = Dispatchers.IO,
+                        mainDispatcher = mainDispatcher,
+                        backgroundDispatcher = bgDispatcher,
                     )
                     .also { clockRegistryProvider = it })
             .get()
@@ -372,7 +391,7 @@
                             secureSettingsRepository = getSecureSettingsRepository(appContext),
                             registry = getClockRegistry(appContext),
                             scope = getApplicationCoroutineScope(),
-                            mainDispatcher = Dispatchers.Main,
+                            mainDispatcher = mainDispatcher,
                         ),
                     snapshotRestorer = { getClockPickerSnapshotRestorer(appContext) },
                 )
@@ -394,7 +413,7 @@
         interactor: ClockPickerInteractor,
     ): ClockCarouselViewModel.Factory {
         return clockCarouselViewModelFactory
-            ?: ClockCarouselViewModel.Factory(interactor, Dispatchers.IO).also {
+            ?: ClockCarouselViewModel.Factory(interactor, bgDispatcher).also {
                 clockCarouselViewModelFactory = it
             }
     }
@@ -489,7 +508,7 @@
             ?: DarkModeSnapshotRestorer(
                     context = appContext,
                     manager = appContext.getSystemService(Context.UI_MODE_SERVICE) as UiModeManager,
-                    backgroundDispatcher = Dispatchers.IO,
+                    backgroundDispatcher = bgDispatcher,
                 )
                 .also { darkModeSnapshotRestorer = it }
     }
@@ -537,7 +556,7 @@
                 .also { clockSettingsViewModelFactory = it }
     }
 
-    override fun getClockDescriptionUtils(): ClockDescriptionUtils {
+    override fun getClockDescriptionUtils(resources: Resources): ClockDescriptionUtils {
         return clockDescriptionUtils
             ?: ThemePickerClockDescriptionUtils().also { clockDescriptionUtils = it }
     }
@@ -553,7 +572,7 @@
                 .also { gridScreenViewModelFactory = it }
     }
 
-    private fun getGridInteractor(
+    fun getGridInteractor(
         context: Context,
     ): GridInteractor {
         val appContext = context.applicationContext
@@ -564,7 +583,9 @@
                         GridRepositoryImpl(
                             applicationScope = getApplicationCoroutineScope(),
                             manager = GridOptionsManager.getInstance(context),
-                            backgroundDispatcher = Dispatchers.IO,
+                            backgroundDispatcher = bgDispatcher,
+                            isGridApplyButtonEnabled =
+                                BaseFlags.get().isGridApplyButtonEnabled(appContext),
                         ),
                     snapshotRestorer = { getGridSnapshotRestorer(appContext) },
                 )
diff --git a/src/com/android/customization/picker/HorizontalTouchMovementAwareNestedScrollView.kt b/src/com/android/customization/picker/HorizontalTouchMovementAwareNestedScrollView.kt
deleted file mode 100644
index 06cf753..0000000
--- a/src/com/android/customization/picker/HorizontalTouchMovementAwareNestedScrollView.kt
+++ /dev/null
@@ -1,64 +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.picker
-
-import android.content.Context
-import android.util.AttributeSet
-import android.view.MotionEvent
-import android.view.ViewConfiguration
-import androidx.core.widget.NestedScrollView
-import kotlin.math.abs
-
-/**
- * This nested scroll view will detect horizontal touch movements and stop vertical scrolls when a
- * horizontal touch movement is detected.
- */
-class HorizontalTouchMovementAwareNestedScrollView(context: Context, attrs: AttributeSet?) :
-    NestedScrollView(context, attrs) {
-
-    private var startXPosition = 0f
-    private var startYPosition = 0f
-    private var isHorizontalTouchMovement = false
-
-    override fun onInterceptTouchEvent(event: MotionEvent): Boolean {
-        when (event.action) {
-            MotionEvent.ACTION_DOWN -> {
-                startXPosition = event.x
-                startYPosition = event.y
-                isHorizontalTouchMovement = false
-            }
-            MotionEvent.ACTION_MOVE -> {
-                val xMoveDistance = abs(event.x - startXPosition)
-                val yMoveDistance = abs(event.y - startYPosition)
-                if (
-                    !isHorizontalTouchMovement &&
-                        xMoveDistance > yMoveDistance &&
-                        xMoveDistance > ViewConfiguration.get(context).scaledTouchSlop
-                ) {
-                    isHorizontalTouchMovement = true
-                }
-            }
-            else -> {}
-        }
-        return if (isHorizontalTouchMovement) {
-            // We only want to intercept the touch event when the touch moves more vertically than
-            // horizontally. So we return false.
-            false
-        } else {
-            super.onInterceptTouchEvent(event)
-        }
-    }
-}
diff --git a/src/com/android/customization/picker/clock/data/repository/ClockPickerRepositoryImpl.kt b/src/com/android/customization/picker/clock/data/repository/ClockPickerRepositoryImpl.kt
index be6c6cb..aeed45e 100644
--- a/src/com/android/customization/picker/clock/data/repository/ClockPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/clock/data/repository/ClockPickerRepositoryImpl.kt
@@ -33,6 +33,7 @@
 import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.callbackFlow
+import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.mapLatest
@@ -56,7 +57,7 @@
                         registry
                             .getClocks()
                             .filter { "NOT_IN_USE" !in it.clockId }
-                            .map { it.toModel() }
+                            .map { it.toModel(isSelected = it.clockId == registry.currentClockId) }
                     trySend(allClocks)
                 }
 
@@ -90,6 +91,7 @@
                             .getClocks()
                             .find { clockMetadata -> clockMetadata.clockId == currentClockId }
                             ?.toModel(
+                                isSelected = true,
                                 selectedColorId = metadata?.getSelectedColorId(),
                                 colorTone = metadata?.getColorTone()
                                         ?: ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS,
@@ -146,9 +148,10 @@
             )
             .map { setting -> setting == 1 }
             .map { isDynamic -> if (isDynamic) ClockSize.DYNAMIC else ClockSize.SMALL }
+            .distinctUntilChanged()
             .shareIn(
                 scope = scope,
-                started = SharingStarted.WhileSubscribed(),
+                started = SharingStarted.Eagerly,
                 replay = 1,
             )
 
@@ -176,6 +179,7 @@
 
     /** By default, [ClockMetadataModel] has no color information unless specified. */
     private fun ClockMetadata.toModel(
+        isSelected: Boolean,
         selectedColorId: String? = null,
         @IntRange(from = 0, to = 100) colorTone: Int = 0,
         @ColorInt seedColor: Int? = null,
@@ -183,6 +187,7 @@
         return ClockMetadataModel(
             clockId = clockId,
             name = name,
+            isSelected = isSelected,
             selectedColorId = selectedColorId,
             colorToneProgress = colorTone,
             seedColor = seedColor,
diff --git a/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt b/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt
index 652ffdd..b197edf 100644
--- a/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt
+++ b/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt
@@ -29,6 +29,7 @@
 import com.android.systemui.shared.plugins.PluginManagerImpl
 import com.android.systemui.shared.plugins.PluginPrefs
 import com.android.systemui.shared.system.UncaughtExceptionPreHandlerManager_Factory
+import com.android.wallpaper.module.InjectorProvider
 import java.util.concurrent.Executors
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
@@ -55,6 +56,8 @@
             DefaultClockProvider(context, LayoutInflater.from(context), context.resources),
             keepAllLoaded = true,
             subTag = "Picker",
+            isTransitClockEnabled =
+                InjectorProvider.getInjector().getFlags().isTransitClockEnabled(context)
         )
     }
 
diff --git a/src/com/android/customization/picker/clock/shared/model/ClockMetadataModel.kt b/src/com/android/customization/picker/clock/shared/model/ClockMetadataModel.kt
index bd87ba6..2522507 100644
--- a/src/com/android/customization/picker/clock/shared/model/ClockMetadataModel.kt
+++ b/src/com/android/customization/picker/clock/shared/model/ClockMetadataModel.kt
@@ -24,6 +24,7 @@
 data class ClockMetadataModel(
     val clockId: String,
     val name: String,
+    val isSelected: Boolean,
     val selectedColorId: String?,
     @IntRange(from = 0, to = 100) val colorToneProgress: Int,
     @ColorInt val seedColor: Int?,
diff --git a/src/com/android/customization/picker/clock/ui/binder/ClockCarouselViewBinder.kt b/src/com/android/customization/picker/clock/ui/binder/ClockCarouselViewBinder.kt
index 89fac89..7880d82 100644
--- a/src/com/android/customization/picker/clock/ui/binder/ClockCarouselViewBinder.kt
+++ b/src/com/android/customization/picker/clock/ui/binder/ClockCarouselViewBinder.kt
@@ -16,8 +16,6 @@
 package com.android.customization.picker.clock.ui.binder
 
 import android.content.Context
-import android.view.ViewGroup
-import android.widget.FrameLayout
 import androidx.core.view.isVisible
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleEventObserver
@@ -27,7 +25,6 @@
 import com.android.customization.picker.clock.ui.view.ClockCarouselView
 import com.android.customization.picker.clock.ui.view.ClockViewFactory
 import com.android.customization.picker.clock.ui.viewmodel.ClockCarouselViewModel
-import com.android.wallpaper.R
 import com.android.wallpaper.picker.customization.ui.section.ScreenPreviewClickView
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.launch
@@ -38,7 +35,6 @@
     fun bind(
         context: Context,
         carouselView: ClockCarouselView,
-        singleClockView: ViewGroup,
         screenPreviewClickView: ScreenPreviewClickView,
         viewModel: ClockCarouselViewModel,
         clockViewFactory: ClockViewFactory,
@@ -46,6 +42,7 @@
         isTwoPaneAndSmallWidth: Boolean,
     ) {
         carouselView.setClockViewFactory(clockViewFactory)
+        carouselView.isVisible = true
         clockViewFactory.updateRegionDarkness()
         val carouselAccessibilityDelegate =
             CarouselAccessibilityDelegate(
@@ -61,12 +58,8 @@
             )
         screenPreviewClickView.accessibilityDelegate = carouselAccessibilityDelegate
 
-        val singleClockHostView =
-            singleClockView.requireViewById<FrameLayout>(R.id.single_clock_host_view)
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
-                launch { viewModel.isCarouselVisible.collect { carouselView.isVisible = it } }
-
                 launch {
                     combine(viewModel.selectedClockSize, viewModel.allClocks, ::Pair).collect {
                         (size, allClocks) ->
@@ -98,21 +91,6 @@
                 launch {
                     viewModel.seedColor.collect { clockViewFactory.updateColorForAllClocks(it) }
                 }
-
-                launch {
-                    viewModel.isSingleClockViewVisible.collect { singleClockView.isVisible = it }
-                }
-
-                launch {
-                    viewModel.clockId.collect { clockId ->
-                        singleClockHostView.removeAllViews()
-                        val clockView = clockViewFactory.getLargeView(clockId)
-                        // The clock view might still be attached to an existing parent. Detach
-                        // before adding to another parent.
-                        (clockView.parent as? ViewGroup)?.removeView(clockView)
-                        singleClockHostView.addView(clockView)
-                    }
-                }
             }
         }
 
diff --git a/src/com/android/customization/picker/clock/ui/binder/ClockSettingsBinder.kt b/src/com/android/customization/picker/clock/ui/binder/ClockSettingsBinder.kt
index 4f4bd1b..6e745d5 100644
--- a/src/com/android/customization/picker/clock/ui/binder/ClockSettingsBinder.kt
+++ b/src/com/android/customization/picker/clock/ui/binder/ClockSettingsBinder.kt
@@ -33,6 +33,7 @@
 import androidx.recyclerview.widget.RecyclerView
 import com.android.customization.picker.clock.shared.ClockSize
 import com.android.customization.picker.clock.ui.adapter.ClockSettingsTabAdapter
+import com.android.customization.picker.clock.ui.view.ClockCarouselView
 import com.android.customization.picker.clock.ui.view.ClockHostView
 import com.android.customization.picker.clock.ui.view.ClockSizeRadioButtonGroup
 import com.android.customization.picker.clock.ui.view.ClockViewFactory
@@ -199,7 +200,7 @@
                                     sizeOptions.radioButtonDynamic.isChecked = false
                                     sizeOptions.radioButtonSmall.isChecked = true
                                     clockHostView.doOnPreDraw {
-                                        it.pivotX = 0F
+                                        it.pivotX = ClockCarouselView.getCenteredHostViewPivotX(it)
                                         it.pivotY = 0F
                                     }
                                 }
diff --git a/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragment.kt b/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragment.kt
deleted file mode 100644
index f138d6a..0000000
--- a/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragment.kt
+++ /dev/null
@@ -1,93 +0,0 @@
-package com.android.customization.picker.clock.ui.fragment
-
-import android.content.Context
-import android.os.Bundle
-import android.util.TypedValue
-import android.view.ContextThemeWrapper
-import android.view.LayoutInflater
-import android.view.View
-import android.view.ViewGroup
-import android.view.ViewGroup.LayoutParams.MATCH_PARENT
-import android.view.ViewGroup.LayoutParams.WRAP_CONTENT
-import android.widget.FrameLayout
-import android.widget.TextView
-import android.widget.Toast
-import androidx.core.view.setPadding
-import androidx.recyclerview.widget.LinearLayoutManager
-import androidx.recyclerview.widget.RecyclerView
-import com.android.customization.module.ThemePickerInjector
-import com.android.internal.annotations.VisibleForTesting
-import com.android.systemui.plugins.ClockMetadata
-import com.android.systemui.shared.clocks.ClockRegistry
-import com.android.wallpaper.R
-import com.android.wallpaper.module.InjectorProvider
-import com.android.wallpaper.picker.AppbarFragment
-
-class ClockCustomDemoFragment : AppbarFragment() {
-    @VisibleForTesting lateinit var recyclerView: RecyclerView
-    @VisibleForTesting lateinit var clockRegistry: ClockRegistry
-
-    override fun onCreateView(
-        inflater: LayoutInflater,
-        container: ViewGroup?,
-        savedInstanceState: Bundle?
-    ): View {
-        val view = inflater.inflate(R.layout.fragment_clock_custom_picker_demo, container, false)
-        setUpToolbar(view)
-        clockRegistry =
-            (InjectorProvider.getInjector() as ThemePickerInjector).getClockRegistry(
-                requireContext(),
-            )
-        val listInUse = clockRegistry.getClocks().filter { "NOT_IN_USE" !in it.clockId }
-
-        recyclerView = view.requireViewById(R.id.clock_preview_card_list_demo)
-        recyclerView.layoutManager = LinearLayoutManager(context, RecyclerView.VERTICAL, false)
-        recyclerView.adapter =
-            ClockRecyclerAdapter(listInUse, requireContext()) {
-                clockRegistry.currentClockId = it.clockId
-                Toast.makeText(context, "${it.name} selected", Toast.LENGTH_SHORT).show()
-            }
-        return view
-    }
-
-    override fun getDefaultTitle(): CharSequence {
-        return getString(R.string.clock_title)
-    }
-
-    internal class ClockRecyclerAdapter(
-        val list: List<ClockMetadata>,
-        val context: Context,
-        val onClockSelected: (ClockMetadata) -> Unit
-    ) : RecyclerView.Adapter<ClockRecyclerAdapter.ViewHolder>() {
-        class ViewHolder(val view: View, val textView: TextView, val onItemClicked: (Int) -> Unit) :
-            RecyclerView.ViewHolder(view) {
-            init {
-                itemView.setOnClickListener { onItemClicked(absoluteAdapterPosition) }
-            }
-        }
-
-        override fun onCreateViewHolder(viewGroup: ViewGroup, viewType: Int): ViewHolder {
-            val rootView = FrameLayout(viewGroup.context)
-            val textView =
-                TextView(ContextThemeWrapper(viewGroup.context, R.style.SectionTitleTextStyle))
-            textView.setPadding(ITEM_PADDING)
-            rootView.addView(textView)
-            val outValue = TypedValue()
-            context.theme.resolveAttribute(android.R.attr.selectableItemBackground, outValue, true)
-            rootView.setBackgroundResource(outValue.resourceId)
-            val lp = RecyclerView.LayoutParams(MATCH_PARENT, WRAP_CONTENT)
-            rootView.layoutParams = lp
-            return ViewHolder(rootView, textView) { onClockSelected(list[it]) }
-        }
-
-        override fun onBindViewHolder(viewHolder: ViewHolder, position: Int) {
-            viewHolder.textView.text = list[position].name
-        }
-
-        override fun getItemCount() = list.size
-
-        companion object {
-            val ITEM_PADDING = 40
-        }
-    }
-}
diff --git a/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt b/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
index f4684d8..4805f37 100644
--- a/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
+++ b/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
@@ -20,8 +20,12 @@
 import android.view.View
 import android.view.ViewGroup
 import androidx.cardview.widget.CardView
+import androidx.core.content.ContextCompat
+import androidx.core.view.isVisible
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.get
+import androidx.transition.Transition
+import androidx.transition.doOnStart
 import com.android.customization.module.ThemePickerInjector
 import com.android.customization.picker.clock.ui.binder.ClockSettingsBinder
 import com.android.systemui.shared.clocks.shared.model.ClockPreviewConstants
@@ -130,6 +134,8 @@
             viewLifecycleOwner,
         )
 
+        (returnTransition as? Transition)?.doOnStart { lockScreenView.isVisible = false }
+
         return view
     }
 
@@ -140,4 +146,8 @@
     override fun getToolbarColorId(): Int {
         return android.R.color.transparent
     }
+
+    override fun getToolbarTextColor(): Int {
+        return ContextCompat.getColor(requireContext(), R.color.system_on_surface)
+    }
 }
diff --git a/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt b/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt
index 8764e54..2d18ab3 100644
--- a/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt
+++ b/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt
@@ -71,17 +71,31 @@
         clockViewFactory = factory
     }
 
+    // This function is for the custom accessibility action to trigger a transition to the next
+    // carousel item. If the current item is the last item in the carousel, the next item
+    // will be the first item.
     fun transitionToNext() {
-        val index = (carousel.currentIndex + 1) % carousel.count
-        if (index < carousel.count && index > 0) {
-            carousel.transitionToIndex(index, 0)
+        if (carousel.count != 0) {
+            val index = (carousel.currentIndex + 1) % carousel.count
+            carousel.jumpToIndex(index)
+            // Explicitly called this since using transitionToIndex(index) leads to
+            // race-condition between announcement of content description of the correct clock-face
+            // and the selection of clock face itself
+            adapter.onNewItem(index)
         }
     }
 
+    // This function is for the custom accessibility action to trigger a transition to
+    // the previous carousel item. If the current item is the first item in the carousel,
+    // the previous item will be the last item.
     fun transitionToPrevious() {
-        val index = (carousel.currentIndex - 1) % carousel.count
-        if (index < carousel.count && index > 0) {
-            carousel.transitionToIndex(index, 0)
+        if (carousel.count != 0) {
+            val index = (carousel.currentIndex + carousel.count - 1) % carousel.count
+            carousel.jumpToIndex(index)
+            // Explicitly called this since using transitionToIndex(index) leads to
+            // race-condition between announcement of content description of the correct clock-face
+            // and the selection of clock face itself
+            adapter.onNewItem(index)
         }
     }
 
@@ -100,8 +114,15 @@
         }
 
         adapter = ClockCarouselAdapter(clockSize, clocks, clockViewFactory, onClockSelected)
+        carousel.isInfinite = clocks.size >= MIN_CLOCKS_TO_ENABLE_INFINITE_CAROUSEL
         carousel.setAdapter(adapter)
-        carousel.refresh()
+        val indexOfSelectedClock =
+            clocks
+                .indexOfFirst { it.isSelected }
+                // If not found, default to the first clock as selected:
+                .takeIf { it != -1 }
+                ?: 0
+        carousel.jumpToIndex(indexOfSelectedClock)
         motionLayout.setTransitionListener(
             object : MotionLayout.TransitionListener {
 
@@ -211,11 +232,13 @@
                         }
                             ?: return
                     offCenterClockHostView.doOnPreDraw {
-                        it.pivotX = progress * it.width / 2
+                        it.pivotX =
+                            progress * it.width / 2 + (1 - progress) * getCenteredHostViewPivotX(it)
                         it.pivotY = progress * it.height / 2
                     }
                     toCenterClockHostView.doOnPreDraw {
-                        it.pivotX = (1 - progress) * it.width / 2
+                        it.pivotX =
+                            (1 - progress) * it.width / 2 + progress * getCenteredHostViewPivotX(it)
                         it.pivotY = (1 - progress) * it.height / 2
                     }
                     offCenterClockFrame.translationX =
@@ -265,9 +288,11 @@
     fun setSelectedClockIndex(
         index: Int,
     ) {
-        // jumpToIndex to the same position can cause the views unnecessarily populate again.
-        // Only call jumpToIndex when the jump-to index is different from the current carousel.
-        if (index != carousel.currentIndex) {
+        // 1. setUpClockCarouselView() can possibly not be called before setSelectedClockIndex().
+        //    We need to check if index out of bound.
+        // 2. jumpToIndex() to the same position can cause the views unnecessarily populate again.
+        //    We only call jumpToIndex when the index is different from the current carousel.
+        if (index < carousel.count && index != carousel.currentIndex) {
             carousel.jumpToIndex(index)
         }
     }
@@ -417,7 +442,7 @@
         ) {
             clockHostView.doOnPreDraw {
                 if (isMiddleView) {
-                    it.pivotX = 0F
+                    it.pivotX = getCenteredHostViewPivotX(it)
                     it.pivotY = 0F
                     clockView.translationX = 0F
                     clockView.translationY = 0F
@@ -446,6 +471,8 @@
     }
 
     companion object {
+        // The carousel needs to have at least 5 different clock faces to be infinite
+        const val MIN_CLOCKS_TO_ENABLE_INFINITE_CAROUSEL = 5
         const val CLOCK_CAROUSEL_VIEW_SCALE = 0.5f
 
         val itemViewIds =
@@ -507,6 +534,10 @@
             return rootViewId == R.id.item_view_2
         }
 
+        fun getCenteredHostViewPivotX(hostView: View): Float {
+            return if (hostView.isLayoutRtl) hostView.width.toFloat() else 0F
+        }
+
         private fun getTranslationDistance(
             hostLength: Int,
             frameLength: Int,
diff --git a/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselItemViewModel.kt b/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselItemViewModel.kt
index 708fa2f..9811426 100644
--- a/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselItemViewModel.kt
+++ b/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselItemViewModel.kt
@@ -20,15 +20,14 @@
 import com.android.wallpaper.R
 import com.android.wallpaper.module.InjectorProvider
 
-class ClockCarouselItemViewModel(val clockId: String) {
+class ClockCarouselItemViewModel(val clockId: String, val isSelected: Boolean) {
 
     /** Description for accessibility purposes when a clock is selected. */
     fun getContentDescription(resources: Resources): String {
         val clockContent =
             (InjectorProvider.getInjector() as? CustomizationInjector)
-                ?.getClockDescriptionUtils()
-                ?.getDescriptionResId(clockId)
-                ?.let { resources.getString(it) }
+                ?.getClockDescriptionUtils(resources)
+                ?.getDescription(clockId)
                 ?: ""
         return resources.getString(R.string.select_clock_action_description, clockContent)
     }
diff --git a/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModel.kt b/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModel.kt
index a4f9cc4..b2dfa61 100644
--- a/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModel.kt
+++ b/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModel.kt
@@ -27,7 +27,6 @@
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.mapLatest
@@ -39,8 +38,7 @@
  * Clock carousel view model that provides data for the carousel of clock previews. When there is
  * only one item, we should show a single clock preview instead of a carousel.
  */
-class ClockCarouselViewModel
-constructor(
+class ClockCarouselViewModel(
     private val interactor: ClockPickerInteractor,
     private val backgroundDispatcher: CoroutineDispatcher,
 ) : ViewModel() {
@@ -50,7 +48,7 @@
             .mapLatest { allClocks ->
                 // Delay to avoid the case that the full list of clocks is not initiated.
                 delay(CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
-                allClocks.map { ClockCarouselItemViewModel(it.clockId) }
+                allClocks.map { ClockCarouselItemViewModel(it.clockId, it.isSelected) }
             }
             .stateIn(viewModelScope, SharingStarted.Eagerly, emptyList())
 
@@ -58,8 +56,6 @@
 
     val seedColor: Flow<Int?> = interactor.seedColor
 
-    val isCarouselVisible: Flow<Boolean> = allClocks.map { it.size > 1 }.distinctUntilChanged()
-
     @OptIn(ExperimentalCoroutinesApi::class)
     val selectedIndex: Flow<Int> =
         allClocks
@@ -77,15 +73,6 @@
             }
             .mapNotNull { it }
 
-    // Handle the case when there is only one clock in the carousel
-    val isSingleClockViewVisible: Flow<Boolean> =
-        allClocks.map { it.size == 1 }.distinctUntilChanged()
-
-    val clockId: Flow<String> =
-        allClocks
-            .map { allClockIds -> if (allClockIds.size == 1) allClockIds[0].clockId else null }
-            .mapNotNull { it }
-
     private var setSelectedClockJob: Job? = null
     fun setSelectedClock(clockId: String) {
         setSelectedClockJob?.cancel()
diff --git a/src/com/android/customization/picker/clock/utils/ClockDescriptionUtils.kt b/src/com/android/customization/picker/clock/utils/ClockDescriptionUtils.kt
index 9a0b66f..28ea4a3 100644
--- a/src/com/android/customization/picker/clock/utils/ClockDescriptionUtils.kt
+++ b/src/com/android/customization/picker/clock/utils/ClockDescriptionUtils.kt
@@ -15,14 +15,12 @@
  */
 package com.android.customization.picker.clock.utils
 
-import androidx.annotation.StringRes
-
 /** Provides clock description for accessibility purposes. */
 interface ClockDescriptionUtils {
 
     /**
-     * TODO (b/287507746) : Migrate description res ID to system UI or a shared library, instead of
-     * preserving the clock description at the Wallpaper Picker side.
+     * TODO (b/287507746) : Migrate the clock description to system UI or a shared library, instead
+     * of preserving at the Wallpaper Picker side.
      */
-    @StringRes fun getDescriptionResId(clockId: String): Int
+    fun getDescription(clockId: String): String
 }
diff --git a/src/com/android/customization/picker/clock/utils/ThemePickerClockDescriptionUtils.kt b/src/com/android/customization/picker/clock/utils/ThemePickerClockDescriptionUtils.kt
index 43f19b3..a04ebff 100644
--- a/src/com/android/customization/picker/clock/utils/ThemePickerClockDescriptionUtils.kt
+++ b/src/com/android/customization/picker/clock/utils/ThemePickerClockDescriptionUtils.kt
@@ -15,12 +15,8 @@
  */
 package com.android.customization.picker.clock.utils
 
-import androidx.annotation.StringRes
-import com.android.wallpaper.R
-
 class ThemePickerClockDescriptionUtils : ClockDescriptionUtils {
-    @StringRes
-    override fun getDescriptionResId(clockId: String): Int {
-        return R.string.clock_title
+    override fun getDescription(clockId: String): String {
+        return ""
     }
 }
diff --git a/src/com/android/customization/picker/color/ColorSectionView.java b/src/com/android/customization/picker/color/ColorSectionView.java
deleted file mode 100644
index b8ba2e4..0000000
--- a/src/com/android/customization/picker/color/ColorSectionView.java
+++ /dev/null
@@ -1,33 +0,0 @@
-/*
- * Copyright (C) 2022 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;
-
-import android.content.Context;
-import android.util.AttributeSet;
-
-import androidx.annotation.Nullable;
-
-import com.android.wallpaper.picker.SectionView;
-
-/**
- * The class inherits from {@link SectionView} as the view representing the color section of the
- * customization picker.
- */
-public final class ColorSectionView extends SectionView {
-    public ColorSectionView(Context context, @Nullable AttributeSet attrs) {
-        super(context, attrs);
-    }
-}
diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
index 7cf9fd0..fccaa65 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
@@ -19,12 +19,15 @@
 import com.android.customization.picker.color.shared.model.ColorOptionModel
 import com.android.customization.picker.color.shared.model.ColorType
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.StateFlow
 
 /**
  * Abstracts access to application state related to functionality for selecting, picking, or setting
  * system color.
  */
 interface ColorPickerRepository {
+    /** Whether the system color is in the process of being updated */
+    val isApplyingSystemColor: StateFlow<Boolean>
 
     /** List of wallpaper and preset color options on the device, categorized by Color Type */
     val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>>
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 41ef3a5..ce5c0c4 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -27,7 +27,9 @@
 import com.android.wallpaper.model.WallpaperColorsModel
 import com.android.wallpaper.model.WallpaperColorsViewModel
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.suspendCancellableCoroutine
@@ -43,9 +45,17 @@
         wallpaperColorsViewModel.homeWallpaperColors
     private val lockWallpaperColors: StateFlow<WallpaperColorsModel?> =
         wallpaperColorsViewModel.lockWallpaperColors
+    private var selectedColorOption: MutableStateFlow<ColorOptionModel> =
+        MutableStateFlow(getCurrentColorOption())
+
+    private val _isApplyingSystemColor = MutableStateFlow(false)
+    override val isApplyingSystemColor = _isApplyingSystemColor.asStateFlow()
 
     override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
-        combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
+        combine(homeWallpaperColors, lockWallpaperColors, selectedColorOption) {
+                homeColors,
+                lockColors,
+                _ ->
                 homeColors to lockColors
             }
             .map { (homeColors, lockColors) ->
@@ -109,17 +119,21 @@
                 }
             }
 
-    override suspend fun select(colorOptionModel: ColorOptionModel) =
+    override suspend fun select(colorOptionModel: ColorOptionModel) {
+        _isApplyingSystemColor.value = true
         suspendCancellableCoroutine { continuation ->
             colorManager.apply(
                 colorOptionModel.colorOption,
                 object : CustomizationManager.Callback {
                     override fun onSuccess() {
+                        _isApplyingSystemColor.value = false
+                        selectedColorOption.value = colorOptionModel
                         continuation.resumeWith(Result.success(Unit))
                     }
 
                     override fun onError(throwable: Throwable?) {
                         Log.w(TAG, "Apply theme with error", throwable)
+                        _isApplyingSystemColor.value = false
                         continuation.resumeWith(
                             Result.failure(throwable ?: Throwable("Error loading theme bundles"))
                         )
@@ -127,6 +141,7 @@
                 }
             )
         }
+    }
 
     override fun getCurrentColorOption(): ColorOptionModel {
         val overlays = colorManager.currentOverlays
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 714129d..bb2ef9d 100644
--- a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
@@ -29,6 +29,9 @@
 
 class FakeColorPickerRepository(private val context: Context) : ColorPickerRepository {
 
+    private val _isApplyingSystemColor = MutableStateFlow(false)
+    override val isApplyingSystemColor = _isApplyingSystemColor.asStateFlow()
+
     private lateinit var selectedColorOption: ColorOptionModel
 
     private val _colorOptions =
diff --git a/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt b/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
index 8c7a4b7..d3b2eba 100644
--- a/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
+++ b/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
@@ -26,6 +26,8 @@
     private val repository: ColorPickerRepository,
     private val snapshotRestorer: Provider<ColorPickerSnapshotRestorer>,
 ) {
+    val isApplyingSystemColor = repository.isApplyingSystemColor
+
     /**
      * The newly selected color option for overwriting the current active option during an
      * optimistic update, the value is set to null when update fails
diff --git a/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt b/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt
index 2daefe4..ad81614 100644
--- a/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt
+++ b/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt
@@ -112,22 +112,20 @@
             val optionSelectedView = itemView.requireViewById<ImageView>(R.id.option_selected)
 
             lifecycleOwner.lifecycleScope.launch {
-                lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
-                    launch {
-                        item.isSelected.collect { isSelected ->
-                            optionSelectedView.isVisible = isSelected
-                        }
+                launch {
+                    item.isSelected.collect { isSelected ->
+                        optionSelectedView.isVisible = isSelected
                     }
-                    launch {
-                        item.onClicked.collect { onClicked ->
-                            itemView.setOnClickListener(
-                                if (onClicked != null) {
-                                    View.OnClickListener { onClicked.invoke() }
-                                } else {
-                                    null
-                                }
-                            )
-                        }
+                }
+                launch {
+                    item.onClicked.collect { onClicked ->
+                        itemView.setOnClickListener(
+                            if (onClicked != null) {
+                                View.OnClickListener { onClicked.invoke() }
+                            } else {
+                                null
+                            }
+                        )
                     }
                 }
             }
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 78bfa43..4ef29d6 100644
--- a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
+++ b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
@@ -22,9 +22,13 @@
 import android.view.ViewGroup
 import android.widget.FrameLayout
 import androidx.cardview.widget.CardView
+import androidx.core.content.ContextCompat
+import androidx.core.view.isVisible
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.get
 import androidx.lifecycle.lifecycleScope
+import androidx.transition.Transition
+import androidx.transition.doOnStart
 import com.android.customization.model.mode.DarkModeSectionController
 import com.android.customization.module.ThemePickerInjector
 import com.android.customization.picker.color.ui.binder.ColorPickerBinder
@@ -92,53 +96,58 @@
 
         savedInstanceState?.let { binding?.restoreInstanceState(it) }
 
-        ScreenPreviewBinder.bind(
-            activity = requireActivity(),
-            previewView = lockScreenView,
-            viewModel =
-                ScreenPreviewViewModel(
-                    previewUtils =
-                        PreviewUtils(
-                            context = requireContext(),
-                            authority =
-                                requireContext()
-                                    .getString(
-                                        R.string.lock_screen_preview_provider_authority,
-                                    ),
-                        ),
-                    wallpaperInfoProvider = { forceReload ->
-                        suspendCancellableCoroutine { continuation ->
-                            wallpaperInfoFactory.createCurrentWallpaperInfos(
-                                { homeWallpaper, lockWallpaper, _ ->
-                                    lifecycleScope.launch {
-                                        if (
-                                            wcViewModel.lockWallpaperColors.value
-                                                is WallpaperColorsModel.Loading
-                                        ) {
-                                            loadInitialColors(
-                                                wallpaperManager,
-                                                wcViewModel,
-                                                CustomizationSections.Screen.LOCK_SCREEN
-                                            )
+        val lockScreenPreviewBinder =
+            ScreenPreviewBinder.bind(
+                activity = requireActivity(),
+                previewView = lockScreenView,
+                viewModel =
+                    ScreenPreviewViewModel(
+                        previewUtils =
+                            PreviewUtils(
+                                context = requireContext(),
+                                authority =
+                                    requireContext()
+                                        .getString(
+                                            R.string.lock_screen_preview_provider_authority,
+                                        ),
+                            ),
+                        wallpaperInfoProvider = { forceReload ->
+                            suspendCancellableCoroutine { continuation ->
+                                wallpaperInfoFactory.createCurrentWallpaperInfos(
+                                    { homeWallpaper, lockWallpaper, _ ->
+                                        lifecycleScope.launch {
+                                            if (
+                                                wcViewModel.lockWallpaperColors.value
+                                                    is WallpaperColorsModel.Loading
+                                            ) {
+                                                loadInitialColors(
+                                                    wallpaperManager,
+                                                    wcViewModel,
+                                                    CustomizationSections.Screen.LOCK_SCREEN
+                                                )
+                                            }
                                         }
-                                    }
-                                    continuation.resume(lockWallpaper ?: homeWallpaper, null)
-                                },
-                                forceReload,
-                            )
-                        }
-                    },
-                    onWallpaperColorChanged = { colors ->
-                        wcViewModel.setLockWallpaperColors(colors)
-                    },
-                    wallpaperInteractor = injector.getWallpaperInteractor(requireContext()),
-                    screen = CustomizationSections.Screen.LOCK_SCREEN,
-                ),
-            lifecycleOwner = this,
-            offsetToStart =
-                displayUtils.isSingleDisplayOrUnfoldedHorizontalHinge(requireActivity()),
-            onWallpaperPreviewDirty = { activity?.recreate() },
-        )
+                                        continuation.resume(lockWallpaper ?: homeWallpaper, null)
+                                    },
+                                    forceReload,
+                                )
+                            }
+                        },
+                        onWallpaperColorChanged = { colors ->
+                            wcViewModel.setLockWallpaperColors(colors)
+                        },
+                        wallpaperInteractor = injector.getWallpaperInteractor(requireContext()),
+                        screen = CustomizationSections.Screen.LOCK_SCREEN,
+                    ),
+                lifecycleOwner = this,
+                offsetToStart =
+                    displayUtils.isSingleDisplayOrUnfoldedHorizontalHinge(requireActivity()),
+                onWallpaperPreviewDirty = { activity?.recreate() },
+            )
+        val shouldMirrorHomePreview =
+            wallpaperManager.getWallpaperInfo(WallpaperManager.FLAG_SYSTEM) != null &&
+                wallpaperManager.getWallpaperId(WallpaperManager.FLAG_LOCK) < 0
+        val mirrorSurface = if (shouldMirrorHomePreview) lockScreenPreviewBinder.surface() else null
         ScreenPreviewBinder.bind(
             activity = requireActivity(),
             previewView = homeScreenView,
@@ -185,6 +194,7 @@
             offsetToStart =
                 displayUtils.isSingleDisplayOrUnfoldedHorizontalHinge(requireActivity()),
             onWallpaperPreviewDirty = { activity?.recreate() },
+            mirrorSurface = mirrorSurface,
         )
         val darkModeToggleContainerView: FrameLayout =
             view.requireViewById(R.id.dark_mode_toggle_container)
@@ -197,6 +207,12 @@
                 .createView(requireContext())
         darkModeSectionView.background = null
         darkModeToggleContainerView.addView(darkModeSectionView)
+
+        (returnTransition as? Transition)?.doOnStart {
+            lockScreenView.isVisible = false
+            homeScreenView.isVisible = false
+        }
+
         return view
     }
 
@@ -236,4 +252,8 @@
     override fun getToolbarColorId(): Int {
         return android.R.color.transparent
     }
+
+    override fun getToolbarTextColor(): Int {
+        return ContextCompat.getColor(requireContext(), R.color.system_on_surface)
+    }
 }
diff --git a/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt b/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
index a828f83..71dfe1d 100644
--- a/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
+++ b/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
@@ -23,7 +23,6 @@
 import android.view.TouchDelegate
 import android.view.View
 import android.view.View.OnAttachStateChangeListener
-import android.view.ViewGroup
 import android.view.ViewStub
 import androidx.activity.ComponentActivity
 import androidx.constraintlayout.helper.widget.Carousel
@@ -39,7 +38,9 @@
 import com.android.customization.picker.clock.ui.view.ClockCarouselView
 import com.android.customization.picker.clock.ui.view.ClockViewFactory
 import com.android.customization.picker.clock.ui.viewmodel.ClockCarouselViewModel
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.wallpaper.R
+import com.android.wallpaper.model.CustomizationSectionController
 import com.android.wallpaper.model.CustomizationSectionController.CustomizationSectionNavigationController
 import com.android.wallpaper.model.WallpaperColorsViewModel
 import com.android.wallpaper.model.WallpaperPreviewNavigator
@@ -49,6 +50,7 @@
 import com.android.wallpaper.picker.customization.ui.section.ScreenPreviewClickView
 import com.android.wallpaper.picker.customization.ui.section.ScreenPreviewSectionController
 import com.android.wallpaper.picker.customization.ui.section.ScreenPreviewView
+import com.android.wallpaper.picker.customization.ui.viewmodel.CustomizationPickerViewModel
 import com.android.wallpaper.util.DisplayUtils
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.launch
@@ -70,8 +72,10 @@
     private val navigationController: CustomizationSectionNavigationController,
     wallpaperInteractor: WallpaperInteractor,
     themedIconInteractor: ThemedIconInteractor,
+    colorPickerInteractor: ColorPickerInteractor,
     wallpaperManager: WallpaperManager,
     private val isTwoPaneAndSmallWidth: Boolean,
+    customizationPickerViewModel: CustomizationPickerViewModel,
 ) :
     PreviewWithThemeSectionController(
         activity,
@@ -83,8 +87,10 @@
         wallpaperPreviewNavigator,
         wallpaperInteractor,
         themedIconInteractor,
+        colorPickerInteractor,
         wallpaperManager,
         isTwoPaneAndSmallWidth,
+        customizationPickerViewModel,
     ) {
 
     private val viewModel =
@@ -98,8 +104,11 @@
 
     override val hideLockScreenClockPreview = true
 
-    override fun createView(context: Context): ScreenPreviewView {
-        val view = super.createView(context)
+    override fun createView(
+        context: Context,
+        params: CustomizationSectionController.ViewCreationParams,
+    ): ScreenPreviewView {
+        val view = super.createView(context, params)
         if (screen == CustomizationSections.Screen.LOCK_SCREEN) {
             val screenPreviewClickView: ScreenPreviewClickView =
                 view.requireViewById(R.id.screen_preview_click_view)
@@ -146,12 +155,6 @@
                 guidelineEnd.layoutParams = layoutParams
             }
 
-            // TODO (b/270716937) We should handle the single clock case in the clock carousel
-            // itself
-            val singleClockViewStub: ViewStub = view.requireViewById(R.id.single_clock_view_stub)
-            singleClockViewStub.layoutResource = R.layout.single_clock_view
-            val singleClockView = singleClockViewStub.inflate() as ViewGroup
-
             /**
              * Only bind after [Carousel.onAttachedToWindow]. This is to avoid the race condition
              * that the flow emits before attached to window where [Carousel.mMotionLayout] is still
@@ -167,7 +170,6 @@
                                 ClockCarouselViewBinder.bind(
                                     context = context,
                                     carouselView = carouselView,
-                                    singleClockView = singleClockView,
                                     screenPreviewClickView = screenPreviewClickView,
                                     viewModel = viewModel,
                                     clockViewFactory = clockViewFactory,
diff --git a/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt b/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
index 56c6c30..c4d6be4 100644
--- a/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
+++ b/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
@@ -22,6 +22,7 @@
 import android.content.Context
 import androidx.lifecycle.LifecycleOwner
 import com.android.customization.model.themedicon.domain.interactor.ThemedIconInteractor
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.customization.picker.preview.ui.viewmodel.PreviewWithThemeViewModel
 import com.android.wallpaper.R
 import com.android.wallpaper.model.WallpaperColorsViewModel
@@ -30,6 +31,7 @@
 import com.android.wallpaper.module.CustomizationSections
 import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
 import com.android.wallpaper.picker.customization.ui.section.ScreenPreviewSectionController
+import com.android.wallpaper.picker.customization.ui.viewmodel.CustomizationPickerViewModel
 import com.android.wallpaper.picker.customization.ui.viewmodel.ScreenPreviewViewModel
 import com.android.wallpaper.util.DisplayUtils
 import com.android.wallpaper.util.PreviewUtils
@@ -49,8 +51,10 @@
     wallpaperPreviewNavigator: WallpaperPreviewNavigator,
     private val wallpaperInteractor: WallpaperInteractor,
     private val themedIconInteractor: ThemedIconInteractor,
+    private val colorPickerInteractor: ColorPickerInteractor,
     wallpaperManager: WallpaperManager,
     isTwoPaneAndSmallWidth: Boolean,
+    customizationPickerViewModel: CustomizationPickerViewModel,
 ) :
     ScreenPreviewSectionController(
         activity,
@@ -62,7 +66,8 @@
         wallpaperPreviewNavigator,
         wallpaperInteractor,
         wallpaperManager,
-        isTwoPaneAndSmallWidth
+        isTwoPaneAndSmallWidth,
+        customizationPickerViewModel,
     ) {
     override fun createScreenPreviewViewModel(context: Context): ScreenPreviewViewModel {
         return PreviewWithThemeViewModel(
@@ -114,6 +119,7 @@
             initialExtrasProvider = { getInitialExtras(isOnLockScreen) },
             wallpaperInteractor = wallpaperInteractor,
             themedIconInteractor = themedIconInteractor,
+            colorPickerInteractor = colorPickerInteractor,
             screen = screen,
         )
     }
diff --git a/src/com/android/customization/picker/preview/ui/viewmodel/PreviewWithThemeViewModel.kt b/src/com/android/customization/picker/preview/ui/viewmodel/PreviewWithThemeViewModel.kt
index 435878d..83f986d 100644
--- a/src/com/android/customization/picker/preview/ui/viewmodel/PreviewWithThemeViewModel.kt
+++ b/src/com/android/customization/picker/preview/ui/viewmodel/PreviewWithThemeViewModel.kt
@@ -20,12 +20,14 @@
 import android.app.WallpaperColors
 import android.os.Bundle
 import com.android.customization.model.themedicon.domain.interactor.ThemedIconInteractor
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.wallpaper.model.WallpaperInfo
 import com.android.wallpaper.module.CustomizationSections
 import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
 import com.android.wallpaper.picker.customization.ui.viewmodel.ScreenPreviewViewModel
 import com.android.wallpaper.util.PreviewUtils
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.combine
 
 /** A ThemePicker version of the [ScreenPreviewViewModel] */
 class PreviewWithThemeViewModel(
@@ -35,6 +37,7 @@
     onWallpaperColorChanged: (WallpaperColors?) -> Unit = {},
     wallpaperInteractor: WallpaperInteractor,
     private val themedIconInteractor: ThemedIconInteractor? = null,
+    colorPickerInteractor: ColorPickerInteractor? = null,
     screen: CustomizationSections.Screen,
 ) :
     ScreenPreviewViewModel(
@@ -46,4 +49,16 @@
         screen,
     ) {
     override fun workspaceUpdateEvents(): Flow<Boolean>? = themedIconInteractor?.isActivated
+
+    private val wallpaperIsLoading = super.isLoading
+
+    override val isLoading: Flow<Boolean> =
+        colorPickerInteractor?.let {
+            combine(wallpaperIsLoading, colorPickerInteractor.isApplyingSystemColor) {
+                wallpaperIsLoading,
+                colorIsLoading ->
+                wallpaperIsLoading || colorIsLoading
+            }
+        }
+            ?: wallpaperIsLoading
 }
diff --git a/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordancePickerBinder.kt b/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordancePickerBinder.kt
index 68367c8..091f484 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordancePickerBinder.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordancePickerBinder.kt
@@ -37,6 +37,7 @@
 import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
 import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.collectIndexed
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flatMapLatest
@@ -99,13 +100,18 @@
                                 selectedFlags.indexOfFirst { it }
                             }
                         }
-                        .collect { selectedPosition ->
+                        .collectIndexed { index, selectedPosition ->
                             // Scroll the view to show the first selected affordance.
                             if (selectedPosition != -1) {
                                 // We use "post" because we need to give the adapter item a pass to
                                 // update the view.
                                 affordancesView.post {
-                                    affordancesView.smoothScrollToPosition(selectedPosition)
+                                    if (index == 0) {
+                                        // don't animate on initial collection
+                                        affordancesView.scrollToPosition(selectedPosition)
+                                    } else {
+                                        affordancesView.smoothScrollToPosition(selectedPosition)
+                                    }
                                 }
                             }
                         }
diff --git a/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordanceSectionViewBinder.kt b/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordanceSectionViewBinder.kt
index 28ad51a..7e1f4d3 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordanceSectionViewBinder.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordanceSectionViewBinder.kt
@@ -48,7 +48,7 @@
 
         lifecycleOwner.lifecycleScope.launch {
             viewModel.summary
-                .flowWithLifecycle(lifecycleOwner.lifecycle, Lifecycle.State.RESUMED)
+                .flowWithLifecycle(lifecycleOwner.lifecycle, Lifecycle.State.STARTED)
                 .collectLatest { summary ->
                     TextViewBinder.bind(
                         view = descriptionView,
diff --git a/src/com/android/customization/picker/quickaffordance/ui/fragment/KeyguardQuickAffordancePickerFragment.kt b/src/com/android/customization/picker/quickaffordance/ui/fragment/KeyguardQuickAffordancePickerFragment.kt
index d5f0d33..467e5a0 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/fragment/KeyguardQuickAffordancePickerFragment.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/fragment/KeyguardQuickAffordancePickerFragment.kt
@@ -21,8 +21,12 @@
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
+import androidx.core.content.ContextCompat
+import androidx.core.view.isVisible
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.get
+import androidx.transition.Transition
+import androidx.transition.doOnStart
 import com.android.customization.module.ThemePickerInjector
 import com.android.customization.picker.quickaffordance.ui.binder.KeyguardQuickAffordancePickerBinder
 import com.android.customization.picker.quickaffordance.ui.binder.KeyguardQuickAffordancePreviewBinder
@@ -30,9 +34,7 @@
 import com.android.wallpaper.R
 import com.android.wallpaper.module.InjectorProvider
 import com.android.wallpaper.picker.AppbarFragment
-import kotlinx.coroutines.ExperimentalCoroutinesApi
 
-@OptIn(ExperimentalCoroutinesApi::class)
 class KeyguardQuickAffordancePickerFragment : AppbarFragment() {
     companion object {
         const val DESTINATION_ID = "quick_affordances"
@@ -77,6 +79,12 @@
             viewModel = viewModel,
             lifecycleOwner = this,
         )
+        postponeEnterTransition()
+        view.post { startPostponedEnterTransition() }
+        (returnTransition as? Transition)?.doOnStart {
+            // Hide preview during exit transition animation
+            view?.findViewById<View>(R.id.preview)?.isVisible = false
+        }
         return view
     }
 
@@ -87,4 +95,8 @@
     override fun getToolbarColorId(): Int {
         return android.R.color.transparent
     }
+
+    override fun getToolbarTextColor(): Int {
+        return ContextCompat.getColor(requireContext(), R.color.system_on_surface)
+    }
 }
diff --git a/src/com/android/customization/widget/GridTileDrawable.java b/src/com/android/customization/widget/GridTileDrawable.java
index 83cd0b5..b9d2036 100644
--- a/src/com/android/customization/widget/GridTileDrawable.java
+++ b/src/com/android/customization/widget/GridTileDrawable.java
@@ -55,8 +55,8 @@
         for (int r = 0; r < mRows; r++) {
             for (int c = 0; c < mCols; c++) {
                 int saveCount = canvas.save();
-                float x = (float) ((r * size / mRows) + SPACE_BETWEEN_ICONS);
-                float y = (float) ((c * size / mCols) + SPACE_BETWEEN_ICONS);
+                float y = (float) ((r * size / mRows) + SPACE_BETWEEN_ICONS);
+                float x = (float) ((c * size / mCols) + SPACE_BETWEEN_ICONS);
                 canvas.translate(x, y);
                 canvas.drawPath(mTransformedPath, mPaint);
                 canvas.restoreToCount(saveCount);
diff --git a/src/com/android/customization/picker/CustomizationPickerApplication.java b/src_override/com/android/customization/picker/CustomizationPickerApplication.java
similarity index 72%
rename from src/com/android/customization/picker/CustomizationPickerApplication.java
rename to src_override/com/android/customization/picker/CustomizationPickerApplication.java
index 178cfbf..08303d9 100644
--- a/src/com/android/customization/picker/CustomizationPickerApplication.java
+++ b/src_override/com/android/customization/picker/CustomizationPickerApplication.java
@@ -20,12 +20,23 @@
 import com.android.customization.module.ThemePickerInjector;
 import com.android.wallpaper.module.InjectorProvider;
 
-public class CustomizationPickerApplication extends Application {
+import dagger.hilt.android.HiltAndroidApp;
+
+import javax.inject.Inject;
+
+/**
+ * Application subclass that initializes the injector.
+ */
+@HiltAndroidApp(Application.class)
+public class CustomizationPickerApplication extends Hilt_CustomizationPickerApplication {
+
+    @Inject ThemePickerInjector mInjector;
+
     @Override
     public void onCreate() {
         super.onCreate();
 
         // Initialize the injector.
-        InjectorProvider.setInjector(new ThemePickerInjector());
+        InjectorProvider.setInjector(mInjector);
     }
 }
diff --git a/tests/common/Android.bp b/tests/common/Android.bp
new file mode 100644
index 0000000..2904fad
--- /dev/null
+++ b/tests/common/Android.bp
@@ -0,0 +1,39 @@
+// 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.
+//
+
+//
+// Build rule for WallpaperPicker2 tests
+//
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+android_library {
+    name: "ThemePickerTestLib",
+
+    defaults: ["ThemePicker_defaults"],
+    srcs: [
+        "src/com/android/customization/testing/**/*.java",
+        "src/com/android/customization/testing/**/*.kt",
+    ],
+    static_libs: [
+        "WallpaperPicker2TestLib",
+        "androidx.annotation_annotation",
+        "kotlinx_coroutines_test",
+        "truth-prebuilt",
+    ],
+
+    platform_apis: true,
+}
diff --git a/tests/common/AndroidManifest.xml b/tests/common/AndroidManifest.xml
new file mode 100644
index 0000000..81a2f03
--- /dev/null
+++ b/tests/common/AndroidManifest.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- 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.
+-->
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+    package="com.android.wallpaper">
+
+  <application>
+  </application>
+</manifest>
diff --git a/tests/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java b/tests/common/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java
similarity index 61%
rename from tests/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java
rename to tests/common/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java
index bcf5a5f..81890f0 100644
--- a/tests/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java
+++ b/tests/common/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java
@@ -1,3 +1,18 @@
+/*
+ * 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.
+ */
 package com.android.customization.testing;
 
 import android.content.Context;
diff --git a/tests/common/src/com/android/customization/testing/TestDrawableLayerResolver.java b/tests/common/src/com/android/customization/testing/TestDrawableLayerResolver.java
new file mode 100644
index 0000000..8b16299
--- /dev/null
+++ b/tests/common/src/com/android/customization/testing/TestDrawableLayerResolver.java
@@ -0,0 +1,31 @@
+/*
+ * 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.
+ */
+package com.android.customization.testing;
+
+import android.graphics.drawable.Drawable;
+import android.graphics.drawable.LayerDrawable;
+
+import com.android.wallpaper.module.DrawableLayerResolver;
+
+/**
+ * Test implementation of {@link DrawableLayerResolver}.
+ */
+public class TestDrawableLayerResolver implements DrawableLayerResolver {
+    @Override
+    public Drawable resolveLayer(LayerDrawable layerDrawable) {
+        return layerDrawable.getDrawable(0);
+    }
+}
diff --git a/tests/common/src/com/android/customization/testing/TestPackageStatusNotifier.java b/tests/common/src/com/android/customization/testing/TestPackageStatusNotifier.java
new file mode 100644
index 0000000..2aadae8
--- /dev/null
+++ b/tests/common/src/com/android/customization/testing/TestPackageStatusNotifier.java
@@ -0,0 +1,33 @@
+/*
+ * 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.
+ */
+package com.android.customization.testing;
+
+import com.android.wallpaper.module.PackageStatusNotifier;
+
+/**
+ * Test implementation of {@link PackageStatusNotifier}.
+ */
+public class TestPackageStatusNotifier implements PackageStatusNotifier {
+    @Override
+    public void addListener(Listener listener, String action) {
+        // Do nothing
+    }
+
+    @Override
+    public void removeListener(Listener listener) {
+        // Do nothing
+    }
+}
diff --git a/tests/src/com/android/customization/testing/TestPluginManager.kt b/tests/common/src/com/android/customization/testing/TestPluginManager.kt
similarity index 100%
rename from tests/src/com/android/customization/testing/TestPluginManager.kt
rename to tests/common/src/com/android/customization/testing/TestPluginManager.kt
diff --git a/tests/src/com/android/customization/testing/TestThemeManager.java b/tests/common/src/com/android/customization/testing/TestThemeManager.java
similarity index 60%
rename from tests/src/com/android/customization/testing/TestThemeManager.java
rename to tests/common/src/com/android/customization/testing/TestThemeManager.java
index c4d25fb..5175b24 100644
--- a/tests/src/com/android/customization/testing/TestThemeManager.java
+++ b/tests/common/src/com/android/customization/testing/TestThemeManager.java
@@ -1,3 +1,18 @@
+/*
+ * 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.
+ */
 package com.android.customization.testing;
 
 import androidx.fragment.app.FragmentActivity;
diff --git a/tests/src/com/android/customization/testing/TestThemesUserEventLogger.java b/tests/common/src/com/android/customization/testing/TestThemesUserEventLogger.java
similarity index 63%
rename from tests/src/com/android/customization/testing/TestThemesUserEventLogger.java
rename to tests/common/src/com/android/customization/testing/TestThemesUserEventLogger.java
index 2bb2082..22a5b94 100644
--- a/tests/src/com/android/customization/testing/TestThemesUserEventLogger.java
+++ b/tests/common/src/com/android/customization/testing/TestThemesUserEventLogger.java
@@ -1,3 +1,18 @@
+/*
+ * 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.
+ */
 package com.android.customization.testing;
 
 import com.android.customization.model.color.ColorOption;
diff --git a/tests/robotests/Android.bp b/tests/robotests/Android.bp
index 4416e1c..ed684f4 100644
--- a/tests/robotests/Android.bp
+++ b/tests/robotests/Android.bp
@@ -2,13 +2,23 @@
 package {
     default_applicable_licenses: ["Android-Apache-2.0"],
 }
+
 android_robolectric_test {
     name: "ThemePickerRoboTests",
     srcs: [
         "src/**/*.java",
         "src/**/*.kt",
     ],
+    // TODO(b/291104503) Enable this test
+    exclude_srcs: ["src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt"],
     java_resource_dirs: ["config"],
+    static_libs: [
+        "WallpaperPicker2TestLib",
+        "androidx.test.rules",
+        "junit",
+        "kotlinx_coroutines_test",
+        "truth-prebuilt",
+    ],
     libs: [
         "androidx.test.core",
         "androidx.test.runner",
diff --git a/tests/robotests/src/com/android/customization/model/color/ColorSectionControllerTest.java b/tests/robotests/src/com/android/customization/model/color/ColorSectionControllerTest.java
deleted file mode 100644
index 820e641..0000000
--- a/tests/robotests/src/com/android/customization/model/color/ColorSectionControllerTest.java
+++ /dev/null
@@ -1,62 +0,0 @@
-/*
- * Copyright (C) 2022 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.google.common.truth.Truth.assertThat;
-
-import androidx.appcompat.app.AppCompatActivity;
-
-import com.android.wallpaper.model.WallpaperColorsViewModel;
-
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.MockitoAnnotations;
-import org.robolectric.Robolectric;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.annotation.Config;
-
-/**
- * Tests of {@link ColorSectionController}.
- */
-@RunWith(RobolectricTestRunner.class)
-@Config(manifest = Config.NONE)
-public final class ColorSectionControllerTest {
-
-    private AppCompatActivity mActivity;
-    private ColorSectionController mColorSectionController;
-
-    /**
-     * Set up the test case.
-     */
-    @Before
-    public void setUp() {
-        MockitoAnnotations.initMocks(this);
-        mActivity = Robolectric.buildActivity(AppCompatActivity.class).create().get();
-        mColorSectionController = new ColorSectionController(mActivity,
-                new WallpaperColorsViewModel(), mActivity, null);
-    }
-
-    /**
-     * isAvailable()'s test.
-     */
-    @Test
-    @Config(manifest = Config.NONE)
-    public void isAvailable_nullContext_shouldReturnFalse() {
-        assertThat(mColorSectionController.isAvailable(/* context= */ null)).isFalse();
-    }
-}
-
diff --git a/tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt b/tests/robotests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
similarity index 96%
rename from tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
rename to tests/robotests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
index 5953937..0c80e28 100644
--- a/tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
+++ b/tests/robotests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
@@ -17,6 +17,7 @@
 
 package com.android.customization.model.grid.data.repository
 
+import com.android.customization.model.grid.GridOption
 import com.android.customization.model.grid.shared.model.GridOptionItemModel
 import com.android.customization.model.grid.shared.model.GridOptionItemsModel
 import kotlinx.coroutines.CoroutineScope
@@ -51,6 +52,8 @@
         return options
     }
 
+    override fun getSelectedOption(): GridOption? = null
+
     fun setOptions(
         count: Int,
         selectedIndex: Int = 0,
diff --git a/tests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt b/tests/robotests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt
similarity index 100%
rename from tests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt
rename to tests/robotests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt
diff --git a/tests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt b/tests/robotests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt
similarity index 100%
rename from tests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt
rename to tests/robotests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt
diff --git a/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt b/tests/robotests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
similarity index 100%
rename from tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
rename to tests/robotests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
diff --git a/tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt b/tests/robotests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
similarity index 100%
rename from tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
rename to tests/robotests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
diff --git a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt b/tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
similarity index 98%
rename from tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
rename to tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
index cbf1365..d4f24ee 100644
--- a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
@@ -32,11 +32,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class ColorPickerInteractorTest {
     private lateinit var underTest: ColorPickerInteractor
     private lateinit var repository: FakeColorPickerRepository
diff --git a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt b/tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
similarity index 98%
rename from tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
rename to tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
index 71a8f23..5f3e39e 100644
--- a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
@@ -34,11 +34,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class ColorPickerSnapshotRestorerTest {
 
     private lateinit var underTest: ColorPickerSnapshotRestorer
diff --git a/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt b/tests/robotests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
similarity index 98%
rename from tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
rename to tests/robotests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
index a29d559..9968c5f 100644
--- a/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
@@ -43,11 +43,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class ColorPickerViewModelTest {
     private lateinit var underTest: ColorPickerViewModel
     private lateinit var repository: FakeColorPickerRepository
diff --git a/tests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
similarity index 97%
rename from tests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
rename to tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
index 3f22ced..35dbadd 100644
--- a/tests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
@@ -35,11 +35,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class KeyguardQuickAffordancePickerRepositoryTest {
 
     private lateinit var underTest: KeyguardQuickAffordancePickerRepository
diff --git a/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
similarity index 98%
rename from tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
rename to tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
index fea94dc..efe9f64 100644
--- a/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
@@ -39,11 +39,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class KeyguardQuickAffordancePickerInteractorTest {
 
     private lateinit var underTest: KeyguardQuickAffordancePickerInteractor
diff --git a/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
similarity index 98%
rename from tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
rename to tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
index 3f10674..f71bfc7 100644
--- a/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
@@ -19,8 +19,8 @@
 
 import android.content.Context
 import android.content.Intent
+import androidx.test.core.app.ApplicationProvider
 import androidx.test.filters.SmallTest
-import androidx.test.platform.app.InstrumentationRegistry
 import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordanceSnapshotRestorer
@@ -57,11 +57,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class KeyguardQuickAffordancePickerViewModelTest {
 
     private lateinit var underTest: KeyguardQuickAffordancePickerViewModel
@@ -75,7 +75,7 @@
     @Before
     fun setUp() {
         InjectorProvider.setInjector(TestInjector())
-        context = InstrumentationRegistry.getInstrumentation().targetContext
+        context = ApplicationProvider.getApplicationContext()
         val testDispatcher = StandardTestDispatcher()
         testScope = TestScope(testDispatcher)
         Dispatchers.setMain(testDispatcher)
diff --git a/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt b/tests/robotests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt
similarity index 100%
rename from tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt
rename to tests/robotests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt
diff --git a/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt b/tests/robotests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt
similarity index 100%
rename from tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt
rename to tests/robotests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt
diff --git a/tests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt b/tests/robotests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
similarity index 84%
rename from tests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
rename to tests/robotests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
index bf2766d..95d7e35 100644
--- a/tests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
+++ b/tests/robotests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
@@ -45,11 +45,12 @@
             val selectedClock = fakeClocks.find { clock -> clock.clockId == selectedClockId }
             checkNotNull(selectedClock)
             ClockMetadataModel(
-                selectedClock.clockId,
-                selectedClock.name,
-                selectedColor,
-                colorTone,
-                seedColor,
+                clockId = selectedClock.clockId,
+                name = selectedClock.name,
+                isSelected = true,
+                selectedColorId = selectedColor,
+                colorToneProgress = colorTone,
+                seedColor = seedColor,
             )
         }
 
@@ -81,10 +82,10 @@
         const val CLOCK_ID_3 = "clock3"
         val fakeClocks =
             listOf(
-                ClockMetadataModel(CLOCK_ID_0, "clock0", null, 50, null),
-                ClockMetadataModel(CLOCK_ID_1, "clock1", null, 50, null),
-                ClockMetadataModel(CLOCK_ID_2, "clock2", null, 50, null),
-                ClockMetadataModel(CLOCK_ID_3, "clock3", null, 50, null),
+                ClockMetadataModel(CLOCK_ID_0, "clock0", true, null, 50, null),
+                ClockMetadataModel(CLOCK_ID_1, "clock1", false, null, 50, null),
+                ClockMetadataModel(CLOCK_ID_2, "clock2", false, null, 50, null),
+                ClockMetadataModel(CLOCK_ID_3, "clock3", false, null, 50, null),
             )
         const val CLOCK_COLOR_ID = "RED"
         const val CLOCK_COLOR_TONE_PROGRESS = 87
diff --git a/tests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt b/tests/robotests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
similarity index 97%
rename from tests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
rename to tests/robotests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
index 1a7ebb5..c8e39be 100644
--- a/tests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
@@ -17,11 +17,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class ClockPickerInteractorTest {
 
     private lateinit var underTest: ClockPickerInteractor
diff --git a/tests/robotests/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragmentTest.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragmentTest.kt
deleted file mode 100644
index 0a54312..0000000
--- a/tests/robotests/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragmentTest.kt
+++ /dev/null
@@ -1,75 +0,0 @@
-package com.android.customization.picker.clock.ui.fragment
-
-import android.view.View
-import androidx.appcompat.app.AppCompatActivity
-import androidx.recyclerview.widget.LinearLayoutManager
-import androidx.recyclerview.widget.RecyclerView
-import com.android.systemui.plugins.ClockMetadata
-import com.android.systemui.plugins.ClockSettings
-import com.android.systemui.plugins.PluginManager
-import com.android.systemui.shared.clocks.ClockRegistry
-import org.junit.Assert
-import org.junit.Before
-import org.junit.Ignore
-import org.junit.Test
-import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.Mockito.verify
-import org.mockito.Mockito.`when` as whenever
-import org.mockito.MockitoAnnotations
-import org.robolectric.Robolectric
-import org.robolectric.RobolectricTestRunner
-import org.robolectric.annotation.Config
-
-/** Tests of [ClockCustomDemoFragment]. */
-@RunWith(RobolectricTestRunner::class)
-@Config(manifest = Config.NONE)
-@Ignore("b/270606895")
-class ClockCustomDemoFragmentTest {
-    private lateinit var mActivity: AppCompatActivity
-    private var mClockCustomDemoFragment: ClockCustomDemoFragment? = null
-    @Mock private lateinit var registry: ClockRegistry
-    @Mock private lateinit var mockPluginManager: PluginManager
-
-    private var settingValue: ClockSettings? = null
-
-    @Before
-    fun setUp() {
-        MockitoAnnotations.initMocks(this)
-        mActivity = Robolectric.buildActivity(AppCompatActivity::class.java).get()
-        mClockCustomDemoFragment = ClockCustomDemoFragment()
-        whenever(registry.getClocks())
-            .thenReturn(
-                listOf(
-                    ClockMetadata("CLOCK_1", "Clock 1"),
-                    ClockMetadata("CLOCK_2", "Clock 2"),
-                    ClockMetadata("CLOCK_NOT_IN_USE", "Clock not in use")
-                )
-            )
-
-        mClockCustomDemoFragment!!.clockRegistry = registry
-        mClockCustomDemoFragment!!.recyclerView = RecyclerView(mActivity)
-        mClockCustomDemoFragment!!.recyclerView.layoutManager =
-            LinearLayoutManager(mActivity, RecyclerView.VERTICAL, false)
-    }
-
-    @Test
-    fun testItemCount_getCorrectClockCount() {
-        Assert.assertEquals(3, mClockCustomDemoFragment!!.recyclerView.adapter!!.itemCount)
-    }
-
-    @Test
-    fun testClick_setCorrectClockId() {
-        mClockCustomDemoFragment!!
-            .recyclerView
-            .measure(View.MeasureSpec.UNSPECIFIED, View.MeasureSpec.UNSPECIFIED)
-        mClockCustomDemoFragment!!.recyclerView.layout(0, 0, 100, 10000)
-        val testPosition = 1
-        mClockCustomDemoFragment!!
-            .recyclerView
-            .findViewHolderForAdapterPosition(testPosition)
-            ?.itemView
-            ?.performClick()
-        verify(registry).currentClockId = "CLOCK_1"
-    }
-}
diff --git a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
similarity index 70%
rename from tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
rename to tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
index c5eb796..ca6f8c7 100644
--- a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
@@ -37,22 +37,23 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class ClockCarouselViewModelTest {
     private val repositoryWithMultipleClocks by lazy { FakeClockPickerRepository() }
     private val repositoryWithSingleClock by lazy {
         FakeClockPickerRepository(
             listOf(
                 ClockMetadataModel(
-                    "clock0",
-                    "clock0",
-                    null,
-                    ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS,
-                    null,
+                    clockId = "clock0",
+                    name = "clock0",
+                    isSelected = true,
+                    selectedColorId = null,
+                    colorToneProgress = ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS,
+                    seedColor = null,
                 ),
             )
         )
@@ -87,38 +88,6 @@
         assertThat(observedSelectedIndex()).isEqualTo(2)
     }
 
-    @Test
-    fun multipleClockCase() = runTest {
-        underTest =
-            ClockCarouselViewModel(
-                getClockPickerInteractor(repositoryWithMultipleClocks),
-                testDispatcher
-            )
-        val observedIsCarouselVisible = collectLastValue(underTest.isCarouselVisible)
-        val observedIsSingleClockViewVisible = collectLastValue(underTest.isSingleClockViewVisible)
-
-        advanceTimeBy(ClockCarouselViewModel.CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
-
-        assertThat(observedIsCarouselVisible()).isTrue()
-        assertThat(observedIsSingleClockViewVisible()).isFalse()
-    }
-
-    @Test
-    fun singleClockCase() = runTest {
-        underTest =
-            ClockCarouselViewModel(
-                getClockPickerInteractor(repositoryWithSingleClock),
-                testDispatcher
-            )
-        val observedIsCarouselVisible = collectLastValue(underTest.isCarouselVisible)
-        val observedIsSingleClockViewVisible = collectLastValue(underTest.isSingleClockViewVisible)
-
-        advanceTimeBy(ClockCarouselViewModel.CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
-
-        assertThat(observedIsCarouselVisible()).isFalse()
-        assertThat(observedIsSingleClockViewVisible()).isTrue()
-    }
-
     private fun getClockPickerInteractor(repository: ClockPickerRepository): ClockPickerInteractor {
         return ClockPickerInteractor(
                 repository = repository,
diff --git a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
similarity index 97%
rename from tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
rename to tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
index 293e393..19a704c 100644
--- a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
@@ -36,11 +36,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class ClockSectionViewModelTest {
 
     private lateinit var clockColorMap: Map<String, ClockColorViewModel>
diff --git a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
similarity index 98%
rename from tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
rename to tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
index f58baf8..f09e977 100644
--- a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
@@ -28,11 +28,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class ClockSettingsViewModelTest {
 
     private lateinit var context: Context
diff --git a/tests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt b/tests/robotests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt
similarity index 100%
rename from tests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt
rename to tests/robotests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt
diff --git a/tests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
similarity index 97%
rename from tests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
rename to tests/robotests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
index 63334f2..5c3544a 100644
--- a/tests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
@@ -37,11 +37,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class NotificationSectionViewModelTest {
 
     private lateinit var underTest: NotificationSectionViewModel
diff --git a/tests/src/com/android/customization/testing/TestDrawableLayerResolver.java b/tests/src/com/android/customization/testing/TestDrawableLayerResolver.java
deleted file mode 100644
index e507221..0000000
--- a/tests/src/com/android/customization/testing/TestDrawableLayerResolver.java
+++ /dev/null
@@ -1,16 +0,0 @@
-package com.android.customization.testing;
-
-import android.graphics.drawable.Drawable;
-import android.graphics.drawable.LayerDrawable;
-
-import com.android.wallpaper.module.DrawableLayerResolver;
-
-/**
- * Test implementation of {@link DrawableLayerResolver}.
- */
-public class TestDrawableLayerResolver implements DrawableLayerResolver {
-    @Override
-    public Drawable resolveLayer(LayerDrawable layerDrawable) {
-        return layerDrawable.getDrawable(0);
-    }
-}
diff --git a/tests/src/com/android/customization/testing/TestPackageStatusNotifier.java b/tests/src/com/android/customization/testing/TestPackageStatusNotifier.java
deleted file mode 100644
index 1e6a1a2..0000000
--- a/tests/src/com/android/customization/testing/TestPackageStatusNotifier.java
+++ /dev/null
@@ -1,18 +0,0 @@
-package com.android.customization.testing;
-
-import com.android.wallpaper.module.PackageStatusNotifier;
-
-/**
- * Test implementation of {@link PackageStatusNotifier}.
- */
-public class TestPackageStatusNotifier implements PackageStatusNotifier {
-    @Override
-    public void addListener(Listener listener, String action) {
-        // Do nothing
-    }
-
-    @Override
-    public void removeListener(Listener listener) {
-        // Do nothing
-    }
-}