diff --git a/res/layout/clock_style_option2.xml b/res/layout/clock_style_option2.xml
new file mode 100644
index 0000000..8bb60d1
--- /dev/null
+++ b/res/layout/clock_style_option2.xml
@@ -0,0 +1,46 @@
+<?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.
+-->
+<!-- Content description is set programmatically on the parent FrameLayout -->
+<FrameLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="80dp"
+    android:layout_height="80dp"
+    android:clipToPadding="false"
+    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:importantForAccessibility="no" />
+
+    <ImageView
+        android:id="@+id/foreground"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        android:layout_margin="@dimen/floating_sheet_clock_style_thumbnail_margin"
+        android:src="@drawable/ic_clock_24px" />
+
+    <ImageView
+        android:id="@+id/edit_icon"
+        android:layout_width="@dimen/floating_sheet_clock_edit_icon_size"
+        android:layout_height="@dimen/floating_sheet_clock_edit_icon_size"
+        android:layout_marginTop="-16dp"
+        android:layout_marginStart="50dp"
+        android:src="@drawable/edit_icon"
+        android:importantForAccessibility="no" />
+</FrameLayout>
+
diff --git a/res/layout/floating_sheet_clock.xml b/res/layout/floating_sheet_clock.xml
index ec2e6c1..cdaf3de 100644
--- a/res/layout/floating_sheet_clock.xml
+++ b/res/layout/floating_sheet_clock.xml
@@ -46,6 +46,7 @@
                 app:layout_constraintEnd_toEndOf="parent"
                 app:layout_constraintTop_toTopOf="parent"
                 app:layout_constraintBottom_toTopOf="@+id/clock_style_clock_size_title"
+                android:layout_marginBottom="8dp"
                 android:clipToPadding="false"
                 android:clipChildren="false">
 
@@ -64,10 +65,15 @@
                     android:layout_height="@dimen/floating_sheet_clock_style_option_size"
                     android:visibility="invisible" />
 
+                <!--
+                TODO (b/377528523): We intentionally disable over scroll mode since it will clip the
+                  edit icon when in over scroll.
+                -->
                 <androidx.recyclerview.widget.RecyclerView
                     android:id="@+id/clock_style_list"
                     android:layout_width="match_parent"
                     android:layout_height="wrap_content"
+                    android:overScrollMode="never"
                     android:clipChildren="false"
                     android:clipToPadding="false"/>
             </FrameLayout>
@@ -167,7 +173,6 @@
             android:layout_width="match_parent"
             android:layout_height="wrap_content"
             android:paddingHorizontal="@dimen/floating_sheet_content_horizontal_padding"
-            android:paddingVertical="@dimen/floating_sheet_content_vertical_padding"
             android:orientation="vertical"
             android:clipToPadding="false"
             android:clipChildren="false">
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index b634df0..a022d00 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -65,13 +65,22 @@
 import com.android.wallpaper.config.BaseFlags
 import com.android.wallpaper.module.CustomizationSections
 import com.android.wallpaper.module.FragmentFactory
+import com.android.wallpaper.module.NetworkStatusNotifier
+import com.android.wallpaper.module.PartnerProvider
 import com.android.wallpaper.module.WallpaperPicker2Injector
+import com.android.wallpaper.module.WallpaperPreferences
+import com.android.wallpaper.module.logging.UserEventLogger
+import com.android.wallpaper.network.Requester
 import com.android.wallpaper.picker.CustomizationPickerActivity
+import com.android.wallpaper.picker.category.wrapper.WallpaperCategoryWrapper
+import com.android.wallpaper.picker.customization.data.content.WallpaperClient
 import com.android.wallpaper.picker.customization.data.repository.WallpaperColorsRepository
 import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
 import com.android.wallpaper.picker.di.modules.BackgroundDispatcher
 import com.android.wallpaper.picker.di.modules.MainDispatcher
 import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.system.UiModeManagerWrapper
+import com.android.wallpaper.util.DisplayUtils
 import dagger.Lazy
 import javax.inject.Inject
 import javax.inject.Singleton
@@ -85,7 +94,46 @@
     @MainDispatcher private val mainScope: CoroutineScope,
     @BackgroundDispatcher private val bgScope: CoroutineScope,
     @BackgroundDispatcher private val bgDispatcher: CoroutineDispatcher,
