Remove more wrappers

- WallpaperManagerWrapper
- UserManagerWrapper

Bug: 76167422
Test: robotests
Change-Id: I3ca866089082e01f1e9d7a4554387962f45a5437
diff --git a/src/com/android/settings/applications/appinfo/PictureInPictureSettings.java b/src/com/android/settings/applications/appinfo/PictureInPictureSettings.java
index e2caa41..ef2c7a9 100644
--- a/src/com/android/settings/applications/appinfo/PictureInPictureSettings.java
+++ b/src/com/android/settings/applications/appinfo/PictureInPictureSettings.java
@@ -40,7 +40,6 @@
 import com.android.settings.applications.AppInfoBase;
 import com.android.settings.notification.EmptyTextSettings;
 import com.android.settings.widget.AppPreference;
-import com.android.settings.wrapper.UserManagerWrapper;
 import com.android.settingslib.wrapper.PackageManagerWrapper;
 
 import java.text.Collator;
@@ -88,7 +87,7 @@
 
     private Context mContext;
     private PackageManagerWrapper mPackageManager;
-    private UserManagerWrapper mUserManager;
+    private UserManager mUserManager;
     private IconDrawableFactory mIconDrawableFactory;
 
     /**
@@ -119,7 +118,7 @@
         // Do nothing
     }
 
-    public PictureInPictureSettings(PackageManagerWrapper pm, UserManagerWrapper um) {
+    public PictureInPictureSettings(PackageManagerWrapper pm, UserManager um) {
         mPackageManager = pm;
         mUserManager = um;
     }
@@ -130,7 +129,7 @@
 
         mContext = getActivity();
         mPackageManager = new PackageManagerWrapper(mContext.getPackageManager());
-        mUserManager = new UserManagerWrapper(mContext.getSystemService(UserManager.class));
+        mUserManager = (UserManager) mContext.getSystemService(Context.USER_SERVICE);
         mIconDrawableFactory = IconDrawableFactory.newInstance(mContext);
     }
 
diff --git a/src/com/android/settings/connecteddevice/usb/UsbBackend.java b/src/com/android/settings/connecteddevice/usb/UsbBackend.java
index 333449f..54811c8 100644
--- a/src/com/android/settings/connecteddevice/usb/UsbBackend.java
+++ b/src/com/android/settings/connecteddevice/usb/UsbBackend.java
@@ -22,11 +22,10 @@
 import android.hardware.usb.UsbPort;
 import android.hardware.usb.UsbPortStatus;
 import android.net.ConnectivityManager;
+import android.os.UserHandle;
 import android.os.UserManager;
 import android.support.annotation.VisibleForTesting;
 
-import com.android.settings.wrapper.UserManagerWrapper;
-
 /**
  * Provides access to underlying system USB functionality.
  */
@@ -50,17 +49,17 @@
     private UsbPortStatus mPortStatus;
 
     public UsbBackend(Context context) {
-        this(context, new UserManagerWrapper(UserManager.get(context)));
+        this(context, (UserManager) context.getSystemService(Context.USER_SERVICE));
     }
 
     @VisibleForTesting
