diff --git a/tests/robotests/Android.bp b/tests/robotests/Android.bp
new file mode 100644
index 0000000..e0a37c2
--- /dev/null
+++ b/tests/robotests/Android.bp
@@ -0,0 +1,17 @@
+
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+android_robolectric_test {
+    name: "ThemePickerRoboTests",
+    srcs: [
+        "src/**/*.java",
+        "src/**/*.kt",
+    ],
+    java_resource_dirs: ["config"],
+    libs: [
+        "androidx.test.core",
+        "androidx.test.runner",
+    ],
+    instrumentation_for: "ThemePicker",
+}
diff --git a/tests/robotests/AndroidManifest.xml b/tests/robotests/AndroidManifest.xml
new file mode 100644
index 0000000..753aa9e
--- /dev/null
+++ b/tests/robotests/AndroidManifest.xml
@@ -0,0 +1,20 @@
+<!--
+  ~ Copyright (C) 2022 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+    package="com.android.customization">
+    <application/>
+</manifest>
diff --git a/tests/robotests/config/robolectric.properties b/tests/robotests/config/robolectric.properties
new file mode 100644
index 0000000..fab7251
--- /dev/null
+++ b/tests/robotests/config/robolectric.properties
@@ -0,0 +1 @@
+sdk=NEWEST_SDK
diff --git a/tests/robotests/res/values/overlayable_icons_test.xml b/tests/robotests/res/values/overlayable_icons_test.xml
new file mode 100644
index 0000000..73cffe1
--- /dev/null
+++ b/tests/robotests/res/values/overlayable_icons_test.xml
@@ -0,0 +1,34 @@
+<!--
+   Copyright (C) 2019 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.
+-->
+<resources>
+  <!-- overlayable_icons references all of the drawables in this package
+       that are being overlayed by resource overlays. If you remove/rename
+       any of these resources, you must also change the resource overlay icons.-->
+  <array name="overlayable_icons">
+    <item>@drawable/ic_add_24px</item>
+    <item>@drawable/ic_close_24px</item>
+    <item>@drawable/ic_colorize_24px</item>
+    <item>@drawable/ic_delete_24px</item>
+    <item>@drawable/ic_font</item>
+    <item>@drawable/ic_nav_clock</item>
+    <item>@drawable/ic_nav_grid</item>
+    <item>@drawable/ic_nav_theme</item>
+    <item>@drawable/ic_nav_wallpaper</item>
+    <item>@drawable/ic_shapes_24px</item>
+    <item>@drawable/ic_tune</item>
+    <item>@drawable/ic_wifi_24px</item>
+  </array>
+</resources>
diff --git a/tests/robotests/src/com/android/customization/model/clock/BaseClockManagerTest.java b/tests/robotests/src/com/android/customization/model/clock/BaseClockManagerTest.java
new file mode 100644
index 0000000..c96e7f8
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/model/clock/BaseClockManagerTest.java
@@ -0,0 +1,111 @@
+/*
+ * Copyright (C) 2019 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.clock;
+
+import static junit.framework.TestCase.assertTrue;
+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.verify;
+import static org.mockito.Mockito.when;
+
+import androidx.annotation.Nullable;
+
+import com.android.customization.model.CustomizationManager.Callback;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+
+@RunWith(RobolectricTestRunner.class)
+public class BaseClockManagerTest {
+
+    private static final String CURRENT_CLOCK = "current_clock";
+
+    @Mock ClockProvider mProvider;
+    private TestClockManager mManager;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mManager = new TestClockManager(mProvider);
+    }
+
+    @Test
+    public void testIsAvailable() {
+        // GIVEN that the ClockProvider is available
+        when(mProvider.isAvailable()).thenReturn(true);
+        // THEN the BaseClockManager is true
+        assertTrue(mManager.isAvailable());
+    }
+
+    @Test
+    public void testApply() {
+        final String id = "id";
+        Clockface clock = new Clockface.Builder().setId(id).build();
+
+        mManager.apply(clock, new Callback() {
+            @Override
+            public void onSuccess() {
+                //Nothing to do here, the test passed
+            }
+            @Override
+            public void onError(@Nullable Throwable throwable) {
+                fail("onError was called when grid had been applied successfully");
+            }
+        });
+
+        assertEquals(id, mManager.getClockId());
+    }
+
+    @Test
+    public void testFetch() {
+        mManager.fetchOptions(null, false);
+        verify(mProvider).fetch(eq(null), anyBoolean());
+    }
+
+    /**
+     * Testable BaseClockManager that provides basic implementations of abstract methods.
+     */
+    private static final class TestClockManager extends BaseClockManager {
+
+        private String mClockId;
+
+        TestClockManager(ClockProvider provider) {
+            super(provider);
+        }
+
+        String getClockId() {
+            return mClockId;
+        }
+
+        @Override
+        protected void handleApply(Clockface option, Callback callback) {
+            mClockId = option.getId();
+            callback.onSuccess();
+        }
+
+        @Override
+        protected String lookUpCurrentClock() {
+            return CURRENT_CLOCK;
+        }
+    }
+}
diff --git a/tests/robotests/src/com/android/customization/model/clock/ClockManagerTest.java b/tests/robotests/src/com/android/customization/model/clock/ClockManagerTest.java
new file mode 100644
index 0000000..574548a
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/model/clock/ClockManagerTest.java
@@ -0,0 +1,134 @@
+/*
+ * Copyright (C) 2019 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.clock;
+
+import static junit.framework.TestCase.fail;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.verify;
+
+import android.content.ContentResolver;
+import android.provider.Settings.Secure;
+
+import androidx.annotation.Nullable;
+
+import com.android.customization.model.CustomizationManager.Callback;
+import com.android.customization.module.ThemesUserEventLogger;
+
+import org.json.JSONException;
+import org.json.JSONObject;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+
+@RunWith(RobolectricTestRunner.class)
+public class ClockManagerTest {
+
+    private static final String CLOCK_ID = "id";
+    private static final String CLOCK_FIELD = "clock";
+    private static final String CLOCK_FACE_SETTING = "fake_clock_face_setting";
+
+    @Mock ClockProvider mProvider;
+    @Mock ThemesUserEventLogger mLogger;
+    private ContentResolver mContentResolver;
+    private ClockManager mManager;
+    @Mock private Clockface mMockClockface;
+    @Mock private Callback mMockCallback;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mContentResolver = RuntimeEnvironment.application.getContentResolver();
+        mManager = new ClockManager(mContentResolver, mProvider, mLogger);
+    }
+
+    @Test
+    public void testApply() throws JSONException {
+        Clockface clock = new Clockface.Builder().setId(CLOCK_ID).build();
+
+        mManager.apply(clock, new Callback() {
+            @Override
+            public void onSuccess() {
+                //Nothing to do here, the test passed
+            }
+
+            @Override
+            public void onError(@Nullable Throwable throwable) {
+                fail("onError was called when grid had been applied successfully");
+            }
+        });
+
+        // THEN the clock id is written to secure settings.
+        JSONObject json =
+                new JSONObject(Secure.getString(mContentResolver, ClockManager.CLOCK_FACE_SETTING));
+        assertEquals(CLOCK_ID, json.getString(CLOCK_FIELD));
+        // AND the event is logged
+        verify(mLogger).logClockApplied(clock);
+    }
+
+    @Test
+    public void testApply_whenJSONExceptionOccurs_callsOnError() {
+        doAnswer((invocation) -> {
+            throw new JSONException("Fake Test Excepton");
+        }).when(mMockClockface).getId();
+
+        mManager.apply(mMockClockface, mMockCallback);
+
+        verify(mMockCallback).onError(null);
+    }
+
+    @Test
+    public void testGetCurrentClock_returnsClockId() throws JSONException {
+        // Secure settings contains a clock id
+        JSONObject json = new JSONObject().put(CLOCK_FIELD, CLOCK_ID);
+        Secure.putString(mContentResolver, ClockManager.CLOCK_FACE_SETTING, json.toString());
+
+        // The current clock is that id
+        assertEquals(CLOCK_ID, mManager.getCurrentClock());
+    }
+
+    @Test
+    public void testGetCurrentClock_whenJSONExceptionOccurs_returnsClockFaceSetting() {
+        // Secure settings contains a clock face setting with invalid format to cause JSONException.
+        Secure.putString(mContentResolver, ClockManager.CLOCK_FACE_SETTING, CLOCK_FACE_SETTING);
+
+        // The current clock is the clock face setting which is saved in secure settings.
+        assertEquals(CLOCK_FACE_SETTING, mManager.getCurrentClock());
+    }
+
+    @Test
+    public void testGetCurrentClock_withNullIdInSecureSettings_returnsNullId() {
+        // Secure settings contains a null clock id
+        Secure.putString(mContentResolver, ClockManager.CLOCK_FACE_SETTING, /* value= */ null);
+
+        // The current clock is null
+        assertEquals(null, mManager.getCurrentClock());
+    }
+
+    @Test
+    public void testGetCurrentClock_withEmptyIdInSecureSettings_returnsEmptyId() {
+        // Secure settings contains an empty clock id
+        Secure.putString(mContentResolver, ClockManager.CLOCK_FACE_SETTING, /* value= */ "");
+
+        // The current clock is empty
+        assertEquals("", mManager.getCurrentClock());
+    }
+}
diff --git a/tests/robotests/src/com/android/customization/model/color/ColorCustomizationManagerTest.kt b/tests/robotests/src/com/android/customization/model/color/ColorCustomizationManagerTest.kt
new file mode 100644
index 0000000..80d01c6
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/model/color/ColorCustomizationManagerTest.kt
@@ -0,0 +1,194 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.customization.model.color
+
+import android.app.WallpaperColors
+import android.graphics.Color
+import com.android.customization.model.CustomizationManager
+import com.android.customization.model.ResourceConstants.OVERLAY_CATEGORY_COLOR
+import com.android.customization.model.ResourceConstants.OVERLAY_CATEGORY_SYSTEM_PALETTE
+import com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_HOME
+import com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_PRESET
+import com.android.customization.model.color.ColorOptionsProvider.OVERLAY_COLOR_BOTH
+import com.android.customization.model.color.ColorOptionsProvider.OVERLAY_COLOR_INDEX
+import com.android.customization.model.color.ColorOptionsProvider.OVERLAY_COLOR_SOURCE
+import com.android.customization.model.theme.OverlayManagerCompat
+import com.android.systemui.monet.Style
+import com.google.common.truth.Truth.assertThat
+import org.json.JSONObject
+import org.junit.Before
+import org.junit.Ignore
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.anyInt
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
+import org.mockito.junit.MockitoJUnit
+import org.mockito.junit.MockitoRule
+import org.robolectric.RobolectricTestRunner
+import org.robolectric.RuntimeEnvironment
+
+/** Tests of {@link ColorCustomizationManager}. */
+// TODO(b/222433744): most of these tests are failing due to the manager apk missing in the image
+@RunWith(RobolectricTestRunner::class)
+class ColorCustomizationManagerTest {
+
+    @get:Rule val rule: MockitoRule = MockitoJUnit.rule()
+
+    @Mock private lateinit var provider: ColorOptionsProvider
+    @Mock private lateinit var mockOM: OverlayManagerCompat
+
+    private lateinit var manager: ColorCustomizationManager
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        val application = RuntimeEnvironment.application
+        manager = ColorCustomizationManager(provider, application.contentResolver, mockOM)
+    }
+
+    @Test
+    @Ignore("b/260925899")
+    fun testParseSettings() {
+        val source = COLOR_SOURCE_HOME
+        val style = Style.SPRITZ
+        val someColor = "aabbcc"
+        val someOtherColor = "bbccdd"
+        val settings =
+            mapOf(
+                OVERLAY_CATEGORY_SYSTEM_PALETTE to someColor,
+                OVERLAY_CATEGORY_COLOR to someOtherColor,
+                OVERLAY_COLOR_SOURCE to source,
+                ColorOption.TIMESTAMP_FIELD to "12345"
+            )
+        val json = JSONObject(settings).toString()
+
+        manager.parseSettings(json)
+
+        assertThat(manager.currentColorSource).isEqualTo(source)
+        assertThat(manager.currentStyle).isEqualTo(style)
+        assertThat(manager.currentOverlays.size).isEqualTo(2)
+        assertThat(manager.currentOverlays.get(OVERLAY_CATEGORY_COLOR)).isEqualTo(someOtherColor)
+        assertThat(manager.currentOverlays.get(OVERLAY_CATEGORY_SYSTEM_PALETTE))
+            .isEqualTo(someColor)
+    }
+
+    @Test
+    @Ignore("b/260925899")
+    fun apply_ColorBundle_index() {
+        testApplyColorBundle(1, "1")
+        testApplyColorBundle(2, "2")
+        testApplyColorBundle(3, "3")
+        testApplyColorBundle(4, "4")
+    }
+
+    private fun testApplyColorBundle(index: Int, value: String) {
+        manager.apply(
+            getColorBundle(index),
+            object : CustomizationManager.Callback {
+                override fun onSuccess() {}
+                override fun onError(throwable: Throwable?) {}
+            }
+        )
+
+        val overlaysJson = JSONObject(manager.storedOverlays)
+
+        assertThat(overlaysJson.getString(OVERLAY_COLOR_INDEX)).isEqualTo(value)
+    }
+
+    private fun getColorBundle(index: Int): ColorBundle {
+        return ColorBundle(
+            "fake color",
+            mapOf("fake_package" to "fake_color"),
+            /* isDefault= */ false,
+            null,
+            /* index= */ index,
+            null
+        )
+    }
+
+    @Test
+    @Ignore("b/260925899")
+    fun apply_ColorSeed_index() {
+        testApplyColorSeed(1, "1")
+        testApplyColorSeed(2, "2")
+        testApplyColorSeed(3, "3")
+        testApplyColorSeed(4, "4")
+    }
+
+    private fun testApplyColorSeed(index: Int, value: String) {
+        manager.apply(
+            getColorSeed(index),
+            object : CustomizationManager.Callback {
+                override fun onSuccess() {}
+                override fun onError(throwable: Throwable?) {}
+            }
+        )
+
+        val overlaysJson = JSONObject(manager.storedOverlays)
+        assertThat(overlaysJson.getString(OVERLAY_COLOR_INDEX)).isEqualTo(value)
+    }
+
+    private fun getColorSeed(index: Int): ColorSeedOption {
+        return ColorSeedOption(
+            "fake color",
+            mapOf("fake_package" to "fake_color"),
+            /* isDefault= */ false,
+            COLOR_SOURCE_PRESET,
+            null,
+            index,
+            null
+        )
+    }
+
+    @Test
+    @Ignore("b/260925899")
+    fun testApply_colorSeedFromWallpaperBoth_shouldReturnBothValue() {
+        val wallpaperColor = WallpaperColors(Color.valueOf(Color.RED), null, null)
+        manager.setWallpaperColors(wallpaperColor, wallpaperColor)
+
+        manager.apply(
+            getColorSeed(anyInt()),
+            object : CustomizationManager.Callback {
+                override fun onSuccess() {}
+                override fun onError(throwable: Throwable?) {}
+            }
+        )
+
+        val overlaysJson = JSONObject(manager.storedOverlays)
+        assertThat(overlaysJson.getString(OVERLAY_COLOR_BOTH)).isEqualTo("1")
+    }
+
+    @Test
+    @Ignore("b/260925899")
+    fun testApply_colorSeedFromWallpaperDifferent_shouldReturnNonBothValue() {
+        val wallpaperColor1 = WallpaperColors(Color.valueOf(Color.RED), null, null)
+        val wallpaperColor2 = WallpaperColors(Color.valueOf(Color.BLUE), null, null)
+        manager.setWallpaperColors(wallpaperColor1, wallpaperColor2)
+
+        manager.apply(
+            getColorSeed(anyInt()),
+            object : CustomizationManager.Callback {
+                override fun onSuccess() {}
+                override fun onError(throwable: Throwable?) {}
+            }
+        )
+
+        val overlaysJson = JSONObject(manager.storedOverlays)
+        assertThat(overlaysJson.getString(OVERLAY_COLOR_BOTH)).isEqualTo("0")
+    }
+}
diff --git a/tests/robotests/src/com/android/customization/model/color/ColorOptionTest.kt b/tests/robotests/src/com/android/customization/model/color/ColorOptionTest.kt
new file mode 100644
index 0000000..0431c19
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/model/color/ColorOptionTest.kt
@@ -0,0 +1,215 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.customization.model.color
+
+import com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_HOME
+import com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_LOCK
+import com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_PRESET
+import com.android.systemui.monet.Style
+import com.google.common.truth.Truth.assertThat
+import org.junit.Assert.assertEquals
+import org.junit.Ignore
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.`when`
+import org.mockito.junit.MockitoJUnit
+import org.mockito.junit.MockitoRule
+import org.robolectric.RobolectricTestRunner
+
+/** Tests of {@link ColorOption}. */
+@RunWith(RobolectricTestRunner::class)
+class ColorOptionTest {
+
+    @get:Rule val rule: MockitoRule = MockitoJUnit.rule()
+
+    @Mock private lateinit var manager: ColorCustomizationManager
+
+    @Test
+    fun colorOption_Source_Preset() {
+        val bundleOption: ColorOption =
+            ColorBundle(
+                "fake color",
+                mapOf("fake_package" to "fake_color"),
+                false,
+                null,
+                /* index= */ 0,
+                null
+            )
+        assertEquals(COLOR_SOURCE_PRESET, bundleOption.source)
+    }
+
+    @Test
+    fun colorOption_bundle_index() {
+        testBundleOptionIndex(1)
+        testBundleOptionIndex(2)
+        testBundleOptionIndex(3)
+        testBundleOptionIndex(4)
+    }
+
+    private fun testBundleOptionIndex(index: Int) {
+        val bundleOption: ColorBundle =
+            ColorBundle(
+                "fake color",
+                mapOf("fake_package" to "fake_color"),
+                false,
+                null,
+                /* index= */ index,
+                null
+            )
+        assertThat(bundleOption.index).isEqualTo(index)
+    }
+
+    @Test
+    fun colorOption_Source_Seed() {
+        testSeedOptionSource(COLOR_SOURCE_HOME)
+        testSeedOptionSource(COLOR_SOURCE_LOCK)
+    }
+
+    private fun testSeedOptionSource(source: String) {
+        val seedOption: ColorOption =
+            ColorSeedOption(
+                "fake color",
+                mapOf("fake_package" to "fake_color"),
+                false,
+                source,
+                null,
+                /* index= */ 0,
+                null
+            )
+        assertThat(seedOption.source).isEqualTo(source)
+    }
+
+    @Test
+    fun colorOption_seed_style() {
+        testSeedOptionStyle(Style.TONAL_SPOT)
+        testSeedOptionStyle(Style.SPRITZ)
+        testSeedOptionStyle(Style.VIBRANT)
+        testSeedOptionStyle(Style.EXPRESSIVE)
+    }
+
+    private fun testSeedOptionStyle(style: Style) {
+        val seedOption: ColorOption =
+            ColorSeedOption(
+                "fake color",
+                mapOf("fake_package" to "fake_color"),
+                /* isDefault= */ false,
+                "fake_source",
+                style,
+                0,
+                null
+            )
+        assertThat(seedOption.style).isEqualTo(style)
+    }
+
+    @Test
+    fun colorOption_seed_index() {
+        testSeedOptionIndex(1)
+        testSeedOptionIndex(2)
+        testSeedOptionIndex(3)
+        testSeedOptionIndex(4)
+    }
+
+    private fun testSeedOptionIndex(index: Int) {
+        val seedOption: ColorOption =
+            ColorSeedOption(
+                "fake color",
+                mapOf("fake_package" to "fake_color"),
+                /* isDefault= */ false,
+                "fake_source",
+                Style.TONAL_SPOT,
+                index,
+                /* previewInfo= */ null
+            )
+        assertThat(seedOption.index).isEqualTo(index)
+    }
+
+    private fun setUpSeedOption(
+        isDefault: Boolean,
+        source: String = "some_source"
+    ): ColorSeedOption {
+        val overlays =
+            if (isDefault) {
+                HashMap()
+            } else {
+                mapOf("package" to "value", "otherPackage" to "otherValue")
+            }
+        `when`(manager.currentOverlays).thenReturn(overlays)
+        return ColorSeedOption(
+            "seed",
+            overlays,
+            isDefault,
+            source,
+            Style.TONAL_SPOT,
+            /* index= */ 0,
+            /* previewInfo= */ null
+        )
+    }
+
+    @Test
+    fun seedOption_isActive_notDefault_SourceSet() {
+        val source = "some_source"
+        val seedOption = setUpSeedOption(false, source)
+        `when`(manager.currentColorSource).thenReturn(source)
+
+        assertThat(seedOption.isActive(manager)).isTrue()
+    }
+
+    @Test
+    fun seedOption_isActive_notDefault_NoSource() {
+        val seedOption = setUpSeedOption(false)
+        `when`(manager.currentColorSource).thenReturn(null)
+
+        assertThat(seedOption.isActive(manager)).isTrue()
+    }
+
+    @Test
+    fun seedOption_isActive_notDefault_differentSource() {
+        val seedOption = setUpSeedOption(false)
+        `when`(manager.currentColorSource).thenReturn("some_other_source")
+
+        assertThat(seedOption.isActive(manager)).isFalse()
+    }
+
+    @Test
+    fun seedOption_isActive_default_emptyJson() {
+        val seedOption = setUpSeedOption(true)
+        `when`(manager.storedOverlays).thenReturn("")
+
+        assertThat(seedOption.isActive(manager)).isTrue()
+    }
+
+    @Test
+    fun seedOption_isActive_default_nonEmptyJson() {
+        val seedOption = setUpSeedOption(true)
+
+        `when`(manager.storedOverlays).thenReturn("{non-empty-json}")
+
+        // Should still be Active because overlays is empty
+        assertThat(seedOption.isActive(manager)).isTrue()
+    }
+
+    @Test
+    @Ignore("b/260925899")
+    fun seedOption_isActive_default_nonEmptyOverlays() {
+        val seedOption = setUpSeedOption(true)
+
+        `when`(manager.currentOverlays).thenReturn(mapOf("a" to "b"))
+        // TODO(b/222433744): failing as it's true
+        assertThat(seedOption.isActive(manager)).isFalse()
+    }
+}
diff --git a/tests/robotests/src/com/android/customization/model/color/ColorSectionControllerTest.java b/tests/robotests/src/com/android/customization/model/color/ColorSectionControllerTest.java
new file mode 100644
index 0000000..820e641
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/model/color/ColorSectionControllerTest.java
@@ -0,0 +1,62 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.customization.model.color;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import androidx.appcompat.app.AppCompatActivity;
+
+import com.android.wallpaper.model.WallpaperColorsViewModel;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.Robolectric;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.annotation.Config;
+
+/**
+ * Tests of {@link ColorSectionController}.
+ */
+@RunWith(RobolectricTestRunner.class)
+@Config(manifest = Config.NONE)
+public final class ColorSectionControllerTest {
+
+    private AppCompatActivity mActivity;
+    private ColorSectionController mColorSectionController;
+
+    /**
+     * Set up the test case.
+     */
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mActivity = Robolectric.buildActivity(AppCompatActivity.class).create().get();
+        mColorSectionController = new ColorSectionController(mActivity,
+                new WallpaperColorsViewModel(), mActivity, null);
+    }
+
+    /**
+     * isAvailable()'s test.
+     */
+    @Test
+    @Config(manifest = Config.NONE)
+    public void isAvailable_nullContext_shouldReturnFalse() {
+        assertThat(mColorSectionController.isAvailable(/* context= */ null)).isFalse();
+    }
+}
+
diff --git a/tests/robotests/src/com/android/customization/model/grid/GridOptionsManagerTest.java b/tests/robotests/src/com/android/customization/model/grid/GridOptionsManagerTest.java
new file mode 100644
index 0000000..89ca676
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/model/grid/GridOptionsManagerTest.java
@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2019 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;
+
+import static junit.framework.TestCase.fail;
+
+import static org.mockito.Matchers.anyBoolean;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import androidx.annotation.Nullable;
+
+import com.android.customization.model.CustomizationManager.Callback;
+import com.android.customization.module.ThemesUserEventLogger;
+
+import org.junit.Before;
+import org.junit.Ignore;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+
+@RunWith(RobolectricTestRunner.class)
+public class GridOptionsManagerTest {
+
+    @Mock LauncherGridOptionsProvider mProvider;
+    @Mock ThemesUserEventLogger mThemesUserEventLogger;
+    private GridOptionsManager mManager;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mManager = new GridOptionsManager(mProvider, mThemesUserEventLogger);
+    }
+
+    @Test
+    public void testApply() {
+        String gridName = "testName";
+        GridOption grid = new GridOption("testTitle", gridName, false, 2, 2, null, 1, "");
+        when(mProvider.applyGrid(gridName)).thenReturn(1);
+
+        mManager.apply(grid, new Callback() {
+            @Override
+            public void onSuccess() {
+                //Nothing to do here, the test passed
+            }
+
+            @Override
+            public void onError(@Nullable Throwable throwable) {
+                fail("onError was called when grid had been applied successfully");
+            }
+        });
+    }
+
+    @Test
+    @Ignore("b/260925899")
+    public void testFetch_backgroundThread() {
+        mManager.fetchOptions(null, false);
+        verify(mProvider).fetch(anyBoolean());
+    }
+}
diff --git a/tests/robotests/src/com/android/customization/model/theme/ThemeManagerTest.java b/tests/robotests/src/com/android/customization/model/theme/ThemeManagerTest.java
new file mode 100644
index 0000000..cfb8a33
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/model/theme/ThemeManagerTest.java
@@ -0,0 +1,222 @@
+/*
+ * Copyright (C) 2019 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.theme;
+
+import static com.android.customization.model.ResourceConstants.ANDROID_PACKAGE;
+import static com.android.customization.model.ResourceConstants.OVERLAY_CATEGORY_COLOR;
+import static com.android.customization.model.ResourceConstants.OVERLAY_CATEGORY_FONT;
+import static com.android.customization.model.ResourceConstants.OVERLAY_CATEGORY_ICON_ANDROID;
+import static com.android.customization.model.ResourceConstants.OVERLAY_CATEGORY_ICON_SETTINGS;
+import static com.android.customization.model.ResourceConstants.OVERLAY_CATEGORY_ICON_SYSUI;
+import static com.android.customization.model.ResourceConstants.OVERLAY_CATEGORY_SHAPE;
+import static com.android.customization.model.ResourceConstants.SETTINGS_PACKAGE;
+import static com.android.customization.model.ResourceConstants.SYSUI_PACKAGE;
+import static com.android.customization.model.ResourceConstants.THEME_SETTING;
+
+import static junit.framework.Assert.assertFalse;
+import static junit.framework.Assert.assertTrue;
+import static junit.framework.TestCase.assertEquals;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.provider.Settings;
+
+import androidx.annotation.Nullable;
+import androidx.fragment.app.FragmentActivity;
+
+import com.android.customization.model.CustomizationManager.Callback;
+import com.android.customization.model.CustomizationManager.OptionsFetchedListener;
+import com.android.customization.model.theme.custom.CustomTheme;
+import com.android.customization.module.ThemesUserEventLogger;
+import com.android.customization.testutils.OverlayManagerMocks;
+
+import org.json.JSONObject;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.Robolectric;
+import org.robolectric.RobolectricTestRunner;
+
+@RunWith(RobolectricTestRunner.class)
+public class ThemeManagerTest {
+
+    @Mock OverlayManagerCompat mMockOm;
+    @Mock ThemesUserEventLogger mThemesUserEventLogger;
+    @Mock ThemeBundleProvider mThemeBundleProvider;
+    private OverlayManagerMocks mMockOmHelper;
+    private ThemeManager mThemeManager;
+    private FragmentActivity mActivity;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        FragmentActivity activity = Robolectric.buildActivity(FragmentActivity.class).get();
+        mActivity = spy(activity);
+        mMockOmHelper = new OverlayManagerMocks();
+        mMockOmHelper.setUpMock(mMockOm);
+        mThemeManager = new ThemeManager(mThemeBundleProvider, activity,
+                mMockOm, mThemesUserEventLogger);
+    }
+
+    @After
+    public void cleanUp() {
+        mMockOmHelper.clearOverlays();
+    }
+
+    @Test
+    public void apply_WithDefaultTheme_StoresEmptyJsonString() {
+        mMockOmHelper.addOverlay("test.package.name_color", ANDROID_PACKAGE,
+                OVERLAY_CATEGORY_COLOR, true, 0);
+        mMockOmHelper.addOverlay("test.package.name_font", ANDROID_PACKAGE,
+                OVERLAY_CATEGORY_FONT, true, 0);
+        mMockOmHelper.addOverlay("test.package.name_shape", ANDROID_PACKAGE,
+                OVERLAY_CATEGORY_SHAPE, true, 0);
+        mMockOmHelper.addOverlay("test.package.name_icon", ANDROID_PACKAGE,
+                OVERLAY_CATEGORY_ICON_ANDROID, true, 0);
+        mMockOmHelper.addOverlay("test.package.name_settings", SETTINGS_PACKAGE,
+                OVERLAY_CATEGORY_ICON_SETTINGS, true, 0);
+        mMockOmHelper.addOverlay("test.package.name_sysui", SYSUI_PACKAGE,
+                OVERLAY_CATEGORY_ICON_SYSUI, true, 0);
+        mMockOmHelper.addOverlay("test.package.name_themepicker", mActivity.getPackageName(),
+                OVERLAY_CATEGORY_ICON_SYSUI, true, 0);
+
+        ThemeBundle defaultTheme = new ThemeBundle.Builder().asDefault().build(mActivity);
+
+        applyTheme(defaultTheme);
+
+        assertEquals("Secure Setting should be empty JSON string after applying default theme",
+                new JSONObject().toString(),
+                Settings.Secure.getString(mActivity.getContentResolver(), THEME_SETTING));
+    }
+
+    @Test
+    public void apply_WithOverlayTheme_StoresSerializedPackagesWithTimestamp() {
+        ThemeBundle theme = getOverlayTheme();
+        final String serializedPackagesWithTimestamp = theme.getSerializedPackagesWithTimestamp();
+
+        theme = spy(theme);
+        // Makes it return the fixed serializedPackagesWithTimestamp to test. Since we will get
+        // fresh time every time, it's hard to compare for testing.
+        when(theme.getSerializedPackagesWithTimestamp())
+                .thenReturn(serializedPackagesWithTimestamp);
+
+        applyTheme(theme);
+
+        assertEquals("Secure Setting should be the overlay packages after applying theme",
+                serializedPackagesWithTimestamp,
+                Settings.Secure.getString(mActivity.getContentResolver(), THEME_SETTING));
+    }
+
+    @Test
+    public void isAvailable_ThemeBundleProviderAndOverlayManagerAreAvailable_ReturnsTrue() {
+        when(mThemeBundleProvider.isAvailable()).thenReturn(true);
+        when(mMockOm.isAvailable()).thenReturn(true);
+
+        assertTrue(mThemeManager.isAvailable());
+    }
+
+    @Test
+    public void isAvailable_ThemeBundleProviderOrOverlayManagerIsAvailable_ReturnsFalse() {
+        when(mThemeBundleProvider.isAvailable()).thenReturn(false);
+        when(mMockOm.isAvailable()).thenReturn(true);
+        assertFalse(mThemeManager.isAvailable());
+
+        when(mThemeBundleProvider.isAvailable()).thenReturn(true);
+        when(mMockOm.isAvailable()).thenReturn(false);
+        assertFalse(mThemeManager.isAvailable());
+    }
+
+    @Test
+    public void fetchOptions_ThemeBundleProviderFetches() {
+        OptionsFetchedListener listener = mock(OptionsFetchedListener.class);
+
+        mThemeManager.fetchOptions(listener, false);
+
+        verify(mThemeBundleProvider).fetch(listener, false);
+    }
+
+    @Test
+    public void removeCustomTheme_ThemeBundleProviderRemovesCustomTheme() {
+        CustomTheme customTheme = mock(CustomTheme.class);
+        mThemeManager.removeCustomTheme(customTheme);
+
+        verify(mThemeBundleProvider).removeCustomTheme(customTheme);
+    }
+
+    @Test
+    public void findThemeByPackages_ThemeBundleProviderFindsEquivalent() {
+        CustomTheme theme = mock(CustomTheme.class);
+        mThemeManager.findThemeByPackages(theme);
+
+        verify(mThemeBundleProvider).findEquivalent(theme);
+    }
+
+    @Test
+    public void storeEmptyTheme_SettingsSecureStoresEmptyTheme() {
+        mThemeManager.storeEmptyTheme();
+
+        assertEquals(
+                new JSONObject().toString(),
+                Settings.Secure.getString(mActivity.getContentResolver(), THEME_SETTING));
+    }
+
+    @Test
+    public void getStoredOverlays_GetsFromSettingsSecureWithExpectedName() {
+        ThemeBundle theme = getOverlayTheme();
+
+        applyTheme(theme);
+
+        assertEquals(
+                Settings.Secure.getString(mActivity.getContentResolver(), THEME_SETTING),
+                mThemeManager.getStoredOverlays());
+    }
+
+    private ThemeBundle getOverlayTheme() {
+        final String bundleColorPackage = "test.package.name_color";
+        final String bundleFontPackage = "test.package.name_font";
+        final String otherPackage = "other.package.name_font";
+
+        mMockOmHelper.addOverlay(bundleColorPackage, ANDROID_PACKAGE,
+                OVERLAY_CATEGORY_COLOR, false, 0);
+        mMockOmHelper.addOverlay(bundleFontPackage, ANDROID_PACKAGE,
+                OVERLAY_CATEGORY_FONT, false, 0);
+        mMockOmHelper.addOverlay(otherPackage, ANDROID_PACKAGE,
+                OVERLAY_CATEGORY_FONT, false, 0);
+
+        return new ThemeBundle.Builder()
+                .addOverlayPackage(OVERLAY_CATEGORY_COLOR, bundleColorPackage)
+                .addOverlayPackage(OVERLAY_CATEGORY_FONT, bundleFontPackage)
+                .build(mActivity);
+    }
+
+    private void applyTheme(ThemeBundle theme) {
+        mThemeManager.apply(theme, new Callback() {
+            @Override
+            public void onSuccess() {
+            }
+
+            @Override
+            public void onError(@Nullable Throwable throwable) {
+            }
+        });
+    }
+}
diff --git a/tests/robotests/src/com/android/customization/testutils/Condition.java b/tests/robotests/src/com/android/customization/testutils/Condition.java
new file mode 100644
index 0000000..f013555
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/testutils/Condition.java
@@ -0,0 +1,20 @@
+/*
+ * Copyright (C) 2019 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.testutils;
+
+public interface Condition {
+    boolean isTrue() throws Throwable;
+}
diff --git a/tests/robotests/src/com/android/customization/testutils/OverlayManagerMocks.java b/tests/robotests/src/com/android/customization/testutils/OverlayManagerMocks.java
new file mode 100644
index 0000000..79b2e43
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/testutils/OverlayManagerMocks.java
@@ -0,0 +1,135 @@
+/*
+ * Copyright (C) 2019 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.testutils;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.when;
+
+import android.content.om.OverlayManager;
+import android.text.TextUtils;
+
+import com.android.customization.model.theme.OverlayManagerCompat;
+
+import org.mockito.stubbing.Answer;
+
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+/**
+ * Helper class to provide mock implementation for OverlayManager, to use, create a Mockito Mock
+ * for OverlayManager and call {@link #setUpMock(OverlayManager)} with it, then use
+ * {@link #addOverlay(String, String, String, boolean, int)} to add fake OverlayInfo to be returned
+ * by the mocked OverlayManager.
+ */
+public class OverlayManagerMocks {
+    private static class MockOverlay {
+        final String mPackageName;
+        final String mTargetPackage;
+        final String mCategory;
+
+        MockOverlay(String packageName, String targetPackage, String category) {
+            this.mPackageName = packageName;
+            this.mTargetPackage = targetPackage;
+            this.mCategory = category;
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            return obj instanceof MockOverlay
+                    && TextUtils.equals(((MockOverlay) obj).mPackageName, mPackageName)
+                    && TextUtils.equals(((MockOverlay) obj).mTargetPackage, mTargetPackage)
+                    && TextUtils.equals(((MockOverlay) obj).mCategory, mCategory);
+        }
+    }
+
+    private Set<MockOverlay> mAllOverlays = new HashSet<>();
+    private Set<MockOverlay> mEnabledOverlays = new HashSet<>();
+
+    private boolean setEnabled(String packageName, boolean enable, int userId) {
+        if (packageName == null) {
+            return false;
+        }
+        Set<MockOverlay> packageOverlays = mAllOverlays.stream()
+                .filter(mockOverlay -> mockOverlay.mPackageName.equals(packageName)).collect(
+                Collectors.toSet());
+        if (packageOverlays.isEmpty()) {
+            return false;
+        }
+        if (enable) {
+            mEnabledOverlays.addAll(packageOverlays);
+        } else {
+            mEnabledOverlays.removeAll(packageOverlays);
+        }
+        return true;
+    }
+
+    public void addOverlay(String packageName, String targetPackage, String category,
+            boolean enabled, int userId) {
+        MockOverlay overlay = new MockOverlay(packageName, targetPackage, category);
+        mAllOverlays.add(overlay);
+        if (enabled) {
+            mEnabledOverlays.add(overlay);
+        }
+    }
+
+    public void clearOverlays() {
+        mAllOverlays.clear();
+        mEnabledOverlays.clear();
+    }
+
+    public void setUpMock(OverlayManagerCompat mockOverlayManager) {
+        when(mockOverlayManager.getEnabledPackageName(anyString(), anyString())).then(
+                (Answer<String>) inv ->
+                        mEnabledOverlays.stream().filter(
+                                mockOverlay ->
+                                        mockOverlay.mTargetPackage.equals(inv.getArgument(0))
+                                            && mockOverlay.mCategory.equals(inv.getArgument(1)))
+                                .map(overlay -> overlay.mPackageName).findFirst().orElse(null));
+
+
+        when(mockOverlayManager.disableOverlay(anyString(), anyInt())).then(
+                (Answer<Boolean>) invocation ->
+                        setEnabled(invocation.getArgument(0),
+                                false,
+                                invocation.getArgument(1)));
+
+        when(mockOverlayManager.setEnabledExclusiveInCategory(anyString(), anyInt())).then(
+                (Answer<Boolean>) invocation ->
+                        setEnabled(
+                                invocation.getArgument(0),
+                                true,
+                                invocation.getArgument(1)));
+
+        when(mockOverlayManager.getEnabledOverlaysForTargets(any())).then(
+                (Answer<Map<String, String>>) inv ->
+                        mEnabledOverlays.stream().filter(
+                                overlay ->
+                                        Arrays.asList(inv.getArguments())
+                                                .contains(overlay.mTargetPackage))
+                                .collect(Collectors.toMap(
+                                        overlay ->
+                                                overlay.mCategory,
+                                        (Function<MockOverlay, String>) overlay ->
+                                                overlay.mPackageName))
+        );
+    }
+}
diff --git a/tests/robotests/src/com/android/customization/testutils/Wait.java b/tests/robotests/src/com/android/customization/testutils/Wait.java
new file mode 100644
index 0000000..54650ba
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/testutils/Wait.java
@@ -0,0 +1,58 @@
+/*
+ * Copyright (C) 2019 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.testutils;
+
+
+import android.os.SystemClock;
+
+import org.junit.Assert;
+
+/**
+ * A utility class for waiting for a condition to be true.
+ */
+public class Wait {
+
+    private static final long DEFAULT_SLEEP_MS = 200;
+
+    public static void atMost(String message, Condition condition, long timeout) {
+        atMost(message, condition, timeout, DEFAULT_SLEEP_MS);
+    }
+
+    public static void atMost(String message, Condition condition, long timeout, long sleepMillis) {
+        long endTime = SystemClock.uptimeMillis() + timeout;
+        while (SystemClock.uptimeMillis() < endTime) {
+            try {
+                if (condition.isTrue()) {
+                    return;
+                }
+            } catch (Throwable t) {
+                throw new RuntimeException(t);
+            }
+            SystemClock.sleep(sleepMillis);
+        }
+
+        // Check once more before returning false.
+        try {
+            if (condition.isTrue()) {
+                return;
+            }
+        } catch (Throwable t) {
+            throw new RuntimeException(t);
+        }
+        Assert.fail(message);
+    }
+}
+