Merge "Fix wallpaper picker color for contrast issues (2/2)" into udc-dev
diff --git a/res/layout/clock_size_radio_button_group.xml b/res/layout/clock_size_radio_button_group.xml
index 30e39ad..4264007 100644
--- a/res/layout/clock_size_radio_button_group.xml
+++ b/res/layout/clock_size_radio_button_group.xml
@@ -57,7 +57,6 @@
         android:id="@+id/button_container_small"
         android:layout_width="match_parent"
         android:layout_height="wrap_content"
-        android:minHeight="48dp"
         android:orientation="horizontal">
 
         <RadioButton
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index 16b5de7..254a055 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -180,6 +180,4 @@
     <dimen name="small_clock_height">114dp</dimen>
     <dimen name="small_clock_padding_top">28dp</dimen>
     <dimen name="clock_padding_start">28dp</dimen>
-
-    <dimen name="tab_touch_delegate_height_padding">8dp</dimen>
 </resources>
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index 66814c5..98c6361 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -16,6 +16,7 @@
 package com.android.customization.module
 
 import android.app.UiModeManager
+import android.app.WallpaperManager
 import android.content.Context
 import android.content.Intent
 import android.net.Uri
@@ -45,6 +46,7 @@
 import com.android.customization.picker.clock.data.repository.ClockPickerRepositoryImpl
 import com.android.customization.picker.clock.data.repository.ClockRegistryProvider
 import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
+import com.android.customization.picker.clock.domain.interactor.ClockPickerSnapshotRestorer
 import com.android.customization.picker.clock.ui.view.ClockViewFactory
 import com.android.customization.picker.clock.ui.viewmodel.ClockCarouselViewModel
 import com.android.customization.picker.clock.ui.viewmodel.ClockSectionViewModel
@@ -98,15 +100,11 @@
     private var keyguardQuickAffordanceSnapshotRestorer: KeyguardQuickAffordanceSnapshotRestorer? =
         null
     private var notificationsSnapshotRestorer: NotificationsSnapshotRestorer? = null
-    /**
-     * Mapping from LifeCycleOwner's hashcode to ClockRegistry as we need to keep different
-     * ClockRegistries per LifeCycle to ensure proper cleanup
-     */
-    private var clockRegistries: MutableMap<Int, ClockRegistry> = HashMap()
     private var clockPickerInteractor: ClockPickerInteractor? = null
     private var clockSectionViewModel: ClockSectionViewModel? = null
     private var clockCarouselViewModelFactory: ClockCarouselViewModel.Factory? = null
     private var clockViewFactories: MutableMap<Int, ClockViewFactory> = HashMap()
+    private var clockPickerSnapshotRestorer: ClockPickerSnapshotRestorer? = null
     private var notificationsInteractor: NotificationsInteractor? = null
     private var notificationSectionViewModelFactory: NotificationSectionViewModel.Factory? = null
     private var colorPickerInteractor: ColorPickerInteractor? = null
@@ -120,6 +118,7 @@
     private var gridInteractor: GridInteractor? = null
     private var gridSnapshotRestorer: GridSnapshotRestorer? = null
     private var gridScreenViewModelFactory: GridScreenViewModel.Factory? = null
+    private var clockRegistryProvider: ClockRegistryProvider? = null
 
     override fun getCustomizationSections(activity: ComponentActivity): CustomizationSections {
         return customizationSections
@@ -195,18 +194,26 @@
         return fragmentFactory ?: ThemePickerFragmentFactory().also { fragmentFactory }
     }
 