-    public UsbBackend(Context context, UserManagerWrapper userManagerWrapper) {
+    public UsbBackend(Context context, UserManager userManager) {
         mUsbManager = context.getSystemService(UsbManager.class);
 
-        mFileTransferRestricted = userManagerWrapper.isUsbFileTransferRestricted();
-        mFileTransferRestrictedBySystem = userManagerWrapper.isUsbFileTransferRestrictedBySystem();
-        mTetheringRestricted = userManagerWrapper.isUsbTetheringRestricted();
-        mTetheringRestrictedBySystem = userManagerWrapper.isUsbTetheringRestrictedBySystem();
+        mFileTransferRestricted = isUsbFileTransferRestricted(userManager);
+        mFileTransferRestrictedBySystem = isUsbFileTransferRestrictedBySystem(userManager);
+        mTetheringRestricted = isUsbTetheringRestricted(userManager);
+        mTetheringRestrictedBySystem = isUsbTetheringRestrictedBySystem(userManager);
 
         mMidiSupported = context.getPackageManager().hasSystemFeature(PackageManager.FEATURE_MIDI);
         ConnectivityManager cm =
@@ -172,6 +171,24 @@
         return Integer.parseInt(role);
     }
 
+    private static boolean isUsbFileTransferRestricted(UserManager userManager) {
+        return userManager.hasUserRestriction(UserManager.DISALLOW_USB_FILE_TRANSFER);
+    }
+
+    private static boolean isUsbTetheringRestricted(UserManager userManager) {
+        return userManager.hasUserRestriction(UserManager.DISALLOW_CONFIG_TETHERING);
+    }
+
+    private static boolean isUsbFileTransferRestrictedBySystem(UserManager userManager) {
+        return userManager.hasBaseUserRestriction(
+                UserManager.DISALLOW_USB_FILE_TRANSFER, UserHandle.of(UserHandle.myUserId()));
+    }
+
+    private static boolean isUsbTetheringRestrictedBySystem(UserManager userManager) {
+        return userManager.hasBaseUserRestriction(
+                UserManager.DISALLOW_CONFIG_TETHERING, UserHandle.of(UserHandle.myUserId()));
+    }
+
     private boolean areFunctionDisallowed(long functions) {
         return (mFileTransferRestricted && ((functions & UsbManager.FUNCTION_MTP) != 0
                 || (functions & UsbManager.FUNCTION_PTP) != 0))
diff --git a/src/com/android/settings/deviceinfo/StorageDashboardFragment.java b/src/com/android/settings/deviceinfo/StorageDashboardFragment.java
index ead941f..8e56b12 100644
--- a/src/com/android/settings/deviceinfo/StorageDashboardFragment.java
+++ b/src/com/android/settings/deviceinfo/StorageDashboardFragment.java
@@ -47,7 +47,6 @@
 import com.android.settings.search.BaseSearchIndexProvider;
 import com.android.settings.search.Indexable;
 import com.android.settings.widget.EntityHeaderController;
-import com.android.settings.wrapper.UserManagerWrapper;
 import com.android.settingslib.applications.StorageStatsSource;
 import com.android.settingslib.core.AbstractPreferenceController;
 import com.android.settingslib.deviceinfo.PrivateStorageInfo;
@@ -185,8 +184,7 @@
                 mVolume, new StorageManagerVolumeProvider(sm));
         controllers.add(mPreferenceController);
 
-        UserManagerWrapper userManager =
-                new UserManagerWrapper(context.getSystemService(UserManager.class));
+        final UserManager userManager = context.getSystemService(UserManager.class);
         mSecondaryUsers = SecondaryUserController.getSecondaryUserControllers(context, userManager);
         controllers.addAll(mSecondaryUsers);
 
@@ -235,8 +233,7 @@
                 public List<AbstractPreferenceController> createPreferenceControllers(
                         Context context) {
                     final StorageManager sm = context.getSystemService(StorageManager.class);
-                    final UserManagerWrapper userManager =
-                            new UserManagerWrapper(context.getSystemService(UserManager.class));
+                    final UserManager userManager = context.getSystemService(UserManager.class);
                     final List<AbstractPreferenceController> controllers = new ArrayList<>();
                     controllers.add(new StorageSummaryDonutPreferenceController(context));
                     controllers.add(new StorageItemPreferenceController(context, null /* host */,
@@ -251,9 +248,8 @@
     @Override
     public Loader<SparseArray<StorageAsyncLoader.AppsStorageResult>> onCreateLoader(int id,
             Bundle args) {
-        Context context = getContext();
-        return new StorageAsyncLoader(context,
-                new UserManagerWrapper(context.getSystemService(UserManager.class)),
+        final Context context = getContext();
+        return new StorageAsyncLoader(context, context.getSystemService(UserManager.class),
                 mVolume.fsUuid,
                 new StorageStatsSource(context),
                 new PackageManagerWrapper(context.getPackageManager()));
diff --git a/src/com/android/settings/deviceinfo/StorageProfileFragment.java b/src/com/android/settings/deviceinfo/StorageProfileFragment.java
index 3f75f5c..a312a81 100644
--- a/src/com/android/settings/deviceinfo/StorageProfileFragment.java
+++ b/src/com/android/settings/deviceinfo/StorageProfileFragment.java
@@ -34,7 +34,6 @@
 import com.android.settings.deviceinfo.storage.StorageAsyncLoader;
 import com.android.settings.deviceinfo.storage.StorageAsyncLoader.AppsStorageResult;
 import com.android.settings.deviceinfo.storage.StorageItemPreferenceController;
-import com.android.settings.wrapper.UserManagerWrapper;
 import com.android.settingslib.applications.StorageStatsSource;
 import com.android.settingslib.core.AbstractPreferenceController;
 import com.android.settingslib.deviceinfo.StorageManagerVolumeProvider;
@@ -114,9 +113,9 @@
 
     @Override
     public Loader<SparseArray<AppsStorageResult>> onCreateLoader(int id, Bundle args) {
-        Context context = getContext();
+        final Context context = getContext();
         return new StorageAsyncLoader(context,
-                new UserManagerWrapper(context.getSystemService(UserManager.class)),
+                context.getSystemService(UserManager.class),
                 mVolume.fsUuid,
                 new StorageStatsSource(context),
                 new PackageManagerWrapper(context.getPackageManager()));
diff --git a/src/com/android/settings/deviceinfo/storage/SecondaryUserController.java b/src/com/android/settings/deviceinfo/storage/SecondaryUserController.java
index 4aeb782..d63436f 100644
--- a/src/com/android/settings/deviceinfo/storage/SecondaryUserController.java
+++ b/src/com/android/settings/deviceinfo/storage/SecondaryUserController.java
@@ -19,6 +19,7 @@
 import android.content.Context;
 import android.content.pm.UserInfo;
 import android.graphics.drawable.Drawable;
+import android.os.UserManager;
 import android.support.annotation.NonNull;
 import android.support.annotation.Nullable;
 import android.support.annotation.VisibleForTesting;
@@ -29,7 +30,6 @@
 import com.android.settings.Utils;
 import com.android.settings.core.PreferenceControllerMixin;
 import com.android.settings.deviceinfo.StorageItemPreference;
-import com.android.settings.wrapper.UserManagerWrapper;
 import com.android.settingslib.core.AbstractPreferenceController;
 
 import java.util.ArrayList;
@@ -61,7 +61,7 @@
      * @param userManager UserManagerWrapper for figuring out which controllers to add.
      */
     public static List<AbstractPreferenceController> getSecondaryUserControllers(
-            Context context, UserManagerWrapper userManager) {
+            Context context, UserManager userManager) {
         List<AbstractPreferenceController> controllers = new ArrayList<>();
         UserInfo primaryUser = userManager.getPrimaryUser();
         boolean addedUser = false;
@@ -74,8 +74,7 @@
 
             if (info == null || Utils.isProfileOf(primaryUser, info)) {
                 controllers.add(
-                        new UserProfileController(
-                                context, info, userManager, USER_PROFILE_INSERTION_LOCATION));
+                        new UserProfileController(context, info, USER_PROFILE_INSERTION_LOCATION));
                 continue;
             }
 
diff --git a/src/com/android/settings/deviceinfo/storage/StorageAsyncLoader.java b/src/com/android/settings/deviceinfo/storage/StorageAsyncLoader.java
index 0b9b697..affcbc9 100644
--- a/src/com/android/settings/deviceinfo/storage/StorageAsyncLoader.java
+++ b/src/com/android/settings/deviceinfo/storage/StorageAsyncLoader.java
@@ -26,11 +26,11 @@
 import android.content.pm.PackageManager.NameNotFoundException;
 import android.content.pm.UserInfo;
 import android.os.UserHandle;
+import android.os.UserManager;
 import android.util.ArraySet;
 import android.util.Log;
 import android.util.SparseArray;
 
-import com.android.settings.wrapper.UserManagerWrapper;
 import com.android.settingslib.applications.StorageStatsSource;
 import com.android.settingslib.utils.AsyncLoader;
 import com.android.settingslib.wrapper.PackageManagerWrapper;
@@ -46,7 +46,7 @@
  */
 public class StorageAsyncLoader
         extends AsyncLoader<SparseArray<StorageAsyncLoader.AppsStorageResult>> {
-    private UserManagerWrapper mUserManager;
+    private UserManager mUserManager;
     private static final String TAG = "StorageAsyncLoader";
 
     private String mUuid;
@@ -54,7 +54,7 @@
     private PackageManagerWrapper mPackageManager;
     private ArraySet<String> mSeenPackages;
 
-    public StorageAsyncLoader(Context context, UserManagerWrapper userManager,
+    public StorageAsyncLoader(Context context, UserManager userManager,
             String uuid, StorageStatsSource source, PackageManagerWrapper pm) {
         super(context);
         mUserManager = userManager;
diff --git a/src/com/android/settings/deviceinfo/storage/UserProfileController.java b/src/com/android/settings/deviceinfo/storage/UserProfileController.java
index 081a0be..510b0c0 100644
--- a/src/com/android/settings/deviceinfo/storage/UserProfileController.java
+++ b/src/com/android/settings/deviceinfo/storage/UserProfileController.java
@@ -32,7 +32,6 @@
 import com.android.settings.core.SubSettingLauncher;
 import com.android.settings.deviceinfo.StorageItemPreference;
 import com.android.settings.deviceinfo.StorageProfileFragment;
-import com.android.settings.wrapper.UserManagerWrapper;
 import com.android.settingslib.core.AbstractPreferenceController;
 
 /**
@@ -44,16 +43,13 @@
         UserIconLoader.UserIconHandler {
     private static final String PREFERENCE_KEY_BASE = "pref_profile_";
     private StorageItemPreference mStoragePreference;
-    private UserManagerWrapper mUserManager;
     private UserInfo mUser;
     private long mTotalSizeBytes;
     private final int mPreferenceOrder;
 
-    public UserProfileController(
-            Context context, UserInfo info, UserManagerWrapper userManager, int preferenceOrder) {
+    public UserProfileController(Context context, UserInfo info, int preferenceOrder) {
         super(context);
         mUser = Preconditions.checkNotNull(info);
-        mUserManager = userManager;
         mPreferenceOrder = preferenceOrder;
     }
 
diff --git a/src/com/android/settings/wallpaper/WallpaperSuggestionActivity.java b/src/com/android/settings/wallpaper/WallpaperSuggestionActivity.java
index 0874cea..22835f3 100644
--- a/src/com/android/settings/wallpaper/WallpaperSuggestionActivity.java
+++ b/src/com/android/settings/wallpaper/WallpaperSuggestionActivity.java
@@ -27,7 +27,6 @@
 import com.android.internal.logging.nano.MetricsProto;
 import com.android.settings.R;
 import com.android.settings.core.SubSettingLauncher;
-import com.android.settings.wrapper.WallpaperManagerWrapper;
 
 public class WallpaperSuggestionActivity extends Activity {
 
@@ -61,9 +60,16 @@
 
     @VisibleForTesting
     public static boolean isSuggestionComplete(Context context) {
-        final WallpaperManagerWrapper manager = new WallpaperManagerWrapper(context);
-        return manager.isWallpaperServiceEnabled() ? manager.getWallpaperId(
-                WallpaperManager.FLAG_SYSTEM) > 0 : false;
+        if (!isWallpaperServiceEnabled(context)) {
+            return true;
+        }
+        final WallpaperManager manager = (WallpaperManager) context.getSystemService(
+                WALLPAPER_SERVICE);
+        return manager.getWallpaperId(WallpaperManager.FLAG_SYSTEM) > 0;
     }
 
+    private static boolean isWallpaperServiceEnabled(Context context) {
+        return context.getResources().getBoolean(
+                com.android.internal.R.bool.config_enableWallpaperService);
+    }
 }
diff --git a/src/com/android/settings/wrapper/UserManagerWrapper.java b/src/com/android/settings/wrapper/UserManagerWrapper.java
deleted file mode 100644
index 8c3a01d..0000000
--- a/src/com/android/settings/wrapper/UserManagerWrapper.java
+++ /dev/null
@@ -1,67 +0,0 @@
-/*
- * Copyright (C) 2017 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.pm.UserInfo;
-import android.os.UserHandle;
-import android.os.UserManager;
-
-import java.util.List;
-
-/**
- * This class replicates a subset of the android.os.UserManager. The class
- * exists so that we can use a thin wrapper around the UserManager in production code and a mock in
- * tests. We cannot directly mock or shadow the UserManager, because some of the methods we rely on
- * are newer than the API version supported by Robolectric or are hidden.
- */
-public class UserManagerWrapper {
-    private UserManager mUserManager;
-
-    public UserManagerWrapper(UserManager userManager) {
-        mUserManager = userManager;
-    }
-
-    public UserInfo getPrimaryUser() {
-        return mUserManager.getPrimaryUser();
-    }
-
-    public List<UserInfo> getUsers() {
-        return mUserManager.getUsers();
-    }
-
-    public List<UserInfo> getProfiles(int userHandle) {
-        return mUserManager.getProfiles(userHandle);
-    }
-
-    public boolean isUsbFileTransferRestricted() {
-        return mUserManager.hasUserRestriction(UserManager.DISALLOW_USB_FILE_TRANSFER);
-    }
-
-    public boolean isUsbTetheringRestricted() {
-        return mUserManager.hasUserRestriction(UserManager.DISALLOW_CONFIG_TETHERING);
-    }
-
-    public boolean isUsbFileTransferRestrictedBySystem() {
-        return mUserManager.hasBaseUserRestriction(
-                UserManager.DISALLOW_USB_FILE_TRANSFER, UserHandle.of(UserHandle.myUserId()));
-    }
-
-    public boolean isUsbTetheringRestrictedBySystem() {
-        return mUserManager.hasBaseUserRestriction(
-                UserManager.DISALLOW_CONFIG_TETHERING, UserHandle.of(UserHandle.myUserId()));
-    }
-}
diff --git a/src/com/android/settings/wrapper/WallpaperManagerWrapper.java b/src/com/android/settings/wrapper/WallpaperManagerWrapper.java
deleted file mode 100644
index eabc5bb..0000000
--- a/src/com/android/settings/wrapper/WallpaperManagerWrapper.java
+++ /dev/null
@@ -1,44 +0,0 @@
-/*
- * Copyright (C) 2017 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.app.WallpaperManager;
-import android.content.Context;
-
-public class WallpaperManagerWrapper {
-
-    private final WallpaperManager mWallpaperManager;
-    private final boolean mWallpaperServiceEnabled;
-
-    public WallpaperManagerWrapper(Context context) {
-        mWallpaperServiceEnabled = context.getResources().getBoolean(
-                com.android.internal.R.bool.config_enableWallpaperService);
-        mWallpaperManager = mWallpaperServiceEnabled ? (WallpaperManager) context.getSystemService(
-                Context.WALLPAPER_SERVICE) : null;
-    }
-
-    public boolean isWallpaperServiceEnabled() {
-        return mWallpaperServiceEnabled;
-    }
-
-    public int getWallpaperId(int which) {
-        if (!mWallpaperServiceEnabled) {
-            throw new RuntimeException("This device does not have wallpaper service enabled.");
-        }
-        return mWallpaperManager.getWallpaperId(which);
-    }
-}
diff --git a/tests/robotests/src/com/android/settings/applications/appinfo/PictureInPictureSettingsTest.java b/tests/robotests/src/com/android/settings/applications/appinfo/PictureInPictureSettingsTest.java
index 3ea2a16..87dedad 100644
--- a/tests/robotests/src/com/android/settings/applications/appinfo/PictureInPictureSettingsTest.java
+++ b/tests/robotests/src/com/android/settings/applications/appinfo/PictureInPictureSettingsTest.java
@@ -25,12 +25,13 @@
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageInfo;
 import android.content.pm.UserInfo;
+import android.os.UserManager;
 import android.util.Pair;
 
 import com.android.settings.testutils.FakeFeatureFactory;
 import com.android.settings.testutils.SettingsRobolectricTestRunner;
-import com.android.settings.wrapper.UserManagerWrapper;
 import com.android.settingslib.wrapper.PackageManagerWrapper;
+
 import com.google.common.collect.ImmutableList;
 
 import org.junit.Before;
@@ -52,7 +53,7 @@
     @Mock
     private PackageManagerWrapper mPackageManager;
     @Mock
-    private UserManagerWrapper mUserManager;
+    private UserManager mUserManager;
     private ArrayList<PackageInfo> mPrimaryUserPackages;
     private ArrayList<PackageInfo> mProfileUserPackages;
 
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/usb/UsbBackendTest.java b/tests/robotests/src/com/android/settings/connecteddevice/usb/UsbBackendTest.java
index 0823fe0..567b072 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/usb/UsbBackendTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/usb/UsbBackendTest.java
@@ -18,6 +18,8 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Answers.RETURNS_DEEP_STUBS;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -27,9 +29,10 @@
 import android.hardware.usb.UsbPort;
 import android.hardware.usb.UsbPortStatus;
 import android.net.ConnectivityManager;
+import android.os.UserHandle;
+import android.os.UserManager;
 
 import com.android.settings.testutils.SettingsRobolectricTestRunner;
-import com.android.settings.wrapper.UserManagerWrapper;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -45,7 +48,7 @@
     @Mock
     private UsbManager mUsbManager;
     @Mock
-    private UserManagerWrapper mUserManagerWrapper;
+    private UserManager mUserManager;
     @Mock
     private ConnectivityManager mConnectivityManager;
     @Mock
@@ -68,7 +71,7 @@
 
     @Test
     public void setDataRole_allRolesSupported_shouldSetDataRole() {
-        final UsbBackend usbBackend = new UsbBackend(mContext, mUserManagerWrapper);
+        final UsbBackend usbBackend = new UsbBackend(mContext, mUserManager);
 
         when(mUsbPortStatus
                 .isRoleCombinationSupported(UsbPort.POWER_ROLE_SINK, UsbPort.DATA_ROLE_DEVICE))
@@ -91,7 +94,7 @@
 
     @Test
     public void setDataRole_notAllRolesSupported_shouldSetDataAndPowerRole() {
-        final UsbBackend usbBackend = new UsbBackend(mContext, mUserManagerWrapper);
+        final UsbBackend usbBackend = new UsbBackend(mContext, mUserManager);
 
         when(mUsbPortStatus
                 .isRoleCombinationSupported(UsbPort.POWER_ROLE_SINK, UsbPort.DATA_ROLE_DEVICE))
@@ -109,7 +112,7 @@
 
     @Test
     public void setPowerRole_allRolesSupported_shouldSetPowerRole() {
-        final UsbBackend usbBackend = new UsbBackend(mContext, mUserManagerWrapper);
+        final UsbBackend usbBackend = new UsbBackend(mContext, mUserManager);
 
         when(mUsbPortStatus
                 .isRoleCombinationSupported(UsbPort.POWER_ROLE_SINK, UsbPort.DATA_ROLE_DEVICE))
@@ -133,7 +136,7 @@
 
     @Test
     public void setPowerRole_notAllRolesSupported_shouldSetDataAndPowerRole() {
-        final UsbBackend usbBackend = new UsbBackend(mContext, mUserManagerWrapper);
+        final UsbBackend usbBackend = new UsbBackend(mContext, mUserManager);
 
         when(mUsbPortStatus
                 .isRoleCombinationSupported(UsbPort.POWER_ROLE_SINK, UsbPort.DATA_ROLE_DEVICE))
@@ -151,20 +154,26 @@
 
     @Test
     public void areFunctionsSupported_fileTransferDisallowed_shouldReturnFalse() {
-        when(mUserManagerWrapper.isUsbFileTransferRestricted()).thenReturn(true);
-        when(mUserManagerWrapper.isUsbFileTransferRestrictedBySystem()).thenReturn(true);
+        when(mUserManager.hasUserRestriction(UserManager.DISALLOW_USB_FILE_TRANSFER))
+                .thenReturn(true);
+        when(mUserManager.hasBaseUserRestriction(
+                eq(UserManager.DISALLOW_USB_FILE_TRANSFER), any(UserHandle.class)))
+                .thenReturn(true);
 
-        final UsbBackend usbBackend = new UsbBackend(mContext, mUserManagerWrapper);
+        final UsbBackend usbBackend = new UsbBackend(mContext, mUserManager);
 
         assertThat(usbBackend.areFunctionsSupported(UsbManager.FUNCTION_MTP)).isFalse();
     }
 
     @Test
     public void areFunctionsSupported_fileTransferAllowed_shouldReturnTrue() {
-        when(mUserManagerWrapper.isUsbFileTransferRestricted()).thenReturn(false);
-        when(mUserManagerWrapper.isUsbFileTransferRestrictedBySystem()).thenReturn(false);
+        when(mUserManager.hasUserRestriction(UserManager.DISALLOW_USB_FILE_TRANSFER))
+                .thenReturn(false);
+        when(mUserManager.hasBaseUserRestriction(
+                eq(UserManager.DISALLOW_USB_FILE_TRANSFER), any(UserHandle.class)))
+                .thenReturn(false);
 
-        final UsbBackend usbBackend = new UsbBackend(mContext, mUserManagerWrapper);
+        final UsbBackend usbBackend = new UsbBackend(mContext, mUserManager);
 
         assertThat(usbBackend.areFunctionsSupported(UsbManager.FUNCTION_MTP)).isTrue();
     }
diff --git a/tests/robotests/src/com/android/settings/deviceinfo/storage/SecondaryUserControllerTest.java b/tests/robotests/src/com/android/settings/deviceinfo/storage/SecondaryUserControllerTest.java
index 49f0312..bbe92ad 100644
--- a/tests/robotests/src/com/android/settings/deviceinfo/storage/SecondaryUserControllerTest.java
+++ b/tests/robotests/src/com/android/settings/deviceinfo/storage/SecondaryUserControllerTest.java
@@ -26,13 +26,13 @@
 import android.content.Context;
 import android.content.pm.UserInfo;
 import android.graphics.drawable.Drawable;
+import android.os.UserManager;
 import android.support.v7.preference.Preference;
 import android.support.v7.preference.PreferenceGroup;
 import android.support.v7.preference.PreferenceScreen;
 import android.util.SparseArray;
 
 import com.android.settings.testutils.SettingsRobolectricTestRunner;
-import com.android.settings.wrapper.UserManagerWrapper;
 import com.android.settingslib.applications.StorageStatsSource;
 import com.android.settingslib.core.AbstractPreferenceController;
 import com.android.settingslib.drawable.UserIconDrawable;
@@ -54,7 +54,7 @@
     private static final String TEST_NAME = "Fred";
     private static final String TARGET_PREFERENCE_GROUP_KEY = "pref_secondary_users";
     @Mock
-    private UserManagerWrapper mUserManager;
+    private UserManager mUserManager;
     @Mock
     private PreferenceScreen mScreen;
     @Mock
diff --git a/tests/robotests/src/com/android/settings/deviceinfo/storage/UserProfileControllerTest.java b/tests/robotests/src/com/android/settings/deviceinfo/storage/UserProfileControllerTest.java
index 84456ec..7ee6216 100644
--- a/tests/robotests/src/com/android/settings/deviceinfo/storage/UserProfileControllerTest.java
+++ b/tests/robotests/src/com/android/settings/deviceinfo/storage/UserProfileControllerTest.java
@@ -35,7 +35,6 @@
 import com.android.settings.SubSettings;
 import com.android.settings.deviceinfo.StorageProfileFragment;
 import com.android.settings.testutils.SettingsRobolectricTestRunner;
-import com.android.settings.wrapper.UserManagerWrapper;
 import com.android.settingslib.applications.StorageStatsSource;
 import com.android.settingslib.drawable.UserIconDrawable;
 
@@ -53,8 +52,6 @@
     private static final String TEST_NAME = "Fred";
 
     @Mock
-    private UserManagerWrapper mUserManager;
-    @Mock
     private PreferenceScreen mScreen;
 
     private Context mContext;
@@ -66,7 +63,7 @@
         MockitoAnnotations.initMocks(this);
         mContext = spy(RuntimeEnvironment.application);
         mPrimaryProfile = new UserInfo();
-        mController = new UserProfileController(mContext, mPrimaryProfile, mUserManager, 0);
+        mController = new UserProfileController(mContext, mPrimaryProfile, 0);
         when(mScreen.getContext()).thenReturn(mContext);
         mPrimaryProfile.name = TEST_NAME;
         mPrimaryProfile.id = 10;
@@ -96,7 +93,7 @@
         final Intent intent = intentCaptor.getValue();
         assertThat(intent.getComponent().getClassName()).isEqualTo(SubSettings.class.getName());
         assertThat(intent.getStringExtra(SettingsActivity.EXTRA_SHOW_FRAGMENT))
-            .isEqualTo(StorageProfileFragment.class.getName());
+                .isEqualTo(StorageProfileFragment.class.getName());
     }
 
     @Test
diff --git a/tests/robotests/src/com/android/settings/wallpaper/WallpaperSuggestionActivityTest.java b/tests/robotests/src/com/android/settings/wallpaper/WallpaperSuggestionActivityTest.java
index 6ebbef3..fc708ef 100644
--- a/tests/robotests/src/com/android/settings/wallpaper/WallpaperSuggestionActivityTest.java
+++ b/tests/robotests/src/com/android/settings/wallpaper/WallpaperSuggestionActivityTest.java
@@ -20,15 +20,15 @@
 import static org.mockito.Mockito.when;
 import static org.robolectric.Shadows.shadowOf;
 
+import android.app.WallpaperManager;
 import android.content.Context;
 import android.content.Intent;
-import android.content.pm.PackageManager;
 import android.content.res.Resources;
 
 import com.android.settings.SubSettings;
 import com.android.settings.testutils.SettingsRobolectricTestRunner;
-import com.android.settings.wrapper.WallpaperManagerWrapper;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -40,6 +40,7 @@
 import org.robolectric.annotation.Config;
 import org.robolectric.annotation.Implementation;
 import org.robolectric.annotation.Implements;
+import org.robolectric.annotation.Resetter;
 import org.robolectric.shadows.ShadowActivity;
 import org.robolectric.shadows.ShadowPackageManager;
 
@@ -59,10 +60,15 @@
         mController = Robolectric.buildActivity(WallpaperSuggestionActivity.class);
     }
 
+    @After
+    public void tearDown() {
+        ShadowWallpaperManager.reset();
+    }
+
     @Test
     public void launch_primarySuggestionActivityDoesNotExist_shouldFallback() {
         ShadowPackageManager packageManager =
-            shadowOf(RuntimeEnvironment.application.getPackageManager());
+                shadowOf(RuntimeEnvironment.application.getPackageManager());
         packageManager.removePackage("com.android.settings");
 
         ShadowActivity activity = shadowOf(mController.setup().get());
@@ -74,32 +80,34 @@
     }
 
     @Test
-    public void wallpaperServiceEnabled_no_shouldReturnFalse() {
+    public void wallpaperServiceEnabled_no_shouldReturnTrue() {
         when(mContext.getResources()).thenReturn(mResources);
         when(mResources.getBoolean(com.android.internal.R.bool.config_enableWallpaperService))
-            .thenReturn(false);
-
-        assertThat(WallpaperSuggestionActivity.isSuggestionComplete(mContext)).isFalse();
-    }
-
-    @Test
-    @Config(shadows = WallpaperSuggestionActivityTest.ShadowWallpaperManagerWrapper.class)
-    public void hasWallpaperSet_no_shouldReturnFalse() {
-        ShadowWallpaperManagerWrapper.setWallpaperId(0);
-
-        assertThat(WallpaperSuggestionActivity.isSuggestionComplete(mContext)).isFalse();
-    }
-
-    @Test
-    @Config(shadows = WallpaperSuggestionActivityTest.ShadowWallpaperManagerWrapper.class)
-    public void hasWallpaperSet_yes_shouldReturnTrue() {
-        ShadowWallpaperManagerWrapper.setWallpaperId(100);
+                .thenReturn(false);
 
         assertThat(WallpaperSuggestionActivity.isSuggestionComplete(mContext)).isTrue();
     }
 
-    @Implements(WallpaperManagerWrapper.class)
-    public static class ShadowWallpaperManagerWrapper {
+    @Test
+    @Config(shadows = ShadowWallpaperManager.class)
+    public void hasWallpaperSet_no_shouldReturnFalse() {
+        ShadowWallpaperManager.setWallpaperId(0);
+
+        assertThat(WallpaperSuggestionActivity.isSuggestionComplete(RuntimeEnvironment.application))
+                .isFalse();
+    }
+
+    @Test
+    @Config(shadows = ShadowWallpaperManager.class)
+    public void hasWallpaperSet_yes_shouldReturnTrue() {
+        ShadowWallpaperManager.setWallpaperId(100);
+
+        assertThat(WallpaperSuggestionActivity.isSuggestionComplete(RuntimeEnvironment.application))
+                .isTrue();
+    }
+
+    @Implements(WallpaperManager.class)
+    public static class ShadowWallpaperManager {
 
         private static int sWallpaperId;
 
@@ -107,13 +115,11 @@
             sWallpaperId = id;
         }
 
+        @Resetter
         public static void reset() {
             sWallpaperId = 0;
         }
 
-        public void __constructor__(Context context) {
-        }
-
         @Implementation
         public boolean isWallpaperServiceEnabled() {
             return true;
diff --git a/tests/unit/src/com/android/settings/deviceinfo/storage/StorageAsyncLoaderTest.java b/tests/unit/src/com/android/settings/deviceinfo/storage/StorageAsyncLoaderTest.java
index f622804..d1d2b04 100644
--- a/tests/unit/src/com/android/settings/deviceinfo/storage/StorageAsyncLoaderTest.java
+++ b/tests/unit/src/com/android/settings/deviceinfo/storage/StorageAsyncLoaderTest.java
@@ -17,7 +17,6 @@
 package com.android.settings.deviceinfo.storage;
 
 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;
@@ -31,11 +30,11 @@
 import android.content.pm.UserInfo;
 import android.net.TrafficStats;
 import android.os.UserHandle;
+import android.os.UserManager;
 import android.support.test.filters.SmallTest;
 import android.support.test.runner.AndroidJUnit4;
 import android.util.SparseArray;
 
-import com.android.settings.wrapper.UserManagerWrapper;
 import com.android.settingslib.applications.StorageStatsSource;
 import com.android.settingslib.wrapper.PackageManagerWrapper;
 
@@ -65,7 +64,7 @@
     @Mock
     private PackageManagerWrapper mPackageManager;
     @Mock
-    private UserManagerWrapper mUserManager;
+    private UserManager mUserManager;
     private List<ApplicationInfo> mInfo = new ArrayList<>();
     private List<UserInfo> mUsers;