Merge "Fix the layout of the grid option customization (2/2)" into main
diff --git a/res/layout/grid_option2.xml b/res/layout/grid_option2.xml
new file mode 100644
index 0000000..437b95b
--- /dev/null
+++ b/res/layout/grid_option2.xml
@@ -0,0 +1,63 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  ~
+  -->
+
+<LinearLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:tools="http://schemas.android.com/tools"
+    android:layout_width="@dimen/option_item_size"
+    android:layout_height="wrap_content"
+    android:orientation="vertical"
+    android:gravity="center_horizontal"
+    android:clipChildren="false">
+
+    <FrameLayout
+        android:layout_width="@dimen/option_item_size"
+        android:layout_height="@dimen/option_item_size"
+        android:clipChildren="false">
+
+        <com.android.wallpaper.picker.option.ui.view.OptionItemBackground
+            android:id="@id/background"
+            android:layout_width="match_parent"
+            android:layout_height="match_parent"
+            android:background="@drawable/option_item_background"
+            android:importantForAccessibility="no" />
+
+        <ImageView
+            android:id="@id/foreground"
+            android:layout_width="48dp"
+            android:layout_height="48dp"
+            android:layout_gravity="center" />
+
+    </FrameLayout>
+
+    <View
+        android:layout_width="0dp"
+        android:layout_height="8dp" />
+
+    <TextView
+        android:id="@id/text"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:textColor="@color/system_on_surface"
+        android:singleLine="true"
+        android:ellipsize="end"
+        android:textSize="12sp"
+        android:text="Placeholder for stable size calculation, please do not remove."
+        tools:ignore="HardcodedText" />
+
+</LinearLayout>
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index 0c36753..4f1062f 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -189,6 +189,7 @@
     <dimen name="floating_sheet_tab_clock_font_toolbar_top_margin">16dp</dimen>
     <dimen name="floating_sheet_tab_clock_font_toolbar_bottom_margin">8dp</dimen>
     <dimen name="floating_sheet_list_item_horizontal_space">4dp</dimen>
+    <dimen name="floating_sheet_grid_list_item_horizontal_space">10dp</dimen>
     <dimen name="floating_sheet_list_item_vertical_space">4dp</dimen>
     <dimen name="floating_sheet_clock_style_option_list_margin_bottom">8dp</dimen>
     <dimen name="floating_sheet_clock_style_option_width">80dp</dimen>
diff --git a/src/com/android/wallpaper/customization/ui/binder/ShapeGridFloatingSheetBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ShapeGridFloatingSheetBinder.kt
index 5f292bd..138a253 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ShapeGridFloatingSheetBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ShapeGridFloatingSheetBinder.kt
@@ -43,6 +43,7 @@
 import com.android.wallpaper.picker.customization.ui.view.adapter.FloatingToolbarTabAdapter
 import com.android.wallpaper.picker.customization.ui.viewmodel.ColorUpdateViewModel
 import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter
+import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter2
 import com.android.wallpaper.picker.option.ui.binder.OptionItemBinder
 import java.lang.ref.WeakReference
 import kotlinx.coroutines.CoroutineDispatcher
@@ -97,7 +98,7 @@
 
         val gridContent = view.requireViewById<View>(R.id.app_grid_container)
         val gridOptionListAdapter =
-            createGridOptionItemAdapter(view.context, lifecycleOwner, backgroundDispatcher)
+            createGridOptionItemAdapter(lifecycleOwner, backgroundDispatcher)
         val gridOptionList =
             view.requireViewById<RecyclerView>(R.id.grid_options).also {
                 it.initGridOptionList(view.context, gridOptionListAdapter)
@@ -243,8 +244,7 @@
                         ),
                     itemHorizontalSpacePx =
                         context.resources.getDimensionPixelSize(
-                            com.android.themepicker.R.dimen
-                                .floating_sheet_list_item_horizontal_space
+                            R.dimen.floating_sheet_list_item_horizontal_space
                         ),
                 )
             )