-) : WallpaperPicker2Injector(mainScope), CustomizationInjector {
+    private val colorContrastSectionViewModelFactory: Lazy<ColorContrastSectionViewModel.Factory>,
+    private val keyguardQuickAffordancePickerInteractor:
+        Lazy<KeyguardQuickAffordancePickerInteractor>,
+    private val keyguardQuickAffordanceSnapshotRestorer:
+        Lazy<KeyguardQuickAffordanceSnapshotRestorer>,
+    private val themesUserEventLogger: Lazy<ThemesUserEventLogger>,
+    private val colorPickerInteractor: Lazy<ColorPickerInteractor>,
+    private val colorPickerSnapshotRestorer: Lazy<ColorPickerSnapshotRestorer>,
+    private val clockRegistry: Lazy<ClockRegistry>,
+    private val secureSettingsRepository: Lazy<SecureSettingsRepository>,
+    private val systemSettingsRepository: Lazy<SystemSettingsRepository>,
+    private val clockPickerInteractor: Lazy<ClockPickerInteractor>,
+    private val clockPickerSnapshotRestorer: Lazy<ClockPickerSnapshotRestorer>,
+    displayUtils: Lazy<DisplayUtils>,
+    requester: Lazy<Requester>,
+    networkStatusNotifier: Lazy<NetworkStatusNotifier>,
+    partnerProvider: Lazy<PartnerProvider>,
+    val uiModeManager: Lazy<UiModeManagerWrapper>,
+    userEventLogger: Lazy<UserEventLogger>,
+    injectedWallpaperClient: Lazy<WallpaperClient>,
+    private val injectedWallpaperInteractor: Lazy<WallpaperInteractor>,
+    prefs: Lazy<WallpaperPreferences>,
+    wallpaperColorsRepository: Lazy<WallpaperColorsRepository>,
+    defaultWallpaperCategoryWrapper: Lazy<WallpaperCategoryWrapper>,
+) :
+    WallpaperPicker2Injector(
+        mainScope,
+        displayUtils,
+        requester,
+        networkStatusNotifier,
+        partnerProvider,
+        uiModeManager,
+        userEventLogger,
+        injectedWallpaperClient,
+        injectedWallpaperInteractor,
+        prefs,
+        wallpaperColorsRepository,
+        defaultWallpaperCategoryWrapper,
+    ),
+    CustomizationInjector {
     private var customizationSections: CustomizationSections? = null
     private var keyguardQuickAffordancePickerViewModelFactory:
         KeyguardQuickAffordancePickerViewModel.Factory? =
@@ -106,24 +154,6 @@
     private var gridSnapshotRestorer: GridSnapshotRestorer? = null
     private var gridScreenViewModelFactory: GridScreenViewModel.Factory? = null
 
-    // Injected objects, sorted by type
-    @Inject
-    lateinit var colorContrastSectionViewModelFactory: Lazy<ColorContrastSectionViewModel.Factory>
-    @Inject
-    lateinit var keyguardQuickAffordancePickerInteractor:
-        Lazy<KeyguardQuickAffordancePickerInteractor>
-    @Inject
-    lateinit var keyguardQuickAffordanceSnapshotRestorer:
-        Lazy<KeyguardQuickAffordanceSnapshotRestorer>
-    @Inject lateinit var themesUserEventLogger: Lazy<ThemesUserEventLogger>
-    @Inject lateinit var colorPickerInteractor: Lazy<ColorPickerInteractor>
-    @Inject lateinit var colorPickerSnapshotRestorer: Lazy<ColorPickerSnapshotRestorer>
-    @Inject lateinit var clockRegistry: Lazy<ClockRegistry>
-    @Inject lateinit var secureSettingsRepository: Lazy<SecureSettingsRepository>
-    @Inject lateinit var systemSettingsRepository: Lazy<SystemSettingsRepository>
-    @Inject lateinit var clockPickerInteractor: Lazy<ClockPickerInteractor>
-    @Inject lateinit var clockPickerSnapshotRestorer: Lazy<ClockPickerSnapshotRestorer>
-
     override fun getCustomizationSections(activity: ComponentActivity): CustomizationSections {
         val appContext = activity.applicationContext
         val clockViewFactory = getClockViewFactory(activity)
diff --git a/src/com/android/wallpaper/customization/ui/binder/ClockFloatingSheetBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ClockFloatingSheetBinder.kt
index 2d53c0e..5f982e2 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ClockFloatingSheetBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ClockFloatingSheetBinder.kt
@@ -22,11 +22,13 @@
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
+import android.view.ViewTreeObserver.OnGlobalLayoutListener
+import android.widget.FrameLayout
 import android.widget.ImageView
+import android.widget.LinearLayout
 import android.widget.SeekBar
 import android.widget.Switch
 import android.widget.TextView
-import androidx.core.view.doOnLayout
 import androidx.core.view.isVisible
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
@@ -53,12 +55,15 @@
 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 java.lang.ref.WeakReference
 import kotlin.math.abs
+import kotlinx.coroutines.DisposableHandle
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.launch
 
 object ClockFloatingSheetBinder {
@@ -68,8 +73,8 @@
 
     private val _clockFloatingSheetHeights: MutableStateFlow<ClockFloatingSheetHeightsViewModel?> =
         MutableStateFlow(null)
-    private val clockFloatingSheetHeights: Flow<ClockFloatingSheetHeightsViewModel?> =
-        _clockFloatingSheetHeights.asStateFlow()
+    private val clockFloatingSheetHeights: Flow<ClockFloatingSheetHeightsViewModel> =
+        _clockFloatingSheetHeights.asStateFlow().filterNotNull()
 
     fun bind(
         view: View,
@@ -90,7 +95,7 @@
                 .also { tabs.setAdapter(it) }
 
         val floatingSheetContainer =
-            view.requireViewById<ViewGroup>(R.id.clock_floating_sheet_content_container)
+            view.requireViewById<FrameLayout>(R.id.clock_floating_sheet_content_container)
 
         // Clock style
         val clockStyleContent = view.requireViewById<View>(R.id.clock_floating_sheet_style_content)
@@ -105,10 +110,10 @@
             view.requireViewById<ViewGroup>(R.id.clock_floating_sheet_font_content)
         val clockFontToolbar = view.requireViewById<ViewGroup>(R.id.clock_font_toolbar)
         clockFontToolbar.requireViewById<View>(R.id.clock_font_revert).setOnClickListener {
-            viewModel.revertFontAxes()
+            viewModel.cancelFontAxes()
         }
         clockFontToolbar.requireViewById<View>(R.id.clock_font_apply).setOnClickListener {
-            viewModel.applyFontAxes()
+            viewModel.confirmFontAxes()
         }
 
         // Clock color
@@ -137,16 +142,39 @@
         // Clock size switch
         val clockSizeSwitch = view.requireViewById<Switch>(R.id.clock_style_clock_size_switch)
 
-        view.doOnLayout {
-            if (_clockFloatingSheetHeights.value == null) {
-                _clockFloatingSheetHeights.value =
-                    ClockFloatingSheetHeightsViewModel(
-                        clockStyleContentHeight = clockStyleContent.height,
-                        clockColorContentHeight = clockColorContent.height,
-                        clockFontContentHeight = clockFontContent.height,
-                    )
+        clockStyleContent.viewTreeObserver.addOnGlobalLayoutListener(
+            object : OnGlobalLayoutListener {
+                override fun onGlobalLayout() {
+                    if (clockStyleContent.height != 0) {
+                        _clockFloatingSheetHeights.value =
+                            _clockFloatingSheetHeights.value?.copy(
+                                clockStyleContentHeight = clockStyleContent.height
+                            )
+                                ?: ClockFloatingSheetHeightsViewModel(
+                                    clockStyleContentHeight = clockStyleContent.height
+                                )
+                        clockStyleContent.viewTreeObserver.removeOnGlobalLayoutListener(this)
+                    }
+                }
             }
-        }
+        )
+
+        clockColorContent.viewTreeObserver.addOnGlobalLayoutListener(
+            object : OnGlobalLayoutListener {
+                override fun onGlobalLayout() {
+                    if (clockColorContent.height != 0) {
+                        _clockFloatingSheetHeights.value =
+                            _clockFloatingSheetHeights.value?.copy(
+                                clockColorContentHeight = clockColorContent.height
+                            )
+                                ?: ClockFloatingSheetHeightsViewModel(
+                                    clockColorContentHeight = clockColorContent.height
+                                )
+                        clockColorContent.viewTreeObserver.removeOnGlobalLayoutListener(this)
+                    }
+                }
+            }
+        )
 
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
@@ -155,27 +183,40 @@
                 launch {
                     combine(clockFloatingSheetHeights, viewModel.selectedTab, ::Pair).collect {
                         (heights, selectedTab) ->
-                        heights ?: return@collect
+                        val (clockStyleContentHeight, clockColorContentHeight) = heights
+                        clockStyleContentHeight ?: return@collect
+                        clockColorContentHeight ?: return@collect
 
-                        val targetHeight =
-                            when (selectedTab) {
-                                Tab.STYLE -> heights.clockStyleContentHeight
-                                Tab.COLOR -> heights.clockColorContentHeight
-                                Tab.FONT -> heights.clockFontContentHeight
-                            } +
-                                view.resources.getDimensionPixelSize(
-                                    R.dimen.floating_sheet_content_vertical_padding
-                                ) * 2
+                        if (selectedTab == Tab.STYLE || selectedTab == Tab.COLOR) {
+                            val targetHeight =
+                                when (selectedTab) {
+                                    Tab.STYLE -> clockStyleContentHeight
+                                    Tab.COLOR -> clockColorContentHeight
+                                    else -> 0
+                                } +
+                                    view.resources.getDimensionPixelSize(
+                                        R.dimen.floating_sheet_content_vertical_padding
+                                    ) * 2
 
-                        val animationFloatingSheet =
                             ValueAnimator.ofInt(floatingSheetContainer.height, targetHeight)
-                        animationFloatingSheet.addUpdateListener { valueAnimator ->
-                            val value = valueAnimator.animatedValue as Int
+                                .apply {
+                                    addUpdateListener { valueAnimator ->
+                                        val value = valueAnimator.animatedValue as Int
+                                        floatingSheetContainer.layoutParams =
+                                            floatingSheetContainer.layoutParams.apply {
+                                                height = value
+                                            }
+                                    }
+                                    duration = ANIMATION_DURATION
+                                }
+                                .start()
+                        } else if (selectedTab == Tab.FONT) {
                             floatingSheetContainer.layoutParams =
-                                floatingSheetContainer.layoutParams.apply { height = value }
+                                LinearLayout.LayoutParams(
+                                    LinearLayout.LayoutParams.MATCH_PARENT,
+                                    LinearLayout.LayoutParams.WRAP_CONTENT,
+                                )
                         }
-                        animationFloatingSheet.setDuration(ANIMATION_DURATION)
-                        animationFloatingSheet.start()
 
                         clockStyleContent.isVisible = selectedTab == Tab.STYLE
                         clockColorContent.isVisible = selectedTab == Tab.COLOR
@@ -192,12 +233,6 @@
                     combine(viewModel.previewingClock, viewModel.previewingFontAxisMap, ::Pair)
                         .collect { pair ->
                             val (clock, fontAxisMap) = pair
-                            if (clock == null) {
-                                boundClockId = null
-                                boundEditorViews = mapOf()
-                                clockFontContent.removeAllViews()
-                                return@collect
-                            }
 
                             if (boundClockId != clock.clockId) {
                                 boundEditorViews =
@@ -281,6 +316,7 @@
         viewModel: ClockPickerViewModel,
     ): Map<String, Pair<View, ClockFontAxis>> {
         parent.removeAllViews()
+
         val inflater = LayoutInflater.from(parent.context)
         val axisMap = mutableMapOf<String, Pair<View, ClockFontAxis>>()
         var nextSwitch: View? = null
@@ -342,27 +378,35 @@
                 )
             }
         }
+
         return axisMap
     }
 
     private fun createClockStyleOptionItemAdapter(
         lifecycleOwner: LifecycleOwner
-    ): OptionItemAdapter<ClockStyleModel> =
-        OptionItemAdapter(
-            layoutResourceId = R.layout.clock_style_option,
+    ): OptionItemAdapter2<ClockStyleModel> =
+        OptionItemAdapter2(
+            layoutResourceId = R.layout.clock_style_option2,
             lifecycleOwner = lifecycleOwner,
-            bindIcon = { view: View, style: ClockStyleModel ->
-                (view.findViewById(R.id.clock_icon) as ImageView).setImageDrawable(style.thumbnail)
-                (view.findViewById(R.id.edit_icon) as ImageView).setVisibility(
-                    // TODO(b/364673969): Route isSelected here and hide when unselected
-                    if (style.isEditable) View.VISIBLE else View.GONE
-                )
+            bindPayload = { view: View, styleModel: ClockStyleModel ->
+                view
+                    .findViewById<ImageView>(R.id.foreground)
+                    ?.setImageDrawable(styleModel.thumbnail)
+                val job =
+                    lifecycleOwner.lifecycleScope.launch {
+                        lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                            styleModel.showEditButton.collect {
+                                view.findViewById<ImageView>(R.id.edit_icon)?.isVisible = it
+                            }
+                        }
+                    }
+                return@OptionItemAdapter2 DisposableHandle { job.cancel() }
             },
         )
 
     private fun RecyclerView.initStyleList(
         context: Context,
-        adapter: OptionItemAdapter<ClockStyleModel>,
+        adapter: OptionItemAdapter2<ClockStyleModel>,
     ) {
         this.adapter = adapter
         layoutManager = LinearLayoutManager(context, LinearLayoutManager.HORIZONTAL, false)
diff --git a/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt
index 58a1fd9..7ddcb01 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt
@@ -120,6 +120,9 @@
                     )
                 val night = uiMode and UI_MODE_NIGHT_MASK == UI_MODE_NIGHT_YES
                 ColorOptionIconBinder2.bind(colorOptionIconView, colorIcon, night)
+                // Return null since it does not need the lifecycleOwner to launch any job for later
+                // disposal when rebind.
+                return@OptionItemAdapter2 null
             },
         )
 
diff --git a/src/com/android/wallpaper/customization/ui/binder/ShortcutFloatingSheetBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ShortcutFloatingSheetBinder.kt
index 4c72412..838ef87 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ShortcutFloatingSheetBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ShortcutFloatingSheetBinder.kt
@@ -156,6 +156,9 @@
                 val imageView =
                     itemView.requireViewById<ImageView>(com.android.wallpaper.R.id.foreground)
                 IconViewBinder.bind(imageView, gridIcon)
+                // Return null since it does not need the lifecycleOwner to launch any job for later
+                // disposal when rebind.
+                return@OptionItemAdapter2 null
             },
         )
 
diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
index 18b9584..357f131 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
@@ -16,20 +16,27 @@
 
 package com.android.wallpaper.customization.ui.binder
 
+import android.animation.ValueAnimator
+import android.view.ViewTreeObserver.OnGlobalLayoutListener
 import android.widget.Button
 import android.widget.FrameLayout
 import android.widget.Toolbar
-import androidx.core.view.isVisible
+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.customization.ui.viewmodel.ThemePickerCustomizationOptionsViewModel
+import com.android.wallpaper.customization.ui.viewmodel.ToolbarHeightsViewModel
 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.CustomizationOptionsViewModel
 import javax.inject.Inject
 import javax.inject.Singleton
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.launch
 
 @Singleton
@@ -37,6 +44,9 @@
 @Inject
 constructor(private val defaultToolbarBinder: DefaultToolbarBinder) : ToolbarBinder {
 
+    private val _toolbarHeights: MutableStateFlow<ToolbarHeightsViewModel?> = MutableStateFlow(null)
+    private val toolbarHeights = _toolbarHeights.asStateFlow().filterNotNull()
+
     override fun bind(
         navButton: FrameLayout,
         toolbar: Toolbar,
@@ -60,6 +70,45 @@
             )
         }
 
+        navButton.viewTreeObserver.addOnGlobalLayoutListener(
+            object : OnGlobalLayoutListener {
+                override fun onGlobalLayout() {
+                    if (navButton.height != 0) {
+                        _toolbarHeights.value =
+                            _toolbarHeights.value?.copy(navButtonHeight = navButton.height)
+                                ?: ToolbarHeightsViewModel(navButtonHeight = navButton.height)
+                    }
+                    navButton.viewTreeObserver.removeOnGlobalLayoutListener(this)
+                }
+            }
+        )
+
+        toolbar.viewTreeObserver.addOnGlobalLayoutListener(
+            object : OnGlobalLayoutListener {
+                override fun onGlobalLayout() {
+                    if (toolbar.height != 0) {
+                        _toolbarHeights.value =
+                            _toolbarHeights.value?.copy(toolbarHeight = toolbar.height)
+                                ?: ToolbarHeightsViewModel(toolbarHeight = toolbar.height)
+                    }
+                    navButton.viewTreeObserver.removeOnGlobalLayoutListener(this)
+                }
+            }
+        )
+
+        applyButton.viewTreeObserver.addOnGlobalLayoutListener(
+            object : OnGlobalLayoutListener {
+                override fun onGlobalLayout() {
+                    if (applyButton.height != 0) {
+                        _toolbarHeights.value =
+                            _toolbarHeights.value?.copy(applyButtonHeight = applyButton.height)
+                                ?: ToolbarHeightsViewModel(applyButtonHeight = applyButton.height)
+                    }
+                    applyButton.viewTreeObserver.removeOnGlobalLayoutListener(this)
+                }
+            }
+        )
+
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
                 launch {
@@ -68,10 +117,60 @@
                     }
                 }
 
-                launch { viewModel.isOnApplyVisible.collect { applyButton.isVisible = it } }
+                launch { viewModel.isOnApplyVisible.collect { applyButton.isInvisible = !it } }
 
                 launch { viewModel.isOnApplyEnabled.collect { applyButton.isEnabled = it } }
+
+                launch {
+                    combine(toolbarHeights, viewModel.isToolbarCollapsed, ::Pair).collect {
+                        (toolbarHeights, isToolbarCollapsed) ->
+                        val (navButtonHeight, toolbarHeight, applyButtonHeight) = toolbarHeights
+                        navButtonHeight ?: return@collect
+                        toolbarHeight ?: return@collect
+                        applyButtonHeight ?: return@collect
+
+                        val navButtonToHeight = if (isToolbarCollapsed) 0 else navButtonHeight
+                        val toolbarToHeight = if (isToolbarCollapsed) 0 else toolbarHeight
+                        val applyButtonToHeight = if (isToolbarCollapsed) 0 else applyButtonHeight
+                        ValueAnimator.ofInt(navButton.height, navButtonToHeight)
+                            .apply {
+                                addUpdateListener { valueAnimator ->
+                                    val value = valueAnimator.animatedValue as Int
+                                    navButton.layoutParams =
+                                        navButton.layoutParams.apply { height = value }
+                                }
+                                duration = ANIMATION_DURATION
+                            }
+                            .start()
+
+                        ValueAnimator.ofInt(toolbar.height, toolbarToHeight)
+                            .apply {
+                                addUpdateListener { valueAnimator ->
+                                    val value = valueAnimator.animatedValue as Int
+                                    toolbar.layoutParams =
+                                        toolbar.layoutParams.apply { height = value }
+                                }
+                                duration = ANIMATION_DURATION
+                            }
+                            .start()
+
+                        ValueAnimator.ofInt(applyButton.height, applyButtonToHeight)
+                            .apply {
+                                addUpdateListener { valueAnimator ->
+                                    val value = valueAnimator.animatedValue as Int
+                                    applyButton.layoutParams =
+                                        applyButton.layoutParams.apply { height = value }
+                                }
+                                duration = ANIMATION_DURATION
+                            }
+                            .start()
+                    }
+                }
             }
         }
     }
