Merge tag 'android-13.0.0_r52' of https://android.googlesource.com/platform/packages/apps/ThemePicker into HEADHEADt13.0

Android 13.0.0 Release 52 (TQ3A.230605.012)

Change-Id: I2cea11fa2f1f02fbd3c9d21cfc1697a79d42a5b7

12 files changed, 1339 insertions, 68 deletions

diff --git a/tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt b/tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
new file mode 100644
index 00000000..59539379
--- /dev/null
+++ b/tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
@@ -0,0 +1,91 @@
+/*
+ * 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.model.grid.data.repository
+
+import com.android.customization.model.grid.shared.model.GridOptionItemModel
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.channels.BufferOverflow
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.asSharedFlow
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
+
+class FakeGridRepository(
+    private val scope: CoroutineScope,
+    initialOptionCount: Int,
+    var available: Boolean = true
+) : GridRepository {
+    private val _optionChanges =
+        MutableSharedFlow<Unit>(
+            replay = 1,
+            onBufferOverflow = BufferOverflow.DROP_OLDEST,
+        )
+
+    override suspend fun isAvailable(): Boolean = available
+
+    override fun getOptionChanges(): Flow<Unit> = _optionChanges.asSharedFlow()
+
+    private val selectedOptionIndex = MutableStateFlow(0)
+    private var options: GridOptionItemsModel = createOptions(count = initialOptionCount)
+
+    override suspend fun getOptions(): GridOptionItemsModel {
+        return options
+    }
+
+    fun setOptions(
+        count: Int,
+        selectedIndex: Int = 0,
+    ) {
+        options = createOptions(count, selectedIndex)
+        _optionChanges.tryEmit(Unit)
+    }
+
+    private fun createOptions(
+        count: Int,
+        selectedIndex: Int = 0,
+    ): GridOptionItemsModel {
+        selectedOptionIndex.value = selectedIndex
+        return GridOptionItemsModel.Loaded(
+            options =
+                buildList {
+                    repeat(times = count) { index ->
+                        add(
+                            GridOptionItemModel(
+                                name = "option_$index",
+                                cols = 4,
+                                rows = index * 2,
+                                isSelected =
+                                    selectedOptionIndex
+                                        .map { it == index }
+                                        .stateIn(
+                                            scope = scope,
+                                            started = SharingStarted.Eagerly,
+                                            initialValue = false,
+                                        ),
+                                onSelected = { selectedOptionIndex.value = index },
+                            )
+                        )
+                    }
+                }
+        )
+    }
+}
diff --git a/tests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt b/tests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt
new file mode 100644
index 00000000..f73d5a38
--- /dev/null
+++ b/tests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt
@@ -0,0 +1,146 @@
+/*
+ * 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.model.grid.domain.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.customization.model.grid.data.repository.FakeGridRepository
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class GridInteractorTest {
+
+    private lateinit var underTest: GridInteractor
+    private lateinit var testScope: TestScope
+    private lateinit var repository: FakeGridRepository
+    private lateinit var store: FakeSnapshotStore
+
+    @Before
+    fun setUp() {
+        testScope = TestScope()
+        repository =
+            FakeGridRepository(
+                scope = testScope.backgroundScope,
+                initialOptionCount = 3,
+            )
+        store = FakeSnapshotStore()
+        underTest =
+            GridInteractor(
+                applicationScope = testScope.backgroundScope,
+                repository = repository,
+                snapshotRestorer = {
+                    GridSnapshotRestorer(
+                            interactor = underTest,
+                        )
+                        .apply {
+                            runBlocking {
+                                setUpSnapshotRestorer(
+                                    store = store,
+                                )
+                            }
+                        }
+                },
+            )
+    }
+
+    @Test
+    fun selectingOptionThroughModel_updatesOptions() =
+        testScope.runTest {
+            val options = collectLastValue(underTest.options)
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(3)
+                assertThat(loaded.options[0].isSelected.value).isTrue()
+                assertThat(loaded.options[1].isSelected.value).isFalse()
+                assertThat(loaded.options[2].isSelected.value).isFalse()
+            }
+
+            val storedSnapshot = store.retrieve()
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                loaded.options[1].onSelected()
+            }
+
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(3)
+                assertThat(loaded.options[0].isSelected.value).isFalse()
+                assertThat(loaded.options[1].isSelected.value).isTrue()
+                assertThat(loaded.options[2].isSelected.value).isFalse()
+            }
+            assertThat(store.retrieve()).isNotEqualTo(storedSnapshot)
+        }
+
+    @Test
+    fun selectingOptionThroughSetter_returnsSelectedOptionFromGetter() =
+        testScope.runTest {
+            val options = collectLastValue(underTest.options)
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(3)
+            }
+
+            val storedSnapshot = store.retrieve()
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                underTest.setSelectedOption(loaded.options[1])
+                runCurrent()
+                assertThat(underTest.getSelectedOption()?.name).isEqualTo(loaded.options[1].name)
+                assertThat(store.retrieve()).isNotEqualTo(storedSnapshot)
+            }
+        }
+
+    @Test
+    fun externalUpdates_reloadInvoked() =
+        testScope.runTest {
+            val options = collectLastValue(underTest.options)
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(3)
+            }
+
+            val storedSnapshot = store.retrieve()
+            repository.setOptions(4)
+
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(4)
+            }
+            // External updates do not record a new snapshot with the undo system.
+            assertThat(store.retrieve()).isEqualTo(storedSnapshot)
+        }
+
+    @Test
+    fun unavailableRepository_emptyOptions() =
+        testScope.runTest {
+            repository.available = false
+            val options = collectLastValue(underTest.options)
+            assertThat(options()).isNull()
+        }
+}
diff --git a/tests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt b/tests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt
new file mode 100644
index 00000000..c2712b1d
--- /dev/null
+++ b/tests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt
@@ -0,0 +1,111 @@
+/*
+ * 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.model.grid.domain.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.customization.model.grid.data.repository.FakeGridRepository
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class GridSnapshotRestorerTest {
+
+    private lateinit var underTest: GridSnapshotRestorer
+    private lateinit var testScope: TestScope
+    private lateinit var repository: FakeGridRepository
+    private lateinit var store: FakeSnapshotStore
+
+    @Before
+    fun setUp() {
+        testScope = TestScope()
+        repository =
+            FakeGridRepository(
+                scope = testScope.backgroundScope,
+                initialOptionCount = 4,
+            )
+        underTest =
+            GridSnapshotRestorer(
+                interactor =
+                    GridInteractor(
+                        applicationScope = testScope.backgroundScope,
+                        repository = repository,
+                        snapshotRestorer = { underTest },
+                    )
+            )
+        store = FakeSnapshotStore()
+    }
+
+    @Test
+    fun restoreToSnapshot_noCallsToStore_restoresToInitialSnapshot() =
+        testScope.runTest {
+            runCurrent()
+            val initialSnapshot = underTest.setUpSnapshotRestorer(store = store)
+            assertThat(initialSnapshot.args).isNotEmpty()
+            repository.setOptions(
+                count = 4,
+                selectedIndex = 2,
+            )
+            runCurrent()
+            assertThat(getSelectedIndex()).isEqualTo(2)
+
+            underTest.restoreToSnapshot(initialSnapshot)
+            runCurrent()
+
+            assertThat(getSelectedIndex()).isEqualTo(0)
+        }
+
+    @Test
+    fun restoreToSnapshot_withCallToStore_restoresToInitialSnapshot() =
+        testScope.runTest {
+            runCurrent()
+            val initialSnapshot = underTest.setUpSnapshotRestorer(store = store)
+            assertThat(initialSnapshot.args).isNotEmpty()
+            repository.setOptions(
+                count = 4,
+                selectedIndex = 2,
+            )
+            runCurrent()
+            assertThat(getSelectedIndex()).isEqualTo(2)
+            underTest.store((repository.getOptions() as GridOptionItemsModel.Loaded).options[1])
+            runCurrent()
+
+            underTest.restoreToSnapshot(initialSnapshot)
+            runCurrent()
+
+            assertThat(getSelectedIndex()).isEqualTo(0)
+        }
+
+    private suspend fun getSelectedIndex(): Int {
+        return (repository.getOptions() as? GridOptionItemsModel.Loaded)?.options?.indexOfFirst {
+            optionItem ->
+            optionItem.isSelected.value
+        }
+            ?: -1
+    }
+}
diff --git a/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt b/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
new file mode 100644
index 00000000..58c5d99f
--- /dev/null
+++ b/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
@@ -0,0 +1,110 @@
+/*
+ * 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.model.grid.ui.viewmodel
+
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.customization.model.grid.data.repository.FakeGridRepository
+import com.android.customization.model.grid.domain.interactor.GridInteractor
+import com.android.customization.model.grid.domain.interactor.GridSnapshotRestorer
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Ignore
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class GridScreenViewModelTest {
+
+    private lateinit var underTest: GridScreenViewModel
+    private lateinit var testScope: TestScope
+    private lateinit var interactor: GridInteractor
+    private lateinit var store: FakeSnapshotStore
+
+    @Before
+    fun setUp() {
+        testScope = TestScope()
+        store = FakeSnapshotStore()
+        interactor =
+            GridInteractor(
+                applicationScope = testScope.backgroundScope,
+                repository =
+                    FakeGridRepository(
+                        scope = testScope.backgroundScope,
+                        initialOptionCount = 4,
+                    ),
+                snapshotRestorer = {
+                    GridSnapshotRestorer(
+                            interactor = interactor,
+                        )
+                        .apply { runBlocking { setUpSnapshotRestorer(store) } }
+                }
+            )
+
+        underTest =
+            GridScreenViewModel(
+                context = InstrumentationRegistry.getInstrumentation().targetContext,
+                interactor = interactor,
+            )
+    }
+
+    @Test
+    @Ignore("b/270371382")
+    fun clickOnItem_itGetsSelected() =
+        testScope.runTest {
+            val optionItemsValueProvider = collectLastValue(underTest.optionItems)
+            var optionItemsValue = checkNotNull(optionItemsValueProvider.invoke())
+            assertThat(optionItemsValue).hasSize(4)
+            assertThat(getSelectedIndex(optionItemsValue)).isEqualTo(0)
+            assertThat(getOnClick(optionItemsValue[0])).isNull()
+
+            val item1OnClickedValue = getOnClick(optionItemsValue[1])
+            assertThat(item1OnClickedValue).isNotNull()
+            item1OnClickedValue?.invoke()
+
+            optionItemsValue = checkNotNull(optionItemsValueProvider.invoke())
+            assertThat(optionItemsValue).hasSize(4)
+            assertThat(getSelectedIndex(optionItemsValue)).isEqualTo(1)
+            assertThat(getOnClick(optionItemsValue[0])).isNotNull()
+            assertThat(getOnClick(optionItemsValue[1])).isNull()
+        }
+
+    private fun TestScope.getSelectedIndex(
+        optionItems: List<OptionItemViewModel<GridIconViewModel>>
+    ): Int {
+        return optionItems.indexOfFirst { optionItem ->
+            collectLastValue(optionItem.isSelected).invoke() == true
+        }
+    }
+
+    private fun TestScope.getOnClick(
+        optionItem: OptionItemViewModel<GridIconViewModel>
+    ): (() -> Unit)? {
+        return collectLastValue(optionItem.onClicked).invoke()
+    }
+}
diff --git a/tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt b/tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
new file mode 100644
index 00000000..38067b75
--- /dev/null
+++ b/tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
@@ -0,0 +1,107 @@
+/*
+ * 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.model.mode
+
+import androidx.test.filters.SmallTest
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class DarkModeSnapshotRestorerTest {
+
+    private lateinit var underTest: DarkModeSnapshotRestorer
+    private lateinit var testScope: TestScope
+
+    private var isActive = false
+
+    @Before
+    fun setUp() {
+        val testDispatcher = StandardTestDispatcher()
+        testScope = TestScope(testDispatcher)
+        underTest =
+            DarkModeSnapshotRestorer(
+                backgroundDispatcher = testDispatcher,
+                isActive = { isActive },
+                setActive = { isActive = it },
+            )
+    }
+
+    @Test
+    fun `set up and restore - active`() =
+        testScope.runTest {
+            isActive = true
+
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val storedSnapshot = store.retrieve()
+
+            underTest.restoreToSnapshot(snapshot = storedSnapshot)
+            assertThat(isActive).isTrue()
+        }
+
+    @Test
+    fun `set up and restore - inactive`() =
+        testScope.runTest {
+            isActive = false
+
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val storedSnapshot = store.retrieve()
+
+            underTest.restoreToSnapshot(snapshot = storedSnapshot)
+            assertThat(isActive).isFalse()
+        }
+
+    @Test
+    fun `set up - deactivate - restore to active`() =
+        testScope.runTest {
+            isActive = true
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val initialSnapshot = store.retrieve()
+
+            underTest.store(isActivated = false)
+
+            underTest.restoreToSnapshot(snapshot = initialSnapshot)
+            assertThat(isActive).isTrue()
+        }
+
+    @Test
+    fun `set up - activate - restore to inactive`() =
+        testScope.runTest {
+            isActive = false
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val initialSnapshot = store.retrieve()
+
+            underTest.store(isActivated = true)
+
+            underTest.restoreToSnapshot(snapshot = initialSnapshot)
+            assertThat(isActive).isFalse()
+        }
+}
diff --git a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt b/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
new file mode 100644
index 00000000..cbf1365a
--- /dev/null
+++ b/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
@@ -0,0 +1,118 @@
+/*
+ * 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.model.picker.color.domain.interactor
+
+import android.content.Context
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.customization.picker.color.data.repository.FakeColorPickerRepository
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
+import com.android.customization.picker.color.domain.interactor.ColorPickerSnapshotRestorer
+import com.android.customization.picker.color.shared.model.ColorType
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ColorPickerInteractorTest {
+    private lateinit var underTest: ColorPickerInteractor
+    private lateinit var repository: FakeColorPickerRepository
+    private lateinit var store: FakeSnapshotStore
+
+    private lateinit var context: Context
+
+    @Before
+    fun setUp() {
+        context = InstrumentationRegistry.getInstrumentation().targetContext
+        repository = FakeColorPickerRepository(context = context)
+        store = FakeSnapshotStore()
+        underTest =
+            ColorPickerInteractor(
+                repository = repository,
+                snapshotRestorer = {
+                    ColorPickerSnapshotRestorer(interactor = underTest).apply {
+                        runBlocking { setUpSnapshotRestorer(store = store) }
+                    }
+                },
+            )
+        repository.setOptions(4, 4, ColorType.WALLPAPER_COLOR, 0)
+    }
+
+    @Test
+    fun select() = runTest {
+        val colorOptions = collectLastValue(underTest.colorOptions)
+
+        val wallpaperColorOptionModelBefore = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(2)
+        assertThat(wallpaperColorOptionModelBefore?.isSelected).isFalse()
+
+        wallpaperColorOptionModelBefore?.let { underTest.select(colorOptionModel = it) }
+        val wallpaperColorOptionModelAfter = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(2)
+        assertThat(wallpaperColorOptionModelAfter?.isSelected).isTrue()
+
+        val presetColorOptionModelBefore = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(1)
+        assertThat(presetColorOptionModelBefore?.isSelected).isFalse()
+
+        presetColorOptionModelBefore?.let { underTest.select(colorOptionModel = it) }
+        val presetColorOptionModelAfter = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(1)
+        assertThat(presetColorOptionModelAfter?.isSelected).isTrue()
+    }
+
+    @Test
+    fun snapshotRestorer_updatesSnapshot() = runTest {
+        val colorOptions = collectLastValue(underTest.colorOptions)
+        val wallpaperColorOptionModel0 = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(0)
+        val wallpaperColorOptionModel1 = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(1)
+        assertThat(wallpaperColorOptionModel0?.isSelected).isTrue()
+        assertThat(wallpaperColorOptionModel1?.isSelected).isFalse()
+
+        val storedSnapshot = store.retrieve()
+        wallpaperColorOptionModel1?.let { underTest.select(it) }
+        val wallpaperColorOptionModel0After = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(0)
+        val wallpaperColorOptionModel1After = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(1)
+        assertThat(wallpaperColorOptionModel0After?.isSelected).isFalse()
+        assertThat(wallpaperColorOptionModel1After?.isSelected).isTrue()
+
+        assertThat(store.retrieve()).isNotEqualTo(storedSnapshot)
+    }
+
+    @Test
+    fun snapshotRestorer_doesNotUpdateSnapshotOnExternalUpdates() = runTest {
+        val colorOptions = collectLastValue(underTest.colorOptions)
+        val wallpaperColorOptionModel0 = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(0)
+        val wallpaperColorOptionModel1 = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(1)
+        assertThat(wallpaperColorOptionModel0?.isSelected).isTrue()
+        assertThat(wallpaperColorOptionModel1?.isSelected).isFalse()
+
+        val storedSnapshot = store.retrieve()
+        repository.setOptions(4, 4, ColorType.WALLPAPER_COLOR, 1)
+        val wallpaperColorOptionModel0After = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(0)
+        val wallpaperColorOptionModel1After = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(1)
+        assertThat(wallpaperColorOptionModel0After?.isSelected).isFalse()
+        assertThat(wallpaperColorOptionModel1After?.isSelected).isTrue()
+
+        assertThat(store.retrieve()).isEqualTo(storedSnapshot)
+    }
+}
diff --git a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt b/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
new file mode 100644
index 00000000..71a8f235
--- /dev/null
+++ b/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
@@ -0,0 +1,138 @@
+/*
+ * 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.model.picker.color.domain.interactor
+
+import android.content.Context
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.customization.picker.color.data.repository.FakeColorPickerRepository
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
+import com.android.customization.picker.color.domain.interactor.ColorPickerSnapshotRestorer
+import com.android.customization.picker.color.shared.model.ColorOptionModel
+import com.android.customization.picker.color.shared.model.ColorType
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ColorPickerSnapshotRestorerTest {
+
+    private lateinit var underTest: ColorPickerSnapshotRestorer
+    private lateinit var repository: FakeColorPickerRepository
+    private lateinit var store: FakeSnapshotStore
+
+    private lateinit var context: Context
+
+    @Before
+    fun setUp() {
+        context = InstrumentationRegistry.getInstrumentation().targetContext
+        repository = FakeColorPickerRepository(context = context)
+        underTest =
+            ColorPickerSnapshotRestorer(
+                interactor =
+                    ColorPickerInteractor(
+                        repository = repository,
+                        snapshotRestorer = { underTest },
+                    )
+            )
+        store = FakeSnapshotStore()
+    }
+
+    @Test
+    fun restoreToSnapshot_noCallsToStore_restoresToInitialSnapshot() = runTest {
+        val colorOptions = collectLastValue(repository.colorOptions)
+
+        repository.setOptions(4, 4, ColorType.WALLPAPER_COLOR, 2)
+        val initialSnapshot = underTest.setUpSnapshotRestorer(store = store)
+        assertThat(initialSnapshot.args).isNotEmpty()
+
+        val colorOptionToSelect = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(3)
+        colorOptionToSelect?.let { repository.select(it) }
+        assertState(colorOptions(), ColorType.PRESET_COLOR, 3)
+
+        underTest.restoreToSnapshot(initialSnapshot)
+        assertState(colorOptions(), ColorType.WALLPAPER_COLOR, 2)
+    }
+
+    @Test
+    fun restoreToSnapshot_withCallToStore_restoresToInitialSnapshot() = runTest {
+        val colorOptions = collectLastValue(repository.colorOptions)
+
+        repository.setOptions(4, 4, ColorType.WALLPAPER_COLOR, 2)
+        val initialSnapshot = underTest.setUpSnapshotRestorer(store = store)
+        assertThat(initialSnapshot.args).isNotEmpty()
+
+        val colorOptionToSelect = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(3)
+        colorOptionToSelect?.let { repository.select(it) }
+        assertState(colorOptions(), ColorType.PRESET_COLOR, 3)
+
+        val colorOptionToStore = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(1)
+        colorOptionToStore?.let { underTest.storeSnapshot(colorOptionToStore) }
+
+        underTest.restoreToSnapshot(initialSnapshot)
+        assertState(colorOptions(), ColorType.WALLPAPER_COLOR, 2)
+    }
+
+    private fun assertState(
+        colorOptions: Map<ColorType, List<ColorOptionModel>>?,
+        selectedColorType: ColorType,
+        selectedColorIndex: Int
+    ) {
+        var foundSelectedColorOption = false
+        assertThat(colorOptions).isNotNull()
+        val optionsOfSelectedColorType = colorOptions?.get(selectedColorType)
+        assertThat(optionsOfSelectedColorType).isNotNull()
+        if (optionsOfSelectedColorType != null) {
+            for (i in optionsOfSelectedColorType.indices) {
+                val colorOptionHasSelectedIndex = i == selectedColorIndex
+                Truth.assertWithMessage(
+                        "Expected color option with index \"${i}\" to have" +
+                            " isSelected=$colorOptionHasSelectedIndex but it was" +
+                            " ${optionsOfSelectedColorType[i].isSelected}, num options: ${colorOptions.size}"
+                    )
+                    .that(optionsOfSelectedColorType[i].isSelected)
+                    .isEqualTo(colorOptionHasSelectedIndex)
+                foundSelectedColorOption = foundSelectedColorOption || colorOptionHasSelectedIndex
+            }
+            if (selectedColorIndex == -1) {
+                Truth.assertWithMessage(
+                        "Expected no color options to be selected, but a color option is" +
+                            " selected"
+                    )
+                    .that(foundSelectedColorOption)
+                    .isFalse()
+            } else {
+                Truth.assertWithMessage(
+                        "Expected a color option to be selected, but no color option is" +
+                            " selected"
+                    )
+                    .that(foundSelectedColorOption)
+                    .isTrue()
+            }
+        }
+    }
+}
diff --git a/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt b/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
new file mode 100644
index 00000000..1d9457a2
--- /dev/null
+++ b/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
@@ -0,0 +1,262 @@
+/*
+ * 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.model.picker.color.ui.viewmodel
+
+import android.content.Context
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.customization.picker.color.data.repository.FakeColorPickerRepository
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
+import com.android.customization.picker.color.domain.interactor.ColorPickerSnapshotRestorer
+import com.android.customization.picker.color.shared.model.ColorType
+import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
+import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
+import com.android.customization.picker.color.ui.viewmodel.ColorTypeTabViewModel
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import com.google.common.truth.Truth.assertWithMessage
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.resetMain
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.test.setMain
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ColorPickerViewModelTest {
+    private lateinit var underTest: ColorPickerViewModel
+    private lateinit var repository: FakeColorPickerRepository
+    private lateinit var interactor: ColorPickerInteractor
+    private lateinit var store: FakeSnapshotStore
+
+    private lateinit var context: Context
+    private lateinit var testScope: TestScope
+
+    @Before
+    fun setUp() {
+        context = InstrumentationRegistry.getInstrumentation().targetContext
+        val testDispatcher = StandardTestDispatcher()
+        testScope = TestScope(testDispatcher)
+        Dispatchers.setMain(testDispatcher)
+        repository = FakeColorPickerRepository(context = context)
+        store = FakeSnapshotStore()
+
+        interactor =
+            ColorPickerInteractor(
+                repository = repository,
+                snapshotRestorer = {
+                    ColorPickerSnapshotRestorer(interactor = interactor).apply {
+                        runBlocking { setUpSnapshotRestorer(store = store) }
+                    }
+                },
+            )
+
+        underTest =
+            ColorPickerViewModel.Factory(context = context, interactor = interactor)
+                .create(ColorPickerViewModel::class.java)
+
+        repository.setOptions(4, 4, ColorType.WALLPAPER_COLOR, 0)
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun `Select a color section color`() =
+        testScope.runTest {
+            val colorSectionOptions = collectLastValue(underTest.colorSectionOptions)
+
+            assertColorOptionUiState(
+                colorOptions = colorSectionOptions(),
+                selectedColorOptionIndex = 0
+            )
+
+            selectColorOption(colorSectionOptions, 2)
+            assertColorOptionUiState(
+                colorOptions = colorSectionOptions(),
+                selectedColorOptionIndex = 2
+            )
+
+            selectColorOption(colorSectionOptions, 4)
+            assertColorOptionUiState(
+                colorOptions = colorSectionOptions(),
+                selectedColorOptionIndex = 4
+            )
+        }
+
+    @Test
+    fun `Select a preset color`() =
+        testScope.runTest {
+            val colorTypes = collectLastValue(underTest.colorTypeTabs)
+            val colorOptions = collectLastValue(underTest.colorOptions)
+
+            // Initially, the wallpaper color tab should be selected
+            assertPickerUiState(
+                colorTypes = colorTypes(),
+                colorOptions = colorOptions(),
+                selectedColorTypeText = "Wallpaper colors",
+                selectedColorOptionIndex = 0
+            )
+
+            // Select "Basic colors" tab
+            colorTypes()?.get(ColorType.PRESET_COLOR)?.onClick?.invoke()
+            assertPickerUiState(
+                colorTypes = colorTypes(),
+                colorOptions = colorOptions(),
+                selectedColorTypeText = "Basic colors",
+                selectedColorOptionIndex = -1
+            )
+
+            // Select a color option
+            selectColorOption(colorOptions, 2)
+
+            // Check original option is no longer selected
+            colorTypes()?.get(ColorType.WALLPAPER_COLOR)?.onClick?.invoke()
+            assertPickerUiState(
+                colorTypes = colorTypes(),
+                colorOptions = colorOptions(),
+                selectedColorTypeText = "Wallpaper colors",
+                selectedColorOptionIndex = -1
+            )
+
+            // Check new option is selected
+            colorTypes()?.get(ColorType.PRESET_COLOR)?.onClick?.invoke()
+            assertPickerUiState(
+                colorTypes = colorTypes(),
+                colorOptions = colorOptions(),
+                selectedColorTypeText = "Basic colors",
+                selectedColorOptionIndex = 2
+            )
+        }
+
+    /** Simulates a user selecting the affordance at the given index, if that is clickable. */
+    private fun TestScope.selectColorOption(
+        colorOptions: () -> List<OptionItemViewModel<ColorOptionIconViewModel>>?,
+        index: Int,
+    ) {
+        val onClickedFlow = colorOptions()?.get(index)?.onClicked
+        val onClickedLastValueOrNull: (() -> (() -> Unit)?)? =
+            onClickedFlow?.let { collectLastValue(it) }
+        onClickedLastValueOrNull?.let { onClickedLastValue ->
+            val onClickedOrNull: (() -> Unit)? = onClickedLastValue()
+            onClickedOrNull?.let { onClicked -> onClicked() }
+        }
+    }
+
+    /**
+     * Asserts the entire picker UI state is what is expected. This includes the color type tabs and
+     * the color options list.
+     *
+     * @param colorTypes The observed color type view-models, keyed by ColorType
+     * @param colorOptions The observed color options
+     * @param selectedColorTypeText The text of the color type that's expected to be selected
+     * @param selectedColorOptionIndex The index of the color option that's expected to be selected,
+     *   -1 stands for no color option should be selected
+     */
+    private fun TestScope.assertPickerUiState(
+        colorTypes: Map<ColorType, ColorTypeTabViewModel>?,
+        colorOptions: List<OptionItemViewModel<ColorOptionIconViewModel>>?,
+        selectedColorTypeText: String,
+        selectedColorOptionIndex: Int,
+    ) {
+        assertColorTypeTabUiState(
+            colorTypes = colorTypes,
+            colorTypeId = ColorType.WALLPAPER_COLOR,
+            isSelected = "Wallpaper colors" == selectedColorTypeText,
+        )
+        assertColorTypeTabUiState(
+            colorTypes = colorTypes,
+            colorTypeId = ColorType.PRESET_COLOR,
+            isSelected = "Basic colors" == selectedColorTypeText,
+        )
+        assertColorOptionUiState(colorOptions, selectedColorOptionIndex)
+    }
+
+    /**
+     * Asserts the picker section UI state is what is expected.
+     *
+     * @param colorOptions The observed color options
+     * @param selectedColorOptionIndex The index of the color option that's expected to be selected,
+     *   -1 stands for no color option should be selected
+     */
+    private fun TestScope.assertColorOptionUiState(
+        colorOptions: List<OptionItemViewModel<ColorOptionIconViewModel>>?,
+        selectedColorOptionIndex: Int,
+    ) {
+        var foundSelectedColorOption = false
+        assertThat(colorOptions).isNotNull()
+        if (colorOptions != null) {
+            for (i in colorOptions.indices) {
+                val colorOptionHasSelectedIndex = i == selectedColorOptionIndex
+                val isSelected: Boolean? = collectLastValue(colorOptions[i].isSelected).invoke()
+                assertWithMessage(
+                        "Expected color option with index \"${i}\" to have" +
+                            " isSelected=$colorOptionHasSelectedIndex but it was" +
+                            " ${isSelected}, num options: ${colorOptions.size}"
+                    )
+                    .that(isSelected)
+                    .isEqualTo(colorOptionHasSelectedIndex)
+                foundSelectedColorOption = foundSelectedColorOption || colorOptionHasSelectedIndex
+            }
+            if (selectedColorOptionIndex == -1) {
+                assertWithMessage(
+                        "Expected no color options to be selected, but a color option is" +
+                            " selected"
+                    )
+                    .that(foundSelectedColorOption)
+                    .isFalse()
+            } else {
+                assertWithMessage(
+                        "Expected a color option to be selected, but no color option is" +
+                            " selected"
+                    )
+                    .that(foundSelectedColorOption)
+                    .isTrue()
+            }
+        }
+    }
+
+    /**
+     * Asserts that a color type tab has the correct UI state.
+     *
+     * @param colorTypes The observed color type view-models, keyed by ColorType enum
+     * @param colorTypeId the ID of the color type to assert
+     * @param isSelected Whether that color type should be selected
+     */
+    private fun assertColorTypeTabUiState(
+        colorTypes: Map<ColorType, ColorTypeTabViewModel>?,
+        colorTypeId: ColorType,
+        isSelected: Boolean,
+    ) {
+        val viewModel =
+            colorTypes?.get(colorTypeId) ?: error("No color type with ID \"$colorTypeId\"!")
+        assertThat(viewModel.isSelected).isEqualTo(isSelected)
+    }
+}
diff --git a/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt b/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
index 9a2a0af6..fea94dc8 100644
--- a/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
+++ b/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
@@ -24,6 +24,7 @@ import com.android.customization.picker.quickaffordance.domain.interactor.Keygua
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSelectionModel
 import com.android.systemui.shared.customization.data.content.FakeCustomizationProviderClient
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
+import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.Dispatchers
@@ -69,7 +70,7 @@ class KeyguardQuickAffordancePickerInteractorTest {
                             interactor = underTest,
                             client = client,
                         )