@@ -253,30 +253,23 @@
     }
 
     private fun createGridOptionItemAdapter(
-        context: Context,
         lifecycleOwner: LifecycleOwner,
         backgroundDispatcher: CoroutineDispatcher,
-    ): OptionItemAdapter<GridIconViewModel> =
-        OptionItemAdapter(
-            layoutResourceId = R.layout.grid_option,
+    ): OptionItemAdapter2<GridIconViewModel> =
+        OptionItemAdapter2(
+            layoutResourceId = R.layout.grid_option2,
             lifecycleOwner = lifecycleOwner,
             backgroundDispatcher = backgroundDispatcher,
-            foregroundTintSpec =
-                OptionItemBinder.TintSpec(
-                    selectedColor =
-                        context.getColor(com.android.wallpaper.R.color.system_on_surface),
-                    unselectedColor =
-                        context.getColor(com.android.wallpaper.R.color.system_on_surface),
-                ),
-            bindIcon = { foregroundView: View, gridIcon: GridIconViewModel ->
-                val imageView = foregroundView as? ImageView
+            bindPayload = { view: View, gridIcon: GridIconViewModel ->
+                val imageView = view.findViewById(R.id.foreground) as? ImageView
                 imageView?.let { GridIconViewBinder.bind(imageView, gridIcon) }
+                return@OptionItemAdapter2 null
             },
         )
 
     private fun RecyclerView.initGridOptionList(
         context: Context,
-        adapter: OptionItemAdapter<GridIconViewModel>,
+        adapter: OptionItemAdapter2<GridIconViewModel>,
     ) {
         apply {
             this.layoutManager = LinearLayoutManager(context, RecyclerView.HORIZONTAL, false)
@@ -288,8 +281,7 @@
                         ),
                     itemHorizontalSpacePx =
                         context.resources.getDimensionPixelSize(
-                            com.android.themepicker.R.dimen
-                                .floating_sheet_list_item_horizontal_space
+                            R.dimen.floating_sheet_grid_list_item_horizontal_space
                         ),
                 )
             )
diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
index eaed4e5..46d0346 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
@@ -52,7 +52,6 @@
 import javax.inject.Singleton
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.launch
 
 @Singleton
@@ -259,7 +258,7 @@
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
                 launch {
                     combine(
-                            clockPickerViewModel.previewingClock.filterNotNull(),
+                            clockPickerViewModel.previewingClock,
                             clockPickerViewModel.previewingClockSize,
                         ) { clock, size ->
                             clock to size
diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
index 5ea757e..c17775a 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
@@ -22,16 +22,18 @@
 import android.widget.FrameLayout
 import android.widget.Toolbar
 import androidx.core.graphics.ColorUtils
+import androidx.core.graphics.drawable.DrawableCompat
 import androidx.core.view.isInvisible
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.lifecycleScope
 import androidx.lifecycle.repeatOnLifecycle
-import com.android.wallpaper.R
 import com.android.wallpaper.customization.ui.viewmodel.ThemePickerCustomizationOptionsViewModel
 import com.android.wallpaper.customization.ui.viewmodel.ToolbarHeightsViewModel
+import com.android.wallpaper.picker.customization.ui.binder.ColorUpdateBinder
 import com.android.wallpaper.picker.customization.ui.binder.DefaultToolbarBinder
 import com.android.wallpaper.picker.customization.ui.binder.ToolbarBinder
+import com.android.wallpaper.picker.customization.ui.viewmodel.ColorUpdateViewModel
 import com.android.wallpaper.picker.customization.ui.viewmodel.CustomizationOptionsViewModel
 import javax.inject.Inject
 import javax.inject.Singleton
@@ -39,6 +41,7 @@
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.filterNotNull
+import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
 
 @Singleton
@@ -54,6 +57,7 @@
         toolbar: Toolbar,
         applyButton: Button,
         viewModel: CustomizationOptionsViewModel,
+        colorUpdateViewModel: ColorUpdateViewModel,
         lifecycleOwner: LifecycleOwner,
         onNavBack: () -> Unit,
     ) {
@@ -62,6 +66,7 @@
             toolbar,
             applyButton,
             viewModel,
+            colorUpdateViewModel,
             lifecycleOwner,
             onNavBack,
         )
@@ -111,13 +116,14 @@
             }
         )
 
