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

Android 13.0.0 Release 52 (TQ3A.230605.012)

Change-Id: I2cea11fa2f1f02fbd3c9d21cfc1697a79d42a5b7

26 files changed, 2401 insertions, 254 deletions

diff --git a/tests/robotests/Android.bp b/tests/robotests/Android.bp
index e0a37c29..4416e1c4 100644
--- a/tests/robotests/Android.bp
+++ b/tests/robotests/Android.bp
@@ -13,5 +13,8 @@ android_robolectric_test {
         "androidx.test.core",
         "androidx.test.runner",
     ],
+
+    upstream: true,
+
     instrumentation_for: "ThemePicker",
 }
diff --git a/tests/robotests/src/com/android/customization/model/clock/BaseClockManagerTest.java b/tests/robotests/src/com/android/customization/model/clock/BaseClockManagerTest.java
index c96e7f8b..eeff5317 100644
--- a/tests/robotests/src/com/android/customization/model/clock/BaseClockManagerTest.java
+++ b/tests/robotests/src/com/android/customization/model/clock/BaseClockManagerTest.java
@@ -20,7 +20,7 @@ import static junit.framework.TestCase.fail;
 
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Matchers.anyBoolean;
+import static org.mockito.Mockito.anyBoolean;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
diff --git a/tests/robotests/src/com/android/customization/model/grid/GridOptionsManagerTest.java b/tests/robotests/src/com/android/customization/model/grid/GridOptionsManagerTest.java
index 89ca6761..04ac0241 100644
--- a/tests/robotests/src/com/android/customization/model/grid/GridOptionsManagerTest.java
+++ b/tests/robotests/src/com/android/customization/model/grid/GridOptionsManagerTest.java
@@ -17,7 +17,7 @@ package com.android.customization.model.grid;
 
 import static junit.framework.TestCase.fail;
 
-import static org.mockito.Matchers.anyBoolean;
+import static org.mockito.Mockito.anyBoolean;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
diff --git a/tests/robotests/src/com/android/customization/picker/clock/ClockCustomDemoFragmentTest.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragmentTest.kt
index ad3dd1ce..0a543123 100644
--- a/tests/robotests/src/com/android/customization/picker/clock/ClockCustomDemoFragmentTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragmentTest.kt
@@ -1,23 +1,21 @@
-package com.android.customization.picker.clock
+package com.android.customization.picker.clock.ui.fragment
 
-import android.os.Handler
-import android.os.UserHandle
 import android.view.View
 import androidx.appcompat.app.AppCompatActivity
 import androidx.recyclerview.widget.LinearLayoutManager
 import androidx.recyclerview.widget.RecyclerView
-import com.android.systemui.plugins.ClockId
 import com.android.systemui.plugins.ClockMetadata
-import com.android.systemui.plugins.ClockProvider
-import com.android.systemui.plugins.ClockProviderPlugin
+import com.android.systemui.plugins.ClockSettings
 import com.android.systemui.plugins.PluginManager
 import com.android.systemui.shared.clocks.ClockRegistry
 import org.junit.Assert
 import org.junit.Before
+import org.junit.Ignore
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mock
-import org.mockito.Mockito
+import org.mockito.Mockito.verify
+import org.mockito.Mockito.`when` as whenever
 import org.mockito.MockitoAnnotations
 import org.robolectric.Robolectric
 import org.robolectric.RobolectricTestRunner
@@ -26,55 +24,33 @@ import org.robolectric.annotation.Config
 /** Tests of [ClockCustomDemoFragment]. */
 @RunWith(RobolectricTestRunner::class)
 @Config(manifest = Config.NONE)
+@Ignore("b/270606895")
 class ClockCustomDemoFragmentTest {
     private lateinit var mActivity: AppCompatActivity
     private var mClockCustomDemoFragment: ClockCustomDemoFragment? = null
-    private lateinit var registry: ClockRegistry
+    @Mock private lateinit var registry: ClockRegistry
     @Mock private lateinit var mockPluginManager: PluginManager
-    @Mock private lateinit var mockHandler: Handler
-    @Mock private lateinit var fakePlugin: ClockProviderPlugin
-    @Mock private lateinit var defaultClockProvider: ClockProvider
 
-    private var settingValue: String = ""
+    private var settingValue: ClockSettings? = null
 
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
         mActivity = Robolectric.buildActivity(AppCompatActivity::class.java).get()
         mClockCustomDemoFragment = ClockCustomDemoFragment()
-        Mockito.`when`(defaultClockProvider.getClocks())
-            .thenReturn(listOf(ClockMetadata("DEFAULT", "Default Clock")))
-        registry =
-            object :
-                ClockRegistry(
-                    mActivity,
-                    mockPluginManager,
-                    mockHandler,
-                    isEnabled = true,
-                    userHandle = UserHandle.USER_ALL,
-                    defaultClockProvider = defaultClockProvider
-                ) {
-                override var currentClockId: ClockId
-                    get() = settingValue
-                    set(value) {
-                        settingValue = value
-                    }
-
-                override fun getClocks(): List<ClockMetadata> {
-                    return defaultClockProvider.getClocks() +
-                        listOf(
-                            ClockMetadata("CLOCK_1", "Clock 1"),
-                            ClockMetadata("CLOCK_2", "Clock 2"),
-                            ClockMetadata("CLOCK_NOT_IN_USE", "Clock not in use")
-                        )
-                }
-            }
+        whenever(registry.getClocks())
+            .thenReturn(
+                listOf(
+                    ClockMetadata("CLOCK_1", "Clock 1"),
+                    ClockMetadata("CLOCK_2", "Clock 2"),
+                    ClockMetadata("CLOCK_NOT_IN_USE", "Clock not in use")
+                )
+            )
 
         mClockCustomDemoFragment!!.clockRegistry = registry
         mClockCustomDemoFragment!!.recyclerView = RecyclerView(mActivity)
         mClockCustomDemoFragment!!.recyclerView.layoutManager =
             LinearLayoutManager(mActivity, RecyclerView.VERTICAL, false)
-        mClockCustomDemoFragment!!.pluginListener.onPluginConnected(fakePlugin, mActivity)
     }
 
     @Test
@@ -94,6 +70,6 @@ class ClockCustomDemoFragmentTest {
             .findViewHolderForAdapterPosition(testPosition)
             ?.itemView
             ?.performClick()
-        Assert.assertEquals("CLOCK_1", settingValue)
+        verify(registry).currentClockId = "CLOCK_1"
     }
 }