+
+    companion object {
+        private const val ANIMATION_DURATION = 200L
+    }
 }
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ClockFloatingSheetHeightsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ClockFloatingSheetHeightsViewModel.kt
index 1f86533..a39ee7d 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ClockFloatingSheetHeightsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ClockFloatingSheetHeightsViewModel.kt
@@ -17,7 +17,6 @@
 package com.android.wallpaper.customization.ui.viewmodel
 
 data class ClockFloatingSheetHeightsViewModel(
-    val clockStyleContentHeight: Int,
-    val clockColorContentHeight: Int,
-    val clockFontContentHeight: Int,
+    val clockStyleContentHeight: Int? = null,
+    val clockColorContentHeight: Int? = null,
 )
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
index 11124f9..da3e65c 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
@@ -53,6 +53,7 @@
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.mapLatest
@@ -123,7 +124,7 @@
             overridingClock ?: selectedClock
         }
 
-    data class ClockStyleModel(val thumbnail: Drawable, val isEditable: Boolean)
+    data class ClockStyleModel(val thumbnail: Drawable, val showEditButton: StateFlow<Boolean>)
 
     @OptIn(ExperimentalCoroutinesApi::class)
     val clockStyleOptions: StateFlow<List<OptionItemViewModel<ClockStyleModel>>> =