-        val applyButtonTextColorEnabled =
-            applyButton.resources.getColor(R.color.system_on_primary, null)
-        val applyButtonTextColorDisabled =
-            ColorUtils.setAlphaComponent(
-                applyButton.resources.getColor(R.color.system_on_surface, null),
-                97,
-            ) // 97 for 38% transparent
+        ColorUpdateBinder.bind(
+            setColor = { color ->
+                DrawableCompat.setTint(DrawableCompat.wrap(applyButton.background), color)
+            },
+            color = colorUpdateViewModel.colorPrimary,
+            shouldAnimate = { true },
+            lifecycleOwner = lifecycleOwner,
+        )
 
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
@@ -133,9 +139,22 @@
                     viewModel.isApplyButtonEnabled.collect {
                         applyButton.isEnabled = it
                         applyButton.background.alpha =
-                            if (it) 255 else 31 // 255 for 100%, 31 for 12% transparent,
-                        applyButton.setTextColor(
-                            if (it) applyButtonTextColorEnabled else applyButtonTextColorDisabled
+                            if (it) 255 else 31 // 255 for 100%, 31 for 12% transparent
+                        ColorUpdateBinder.bind(
+                            setColor = { color -> applyButton.setTextColor(color) },
+                            color =
+                                if (it) {
+                                    colorUpdateViewModel.colorOnPrimary
+                                } else {
+                                    colorUpdateViewModel.colorOnSurface.map { color: Int ->
+                                        ColorUtils.setAlphaComponent(
+                                            color,
+                                            97,
+                                        ) // 97 for 38% transparent
+                                    }
+                                },
+                            shouldAnimate = { true },
+                            lifecycleOwner = lifecycleOwner,
                         )
                     }
                 }
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
index 6ca773b..2a1a8c9 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
@@ -57,6 +57,7 @@
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.mapLatest
+import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
 
 /** View model for the clock customization screen. */
@@ -120,8 +121,9 @@
     val selectedClock = clockPickerInteractor.selectedClock
     val previewingClock =
         combine(overridingClock, selectedClock) { overridingClock, selectedClock ->
-            overridingClock ?: selectedClock
-        }
+                (overridingClock ?: selectedClock)
+            }
+            .shareIn(viewModelScope, SharingStarted.WhileSubscribed(), 1)
 
     data class ClockStyleModel(val thumbnail: Drawable, val showEditButton: StateFlow<Boolean>)
 
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
index e7efebd..4029dbe 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
@@ -33,6 +33,8 @@
 import dagger.assisted.AssistedInject
 import dagger.hilt.android.qualifiers.ApplicationContext
 import dagger.hilt.android.scopes.ViewModelScoped
+import kotlin.coroutines.resume
+import kotlinx.coroutines.CancellableContinuation
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -42,6 +44,7 @@
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.suspendCancellableCoroutine
 
 /** Models UI state for a color picker experience. */
 class ColorPickerViewModel2
@@ -57,6 +60,7 @@
     val previewingColorOption = overridingColorOption.asStateFlow()
 
     private val selectedColorTypeTabId = MutableStateFlow<ColorType?>(null)
+    private var onApplyContinuation: CancellableContinuation<Unit>? = null
 
     /** View-models for each color tab. */
     val colorTypeTabs: Flow<List<FloatingToolbarTabViewModel>> =
@@ -165,6 +169,10 @@
                 .toMap()
         }
 