-                        .apply { runBlocking { setUpSnapshotRestorer {} } }
+                        .apply { runBlocking { setUpSnapshotRestorer(FakeSnapshotStore()) } }
                 },
             )
     }
@@ -114,23 +115,6 @@ class KeyguardQuickAffordancePickerInteractorTest {
         }
 
     @Test
-    fun unselect() =
-        testScope.runTest {
-            val selections = collectLastValue(underTest.selections)
-            underTest.select(
-                slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
-                affordanceId = FakeCustomizationProviderClient.AFFORDANCE_1,
-            )
-
-            underTest.unselect(
-                slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
-                affordanceId = FakeCustomizationProviderClient.AFFORDANCE_1,
-            )
-
-            assertThat(selections()).isEmpty()
-        }
-
-    @Test
     fun unselectAll() =
         testScope.runTest {
             client.setSlotCapacity(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END, 3)
diff --git a/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt b/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
index d1214c1a..103ae84d 100644
--- a/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
+++ b/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
@@ -18,6 +18,7 @@
 package com.android.customization.model.picker.quickaffordance.ui.viewmodel
 
 import android.content.Context
+import android.content.Intent
 import androidx.test.filters.SmallTest
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
@@ -26,14 +27,15 @@ import com.android.customization.picker.quickaffordance.domain.interactor.Keygua
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordancePickerViewModel
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordanceSlotViewModel
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordanceSummaryViewModel
-import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordanceViewModel
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
 import com.android.systemui.shared.customization.data.content.FakeCustomizationProviderClient
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
+import com.android.wallpaper.R
 import com.android.wallpaper.module.InjectorProvider
-import com.android.wallpaper.picker.undo.data.repository.UndoRepository
-import com.android.wallpaper.picker.undo.domain.interactor.UndoInteractor
-import com.android.wallpaper.testing.FAKE_RESTORERS
+import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
+import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.TestCurrentWallpaperInfoFactory
 import com.android.wallpaper.testing.TestInjector
 import com.android.wallpaper.testing.collectLastValue
@@ -65,6 +67,8 @@ class KeyguardQuickAffordancePickerViewModelTest {
     private lateinit var client: FakeCustomizationProviderClient
     private lateinit var quickAffordanceInteractor: KeyguardQuickAffordancePickerInteractor
 
+    private var latestStartedActivityIntent: Intent? = null
+
     @Before
     fun setUp() {
         InjectorProvider.setInjector(TestInjector())
@@ -87,21 +91,15 @@ class KeyguardQuickAffordancePickerViewModelTest {
                             interactor = quickAffordanceInteractor,
                             client = client,
                         )
-                        .apply { runBlocking { setUpSnapshotRestorer {} } }
+                        .apply { runBlocking { setUpSnapshotRestorer(FakeSnapshotStore()) } }
                 },
             )
-        val undoInteractor =
-            UndoInteractor(
-                scope = testScope.backgroundScope,
-                repository = UndoRepository(),
-                restorerByOwnerId = FAKE_RESTORERS,
-            )
         underTest =
             KeyguardQuickAffordancePickerViewModel.Factory(
                     context = context,
                     quickAffordanceInteractor = quickAffordanceInteractor,
-                    undoInteractor = undoInteractor,
                     wallpaperInfoFactory = TestCurrentWallpaperInfoFactory(context),
+                    activityStarter = { intent -> latestStartedActivityIntent = intent },
                 )
                 .create(KeyguardQuickAffordancePickerViewModel::class.java)
     }
