diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DatabaseUtilsTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DatabaseUtilsTest.java
index 722f869bd7f..154497e8653 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DatabaseUtilsTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DatabaseUtilsTest.java
@@ -45,7 +45,6 @@ import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
-import org.robolectric.Shadows;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -56,8 +55,8 @@ public final class DatabaseUtilsTest {
 
     private Context mContext;
 
-    @Mock
-    private PackageManager mPackageManager;
+    @Mock private PackageManager mPackageManager;
+    @Mock private UserManager mUserManager;
     @Mock private ContentResolver mMockContentResolver;
     @Mock private ContentResolver mMockContentResolver2;
     @Mock private BatteryUsageStats mBatteryUsageStats;
@@ -383,7 +382,11 @@ public final class DatabaseUtilsTest {
         doReturn("com.fake.package").when(mContext).getPackageName();
         doReturn(mMockContext).when(mContext).createPackageContextAsUser(
                 "com.fake.package", /*flags=*/ 0, UserHandle.OWNER);
-        BatteryTestUtils.setWorkProfile(mContext);
+        doReturn(UserHandle.CURRENT).when(mContext).getUser();
+        doReturn(mUserManager).when(mContext).getSystemService(UserManager.class);
+        doReturn(true).when(mUserManager).isManagedProfile();
+        doReturn(UserHandle.SYSTEM).when(mUserManager).getProfileParent(UserHandle.CURRENT);
+
         DatabaseUtils.sFakeBatteryStateSupplier = () -> getMatrixCursor();
 
         final Map<Long, Map<String, BatteryHistEntry>> batteryHistMap =