Add initial Robolectric tests

Set up the infrastructure for tests and add a few initial Robolectric
tests for ThemePicker.

Bug: 118758604

Change-Id: Ie6f34c840d31d24349fd12355f93369f51168802
diff --git a/robolectric_tests/Android.mk b/robolectric_tests/Android.mk
new file mode 100644
index 0000000..3037a47
--- /dev/null
+++ b/robolectric_tests/Android.mk
@@ -0,0 +1,55 @@
+# 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.
+
+#############################################
+# ThenePicker Robolectric test target.      #
+#############################################
+LOCAL_PATH := $(call my-dir)
+include $(CLEAR_VARS)
+
+LOCAL_MODULE := ThemePickerRoboTests
+LOCAL_SDK_VERSION := system_current
+LOCAL_SRC_FILES := $(call all-java-files-under, src)
+LOCAL_STATIC_JAVA_LIBRARIES := \
+    androidx.test.runner \
+    androidx.test.rules \
+    mockito-robolectric-prebuilt \
+    truth-prebuilt
+LOCAL_JAVA_LIBRARIES := \
+    platform-robolectric-3.6.2-prebuilt
+
+LOCAL_JAVA_RESOURCE_DIRS := config
+
+LOCAL_INSTRUMENTATION_FOR := ThemePicker
+LOCAL_MODULE_TAGS := optional
+
+include $(BUILD_STATIC_JAVA_LIBRARY)
+
+############################################
+# Target to run the previous target.       #
+############################################
+include $(CLEAR_VARS)
+
+LOCAL_MODULE := RunThemePickerRoboTests
+LOCAL_SDK_VERSION := system_current
+LOCAL_JAVA_LIBRARIES := \
+    ThemePickerRoboTests
+
+LOCAL_TEST_PACKAGE := ThemePicker
+
+LOCAL_INSTRUMENT_SOURCE_DIRS := $(dir $(LOCAL_PATH))../src \
+
+LOCAL_ROBOTEST_TIMEOUT := 36000
+
+include prebuilts/misc/common/robolectric/3.6.2/run_robotests.mk
diff --git a/robolectric_tests/config/robolectric.properties b/robolectric_tests/config/robolectric.properties
new file mode 100644
index 0000000..197c393
--- /dev/null
+++ b/robolectric_tests/config/robolectric.properties
@@ -0,0 +1,2 @@
+manifest=vendor/unbundled_google/packages/WallpaperPickerGoogle/AndroidManifest.xml
+sdk=27
diff --git a/robolectric_tests/robolectric_gradle_config/robolectric.properties b/robolectric_tests/robolectric_gradle_config/robolectric.properties
new file mode 100644
index 0000000..926e354
--- /dev/null
+++ b/robolectric_tests/robolectric_gradle_config/robolectric.properties
@@ -0,0 +1,3 @@
+# Do not include the manifest definition in this version
+# as it is specified directly in the gradle file
+sdk=27
\ No newline at end of file
diff --git a/robolectric_tests/src/com/android/customization/model/grid/GridOptionsManagerTest.java b/robolectric_tests/src/com/android/customization/model/grid/GridOptionsManagerTest.java
new file mode 100644
index 0000000..abbbaa0
--- /dev/null
+++ b/robolectric_tests/src/com/android/customization/model/grid/GridOptionsManagerTest.java
@@ -0,0 +1,73 @@
+/*
+ * 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 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 GridOptionsManagerTest {
+
+    @Mock LauncherGridOptionsProvider mProvider;
+    private GridOptionsManager mManager;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mManager = new GridOptionsManager(mProvider);
+    }
+
+    @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
+    public void testFetch_backgroundThread() {
+        mManager.fetchOptions(null);
+        Robolectric.flushBackgroundThreadScheduler();
+        verify(mProvider).fetch(anyBoolean());
+    }
+}
diff --git a/robolectric_tests/src/com/android/customization/model/theme/ThemeManagerTest.java b/robolectric_tests/src/com/android/customization/model/theme/ThemeManagerTest.java
new file mode 100644
index 0000000..bb5df66
--- /dev/null
+++ b/robolectric_tests/src/com/android/customization/model/theme/ThemeManagerTest.java
@@ -0,0 +1,171 @@
+/*
+ * 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.TestCase.assertEquals;
+import static junit.framework.TestCase.assertTrue;
+
+import static org.junit.Assert.assertFalse;
+import static org.mockito.Matchers.anyInt;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import android.app.Activity;
+import android.provider.Settings;
+
+import androidx.annotation.Nullable;
+import androidx.fragment.app.FragmentActivity;
+
+import com.android.customization.model.CustomizationManager.Callback;
+import com.android.customization.testutils.Condition;
+import com.android.customization.testutils.OverlayManagerMocks;
+import com.android.customization.testutils.Wait;
+import com.android.wallpaper.module.WallpaperSetter;
+
+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;
+
+import java.util.Map;
+
+@RunWith(RobolectricTestRunner.class)
+public class ThemeManagerTest {
+
+    @Mock OverlayManagerCompat mMockOm;
+    @Mock WallpaperSetter mMockWallpaperSetter;
+    private OverlayManagerMocks mMockOmHelper;
+    private ThemeManager mThemeManager;
+    private Activity mActivity;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        Activity activity = Robolectric.buildActivity(FragmentActivity.class).get();
+        mActivity = spy(activity);
+        mMockOmHelper = new OverlayManagerMocks();
+        mMockOmHelper.setUpMock(mMockOm);
+        ThemeBundleProvider provider = mock(ThemeBundleProvider.class);
+        mThemeManager = new ThemeManager(provider, activity, mMockWallpaperSetter, mMockOm);
+    }
+
+    @After
+    public void cleanUp() {
+        mMockOmHelper.clearOverlays();
+    }
+
+    @Test
+    public void testApply_DefaultTheme() {
+        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);
+
+        ThemeBundle defaultTheme = new ThemeBundle.Builder().asDefault().build();
+
+        applyThemeAndWaitForCondition(defaultTheme, "Overlays didn't get disabled", () -> {
+            verify(mMockOm, times(6)).disableOverlay(anyString(), anyInt());
+            return true;
+        });
+
+        assertEquals("Secure Setting should be emtpy after applying default theme",
+                "",
+                Settings.Secure.getString(mActivity.getContentResolver(), THEME_SETTING));
+    }
+
+    @Test
+    public void testApply_NonDefault() {
+        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);
+
+        ThemeBundle theme = new ThemeBundle.Builder()
+                .addOverlayPackage(OVERLAY_CATEGORY_COLOR, bundleColorPackage)
+                .addOverlayPackage(OVERLAY_CATEGORY_FONT, bundleFontPackage)
+                .build();
+
+        applyThemeAndWaitForCondition(theme, "Overlays didn't get enabled", () -> {
+            verify(mMockOm, times(2)).setEnabledExclusiveInCategory(anyString(), anyInt());
+            Map<String, String> overlays = mMockOm.getEnabledOverlaysForTargets(ANDROID_PACKAGE);
+            assertEquals(2, overlays.size());
+            assertTrue(bundleColorPackage  + " should be enabled",
+                    overlays.values().contains(bundleColorPackage));
+            assertTrue(bundleFontPackage  + " should be enabled",
+                    overlays.values().contains(bundleFontPackage));
+            assertFalse(otherPackage  + " should not be enabled",
+                    overlays.values().contains(otherPackage));
+            return true;
+        });
+
+        assertEquals("Secure Setting was not properly set after applying theme",
+                theme.getSerializedPackages(),
+                Settings.Secure.getString(mActivity.getContentResolver(), THEME_SETTING));
+    }
+
+    private void applyThemeAndWaitForCondition(ThemeBundle theme, String message,
+            Condition condition) {
+        boolean[] done = {false, false};
+        mThemeManager.apply(theme, new Callback() {
+            @Override
+            public void onSuccess() {
+                done[0] = true;
+            }
+
+            @Override
+            public void onError(@Nullable Throwable throwable) {
+                done[0] = true;
+                done[1] = true;
+            }
+        });
+        // TODO: refactor these tests so we can get rid of the long wait.
+        Wait.atMost(message, () -> done[0] && condition.isTrue(), 1000);
+        // done[1] is only set to true in the onError callback.
+        assertFalse(done[1]);
+    }
+}
diff --git a/robolectric_tests/src/com/android/customization/testutils/Condition.java b/robolectric_tests/src/com/android/customization/testutils/Condition.java
new file mode 100644
index 0000000..f013555
--- /dev/null
+++ b/robolectric_tests/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/robolectric_tests/src/com/android/customization/testutils/OverlayManagerMocks.java b/robolectric_tests/src/com/android/customization/testutils/OverlayManagerMocks.java
new file mode 100644
index 0000000..ed77224
--- /dev/null
+++ b/robolectric_tests/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 packageName;
+        final String targetPackage;
+        final String category;
+
+        public MockOverlay(String packageName, String targetPackage, String category) {
+            this.packageName = packageName;
+            this.targetPackage = targetPackage;
+            this.category = category;
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            return obj instanceof MockOverlay
+                    && TextUtils.equals(((MockOverlay) obj).packageName, packageName)
+                    && TextUtils.equals(((MockOverlay) obj).targetPackage, targetPackage)
+                    && TextUtils.equals(((MockOverlay) obj).category, category);
+        }
+    }
+
+    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.packageName.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.targetPackage.equals(inv.getArgument(0))
+                                            && mockOverlay.category.equals(inv.getArgument(1)))
+                                .map(overlay -> overlay.packageName).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.targetPackage))
+                                .collect(Collectors.toMap(
+                                        overlay ->
+                                                overlay.category,
+                                        (Function<MockOverlay, String>) overlay ->
+                                                overlay.packageName))
+        );
+    }
+}
\ No newline at end of file
diff --git a/robolectric_tests/src/com/android/customization/testutils/Wait.java b/robolectric_tests/src/com/android/customization/testutils/Wait.java
new file mode 100644
index 0000000..54650ba
--- /dev/null
+++ b/robolectric_tests/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);
+    }
+}
+
diff --git a/src/com/android/customization/model/ResourceConstants.java b/src/com/android/customization/model/ResourceConstants.java
index d750e14..2822f4a 100644
--- a/src/com/android/customization/model/ResourceConstants.java
+++ b/src/com/android/customization/model/ResourceConstants.java
@@ -15,6 +15,8 @@
  */
 package com.android.customization.model;
 