@@ -134,7 +132,7 @@ class KeyguardQuickAffordancePickerViewModelTest {
             )
 
             // Select "affordance 1" for the first slot.
-            quickAffordances()?.get(1)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 1)
             assertPickerUiState(
                 slots = slots(),
                 affordances = quickAffordances(),
@@ -155,7 +153,7 @@ class KeyguardQuickAffordancePickerViewModelTest {
             // First, switch to the second slot:
             slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END)?.onClicked?.invoke()
             // Second, select the "affordance 3" affordance:
-            quickAffordances()?.get(3)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 3)
             assertPickerUiState(
                 slots = slots(),
                 affordances = quickAffordances(),
@@ -174,7 +172,7 @@ class KeyguardQuickAffordancePickerViewModelTest {
             )
 
             // Select a different affordance for the second slot.
-            quickAffordances()?.get(2)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 2)
             assertPickerUiState(
                 slots = slots(),
                 affordances = quickAffordances(),
@@ -200,17 +198,17 @@ class KeyguardQuickAffordancePickerViewModelTest {
             val quickAffordances = collectLastValue(underTest.quickAffordances)
 
             // Select "affordance 1" for the first slot.
-            quickAffordances()?.get(1)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 1)
             // Select an affordance for the second slot.
             // First, switch to the second slot:
             slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END)?.onClicked?.invoke()
             // Second, select the "affordance 3" affordance:
-            quickAffordances()?.get(3)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 3)
 
             // Switch back to the first slot:
             slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START)?.onClicked?.invoke()
             // Select the "none" affordance, which is always in position 0:
-            quickAffordances()?.get(0)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 0)
 
             assertPickerUiState(
                 slots = slots(),
@@ -256,14 +254,17 @@ class KeyguardQuickAffordancePickerViewModelTest {
                 )
 
             // Lets try to select that disabled affordance:
-            quickAffordances()?.get(affordanceIndex + 1)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, affordanceIndex + 1)
 
             // We expect there to be a dialog that should be shown:
-            assertThat(dialog()?.icon).isEqualTo(FakeCustomizationProviderClient.ICON_1)
-            assertThat(dialog()?.instructions).isEqualTo(enablementInstructions)
-            assertThat(dialog()?.actionText).isEqualTo(enablementActionText)
-            assertThat(dialog()?.intent?.`package`).isEqualTo(packageName)
-            assertThat(dialog()?.intent?.action).isEqualTo(action)
+            assertThat(dialog()?.icon)
+                .isEqualTo(Icon.Loaded(FakeCustomizationProviderClient.ICON_1, null))
+            assertThat(dialog()?.title).isEqualTo(Text.Loaded("disabled"))
+            assertThat(dialog()?.message)
+                .isEqualTo(Text.Loaded(enablementInstructions.joinToString("\n")))
+            assertThat(dialog()?.buttons?.size).isEqualTo(1)
+            assertThat(dialog()?.buttons?.first()?.text)
+                .isEqualTo(Text.Loaded(enablementActionText))
 
             // Once we report that the dialog has been dismissed by the user, we expect there to be
             // no
