Make Clock options resettable (1/3)

Moved `ClockPickerSnapshotRestorer` to `ThemePicker`.

`ClockPickerSnapshotRestorer` creates snapshots based on `ClockOptionModel` which contains all clock related options. The original snapshot created from the latest options in the repository, after user made a change, the snapshot is created from the repository and override with the change, this is to prevent race condition where the change has not been written in the `SecureSettings`.

Bug: 274403822
Test: play with Clock options
Change-Id: Ia5475e3b367604b96a3dc6533c4cdbacc1f1fcb5
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index 66814c5..89f0233 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -45,6 +45,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
@@ -107,6 +108,7 @@
     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
@@ -195,18 +197,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 +274,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 +306,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(
@@ -346,11 +372,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 }
     }
@@ -396,20 +425,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 +593,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 +601,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/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/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/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/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())