-    override fun getSnapshotRestorers(context: Context): Map<Int, SnapshotRestorer> {
-        return super<WallpaperPicker2Injector>.getSnapshotRestorers(context).toMutableMap().apply {
-            this[KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER] =
-                getKeyguardQuickAffordanceSnapshotRestorer(context)
-            this[KEY_WALLPAPER_SNAPSHOT_RESTORER] = getWallpaperSnapshotRestorer(context)
-            this[KEY_NOTIFICATIONS_SNAPSHOT_RESTORER] = getNotificationsSnapshotRestorer(context)
-            this[KEY_DARK_MODE_SNAPSHOT_RESTORER] = getDarkModeSnapshotRestorer(context)
-            this[KEY_THEMED_ICON_SNAPSHOT_RESTORER] = getThemedIconSnapshotRestorer(context)
-            this[KEY_APP_GRID_SNAPSHOT_RESTORER] = getGridSnapshotRestorer(context)
-            this[KEY_COLOR_PICKER_SNAPSHOT_RESTORER] =
-                getColorPickerSnapshotRestorer(context, getWallpaperColorsViewModel())
-        }
+    override fun getSnapshotRestorers(
+        context: Context,
+        lifecycleOwner: LifecycleOwner
+    ): Map<Int, SnapshotRestorer> {
+        return super<WallpaperPicker2Injector>.getSnapshotRestorers(context, lifecycleOwner)
+            .toMutableMap()
+            .apply {
+                this[KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER] =
+                    getKeyguardQuickAffordanceSnapshotRestorer(context)
+                this[KEY_WALLPAPER_SNAPSHOT_RESTORER] = getWallpaperSnapshotRestorer(context)
+                this[KEY_NOTIFICATIONS_SNAPSHOT_RESTORER] =
+                    getNotificationsSnapshotRestorer(context)
+                this[KEY_DARK_MODE_SNAPSHOT_RESTORER] = getDarkModeSnapshotRestorer(context)
+                this[KEY_THEMED_ICON_SNAPSHOT_RESTORER] = getThemedIconSnapshotRestorer(context)
+                this[KEY_APP_GRID_SNAPSHOT_RESTORER] = getGridSnapshotRestorer(context)
+                this[KEY_COLOR_PICKER_SNAPSHOT_RESTORER] =
+                    getColorPickerSnapshotRestorer(context, getWallpaperColorsViewModel())
+                this[KEY_CLOCKS_SNAPSHOT_RESTORER] =
+                    getClockPickerSnapshotRestorer(context, lifecycleOwner)
+            }
     }
 
     override fun getCustomizationPreferences(context: Context): CustomizationPreferences {
@@ -264,16 +271,6 @@
                 .also { keyguardQuickAffordancePickerViewModelFactory = it }
     }
 
-    fun getNotificationSectionViewModelFactory(
-        context: Context,
-    ): NotificationSectionViewModel.Factory {
-        return notificationSectionViewModelFactory
-            ?: NotificationSectionViewModel.Factory(
-                    interactor = getNotificationsInteractor(context),
-                )
-                .also { notificationSectionViewModelFactory = it }
-    }
-
     private fun getKeyguardQuickAffordancePickerInteractorImpl(
         context: Context
     ): KeyguardQuickAffordancePickerInteractor {
@@ -306,6 +303,32 @@
                 .also { keyguardQuickAffordanceSnapshotRestorer = it }
     }
 
+    fun getNotificationSectionViewModelFactory(
+        context: Context,
+    ): NotificationSectionViewModel.Factory {
+        return notificationSectionViewModelFactory
+            ?: NotificationSectionViewModel.Factory(
+                    interactor = getNotificationsInteractor(context),
+                )
+                .also { notificationSectionViewModelFactory = it }
+    }
+
+    private fun getNotificationsInteractor(
+        context: Context,
+    ): NotificationsInteractor {
+        return notificationsInteractor
+            ?: NotificationsInteractor(
+                    repository =
+                        NotificationsRepository(
+                            scope = getApplicationCoroutineScope(),
+                            backgroundDispatcher = Dispatchers.IO,
+                            secureSettingsRepository = getSecureSettingsRepository(context),
+                        ),
+                    snapshotRestorer = { getNotificationsSnapshotRestorer(context) },
+                )
+                .also { notificationsInteractor = it }
+    }
+
     private fun getNotificationsSnapshotRestorer(context: Context): NotificationsSnapshotRestorer {
         return notificationsSnapshotRestorer
             ?: NotificationsSnapshotRestorer(
@@ -318,26 +341,15 @@
     }
 
     override fun getClockRegistry(context: Context, lifecycleOwner: LifecycleOwner): ClockRegistry {
-        return clockRegistries[lifecycleOwner.hashCode()]
-            ?: ClockRegistryProvider(
-                    context = context,
-                    coroutineScope = getApplicationCoroutineScope(),
-                    mainDispatcher = Dispatchers.Main,
-                    backgroundDispatcher = Dispatchers.IO,
-                )
-                .get()
-                .also {
-                    clockRegistries[lifecycleOwner.hashCode()] = it
-                    lifecycleOwner.lifecycle.addObserver(
-                        object : DefaultLifecycleObserver {
-                            override fun onDestroy(owner: LifecycleOwner) {
-                                super.onDestroy(owner)
-                                clockRegistries[lifecycleOwner.hashCode()]?.unregisterListeners()
-                                clockRegistries.remove(lifecycleOwner.hashCode())
-                            }
-                        }
+        return (clockRegistryProvider
+                ?: ClockRegistryProvider(
+                        context = context,
+                        coroutineScope = getApplicationCoroutineScope(),
+                        mainDispatcher = Dispatchers.Main,
+                        backgroundDispatcher = Dispatchers.IO,
                     )
-                }
+                    .also { clockRegistryProvider = it })
+            .getForOwner(lifecycleOwner)
     }
 
     override fun getClockPickerInteractor(
@@ -346,11 +358,14 @@
     ): ClockPickerInteractor {
         return clockPickerInteractor
             ?: ClockPickerInteractor(
-                    ClockPickerRepositoryImpl(
-                        secureSettingsRepository = getSecureSettingsRepository(context),
-                        registry = getClockRegistry(context, lifecycleOwner),
-                        scope = getApplicationCoroutineScope(),
-                    ),
+                    repository =
+                        ClockPickerRepositoryImpl(
+                            secureSettingsRepository = getSecureSettingsRepository(context),
+                            registry = getClockRegistry(context, lifecycleOwner),
+                            scope = getApplicationCoroutineScope(),
+                            mainDispatcher = Dispatchers.Main,
+                        ),
+                    snapshotRestorer = { getClockPickerSnapshotRestorer(context, lifecycleOwner) },
                 )
                 .also { clockPickerInteractor = it }
     }
@@ -380,7 +395,11 @@
                     activity.applicationContext,
                     ScreenSizeCalculator.getInstance()
                         .getScreenSize(activity.windowManager.defaultDisplay),
-                    getClockRegistry(activity.applicationContext, activity),
+                    WallpaperManager.getInstance(activity.applicationContext),
+                    getClockRegistry(
+                        context = activity.applicationContext,
+                        lifecycleOwner = activity,
+                    ),
                 )
                 .also {
                     clockViewFactories[activityHashCode] = it
@@ -396,20 +415,14 @@
                 }
     }
 
-    private fun getNotificationsInteractor(
+    private fun getClockPickerSnapshotRestorer(
         context: Context,
-    ): NotificationsInteractor {
-        return notificationsInteractor
-            ?: NotificationsInteractor(
-                    repository =
-                        NotificationsRepository(
-                            scope = getApplicationCoroutineScope(),
-                            backgroundDispatcher = Dispatchers.IO,
-                            secureSettingsRepository = getSecureSettingsRepository(context),
-                        ),
-                    snapshotRestorer = { getNotificationsSnapshotRestorer(context) },
-                )
-                .also { notificationsInteractor = it }
+        lifecycleOwner: LifecycleOwner
+    ): ClockPickerSnapshotRestorer {
+        return clockPickerSnapshotRestorer
+            ?: ClockPickerSnapshotRestorer(getClockPickerInteractor(context, lifecycleOwner)).also {
+                clockPickerSnapshotRestorer = it
+            }
     }
 
     override fun getColorPickerInteractor(
@@ -570,6 +583,7 @@
         private val KEY_APP_GRID_SNAPSHOT_RESTORER = KEY_THEMED_ICON_SNAPSHOT_RESTORER + 1
         @JvmStatic
         private val KEY_COLOR_PICKER_SNAPSHOT_RESTORER = KEY_APP_GRID_SNAPSHOT_RESTORER + 1
+        @JvmStatic private val KEY_CLOCKS_SNAPSHOT_RESTORER = KEY_COLOR_PICKER_SNAPSHOT_RESTORER + 1
 
         /**
          * When this injector is overridden, this is the minimal value that should be used by
@@ -577,6 +591,6 @@
          *
          * It should always be greater than the biggest restorer key.
          */
-        @JvmStatic protected val MIN_SNAPSHOT_RESTORER_KEY = KEY_COLOR_PICKER_SNAPSHOT_RESTORER + 1
+        @JvmStatic protected val MIN_SNAPSHOT_RESTORER_KEY = KEY_CLOCKS_SNAPSHOT_RESTORER + 1
     }
 }
diff --git a/src/com/android/customization/picker/clock/data/repository/ClockPickerRepository.kt b/src/com/android/customization/picker/clock/data/repository/ClockPickerRepository.kt
index cb2c86e..57f77b0 100644
--- a/src/com/android/customization/picker/clock/data/repository/ClockPickerRepository.kt
+++ b/src/com/android/customization/picker/clock/data/repository/ClockPickerRepository.kt
@@ -42,7 +42,7 @@
      * @param colorToneProgress color tone from 0 to 100 to apply to the selected color
      * @param seedColor the actual clock color after blending the selected color and color tone
      */
-    fun setClockColor(
+    suspend fun setClockColor(
         selectedColorId: String?,
         @IntRange(from = 0, to = 100) colorToneProgress: Int,
         @ColorInt seedColor: Int?,
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 747f174..be6c6cb 100644
--- a/src/com/android/customization/picker/clock/data/repository/ClockPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/clock/data/repository/ClockPickerRepositoryImpl.kt
@@ -24,6 +24,7 @@
 import com.android.systemui.plugins.ClockMetadata
 import com.android.systemui.shared.clocks.ClockRegistry
 import com.android.wallpaper.settings.data.repository.SecureSettingsRepository
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.channels.awaitClose
@@ -32,6 +33,7 @@
 import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.callbackFlow
+import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.mapLatest
 import kotlinx.coroutines.flow.mapNotNull
@@ -43,6 +45,7 @@
     private val secureSettingsRepository: SecureSettingsRepository,
     private val registry: ClockRegistry,
     scope: CoroutineScope,
+    mainDispatcher: CoroutineDispatcher,
 ) : ClockPickerRepository {
 
     @OptIn(ExperimentalCoroutinesApi::class)
@@ -67,6 +70,7 @@
                 send()
                 awaitClose { registry.unregisterClockChangeListener(listener) }
             }
+            .flowOn(mainDispatcher)
             .mapLatest { allClocks ->
                 // Loading list of clock plugins can cause many consecutive calls of
                 // onAvailableClocksChanged(). We only care about the final fully-initiated clock
@@ -108,6 +112,7 @@
                 send()
                 awaitClose { registry.unregisterClockChangeListener(listener) }
             }
+            .flowOn(mainDispatcher)
             .mapNotNull { it }
 
     override suspend fun setSelectedClock(clockId: String) {
@@ -118,7 +123,7 @@
         }
     }
 
-    override fun setClockColor(
+    override suspend fun setClockColor(
         selectedColorId: String?,
         @IntRange(from = 0, to = 100) colorToneProgress: Int,
         @ColorInt seedColor: Int?,
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 e359117..52c3c4e 100644
--- a/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt
+++ b/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt
@@ -19,6 +19,8 @@
 import android.content.ComponentName
 import android.content.Context
 import android.view.LayoutInflater
+import androidx.lifecycle.DefaultLifecycleObserver
+import androidx.lifecycle.LifecycleOwner
 import com.android.systemui.plugins.Plugin
 import com.android.systemui.plugins.PluginManager
 import com.android.systemui.shared.clocks.ClockRegistry
@@ -43,6 +45,7 @@
     private val mainDispatcher: CoroutineDispatcher,
     private val backgroundDispatcher: CoroutineDispatcher,
 ) {
+    private val lifecycleOwners = mutableSetOf<Int>()
     private val pluginManager: PluginManager by lazy { createPluginManager(context) }
     private val clockRegistry: ClockRegistry by lazy {
         ClockRegistry(
@@ -60,10 +63,32 @@
             .apply { registerListeners() }
     }
 
-    fun get(): ClockRegistry {
+    fun getForOwner(lifecycleOwner: LifecycleOwner): ClockRegistry {
+        registerLifecycleOwner(lifecycleOwner)
         return clockRegistry
     }
 
+    private fun registerLifecycleOwner(lifecycleOwner: LifecycleOwner) {
+        lifecycleOwners.add(lifecycleOwner.hashCode())
+
+        lifecycleOwner.lifecycle.addObserver(
+            object : DefaultLifecycleObserver {
+                override fun onDestroy(owner: LifecycleOwner) {
+                    super.onDestroy(owner)
+                    unregisterLifecycleOwner(owner)
+                }
+            }
+        )
+    }
+
+    private fun unregisterLifecycleOwner(lifecycleOwner: LifecycleOwner) {
+        lifecycleOwners.remove(lifecycleOwner.hashCode())
+
+        if (lifecycleOwners.isEmpty()) {
+            clockRegistry.unregisterListeners()
+        }
+    }
+
     private fun createPluginManager(context: Context): PluginManager {
         val privilegedPlugins = listOf<String>()
         val isDebugDevice = true
diff --git a/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractor.kt b/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractor.kt
index 91b2773..30887e5 100644
--- a/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractor.kt
+++ b/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractor.kt
@@ -22,15 +22,21 @@
 import com.android.customization.picker.clock.data.repository.ClockPickerRepository
 import com.android.customization.picker.clock.shared.ClockSize
 import com.android.customization.picker.clock.shared.model.ClockMetadataModel
+import com.android.customization.picker.clock.shared.model.ClockSnapshotModel
+import javax.inject.Provider
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.firstOrNull
 import kotlinx.coroutines.flow.map
 
 /**
  * Interactor for accessing application clock settings, as well as selecting and configuring custom
  * clocks.
  */
-class ClockPickerInteractor(private val repository: ClockPickerRepository) {
+class ClockPickerInteractor(
+    private val repository: ClockPickerRepository,
+    private val snapshotRestorer: Provider<ClockPickerSnapshotRestorer>,
+) {
 
     val allClocks: Flow<List<ClockMetadataModel>> = repository.allClocks
 
@@ -48,18 +54,68 @@
     val selectedClockSize: Flow<ClockSize> = repository.selectedClockSize
 
     suspend fun setSelectedClock(clockId: String) {
-        repository.setSelectedClock(clockId)
+        // Use the [clockId] to override saved clock id, since it might not be updated in time
+        setClockOption(ClockSnapshotModel(clockId = clockId))
     }
 
-    fun setClockColor(
+    suspend fun setClockColor(
         selectedColorId: String?,
         @IntRange(from = 0, to = 100) colorToneProgress: Int,
         @ColorInt seedColor: Int?,
     ) {
-        repository.setClockColor(selectedColorId, colorToneProgress, seedColor)
+        // Use the color to override saved color, since it might not be updated in time
+        setClockOption(
+            ClockSnapshotModel(
+                selectedColorId = selectedColorId,
+                colorToneProgress = colorToneProgress,
+                seedColor = seedColor,
+            )
+        )
     }
 
     suspend fun setClockSize(size: ClockSize) {
-        repository.setClockSize(size)
+        // Use the [ClockSize] to override saved clock size, since it might not be updated in time
+        setClockOption(ClockSnapshotModel(clockSize = size))
+    }
+
+    suspend fun setClockOption(clockSnapshotModel: ClockSnapshotModel) {
+        // [ClockCarouselViewModel] is monitoring the [ClockPickerInteractor.setSelectedClock] job,
+        // so it needs to finish last.
+        storeCurrentClockOption(clockSnapshotModel)
+
+        clockSnapshotModel.clockSize?.let { repository.setClockSize(it) }
+        clockSnapshotModel.colorToneProgress?.let {
+            repository.setClockColor(
+                selectedColorId = clockSnapshotModel.selectedColorId,
+                colorToneProgress = clockSnapshotModel.colorToneProgress,
+                seedColor = clockSnapshotModel.seedColor
+            )
+        }
+        clockSnapshotModel.clockId?.let { repository.setSelectedClock(it) }
+    }
+
+    /**
+     * Gets the [ClockSnapshotModel] from the storage and override with [latestOption].
+     *
+     * The storage might be in the middle of a write, and not reflecting the user's options, always
+     * pass in a [ClockSnapshotModel] if we know it's the latest option from a user's point of view.
+     *
+     * [selectedColorId] and [seedColor] have null state collide with nullable type, but we know
+     * they are presented whenever there's a [colorToneProgress].
+     */
+    suspend fun getCurrentClockToRestore(latestOption: ClockSnapshotModel? = null) =
+        ClockSnapshotModel(
+            clockId = latestOption?.clockId ?: selectedClockId.firstOrNull(),
+            clockSize = latestOption?.clockSize ?: selectedClockSize.firstOrNull(),
+            colorToneProgress = latestOption?.colorToneProgress ?: colorToneProgress.firstOrNull(),
+            selectedColorId = latestOption?.colorToneProgress?.let { latestOption.selectedColorId }
+                    ?: selectedColorId.firstOrNull(),
+            seedColor = latestOption?.colorToneProgress?.let { latestOption.seedColor }
+                    ?: seedColor.firstOrNull(),
+        )
+
+    private suspend fun storeCurrentClockOption(clockSnapshotModel: ClockSnapshotModel) {
+        val option = getCurrentClockToRestore(clockSnapshotModel)
+        snapshotRestorer.get().storeSnapshot(option)
     }
 }
diff --git a/src/com/android/customization/picker/clock/domain/interactor/ClockPickerSnapshotRestorer.kt b/src/com/android/customization/picker/clock/domain/interactor/ClockPickerSnapshotRestorer.kt
new file mode 100644
index 0000000..ecaf10f
--- /dev/null
+++ b/src/com/android/customization/picker/clock/domain/interactor/ClockPickerSnapshotRestorer.kt
@@ -0,0 +1,94 @@
+/*
+ * 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.clock.domain.interactor
+
+import android.text.TextUtils
+import android.util.Log
+import com.android.customization.picker.clock.shared.model.ClockSnapshotModel
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
+import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
+
+/** Handles state restoration for clocks. */
+class ClockPickerSnapshotRestorer(private val interactor: ClockPickerInteractor) :
+    SnapshotRestorer {
+    private var snapshotStore: SnapshotStore = SnapshotStore.NOOP
+    private var originalOption: ClockSnapshotModel? = null
+
+    override suspend fun setUpSnapshotRestorer(
+        store: SnapshotStore,
+    ): RestorableSnapshot {
+        snapshotStore = store
+        originalOption = interactor.getCurrentClockToRestore()
+        return snapshot(originalOption)
+    }
+
+    override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
+        originalOption?.let { optionToRestore ->
+            if (
+                TextUtils.isEmpty(optionToRestore.clockId) ||
+                    optionToRestore.clockId != snapshot.args[KEY_CLOCK_ID] ||
+                    optionToRestore.clockSize?.toString() != snapshot.args[KEY_CLOCK_SIZE] ||
+                    optionToRestore.colorToneProgress?.toString() !=
+                        snapshot.args[KEY_COLOR_TONE_PROGRESS] ||
+                    optionToRestore.seedColor?.toString() != snapshot.args[KEY_SEED_COLOR] ||
+                    optionToRestore.selectedColorId != snapshot.args[KEY_COLOR_ID]
+            ) {
+                Log.wtf(
+                    TAG,
+                    """ Original clock option does not match snapshot option to restore to. The
+                        | current implementation doesn't support undo, only a reset back to the
+                        | original clock option."""
+                        .trimMargin(),
+                )
+            }
+
+            interactor.setClockOption(optionToRestore)
+        }
+    }
+
+    fun storeSnapshot(clockSnapshotModel: ClockSnapshotModel) {
+        snapshotStore.store(snapshot(clockSnapshotModel))
+    }
+
+    private fun snapshot(clockSnapshotModel: ClockSnapshotModel? = null): RestorableSnapshot {
+        val options =
+            if (clockSnapshotModel == null) emptyMap()
+            else
+                buildMap {
+                    clockSnapshotModel.clockId?.let { put(KEY_CLOCK_ID, it) }
+                    clockSnapshotModel.clockSize?.let { put(KEY_CLOCK_SIZE, it.toString()) }
+                    clockSnapshotModel.selectedColorId?.let { put(KEY_COLOR_ID, it) }
+                    clockSnapshotModel.colorToneProgress?.let {
+                        put(KEY_COLOR_TONE_PROGRESS, it.toString())
+                    }
+                    clockSnapshotModel.seedColor?.let { put(KEY_SEED_COLOR, it.toString()) }
+                }
+
+        return RestorableSnapshot(options)
+    }
+
+    companion object {
+        private const val TAG = "ClockPickerSnapshotRestorer"
+        private const val KEY_CLOCK_ID = "clock_id"
+        private const val KEY_CLOCK_SIZE = "clock_size"
+        private const val KEY_COLOR_ID = "color_id"
+        private const val KEY_COLOR_TONE_PROGRESS = "color_tone_progress"
+        private const val KEY_SEED_COLOR = "seed_color"
+    }
+}
diff --git a/src/com/android/customization/picker/clock/domain/interactor/ClocksSnapshotRestorer.kt b/src/com/android/customization/picker/clock/domain/interactor/ClocksSnapshotRestorer.kt
deleted file mode 100644
index 7bb3232..0000000
--- a/src/com/android/customization/picker/clock/domain/interactor/ClocksSnapshotRestorer.kt
+++ /dev/null
@@ -1,36 +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.clock.domain.interactor
-
-import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
-import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
-import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
-
-/** Handles state restoration for clocks. */
-class ClocksSnapshotRestorer : SnapshotRestorer {
-    override suspend fun setUpSnapshotRestorer(
-        store: SnapshotStore,
-    ): RestorableSnapshot {
-        // TODO(b/262924055): implement as part of the clock settings screen.
-        return RestorableSnapshot(mapOf())
-    }
-
-    override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
-        // TODO(b/262924055): implement as part of the clock settings screen.
-    }
-}
diff --git a/src/com/android/customization/picker/clock/shared/model/ClockSnapshotModel.kt b/src/com/android/customization/picker/clock/shared/model/ClockSnapshotModel.kt
new file mode 100644
index 0000000..942cc59
--- /dev/null
+++ b/src/com/android/customization/picker/clock/shared/model/ClockSnapshotModel.kt
@@ -0,0 +1,31 @@
+/*
+ * 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.clock.shared.model
+
+import androidx.annotation.ColorInt
+import androidx.annotation.IntRange
+import com.android.customization.picker.clock.shared.ClockSize
+
+/** Models application state for a clock option in a picker experience. */
+data class ClockSnapshotModel(
+    val clockId: String? = null,
+    val clockSize: ClockSize? = null,
+    val selectedColorId: String? = null,
+    @IntRange(from = 0, to = 100) val colorToneProgress: Int? = null,
+    @ColorInt val seedColor: Int? = null,
+)
diff --git a/src/com/android/customization/picker/clock/ui/adapter/ClockSettingsTabAdapter.kt b/src/com/android/customization/picker/clock/ui/adapter/ClockSettingsTabAdapter.kt
index d0e6f18..746fdb3 100644
--- a/src/com/android/customization/picker/clock/ui/adapter/ClockSettingsTabAdapter.kt
+++ b/src/com/android/customization/picker/clock/ui/adapter/ClockSettingsTabAdapter.kt
@@ -16,17 +16,13 @@
  */
 package com.android.customization.picker.clock.ui.adapter
 
-import android.graphics.Rect
 import android.view.LayoutInflater
-import android.view.TouchDelegate
 import android.view.View
 import android.view.ViewGroup
 import android.widget.TextView
 import androidx.recyclerview.widget.RecyclerView
 import com.android.customization.picker.clock.ui.viewmodel.ClockSettingsTabViewModel
 import com.android.wallpaper.R
-import com.android.wallpaper.util.ViewUtils.setupTouchDelegate
-
 
 /** Adapter for the tab recycler view on the clock settings screen. */
 class ClockSettingsTabAdapter : RecyclerView.Adapter<ClockSettingsTabAdapter.ViewHolder>() {
@@ -58,10 +54,6 @@
         val item = items[position]
         holder.itemView.isSelected = item.isSelected
         holder.textView.text = item.name
-        holder.textView.setupTouchDelegate(
-            parentView = holder.itemView,
-            heightRes = R.dimen.tab_touch_delegate_height_padding
-        )
         holder.textView.setOnClickListener(
             if (item.onClicked != null) {
                 View.OnClickListener { item.onClicked.invoke() }
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 e703517..270931c 100644
--- a/src/com/android/customization/picker/clock/ui/binder/ClockCarouselViewBinder.kt
+++ b/src/com/android/customization/picker/clock/ui/binder/ClockCarouselViewBinder.kt
@@ -41,6 +41,7 @@
         lifecycleOwner: LifecycleOwner,
     ) {
         carouselView.setClockViewFactory(clockViewFactory)
+        clockViewFactory.updateRegionDarkness()
         val singleClockHostView =
             singleClockView.requireViewById<FrameLayout>(R.id.single_clock_host_view)
         lifecycleOwner.lifecycleScope.launch {
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 671a7ae..d8c5dce 100644
--- a/src/com/android/customization/picker/clock/ui/binder/ClockSettingsBinder.kt
+++ b/src/com/android/customization/picker/clock/ui/binder/ClockSettingsBinder.kt
@@ -75,7 +75,9 @@
 
                 override fun onStartTrackingTouch(seekBar: SeekBar?) = Unit
                 override fun onStopTrackingTouch(seekBar: SeekBar?) {
-                    seekBar?.progress?.let { viewModel.onSliderProgressStop(it) }
+                    seekBar?.progress?.let {
+                        lifecycleOwner.lifecycleScope.launch { viewModel.onSliderProgressStop(it) }
+                    }
                 }
             }
         )
diff --git a/src/com/android/customization/picker/clock/ui/view/ClockViewFactory.kt b/src/com/android/customization/picker/clock/ui/view/ClockViewFactory.kt
index 1aa6206..3d490d4 100644
--- a/src/com/android/customization/picker/clock/ui/view/ClockViewFactory.kt
+++ b/src/com/android/customization/picker/clock/ui/view/ClockViewFactory.kt
@@ -15,11 +15,12 @@
  */
 package com.android.customization.picker.clock.ui.view
 
+import android.app.WallpaperColors
+import android.app.WallpaperManager
 import android.content.Context
 import android.content.res.Resources
 import android.graphics.Point
 import android.graphics.Rect
-import android.util.TypedValue
 import android.view.View
 import android.widget.FrameLayout
 import androidx.annotation.ColorInt
@@ -39,6 +40,7 @@
 class ClockViewFactory(
     private val appContext: Context,
     val screenSize: Point,
+    private val wallpaperManager: WallpaperManager,
     private val registry: ClockRegistry,
 ) {
     private val resources = appContext.resources
@@ -89,9 +91,21 @@
     }
 
     fun updateColor(clockId: String, @ColorInt seedColor: Int?) {
-        return (clockControllers[clockId] ?: initClockController(clockId))
-            .events
-            .onSeedColorChanged(seedColor)
+        clockControllers[clockId]?.events?.onSeedColorChanged(seedColor)
+    }
+
+    fun updateRegionDarkness() {
+        val isRegionDark = isLockscreenWallpaperDark()
+        clockControllers.values.forEach {
+            it.largeClock.events.onRegionDarknessChanged(isRegionDark)
+            it.smallClock.events.onRegionDarknessChanged(isRegionDark)
+        }
+    }
+
+    private fun isLockscreenWallpaperDark(): Boolean {
+        val colors = wallpaperManager.getWallpaperColors(WallpaperManager.FLAG_LOCK)
+        return (colors?.colorHints?.and(WallpaperColors.HINT_SUPPORTS_DARK_TEXT)) !=
+            WallpaperColors.HINT_SUPPORTS_DARK_TEXT
     }
 
     fun updateTimeFormat(clockId: String) {
@@ -136,18 +150,16 @@
             registry.createExampleClock(clockId).also { it?.initialize(resources, 0f, 0f) }
         checkNotNull(controller)
 
-        // Configure light/dark theme
-        val isLightTheme = TypedValue()
-        appContext.theme.resolveAttribute(android.R.attr.isLightTheme, isLightTheme, true)
-        val isRegionDark = isLightTheme.data == 0
-        controller.largeClock.events.onRegionDarknessChanged(isRegionDark)
-        // Configure font size
+        val isWallpaperDark = isLockscreenWallpaperDark()
+        // Initialize large clock
+        controller.largeClock.events.onRegionDarknessChanged(isWallpaperDark)
         controller.largeClock.events.onFontSettingChanged(
             resources.getDimensionPixelSize(R.dimen.large_clock_text_size).toFloat()
         )
         controller.largeClock.events.onTargetRegionChanged(getLargeClockRegion())
 
-        controller.smallClock.events.onRegionDarknessChanged(isRegionDark)
+        // Initialize small clock
+        controller.smallClock.events.onRegionDarknessChanged(isWallpaperDark)
         controller.smallClock.events.onFontSettingChanged(
             resources.getDimensionPixelSize(R.dimen.small_clock_text_size).toFloat()
         )
diff --git a/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModel.kt b/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModel.kt
index b0ff1db..a498c71 100644
--- a/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModel.kt
+++ b/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModel.kt
@@ -103,7 +103,7 @@
             )
     }
 
-    fun onSliderProgressStop(progress: Int) {
+    suspend fun onSliderProgressStop(progress: Int) {
         val selectedColorId = selectedColorId.value ?: return
         val clockColorViewModel = colorMap[selectedColorId] ?: return
         clockPickerInteractor.setClockColor(
@@ -168,18 +168,20 @@
                                         null
                                     } else {
                                         {
-                                            clockPickerInteractor.setClockColor(
-                                                selectedColorId = colorModel.colorId,
-                                                colorToneProgress = colorToneProgress,
-                                                seedColor =
-                                                    blendColorWithTone(
-                                                        color = colorModel.color,
-                                                        colorTone =
-                                                            colorModel.getColorTone(
-                                                                colorToneProgress,
-                                                            ),
-                                                    ),
-                                            )
+                                            viewModelScope.launch {
+                                                clockPickerInteractor.setClockColor(
+                                                    selectedColorId = colorModel.colorId,
+                                                    colorToneProgress = colorToneProgress,
+                                                    seedColor =
+                                                        blendColorWithTone(
+                                                            color = colorModel.color,
+                                                            colorTone =
+                                                                colorModel.getColorTone(
+                                                                    colorToneProgress,
+                                                                ),
+                                                        ),
+                                                )
+                                            }
                                         }
                                     }
                                 },
@@ -235,11 +237,14 @@
                         null
                     } else {
                         {
-                            clockPickerInteractor.setClockColor(
-                                selectedColorId = null,
-                                colorToneProgress = ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS,
-                                seedColor = null,
-                            )
+                            viewModelScope.launch {
+                                clockPickerInteractor.setClockColor(
+                                    selectedColorId = null,
+                                    colorToneProgress =
+                                        ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS,
+                                    seedColor = null,
+                                )
+                            }
                         }
                     }
                 },
diff --git a/src/com/android/customization/picker/color/ui/adapter/ColorTypeTabAdapter.kt b/src/com/android/customization/picker/color/ui/adapter/ColorTypeTabAdapter.kt
index ac5ad81..bb9f082 100644
--- a/src/com/android/customization/picker/color/ui/adapter/ColorTypeTabAdapter.kt
+++ b/src/com/android/customization/picker/color/ui/adapter/ColorTypeTabAdapter.kt
@@ -17,16 +17,13 @@
 
 package com.android.customization.picker.color.ui.adapter
 
-import android.graphics.Rect
 import android.view.LayoutInflater
-import android.view.TouchDelegate
 import android.view.View
 import android.view.ViewGroup
 import android.widget.TextView
 import androidx.recyclerview.widget.RecyclerView
 import com.android.customization.picker.color.ui.viewmodel.ColorTypeTabViewModel
 import com.android.wallpaper.R
-import com.android.wallpaper.util.ViewUtils.setupTouchDelegate
 
 /** Adapts between color type items and views. */
 class ColorTypeTabAdapter : RecyclerView.Adapter<ColorTypeTabAdapter.ViewHolder>() {
@@ -58,10 +55,6 @@
         val item = items[position]
         holder.itemView.isSelected = item.isSelected
         holder.textView.text = item.name
-        holder.textView.setupTouchDelegate(
-            parentView = holder.itemView,
-            heightRes = R.dimen.tab_touch_delegate_height_padding
-        )
         holder.textView.setOnClickListener(
             if (item.onClick != null) {
                 View.OnClickListener { item.onClick.invoke() }
diff --git a/src/com/android/customization/picker/quickaffordance/ui/adapter/SlotTabAdapter.kt b/src/com/android/customization/picker/quickaffordance/ui/adapter/SlotTabAdapter.kt
index b34ea1b..5203ed3 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/adapter/SlotTabAdapter.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/adapter/SlotTabAdapter.kt
@@ -17,16 +17,13 @@
 
 package com.android.customization.picker.quickaffordance.ui.adapter
 
-import android.graphics.Rect
 import android.view.LayoutInflater
-import android.view.TouchDelegate
 import android.view.View
 import android.view.ViewGroup
 import android.widget.TextView
 import androidx.recyclerview.widget.RecyclerView
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordanceSlotViewModel
 import com.android.wallpaper.R
-import com.android.wallpaper.util.ViewUtils.setupTouchDelegate
 
 /** Adapts between lock screen quick affordance slot items and views. */
 class SlotTabAdapter : RecyclerView.Adapter<SlotTabAdapter.ViewHolder>() {
@@ -58,10 +55,6 @@
         val item = items[position]
         holder.itemView.isSelected = item.isSelected
         holder.textView.text = item.name
-        holder.textView.setupTouchDelegate(
-            parentView = holder.itemView,
-            heightRes = R.dimen.tab_touch_delegate_height_padding
-        )
         holder.textView.setOnClickListener(
             if (item.onClicked != null) {
                 View.OnClickListener { item.onClicked.invoke() }
diff --git a/tests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt b/tests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
index 38bf25a..bf2766d 100644
--- a/tests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
+++ b/tests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
@@ -60,7 +60,7 @@
         selectedClockId.value = clockId
     }
 
-    override fun setClockColor(
+    override suspend fun setClockColor(
         selectedColorId: String?,
         @IntRange(from = 0, to = 100) colorToneProgress: Int,
         @ColorInt seedColor: Int?,
diff --git a/tests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt b/tests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
index cd41d7d..1a7ebb5 100644
--- a/tests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
+++ b/tests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
@@ -3,10 +3,12 @@
 import androidx.test.filters.SmallTest
 import com.android.customization.picker.clock.data.repository.FakeClockPickerRepository
 import com.android.customization.picker.clock.shared.ClockSize
+import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.resetMain
 import kotlinx.coroutines.test.runTest
@@ -28,7 +30,15 @@
     fun setUp() {
         val testDispatcher = StandardTestDispatcher()
         Dispatchers.setMain(testDispatcher)
-        underTest = ClockPickerInteractor(FakeClockPickerRepository())
+        underTest =
+            ClockPickerInteractor(
+                repository = FakeClockPickerRepository(),
+                snapshotRestorer = {
+                    ClockPickerSnapshotRestorer(interactor = underTest).apply {
+                        runBlocking { setUpSnapshotRestorer(store = FakeSnapshotStore()) }
+                    }
+                },
+            )
     }
 
     @After
diff --git a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt b/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
index d3e458f..c5eb796 100644
--- a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
+++ b/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
@@ -16,14 +16,18 @@
 package com.android.customization.picker.clock.ui.viewmodel
 
 import androidx.test.filters.SmallTest
+import com.android.customization.picker.clock.data.repository.ClockPickerRepository
 import com.android.customization.picker.clock.data.repository.FakeClockPickerRepository
 import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
+import com.android.customization.picker.clock.domain.interactor.ClockPickerSnapshotRestorer
 import com.android.customization.picker.clock.shared.model.ClockMetadataModel
+import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.advanceTimeBy
 import kotlinx.coroutines.test.resetMain
@@ -55,6 +59,7 @@
     }
     private lateinit var testDispatcher: CoroutineDispatcher
     private lateinit var underTest: ClockCarouselViewModel
+    private lateinit var interactor: ClockPickerInteractor
 
     @Before
     fun setUp() {
@@ -71,12 +76,14 @@
     fun setSelectedClock() = runTest {
         underTest =
             ClockCarouselViewModel(
-                ClockPickerInteractor(repositoryWithMultipleClocks),
-                testDispatcher,
+                getClockPickerInteractor(repositoryWithMultipleClocks),
+                testDispatcher
             )
         val observedSelectedIndex = collectLastValue(underTest.selectedIndex)
         advanceTimeBy(ClockCarouselViewModel.CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
+
         underTest.setSelectedClock(FakeClockPickerRepository.fakeClocks[2].clockId)
+
         assertThat(observedSelectedIndex()).isEqualTo(2)
     }
 
@@ -84,12 +91,14 @@
     fun multipleClockCase() = runTest {
         underTest =
             ClockCarouselViewModel(
-                ClockPickerInteractor(repositoryWithMultipleClocks),
-                testDispatcher,
+                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()
     }
@@ -98,13 +107,27 @@
     fun singleClockCase() = runTest {
         underTest =
             ClockCarouselViewModel(
-                ClockPickerInteractor(repositoryWithSingleClock),
-                testDispatcher,
+                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,
+                snapshotRestorer = {
+                    ClockPickerSnapshotRestorer(interactor = interactor).apply {
+                        runBlocking { setUpSnapshotRestorer(store = FakeSnapshotStore()) }
+                    }
+                }
+            )
+            .also { interactor = it }
+    }
 }
diff --git a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt b/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
index 573777d..293e393 100644
--- a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
+++ b/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
@@ -19,12 +19,15 @@
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.customization.picker.clock.data.repository.FakeClockPickerRepository
 import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
+import com.android.customization.picker.clock.domain.interactor.ClockPickerSnapshotRestorer
 import com.android.customization.picker.clock.shared.ClockSize
 import com.android.customization.picker.clock.shared.model.ClockMetadataModel
+import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.resetMain
 import kotlinx.coroutines.test.runTest
@@ -50,7 +53,15 @@
         Dispatchers.setMain(testDispatcher)
         val context = InstrumentationRegistry.getInstrumentation().targetContext
         clockColorMap = ClockColorViewModel.getPresetColorMap(context.resources)
-        interactor = ClockPickerInteractor(FakeClockPickerRepository())
+        interactor =
+            ClockPickerInteractor(
+                repository = FakeClockPickerRepository(),
+                snapshotRestorer = {
+                    ClockPickerSnapshotRestorer(interactor = interactor).apply {
+                        runBlocking { setUpSnapshotRestorer(store = FakeSnapshotStore()) }
+                    }
+                },
+            )
         underTest =
             ClockSectionViewModel(
                 context,
@@ -68,6 +79,7 @@
         val colorGrey = clockColorMap.values.first()
         val observedSelectedClockColorAndSizeText =
             collectLastValue(underTest.selectedClockColorAndSizeText)
+
         interactor.setClockColor(
             colorGrey.colorId,
             ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS,
@@ -77,6 +89,7 @@
             )
         )
         interactor.setClockSize(ClockSize.DYNAMIC)
+
         assertThat(observedSelectedClockColorAndSizeText()).isEqualTo("Grey, dynamic")
     }
 }
diff --git a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt b/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
index a329bb3..f58baf8 100644
--- a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
+++ b/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
@@ -5,6 +5,7 @@
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.customization.picker.clock.data.repository.FakeClockPickerRepository
 import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
+import com.android.customization.picker.clock.domain.interactor.ClockPickerSnapshotRestorer
 import com.android.customization.picker.clock.shared.ClockSize
 import com.android.customization.picker.clock.shared.model.ClockMetadataModel
 import com.android.customization.picker.color.data.repository.FakeColorPickerRepository
@@ -57,7 +58,15 @@
         Dispatchers.setMain(testDispatcher)
         context = InstrumentationRegistry.getInstrumentation().targetContext
         testScope = TestScope(testDispatcher)
-        clockPickerInteractor = ClockPickerInteractor(FakeClockPickerRepository())
+        clockPickerInteractor =
+            ClockPickerInteractor(
+                repository = FakeClockPickerRepository(),
+                snapshotRestorer = {
+                    ClockPickerSnapshotRestorer(interactor = clockPickerInteractor).apply {
+                        runBlocking { setUpSnapshotRestorer(store = FakeSnapshotStore()) }
+                    }
+                },
+            )
         colorPickerInteractor =
             ColorPickerInteractor(
                 repository = FakeColorPickerRepository(context = context),
@@ -160,7 +169,7 @@
         underTest.onSliderProgressChanged(targetProgress1)
         assertThat(observedSliderProgress()).isEqualTo(targetProgress1)
         val targetProgress2 = 55
-        underTest.onSliderProgressStop(targetProgress2)
+        testScope.launch { underTest.onSliderProgressStop(targetProgress2) }
         assertThat(observedSliderProgress()).isEqualTo(targetProgress2)
         val expectedSelectedColorModel = colorMap.values.first() // RED
         assertThat(observedSeedColor())