diff --git a/src/com/android/settings/development/EmulateDisplayCutoutPreferenceController.java b/src/com/android/settings/development/EmulateDisplayCutoutPreferenceController.java
index 9fa323a..973868b 100644
--- a/src/com/android/settings/development/EmulateDisplayCutoutPreferenceController.java
+++ b/src/com/android/settings/development/EmulateDisplayCutoutPreferenceController.java
@@ -19,18 +19,21 @@
 import static android.os.UserHandle.USER_SYSTEM;
 
 import android.content.Context;
+import android.content.om.IOverlayManager;
+import android.content.om.OverlayInfo;
 import android.content.pm.PackageManager;
+import android.os.RemoteException;
+import android.os.ServiceManager;
+import android.text.TextUtils;
+import android.view.DisplayCutout;
+
 import android.support.annotation.VisibleForTesting;
 import android.support.v7.preference.ListPreference;
 import android.support.v7.preference.Preference;
 import android.support.v7.preference.PreferenceScreen;
-import android.text.TextUtils;
-import android.view.DisplayCutout;
 
 import com.android.settings.R;
 import com.android.settings.core.PreferenceControllerMixin;
-import com.android.settings.wrapper.OverlayManagerWrapper;
-import com.android.settings.wrapper.OverlayManagerWrapper.OverlayInfo;
 import com.android.settingslib.development.DeveloperOptionsPreferenceController;
 
 import java.util.List;
@@ -41,7 +44,7 @@
 
     private static final String KEY = "display_cutout_emulation";
 
-    private final OverlayManagerWrapper mOverlayManager;
+    private final IOverlayManager mOverlayManager;
     private final boolean mAvailable;
 
     private ListPreference mPreference;
@@ -49,7 +52,7 @@
 
     @VisibleForTesting
     EmulateDisplayCutoutPreferenceController(Context context, PackageManager packageManager,
-            OverlayManagerWrapper overlayManager) {
+            IOverlayManager overlayManager) {
         super(context);
         mOverlayManager = overlayManager;
         mPackageManager = packageManager;
@@ -57,7 +60,8 @@
     }
 
     public EmulateDisplayCutoutPreferenceController(Context context) {
-        this(context, context.getPackageManager(), new OverlayManagerWrapper());
+        this(context, context.getPackageManager(), IOverlayManager.Stub
+                .asInterface(ServiceManager.getService(Context.OVERLAY_SERVICE)));
     }
 
     @Override
@@ -102,10 +106,14 @@
         }
 
         final boolean result;