diff --git a/tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt b/tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
new file mode 100644
index 00000000..59539379
--- /dev/null
+++ b/tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
@@ -0,0 +1,91 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.data.repository
+
+import com.android.customization.model.grid.shared.model.GridOptionItemModel
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.channels.BufferOverflow
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.asSharedFlow
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
+
+class FakeGridRepository(
+    private val scope: CoroutineScope,
+    initialOptionCount: Int,
+    var available: Boolean = true
+) : GridRepository {
+    private val _optionChanges =
+        MutableSharedFlow<Unit>(
+            replay = 1,
+            onBufferOverflow = BufferOverflow.DROP_OLDEST,
+        )
+
+    override suspend fun isAvailable(): Boolean = available
+
+    override fun getOptionChanges(): Flow<Unit> = _optionChanges.asSharedFlow()
+
+    private val selectedOptionIndex = MutableStateFlow(0)
+    private var options: GridOptionItemsModel = createOptions(count = initialOptionCount)
+
+    override suspend fun getOptions(): GridOptionItemsModel {
+        return options
+    }
+
+    fun setOptions(
+        count: Int,
+        selectedIndex: Int = 0,
+    ) {
+        options = createOptions(count, selectedIndex)
+        _optionChanges.tryEmit(Unit)
+    }
+
+    private fun createOptions(
+        count: Int,
+        selectedIndex: Int = 0,
+    ): GridOptionItemsModel {
+        selectedOptionIndex.value = selectedIndex
+        return GridOptionItemsModel.Loaded(
+            options =
+                buildList {
+                    repeat(times = count) { index ->
+                        add(
+                            GridOptionItemModel(
+                                name = "option_$index",
+                                cols = 4,
+                                rows = index * 2,
+                                isSelected =
+                                    selectedOptionIndex
+                                        .map { it == index }
+                                        .stateIn(
+                                            scope = scope,
+                                            started = SharingStarted.Eagerly,
+                                            initialValue = false,
+                                        ),
+                                onSelected = { selectedOptionIndex.value = index },
+                            )
+                        )
+                    }
+                }
+        )
+    }
+}
diff --git a/tests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt b/tests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt
new file mode 100644
index 00000000..f73d5a38
--- /dev/null
+++ b/tests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt
@@ -0,0 +1,146 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.domain.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.customization.model.grid.data.repository.FakeGridRepository
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class GridInteractorTest {
+
+    private lateinit var underTest: GridInteractor
+    private lateinit var testScope: TestScope
+    private lateinit var repository: FakeGridRepository
+    private lateinit var store: FakeSnapshotStore
+
+    @Before
+    fun setUp() {
+        testScope = TestScope()
+        repository =
+            FakeGridRepository(
+                scope = testScope.backgroundScope,
+                initialOptionCount = 3,
+            )
+        store = FakeSnapshotStore()
+        underTest =
+            GridInteractor(
+                applicationScope = testScope.backgroundScope,
+                repository = repository,
+                snapshotRestorer = {
+                    GridSnapshotRestorer(
+                            interactor = underTest,
+                        )
+                        .apply {
+                            runBlocking {
+                                setUpSnapshotRestorer(
+                                    store = store,
+                                )
+                            }
+                        }
+                },
+            )
+    }
+
+    @Test
+    fun selectingOptionThroughModel_updatesOptions() =
+        testScope.runTest {
+            val options = collectLastValue(underTest.options)
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(3)
+                assertThat(loaded.options[0].isSelected.value).isTrue()
+                assertThat(loaded.options[1].isSelected.value).isFalse()
+                assertThat(loaded.options[2].isSelected.value).isFalse()
+            }
+
+            val storedSnapshot = store.retrieve()
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                loaded.options[1].onSelected()
+            }
+
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(3)
+                assertThat(loaded.options[0].isSelected.value).isFalse()
+                assertThat(loaded.options[1].isSelected.value).isTrue()
+                assertThat(loaded.options[2].isSelected.value).isFalse()
+            }
+            assertThat(store.retrieve()).isNotEqualTo(storedSnapshot)
+        }
+
+    @Test
+    fun selectingOptionThroughSetter_returnsSelectedOptionFromGetter() =
+        testScope.runTest {
+            val options = collectLastValue(underTest.options)
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(3)
+            }
+
+            val storedSnapshot = store.retrieve()
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                underTest.setSelectedOption(loaded.options[1])
+                runCurrent()
+                assertThat(underTest.getSelectedOption()?.name).isEqualTo(loaded.options[1].name)
+                assertThat(store.retrieve()).isNotEqualTo(storedSnapshot)
+            }
+        }
+
+    @Test
+    fun externalUpdates_reloadInvoked() =
+        testScope.runTest {
+            val options = collectLastValue(underTest.options)
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(3)
+            }
+
+            val storedSnapshot = store.retrieve()
+            repository.setOptions(4)
+
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(4)
+            }
+            // External updates do not record a new snapshot with the undo system.
+            assertThat(store.retrieve()).isEqualTo(storedSnapshot)
+        }
+
+    @Test
+    fun unavailableRepository_emptyOptions() =
+        testScope.runTest {
+            repository.available = false
+            val options = collectLastValue(underTest.options)
+            assertThat(options()).isNull()
+        }
+}
diff --git a/tests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt b/tests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt
new file mode 100644
index 00000000..c2712b1d
--- /dev/null
+++ b/tests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt
@@ -0,0 +1,111 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.domain.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.customization.model.grid.data.repository.FakeGridRepository
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class GridSnapshotRestorerTest {
+
+    private lateinit var underTest: GridSnapshotRestorer
+    private lateinit var testScope: TestScope
+    private lateinit var repository: FakeGridRepository
+    private lateinit var store: FakeSnapshotStore
+
+    @Before
+    fun setUp() {
+        testScope = TestScope()
+        repository =
+            FakeGridRepository(
+                scope = testScope.backgroundScope,
+                initialOptionCount = 4,
+            )
+        underTest =
+            GridSnapshotRestorer(
+                interactor =
+                    GridInteractor(
+                        applicationScope = testScope.backgroundScope,
+                        repository = repository,
+                        snapshotRestorer = { underTest },
+                    )
+            )
+        store = FakeSnapshotStore()
+    }
+
+    @Test
+    fun restoreToSnapshot_noCallsToStore_restoresToInitialSnapshot() =
+        testScope.runTest {
+            runCurrent()
+            val initialSnapshot = underTest.setUpSnapshotRestorer(store = store)
+            assertThat(initialSnapshot.args).isNotEmpty()
+            repository.setOptions(
+                count = 4,
+                selectedIndex = 2,
+            )
+            runCurrent()
+            assertThat(getSelectedIndex()).isEqualTo(2)
+
+            underTest.restoreToSnapshot(initialSnapshot)
+            runCurrent()
+
+            assertThat(getSelectedIndex()).isEqualTo(0)
+        }
+
+    @Test
+    fun restoreToSnapshot_withCallToStore_restoresToInitialSnapshot() =
+        testScope.runTest {
+            runCurrent()
+            val initialSnapshot = underTest.setUpSnapshotRestorer(store = store)
+            assertThat(initialSnapshot.args).isNotEmpty()
+            repository.setOptions(
+                count = 4,
+                selectedIndex = 2,
+            )
+            runCurrent()
+            assertThat(getSelectedIndex()).isEqualTo(2)
+            underTest.store((repository.getOptions() as GridOptionItemsModel.Loaded).options[1])
+            runCurrent()
+
+            underTest.restoreToSnapshot(initialSnapshot)
+            runCurrent()
+
+            assertThat(getSelectedIndex()).isEqualTo(0)
+        }
+
+    private suspend fun getSelectedIndex(): Int {
+        return (repository.getOptions() as? GridOptionItemsModel.Loaded)?.options?.indexOfFirst {
+            optionItem ->
+            optionItem.isSelected.value
+        }
+            ?: -1
+    }
+}
diff --git a/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt b/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
new file mode 100644
index 00000000..58c5d99f
--- /dev/null
+++ b/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
@@ -0,0 +1,110 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.ui.viewmodel
+
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.customization.model.grid.data.repository.FakeGridRepository
+import com.android.customization.model.grid.domain.interactor.GridInteractor
+import com.android.customization.model.grid.domain.interactor.GridSnapshotRestorer
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Ignore
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class GridScreenViewModelTest {
+
+    private lateinit var underTest: GridScreenViewModel
+    private lateinit var testScope: TestScope
+    private lateinit var interactor: GridInteractor
+    private lateinit var store: FakeSnapshotStore
+
+    @Before
+    fun setUp() {
+        testScope = TestScope()
+        store = FakeSnapshotStore()
+        interactor =
+            GridInteractor(
+                applicationScope = testScope.backgroundScope,
+                repository =
+                    FakeGridRepository(
+                        scope = testScope.backgroundScope,
+                        initialOptionCount = 4,
+                    ),
+                snapshotRestorer = {
+                    GridSnapshotRestorer(
+                            interactor = interactor,
+                        )
+                        .apply { runBlocking { setUpSnapshotRestorer(store) } }
+                }
+            )
+
+        underTest =
+            GridScreenViewModel(
+                context = InstrumentationRegistry.getInstrumentation().targetContext,
+                interactor = interactor,
+            )
+    }
+
+    @Test
+    @Ignore("b/270371382")
+    fun clickOnItem_itGetsSelected() =
+        testScope.runTest {
+            val optionItemsValueProvider = collectLastValue(underTest.optionItems)
+            var optionItemsValue = checkNotNull(optionItemsValueProvider.invoke())
+            assertThat(optionItemsValue).hasSize(4)
+            assertThat(getSelectedIndex(optionItemsValue)).isEqualTo(0)
+            assertThat(getOnClick(optionItemsValue[0])).isNull()
+
+            val item1OnClickedValue = getOnClick(optionItemsValue[1])
+            assertThat(item1OnClickedValue).isNotNull()
+            item1OnClickedValue?.invoke()
+
+            optionItemsValue = checkNotNull(optionItemsValueProvider.invoke())
+            assertThat(optionItemsValue).hasSize(4)
+            assertThat(getSelectedIndex(optionItemsValue)).isEqualTo(1)
+            assertThat(getOnClick(optionItemsValue[0])).isNotNull()
+            assertThat(getOnClick(optionItemsValue[1])).isNull()
+        }
+
+    private fun TestScope.getSelectedIndex(
+        optionItems: List<OptionItemViewModel<GridIconViewModel>>
+    ): Int {
+        return optionItems.indexOfFirst { optionItem ->
+            collectLastValue(optionItem.isSelected).invoke() == true
+        }
+    }
+
+    private fun TestScope.getOnClick(
+        optionItem: OptionItemViewModel<GridIconViewModel>
+    ): (() -> Unit)? {
+        return collectLastValue(optionItem.onClicked).invoke()
+    }
+}
diff --git a/tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt b/tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
new file mode 100644
index 00000000..38067b75
--- /dev/null
+++ b/tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
@@ -0,0 +1,107 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.mode
+
+import androidx.test.filters.SmallTest
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class DarkModeSnapshotRestorerTest {
+
+    private lateinit var underTest: DarkModeSnapshotRestorer
+    private lateinit var testScope: TestScope
+
+    private var isActive = false
+
+    @Before
+    fun setUp() {
+        val testDispatcher = StandardTestDispatcher()
+        testScope = TestScope(testDispatcher)
+        underTest =
+            DarkModeSnapshotRestorer(
+                backgroundDispatcher = testDispatcher,
+                isActive = { isActive },
+                setActive = { isActive = it },
+            )
+    }
+
+    @Test
+    fun `set up and restore - active`() =
+        testScope.runTest {
+            isActive = true
+
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val storedSnapshot = store.retrieve()
+
+            underTest.restoreToSnapshot(snapshot = storedSnapshot)
+            assertThat(isActive).isTrue()
+        }
+
+    @Test
+    fun `set up and restore - inactive`() =
+        testScope.runTest {
+            isActive = false
+
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val storedSnapshot = store.retrieve()
+
+            underTest.restoreToSnapshot(snapshot = storedSnapshot)
+            assertThat(isActive).isFalse()
+        }
+
+    @Test
+    fun `set up - deactivate - restore to active`() =
+        testScope.runTest {
+            isActive = true
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val initialSnapshot = store.retrieve()
+
+            underTest.store(isActivated = false)
+
+            underTest.restoreToSnapshot(snapshot = initialSnapshot)
+            assertThat(isActive).isTrue()
+        }
+
+    @Test
+    fun `set up - activate - restore to inactive`() =
+        testScope.runTest {
+            isActive = false
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val initialSnapshot = store.retrieve()
+
+            underTest.store(isActivated = true)
+
+            underTest.restoreToSnapshot(snapshot = initialSnapshot)
+            assertThat(isActive).isFalse()
+        }
+}
diff --git a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt b/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
new file mode 100644
index 00000000..cbf1365a
--- /dev/null
+++ b/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
@@ -0,0 +1,118 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+package com.android.customization.model.picker.color.domain.interactor
+
+import android.content.Context
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.customization.picker.color.data.repository.FakeColorPickerRepository
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
+import com.android.customization.picker.color.domain.interactor.ColorPickerSnapshotRestorer
+import com.android.customization.picker.color.shared.model.ColorType
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ColorPickerInteractorTest {
+    private lateinit var underTest: ColorPickerInteractor
+    private lateinit var repository: FakeColorPickerRepository
+    private lateinit var store: FakeSnapshotStore
+
+    private lateinit var context: Context
+
+    @Before
+    fun setUp() {
+        context = InstrumentationRegistry.getInstrumentation().targetContext
+        repository = FakeColorPickerRepository(context = context)
+        store = FakeSnapshotStore()
+        underTest =
+            ColorPickerInteractor(
+                repository = repository,
+                snapshotRestorer = {
+                    ColorPickerSnapshotRestorer(interactor = underTest).apply {
+                        runBlocking { setUpSnapshotRestorer(store = store) }
+                    }
+                },
+            )
+        repository.setOptions(4, 4, ColorType.WALLPAPER_COLOR, 0)
+    }
+
+    @Test
+    fun select() = runTest {
+        val colorOptions = collectLastValue(underTest.colorOptions)
+
+        val wallpaperColorOptionModelBefore = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(2)
+        assertThat(wallpaperColorOptionModelBefore?.isSelected).isFalse()
+
+        wallpaperColorOptionModelBefore?.let { underTest.select(colorOptionModel = it) }
+        val wallpaperColorOptionModelAfter = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(2)
+        assertThat(wallpaperColorOptionModelAfter?.isSelected).isTrue()
+
+        val presetColorOptionModelBefore = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(1)
+        assertThat(presetColorOptionModelBefore?.isSelected).isFalse()
+
+        presetColorOptionModelBefore?.let { underTest.select(colorOptionModel = it) }
+        val presetColorOptionModelAfter = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(1)
+        assertThat(presetColorOptionModelAfter?.isSelected).isTrue()
+    }
+
+    @Test
+    fun snapshotRestorer_updatesSnapshot() = runTest {
+        val colorOptions = collectLastValue(underTest.colorOptions)
+        val wallpaperColorOptionModel0 = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(0)
+        val wallpaperColorOptionModel1 = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(1)
+        assertThat(wallpaperColorOptionModel0?.isSelected).isTrue()
+        assertThat(wallpaperColorOptionModel1?.isSelected).isFalse()
+
+        val storedSnapshot = store.retrieve()
+        wallpaperColorOptionModel1?.let { underTest.select(it) }
+        val wallpaperColorOptionModel0After = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(0)
+        val wallpaperColorOptionModel1After = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(1)
+        assertThat(wallpaperColorOptionModel0After?.isSelected).isFalse()
+        assertThat(wallpaperColorOptionModel1After?.isSelected).isTrue()
+
+        assertThat(store.retrieve()).isNotEqualTo(storedSnapshot)
+    }
+
+    @Test
+    fun snapshotRestorer_doesNotUpdateSnapshotOnExternalUpdates() = runTest {
+        val colorOptions = collectLastValue(underTest.colorOptions)
+        val wallpaperColorOptionModel0 = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(0)
+        val wallpaperColorOptionModel1 = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(1)
+        assertThat(wallpaperColorOptionModel0?.isSelected).isTrue()
+        assertThat(wallpaperColorOptionModel1?.isSelected).isFalse()
+
+        val storedSnapshot = store.retrieve()
+        repository.setOptions(4, 4, ColorType.WALLPAPER_COLOR, 1)
+        val wallpaperColorOptionModel0After = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(0)
+        val wallpaperColorOptionModel1After = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(1)
+        assertThat(wallpaperColorOptionModel0After?.isSelected).isFalse()
+        assertThat(wallpaperColorOptionModel1After?.isSelected).isTrue()
+
+        assertThat(store.retrieve()).isEqualTo(storedSnapshot)
+    }
+}
diff --git a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt b/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
new file mode 100644
index 00000000..71a8f235
--- /dev/null
+++ b/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
@@ -0,0 +1,138 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.picker.color.domain.interactor
+
+import android.content.Context
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.customization.picker.color.data.repository.FakeColorPickerRepository
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
+import com.android.customization.picker.color.domain.interactor.ColorPickerSnapshotRestorer
+import com.android.customization.picker.color.shared.model.ColorOptionModel
+import com.android.customization.picker.color.shared.model.ColorType
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ColorPickerSnapshotRestorerTest {
+
+    private lateinit var underTest: ColorPickerSnapshotRestorer
+    private lateinit var repository: FakeColorPickerRepository
+    private lateinit var store: FakeSnapshotStore
+
+    private lateinit var context: Context
+
+    @Before
+    fun setUp() {
+        context = InstrumentationRegistry.getInstrumentation().targetContext
+        repository = FakeColorPickerRepository(context = context)
+        underTest =
+            ColorPickerSnapshotRestorer(
+                interactor =
+                    ColorPickerInteractor(
+                        repository = repository,
+                        snapshotRestorer = { underTest },
+                    )
+            )
+        store = FakeSnapshotStore()
+    }
+
+    @Test
+    fun restoreToSnapshot_noCallsToStore_restoresToInitialSnapshot() = runTest {
+        val colorOptions = collectLastValue(repository.colorOptions)
+
+        repository.setOptions(4, 4, ColorType.WALLPAPER_COLOR, 2)
+        val initialSnapshot = underTest.setUpSnapshotRestorer(store = store)
+        assertThat(initialSnapshot.args).isNotEmpty()
+
+        val colorOptionToSelect = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(3)
+        colorOptionToSelect?.let { repository.select(it) }
+        assertState(colorOptions(), ColorType.PRESET_COLOR, 3)
+
+        underTest.restoreToSnapshot(initialSnapshot)
+        assertState(colorOptions(), ColorType.WALLPAPER_COLOR, 2)
+    }
+
+    @Test
+    fun restoreToSnapshot_withCallToStore_restoresToInitialSnapshot() = runTest {
+        val colorOptions = collectLastValue(repository.colorOptions)
+
+        repository.setOptions(4, 4, ColorType.WALLPAPER_COLOR, 2)
+        val initialSnapshot = underTest.setUpSnapshotRestorer(store = store)
+        assertThat(initialSnapshot.args).isNotEmpty()
+
+        val colorOptionToSelect = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(3)
+        colorOptionToSelect?.let { repository.select(it) }
+        assertState(colorOptions(), ColorType.PRESET_COLOR, 3)
+
+        val colorOptionToStore = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(1)
+        colorOptionToStore?.let { underTest.storeSnapshot(colorOptionToStore) }
+
+        underTest.restoreToSnapshot(initialSnapshot)
+        assertState(colorOptions(), ColorType.WALLPAPER_COLOR, 2)
+    }
+
+    private fun assertState(
+        colorOptions: Map<ColorType, List<ColorOptionModel>>?,
+        selectedColorType: ColorType,
+        selectedColorIndex: Int
+    ) {
+        var foundSelectedColorOption = false
+        assertThat(colorOptions).isNotNull()
+        val optionsOfSelectedColorType = colorOptions?.get(selectedColorType)
+        assertThat(optionsOfSelectedColorType).isNotNull()
+        if (optionsOfSelectedColorType != null) {
+            for (i in optionsOfSelectedColorType.indices) {
+                val colorOptionHasSelectedIndex = i == selectedColorIndex
+                Truth.assertWithMessage(
+                        "Expected color option with index \"${i}\" to have" +
+                            " isSelected=$colorOptionHasSelectedIndex but it was" +
+                            " ${optionsOfSelectedColorType[i].isSelected}, num options: ${colorOptions.size}"
+                    )
+                    .that(optionsOfSelectedColorType[i].isSelected)
+                    .isEqualTo(colorOptionHasSelectedIndex)
+                foundSelectedColorOption = foundSelectedColorOption || colorOptionHasSelectedIndex
+            }
+            if (selectedColorIndex == -1) {
+                Truth.assertWithMessage(
+                        "Expected no color options to be selected, but a color option is" +
+                            " selected"
+                    )
+                    .that(foundSelectedColorOption)
+                    .isFalse()
+            } else {
+                Truth.assertWithMessage(
+                        "Expected a color option to be selected, but no color option is" +
+                            " selected"
+                    )
+                    .that(foundSelectedColorOption)
+                    .isTrue()
+            }
+        }
+    }
+}
diff --git a/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt b/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
new file mode 100644
index 00000000..1d9457a2
--- /dev/null
+++ b/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
@@ -0,0 +1,262 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+package com.android.customization.model.picker.color.ui.viewmodel
+
+import android.content.Context
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.customization.picker.color.data.repository.FakeColorPickerRepository
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
+import com.android.customization.picker.color.domain.interactor.ColorPickerSnapshotRestorer
+import com.android.customization.picker.color.shared.model.ColorType
+import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
+import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
+import com.android.customization.picker.color.ui.viewmodel.ColorTypeTabViewModel
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import com.google.common.truth.Truth.assertWithMessage
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.resetMain
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.test.setMain
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ColorPickerViewModelTest {
+    private lateinit var underTest: ColorPickerViewModel
+    private lateinit var repository: FakeColorPickerRepository
+    private lateinit var interactor: ColorPickerInteractor
+    private lateinit var store: FakeSnapshotStore
+
+    private lateinit var context: Context
+    private lateinit var testScope: TestScope
+
+    @Before
+    fun setUp() {
+        context = InstrumentationRegistry.getInstrumentation().targetContext
+        val testDispatcher = StandardTestDispatcher()
+        testScope = TestScope(testDispatcher)
+        Dispatchers.setMain(testDispatcher)
+        repository = FakeColorPickerRepository(context = context)
+        store = FakeSnapshotStore()
+
+        interactor =
+            ColorPickerInteractor(
+                repository = repository,
+                snapshotRestorer = {
+                    ColorPickerSnapshotRestorer(interactor = interactor).apply {
+                        runBlocking { setUpSnapshotRestorer(store = store) }
+                    }
+                },
+            )
+
+        underTest =
+            ColorPickerViewModel.Factory(context = context, interactor = interactor)
+                .create(ColorPickerViewModel::class.java)
+
+        repository.setOptions(4, 4, ColorType.WALLPAPER_COLOR, 0)
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun `Select a color section color`() =
+        testScope.runTest {
+            val colorSectionOptions = collectLastValue(underTest.colorSectionOptions)
+
+            assertColorOptionUiState(
+                colorOptions = colorSectionOptions(),
+                selectedColorOptionIndex = 0
+            )
+
+            selectColorOption(colorSectionOptions, 2)
+            assertColorOptionUiState(
+                colorOptions = colorSectionOptions(),
+                selectedColorOptionIndex = 2
+            )
+
+            selectColorOption(colorSectionOptions, 4)
+            assertColorOptionUiState(
+                colorOptions = colorSectionOptions(),
+                selectedColorOptionIndex = 4
+            )
+        }
+
+    @Test
+    fun `Select a preset color`() =
+        testScope.runTest {
+            val colorTypes = collectLastValue(underTest.colorTypeTabs)
+            val colorOptions = collectLastValue(underTest.colorOptions)
+
+            // Initially, the wallpaper color tab should be selected
+            assertPickerUiState(
+                colorTypes = colorTypes(),
+                colorOptions = colorOptions(),
+                selectedColorTypeText = "Wallpaper colors",
+                selectedColorOptionIndex = 0
+            )
+
+            // Select "Basic colors" tab
+            colorTypes()?.get(ColorType.PRESET_COLOR)?.onClick?.invoke()
+            assertPickerUiState(
+                colorTypes = colorTypes(),
+                colorOptions = colorOptions(),
+                selectedColorTypeText = "Basic colors",
+                selectedColorOptionIndex = -1
+            )
+
+            // Select a color option
+            selectColorOption(colorOptions, 2)
+
+            // Check original option is no longer selected
+            colorTypes()?.get(ColorType.WALLPAPER_COLOR)?.onClick?.invoke()
+            assertPickerUiState(
+                colorTypes = colorTypes(),
+                colorOptions = colorOptions(),
+                selectedColorTypeText = "Wallpaper colors",
+                selectedColorOptionIndex = -1
+            )
+
+            // Check new option is selected
+            colorTypes()?.get(ColorType.PRESET_COLOR)?.onClick?.invoke()
+            assertPickerUiState(
+                colorTypes = colorTypes(),
+                colorOptions = colorOptions(),
+                selectedColorTypeText = "Basic colors",
+                selectedColorOptionIndex = 2
+            )
+        }
+
+    /** Simulates a user selecting the affordance at the given index, if that is clickable. */
+    private fun TestScope.selectColorOption(
+        colorOptions: () -> List<OptionItemViewModel<ColorOptionIconViewModel>>?,
+        index: Int,
+    ) {
+        val onClickedFlow = colorOptions()?.get(index)?.onClicked
+        val onClickedLastValueOrNull: (() -> (() -> Unit)?)? =
+            onClickedFlow?.let { collectLastValue(it) }
+        onClickedLastValueOrNull?.let { onClickedLastValue ->
+            val onClickedOrNull: (() -> Unit)? = onClickedLastValue()
+            onClickedOrNull?.let { onClicked -> onClicked() }
+        }
+    }
+
+    /**
+     * Asserts the entire picker UI state is what is expected. This includes the color type tabs and
+     * the color options list.
+     *
+     * @param colorTypes The observed color type view-models, keyed by ColorType
+     * @param colorOptions The observed color options
+     * @param selectedColorTypeText The text of the color type that's expected to be selected
+     * @param selectedColorOptionIndex The index of the color option that's expected to be selected,
+     *   -1 stands for no color option should be selected
+     */
+    private fun TestScope.assertPickerUiState(
+        colorTypes: Map<ColorType, ColorTypeTabViewModel>?,
+        colorOptions: List<OptionItemViewModel<ColorOptionIconViewModel>>?,
+        selectedColorTypeText: String,
+        selectedColorOptionIndex: Int,
+    ) {
+        assertColorTypeTabUiState(
+            colorTypes = colorTypes,
+            colorTypeId = ColorType.WALLPAPER_COLOR,
+            isSelected = "Wallpaper colors" == selectedColorTypeText,
+        )
+        assertColorTypeTabUiState(
+            colorTypes = colorTypes,
+            colorTypeId = ColorType.PRESET_COLOR,
+            isSelected = "Basic colors" == selectedColorTypeText,
+        )
+        assertColorOptionUiState(colorOptions, selectedColorOptionIndex)
+    }
+
+    /**
+     * Asserts the picker section UI state is what is expected.
+     *
+     * @param colorOptions The observed color options
+     * @param selectedColorOptionIndex The index of the color option that's expected to be selected,
+     *   -1 stands for no color option should be selected
+     */
+    private fun TestScope.assertColorOptionUiState(
+        colorOptions: List<OptionItemViewModel<ColorOptionIconViewModel>>?,
+        selectedColorOptionIndex: Int,
+    ) {
+        var foundSelectedColorOption = false
+        assertThat(colorOptions).isNotNull()
+        if (colorOptions != null) {
+            for (i in colorOptions.indices) {
+                val colorOptionHasSelectedIndex = i == selectedColorOptionIndex
+                val isSelected: Boolean? = collectLastValue(colorOptions[i].isSelected).invoke()
+                assertWithMessage(
+                        "Expected color option with index \"${i}\" to have" +
+                            " isSelected=$colorOptionHasSelectedIndex but it was" +
+                            " ${isSelected}, num options: ${colorOptions.size}"
+                    )
+                    .that(isSelected)
+                    .isEqualTo(colorOptionHasSelectedIndex)
+                foundSelectedColorOption = foundSelectedColorOption || colorOptionHasSelectedIndex
+            }
+            if (selectedColorOptionIndex == -1) {
+                assertWithMessage(
+                        "Expected no color options to be selected, but a color option is" +
+                            " selected"
+                    )
+                    .that(foundSelectedColorOption)
+                    .isFalse()
+            } else {
+                assertWithMessage(
+                        "Expected a color option to be selected, but no color option is" +
+                            " selected"
+                    )
+                    .that(foundSelectedColorOption)
+                    .isTrue()
+            }
+        }
+    }
+
+    /**
+     * Asserts that a color type tab has the correct UI state.
+     *
+     * @param colorTypes The observed color type view-models, keyed by ColorType enum
+     * @param colorTypeId the ID of the color type to assert
+     * @param isSelected Whether that color type should be selected
+     */
+    private fun assertColorTypeTabUiState(
+        colorTypes: Map<ColorType, ColorTypeTabViewModel>?,
+        colorTypeId: ColorType,
+        isSelected: Boolean,
+    ) {
+        val viewModel =
+            colorTypes?.get(colorTypeId) ?: error("No color type with ID \"$colorTypeId\"!")
+        assertThat(viewModel.isSelected).isEqualTo(isSelected)
+    }
+}
diff --git a/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt b/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
index 9a2a0af6..fea94dc8 100644
--- a/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
+++ b/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
@@ -24,6 +24,7 @@ import com.android.customization.picker.quickaffordance.domain.interactor.Keygua
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSelectionModel
 import com.android.systemui.shared.customization.data.content.FakeCustomizationProviderClient
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
+import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.Dispatchers
@@ -69,7 +70,7 @@ class KeyguardQuickAffordancePickerInteractorTest {
                             interactor = underTest,
                             client = client,
                         )
-                        .apply { runBlocking { setUpSnapshotRestorer {} } }
+                        .apply { runBlocking { setUpSnapshotRestorer(FakeSnapshotStore()) } }
                 },
             )
     }
