diff --git a/tests/robotests/Android.bp b/tests/robotests/Android.bp
index 4416e1c..ed684f4 100644
--- a/tests/robotests/Android.bp
+++ b/tests/robotests/Android.bp
@@ -2,13 +2,23 @@
 package {
     default_applicable_licenses: ["Android-Apache-2.0"],
 }
+
 android_robolectric_test {
     name: "ThemePickerRoboTests",
     srcs: [
         "src/**/*.java",
         "src/**/*.kt",
     ],
+    // TODO(b/291104503) Enable this test
+    exclude_srcs: ["src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt"],
     java_resource_dirs: ["config"],
+    static_libs: [
+        "WallpaperPicker2TestLib",
+        "androidx.test.rules",
+        "junit",
+        "kotlinx_coroutines_test",
+        "truth-prebuilt",
+    ],
     libs: [
         "androidx.test.core",
         "androidx.test.runner",
diff --git a/tests/robotests/src/com/android/customization/model/color/ColorSectionControllerTest.java b/tests/robotests/src/com/android/customization/model/color/ColorSectionControllerTest.java
index 820e641..4e1a36a 100644
--- a/tests/robotests/src/com/android/customization/model/color/ColorSectionControllerTest.java
+++ b/tests/robotests/src/com/android/customization/model/color/ColorSectionControllerTest.java
@@ -22,6 +22,7 @@
 import com.android.wallpaper.model.WallpaperColorsViewModel;
 
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.MockitoAnnotations;
@@ -32,6 +33,7 @@
 /**
  * Tests of {@link ColorSectionController}.
  */
+@Ignore("b/290798811")
 @RunWith(RobolectricTestRunner.class)
 @Config(manifest = Config.NONE)
 public final class ColorSectionControllerTest {
diff --git a/tests/robotests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt b/tests/robotests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
new file mode 100644
index 0000000..5953937
--- /dev/null
+++ b/tests/robotests/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/robotests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt b/tests/robotests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt
new file mode 100644
index 0000000..f73d5a3
--- /dev/null
+++ b/tests/robotests/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/robotests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt b/tests/robotests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt
new file mode 100644
index 0000000..c2712b1
--- /dev/null
+++ b/tests/robotests/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/robotests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt b/tests/robotests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
new file mode 100644
index 0000000..58c5d99
--- /dev/null
+++ b/tests/robotests/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/robotests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt b/tests/robotests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
new file mode 100644
index 0000000..38067b7
--- /dev/null
+++ b/tests/robotests/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/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt b/tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
new file mode 100644
index 0000000..d4f24ee
--- /dev/null
+++ b/tests/robotests/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.robolectric.RobolectricTestRunner
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::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/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt b/tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
new file mode 100644
index 0000000..5f3e39e
--- /dev/null
+++ b/tests/robotests/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.robolectric.RobolectricTestRunner
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::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/robotests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt b/tests/robotests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
new file mode 100644
index 0000000..9968c5f
--- /dev/null
+++ b/tests/robotests/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.robolectric.RobolectricTestRunner
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::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()
+        Dispatchers.setMain(testDispatcher)
+        testScope = TestScope(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/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
new file mode 100644
index 0000000..35dbadd
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
@@ -0,0 +1,100 @@
+/*
+ * Copyright (C) 2022 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.quickaffordance.data.repository
+
+import androidx.test.filters.SmallTest
+import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
+import com.android.systemui.shared.customization.data.content.CustomizationProviderContract
+import com.android.systemui.shared.customization.data.content.FakeCustomizationProviderClient
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.toList
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+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.robolectric.RobolectricTestRunner
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::class)
+class KeyguardQuickAffordancePickerRepositoryTest {
+
+    private lateinit var underTest: KeyguardQuickAffordancePickerRepository
+
+    private lateinit var testScope: TestScope
+    private lateinit var client: FakeCustomizationProviderClient
+
+    @Before
+    fun setUp() {
+        client = FakeCustomizationProviderClient()
+        val coroutineDispatcher = UnconfinedTestDispatcher()
+        testScope = TestScope(coroutineDispatcher)
+        Dispatchers.setMain(coroutineDispatcher)
+
+        underTest =
+            KeyguardQuickAffordancePickerRepository(
+                client = client,
+                backgroundDispatcher = coroutineDispatcher,
+            )
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun `isFeatureEnabled - enabled`() =
+        testScope.runTest {
+            client.setFlag(
+                CustomizationProviderContract.FlagsTable
+                    .FLAG_NAME_CUSTOM_LOCK_SCREEN_QUICK_AFFORDANCES_ENABLED,
+                true,
+            )
+            val values = mutableListOf<Boolean>()
+            val job = launch { underTest.isFeatureEnabled.toList(values) }
+
+            assertThat(values.last()).isTrue()
+
+            job.cancel()
+        }
+
+    @Test
+    fun `isFeatureEnabled - not enabled`() =
+        testScope.runTest {
+            client.setFlag(
+                CustomizationProviderContract.FlagsTable
+                    .FLAG_NAME_CUSTOM_LOCK_SCREEN_QUICK_AFFORDANCES_ENABLED,
+                false,
+            )
+            val values = mutableListOf<Boolean>()
+            val job = launch { underTest.isFeatureEnabled.toList(values) }
+
+            assertThat(values.last()).isFalse()
+
+            job.cancel()
+        }
+}
diff --git a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
new file mode 100644
index 0000000..efe9f64
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
@@ -0,0 +1,141 @@
+/*
+ * Copyright (C) 2022 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.quickaffordance.domain.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
+import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
+import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordanceSnapshotRestorer
+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
+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.robolectric.RobolectricTestRunner
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::class)
+class KeyguardQuickAffordancePickerInteractorTest {
+
+    private lateinit var underTest: KeyguardQuickAffordancePickerInteractor
+
+    private lateinit var testScope: TestScope
+    private lateinit var client: FakeCustomizationProviderClient
+
+    @Before
+    fun setUp() {
+        val testDispatcher = StandardTestDispatcher()
+        testScope = TestScope(testDispatcher)
+        Dispatchers.setMain(testDispatcher)
+        client = FakeCustomizationProviderClient()
+        underTest =
+            KeyguardQuickAffordancePickerInteractor(
+                repository =
+                    KeyguardQuickAffordancePickerRepository(
+                        client = client,
+                        backgroundDispatcher = testDispatcher,
+                    ),
+                client = client,
+                snapshotRestorer = {
+                    KeyguardQuickAffordanceSnapshotRestorer(
+                            interactor = underTest,
+                            client = client,
+                        )
+                        .apply { runBlocking { setUpSnapshotRestorer(FakeSnapshotStore()) } }
+                },
+            )
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun select() =
+        testScope.runTest {
+            val selections = collectLastValue(underTest.selections)
+
+            underTest.select(
+                slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
+                affordanceId = FakeCustomizationProviderClient.AFFORDANCE_1,
+            )
+            assertThat(selections())
+                .isEqualTo(
+                    listOf(
+                        KeyguardQuickAffordancePickerSelectionModel(
+                            slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
+                            affordanceId = FakeCustomizationProviderClient.AFFORDANCE_1,
+                        ),
+                    )
+                )
+
+            underTest.select(
+                slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
+                affordanceId = FakeCustomizationProviderClient.AFFORDANCE_2,
+            )
+            assertThat(selections())
+                .isEqualTo(
+                    listOf(
+                        KeyguardQuickAffordancePickerSelectionModel(
+                            slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
+                            affordanceId = FakeCustomizationProviderClient.AFFORDANCE_2,
+                        ),
+                    )
+                )
+        }
+
+    @Test
+    fun unselectAll() =
+        testScope.runTest {
+            client.setSlotCapacity(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END, 3)
+            val selections = collectLastValue(underTest.selections)
+            underTest.select(
+                slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END,
+                affordanceId = FakeCustomizationProviderClient.AFFORDANCE_1,
+            )
+            underTest.select(
+                slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END,
+                affordanceId = FakeCustomizationProviderClient.AFFORDANCE_2,
+            )
+            underTest.select(
+                slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END,
+                affordanceId = FakeCustomizationProviderClient.AFFORDANCE_3,
+            )
+
+            underTest.unselectAll(
+                slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END,
+            )
+
+            assertThat(selections()).isEmpty()
+        }
+}
diff --git a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
new file mode 100644
index 0000000..3f10674
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
@@ -0,0 +1,513 @@
+/*
+ * Copyright (C) 2022 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.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
+import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
+import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordanceSnapshotRestorer
+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.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.common.icon.ui.viewmodel.Icon
+import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
+import com.android.wallpaper.picker.customization.data.repository.WallpaperRepository
+import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.android.wallpaper.testing.FakeWallpaperClient
+import com.android.wallpaper.testing.TestCurrentWallpaperInfoFactory
+import com.android.wallpaper.testing.TestInjector
+import com.android.wallpaper.testing.TestWallpaperPreferences
+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 KeyguardQuickAffordancePickerViewModelTest {
+
+    private lateinit var underTest: KeyguardQuickAffordancePickerViewModel
+
+    private lateinit var context: Context
+    private lateinit var testScope: TestScope
+    private lateinit var client: FakeCustomizationProviderClient
+    private lateinit var quickAffordanceInteractor: KeyguardQuickAffordancePickerInteractor
+    private lateinit var wallpaperInteractor: WallpaperInteractor
+
+    @Before
+    fun setUp() {
+        InjectorProvider.setInjector(TestInjector())
+        context = InstrumentationRegistry.getInstrumentation().targetContext
+        val testDispatcher = StandardTestDispatcher()
+        testScope = TestScope(testDispatcher)
+        Dispatchers.setMain(testDispatcher)
+        client = FakeCustomizationProviderClient()
+
+        quickAffordanceInteractor =
+            KeyguardQuickAffordancePickerInteractor(
+                repository =
+                    KeyguardQuickAffordancePickerRepository(
+                        client = client,
+                        backgroundDispatcher = testDispatcher,
+                    ),
+                client = client,
+                snapshotRestorer = {
+                    KeyguardQuickAffordanceSnapshotRestorer(
+                            interactor = quickAffordanceInteractor,
+                            client = client,
+                        )
+                        .apply { runBlocking { setUpSnapshotRestorer(FakeSnapshotStore()) } }
+                },
+            )
+        wallpaperInteractor =
+            WallpaperInteractor(
+                repository =
+                    WallpaperRepository(
+                        scope = testScope.backgroundScope,
+                        client = FakeWallpaperClient(),
+                        wallpaperPreferences = TestWallpaperPreferences(),
+                        backgroundDispatcher = testDispatcher,
+                    ),
+            )
+        underTest =
+            KeyguardQuickAffordancePickerViewModel.Factory(
+                    context = context,
+                    quickAffordanceInteractor = quickAffordanceInteractor,
+                    wallpaperInteractor = wallpaperInteractor,
+                    wallpaperInfoFactory = TestCurrentWallpaperInfoFactory(context),
+                )
+                .create(KeyguardQuickAffordancePickerViewModel::class.java)
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun `Select an affordance for each side`() =
+        testScope.runTest {
+            val slots = collectLastValue(underTest.slots)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+
+            // Initially, the first slot is selected with the "none" affordance selected.
+            assertPickerUiState(
+                slots = slots(),
+                affordances = quickAffordances(),
+                selectedSlotText = "Left button",
+                selectedAffordanceText = "None",
+            )
+            assertPreviewUiState(
+                slots = slots(),
+                expectedAffordanceNameBySlotId =
+                    mapOf(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to null,
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to null,
+                    ),
+            )
+
+            // Select "affordance 1" for the first slot.
+            selectAffordance(quickAffordances, 1)
+            assertPickerUiState(
+                slots = slots(),
+                affordances = quickAffordances(),
+                selectedSlotText = "Left button",
+                selectedAffordanceText = FakeCustomizationProviderClient.AFFORDANCE_1,
+            )
+            assertPreviewUiState(
+                slots = slots(),
+                expectedAffordanceNameBySlotId =
+                    mapOf(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
+                            FakeCustomizationProviderClient.AFFORDANCE_1,
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to null,
+                    ),
+            )
+
+            // 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:
+            selectAffordance(quickAffordances, 3)
+            assertPickerUiState(
+                slots = slots(),
+                affordances = quickAffordances(),
+                selectedSlotText = "Right button",
+                selectedAffordanceText = FakeCustomizationProviderClient.AFFORDANCE_3,
+            )
+            assertPreviewUiState(
+                slots = slots(),
+                expectedAffordanceNameBySlotId =
+                    mapOf(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
+                            FakeCustomizationProviderClient.AFFORDANCE_1,
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to
+                            FakeCustomizationProviderClient.AFFORDANCE_3,
+                    ),
+            )
+
+            // Select a different affordance for the second slot.
+            selectAffordance(quickAffordances, 2)
+            assertPickerUiState(
+                slots = slots(),
+                affordances = quickAffordances(),
+                selectedSlotText = "Right button",
+                selectedAffordanceText = FakeCustomizationProviderClient.AFFORDANCE_2,
+            )
+            assertPreviewUiState(
+                slots = slots(),
+                expectedAffordanceNameBySlotId =
+                    mapOf(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
+                            FakeCustomizationProviderClient.AFFORDANCE_1,
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to
+                            FakeCustomizationProviderClient.AFFORDANCE_2,
+                    ),
+            )
+        }
+
+    @Test
+    fun `Unselect - AKA selecting the none affordance - on one side`() =
+        testScope.runTest {
+            val slots = collectLastValue(underTest.slots)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+
+            // Select "affordance 1" for the first slot.
+            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:
+            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:
+            selectAffordance(quickAffordances, 0)
+
+            assertPickerUiState(
+                slots = slots(),
+                affordances = quickAffordances(),
+                selectedSlotText = "Left button",
+                selectedAffordanceText = "None",
+            )
+            assertPreviewUiState(
+                slots = slots(),
+                expectedAffordanceNameBySlotId =
+                    mapOf(
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to null,
+                        KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to
+                            FakeCustomizationProviderClient.AFFORDANCE_3,
+                    ),
+            )
+        }
+
+    @Test
+    fun `Show enablement dialog when selecting a disabled affordance`() =
+        testScope.runTest {
+            val slots = collectLastValue(underTest.slots)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+            val dialog = collectLastValue(underTest.dialog)
+            val activityStartRequest = collectLastValue(underTest.activityStartRequests)
+
+            val enablementExplanation = "enablementExplanation"
+            val enablementActionText = "enablementActionText"
+            val packageName = "packageName"
+            val action = "action"
+            val enablementActionIntent = Intent(action).apply { `package` = packageName }
+            // Lets add a disabled affordance to the picker:
+            val affordanceIndex =
+                client.addAffordance(
+                    CustomizationProviderClient.Affordance(
+                        id = "disabled",
+                        name = "disabled",
+                        iconResourceId = 1,
+                        isEnabled = false,
+                        enablementExplanation = enablementExplanation,
+                        enablementActionText = enablementActionText,
+                        enablementActionIntent = enablementActionIntent,
+                    )
+                )
+
+            // Lets try to select that disabled affordance:
+            selectAffordance(quickAffordances, affordanceIndex + 1)
+
+            // We expect there to be a dialog that should be shown:
+            assertThat(dialog()?.icon)
+                .isEqualTo(Icon.Loaded(FakeCustomizationProviderClient.ICON_1, null))
+            assertThat(dialog()?.headline)
+                .isEqualTo(Text.Resource(R.string.keyguard_affordance_enablement_dialog_headline))
+            assertThat(dialog()?.message).isEqualTo(Text.Loaded(enablementExplanation))
+            assertThat(dialog()?.buttons?.size).isEqualTo(2)
+            assertThat(dialog()?.buttons?.first()?.text).isEqualTo(Text.Resource(R.string.cancel))
+            assertThat(dialog()?.buttons?.get(1)?.text).isEqualTo(Text.Loaded(enablementActionText))
+
+            // When the button is clicked, we expect an intent of the given enablement action
+            // component name to be emitted.
+            dialog()?.buttons?.get(1)?.onClicked?.invoke()
+            assertThat(activityStartRequest()?.`package`).isEqualTo(packageName)
+            assertThat(activityStartRequest()?.action).isEqualTo(action)
+
+            // Once we report that the activity was started, the activity start request should be
+            // nullified.
+            underTest.onActivityStarted()
+            assertThat(activityStartRequest()).isNull()
+
+            // Once we report that the dialog has been dismissed by the user, we expect there to be
+            // no dialog to be shown:
+            underTest.onDialogDismissed()
+            assertThat(dialog()).isNull()
+        }
+
+    @Test
+    fun `Start settings activity when long-pressing an affordance`() =
+        testScope.runTest {
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+            val activityStartRequest = collectLastValue(underTest.activityStartRequests)
+
+            // 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(activityStartRequest()).isEqualTo(configureIntent)
+            // Once we report that the activity was started, the activity start request should be
+            // nullified.
+            underTest.onActivityStarted()
+            assertThat(activityStartRequest()).isNull()
+        }
+
+    @Test
+    fun `summary - affordance selected in both bottom-start and bottom-end`() =
+        testScope.runTest {
+            val slots = collectLastValue(underTest.slots)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+            val summary = collectLastValue(underTest.summary)
+
+            // Select "affordance 1" for the first slot.
+            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:
+            selectAffordance(quickAffordances, 3)
+
+            assertThat(summary())
+                .isEqualTo(
+                    KeyguardQuickAffordanceSummaryViewModel(
+                        description =
+                            Text.Loaded(
+                                "${FakeCustomizationProviderClient.AFFORDANCE_1}," +
+                                    " ${FakeCustomizationProviderClient.AFFORDANCE_3}"
+                            ),
+                        icon1 = Icon.Loaded(FakeCustomizationProviderClient.ICON_1, null),
+                        icon2 = Icon.Loaded(FakeCustomizationProviderClient.ICON_3, null),
+                    )
+                )
+        }
+
+    @Test
+    fun `summary - affordance selected only on bottom-start`() =
+        testScope.runTest {
+            val slots = collectLastValue(underTest.slots)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+            val summary = collectLastValue(underTest.summary)
+
+            // Select "affordance 1" for the first slot.
+            selectAffordance(quickAffordances, 1)
+
+            assertThat(summary())
+                .isEqualTo(
+                    KeyguardQuickAffordanceSummaryViewModel(
+                        description = Text.Loaded(FakeCustomizationProviderClient.AFFORDANCE_1),
+                        icon1 = Icon.Loaded(FakeCustomizationProviderClient.ICON_1, null),
+                        icon2 = null,
+                    )
+                )
+        }
+
+    @Test
+    fun `summary - affordance selected only on bottom-end`() =
+        testScope.runTest {
+            val slots = collectLastValue(underTest.slots)
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+            val summary = collectLastValue(underTest.summary)
+
+            // 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:
+            selectAffordance(quickAffordances, 3)
+
+            assertThat(summary())
+                .isEqualTo(
+                    KeyguardQuickAffordanceSummaryViewModel(
+                        description = Text.Loaded(FakeCustomizationProviderClient.AFFORDANCE_3),
+                        icon1 = null,
+                        icon2 = Icon.Loaded(FakeCustomizationProviderClient.ICON_3, null),
+                    )
+                )
+        }
+
+    @Test
+    fun `summary - no affordances selected`() =
+        testScope.runTest {
+            val summary = collectLastValue(underTest.summary)
+
+            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.
+     *
+     * @param slots The observed slot view-models, keyed by slot ID
+     * @param affordances The observed affordances
+     * @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 TestScope.assertPickerUiState(
+        slots: Map<String, KeyguardQuickAffordanceSlotViewModel>?,
+        affordances: List<OptionItemViewModel<Icon>>?,
+        selectedSlotText: String,
+        selectedAffordanceText: String,
+    ) {
+        assertSlotTabUiState(
+            slots = slots,
+            slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
+            isSelected = "Left button" == selectedSlotText,
+        )
+        assertSlotTabUiState(
+            slots = slots,
+            slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END,
+            isSelected = "Right button" == selectedSlotText,
+        )
+
+        var foundSelectedAffordance = false
+        assertThat(affordances).isNotNull()
+        affordances?.forEach { affordance ->
+            val nameMatchesSelectedName =
+                Text.evaluationEquals(
+                    context,
+                    affordance.text,
+                    Text.Loaded(selectedAffordanceText),
+                )
+            val isSelected: Boolean? = collectLastValue(affordance.isSelected).invoke()
+            assertWithMessage(
+                    "Expected affordance with name \"${affordance.text}\" to have" +
+                        " isSelected=$nameMatchesSelectedName but it was $isSelected"
+                )
+                .that(isSelected)
+                .isEqualTo(nameMatchesSelectedName)
+            foundSelectedAffordance = foundSelectedAffordance || nameMatchesSelectedName
+        }
+        assertWithMessage("No affordance is selected!").that(foundSelectedAffordance).isTrue()
+    }
+
+    /**
+     * Asserts that a slot tab has the correct UI state.
+     *
+     * @param slots The observed slot view-models, keyed by slot ID
+     * @param slotId the ID of the slot to assert
+     * @param isSelected Whether that slot should be selected
+     */
+    private fun assertSlotTabUiState(
+        slots: Map<String, KeyguardQuickAffordanceSlotViewModel>?,
+        slotId: String,
+        isSelected: Boolean,
+    ) {
+        val viewModel = slots?.get(slotId) ?: error("No slot with ID \"$slotId\"!")
+        assertThat(viewModel.isSelected).isEqualTo(isSelected)
+    }
+
+    /**
+     * Asserts the UI state of the preview.
+     *
+     * @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
+     */
+    private fun assertPreviewUiState(
+        slots: Map<String, KeyguardQuickAffordanceSlotViewModel>?,
+        expectedAffordanceNameBySlotId: Map<String, String?>,
+    ) {
+        assertThat(slots).isNotNull()
+        slots?.forEach { (slotId, slotViewModel) ->
+            val expectedAffordanceName = expectedAffordanceNameBySlotId[slotId]
+            val actualAffordanceName = slotViewModel.selectedQuickAffordances.firstOrNull()?.text
+            assertWithMessage(
+                    "At slotId=\"$slotId\", expected affordance=\"$expectedAffordanceName\" but" +
+                        " was \"${actualAffordanceName?.asString(context)}\"!"
+                )
+                .that(actualAffordanceName?.asString(context))
+                .isEqualTo(expectedAffordanceName)
+        }
+    }
+}
diff --git a/tests/robotests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt b/tests/robotests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt
new file mode 100644
index 0000000..e6e30c3
--- /dev/null
+++ b/tests/robotests/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/robotests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt b/tests/robotests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt
new file mode 100644
index 0000000..df1fd20
--- /dev/null
+++ b/tests/robotests/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()
+    }
+}
diff --git a/tests/robotests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt b/tests/robotests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
new file mode 100644
index 0000000..95d7e35
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.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.data.repository
+
+import android.graphics.Color
+import androidx.annotation.ColorInt
+import androidx.annotation.IntRange
+import com.android.customization.picker.clock.data.repository.FakeClockPickerRepository.Companion.fakeClocks
+import com.android.customization.picker.clock.shared.ClockSize
+import com.android.customization.picker.clock.shared.model.ClockMetadataModel
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.combine
+
+/** By default [FakeClockPickerRepository] uses [fakeClocks]. */
+open class FakeClockPickerRepository(clocks: List<ClockMetadataModel> = fakeClocks) :
+    ClockPickerRepository {
+    override val allClocks: Flow<List<ClockMetadataModel>> = MutableStateFlow(clocks).asStateFlow()
+
+    private val selectedClockId = MutableStateFlow(fakeClocks[0].clockId)
+    @ColorInt private val selectedColorId = MutableStateFlow<String?>(null)
+    private val colorTone = MutableStateFlow(ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS)
+    @ColorInt private val seedColor = MutableStateFlow<Int?>(null)
+    override val selectedClock: Flow<ClockMetadataModel> =
+        combine(
+            selectedClockId,
+            selectedColorId,
+            colorTone,
+            seedColor,
+        ) { selectedClockId, selectedColor, colorTone, seedColor ->
+            val selectedClock = fakeClocks.find { clock -> clock.clockId == selectedClockId }
+            checkNotNull(selectedClock)
+            ClockMetadataModel(
+                clockId = selectedClock.clockId,
+                name = selectedClock.name,
+                isSelected = true,
+                selectedColorId = selectedColor,
+                colorToneProgress = colorTone,
+                seedColor = seedColor,
+            )
+        }
+
+    private val _selectedClockSize = MutableStateFlow(ClockSize.SMALL)
+    override val selectedClockSize: Flow<ClockSize> = _selectedClockSize.asStateFlow()
+
+    override suspend fun setSelectedClock(clockId: String) {
+        selectedClockId.value = clockId
+    }
+
+    override suspend fun setClockColor(
+        selectedColorId: String?,
+        @IntRange(from = 0, to = 100) colorToneProgress: Int,
+        @ColorInt seedColor: Int?,
+    ) {
+        this.selectedColorId.value = selectedColorId
+        this.colorTone.value = colorToneProgress
+        this.seedColor.value = seedColor
+    }
+
+    override suspend fun setClockSize(size: ClockSize) {
+        _selectedClockSize.value = size
+    }
+
+    companion object {
+        const val CLOCK_ID_0 = "clock0"
+        const val CLOCK_ID_1 = "clock1"
+        const val CLOCK_ID_2 = "clock2"
+        const val CLOCK_ID_3 = "clock3"
+        val fakeClocks =
+            listOf(
+                ClockMetadataModel(CLOCK_ID_0, "clock0", true, null, 50, null),
+                ClockMetadataModel(CLOCK_ID_1, "clock1", false, null, 50, null),
+                ClockMetadataModel(CLOCK_ID_2, "clock2", false, null, 50, null),
+                ClockMetadataModel(CLOCK_ID_3, "clock3", false, null, 50, null),
+            )
+        const val CLOCK_COLOR_ID = "RED"
+        const val CLOCK_COLOR_TONE_PROGRESS = 87
+        const val SEED_COLOR = Color.RED
+    }
+}
diff --git a/tests/robotests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt b/tests/robotests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
new file mode 100644
index 0000000..c8e39be
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
@@ -0,0 +1,83 @@
+package com.android.customization.picker.clock.domain.interactor
+
+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
+import kotlinx.coroutines.test.setMain
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::class)
+class ClockPickerInteractorTest {
+
+    private lateinit var underTest: ClockPickerInteractor
+
+    @Before
+    fun setUp() {
+        val testDispatcher = StandardTestDispatcher()
+        Dispatchers.setMain(testDispatcher)
+        underTest =
+            ClockPickerInteractor(
+                repository = FakeClockPickerRepository(),
+                snapshotRestorer = {
+                    ClockPickerSnapshotRestorer(interactor = underTest).apply {
+                        runBlocking { setUpSnapshotRestorer(store = FakeSnapshotStore()) }
+                    }
+                },
+            )
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun setSelectedClock() = runTest {
+        val observedSelectedClockId = collectLastValue(underTest.selectedClockId)
+        underTest.setSelectedClock(FakeClockPickerRepository.fakeClocks[1].clockId)
+        Truth.assertThat(observedSelectedClockId())
+            .isEqualTo(FakeClockPickerRepository.fakeClocks[1].clockId)
+    }
+
+    @Test
+    fun setClockSize() = runTest {
+        val observedClockSize = collectLastValue(underTest.selectedClockSize)
+        underTest.setClockSize(ClockSize.DYNAMIC)
+        Truth.assertThat(observedClockSize()).isEqualTo(ClockSize.DYNAMIC)
+
+        underTest.setClockSize(ClockSize.SMALL)
+        Truth.assertThat(observedClockSize()).isEqualTo(ClockSize.SMALL)
+    }
+
+    @Test
+    fun setColor() = runTest {
+        val observedSelectedColor = collectLastValue(underTest.selectedColorId)
+        val observedColorToneProgress = collectLastValue(underTest.colorToneProgress)
+        val observedSeedColor = collectLastValue(underTest.seedColor)
+        underTest.setClockColor(
+            FakeClockPickerRepository.CLOCK_COLOR_ID,
+            FakeClockPickerRepository.CLOCK_COLOR_TONE_PROGRESS,
+            FakeClockPickerRepository.SEED_COLOR,
+        )
+        Truth.assertThat(observedSelectedColor())
+            .isEqualTo(FakeClockPickerRepository.CLOCK_COLOR_ID)
+        Truth.assertThat(observedColorToneProgress())
+            .isEqualTo(FakeClockPickerRepository.CLOCK_COLOR_TONE_PROGRESS)
+        Truth.assertThat(observedSeedColor()).isEqualTo(FakeClockPickerRepository.SEED_COLOR)
+    }
+}
diff --git a/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
new file mode 100644
index 0000000..1b1eb9a
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
@@ -0,0 +1,134 @@
+/*
+ * 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.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
+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.robolectric.RobolectricTestRunner
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::class)
+class ClockCarouselViewModelTest {
+    private val repositoryWithMultipleClocks by lazy { FakeClockPickerRepository() }
+    private val repositoryWithSingleClock by lazy {
+        FakeClockPickerRepository(
+            listOf(
+                ClockMetadataModel(
+                    clockId = "clock0",
+                    name = "clock0",
+                    isSelected = true,
+                    selectedColorId = null,
+                    colorToneProgress = ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS,
+                    seedColor = null,
+                ),
+            )
+        )
+    }
+    private lateinit var testDispatcher: CoroutineDispatcher
+    private lateinit var underTest: ClockCarouselViewModel
+    private lateinit var interactor: ClockPickerInteractor
+
+    @Before
+    fun setUp() {
+        testDispatcher = StandardTestDispatcher()
+        Dispatchers.setMain(testDispatcher)
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun setSelectedClock() = runTest {
+        underTest =
+            ClockCarouselViewModel(
+                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)
+    }
+
+    @Test
+    fun multipleClockCase() = runTest {
+        underTest =
+            ClockCarouselViewModel(
+                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()
+    }
+
+    @Test
+    fun singleClockCase() = runTest {
+        underTest =
+            ClockCarouselViewModel(
+                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/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
new file mode 100644
index 0000000..19a704c
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
@@ -0,0 +1,95 @@
+/*
+ * 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.ui.viewmodel
+
+import androidx.test.filters.SmallTest
+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
+import kotlinx.coroutines.test.setMain
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::class)
+class ClockSectionViewModelTest {
+
+    private lateinit var clockColorMap: Map<String, ClockColorViewModel>
+    private lateinit var interactor: ClockPickerInteractor
+    private lateinit var underTest: ClockSectionViewModel
+
+    @Before
+    fun setUp() {
+        val testDispatcher = StandardTestDispatcher()
+        Dispatchers.setMain(testDispatcher)
+        val context = InstrumentationRegistry.getInstrumentation().targetContext
+        clockColorMap = ClockColorViewModel.getPresetColorMap(context.resources)
+        interactor =
+            ClockPickerInteractor(
+                repository = FakeClockPickerRepository(),
+                snapshotRestorer = {
+                    ClockPickerSnapshotRestorer(interactor = interactor).apply {
+                        runBlocking { setUpSnapshotRestorer(store = FakeSnapshotStore()) }
+                    }
+                },
+            )
+        underTest =
+            ClockSectionViewModel(
+                context,
+                interactor,
+            )
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun setSelectedClock() = runTest {
+        val colorGrey = clockColorMap.values.first()
+        val observedSelectedClockColorAndSizeText =
+            collectLastValue(underTest.selectedClockColorAndSizeText)
+
+        interactor.setClockColor(
+            colorGrey.colorId,
+            ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS,
+            ClockSettingsViewModel.blendColorWithTone(
+                colorGrey.color,
+                colorGrey.getColorTone(ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS),
+            )
+        )
+        interactor.setClockSize(ClockSize.DYNAMIC)
+
+        assertThat(observedSelectedClockColorAndSizeText()).isEqualTo("Grey, dynamic")
+    }
+}
diff --git a/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
new file mode 100644
index 0000000..f09e977
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
@@ -0,0 +1,210 @@
+package com.android.customization.picker.clock.ui.viewmodel
+
+import android.content.Context
+import androidx.test.filters.SmallTest
+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
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
+import com.android.customization.picker.color.domain.interactor.ColorPickerSnapshotRestorer
+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.launch
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.advanceTimeBy
+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.robolectric.RobolectricTestRunner
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::class)
+class ClockSettingsViewModelTest {
+
+    private lateinit var context: Context
+    private lateinit var testScope: TestScope
+    private lateinit var colorPickerInteractor: ColorPickerInteractor
+    private lateinit var clockPickerInteractor: ClockPickerInteractor
+    private lateinit var underTest: ClockSettingsViewModel
+    private lateinit var colorMap: Map<String, ClockColorViewModel>
+    // We make the condition that CLOCK_ID_3 is not reactive to tone
+    private val getIsReactiveToTone: (clockId: String?) -> Boolean = { clockId ->
+        when (clockId) {
+            FakeClockPickerRepository.CLOCK_ID_0 -> true
+            FakeClockPickerRepository.CLOCK_ID_1 -> true
+            FakeClockPickerRepository.CLOCK_ID_2 -> true
+            FakeClockPickerRepository.CLOCK_ID_3 -> false
+            else -> false
+        }
+    }
+
+    @Before
+    fun setUp() {
+        val testDispatcher = StandardTestDispatcher()
+        Dispatchers.setMain(testDispatcher)
+        context = InstrumentationRegistry.getInstrumentation().targetContext
+        testScope = TestScope(testDispatcher)
+        clockPickerInteractor =
+            ClockPickerInteractor(
+                repository = FakeClockPickerRepository(),
+                snapshotRestorer = {
+                    ClockPickerSnapshotRestorer(interactor = clockPickerInteractor).apply {
+                        runBlocking { setUpSnapshotRestorer(store = FakeSnapshotStore()) }
+                    }
+                },
+            )
+        colorPickerInteractor =
+            ColorPickerInteractor(
+                repository = FakeColorPickerRepository(context = context),
+                snapshotRestorer = {
+                    ColorPickerSnapshotRestorer(interactor = colorPickerInteractor).apply {
+                        runBlocking { setUpSnapshotRestorer(store = FakeSnapshotStore()) }
+                    }
+                },
+            )
+        underTest =
+            ClockSettingsViewModel.Factory(
+                    context = context,
+                    clockPickerInteractor = clockPickerInteractor,
+                    colorPickerInteractor = colorPickerInteractor,
+                    getIsReactiveToTone = getIsReactiveToTone,
+                )
+                .create(ClockSettingsViewModel::class.java)
+        colorMap = ClockColorViewModel.getPresetColorMap(context.resources)
+
+        testScope.launch {
+            clockPickerInteractor.setSelectedClock(FakeClockPickerRepository.CLOCK_ID_0)
+        }
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun clickOnColorSettingsTab() = runTest {
+        val tabs = collectLastValue(underTest.tabs)
+        assertThat(tabs()?.get(0)?.name).isEqualTo("Color")
+        assertThat(tabs()?.get(0)?.isSelected).isTrue()
+        assertThat(tabs()?.get(1)?.name).isEqualTo("Size")
+        assertThat(tabs()?.get(1)?.isSelected).isFalse()
+
+        tabs()?.get(1)?.onClicked?.invoke()
+        assertThat(tabs()?.get(0)?.isSelected).isFalse()
+        assertThat(tabs()?.get(1)?.isSelected).isTrue()
+    }
+
+    @Test
+    fun setSelectedColor() = runTest {
+        val observedClockColorOptions = collectLastValue(underTest.colorOptions)
+        val observedSelectedColorOptionPosition =
+            collectLastValue(underTest.selectedColorOptionPosition)
+        val observedSliderProgress = collectLastValue(underTest.sliderProgress)
+        val observedSeedColor = collectLastValue(underTest.seedColor)
+        // Advance COLOR_OPTIONS_EVENT_UPDATE_DELAY_MILLIS since there is a delay from colorOptions
+        advanceTimeBy(ClockSettingsViewModel.COLOR_OPTIONS_EVENT_UPDATE_DELAY_MILLIS)
+        val option0IsSelected = collectLastValue(observedClockColorOptions()!![0].isSelected)
+        val option0OnClicked = collectLastValue(observedClockColorOptions()!![0].onClicked)
+        assertThat(option0IsSelected()).isTrue()
+        assertThat(option0OnClicked()).isNull()
+        assertThat(observedSelectedColorOptionPosition()).isEqualTo(0)
+
+        val option1OnClickedBefore = collectLastValue(observedClockColorOptions()!![1].onClicked)
+        option1OnClickedBefore()?.invoke()
+        // Advance COLOR_OPTIONS_EVENT_UPDATE_DELAY_MILLIS since there is a delay from colorOptions
+        advanceTimeBy(ClockSettingsViewModel.COLOR_OPTIONS_EVENT_UPDATE_DELAY_MILLIS)
+        val option1IsSelected = collectLastValue(observedClockColorOptions()!![1].isSelected)
+        val option1OnClickedAfter = collectLastValue(observedClockColorOptions()!![1].onClicked)
+        assertThat(option1IsSelected()).isTrue()
+        assertThat(option1OnClickedAfter()).isNull()
+        assertThat(observedSelectedColorOptionPosition()).isEqualTo(1)
+        assertThat(observedSliderProgress())
+            .isEqualTo(ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS)
+        val expectedSelectedColorModel = colorMap.values.first() // RED
+        assertThat(observedSeedColor())
+            .isEqualTo(
+                ClockSettingsViewModel.blendColorWithTone(
+                    expectedSelectedColorModel.color,
+                    expectedSelectedColorModel.getColorTone(
+                        ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS
+                    ),
+                )
+            )
+    }
+
+    @Test
+    fun setColorTone() = runTest {
+        val observedClockColorOptions = collectLastValue(underTest.colorOptions)
+        val observedIsSliderEnabled = collectLastValue(underTest.isSliderEnabled)
+        val observedSliderProgress = collectLastValue(underTest.sliderProgress)
+        val observedSeedColor = collectLastValue(underTest.seedColor)
+        // Advance COLOR_OPTIONS_EVENT_UPDATE_DELAY_MILLIS since there is a delay from colorOptions
+        advanceTimeBy(ClockSettingsViewModel.COLOR_OPTIONS_EVENT_UPDATE_DELAY_MILLIS)
+        val option0IsSelected = collectLastValue(observedClockColorOptions()!![0].isSelected)
+        assertThat(option0IsSelected()).isTrue()
+        assertThat(observedIsSliderEnabled()).isFalse()
+
+        val option1OnClicked = collectLastValue(observedClockColorOptions()!![1].onClicked)
+        option1OnClicked()?.invoke()
+
+        // Advance COLOR_OPTIONS_EVENT_UPDATE_DELAY_MILLIS since there is a delay from colorOptions
+        advanceTimeBy(ClockSettingsViewModel.COLOR_OPTIONS_EVENT_UPDATE_DELAY_MILLIS)
+        assertThat(observedIsSliderEnabled()).isTrue()
+        val targetProgress1 = 99
+        underTest.onSliderProgressChanged(targetProgress1)
+        assertThat(observedSliderProgress()).isEqualTo(targetProgress1)
+        val targetProgress2 = 55
+        testScope.launch { underTest.onSliderProgressStop(targetProgress2) }
+        assertThat(observedSliderProgress()).isEqualTo(targetProgress2)
+        val expectedSelectedColorModel = colorMap.values.first() // RED
+        assertThat(observedSeedColor())
+            .isEqualTo(
+                ClockSettingsViewModel.blendColorWithTone(
+                    expectedSelectedColorModel.color,
+                    expectedSelectedColorModel.getColorTone(targetProgress2),
+                )
+            )
+    }
+
+    @Test
+    fun setClockSize() = runTest {
+        val observedClockSize = collectLastValue(underTest.selectedClockSize)
+        underTest.setClockSize(ClockSize.DYNAMIC)
+        assertThat(observedClockSize()).isEqualTo(ClockSize.DYNAMIC)
+
+        underTest.setClockSize(ClockSize.SMALL)
+        assertThat(observedClockSize()).isEqualTo(ClockSize.SMALL)
+    }
+
+    @Test
+    fun getIsReactiveToTone() = runTest {
+        val observedClockColorOptions = collectLastValue(underTest.colorOptions)
+        val isSliderEnabled = collectLastValue(underTest.isSliderEnabled)
+        // Advance COLOR_OPTIONS_EVENT_UPDATE_DELAY_MILLIS since there is a delay from colorOptions
+        advanceTimeBy(ClockSettingsViewModel.COLOR_OPTIONS_EVENT_UPDATE_DELAY_MILLIS)
+        val option1OnClicked = collectLastValue(observedClockColorOptions()!![1].onClicked)
+        option1OnClicked()?.invoke()
+
+        clockPickerInteractor.setSelectedClock(FakeClockPickerRepository.CLOCK_ID_0)
+        assertThat(isSliderEnabled()).isTrue()
+
+        // We make the condition that CLOCK_ID_0 is not reactive to tone
+        clockPickerInteractor.setSelectedClock(FakeClockPickerRepository.CLOCK_ID_3)
+        assertThat(isSliderEnabled()).isFalse()
+    }
+}
diff --git a/tests/robotests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt b/tests/robotests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt
new file mode 100644
index 0000000..be799db
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.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.notifications.data.repository
+
+import android.provider.Settings
+import androidx.test.filters.SmallTest
+import com.android.customization.picker.notifications.shared.model.NotificationSettingsModel
+import com.android.wallpaper.testing.FakeSecureSettingsRepository
+import com.android.wallpaper.testing.collectLastValue
+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 NotificationsRepositoryTest {
+
+    private lateinit var underTest: NotificationsRepository
+
+    private lateinit var testScope: TestScope
+    private lateinit var secureSettingsRepository: FakeSecureSettingsRepository
+
+    @Before
+    fun setUp() {
+        val testDispatcher = StandardTestDispatcher()
+        testScope = TestScope(testDispatcher)
+        secureSettingsRepository = FakeSecureSettingsRepository()
+
+        underTest =
+            NotificationsRepository(
+                scope = testScope.backgroundScope,
+                backgroundDispatcher = testDispatcher,
+                secureSettingsRepository = secureSettingsRepository,
+            )
+    }
+
+    @Test
+    fun settings() =
+        testScope.runTest {
+            val settings = collectLastValue(underTest.settings)
+
+            secureSettingsRepository.set(
+                name = Settings.Secure.LOCK_SCREEN_SHOW_NOTIFICATIONS,
+                value = 1,
+            )
+            assertThat(settings())
+                .isEqualTo(NotificationSettingsModel(isShowNotificationsOnLockScreenEnabled = true))
+
+            secureSettingsRepository.set(
+                name = Settings.Secure.LOCK_SCREEN_SHOW_NOTIFICATIONS,
+                value = 0,
+            )
+            assertThat(settings())
+                .isEqualTo(
+                    NotificationSettingsModel(isShowNotificationsOnLockScreenEnabled = false)
+                )
+        }
+
+    @Test
+    fun setSettings() =
+        testScope.runTest {
+            val settings = collectLastValue(underTest.settings)
+
+            val model1 = NotificationSettingsModel(isShowNotificationsOnLockScreenEnabled = true)
+            underTest.setSettings(model1)
+            assertThat(settings()).isEqualTo(model1)
+
+            val model2 = NotificationSettingsModel(isShowNotificationsOnLockScreenEnabled = false)
+            underTest.setSettings(model2)
+            assertThat(settings()).isEqualTo(model2)
+        }
+}
diff --git a/tests/robotests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
new file mode 100644
index 0000000..5c3544a
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
@@ -0,0 +1,97 @@
+/*
+ * 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.notifications.ui.viewmodel
+
+import androidx.test.filters.SmallTest
+import com.android.customization.picker.notifications.data.repository.NotificationsRepository
+import com.android.customization.picker.notifications.domain.interactor.NotificationsInteractor
+import com.android.customization.picker.notifications.domain.interactor.NotificationsSnapshotRestorer
+import com.android.wallpaper.testing.FakeSecureSettingsRepository
+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.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+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.robolectric.RobolectricTestRunner
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::class)
+class NotificationSectionViewModelTest {
+
+    private lateinit var underTest: NotificationSectionViewModel
+
+    private lateinit var testScope: TestScope
+    private lateinit var interactor: NotificationsInteractor
+
+    @Before
+    fun setUp() {
+        val testDispatcher = UnconfinedTestDispatcher()
+        Dispatchers.setMain(testDispatcher)
+        testScope = TestScope(testDispatcher)
+        interactor =
+            NotificationsInteractor(
+                repository =
+                    NotificationsRepository(
+                        scope = testScope.backgroundScope,
+                        backgroundDispatcher = testDispatcher,
+                        secureSettingsRepository = FakeSecureSettingsRepository(),
+                    ),
+                snapshotRestorer = {
+                    NotificationsSnapshotRestorer(
+                            interactor = interactor,
+                        )
+                        .apply { runBlocking { setUpSnapshotRestorer(FakeSnapshotStore()) } }
+                },
+            )
+
+        underTest =
+            NotificationSectionViewModel(
+                interactor = interactor,
+            )
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun `toggles back and forth`() =
+        testScope.runTest {
+            val isSwitchOn = collectLastValue(underTest.isSwitchOn)
+
+            val initialIsSwitchOn = isSwitchOn()
+
+            underTest.onClicked()
+            assertThat(isSwitchOn()).isNotEqualTo(initialIsSwitchOn)
+
+            underTest.onClicked()
+            assertThat(isSwitchOn()).isEqualTo(initialIsSwitchOn)
+        }
+}