@@ -131,40 +132,10 @@
             .mapLatest { allClocks ->
                 // Delay to avoid the case that the full list of clocks is not initiated.
                 delay(CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
-                allClocks.map { clockModel ->
-                    val isSelectedFlow =
-                        previewingClock
-                            .map { it.clockId == clockModel.clockId }
-                            .stateIn(viewModelScope)
-                    val contentDescription =
-                        resources.getString(
-                            R.string.select_clock_action_description,
-                            clockModel.description,
-                        )
-                    OptionItemViewModel<ClockStyleModel>(
-                        key = MutableStateFlow(clockModel.clockId) as StateFlow<String>,
-                        payload =
-                            ClockStyleModel(
-                                clockModel.thumbnail,
-                                isEditable = !clockModel.fontAxes.isEmpty(),
-                            ),
-                        text = Text.Loaded(contentDescription),
-                        isTextUserVisible = false,
-                        isSelected = isSelectedFlow,
-                        onClicked =
-                            isSelectedFlow.map { isSelected ->
-                                if (isSelected) {
-                                    fun() {
-                                        _selectedTab.value = Tab.FONT
-                                    }
-                                } else {
-                                    fun() {
-                                        overridingClock.value = clockModel
-                                        overrideFontAxisMap.value = null
-                                    }
-                                }
-                            },
-                    )
+                val allClockMap = allClocks.groupBy { it.fontAxes.isNotEmpty() }
+                buildList {
+                    allClockMap[true]?.map { add(it.toOption(resources)) }
+                    allClockMap[false]?.map { add(it.toOption(resources)) }
                 }
             }
             // makes sure that the operations above this statement are executed on I/O dispatcher
@@ -173,27 +144,65 @@
             .flowOn(backgroundDispatcher.limitedParallelism(1))
             .stateIn(viewModelScope, SharingStarted.Eagerly, emptyList())
 
+    private suspend fun ClockMetadataModel.toOption(
+        resources: Resources
+    ): OptionItemViewModel<ClockStyleModel> {
+        val isSelectedFlow = previewingClock.map { it.clockId == clockId }.stateIn(viewModelScope)
+        val isEditable = fontAxes.isNotEmpty()
+        val showEditButton = isSelectedFlow.map { it && isEditable }.stateIn(viewModelScope)
+        val contentDescription =
+            resources.getString(R.string.select_clock_action_description, description)
+        return OptionItemViewModel<ClockStyleModel>(
+            key = MutableStateFlow(clockId) as StateFlow<String>,
+            payload = ClockStyleModel(thumbnail = thumbnail, showEditButton = showEditButton),
+            text = Text.Loaded(contentDescription),
+            isTextUserVisible = false,
+            isSelected = isSelectedFlow,
+            onClicked =
+                isSelectedFlow.map { isSelected ->
+                    if (isSelected && isEditable) {
+                        fun() {
+                            _selectedTab.value = Tab.FONT
+                        }
+                    } else {
+                        fun() {
+                            overridingClock.value = this
+                            overrideFontAxisMap.value = null
+                        }
+                    }
+                },
+        )
+    }
+
     // Clock Font Axis Editor
     private val overrideFontAxisMap = MutableStateFlow<Map<String, Float>?>(null)
+    private val isFontAxisMapEdited = overrideFontAxisMap.map { it != null }
+    private val selectedClockFontAxes =
+        previewingClock
+            .map { clock -> clock.fontAxes.associate { it.key to it.currentValue } }
+            .stateIn(viewModelScope, SharingStarted.Eagerly, null)
     val previewingFontAxisMap =
-        combine(overrideFontAxisMap, previewingClock) { overrideAxes, previewingClock ->
-                overrideAxes ?: previewingClock.fontAxes.associate { it.key to it.currentValue }
+        combine(overrideFontAxisMap, selectedClockFontAxes.filterNotNull()) {
+                overrideAxes,
+                selectedFontAxes ->
+                selectedFontAxes.toMutableMap().let { mutableMap ->
+                    overrideAxes?.forEach { (key, value) -> mutableMap[key] = value }
+                    mutableMap.toMap()
+                }
             }
             .stateIn(viewModelScope, SharingStarted.Eagerly, emptyMap())
 
-    private val isFontAxisMapEdited = overrideFontAxisMap.map { it != null }
-
     fun updatePreviewFontAxis(key: String, value: Float) {
-        val axisMap = previewingFontAxisMap.value.toMutableMap()
+        val axisMap = (overrideFontAxisMap.value?.toMutableMap() ?: mutableMapOf())
         axisMap[key] = value
-        overrideFontAxisMap.value = axisMap
+        overrideFontAxisMap.value = axisMap.toMap()
     }
 
-    fun applyFontAxes() {
+    fun confirmFontAxes() {
         _selectedTab.value = Tab.STYLE
     }
 
-    fun revertFontAxes() {
+    fun cancelFontAxes() {
         overrideFontAxisMap.value = null
         _selectedTab.value = Tab.STYLE
     }
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
index b2e2039..54ae132 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
@@ -30,6 +30,7 @@
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.flow.map
@@ -171,6 +172,14 @@
 
     val isOnApplyVisible: Flow<Boolean> = selectedOption.map { it != null }
 
+    val isToolbarCollapsed: Flow<Boolean> =
+        combine(selectedOption, clockPickerViewModel.selectedTab) { selectedOption, selectedTab ->
+                selectedOption ==
+                    ThemePickerCustomizationOptionUtil.ThemePickerLockCustomizationOption.CLOCK &&
+                    selectedTab == ClockPickerViewModel.Tab.FONT
+            }
+            .distinctUntilChanged()
+
     @ViewModelScoped
     @AssistedFactory
     interface Factory : CustomizationOptionsViewModelFactory {
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ToolbarHeightsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ToolbarHeightsViewModel.kt
new file mode 100644
index 0000000..0d859da
--- /dev/null
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ToolbarHeightsViewModel.kt
@@ -0,0 +1,23 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wallpaper.customization.ui.viewmodel
+
+data class ToolbarHeightsViewModel(
+    val navButtonHeight: Int? = null,
+    val toolbarHeight: Int? = null,
+    val applyButtonHeight: Int? = null,
+)
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 bf612c2..0656696 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
@@ -16,7 +16,6 @@
 
 package com.android.wallpaper.picker.common.preview.ui.binder
 
-import android.app.WallpaperManager
 import android.os.Bundle
 import android.os.Message
 import androidx.core.os.bundleOf
@@ -43,13 +42,13 @@
 import com.android.wallpaper.picker.customization.ui.viewmodel.CustomizationOptionsViewModel
 import javax.inject.Inject
 import javax.inject.Singleton
+import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.launch
 
 @Singleton
 class ThemePickerWorkspaceCallbackBinder
 @Inject
 constructor(
-    private val wallpaperManager: WallpaperManager,
     private val defaultWorkspaceCallbackBinder: DefaultWorkspaceCallbackBinder,
     private val materialColorsGenerator: MaterialColorsGenerator,
 ) : WorkspaceCallbackBinder {
@@ -157,25 +156,28 @@
                         }
 
                         launch {
-                            viewModel.colorPickerViewModel2.previewingColorOption.collect {
-                                if (it == null) {
-                                    workspaceCallback.sendMessage(MESSAGE_ID_UPDATE_COLOR, Bundle())
-                                    return@collect
-                                }
-                                val seedColor = it.colorOption.seedColor
-                                val (ids, colors) =
-                                    materialColorsGenerator.generate(
-                                        seedColor,
-                                        it.colorOption.style,
-                                    )
-                                workspaceCallback.sendMessage(
-                                    MESSAGE_ID_UPDATE_COLOR,
-                                    Bundle().apply {
-                                        putIntArray(KEY_COLOR_RESOURCE_IDS, ids)
-                                        putIntArray(KEY_COLOR_VALUES, colors)
-                                    },
+                            combine(
+                                    viewModel.colorPickerViewModel2.previewingColorOption,
+                                    viewModel.darkModeViewModel.previewingIsDarkMode,
+                                    ::Pair,
                                 )
-                            }
+                                .collect { (colorModel, darkMode) ->
+                                    val bundle =
+                                        Bundle().apply {
+                                            if (colorModel != null) {
+                                                val (ids, colors) =
+                                                    materialColorsGenerator.generate(
+                                                        colorModel.colorOption.seedColor,
+                                                        colorModel.colorOption.style,
+                                                    )
+                                                putIntArray(KEY_COLOR_RESOURCE_IDS, ids)
+                                                putIntArray(KEY_COLOR_VALUES, colors)
+                                            }
+
+                                            putBoolean(KEY_DARK_MODE, darkMode)
+                                        }
+                                    workspaceCallback.sendMessage(MESSAGE_ID_UPDATE_COLOR, bundle)
+                                }
                         }
                     }
                 }
@@ -189,5 +191,6 @@
         const val MESSAGE_ID_UPDATE_COLOR = 856
         const val KEY_COLOR_RESOURCE_IDS: String = "color_resource_ids"
         const val KEY_COLOR_VALUES: String = "color_values"
+        const val KEY_DARK_MODE: String = "use_dark_mode"
     }
 }
diff --git a/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt b/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
index 4e97599..a550119 100644
--- a/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
+++ b/tests/common/src/com/android/customization/testing/TestCustomizationInjector.kt
@@ -14,10 +14,17 @@
 import com.android.customization.picker.clock.ui.viewmodel.ClockSettingsViewModel
 import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
+import com.android.wallpaper.module.NetworkStatusNotifier
+import com.android.wallpaper.module.PartnerProvider
+import com.android.wallpaper.module.WallpaperPreferences
 import com.android.wallpaper.module.logging.UserEventLogger
+import com.android.wallpaper.network.Requester
 import com.android.wallpaper.picker.category.wrapper.WallpaperCategoryWrapper
 import com.android.wallpaper.picker.customization.data.repository.WallpaperColorsRepository
+import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
+import com.android.wallpaper.testing.FakeWallpaperClient
 import com.android.wallpaper.testing.TestInjector
+import com.android.wallpaper.util.DisplayUtils
 import javax.inject.Inject
 import javax.inject.Singleton
 
@@ -27,7 +34,27 @@
 constructor(
     private val customPrefs: TestDefaultCustomizationPreferences,
     private val themesUserEventLogger: ThemesUserEventLogger,
-) : TestInjector(themesUserEventLogger), CustomizationInjector {
+    displayUtils: DisplayUtils,
+    requester: Requester,
+    networkStatusNotifier: NetworkStatusNotifier,
+    partnerProvider: PartnerProvider,
+    wallpaperClient: FakeWallpaperClient,
+    injectedWallpaperInteractor: WallpaperInteractor,
+    prefs: WallpaperPreferences,
+    private val fakeWallpaperCategoryWrapper: WallpaperCategoryWrapper,
+) :
+    TestInjector(
+        themesUserEventLogger,
+        displayUtils,
+        requester,
+        networkStatusNotifier,
+        partnerProvider,
+        wallpaperClient,
+        injectedWallpaperInteractor,
+        prefs,
+        fakeWallpaperCategoryWrapper,
+    ),
+    CustomizationInjector {
     /////////////////
     // CustomizationInjector implementations
     /////////////////
@@ -82,6 +109,6 @@
     }
 
     override fun getWallpaperCategoryWrapper(): WallpaperCategoryWrapper {
-        return super.fakeWallpaperCategoryWrapper
+        return fakeWallpaperCategoryWrapper
     }
 }
diff --git a/tests/robotests/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
index 870d9f5..63c27ed 100644
--- a/tests/robotests/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
@@ -33,6 +33,11 @@
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
 import com.android.themepicker.R
 import com.android.wallpaper.module.InjectorProvider
+import com.android.wallpaper.module.NetworkStatusNotifier
+import com.android.wallpaper.module.PartnerProvider
+import com.android.wallpaper.module.WallpaperPreferences
+import com.android.wallpaper.network.Requester
+import com.android.wallpaper.picker.category.wrapper.WallpaperCategoryWrapper
 import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
 import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
 import com.android.wallpaper.picker.customization.data.repository.WallpaperRepository
@@ -43,6 +48,8 @@
 import com.android.wallpaper.testing.TestInjector
 import com.android.wallpaper.testing.TestWallpaperPreferences
 import com.android.wallpaper.testing.collectLastValue
+import com.android.wallpaper.util.DisplayUtils
+import com.android.wallpaper.util.DisplaysProvider
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.Dispatchers
@@ -56,6 +63,7 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.Mockito.mock
 import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
@@ -75,7 +83,6 @@
 
     @Before
     fun setUp() {
-        InjectorProvider.setInjector(TestInjector(logger))
         context = ApplicationProvider.getApplicationContext()
         val testDispatcher = StandardTestDispatcher()
         testScope = TestScope(testDispatcher)
@@ -100,8 +107,21 @@
                         client = FakeWallpaperClient(),
                         wallpaperPreferences = TestWallpaperPreferences(),
                         backgroundDispatcher = testDispatcher,
-                    ),
+                    )
             )