@@ -114,23 +115,6 @@ class KeyguardQuickAffordancePickerInteractorTest {
         }
 
     @Test
-    fun unselect() =
-        testScope.runTest {
-            val selections = collectLastValue(underTest.selections)
-            underTest.select(
-                slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
-                affordanceId = FakeCustomizationProviderClient.AFFORDANCE_1,
-            )
-
-            underTest.unselect(
-                slotId = KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START,
-                affordanceId = FakeCustomizationProviderClient.AFFORDANCE_1,
-            )
-
-            assertThat(selections()).isEmpty()
-        }
-
-    @Test
     fun unselectAll() =
         testScope.runTest {
             client.setSlotCapacity(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END, 3)
diff --git a/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt b/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
index d1214c1a..103ae84d 100644
--- a/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
+++ b/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
@@ -18,6 +18,7 @@
 package com.android.customization.model.picker.quickaffordance.ui.viewmodel
 
 import android.content.Context
+import android.content.Intent
 import androidx.test.filters.SmallTest
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
@@ -26,14 +27,15 @@ import com.android.customization.picker.quickaffordance.domain.interactor.Keygua
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordancePickerViewModel
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordanceSlotViewModel
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordanceSummaryViewModel
-import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordanceViewModel
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
 import com.android.systemui.shared.customization.data.content.FakeCustomizationProviderClient
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
+import com.android.wallpaper.R
 import com.android.wallpaper.module.InjectorProvider
-import com.android.wallpaper.picker.undo.data.repository.UndoRepository
-import com.android.wallpaper.picker.undo.domain.interactor.UndoInteractor
-import com.android.wallpaper.testing.FAKE_RESTORERS
+import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
+import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.TestCurrentWallpaperInfoFactory
 import com.android.wallpaper.testing.TestInjector
 import com.android.wallpaper.testing.collectLastValue
@@ -65,6 +67,8 @@ class KeyguardQuickAffordancePickerViewModelTest {
     private lateinit var client: FakeCustomizationProviderClient
     private lateinit var quickAffordanceInteractor: KeyguardQuickAffordancePickerInteractor
 
+    private var latestStartedActivityIntent: Intent? = null
+
     @Before
     fun setUp() {
         InjectorProvider.setInjector(TestInjector())
@@ -87,21 +91,15 @@ class KeyguardQuickAffordancePickerViewModelTest {
                             interactor = quickAffordanceInteractor,
                             client = client,
                         )
-                        .apply { runBlocking { setUpSnapshotRestorer {} } }
+                        .apply { runBlocking { setUpSnapshotRestorer(FakeSnapshotStore()) } }
                 },
             )
-        val undoInteractor =
-            UndoInteractor(
-                scope = testScope.backgroundScope,
-                repository = UndoRepository(),
-                restorerByOwnerId = FAKE_RESTORERS,
-            )
         underTest =
             KeyguardQuickAffordancePickerViewModel.Factory(
                     context = context,
                     quickAffordanceInteractor = quickAffordanceInteractor,
-                    undoInteractor = undoInteractor,
                     wallpaperInfoFactory = TestCurrentWallpaperInfoFactory(context),
+                    activityStarter = { intent -> latestStartedActivityIntent = intent },
                 )
                 .create(KeyguardQuickAffordancePickerViewModel::class.java)
     }