+import android.provider.Settings.Secure;
+
 /**
  * Holds common strings used to reference system resources.
  */
@@ -40,4 +42,18 @@
      */
     String CONFIG_ICON_MASK = "config_icon_mask";
 
+    /**
+     * Overlay Categories that theme picker handles.
+     */
+    String OVERLAY_CATEGORY_COLOR = "android.theme.customization.accent_color";
+    String OVERLAY_CATEGORY_FONT = "android.theme.customization.font";
+    String OVERLAY_CATEGORY_SHAPE = "android.theme.customization.adaptive_icon_shape";
+    String OVERLAY_CATEGORY_ICON_ANDROID = "android.theme.customization.icon_pack.android";
+    String OVERLAY_CATEGORY_ICON_SETTINGS = "android.theme.customization.icon_pack.settings";
+    String OVERLAY_CATEGORY_ICON_SYSUI = "android.theme.customization.icon_pack.systemui";
+
+    /**
+     * Secure Setting used to store the currently set theme.
+     */
+    String THEME_SETTING = Secure.THEME_CUSTOMIZATION_OVERLAY_PACKAGES;
 }
diff --git a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
index a824b28..d0b8d01 100644
--- a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
+++ b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
@@ -57,16 +57,14 @@
     private final ProviderInfo mProviderInfo;
     private List<GridOption> mOptions;
 