+    /**
+     * This function suspends until onApplyComplete is called to accommodate for configuration
+     * change updates, which are applied with a latency.
+     */
     val onApply: Flow<(suspend () -> Unit)?> =
         previewingColorOption.map { previewingColorOption ->
             previewingColorOption?.let {
@@ -173,6 +181,12 @@
                 } else {
                     {
                         interactor.select(it)
+                        // Suspend until onApplyComplete is called, e.g. on configuration change
+                        suspendCancellableCoroutine { continuation: CancellableContinuation<Unit> ->
+                            onApplyContinuation?.cancel()
+                            onApplyContinuation = continuation
+                            continuation.invokeOnCancellation { onApplyContinuation = null }
+                        }
                         logger.logThemeColorApplied(
                             previewingColorOption.colorOption.sourceForLogging,
                             previewingColorOption.colorOption.styleForLogging,
@@ -187,6 +201,12 @@
         overridingColorOption.value = null
     }
 
+    /** Resumes the onApply function if apply is in progress, otherwise no-op */
+    fun onApplyComplete() {
+        onApplyContinuation?.resume(Unit)
+        onApplyContinuation = null
+    }
+
     /** The list of all available color options for the selected Color Type. */
     val colorOptions: Flow<List<OptionItemViewModel2<ColorOptionIconViewModel>>> =
         combine(allColorOptions, selectedColorTypeTabId) {
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
index 7f3c4cb..1e19e80 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
@@ -29,6 +29,7 @@
 import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
 import com.android.wallpaper.picker.customization.ui.viewmodel.FloatingToolbarTabViewModel
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel2
 import dagger.assisted.Assisted
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
@@ -126,7 +127,7 @@
             overridingGridOptionKey ?: selectedGridOption.key.value
         }
 
-    val gridOptions: Flow<List<OptionItemViewModel<GridIconViewModel>>> =
+    val gridOptions: Flow<List<OptionItemViewModel2<GridIconViewModel>>> =
         interactor.gridOptions
             .filterNotNull()
             .map { gridOptions -> gridOptions.map { toGridOptionItemViewModel(it) } }
@@ -184,7 +185,7 @@
 
     private fun toGridOptionItemViewModel(
         option: GridOptionModel
-    ): OptionItemViewModel<GridIconViewModel> {
+    ): OptionItemViewModel2<GridIconViewModel> {
         val iconShapePath =
             context.resources.getString(
                 Resources.getSystem()
@@ -203,7 +204,7 @@
                     initialValue = false,
                 )
 
-        return OptionItemViewModel(
+        return OptionItemViewModel2(
             key = MutableStateFlow(option.key),
             payload =
                 GridIconViewModel(columns = option.cols, rows = option.rows, path = iconShapePath),
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
index 62a96fe..6bc6180 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
@@ -27,6 +27,7 @@
 import dagger.hilt.android.scopes.ViewModelScoped
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.combine
@@ -59,6 +60,8 @@
     val shapeGridPickerViewModel =
         shapeGridPickerViewModelFactory.create(viewModelScope = viewModelScope)
 
+    private var onApplyJob: Job? = null
+
     override val selectedOption = defaultCustomizationOptionsViewModel.selectedOption
 
     override fun handleBackPressed(): Boolean {
@@ -167,9 +170,14 @@
             .map { onApply ->
                 if (onApply != null) {
                     fun(onComplete: () -> Unit) {
-                        viewModelScope.launch {
-                            onApply()
-                            onComplete()
+                        // Prevent double apply
+                        if (onApplyJob?.isActive != true) {
+                            onApplyJob =
+                                viewModelScope.launch {
+                                    onApply()
+                                    onComplete()
+                                    onApplyJob = null
+                                }
                         }
                     }
                 } else {
diff --git a/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt b/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
index 6952d91..e8c5b15 100644
--- a/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
+++ b/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
@@ -25,13 +25,21 @@
 import androidx.lifecycle.repeatOnLifecycle
 import com.android.customization.model.grid.DefaultShapeGridManager.Companion.COL_GRID_NAME
 import com.android.customization.model.grid.DefaultShapeGridManager.Companion.COL_SHAPE_KEY
+import com.android.customization.picker.clock.shared.ClockSize
+import com.android.customization.picker.clock.ui.view.ClockViewFactory
 import com.android.customization.picker.color.data.util.MaterialColorsGenerator
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.CLOCK_SIZE_DYNAMIC
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.CLOCK_SIZE_SMALL
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.KEY_CLOCK_SIZE
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.KEY_HIDE_SMART_SPACE
 import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.KEY_INITIALLY_SELECTED_SLOT_ID
 import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.KEY_QUICK_AFFORDANCE_ID
 import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.KEY_SLOT_ID
 import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.MESSAGE_ID_DEFAULT_PREVIEW
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.MESSAGE_ID_HIDE_SMART_SPACE
+import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.MESSAGE_ID_PREVIEW_CLOCK_SIZE
 import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.MESSAGE_ID_PREVIEW_QUICK_AFFORDANCE_SELECTED
 import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.MESSAGE_ID_SLOT_SELECTED
 import com.android.systemui.shared.quickaffordance.shared.model.KeyguardPreviewConstants.MESSAGE_ID_START_CUSTOMIZING_QUICK_AFFORDANCES
@@ -60,6 +68,7 @@
         colorUpdateViewModel: ColorUpdateViewModel,
         screen: Screen,
         lifecycleOwner: LifecycleOwner,
+        clockViewFactory: ClockViewFactory,
     ) {
         defaultWorkspaceCallbackBinder.bind(
             workspaceCallback = workspaceCallback,
@@ -67,6 +76,7 @@
             colorUpdateViewModel = colorUpdateViewModel,
             screen = screen,
             lifecycleOwner = lifecycleOwner,
+            clockViewFactory = clockViewFactory,
         )
 
         if (viewModel !is ThemePickerCustomizationOptionsViewModel) {
@@ -135,6 +145,48 @@
                                     }
                                 }
                         }
+
+                        launch {
+                            combine(
+                                    viewModel.clockPickerViewModel.previewingClock,
+                                    viewModel.clockPickerViewModel.previewingClockSize,
+                                    ::Pair,
+                                )
+                                .collect { (previewingClock, previewingClockSize) ->
+                                    val hideSmartspace =
+                                        clockViewFactory
+                                            .getController(previewingClock.clockId)
+                                            .let {
+                                                when (previewingClockSize) {
+                                                    ClockSize.DYNAMIC ->
+                                                        it.largeClock.config
+                                                            .hasCustomWeatherDataDisplay
+                                                    ClockSize.SMALL ->
+                                                        it.smallClock.config
+                                                            .hasCustomWeatherDataDisplay
+                                                }
+                                            }
+                                    workspaceCallback.sendMessage(
+                                        MESSAGE_ID_HIDE_SMART_SPACE,
+                                        Bundle().apply {
+                                            putBoolean(KEY_HIDE_SMART_SPACE, hideSmartspace)
+                                        },
+                                    )
+
+                                    workspaceCallback.sendMessage(
+                                        MESSAGE_ID_PREVIEW_CLOCK_SIZE,
+                                        Bundle().apply {
+                                            putString(
+                                                KEY_CLOCK_SIZE,
+                                                when (previewingClockSize) {
+                                                    ClockSize.DYNAMIC -> CLOCK_SIZE_DYNAMIC
+                                                    ClockSize.SMALL -> CLOCK_SIZE_SMALL
+                                                },
+                                            )
+                                        },
+                                    )
+                                }
+                        }
                     }
                 }
             Screen.HOME_SCREEN ->
@@ -159,13 +211,18 @@
                         }
 
                         launch {
+                            colorUpdateViewModel.systemColorsUpdated.collect {
+                                viewModel.colorPickerViewModel2.onApplyComplete()
+                            }
+                        }
+
+                        launch {
                             combine(
                                     viewModel.colorPickerViewModel2.previewingColorOption,
                                     viewModel.darkModeViewModel.overridingIsDarkMode,
-                                    colorUpdateViewModel.systemColorsUpdated,
-                                    ::Triple,
+                                    ::Pair,
                                 )
-                                .collect { (colorModel, darkMode, _) ->
+                                .collect { (colorModel, darkMode) ->
                                     val bundle =
                                         Bundle().apply {
                                             if (colorModel != null) {
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
index 421d8a0..e4a4884 100644
--- a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
@@ -36,6 +36,7 @@
 import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
@@ -96,6 +97,27 @@
     }
 
     @Test
+    fun onApply_suspendsUntilOnApplyCompleteIsCalled() =
+        testScope.runTest {
+            val colorTypes = collectLastValue(underTest.colorTypeTabs)
+            val colorOptions = collectLastValue(underTest.colorOptions)
+            val onApply = collectLastValue(underTest.onApply)
+
+            // Select "Wallpaper colors" tab
+            colorTypes()?.get(0)?.onClick?.invoke()
+            // Select a color option to preview
+            selectColorOption(colorOptions, 1)
+            // Apply the selected color option
+            val job = testScope.launch { onApply()?.invoke() }
+
+            assertThat(job.isActive).isTrue()
+
+            underTest.onApplyComplete()
+
+            assertThat(job.isActive).isFalse()
+        }
+
+    @Test
     fun onApply_wallpaperColor_shouldLogColor() =
         testScope.runTest {
             repository.setOptions(
@@ -203,7 +225,7 @@
             )
         }
 
-    /** Simulates a user selecting the affordance at the given index, if that is clickable. */
+    /** Simulates a user selecting the color option at the given index. */
     private fun TestScope.selectColorOption(
         colorOptions: () -> List<OptionItemViewModel2<ColorOptionIconViewModel>>?,
         index: Int,
@@ -217,10 +239,11 @@
         }
     }
 
-    /** Simulates a user selecting the affordance at the given index, if that is clickable. */
+    /** Simulates a user applying the color option at the given index, and the apply completes. */
     private suspend fun TestScope.applySelectedColorOption() {
         val onApply = collectLastValue(underTest.onApply)()
-        onApply?.invoke()
+        testScope.launch { onApply?.invoke() }
+        underTest.onApplyComplete()
     }
 
     /**
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModelTest.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModelTest.kt
index c99acca..71ea0d9 100644
--- a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModelTest.kt
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModelTest.kt
@@ -27,6 +27,7 @@
 import com.android.customization.picker.grid.ui.viewmodel.ShapeIconViewModel
 import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel2
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import dagger.hilt.android.qualifiers.ApplicationContext
@@ -159,7 +160,7 @@
         testScope.runTest {
             val selectedGridOption = collectLastValue(underTest.selectedGridOption)
 
-            assertOptionItem(
+            assertGridItem(
                 optionItem = selectedGridOption(),
                 key = "normal",
                 payload = GridIconViewModel(5, 5, iconShapePath),
@@ -183,7 +184,7 @@
             onPracticalOptionClick()?.invoke()
             onApply()?.invoke()
 
-            assertOptionItem(
+            assertGridItem(
                 optionItem = selectedGridOption(),
                 key = "practical",
                 payload = GridIconViewModel(4, 5, iconShapePath),
@@ -199,7 +200,7 @@
         testScope.runTest {
             val optionItems = collectLastValue(underTest.gridOptions)
 
-            assertOptionItem(
+            assertGridItem(
                 optionItem = optionItems()?.get(0),
                 key = "normal",
                 payload = GridIconViewModel(5, 5, iconShapePath),
@@ -208,7 +209,7 @@
                 isSelected = true,
                 isEnabled = true,
             )
-            assertOptionItem(
+            assertGridItem(
                 optionItem = optionItems()?.get(1),
                 key = "practical",
                 payload = GridIconViewModel(4, 5, iconShapePath),
@@ -229,7 +230,7 @@
 
             onPracticalOptionClick()?.invoke()
 
-            assertOptionItem(
+            assertGridItem(
                 optionItem = optionItems()?.get(0),
                 key = "normal",
                 payload = GridIconViewModel(5, 5, iconShapePath),
@@ -238,7 +239,7 @@
                 isSelected = false,
                 isEnabled = true,
             )
-            assertOptionItem(
+            assertGridItem(
                 optionItem = optionItems()?.get(1),
                 key = "practical",
                 payload = GridIconViewModel(4, 5, iconShapePath),
@@ -267,8 +268,8 @@
         assertThat(optionItem.isEnabled).isEqualTo(isEnabled)
     }
 
-    private fun TestScope.assertOptionItem(
-        optionItem: OptionItemViewModel<GridIconViewModel>?,
+    private fun TestScope.assertGridItem(
+        optionItem: OptionItemViewModel2<GridIconViewModel>?,
         key: String,
         payload: GridIconViewModel?,
         text: Text,