@@ -134,7 +132,7 @@ class KeyguardQuickAffordancePickerViewModelTest {
             )
 
             // Select "affordance 1" for the first slot.
-            quickAffordances()?.get(1)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 1)
             assertPickerUiState(
                 slots = slots(),
                 affordances = quickAffordances(),
@@ -155,7 +153,7 @@ class KeyguardQuickAffordancePickerViewModelTest {
             // First, switch to the second slot:
             slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END)?.onClicked?.invoke()
             // Second, select the "affordance 3" affordance:
-            quickAffordances()?.get(3)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 3)
             assertPickerUiState(
                 slots = slots(),
                 affordances = quickAffordances(),
@@ -174,7 +172,7 @@ class KeyguardQuickAffordancePickerViewModelTest {
             )
 
             // Select a different affordance for the second slot.
-            quickAffordances()?.get(2)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 2)
             assertPickerUiState(
                 slots = slots(),
                 affordances = quickAffordances(),
@@ -200,17 +198,17 @@ class KeyguardQuickAffordancePickerViewModelTest {
             val quickAffordances = collectLastValue(underTest.quickAffordances)
 
             // Select "affordance 1" for the first slot.
-            quickAffordances()?.get(1)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 1)
             // Select an affordance for the second slot.
             // First, switch to the second slot:
             slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END)?.onClicked?.invoke()
             // Second, select the "affordance 3" affordance:
-            quickAffordances()?.get(3)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 3)
 
             // Switch back to the first slot:
             slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START)?.onClicked?.invoke()
             // Select the "none" affordance, which is always in position 0:
-            quickAffordances()?.get(0)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 0)
 
             assertPickerUiState(
                 slots = slots(),
@@ -256,14 +254,17 @@ class KeyguardQuickAffordancePickerViewModelTest {
                 )
 
             // Lets try to select that disabled affordance:
-            quickAffordances()?.get(affordanceIndex + 1)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, affordanceIndex + 1)
 
             // We expect there to be a dialog that should be shown:
-            assertThat(dialog()?.icon).isEqualTo(FakeCustomizationProviderClient.ICON_1)
-            assertThat(dialog()?.instructions).isEqualTo(enablementInstructions)
-            assertThat(dialog()?.actionText).isEqualTo(enablementActionText)
-            assertThat(dialog()?.intent?.`package`).isEqualTo(packageName)
-            assertThat(dialog()?.intent?.action).isEqualTo(action)
+            assertThat(dialog()?.icon)
+                .isEqualTo(Icon.Loaded(FakeCustomizationProviderClient.ICON_1, null))
+            assertThat(dialog()?.title).isEqualTo(Text.Loaded("disabled"))
+            assertThat(dialog()?.message)
+                .isEqualTo(Text.Loaded(enablementInstructions.joinToString("\n")))
+            assertThat(dialog()?.buttons?.size).isEqualTo(1)
+            assertThat(dialog()?.buttons?.first()?.text)
+                .isEqualTo(Text.Loaded(enablementActionText))
 
             // Once we report that the dialog has been dismissed by the user, we expect there to be
             // no
@@ -273,6 +274,30 @@ class KeyguardQuickAffordancePickerViewModelTest {
         }
 
     @Test
+    fun `Start settings activity when long-pressing an affordance`() =
+        testScope.runTest {
+            val quickAffordances = collectLastValue(underTest.quickAffordances)
+
+            // Lets add a configurable affordance to the picker:
+            val configureIntent = Intent("some.action")
+            val affordanceIndex =
+                client.addAffordance(
+                    CustomizationProviderClient.Affordance(
+                        id = "affordance",
+                        name = "affordance",
+                        iconResourceId = 1,
+                        isEnabled = true,
+                        configureIntent = configureIntent,
+                    )
+                )
+
+            // Lets try to long-click the affordance:
+            quickAffordances()?.get(affordanceIndex + 1)?.onLongClicked?.invoke()
+
+            assertThat(latestStartedActivityIntent).isEqualTo(configureIntent)
+        }
+
+    @Test
     fun `summary - affordance selected in both bottom-start and bottom-end`() =
         testScope.runTest {
             val slots = collectLastValue(underTest.slots)
@@ -280,21 +305,23 @@ class KeyguardQuickAffordancePickerViewModelTest {
             val summary = collectLastValue(underTest.summary)
 
             // Select "affordance 1" for the first slot.
-            quickAffordances()?.get(1)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 1)
             // Select an affordance for the second slot.
             // First, switch to the second slot:
             slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END)?.onClicked?.invoke()
             // Second, select the "affordance 3" affordance:
-            quickAffordances()?.get(3)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 3)
 
             assertThat(summary())
                 .isEqualTo(
                     KeyguardQuickAffordanceSummaryViewModel(
                         description =
-                            "${FakeCustomizationProviderClient.AFFORDANCE_1}," +
-                                " ${FakeCustomizationProviderClient.AFFORDANCE_3}",
-                        icon1 = FakeCustomizationProviderClient.ICON_1,
-                        icon2 = FakeCustomizationProviderClient.ICON_3,
+                            Text.Loaded(
+                                "${FakeCustomizationProviderClient.AFFORDANCE_1}," +
+                                    " ${FakeCustomizationProviderClient.AFFORDANCE_3}"
+                            ),
+                        icon1 = Icon.Loaded(FakeCustomizationProviderClient.ICON_1, null),
+                        icon2 = Icon.Loaded(FakeCustomizationProviderClient.ICON_3, null),
                     )
                 )
         }
@@ -307,13 +334,13 @@ class KeyguardQuickAffordancePickerViewModelTest {
             val summary = collectLastValue(underTest.summary)
 
             // Select "affordance 1" for the first slot.
-            quickAffordances()?.get(1)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 1)
 
             assertThat(summary())
                 .isEqualTo(
                     KeyguardQuickAffordanceSummaryViewModel(
-                        description = FakeCustomizationProviderClient.AFFORDANCE_1,
-                        icon1 = FakeCustomizationProviderClient.ICON_1,
+                        description = Text.Loaded(FakeCustomizationProviderClient.AFFORDANCE_1),
+                        icon1 = Icon.Loaded(FakeCustomizationProviderClient.ICON_1, null),
                         icon2 = null,
                     )
                 )
@@ -330,14 +357,14 @@ class KeyguardQuickAffordancePickerViewModelTest {
             // First, switch to the second slot:
             slots()?.get(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END)?.onClicked?.invoke()
             // Second, select the "affordance 3" affordance:
-            quickAffordances()?.get(3)?.onClicked?.invoke()
+            selectAffordance(quickAffordances, 3)
 
             assertThat(summary())
                 .isEqualTo(
                     KeyguardQuickAffordanceSummaryViewModel(
-                        description = FakeCustomizationProviderClient.AFFORDANCE_3,
+                        description = Text.Loaded(FakeCustomizationProviderClient.AFFORDANCE_3),
                         icon1 = null,
-                        icon2 = FakeCustomizationProviderClient.ICON_3,
+                        icon2 = Icon.Loaded(FakeCustomizationProviderClient.ICON_3, null),
                     )
                 )
         }