-        if (TextUtils.isEmpty(packageName)) {
-            result = mOverlayManager.setEnabled(currentPackageName, false, USER_SYSTEM);
-        } else {
-            result = mOverlayManager.setEnabledExclusiveInCategory(packageName, USER_SYSTEM);
+        try {
+            if (TextUtils.isEmpty(packageName)) {
+                result = mOverlayManager.setEnabled(currentPackageName, false, USER_SYSTEM);
+            } else {
+                result = mOverlayManager.setEnabledExclusiveInCategory(packageName, USER_SYSTEM);
+            }
+        } catch (RemoteException re) {
+            throw re.rethrowFromSystemServer();
         }
         updateState(mPreference);
         return result;
@@ -145,13 +153,17 @@
     }
 
     private OverlayInfo[] getOverlayInfos() {
-        @SuppressWarnings("unchecked") List<OverlayInfo> overlayInfos =
-                mOverlayManager.getOverlayInfosForTarget("android", USER_SYSTEM);
-        for (int i = overlayInfos.size() - 1; i >= 0; i--) {
-            if (!DisplayCutout.EMULATION_OVERLAY_CATEGORY.equals(
-                    overlayInfos.get(i).category)) {
-                overlayInfos.remove(i);
+        List<OverlayInfo> overlayInfos;
+        try {
+            overlayInfos = mOverlayManager.getOverlayInfosForTarget("android", USER_SYSTEM);
+            for (int i = overlayInfos.size() - 1; i >= 0; i--) {
+                if (!DisplayCutout.EMULATION_OVERLAY_CATEGORY.equals(
+                        overlayInfos.get(i).category)) {
+                    overlayInfos.remove(i);
+                }
             }
+        } catch (RemoteException re) {
+            throw re.rethrowFromSystemServer();
         }
         return overlayInfos.toArray(new OverlayInfo[overlayInfos.size()]);
     }
diff --git a/src/com/android/settings/display/ThemePreferenceController.java b/src/com/android/settings/display/ThemePreferenceController.java
index 1bc49b6..c056834 100644
--- a/src/com/android/settings/display/ThemePreferenceController.java
+++ b/src/com/android/settings/display/ThemePreferenceController.java
@@ -13,7 +13,11 @@
  */
 package com.android.settings.display;
 
+import static com.android.internal.logging.nano.MetricsProto.MetricsEvent.ACTION_THEME;
+
 import android.content.Context;
+import android.content.om.IOverlayManager;
+import android.content.om.OverlayInfo;
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.PackageManager.NameNotFoundException;
@@ -23,13 +27,12 @@
 import android.support.annotation.VisibleForTesting;
 import android.support.v7.preference.ListPreference;
 import android.support.v7.preference.Preference;
+
 import android.text.TextUtils;
 
 import com.android.settings.R;
 import com.android.settings.core.PreferenceControllerMixin;
 import com.android.settings.overlay.FeatureFactory;
-import com.android.settings.wrapper.OverlayManagerWrapper;
-import com.android.settings.wrapper.OverlayManagerWrapper.OverlayInfo;
 import com.android.settingslib.core.AbstractPreferenceController;
 import com.android.settingslib.core.instrumentation.MetricsFeatureProvider;
 
@@ -37,24 +40,22 @@
 import java.util.List;
 import java.util.Objects;
 
-import static com.android.internal.logging.nano.MetricsProto.MetricsEvent.ACTION_THEME;
-
 public class ThemePreferenceController extends AbstractPreferenceController implements
         PreferenceControllerMixin, Preference.OnPreferenceChangeListener {
 
     private static final String KEY_THEME = "theme";
 
     private final MetricsFeatureProvider mMetricsFeatureProvider;
-    private final OverlayManagerWrapper mOverlayService;
+    private final IOverlayManager mOverlayService;
     private final PackageManager mPackageManager;
 
     public ThemePreferenceController(Context context) {
-        this(context, ServiceManager.getService(Context.OVERLAY_SERVICE) != null
-                ? new OverlayManagerWrapper() : null);
+        this(context, IOverlayManager.Stub
+                .asInterface(ServiceManager.getService(Context.OVERLAY_SERVICE)));
     }
 
     @VisibleForTesting
-    ThemePreferenceController(Context context, OverlayManagerWrapper overlayManager) {
+    ThemePreferenceController(Context context, IOverlayManager overlayManager) {
         super(context);
         mOverlayService = overlayManager;
         mPackageManager = context.getPackageManager();
@@ -77,7 +78,7 @@
     @Override
     public void updateState(Preference preference) {
         ListPreference pref = (ListPreference) preference;
-        String[] pkgs = getAvailableThemes();
+        String[] pkgs = getAvailableThemes(false /* currentThemeOnly */);
         CharSequence[] labels = new CharSequence[pkgs.length];
         for (int i = 0; i < pkgs.length; i++) {
             try {
@@ -109,11 +110,15 @@
 
     @Override
     public boolean onPreferenceChange(Preference preference, Object newValue) {
-        String current = getTheme();
+        String current = getCurrentTheme();
         if (Objects.equals(newValue, current)) {
             return true;
         }
-        mOverlayService.setEnabledExclusiveInCategory((String) newValue, UserHandle.myUserId());
+        try {
+            mOverlayService.setEnabledExclusiveInCategory((String) newValue, UserHandle.myUserId());
+        } catch (RemoteException re) {
+            throw re.rethrowFromSystemServer();
+        }
         return true;
     }
 
@@ -129,39 +134,43 @@
         }
     }
 
-    private String getTheme() {
-        List<OverlayInfo> infos = mOverlayService.getOverlayInfosForTarget("android",
-                UserHandle.myUserId());
-        for (int i = 0, size = infos.size(); i < size; i++) {
-            if (infos.get(i).isEnabled() && isTheme(infos.get(i))) {
-                return infos.get(i).packageName;
-            }
-        }
-        return null;
-    }
-
     @Override
     public boolean isAvailable() {
         if (mOverlayService == null) return false;
-        String[] themes = getAvailableThemes();
+        String[] themes = getAvailableThemes(false /* currentThemeOnly */);
         return themes != null && themes.length > 1;
     }
 
 
     @VisibleForTesting
     String getCurrentTheme() {
-        return getTheme();
+        String[] themePackages = getAvailableThemes(true /* currentThemeOnly */);
+        return themePackages.length < 1 ? null : themePackages[0];
     }
 
     @VisibleForTesting
-    String[] getAvailableThemes() {
-        List<OverlayInfo> infos = mOverlayService.getOverlayInfosForTarget("android",
-                UserHandle.myUserId());
-        List<String> pkgs = new ArrayList<>(infos.size());
-        for (int i = 0, size = infos.size(); i < size; i++) {
-            if (isTheme(infos.get(i))) {
-                pkgs.add(infos.get(i).packageName);
+    String[] getAvailableThemes(boolean currentThemeOnly) {
+        List<OverlayInfo> infos;
+        List<String> pkgs;
+        try {
+            infos = mOverlayService.getOverlayInfosForTarget("android", UserHandle.myUserId());
+            pkgs = new ArrayList<>(infos.size());
+            for (int i = 0, size = infos.size(); i < size; i++) {
+                if (isTheme(infos.get(i))) {
+                    if (infos.get(i).isEnabled() && currentThemeOnly) {
+                        return new String[] {infos.get(i).packageName};
+                    } else {
+                        pkgs.add(infos.get(i).packageName);
+                    }
+                }
             }
+        } catch (RemoteException re) {
+            throw re.rethrowFromSystemServer();
+        }
+
+        // Current enabled theme is not found.
+        if (currentThemeOnly) {
+            return new String[0];
         }
         return pkgs.toArray(new String[pkgs.size()]);
     }
diff --git a/src/com/android/settings/wrapper/OverlayManagerWrapper.java b/src/com/android/settings/wrapper/OverlayManagerWrapper.java
deleted file mode 100644
index 6e3c234..0000000
--- a/src/com/android/settings/wrapper/OverlayManagerWrapper.java
+++ /dev/null
@@ -1,102 +0,0 @@
-/*
- * Copyright (C) 2018 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.settings.wrapper;
-
-import android.content.Context;
-import android.content.om.IOverlayManager;
-import android.content.om.OverlayInfo;
-import android.os.RemoteException;
-import android.os.ServiceManager;
-
-import java.util.ArrayList;
-import java.util.List;
-
-public class OverlayManagerWrapper {
-
-    private final IOverlayManager mOverlayManager;
-
-    public OverlayManagerWrapper(IOverlayManager overlayManager) {
-        mOverlayManager = overlayManager;
-    }
-
-    public OverlayManagerWrapper() {
-        this(IOverlayManager.Stub.asInterface(ServiceManager.getService(Context.OVERLAY_SERVICE)));
-    }
-
-    public List<OverlayInfo> getOverlayInfosForTarget(String overlay, int userId) {
-        if (mOverlayManager == null) {
-            return new ArrayList<>();
-        }
-        try {
-            List<android.content.om.OverlayInfo> infos
-                    = mOverlayManager.getOverlayInfosForTarget(overlay, userId);
-            ArrayList<OverlayInfo> result = new ArrayList<>(infos.size());
-            for (int i = 0; i < infos.size(); i++) {
-                result.add(new OverlayInfo(infos.get(i)));
-            }
-            return result;
-        } catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }
-    }
-
-    public boolean setEnabled(String overlay, boolean enabled, int userId) {
-        if (mOverlayManager == null) {
-            return false;
-        }
-        try {
-            return mOverlayManager.setEnabled(overlay, enabled, userId);
-        } catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }
-    }
-
-    public boolean setEnabledExclusiveInCategory(String overlay, int userId) {
-        if (mOverlayManager == null) {
-            return false;
-        }
-        try {
-            return mOverlayManager.setEnabledExclusiveInCategory(overlay, userId);
-        } catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }
-    }
-
-    public static class OverlayInfo {
-
-        public static final String CATEGORY_THEME = android.content.om.OverlayInfo.CATEGORY_THEME;
-        public final String packageName;
-        public final String category;
-        private final boolean mEnabled;
-
-        public OverlayInfo(String packageName, String category, boolean enabled) {
-            this.packageName = packageName;
-            this.category = category;
-            mEnabled = enabled;
-        }
-
-        public OverlayInfo(android.content.om.OverlayInfo info) {
-            mEnabled = info.isEnabled();
-            category = info.category;
-            packageName = info.packageName;
-        }
-
-        public boolean isEnabled() {
-            return mEnabled;
-        }
-    }
-}
diff --git a/tests/robotests/src/com/android/settings/development/EmulateDisplayCutoutPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/development/EmulateDisplayCutoutPreferenceControllerTest.java
index 73c976c..a742988 100644
--- a/tests/robotests/src/com/android/settings/development/EmulateDisplayCutoutPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/development/EmulateDisplayCutoutPreferenceControllerTest.java
@@ -26,14 +26,15 @@
 import static org.mockito.Mockito.when;
 
 import android.content.Context;
+import android.content.om.IOverlayManager;
+import android.content.om.OverlayInfo;
 import android.content.pm.PackageManager;
 import android.support.v7.preference.ListPreference;
 import android.support.v7.preference.PreferenceScreen;
+import android.os.RemoteException;
 import android.view.DisplayCutout;
 
 import com.android.settings.testutils.SettingsRobolectricTestRunner;
-import com.android.settings.wrapper.OverlayManagerWrapper;
-import com.android.settings.wrapper.OverlayManagerWrapper.OverlayInfo;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -41,6 +42,7 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 
 @RunWith(SettingsRobolectricTestRunner.class)
@@ -54,7 +56,7 @@
     @Mock
     private Context mContext;
     @Mock
-    private OverlayManagerWrapper mOverlayManager;
+    private IOverlayManager mOverlayManager;
     @Mock
     private PackageManager mPackageManager;
     @Mock
@@ -64,6 +66,7 @@
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
+        when(mContext.getSystemService(Context.OVERLAY_SERVICE)).thenReturn(mOverlayManager);
         mockCurrentOverlays();
         when(mPackageManager.getApplicationInfo(any(), anyInt()))
             .thenThrow(PackageManager.NameNotFoundException.class);
@@ -72,8 +75,12 @@
     }
 
     Object mockCurrentOverlays(OverlayInfo... overlays) {
-        return when(mOverlayManager.getOverlayInfosForTarget(eq("android"), anyInt()))
-            .thenReturn(Arrays.asList(overlays));
+        try {
+            return when(mOverlayManager.getOverlayInfosForTarget(eq("android"), anyInt()))
+                .thenReturn(Arrays.asList(overlays));
+        } catch (RemoteException re) {
+            return new ArrayList<OverlayInfo>();
+        }
     }
 
     @Test
@@ -146,6 +153,15 @@
     }
 
     private static OverlayInfo createFakeOverlay(String pkg, boolean enabled) {
-        return new OverlayInfo(pkg, DisplayCutout.EMULATION_OVERLAY_CATEGORY, enabled);
+        final int state = (enabled) ? OverlayInfo.STATE_ENABLED : OverlayInfo.STATE_DISABLED;
+
+        return new OverlayInfo(pkg /* packageName */,
+                pkg + ".target" /* targetPackageName */,
+                DisplayCutout.EMULATION_OVERLAY_CATEGORY /* category */,
+                pkg + ".baseCodePath" /* baseCodePath */,
+                state /* state */,
+                0 /* userId */,
+                0 /* priority */,
+                true /* isStatic */);
     }
 }
\ No newline at end of file
diff --git a/tests/robotests/src/com/android/settings/display/ThemePreferenceControllerTest.java b/tests/robotests/src/com/android/settings/display/ThemePreferenceControllerTest.java
index 6a059b3..56e9c2d 100644
--- a/tests/robotests/src/com/android/settings/display/ThemePreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/display/ThemePreferenceControllerTest.java
@@ -16,26 +16,28 @@
 
 package com.android.settings.display;
 
+import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
 import static org.mockito.Matchers.anyString;
 import static org.mockito.Mockito.doReturn;
-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.content.Context;
+import android.content.om.IOverlayManager;
+import android.content.om.OverlayInfo;
 import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.PackageManager.NameNotFoundException;
 import android.support.v7.preference.ListPreference;
 
+
 import com.android.settings.R;
 import com.android.settings.testutils.FakeFeatureFactory;
 import com.android.settings.testutils.SettingsRobolectricTestRunner;
-import com.android.settings.wrapper.OverlayManagerWrapper;
-
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -44,6 +46,8 @@
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RuntimeEnvironment;
 
+import java.util.Arrays;
+
 @RunWith(SettingsRobolectricTestRunner.class)
 public class ThemePreferenceControllerTest {
 
@@ -55,6 +59,8 @@
     private ApplicationInfo mApplicationInfo;
     @Mock
     private ListPreference mPreference;
+    @Mock
+    private IOverlayManager mOverlayManager;
 
     private ThemePreferenceController mController;
 
@@ -67,8 +73,28 @@
         when(mContext.getString(R.string.default_theme))
                 .thenReturn(RuntimeEnvironment.application.getString(R.string.default_theme));
 
-        mController =
-            spy(new ThemePreferenceController(mContext, mock(OverlayManagerWrapper.class)));
+        when(mContext.getSystemService(Context.OVERLAY_SERVICE)).thenReturn(mOverlayManager);
+        mController = spy(new ThemePreferenceController(mContext, mOverlayManager));
+    }
+
+    @Test
+    public void testAvailable_false() throws Exception {
+        when(mPackageManager.getPackageInfo(anyString(), anyInt())).thenReturn(
+                new PackageInfo());
+        when(mOverlayManager.getOverlayInfosForTarget(any(), anyInt()))
+                .thenReturn(Arrays.asList(new OverlayInfo("", "", "", "", 0, 0, 0, false)));
+        assertThat(mController.isAvailable()).isFalse();
+    }
+
+    @Test
+    public void testAvailable_true() throws Exception {
+        when(mPackageManager.getPackageInfo(anyString(), anyInt())).thenReturn(
+                new PackageInfo());
+        when(mOverlayManager.getOverlayInfosForTarget(any(), anyInt()))
+                .thenReturn(Arrays.asList(
+                        new OverlayInfo("", "", OverlayInfo.CATEGORY_THEME, "", 0, 0, 0, true),
+                        new OverlayInfo("", "", OverlayInfo.CATEGORY_THEME, "", 0, 0, 0, true)));
+        assertThat(mController.isAvailable()).isTrue();
     }
 
     @Test
@@ -79,7 +105,7 @@
         final String themeLabel2 = "Theme2";
         final String[] themes = {pkg1, pkg2};
         doReturn("pkg1.theme1").when(mController).getCurrentTheme();
-        doReturn(themes).when(mController).getAvailableThemes();
+        doReturn(themes).when(mController).getAvailableThemes(false /* currentThemeOnly */);
         when(mPackageManager.getApplicationInfo(anyString(), anyInt()).loadLabel(mPackageManager))
                 .thenReturn(themeLabel1)
                 .thenReturn(themeLabel2);
@@ -98,7 +124,7 @@
         final String themeLabel2 = "Theme2";
         final String[] themes = {pkg1, pkg2};
         doReturn(null).when(mController).getCurrentTheme();
-        doReturn(themes).when(mController).getAvailableThemes();
+        doReturn(themes).when(mController).getAvailableThemes(false /* currentThemeOnly */);
         when(mPackageManager.getApplicationInfo(anyString(), anyInt()).loadLabel(mPackageManager))
                 .thenReturn(themeLabel1)
                 .thenReturn(themeLabel2);
@@ -109,4 +135,32 @@
                 .setSummary(RuntimeEnvironment.application.getString(R.string.default_theme));
         verify(mPreference).setValue(null);
     }
+
+    @Test
+    public void getCurrentTheme_withEnabledState() throws Exception {
+        OverlayInfo info1 = new OverlayInfo("com.android.Theme1", "android",
+                OverlayInfo.CATEGORY_THEME, "", OverlayInfo.STATE_ENABLED, 0, 0, true);
+        OverlayInfo info2 = new OverlayInfo("com.android.Theme2", "android",
+                OverlayInfo.CATEGORY_THEME, "", 0, 0, 0, true);
+        when(mOverlayManager.getOverlayInfosForTarget(any(), anyInt())).thenReturn(
+                Arrays.asList(info1, info2));
+        when(mPackageManager.getPackageInfo(anyString(), anyInt())).thenReturn(
+                new PackageInfo());
+
+        assertThat(mController.getCurrentTheme()).isEqualTo(info1.packageName);
+    }
+
+    @Test
+    public void testGetCurrentTheme_withoutEnabledState() throws Exception {
+        OverlayInfo info1 = new OverlayInfo("com.android.Theme1", "android",
+                OverlayInfo.CATEGORY_THEME, "", OverlayInfo.STATE_DISABLED, 0, 0, true);
+        OverlayInfo info2 = new OverlayInfo("com.android.Theme2", "android",
+                OverlayInfo.CATEGORY_THEME, "", 0, 0, 0, true);
+        when(mOverlayManager.getOverlayInfosForTarget(any(), anyInt())).thenReturn(
+                Arrays.asList(info1, info2));
+        when(mPackageManager.getPackageInfo(anyString(), anyInt())).thenReturn(
+                new PackageInfo());
+
+        assertThat(mController.getCurrentTheme()).isNull();
+    }
 }
diff --git a/tests/robotests/src/com/android/settings/location/LocationEnablerTest.java b/tests/robotests/src/com/android/settings/location/LocationEnablerTest.java
index 0c0cd4c..f4973c0 100644
--- a/tests/robotests/src/com/android/settings/location/LocationEnablerTest.java
+++ b/tests/robotests/src/com/android/settings/location/LocationEnablerTest.java
@@ -43,7 +43,6 @@
 import com.android.settings.testutils.SettingsRobolectricTestRunner;
 import com.android.settings.testutils.shadow.ShadowSecureSettings;
 import com.android.settingslib.core.lifecycle.Lifecycle;
-import com.android.settingslib.wrapper.LocationManagerWrapper;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -62,7 +61,7 @@
 @RunWith(SettingsRobolectricTestRunner.class)
 @Config(shadows = {
     ShadowSecureSettings.class,
-    LocationEnablerTest.ShadowLocationManagerWrapper.class})
+    LocationEnablerTest.ShadowLocationManager.class})
 public class LocationEnablerTest {
 
     @Mock
@@ -272,8 +271,8 @@
         return intent -> TextUtils.equals(expected, intent.getAction());
     }
 
-    @Implements(value = LocationManagerWrapper.class)
-    public static class ShadowLocationManagerWrapper {
+    @Implements(value = LocationManager.class)
+    public static class ShadowLocationManager {
 
         @Implementation
         public void setLocationEnabledForUser(boolean enabled, UserHandle userHandle) {
diff --git a/tests/unit/src/com/android/settings/display/ThemePreferenceControllerTest.java b/tests/unit/src/com/android/settings/display/ThemePreferenceControllerTest.java
index 586c448..d4f0049 100644
--- a/tests/unit/src/com/android/settings/display/ThemePreferenceControllerTest.java
+++ b/tests/unit/src/com/android/settings/display/ThemePreferenceControllerTest.java
@@ -17,6 +17,7 @@
 package com.android.settings.display;
 
 import static com.google.common.truth.Truth.assertThat;
+
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
 import static org.mockito.Matchers.anyString;
@@ -26,6 +27,7 @@
 import static org.mockito.Mockito.when;
 
 import android.content.ContextWrapper;
+import android.content.om.IOverlayManager;
 import android.content.om.OverlayInfo;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageInfo;
@@ -35,8 +37,6 @@
 import android.support.test.runner.AndroidJUnit4;
 import android.support.v7.preference.ListPreference;
 
-import com.android.settings.wrapper.OverlayManagerWrapper;
-
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -48,14 +48,14 @@
 @RunWith(AndroidJUnit4.class)
 public class ThemePreferenceControllerTest {
 
-    private OverlayManagerWrapper mMockOverlayManager;
+    private IOverlayManager mMockOverlayManager;
     private ContextWrapper mContext;
     private ThemePreferenceController mPreferenceController;
     private PackageManager mMockPackageManager;
 
     @Before
     public void setup() {
-        mMockOverlayManager = mock(OverlayManagerWrapper.class);
+        mMockOverlayManager = mock(IOverlayManager.class);
         mMockPackageManager = mock(PackageManager.class);
         mContext = new ContextWrapper(InstrumentationRegistry.getTargetContext()) {
             @Override
@@ -69,9 +69,9 @@
     @Test
     public void testUpdateState() throws Exception {
         OverlayInfo info1 = new OverlayInfo("com.android.Theme1", "android",
-                "", "", OverlayInfo.STATE_ENABLED, 0, 0, true);
+                OverlayInfo.CATEGORY_THEME, "", OverlayInfo.STATE_ENABLED, 0, 0, true);
         OverlayInfo info2 = new OverlayInfo("com.android.Theme2", "android",
-                "", "", 0, 0, 0, true);
+                OverlayInfo.CATEGORY_THEME, "", 0, 0, 0, true);
         when(mMockPackageManager.getApplicationInfo(any(), anyInt())).thenAnswer(inv -> {
             ApplicationInfo info = mock(ApplicationInfo.class);
             if ("com.android.Theme1".equals(inv.getArguments()[0])) {
@@ -105,9 +105,9 @@
     @Test
     public void testUpdateState_withStaticOverlay() throws Exception {
         OverlayInfo info1 = new OverlayInfo("com.android.Theme1", "android",
-                "", "", OverlayInfo.STATE_ENABLED, 0, 0, true);
+                OverlayInfo.CATEGORY_THEME, "", OverlayInfo.STATE_ENABLED, 0, 0, true);
         OverlayInfo info2 = new OverlayInfo("com.android.Theme2", "android",
-                "", "", OverlayInfo.STATE_ENABLED, 0, 0, true);
+                OverlayInfo.CATEGORY_THEME, "", OverlayInfo.STATE_ENABLED, 0, 0, true);
         when(mMockPackageManager.getApplicationInfo(any(), anyInt())).thenAnswer(inv -> {
             ApplicationInfo info = mock(ApplicationInfo.class);
             if ("com.android.Theme1".equals(inv.getArguments()[0])) {
@@ -140,29 +140,10 @@
         verify(pref).setValue(eq("com.android.Theme2"));
     }
 
-    @Test
-    public void testAvailable_false() throws Exception {
-        when(mMockPackageManager.getPackageInfo(anyString(), anyInt())).thenReturn(
-                new PackageInfo());
-        when(mMockOverlayManager.getOverlayInfosForTarget(any(), anyInt()))
-                .thenReturn(list(new OverlayInfo("", "", "", "", 0, 0, 0, false)));
-        assertThat(mPreferenceController.isAvailable()).isFalse();
-    }
-
-    @Test
-    public void testAvailable_true() throws Exception {
-        when(mMockPackageManager.getPackageInfo(anyString(), anyInt())).thenReturn(
-                 new PackageInfo());
-        when(mMockOverlayManager.getOverlayInfosForTarget(any(), anyInt()))
-                .thenReturn(list(new OverlayInfo("", "", "", "", 0, 0, 0, true),
-                        new OverlayInfo("", "", "", "", 0, 0, 0, true)));
-        assertThat(mPreferenceController.isAvailable()).isTrue();
-    }
-
-    private ArrayList<OverlayManagerWrapper.OverlayInfo> list(OverlayInfo... infos) {
-        ArrayList<OverlayManagerWrapper.OverlayInfo> list = new ArrayList<>();
+    private ArrayList<OverlayInfo> list(OverlayInfo... infos) {
+        ArrayList<OverlayInfo> list = new ArrayList<>();
         for (OverlayInfo info : infos) {
-            list.add(new OverlayManagerWrapper.OverlayInfo(info));
+            list.add(info);
         }
         return list;
     }