@@ -273,6 +274,30 @@ class KeyguardQuickAffordancePickerViewModelTest {
         }
 
     @Test
+    fun `Start settings activity when long-pressing an affordance`() =
+        testScope.runTest {
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+
+            // Lets add a configurable affordance to the picker:
+            val configureIntent = Intent("some.action")
+            val affordanceIndex =
+                client.addAffordance(
+                    CustomizationProviderClient.Affordance(
+                        id = "affordance",
+                        name = "affordance",
+                        iconResourceId = 1,
+                        isEnabled = true,
+                        configureIntent = configureIntent,
+                    )
+                )
+
+            // Lets try to long-click the affordance:
+            quickAffordances()?.get(affordanceIndex + 1)?.onLongClicked?.invoke()
+
+            assertThat(latestStartedActivityIntent).isEqualTo(configureIntent)
+        }
+
+    @Test
     fun `summary - affordance selected in both bottom-start and bottom-end`() =
         testScope.runTest {
             val slots = collectLastValue(underTest.slots)
@@ -280,21 +305,23 @@ class KeyguardQuickAffordancePickerViewModelTest {
             val summary = collectLastValue(underTest.summary)
 
             // Select "affordance 1" for the first slot.
-            quickAffordances()?.get(1)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 1)
             // Select an affordance for the second slot.
             // First, switch to the second slot:
             slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END)?.onClicked?.invoke()
             // Second, select the "affordance 3" affordance:
-            quickAffordances()?.get(3)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 3)
 
             assertThat(summary())
                 .isEqualTo(
                     KeyguardQuickAffordanceSummaryViewModel(
                         description =
-                            "${FakeCustomizationProviderClient.AFFORDANCE_1}," +
-                                " ${FakeCustomizationProviderClient.AFFORDANCE_3}",
-                        icon1 = FakeCustomizationProviderClient.ICON_1,
-                        icon2 = FakeCustomizationProviderClient.ICON_3,
+                            Text.Loaded(
+                                "${FakeCustomizationProviderClient.AFFORDANCE_1}," +
+                                    " ${FakeCustomizationProviderClient.AFFORDANCE_3}"
+                            ),
+                        icon1 = Icon.Loaded(FakeCustomizationProviderClient.ICON_1, null),
+                        icon2 = Icon.Loaded(FakeCustomizationProviderClient.ICON_3, null),
                     )
                 )
         }
@@ -307,13 +334,13 @@ class KeyguardQuickAffordancePickerViewModelTest {
             val summary = collectLastValue(underTest.summary)
 
             // Select "affordance 1" for the first slot.
-            quickAffordances()?.get(1)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 1)
 
             assertThat(summary())
                 .isEqualTo(
                     KeyguardQuickAffordanceSummaryViewModel(
-                        description = FakeCustomizationProviderClient.AFFORDANCE_1,
-                        icon1 = FakeCustomizationProviderClient.ICON_1,
+                        description = Text.Loaded(FakeCustomizationProviderClient.AFFORDANCE_1),
+                        icon1 = Icon.Loaded(FakeCustomizationProviderClient.ICON_1, null),
                         icon2 = null,
                     )
                 )