+        InjectorProvider.setInjector(
+            TestInjector(
+                logger,
+                DisplayUtils(context, mock(DisplaysProvider::class.java)),
+                mock(Requester::class.java),
+                mock(NetworkStatusNotifier::class.java),
+                mock(PartnerProvider::class.java),
+                FakeWallpaperClient(),
+                wallpaperInteractor,
+                mock(WallpaperPreferences::class.java),
+                mock(WallpaperCategoryWrapper::class.java),
+            )
+        )
         underTest =
             KeyguardQuickAffordancePickerViewModel.Factory(
                     context = context,
@@ -348,12 +368,12 @@
                         icon1 =
                             Icon.Loaded(
                                 FakeCustomizationProviderClient.ICON_1,
-                                Text.Loaded("Left shortcut")
+                                Text.Loaded("Left shortcut"),
                             ),
                         icon2 =
                             Icon.Loaded(
                                 FakeCustomizationProviderClient.ICON_3,
-                                Text.Loaded("Right shortcut")
+                                Text.Loaded("Right shortcut"),
                             ),
                     )
                 )
@@ -376,7 +396,7 @@
                         icon1 =
                             Icon.Loaded(
                                 FakeCustomizationProviderClient.ICON_1,
-                                Text.Loaded("Left shortcut")
+                                Text.Loaded("Left shortcut"),
                             ),
                         icon2 = null,
                     )