-    public LauncherGridOptionsProvider(Context context) {
+    public LauncherGridOptionsProvider(Context context, String authorityMetadataKey) {
         mContext = context;
-        Intent homeIntent =  new Intent(Intent.ACTION_MAIN)
-                .addCategory(Intent.CATEGORY_HOME);
+        Intent homeIntent = new Intent(Intent.ACTION_MAIN).addCategory(Intent.CATEGORY_HOME);
 
         ResolveInfo info = context.getPackageManager().resolveActivity(homeIntent,
                 PackageManager.MATCH_DEFAULT_ONLY | PackageManager.GET_META_DATA);
         if (info != null && info.activityInfo != null && info.activityInfo.metaData != null) {
-            mGridProviderAuthority = info.activityInfo.metaData.getString(
-                    mContext.getString(R.string.grid_control_metadata_name));
+            mGridProviderAuthority = info.activityInfo.metaData.getString(authorityMetadataKey);
         } else {
             mGridProviderAuthority = null;
         }
diff --git a/src/com/android/customization/model/theme/OverlayManagerCompat.java b/src/com/android/customization/model/theme/OverlayManagerCompat.java
new file mode 100644
index 0000000..a32f120
--- /dev/null
+++ b/src/com/android/customization/model/theme/OverlayManagerCompat.java
@@ -0,0 +1,96 @@
+/*
+ * 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 android.content.Context;
+import android.content.om.OverlayInfo;
+import android.content.om.OverlayManager;
+import android.os.UserHandle;
+
+import androidx.annotation.Nullable;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Wrapper over {@link OverlayManager} that abstracts away its internals and can be mocked for
+ * testing.
+ */
+public class OverlayManagerCompat {
+    private final OverlayManager mOverlayManager;
+
+    public OverlayManagerCompat(Context context) {
+        mOverlayManager = context.getSystemService(OverlayManager.class);
+    }
+
+    /**
+     * Enables the overlay provided by the given package for the given user Id
+     * @return true if the operation succeeded
+     */
+    public boolean setEnabledExclusiveInCategory(String packageName, int userId) {
+        return mOverlayManager.setEnabledExclusiveInCategory(packageName, userId);
+    }
+
+    /**
+     * Disables the overlay provided by the given package for the given user Id
+     * @return true if the operation succeeded
+     */
+    public boolean disableOverlay(String packageName, int userId) {
+        return mOverlayManager.setEnabled(packageName, false, userId);
+    }
+
+    /**
+     * @return the package name of the currently enabled overlay for the given target package, in
+     * the given category, or {@code null} if none is currently enabled.
+     */
+    @Nullable
+    public String getEnabledPackageName(String targetPackageName, String category) {
+        List<OverlayInfo> overlayInfos = getOverlayInfosForTarget(targetPackageName,
+                UserHandle.myUserId());
+        for (OverlayInfo overlayInfo : overlayInfos) {
+            if (category.equals(overlayInfo.category) && overlayInfo.isEnabled()) {
+                return overlayInfo.packageName;
+            }
+        }
+        return null;
+    }
+
+    /**
+     * @return a Map of Category -> PackageName of all the overlays enabled for the given target
+     * packages. It might be empty if no overlay is enabled for those targets.
+     */
+    public Map<String, String> getEnabledOverlaysForTargets(String... targetPackages) {
+        Map<String, String> overlays = new HashMap<>();
+        for (String target : targetPackages) {
+            addAllEnabledOverlaysForTarget(overlays, target);
+        }
+        return overlays;
+    }
+
+    private List<OverlayInfo> getOverlayInfosForTarget(String targetPackageName, int userId) {
+        return mOverlayManager.getOverlayInfosForTarget(targetPackageName, userId);
+    }
+
+    private void addAllEnabledOverlaysForTarget(Map<String, String> overlays, String target) {
+        for (OverlayInfo overlayInfo : getOverlayInfosForTarget(target,
+                UserHandle.myUserId())) {
+            if (overlayInfo.isEnabled()) {
+                overlays.put(overlayInfo.category, overlayInfo.packageName);
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/com/android/customization/model/theme/ThemeManager.java b/src/com/android/customization/model/theme/ThemeManager.java
index 31e76ca..a37893b 100644
--- a/src/com/android/customization/model/theme/ThemeManager.java
+++ b/src/com/android/customization/model/theme/ThemeManager.java
@@ -20,8 +20,6 @@
 import static com.android.customization.model.ResourceConstants.SYSUI_PACKAGE;
 
 import android.app.Activity;
-import android.content.om.OverlayInfo;
-import android.content.om.OverlayManager;
 import android.graphics.Point;
 import android.os.UserHandle;
 import android.provider.Settings;
@@ -29,57 +27,43 @@
 import androidx.annotation.Nullable;
 
 import com.android.customization.model.CustomizationManager;
+import com.android.customization.model.ResourceConstants;
 import com.android.wallpaper.asset.Asset;
 import com.android.wallpaper.module.WallpaperPersister;
 import com.android.wallpaper.module.WallpaperPersister.SetWallpaperCallback;
 import com.android.wallpaper.module.WallpaperSetter;
 import com.android.wallpaper.util.WallpaperCropUtils;
 
-import java.util.HashMap;
 import java.util.HashSet;
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
 public class ThemeManager implements CustomizationManager<ThemeBundle> {
 
-    private static final String OVERLAY_CATEGORY_COLOR = "android.theme.customization.accent_color";
-    private static final String OVERLAY_CATEGORY_FONT = "android.theme.customization.font";
-    private static final String OVERLAY_CATEGORY_SHAPE =
-            "android.theme.customization.adaptive_icon_shape";
-    private static final String OVERLAY_CATEGORY_ICON_ANDROID =
-            "android.theme.customization.icon_pack.android";
-    private static final String OVERLAY_CATEGORY_ICON_SETTINGS =
-            "android.theme.customization.icon_pack.settings";
-    private static final String OVERLAY_CATEGORY_ICON_SYSUI =
-            "android.theme.customization.icon_pack.systemui";
-
     private static final Set<String> THEME_CATEGORIES = new HashSet<>();
     static {
-        THEME_CATEGORIES.add(OVERLAY_CATEGORY_COLOR);
-        THEME_CATEGORIES.add(OVERLAY_CATEGORY_FONT);
-        THEME_CATEGORIES.add(OVERLAY_CATEGORY_SHAPE);
-        THEME_CATEGORIES.add(OVERLAY_CATEGORY_ICON_ANDROID);
-        THEME_CATEGORIES.add(OVERLAY_CATEGORY_ICON_SETTINGS);
-        THEME_CATEGORIES.add(OVERLAY_CATEGORY_ICON_SYSUI);
+        THEME_CATEGORIES.add(ResourceConstants.OVERLAY_CATEGORY_COLOR);
+        THEME_CATEGORIES.add(ResourceConstants.OVERLAY_CATEGORY_FONT);
+        THEME_CATEGORIES.add(ResourceConstants.OVERLAY_CATEGORY_SHAPE);
+        THEME_CATEGORIES.add(ResourceConstants.OVERLAY_CATEGORY_ICON_ANDROID);
+        THEME_CATEGORIES.add(ResourceConstants.OVERLAY_CATEGORY_ICON_SETTINGS);
+        THEME_CATEGORIES.add(ResourceConstants.OVERLAY_CATEGORY_ICON_SYSUI);
     };
 
-    //TODO: replace with System.Secure constant
-    private static final String THEME_SETTING = "theme_customization_overlay_packages";
-
 
     private final ThemeBundleProvider mProvider;
-    private final OverlayManager mOverlayManager;
+    private final OverlayManagerCompat mOverlayManagerCompat;
+
     private final WallpaperSetter mWallpaperSetter;
     private final Activity mActivity;
 
     private Map<String, String> mCurrentOverlays;
 
     public ThemeManager(ThemeBundleProvider provider, Activity activity,
-            WallpaperSetter wallpaperSetter) {
+            WallpaperSetter wallpaperSetter, OverlayManagerCompat overlayManagerCompat) {
         mProvider = provider;
         mActivity = activity;
-        mOverlayManager = activity.getSystemService(OverlayManager.class);
+        mOverlayManagerCompat = overlayManagerCompat;
         mWallpaperSetter = wallpaperSetter;
     }
 
@@ -131,22 +115,22 @@
     private void applyOverlays(ThemeBundle theme, Callback callback) {
         boolean allApplied = true;
         if (theme.isDefault()) {
-            allApplied &= disableCurrentOverlay(ANDROID_PACKAGE, OVERLAY_CATEGORY_COLOR);
-            allApplied &= disableCurrentOverlay(ANDROID_PACKAGE, OVERLAY_CATEGORY_FONT);
-            allApplied &= disableCurrentOverlay(ANDROID_PACKAGE, OVERLAY_CATEGORY_SHAPE);
-            allApplied &= disableCurrentOverlay(ANDROID_PACKAGE, OVERLAY_CATEGORY_ICON_ANDROID);
-            allApplied &= disableCurrentOverlay(SYSUI_PACKAGE, OVERLAY_CATEGORY_ICON_SYSUI);
-            allApplied &= disableCurrentOverlay(SETTINGS_PACKAGE, OVERLAY_CATEGORY_ICON_SETTINGS);
+            allApplied &= disableCurrentOverlay(ANDROID_PACKAGE, ResourceConstants.OVERLAY_CATEGORY_COLOR);
+            allApplied &= disableCurrentOverlay(ANDROID_PACKAGE, ResourceConstants.OVERLAY_CATEGORY_FONT);
+            allApplied &= disableCurrentOverlay(ANDROID_PACKAGE, ResourceConstants.OVERLAY_CATEGORY_SHAPE);
+            allApplied &= disableCurrentOverlay(ANDROID_PACKAGE, ResourceConstants.OVERLAY_CATEGORY_ICON_ANDROID);
+            allApplied &= disableCurrentOverlay(SYSUI_PACKAGE, ResourceConstants.OVERLAY_CATEGORY_ICON_SYSUI);
+            allApplied &= disableCurrentOverlay(SETTINGS_PACKAGE, ResourceConstants.OVERLAY_CATEGORY_ICON_SETTINGS);
         } else {
             for (String packageName : theme.getAllPackages()) {
                 if (packageName != null) {
-                    allApplied &= mOverlayManager.setEnabledExclusiveInCategory(packageName,
+                    allApplied &= mOverlayManagerCompat.setEnabledExclusiveInCategory(packageName,
                             UserHandle.myUserId());
                 }
             }
         }
         allApplied &= Settings.Secure.putString(mActivity.getContentResolver(),
-                THEME_SETTING, theme.getSerializedPackages());
+                ResourceConstants.THEME_SETTING, theme.getSerializedPackages());
         mCurrentOverlays = null;
         if (allApplied) {
             callback.onSuccess();
@@ -160,46 +144,26 @@
         mProvider.fetch(callback, false);
     }
 
-    private boolean disableCurrentOverlay(String packageName, String category) {
-        OverlayInfo current = getEnabledOverlayInfo(packageName, category);
-        if (current != null) {
-           return mOverlayManager.setEnabled(current.packageName, false, UserHandle.myUserId());
+    private boolean disableCurrentOverlay(String targetPackage, String category) {
+        String currentPackageName = mOverlayManagerCompat.getEnabledPackageName(targetPackage,
+                category);
+        if (currentPackageName != null) {
+           return mOverlayManagerCompat.disableOverlay(currentPackageName, UserHandle.myUserId());
         }
         return true;
     }
 
-    @Nullable
-    private OverlayInfo getEnabledOverlayInfo(String packageName, String category) {
-        List<OverlayInfo> overlayInfos = mOverlayManager
-                .getOverlayInfosForTarget(packageName, UserHandle.myUserId());
-        for (OverlayInfo overlayInfo : overlayInfos) {
-            if (category.equals(overlayInfo.category) && overlayInfo.isEnabled()) {
-                return overlayInfo;
-            }
-        }
-        return null;
-    }
-
     public Map<String, String> getCurrentOverlays() {
         if (mCurrentOverlays == null) {
-            mCurrentOverlays = new HashMap<>();
-            addAllEnabledOverlaysForPackage(ANDROID_PACKAGE);
-            addAllEnabledOverlaysForPackage(SYSUI_PACKAGE);
-            addAllEnabledOverlaysForPackage(SETTINGS_PACKAGE);
+            mCurrentOverlays = mOverlayManagerCompat.getEnabledOverlaysForTargets(ANDROID_PACKAGE,
+                    SYSUI_PACKAGE, SETTINGS_PACKAGE);
+            mCurrentOverlays.entrySet().removeIf(
+                    categoryAndPackage -> !THEME_CATEGORIES.contains(categoryAndPackage.getKey()));
         }
         return mCurrentOverlays;
     }
 
-    private void addAllEnabledOverlaysForPackage(String targetPackage) {
-        for (OverlayInfo overlayInfo :
-                mOverlayManager.getOverlayInfosForTarget(targetPackage, UserHandle.myUserId())) {
-            if (overlayInfo.isEnabled() && THEME_CATEGORIES.contains(overlayInfo.category)) {
-                mCurrentOverlays.put(overlayInfo.category, overlayInfo.packageName);
-            }
-        }
-    }
-
     public String getStoredOverlays() {
-        return Settings.Secure.getString(mActivity.getContentResolver(), THEME_SETTING);
+        return Settings.Secure.getString(mActivity.getContentResolver(), ResourceConstants.THEME_SETTING);
     }
 }
diff --git a/src/com/android/customization/picker/CustomizationPickerActivity.java b/src/com/android/customization/picker/CustomizationPickerActivity.java
index 7a35b7d..de19aa7 100644
--- a/src/com/android/customization/picker/CustomizationPickerActivity.java
+++ b/src/com/android/customization/picker/CustomizationPickerActivity.java
@@ -40,6 +40,7 @@
 import com.android.customization.model.grid.GridOptionsManager;
 import com.android.customization.model.grid.LauncherGridOptionsProvider;
 import com.android.customization.model.theme.DefaultThemeProvider;
+import com.android.customization.model.theme.OverlayManagerCompat;
 import com.android.customization.model.theme.ThemeBundle;
 import com.android.customization.model.theme.ThemeManager;
 import com.android.customization.picker.clock.ClockFragment;
@@ -159,7 +160,7 @@
         mWallpaperSetter = new WallpaperSetter(injector.getWallpaperPersister(this),
                 injector.getPreferences(this), mUserEventLogger, false);
         ThemeManager themeManager = new ThemeManager(new DefaultThemeProvider(this), this,
-                mWallpaperSetter);
+                mWallpaperSetter, new OverlayManagerCompat(this));
         if (themeManager.isAvailable()) {
             mSections.put(R.id.nav_theme, new ThemeSection(R.id.nav_theme, themeManager));
         }
@@ -171,7 +172,8 @@
         }
         //Grid
         GridOptionsManager gridManager = new GridOptionsManager(
-                new LauncherGridOptionsProvider(this));
+                new LauncherGridOptionsProvider(this,
+                        getString(R.string.grid_control_metadata_name)));
         if (gridManager.isAvailable()) {
             mSections.put(R.id.nav_grid, new GridSection(R.id.nav_grid, gridManager));
         }