@@ -345,15 +372,28 @@ class KeyguardQuickAffordancePickerViewModelTest {
     @Test
     fun `summary - no affordances selected`() =
         testScope.runTest {
-            val slots = collectLastValue(underTest.slots)
-            val quickAffordances = collectLastValue(underTest.quickAffordances)
             val summary = collectLastValue(underTest.summary)
 
-            assertThat(summary()?.description).isEqualTo("None")
+            assertThat(summary()?.description)
+                .isEqualTo(Text.Resource(R.string.keyguard_quick_affordance_none_selected))
             assertThat(summary()?.icon1).isNotNull()
             assertThat(summary()?.icon2).isNull()
         }
 
+    /** Simulates a user selecting the affordance at the given index, if that is clickable. */
+    private fun TestScope.selectAffordance(
+        affordances: () -> List<OptionItemViewModel<Icon>>?,
+        index: Int,
+    ) {
+        val onClickedFlow = affordances()?.get(index)?.onClicked
+        val onClickedLastValueOrNull: (() -> (() -> Unit)?)? =
+            onClickedFlow?.let { collectLastValue(it) }
+        onClickedLastValueOrNull?.let { onClickedLastValue ->
+            val onClickedOrNull: (() -> Unit)? = onClickedLastValue()
+            onClickedOrNull?.let { onClicked -> onClicked() }
+        }
+    }
+
     /**
      * Asserts the entire picker UI state is what is expected. This includes the slot tabs and the
      * affordance list.
@@ -363,9 +403,9 @@ class KeyguardQuickAffordancePickerViewModelTest {
      * @param selectedSlotText The text of the slot that's expected to be selected
      * @param selectedAffordanceText The text of the affordance that's expected to be selected
      */
-    private fun assertPickerUiState(
+    private fun TestScope.assertPickerUiState(
         slots: Map<String, KeyguardQuickAffordanceSlotViewModel>?,
-        affordances: List<KeyguardQuickAffordanceViewModel>?,
+        affordances: List<OptionItemViewModel<Icon>>?,
         selectedSlotText: String,
         selectedAffordanceText: String,
     ) {
@@ -383,12 +423,18 @@ class KeyguardQuickAffordancePickerViewModelTest {
         var foundSelectedAffordance = false
         assertThat(affordances).isNotNull()
         affordances?.forEach { affordance ->
-            val nameMatchesSelectedName = affordance.contentDescription == selectedAffordanceText
+            val nameMatchesSelectedName =
+                Text.evaluationEquals(
+                    context,
+                    affordance.text,
+                    Text.Loaded(selectedAffordanceText),
+                )
+            val isSelected: Boolean? = collectLastValue(affordance.isSelected).invoke()
             assertWithMessage(
-                    "Expected affordance with name \"${affordance.contentDescription}\" to have" +
-                        " isSelected=$nameMatchesSelectedName but it was ${affordance.isSelected}"
+                    "Expected affordance with name \"${affordance.text}\" to have" +
+                        " isSelected=$nameMatchesSelectedName but it was $isSelected"
                 )
-                .that(affordance.isSelected)
+                .that(isSelected)
                 .isEqualTo(nameMatchesSelectedName)
             foundSelectedAffordance = foundSelectedAffordance || nameMatchesSelectedName
         }
@@ -416,7 +462,8 @@ class KeyguardQuickAffordancePickerViewModelTest {
      *
      * @param slots The observed slot view-models, keyed by slot ID
      * @param expectedAffordanceNameBySlotId The expected name of the selected affordance for each
-     * slot ID or `null` if it's expected for there to be no affordance for that slot in the preview
+     *   slot ID or `null` if it's expected for there to be no affordance for that slot in the
+     *   preview
      */
     private fun assertPreviewUiState(
         slots: Map<String, KeyguardQuickAffordanceSlotViewModel>?,
@@ -425,13 +472,12 @@ class KeyguardQuickAffordancePickerViewModelTest {
         assertThat(slots).isNotNull()
         slots?.forEach { (slotId, slotViewModel) ->
             val expectedAffordanceName = expectedAffordanceNameBySlotId[slotId]
-            val actualAffordanceName =
-                slotViewModel.selectedQuickAffordances.firstOrNull()?.contentDescription
+            val actualAffordanceName = slotViewModel.selectedQuickAffordances.firstOrNull()?.text
             assertWithMessage(
                     "At slotId=\"$slotId\", expected affordance=\"$expectedAffordanceName\" but" +
-                        " was \"$actualAffordanceName\"!"
+                        " was \"${actualAffordanceName?.asString(context)}\"!"
                 )
-                .that(actualAffordanceName)
+                .that(actualAffordanceName?.asString(context))
                 .isEqualTo(expectedAffordanceName)
         }
     }
diff --git a/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt b/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt
new file mode 100644
index 00000000..e6e30c31
--- /dev/null
+++ b/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.themedicon.domain.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.customization.model.themedicon.data.repository.ThemeIconRepository
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ThemedIconInteractorTest {
+
+    private lateinit var underTest: ThemedIconInteractor
+
+    @Before
+    fun setUp() {
+        underTest =
+            ThemedIconInteractor(
+                repository = ThemeIconRepository(),
+            )
+    }
+
+    @Test
+    fun `end-to-end`() = runTest {
+        val isActivated = collectLastValue(underTest.isActivated)
+
+        underTest.setActivated(isActivated = true)
+        assertThat(isActivated()).isTrue()
+
+        underTest.setActivated(isActivated = false)
+        assertThat(isActivated()).isFalse()
+    }
+}
diff --git a/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt b/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt
new file mode 100644
index 00000000..df1fd201
--- /dev/null
+++ b/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt
@@ -0,0 +1,102 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.themedicon.domain.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.customization.model.themedicon.data.repository.ThemeIconRepository
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ThemedIconSnapshotRestorerTest {
+
+    private lateinit var underTest: ThemedIconSnapshotRestorer
+    private var isActivated = false
+
+    @Before
+    fun setUp() {
+        isActivated = false
+        underTest =
+            ThemedIconSnapshotRestorer(
+                isActivated = { isActivated },
+                setActivated = { isActivated = it },
+                interactor =
+                    ThemedIconInteractor(
+                        repository = ThemeIconRepository(),
+                    )
+            )
+    }
+
+    @Test
+    fun `set up and restore - active`() = runTest {
+        isActivated = true
+
+        val store = FakeSnapshotStore()
+        store.store(underTest.setUpSnapshotRestorer(store = store))
+        val storedSnapshot = store.retrieve()
+
+        underTest.restoreToSnapshot(snapshot = storedSnapshot)
+        assertThat(isActivated).isTrue()
+    }
+
+    @Test
+    fun `set up and restore - inactive`() = runTest {
+        isActivated = false
+
+        val store = FakeSnapshotStore()
+        store.store(underTest.setUpSnapshotRestorer(store = store))
+        val storedSnapshot = store.retrieve()
+
+        underTest.restoreToSnapshot(snapshot = storedSnapshot)
+        assertThat(isActivated).isFalse()
+    }
+
+    @Test
+    fun `set up - deactivate - restore to active`() = runTest {
+        isActivated = true
+        val store = FakeSnapshotStore()
+        store.store(underTest.setUpSnapshotRestorer(store = store))
+        val initialSnapshot = store.retrieve()
+
+        underTest.store(isActivated = false)
+
+        underTest.restoreToSnapshot(snapshot = initialSnapshot)
+        assertThat(isActivated).isTrue()
+    }
+
+    @Test
+    fun `set up - activate - restore to inactive`() = runTest {
+        isActivated = false
+        val store = FakeSnapshotStore()
+        store.store(underTest.setUpSnapshotRestorer(store = store))
+        val initialSnapshot = store.retrieve()
+
+        underTest.store(isActivated = true)
+
+        underTest.restoreToSnapshot(snapshot = initialSnapshot)
+        assertThat(isActivated).isFalse()
+    }
+}
diff --git a/tests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt b/tests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
new file mode 100644
index 00000000..2ef4e974
--- /dev/null
+++ b/tests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
@@ -0,0 +1,89 @@
+/*
+ * 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(
+                selectedClock.clockId,
+                selectedClock.name,
+                selectedColor,
+                colorTone,
+                seedColor,
+            )
+        }
+
+    private val _selectedClockSize = MutableStateFlow(ClockSize.SMALL)
+    override val selectedClockSize: Flow<ClockSize> = _selectedClockSize.asStateFlow()
+
+    override fun setSelectedClock(clockId: String) {
+        selectedClockId.value = clockId
+    }
+
+    override 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 {
+        val fakeClocks =
+            listOf(
+                ClockMetadataModel("clock0", "clock0", null, 50, null),
+                ClockMetadataModel("clock1", "clock1", null, 50, null),
+                ClockMetadataModel("clock2", "clock2", null, 50, null),
+                ClockMetadataModel("clock3", "clock3", 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/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt b/tests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
new file mode 100644
index 00000000..cd41d7d0
--- /dev/null
+++ b/tests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
@@ -0,0 +1,73 @@
+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.collectLastValue
+import com.google.common.truth.Truth
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+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.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ClockPickerInteractorTest {
+
+    private lateinit var underTest: ClockPickerInteractor
+
+    @Before
+    fun setUp() {
+        val testDispatcher = StandardTestDispatcher()
+        Dispatchers.setMain(testDispatcher)
+        underTest = ClockPickerInteractor(FakeClockPickerRepository())
+    }
+
+    @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/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt b/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
new file mode 100644
index 00000000..63f77bd6
--- /dev/null
+++ b/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
@@ -0,0 +1,120 @@
+/*
+ * 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.FakeClockPickerRepository
+import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
+import com.android.customization.picker.clock.shared.model.ClockMetadataModel
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+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.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ClockCarouselViewModelTest {
+    private val repositoryWithMultipleClocks by lazy { FakeClockPickerRepository() }
+    private val repositoryWithSingleClock by lazy {
+        FakeClockPickerRepository(
+            listOf(
+                ClockMetadataModel(
+                    "clock0",
+                    "clock0",
+                    null,
+                    ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS,
+                    null,
+                ),
+            )
+        )
+    }
+    private lateinit var underTest: ClockCarouselViewModel
+
+    @Before
+    fun setUp() {
+        val testDispatcher = StandardTestDispatcher()
+        Dispatchers.setMain(testDispatcher)
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun setSelectedClock() = runTest {
+        underTest = ClockCarouselViewModel(ClockPickerInteractor(repositoryWithMultipleClocks))
+        val observedSelectedIndex = collectLastValue(underTest.selectedIndex)
+        advanceTimeBy(ClockCarouselViewModel.CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
+        underTest.setSelectedClock(FakeClockPickerRepository.fakeClocks[2].clockId)
+        assertThat(observedSelectedIndex()).isEqualTo(2)
+    }
+
+    @Test
+    fun setShouldShowCarousel() = runTest {
+        underTest = ClockCarouselViewModel(ClockPickerInteractor(repositoryWithMultipleClocks))
+        val observedIsCarouselVisible = collectLastValue(underTest.isCarouselVisible)
+        advanceTimeBy(ClockCarouselViewModel.CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
+        underTest.showClockCarousel(false)
+        assertThat(observedIsCarouselVisible()).isFalse()
+        underTest.showClockCarousel(true)
+        assertThat(observedIsCarouselVisible()).isTrue()
+    }
+
+    @Test
+    fun shouldNotShowCarouselWhenSingleClock() = runTest {
+        underTest = ClockCarouselViewModel(ClockPickerInteractor(repositoryWithSingleClock))
+        val observedIsCarouselVisible = collectLastValue(underTest.isCarouselVisible)
+        advanceTimeBy(ClockCarouselViewModel.CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
+        underTest.showClockCarousel(false)
+        assertThat(observedIsCarouselVisible()).isFalse()
+        underTest.showClockCarousel(true)
+        assertThat(observedIsCarouselVisible()).isFalse()
+    }
+
+    @Test
+    fun setShouldShowSingleClock() = runTest {
+        underTest = ClockCarouselViewModel(ClockPickerInteractor(repositoryWithSingleClock))
+        val observedIsSingleClockViewVisible = collectLastValue(underTest.isSingleClockViewVisible)
+        advanceTimeBy(ClockCarouselViewModel.CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
+        underTest.showClockCarousel(false)
+        assertThat(observedIsSingleClockViewVisible()).isFalse()
+        underTest.showClockCarousel(true)
+        assertThat(observedIsSingleClockViewVisible()).isTrue()
+    }
+
+    @Test
+    fun shouldNotShowSingleClockWhenMultipleClocks() = runTest {
+        underTest = ClockCarouselViewModel(ClockPickerInteractor(repositoryWithMultipleClocks))
+        val observedIsSingleClockViewVisible = collectLastValue(underTest.isSingleClockViewVisible)
+        advanceTimeBy(ClockCarouselViewModel.CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
+        underTest.showClockCarousel(false)
+        assertThat(observedIsSingleClockViewVisible()).isFalse()
+        underTest.showClockCarousel(true)
+        assertThat(observedIsSingleClockViewVisible()).isFalse()
+    }
+}
diff --git a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt b/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
new file mode 100644
index 00000000..61976ad6
--- /dev/null
+++ b/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
@@ -0,0 +1,82 @@
+/*
+ * 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.shared.ClockSize
+import com.android.customization.picker.clock.shared.model.ClockMetadataModel
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+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.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::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(FakeClockPickerRepository())
+        underTest =
+            ClockSectionViewModel(
+                context,
+                interactor,
+            )
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun setSelectedClock() = runTest {
+        val colorRed = clockColorMap.values.first()
+        val observedSelectedClockColorAndSizeText =
+            collectLastValue(underTest.selectedClockColorAndSizeText)
+        interactor.setClockColor(
+            colorRed.colorId,
+            ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS,
+            ClockSettingsViewModel.blendColorWithTone(
+                colorRed.color,
+                colorRed.getColorTone(ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS),
+            )
+        )
+        interactor.setClockSize(ClockSize.DYNAMIC)
+        assertThat(observedSelectedClockColorAndSizeText()).isEqualTo("Red, dynamic")
+    }
+}
diff --git a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt b/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
new file mode 100644
index 00000000..d53288d0
--- /dev/null
+++ b/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
@@ -0,0 +1,156 @@
+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.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.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.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ClockSettingsViewModelTest {
+
+    private lateinit var context: Context
+    private lateinit var colorPickerInteractor: ColorPickerInteractor
+    private lateinit var underTest: ClockSettingsViewModel
+    private lateinit var colorMap: Map<String, ClockColorViewModel>
+
+    @Before
+    fun setUp() {
+        val testDispatcher = StandardTestDispatcher()
+        Dispatchers.setMain(testDispatcher)
+        context = InstrumentationRegistry.getInstrumentation().targetContext
+        colorPickerInteractor =
+            ColorPickerInteractor(
+                repository = FakeColorPickerRepository(context = context),
+                snapshotRestorer = {
+                    ColorPickerSnapshotRestorer(interactor = colorPickerInteractor).apply {
+                        runBlocking { setUpSnapshotRestorer(store = FakeSnapshotStore()) }
+                    }
+                },
+            )
+        underTest =
+            ClockSettingsViewModel.Factory(
+                    context = context,
+                    clockPickerInteractor = ClockPickerInteractor(FakeClockPickerRepository()),
+                    colorPickerInteractor = colorPickerInteractor,
+                )
+                .create(ClockSettingsViewModel::class.java)
+        colorMap = ClockColorViewModel.getPresetColorMap(context.resources)
+    }
+
+    @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)
+        assertThat(observedClockColorOptions()!![0].isSelected).isTrue()
+        assertThat(observedClockColorOptions()!![0].onClick).isNull()
+        assertThat(observedSelectedColorOptionPosition()).isEqualTo(0)
+
+        observedClockColorOptions()!![1].onClick?.invoke()
+        // Advance COLOR_OPTIONS_EVENT_UPDATE_DELAY_MILLIS since there is a delay from colorOptions
+        advanceTimeBy(ClockSettingsViewModel.COLOR_OPTIONS_EVENT_UPDATE_DELAY_MILLIS)
+        assertThat(observedClockColorOptions()!![1].isSelected).isTrue()
+        assertThat(observedClockColorOptions()!![1].onClick).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)
+        assertThat(observedClockColorOptions()!![0].isSelected).isTrue()
+        assertThat(observedIsSliderEnabled()).isFalse()
+
+        observedClockColorOptions()!![1].onClick?.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
+        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)
+    }
+}
diff --git a/tests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt b/tests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt
new file mode 100644
index 00000000..be799db3
--- /dev/null
+++ b/tests/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/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt b/tests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
new file mode 100644
index 00000000..64426094
--- /dev/null
+++ b/tests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
@@ -0,0 +1,101 @@
+/*
+ * 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.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::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 subtitleStringResId = collectLastValue(underTest.subtitleStringResourceId)
+            val isSwitchOn = collectLastValue(underTest.isSwitchOn)
+
+            val initialSubtitleStringRes = subtitleStringResId()
+            val initialIsSwitchOn = isSwitchOn()
+
+            underTest.onClicked()
+            assertThat(subtitleStringResId()).isNotEqualTo(initialSubtitleStringRes)
+            assertThat(isSwitchOn()).isNotEqualTo(initialIsSwitchOn)
+
+            underTest.onClicked()
+            assertThat(subtitleStringResId()).isEqualTo(initialSubtitleStringRes)
+            assertThat(isSwitchOn()).isEqualTo(initialIsSwitchOn)
+        }
+}
diff --git a/tests/src/com/android/customization/testing/TestCustomizationInjector.java b/tests/src/com/android/customization/testing/TestCustomizationInjector.java
deleted file mode 100644
index d6093356..00000000
--- a/tests/src/com/android/customization/testing/TestCustomizationInjector.java
+++ /dev/null
@@ -1,143 +0,0 @@
-package com.android.customization.testing;
-
-import android.content.Context;
-
-import androidx.fragment.app.FragmentActivity;
-
-import com.android.customization.model.theme.OverlayManagerCompat;
-import com.android.customization.model.theme.ThemeBundleProvider;
-import com.android.customization.model.theme.ThemeManager;
-import com.android.customization.module.CustomizationInjector;
-import com.android.customization.module.CustomizationPreferences;
-import com.android.customization.module.ThemesUserEventLogger;
-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.systemui.shared.customization.data.content.CustomizationProviderClient;
-import com.android.systemui.shared.customization.data.content.CustomizationProviderClientImpl;
-import com.android.wallpaper.config.BaseFlags;
-import com.android.wallpaper.module.DrawableLayerResolver;
-import com.android.wallpaper.module.PackageStatusNotifier;
-import com.android.wallpaper.module.UserEventLogger;
-import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer;
-import com.android.wallpaper.testing.TestInjector;
-
-import java.util.HashMap;
-import java.util.Map;
-
-import kotlinx.coroutines.Dispatchers;
-
-/**
- * Test implementation of the dependency injector.
- */
-public class TestCustomizationInjector extends TestInjector implements CustomizationInjector {
-    private CustomizationPreferences mCustomizationPreferences;
-    private ThemeManager mThemeManager;
-    private PackageStatusNotifier mPackageStatusNotifier;
-    private DrawableLayerResolver mDrawableLayerResolver;
-    private UserEventLogger mUserEventLogger;
-    private KeyguardQuickAffordancePickerInteractor mKeyguardQuickAffordancePickerInteractor;
-    private BaseFlags mFlags;
-    private CustomizationProviderClient mCustomizationProviderClient;
-    private KeyguardQuickAffordanceSnapshotRestorer mKeyguardQuickAffordanceSnapshotRestorer;
-
-    @Override
-    public CustomizationPreferences getCustomizationPreferences(Context context) {
-        if (mCustomizationPreferences == null) {
-            mCustomizationPreferences = new TestDefaultCustomizationPreferences(context);
-        }
-        return mCustomizationPreferences;
-    }
-
-    @Override
-    public ThemeManager getThemeManager(
-            ThemeBundleProvider provider,
-            FragmentActivity activity,
-            OverlayManagerCompat overlayManagerCompat,
-            ThemesUserEventLogger logger) {
-        if (mThemeManager == null) {
-            mThemeManager = new TestThemeManager(provider, activity, overlayManagerCompat, logger);
-        }
-        return mThemeManager;
-    }
-
-    @Override
-    public PackageStatusNotifier getPackageStatusNotifier(Context context) {
-        if (mPackageStatusNotifier == null) {
-            mPackageStatusNotifier =  new TestPackageStatusNotifier();
-        }
-        return mPackageStatusNotifier;
-    }
-
-    @Override
-    public DrawableLayerResolver getDrawableLayerResolver() {
-        if (mDrawableLayerResolver == null) {
-            mDrawableLayerResolver = new TestDrawableLayerResolver();
-        }
-        return mDrawableLayerResolver;
-    }
-
-    @Override
-    public UserEventLogger getUserEventLogger(Context unused) {
-        if (mUserEventLogger == null) {
-            mUserEventLogger = new TestThemesUserEventLogger();
-        }
-        return mUserEventLogger;
-    }
-
-    @Override
-    public KeyguardQuickAffordancePickerInteractor getKeyguardQuickAffordancePickerInteractor(
-            Context context) {
-        if (mKeyguardQuickAffordancePickerInteractor == null) {
-            final CustomizationProviderClient client =
-                    new CustomizationProviderClientImpl(context, Dispatchers.getIO());
-            mKeyguardQuickAffordancePickerInteractor = new KeyguardQuickAffordancePickerInteractor(
-                    new KeyguardQuickAffordancePickerRepository(client, Dispatchers.getIO()),
-                    client,
-                    () -> getKeyguardQuickAffordanceSnapshotRestorer(context));
-        }
-        return mKeyguardQuickAffordancePickerInteractor;
-    }
-
-    @Override
-    public BaseFlags getFlags() {
-        if (mFlags == null) {
-            mFlags = new BaseFlags() {};
-        }
-
-        return mFlags;
-    }
-
-    @Override
-    public Map<Integer, SnapshotRestorer> getSnapshotRestorers(Context context) {
-        final Map<Integer, SnapshotRestorer> restorers = new HashMap<>();
-        restorers.put(
-                KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER,
-                getKeyguardQuickAffordanceSnapshotRestorer(context));
-        return restorers;
-    }
-
-    /** Returns the {@link CustomizationProviderClient}. */
-    private CustomizationProviderClient getKeyguardQuickAffordancePickerProviderClient(
-            Context context) {
-        if (mCustomizationProviderClient == null) {
-            mCustomizationProviderClient =
-                    new CustomizationProviderClientImpl(context, Dispatchers.getIO());
-        }
-
-        return mCustomizationProviderClient;
-    }
-
-    private KeyguardQuickAffordanceSnapshotRestorer getKeyguardQuickAffordanceSnapshotRestorer(
-            Context context) {
-        if (mKeyguardQuickAffordanceSnapshotRestorer == null) {
-            mKeyguardQuickAffordanceSnapshotRestorer = new KeyguardQuickAffordanceSnapshotRestorer(
-                    getKeyguardQuickAffordancePickerInteractor(context),
-                    getKeyguardQuickAffordancePickerProviderClient(context));
-        }
-
-        return mKeyguardQuickAffordanceSnapshotRestorer;
-    }
-
-    private static final int KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER = 1;
-}
diff --git a/tests/src/com/android/customization/testing/TestCustomizationInjector.kt b/tests/src/com/android/customization/testing/TestCustomizationInjector.kt
new file mode 100644
index 00000000..b49e6546
--- /dev/null
+++ b/tests/src/com/android/customization/testing/TestCustomizationInjector.kt
@@ -0,0 +1,289 @@
+package com.android.customization.testing
+
+import android.app.Activity
+import android.content.Context
+import android.text.TextUtils
+import androidx.fragment.app.FragmentActivity
+import com.android.customization.model.color.ColorCustomizationManager
+import com.android.customization.model.color.ColorOptionsProvider
+import com.android.customization.model.theme.OverlayManagerCompat
+import com.android.customization.model.theme.ThemeBundleProvider
+import com.android.customization.model.theme.ThemeManager
+import com.android.customization.module.CustomizationInjector
+import com.android.customization.module.CustomizationPreferences
+import com.android.customization.module.ThemesUserEventLogger
+import com.android.customization.picker.clock.data.repository.ClockRegistryProvider
+import com.android.customization.picker.clock.data.repository.FakeClockPickerRepository
+import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
+import com.android.customization.picker.clock.ui.view.ClockViewFactory
+import com.android.customization.picker.clock.ui.viewmodel.ClockCarouselViewModel
+import com.android.customization.picker.clock.ui.viewmodel.ClockSectionViewModel
+import com.android.customization.picker.clock.ui.viewmodel.ClockSettingsViewModel
+import com.android.customization.picker.color.data.repository.ColorPickerRepository
+import com.android.customization.picker.color.data.repository.ColorPickerRepositoryImpl
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
+import com.android.customization.picker.color.domain.interactor.ColorPickerSnapshotRestorer
+import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
+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.systemui.shared.clocks.ClockRegistry
+import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
+import com.android.systemui.shared.customization.data.content.CustomizationProviderClientImpl
+import com.android.wallpaper.config.BaseFlags
+import com.android.wallpaper.model.WallpaperColorsViewModel
+import com.android.wallpaper.module.DrawableLayerResolver
+import com.android.wallpaper.module.PackageStatusNotifier
+import com.android.wallpaper.module.UserEventLogger
+import com.android.wallpaper.picker.customization.data.content.WallpaperClientImpl
+import com.android.wallpaper.picker.customization.data.repository.WallpaperRepository
+import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.testing.TestInjector
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.GlobalScope
+
+/** Test implementation of the dependency injector. */
+class TestCustomizationInjector : TestInjector(), CustomizationInjector {
+    private var customizationPreferences: CustomizationPreferences? = null
+    private var themeManager: ThemeManager? = null
+    private var packageStatusNotifier: PackageStatusNotifier? = null
+    private var drawableLayerResolver: DrawableLayerResolver? = null
+    private var userEventLogger: UserEventLogger? = null
+    private var wallpaperInteractor: WallpaperInteractor? = null
+    private var keyguardQuickAffordancePickerInteractor: KeyguardQuickAffordancePickerInteractor? =
+        null
+    private var flags: BaseFlags? = null
+    private var customizationProviderClient: CustomizationProviderClient? = null
+    private var keyguardQuickAffordanceSnapshotRestorer: KeyguardQuickAffordanceSnapshotRestorer? =
+        null
+    private var clockRegistry: ClockRegistry? = null
+    private var clockPickerInteractor: ClockPickerInteractor? = null
+    private var clockSectionViewModel: ClockSectionViewModel? = null
+    private var clockViewFactory: ClockViewFactory? = null
+    private var colorPickerRepository: ColorPickerRepository? = null
+    private var colorPickerInteractor: ColorPickerInteractor? = null
+    private var colorPickerViewModelFactory: ColorPickerViewModel.Factory? = null
+    private var colorPickerSnapshotRestorer: ColorPickerSnapshotRestorer? = null
+    private var colorCustomizationManager: ColorCustomizationManager? = null
+    private var clockCarouselViewModel: ClockCarouselViewModel? = null
+    private var clockSettingsViewModelFactory: ClockSettingsViewModel.Factory? = null
+
+    override fun getCustomizationPreferences(context: Context): CustomizationPreferences {
+        return customizationPreferences
+            ?: TestDefaultCustomizationPreferences(context).also { customizationPreferences = it }
+    }
+
+    override fun getThemeManager(
+        provider: ThemeBundleProvider,
+        activity: FragmentActivity,
+        overlayManagerCompat: OverlayManagerCompat,
+        logger: ThemesUserEventLogger
+    ): ThemeManager {
+        return themeManager
+            ?: TestThemeManager(provider, activity, overlayManagerCompat, logger).also {
+                themeManager = it
+            }
+    }
+
+    override fun getPackageStatusNotifier(context: Context): PackageStatusNotifier {
+        return packageStatusNotifier
+            ?: TestPackageStatusNotifier().also {
+                packageStatusNotifier = TestPackageStatusNotifier()
+            }
+    }
+
+    override fun getDrawableLayerResolver(): DrawableLayerResolver {
+        return drawableLayerResolver
+            ?: TestDrawableLayerResolver().also { drawableLayerResolver = it }
+    }
+
+    override fun getUserEventLogger(context: Context): UserEventLogger {
+        return userEventLogger ?: TestThemesUserEventLogger().also { userEventLogger = it }
+    }
+
+    override fun getWallpaperInteractor(context: Context): WallpaperInteractor {
+        return wallpaperInteractor
+            ?: WallpaperInteractor(
+                    repository =
+                        WallpaperRepository(
+                            scope = GlobalScope,
+                            client = WallpaperClientImpl(context = context),
+                            backgroundDispatcher = Dispatchers.IO,
+                        ),
+                    shouldHandleReload = {
+                        TextUtils.equals(
+                            getColorCustomizationManager(context).currentColorSource,
+                            ColorOptionsProvider.COLOR_SOURCE_PRESET
+                        )
+                    }
+                )
+                .also { wallpaperInteractor = it }
+    }
+
+    override fun getKeyguardQuickAffordancePickerInteractor(
+        context: Context
+    ): KeyguardQuickAffordancePickerInteractor {
+        return keyguardQuickAffordancePickerInteractor
+            ?: createCustomizationProviderClient(context).also {
+                keyguardQuickAffordancePickerInteractor = it
+            }
+    }
+
+    private fun createCustomizationProviderClient(
+        context: Context
+    ): KeyguardQuickAffordancePickerInteractor {
+        val client: CustomizationProviderClient =
+            CustomizationProviderClientImpl(context, Dispatchers.IO)
+        return KeyguardQuickAffordancePickerInteractor(
+            KeyguardQuickAffordancePickerRepository(client, Dispatchers.IO),
+            client
+        ) {
+            getKeyguardQuickAffordanceSnapshotRestorer(context)
+        }
+    }
+
+    override fun getFlags(): BaseFlags {
+        return flags ?: object : BaseFlags() {}.also { flags = it }
+    }
+
+    override fun getSnapshotRestorers(context: Context): Map<Int, SnapshotRestorer> {
+        val restorers: MutableMap<Int, SnapshotRestorer> = HashMap()
+        restorers[KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER] =
+            getKeyguardQuickAffordanceSnapshotRestorer(context)
+        restorers[KEY_COLOR_PICKER_SNAPSHOT_RESTORER] =
+            getColorPickerSnapshotRestorer(context, getWallpaperColorsViewModel())
+        return restorers
+    }
+
+    /** Returns the [CustomizationProviderClient]. */
+    private fun getKeyguardQuickAffordancePickerProviderClient(
+        context: Context
+    ): CustomizationProviderClient {
+        return customizationProviderClient
+            ?: CustomizationProviderClientImpl(context, Dispatchers.IO).also {
+                customizationProviderClient = it
+            }
+    }
+
+    private fun getKeyguardQuickAffordanceSnapshotRestorer(
+        context: Context
+    ): KeyguardQuickAffordanceSnapshotRestorer {
+        return keyguardQuickAffordanceSnapshotRestorer
+            ?: KeyguardQuickAffordanceSnapshotRestorer(
+                    getKeyguardQuickAffordancePickerInteractor(context),
+                    getKeyguardQuickAffordancePickerProviderClient(context)
+                )
+                .also { keyguardQuickAffordanceSnapshotRestorer = it }
+    }
+
+    override fun getClockRegistry(context: Context): ClockRegistry {
+        return clockRegistry
+            ?: ClockRegistryProvider(context, GlobalScope, Dispatchers.Main, Dispatchers.IO)
+                .get()
+                .also { clockRegistry = it }
+    }
+
+    override fun getClockPickerInteractor(context: Context): ClockPickerInteractor {
+        return clockPickerInteractor
+            ?: ClockPickerInteractor(FakeClockPickerRepository()).also {
+                clockPickerInteractor = it
+            }
+    }
+
+    override fun getClockSectionViewModel(context: Context): ClockSectionViewModel {
+        return clockSectionViewModel
+            ?: ClockSectionViewModel(context, getClockPickerInteractor(context)).also {
+                clockSectionViewModel = it
+            }
+    }
+
+    private fun getColorPickerRepository(
+        context: Context,
+        wallpaperColorsViewModel: WallpaperColorsViewModel,
+    ): ColorPickerRepository {
+        return colorPickerRepository
+            ?: ColorPickerRepositoryImpl(
+                wallpaperColorsViewModel,
+                getColorCustomizationManager(context)
+            )
+    }
+
+    override fun getColorPickerInteractor(
+        context: Context,
+        wallpaperColorsViewModel: WallpaperColorsViewModel,
+    ): ColorPickerInteractor {
+        return colorPickerInteractor
+            ?: ColorPickerInteractor(
+                    repository = getColorPickerRepository(context, wallpaperColorsViewModel),
+                    snapshotRestorer = {
+                        getColorPickerSnapshotRestorer(context, wallpaperColorsViewModel)
+                    },
+                )
+                .also { colorPickerInteractor = it }
+    }
+
+    override fun getColorPickerViewModelFactory(
+        context: Context,
+        wallpaperColorsViewModel: WallpaperColorsViewModel,
+    ): ColorPickerViewModel.Factory {
+        return colorPickerViewModelFactory
+            ?: ColorPickerViewModel.Factory(
+                    context,
+                    getColorPickerInteractor(context, wallpaperColorsViewModel),
+                )
+                .also { colorPickerViewModelFactory = it }
+    }
+
+    private fun getColorPickerSnapshotRestorer(
+        context: Context,
+        wallpaperColorsViewModel: WallpaperColorsViewModel
+    ): ColorPickerSnapshotRestorer {
+        return colorPickerSnapshotRestorer
+            ?: ColorPickerSnapshotRestorer(
+                    getColorPickerInteractor(context, wallpaperColorsViewModel)
+                )
+                .also { colorPickerSnapshotRestorer = it }
+    }
+
+    private fun getColorCustomizationManager(context: Context): ColorCustomizationManager {
+        return colorCustomizationManager
+            ?: ColorCustomizationManager.getInstance(context, OverlayManagerCompat(context)).also {
+                colorCustomizationManager = it
+            }
+    }
+
+    override fun getClockCarouselViewModel(context: Context): ClockCarouselViewModel {
+        return clockCarouselViewModel
+            ?: ClockCarouselViewModel(getClockPickerInteractor(context)).also {
+                clockCarouselViewModel = it
+            }
+    }
+
+    override fun getClockViewFactory(activity: Activity): ClockViewFactory {
+        return clockViewFactory
+            ?: ClockViewFactory(activity, getClockRegistry(activity)).also { clockViewFactory = it }
+    }
+
+    override fun getClockSettingsViewModelFactory(
+        context: Context,
+        wallpaperColorsViewModel: WallpaperColorsViewModel,
+    ): ClockSettingsViewModel.Factory {
+        return clockSettingsViewModelFactory
+            ?: ClockSettingsViewModel.Factory(
+                    context,
+                    getClockPickerInteractor(context),
+                    getColorPickerInteractor(
+                        context,
+                        wallpaperColorsViewModel,
+                    ),
+                )
+                .also { clockSettingsViewModelFactory = it }
+    }
+
+    companion object {
+        private const val KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER = 1
+        private const val KEY_COLOR_PICKER_SNAPSHOT_RESTORER =
+            KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER + 1
+    }
+}
diff --git a/tests/src/com/android/customization/testing/TestPluginManager.kt b/tests/src/com/android/customization/testing/TestPluginManager.kt
new file mode 100644
index 00000000..167d8ddc
--- /dev/null
+++ b/tests/src/com/android/customization/testing/TestPluginManager.kt
@@ -0,0 +1,36 @@
+package com.android.customization.testing
+
+import com.android.systemui.plugins.Plugin
+import com.android.systemui.plugins.PluginListener
+import com.android.systemui.plugins.PluginManager
+
+class TestPluginManager : PluginManager {
+    override fun getPrivilegedPlugins(): Array<String> {
+        return emptyArray()
+    }
+
+    override fun <T : Plugin?> addPluginListener(listener: PluginListener<T>, cls: Class<T>) {}
+    override fun <T : Plugin?> addPluginListener(
+        listener: PluginListener<T>,
+        cls: Class<T>,
+        allowMultiple: Boolean
+    ) {}
+
+    override fun <T : Plugin?> addPluginListener(
+        action: String,
+        listener: PluginListener<T>,
+        cls: Class<T>
+    ) {}
+
+    override fun <T : Plugin?> addPluginListener(
+        action: String,
+        listener: PluginListener<T>,
+        cls: Class<T>,
+        allowMultiple: Boolean
+    ) {}
+
+    override fun removePluginListener(listener: PluginListener<*>?) {}
+    override fun <T> dependsOn(p: Plugin, cls: Class<T>): Boolean {
+        return false
+    }
+}