@@ -404,7 +424,7 @@
                         icon2 =
                             Icon.Loaded(
                                 FakeCustomizationProviderClient.ICON_3,
-                                Text.Loaded("Right shortcut")
+                                Text.Loaded("Right shortcut"),
                             ),
                     )
                 )
@@ -465,11 +485,7 @@
         assertThat(affordances).isNotNull()
         affordances?.forEach { affordance ->
             val nameMatchesSelectedName =
-                Text.evaluationEquals(
-                    context,
-                    affordance.text,
-                    Text.Loaded(selectedAffordanceText),
-                )
+                Text.evaluationEquals(context, affordance.text, Text.Loaded(selectedAffordanceText))
             val isSelected: Boolean? = collectLastValue(affordance.isSelected).invoke()
             assertWithMessage(
                     "Expected affordance with name \"${affordance.text}\" to have" +
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModelTest.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModelTest.kt
index ab64f89..0a8aa53 100644
--- a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModelTest.kt
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModelTest.kt
@@ -243,7 +243,7 @@
         assertThat(selectedTab()).isEqualTo(Tab.FONT)
         assertThat(previewingFontAxes()).isEqualTo(mapOf("key" to 100f))
 
-        underTest.applyFontAxes()
+        underTest.confirmFontAxes()
 
         assertThat(selectedTab()).isEqualTo(Tab.STYLE)
         assertThat(previewingFontAxes()).isEqualTo(mapOf("key" to 100f))
@@ -269,7 +269,7 @@
         assertThat(selectedTab()).isEqualTo(Tab.FONT)
         assertThat(previewingFontAxes()).isEqualTo(mapOf("key" to 100f))
 
-        underTest.revertFontAxes()
+        underTest.cancelFontAxes()
 
         assertThat(selectedTab()).isEqualTo(Tab.STYLE)
         assertThat(previewingFontAxes()).isEqualTo(mapOf("key" to 50f))