@@ -330,14 +357,14 @@ class KeyguardQuickAffordancePickerViewModelTest {
             // First, switch to the second slot:
             slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END)?.onClicked?.invoke()
             // Second, select the "affordance 3" affordance:
-            quickAffordances()?.get(3)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 3)
 
             assertThat(summary())
                 .isEqualTo(
                     KeyguardQuickAffordanceSummaryViewModel(
-                        description = FakeCustomizationProviderClient.AFFORDANCE_3,
+                        description = Text.Loaded(FakeCustomizationProviderClient.AFFORDANCE_3),
                         icon1 = null,
-                        icon2 = FakeCustomizationProviderClient.ICON_3,
+                        icon2 = Icon.Loaded(FakeCustomizationProviderClient.ICON_3, null),
                     )
                 )
         }
@@ -345,15 +372,28 @@ class KeyguardQuickAffordancePickerViewModelTest {
     @Test
     fun `summary - no affordances selected`() =
         testScope.runTest {
-            val slots = collectLastValue(underTest.slots)
-            val quickAffordances = collectLastValue(underTest.quickAffordances)
             val summary = collectLastValue(underTest.summary)
 
-            assertThat(summary()?.description).isEqualTo("None")
+            assertThat(summary()?.description)
+                .isEqualTo(Text.Resource(R.string.keyguard_quick_affordance_none_selected))
             assertThat(summary()?.icon1).isNotNull()
             assertThat(summary()?.icon2).isNull()
         }
 
+    /** Simulates a user selecting the affordance at the given index, if that is clickable. */
+    private fun TestScope.selectAffordance(
+        affordances: () -> List<OptionItemViewModel<Icon>>?,
+        index: Int,
+    ) {
+        val onClickedFlow = affordances()?.get(index)?.onClicked
+        val onClickedLastValueOrNull: (() -> (() -> Unit)?)? =
+            onClickedFlow?.let { collectLastValue(it) }
+        onClickedLastValueOrNull?.let { onClickedLastValue ->
+            val onClickedOrNull: (() -> Unit)? = onClickedLastValue()
+            onClickedOrNull?.let { onClicked -> onClicked() }
+        }
+    }
+
     /**
      * Asserts the entire picker UI state is what is expected. This includes the slot tabs and the
      * affordance list.
@@ -363,9 +403,9 @@ class KeyguardQuickAffordancePickerViewModelTest {
      * @param selectedSlotText The text of the slot that's expected to be selected
      * @param selectedAffordanceText The text of the affordance that's expected to be selected
      */
-    private fun assertPickerUiState(
+    private fun TestScope.assertPickerUiState(
         slots: Map<String, KeyguardQuickAffordanceSlotViewModel>?,
-        affordances: List<KeyguardQuickAffordanceViewModel>?,
+        affordances: List<OptionItemViewModel<Icon>>?,
         selectedSlotText: String,
         selectedAffordanceText: String,
     ) {
@@ -383,12 +423,18 @@ class KeyguardQuickAffordancePickerViewModelTest {
         var foundSelectedAffordance = false
         assertThat(affordances).isNotNull()
         affordances?.forEach { affordance ->
-            val nameMatchesSelectedName = affordance.contentDescription == selectedAffordanceText
+            val nameMatchesSelectedName =
+                Text.evaluationEquals(
+                    context,
+                    affordance.text,
+                    Text.Loaded(selectedAffordanceText),
+                )
+            val isSelected: Boolean? = collectLastValue(affordance.isSelected).invoke()
             assertWithMessage(
-                    "Expected affordance with name \"${affordance.contentDescription}\" to have" +
-                        " isSelected=$nameMatchesSelectedName but it was ${affordance.isSelected}"
+                    "Expected affordance with name \"${affordance.text}\" to have" +
+                        " isSelected=$nameMatchesSelectedName but it was $isSelected"
                 )
-                .that(affordance.isSelected)
+                .that(isSelected)
                 .isEqualTo(nameMatchesSelectedName)
             foundSelectedAffordance = foundSelectedAffordance || nameMatchesSelectedName
         }
@@ -416,7 +462,8 @@ class KeyguardQuickAffordancePickerViewModelTest {
      *
      * @param slots The observed slot view-models, keyed by slot ID
      * @param expectedAffordanceNameBySlotId The expected name of the selected affordance for each
-     * slot ID or `null` if it's expected for there to be no affordance for that slot in the preview
+     *   slot ID or `null` if it's expected for there to be no affordance for that slot in the
+     *   preview
      */
     private fun assertPreviewUiState(
         slots: Map<String, KeyguardQuickAffordanceSlotViewModel>?,
@@ -425,13 +472,12 @@ class KeyguardQuickAffordancePickerViewModelTest {
         assertThat(slots).isNotNull()
         slots?.forEach { (slotId, slotViewModel) ->
             val expectedAffordanceName = expectedAffordanceNameBySlotId[slotId]
-            val actualAffordanceName =
-                slotViewModel.selectedQuickAffordances.firstOrNull()?.contentDescription
+            val actualAffordanceName = slotViewModel.selectedQuickAffordances.firstOrNull()?.text
             assertWithMessage(
                     "At slotId=\"$slotId\", expected affordance=\"$expectedAffordanceName\" but" +
-                        " was \"$actualAffordanceName\"!"
+                        " was \"${actualAffordanceName?.asString(context)}\"!"
                 )
-                .that(actualAffordanceName)
+                .that(actualAffordanceName?.asString(context))
                 .isEqualTo(expectedAffordanceName)
         }
     }
diff --git a/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt b/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt
new file mode 100644
index 00000000..e6e30c31
--- /dev/null
+++ b/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt
@@ -0,0 +1,56 @@
+/*
+ * 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.model.themedicon.domain.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.customization.model.themedicon.data.repository.ThemeIconRepository
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ThemedIconInteractorTest {
+
+    private lateinit var underTest: ThemedIconInteractor
+
+    @Before
+    fun setUp() {
+        underTest =
+            ThemedIconInteractor(
+                repository = ThemeIconRepository(),
+            )
+    }
+
+    @Test
+    fun `end-to-end`() = runTest {
+        val isActivated = collectLastValue(underTest.isActivated)
+
+        underTest.setActivated(isActivated = true)
+        assertThat(isActivated()).isTrue()
+
+        underTest.setActivated(isActivated = false)
+        assertThat(isActivated()).isFalse()
+    }
+}
diff --git a/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt b/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt
new file mode 100644
index 00000000..df1fd201
--- /dev/null
+++ b/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt
@@ -0,0 +1,102 @@
+/*
+ * 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.model.themedicon.domain.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.customization.model.themedicon.data.repository.ThemeIconRepository
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ThemedIconSnapshotRestorerTest {
+
+    private lateinit var underTest: ThemedIconSnapshotRestorer
+    private var isActivated = false
+
+    @Before
+    fun setUp() {
+        isActivated = false
+        underTest =
+            ThemedIconSnapshotRestorer(
+                isActivated = { isActivated },
+                setActivated = { isActivated = it },
+                interactor =
+                    ThemedIconInteractor(
+                        repository = ThemeIconRepository(),
+                    )
+            )
+    }
+
+    @Test
+    fun `set up and restore - active`() = runTest {
+        isActivated = true
+
+        val store = FakeSnapshotStore()
+        store.store(underTest.setUpSnapshotRestorer(store = store))
+        val storedSnapshot = store.retrieve()
+
+        underTest.restoreToSnapshot(snapshot = storedSnapshot)
+        assertThat(isActivated).isTrue()
+    }
+
+    @Test
+    fun `set up and restore - inactive`() = runTest {
+        isActivated = false
+
+        val store = FakeSnapshotStore()
+        store.store(underTest.setUpSnapshotRestorer(store = store))
+        val storedSnapshot = store.retrieve()
+
+        underTest.restoreToSnapshot(snapshot = storedSnapshot)
+        assertThat(isActivated).isFalse()
+    }
+
+    @Test
+    fun `set up - deactivate - restore to active`() = runTest {
+        isActivated = true
+        val store = FakeSnapshotStore()
+        store.store(underTest.setUpSnapshotRestorer(store = store))
+        val initialSnapshot = store.retrieve()
+
+        underTest.store(isActivated = false)
+
+        underTest.restoreToSnapshot(snapshot = initialSnapshot)
+        assertThat(isActivated).isTrue()
+    }
+
+    @Test
+    fun `set up - activate - restore to inactive`() = runTest {
+        isActivated = false
+        val store = FakeSnapshotStore()
+        store.store(underTest.setUpSnapshotRestorer(store = store))
+        val initialSnapshot = store.retrieve()
+
+        underTest.store(isActivated = true)
+
+        underTest.restoreToSnapshot(snapshot = initialSnapshot)
+        assertThat(isActivated).isFalse()
+    }
+}