Add cache strategy for getUsageSource().

Cache usage source into SharedPreferences when phone is booting to avoid
calling it too frequently.
It should be safe because the usage source can only change on reboot.

Bug: 293366011
Test: make RunSettingsRoboTests
Change-Id: I35c07539d294737c5764b03b746cfb39f4ce008d
diff --git a/src/com/android/settings/fuelgauge/batteryusage/BootBroadcastReceiver.java b/src/com/android/settings/fuelgauge/batteryusage/BootBroadcastReceiver.java
index 920670f..6d14e1c 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/BootBroadcastReceiver.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/BootBroadcastReceiver.java
@@ -83,6 +83,10 @@
             recheckIntent.setClass(context, BootBroadcastReceiver.class);
             final long delayedTime = getRescheduleTimeForBootAction(context);
             mHandler.postDelayed(() -> context.sendBroadcast(recheckIntent), delayedTime);
+
+            // Refreshes the usage source from UsageStatsManager when booting.
+            DatabaseUtils.removeUsageSource(context);
+
             BatteryUsageLogUtils.writeLog(context, Action.RECHECK_JOB, "delay:" + delayedTime);
         } else if (ACTION_SETUP_WIZARD_FINISHED.equals(action)) {
             ElapsedTimeUtils.storeSuwFinishedTimestamp(context, System.currentTimeMillis());
diff --git a/src/com/android/settings/fuelgauge/batteryusage/ConvertUtils.java b/src/com/android/settings/fuelgauge/batteryusage/ConvertUtils.java
index 2c98c4b..3fc44cc 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/ConvertUtils.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/ConvertUtils.java
@@ -17,7 +17,6 @@
 
 import android.annotation.IntDef;
 import android.annotation.Nullable;
-import android.app.usage.IUsageStatsManager;
 import android.app.usage.UsageEvents.Event;
 import android.app.usage.UsageStatsManager;
 import android.content.ContentValues;
@@ -27,7 +26,6 @@
 import android.os.BatteryUsageStats;
 import android.os.Build;
 import android.os.LocaleList;
-import android.os.RemoteException;
 import android.os.UserHandle;
 import android.text.TextUtils;
 import android.text.format.DateFormat;
@@ -67,6 +65,12 @@
     public static final int CONSUMER_TYPE_USER_BATTERY = 2;
     public static final int CONSUMER_TYPE_SYSTEM_BATTERY = 3;
 
+    public static final int DEFAULT_USAGE_SOURCE = UsageStatsManager.USAGE_SOURCE_CURRENT_ACTIVITY;
+    public static final int EMPTY_USAGE_SOURCE = -1;
+
+    @VisibleForTesting
+    static int sUsageSource = EMPTY_USAGE_SOURCE;
+
     private ConvertUtils() {
     }
 
@@ -181,8 +185,7 @@
     /** Converts to {@link AppUsageEvent} from {@link Event} */
     @Nullable
     public static AppUsageEvent convertToAppUsageEvent(
-            Context context, final IUsageStatsManager usageStatsManager, final Event event,
-            final long userId) {
+            Context context, final Event event, final long userId) {
         final String packageName = event.getPackageName();
         if (packageName == null) {
             // See b/190609174: Event package names should never be null, but sometimes they are.
@@ -207,7 +210,7 @@
         }
 
         final String effectivePackageName =
-                getEffectivePackageName(usageStatsManager, packageName, taskRootPackageName);
+                getEffectivePackageName(context, packageName, taskRootPackageName);
         try {
             final long uid = context
                     .getPackageManager()
@@ -323,9 +326,8 @@
      */
     @VisibleForTesting
     static String getEffectivePackageName(
-            final IUsageStatsManager usageStatsManager, final String packageName,
-            final String taskRootPackageName) {
-        int usageSource = getUsageSource(usageStatsManager);
+            Context context, final String packageName, final String taskRootPackageName) {
+        final int usageSource = getUsageSource(context);
         switch (usageSource) {
             case UsageStatsManager.USAGE_SOURCE_TASK_ROOT_ACTIVITY:
                 return !TextUtils.isEmpty(taskRootPackageName)
@@ -370,18 +372,11 @@
         }
     }
 
-    /**
-     * Returns what App Usage Observers will consider the source of usage for an activity.
-     *
-     * @see UsageStatsManager#getUsageSource()
-     */
-    private static int getUsageSource(final IUsageStatsManager usageStatsManager) {
-        try {
-            return usageStatsManager.getUsageSource();
-        } catch (RemoteException e) {
-            Log.e(TAG, "Failed to getUsageSource", e);
-            return UsageStatsManager.USAGE_SOURCE_CURRENT_ACTIVITY;
+    private static int getUsageSource(Context context) {
+        if (sUsageSource == EMPTY_USAGE_SOURCE) {
+            sUsageSource = DatabaseUtils.getUsageSource(context);
         }
+        return sUsageSource;
     }
 
     private static AppUsageEventType getAppUsageEventType(final int eventType) {
diff --git a/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java b/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java
index 0f67e6a..1c851fd 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java
@@ -397,8 +397,8 @@
         }
         // Generates the indexed AppUsagePeriod list data for each corresponding time slot for
         // further use.
-        mAppUsagePeriodMap = DataProcessor.generateAppUsagePeriodMap(mRawStartTimestamp,
-                mHourlyBatteryLevelsPerDay, mAppUsageEventList, mBatteryEventList);
+        mAppUsagePeriodMap = DataProcessor.generateAppUsagePeriodMap(
+                mContext, mHourlyBatteryLevelsPerDay, mAppUsageEventList, mBatteryEventList);
     }
 
     private void tryToGenerateFinalDataAndApplyCallback() {
diff --git a/src/com/android/settings/fuelgauge/batteryusage/DataProcessor.java b/src/com/android/settings/fuelgauge/batteryusage/DataProcessor.java
index 6914c30..badc359 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/DataProcessor.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/DataProcessor.java
@@ -32,7 +32,6 @@
 import android.os.BatteryUsageStatsQuery;
 import android.os.Process;
 import android.os.RemoteException;
-import android.os.ServiceManager;
 import android.os.UidBatteryConsumer;
 import android.os.UserBatteryConsumer;
 import android.os.UserHandle;
@@ -78,8 +77,6 @@
     private static final int MIN_DAILY_DATA_SIZE = 2;
     private static final int MIN_TIMESTAMP_DATA_SIZE = 2;
     private static final int MAX_DIFF_SECONDS_OF_UPPER_TIMESTAMP = 5;
-    // Maximum total time value for each hourly slot cumulative data at most 2 hours.
-    private static final float TOTAL_HOURLY_TIME_THRESHOLD = DateUtils.HOUR_IN_MILLIS * 2;
     private static final long MIN_TIME_SLOT = DateUtils.HOUR_IN_MILLIS * 2;
     private static final String MEDIASERVER_PACKAGE_NAME = "mediaserver";
     private static final String ANDROID_CORE_APPS_SHARED_USER_ID = "android.uid.shared";
@@ -111,11 +108,6 @@
     @VisibleForTesting
     static Set<String> sTestSystemAppsPackageNames;
 
-    @VisibleForTesting
-    static IUsageStatsManager sUsageStatsManager =
-            IUsageStatsManager.Stub.asInterface(
-                    ServiceManager.getService(Context.USAGE_STATS_SERVICE));
-
     public static final String CURRENT_TIME_BATTERY_HISTORY_PLACEHOLDER =
             "CURRENT_TIME_BATTERY_HISTORY_PLACEHOLDER";
 
@@ -271,7 +263,7 @@
     @Nullable
     public static Map<Integer, Map<Integer, Map<Long, Map<String, List<AppUsagePeriod>>>>>
             generateAppUsagePeriodMap(
-                    final long rawStartTimestamp,
+                    Context context,
                     final List<BatteryLevelData.PeriodBatteryLevelData> hourlyBatteryLevelsPerDay,
                     final List<AppUsageEvent> appUsageEventList,
                     final List<BatteryEvent> batteryEventList) {
@@ -305,7 +297,7 @@
                 // The value could be null when there is no data in the hourly slot.
                 dailyMap.put(
                         hourlyIndex,
-                        buildAppUsagePeriodList(hourlyAppUsageEventList, batteryEventList,
+                        buildAppUsagePeriodList(context, hourlyAppUsageEventList, batteryEventList,
                                 startTimestamp, endTimestamp));
             }
         }
@@ -346,8 +338,7 @@
                             break;
                         }
                         final AppUsageEvent appUsageEvent =
-                                ConvertUtils.convertToAppUsageEvent(
-                                        context, sUsageStatsManager, event, userId);
+                                ConvertUtils.convertToAppUsageEvent(context, event, userId);
                         if (appUsageEvent != null) {
                             numEventsFetched++;
                             appUsageEventList.add(appUsageEvent);
@@ -661,8 +652,8 @@
     @VisibleForTesting
     @Nullable
     static Map<Long, Map<String, List<AppUsagePeriod>>> buildAppUsagePeriodList(
-            final List<AppUsageEvent> appUsageEvents, final List<BatteryEvent> batteryEventList,
-            final long startTime, final long endTime) {
+            Context context, final List<AppUsageEvent> appUsageEvents,
+            final List<BatteryEvent> batteryEventList, final long startTime, final long endTime) {
         if (appUsageEvents.isEmpty()) {
             return null;
         }
@@ -702,7 +693,7 @@
             final AppUsageEvent firstEvent = usageEvents.get(0);
             final long eventUserId = firstEvent.getUserId();
             final String packageName = getEffectivePackageName(
-                    sUsageStatsManager,
+                    context,
                     firstEvent.getPackageName(),
                     firstEvent.getTaskRootPackageName());
             usageEvents.addAll(deviceEvents);
@@ -975,7 +966,7 @@
         final long startTime = DatabaseUtils.getAppUsageStartTimestampOfUser(
                 context, userID, earliestTimestamp);
         return loadAppUsageEventsForUserFromService(
-                sUsageStatsManager, startTime, now, userID, callingPackage);
+                DatabaseUtils.sUsageStatsManager, startTime, now, userID, callingPackage);
     }
 
     @Nullable
diff --git a/src/com/android/settings/fuelgauge/batteryusage/DatabaseUtils.java b/src/com/android/settings/fuelgauge/batteryusage/DatabaseUtils.java
index 8d1a2f9..ea1f3ed 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/DatabaseUtils.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/DatabaseUtils.java
@@ -15,6 +15,8 @@
  */
 
 package com.android.settings.fuelgauge.batteryusage;
+import android.app.usage.IUsageStatsManager;
+import android.app.usage.UsageStatsManager;
 import android.content.ContentResolver;
 import android.content.ContentValues;
 import android.content.Context;
@@ -28,6 +30,8 @@
 import android.os.BatteryUsageStats;
 import android.os.Handler;
 import android.os.Looper;
+import android.os.RemoteException;
+import android.os.ServiceManager;
 import android.os.SystemClock;
 import android.os.UserManager;
 import android.util.Log;
@@ -63,6 +67,7 @@
     static final int DATA_RETENTION_INTERVAL_DAY = 9;
     static final String KEY_LAST_LOAD_FULL_CHARGE_TIME = "last_load_full_charge_time";
     static final String KEY_LAST_UPLOAD_FULL_CHARGE_TIME = "last_upload_full_charge_time";
+    static final String KEY_LAST_USAGE_SOURCE = "last_usage_source";
 
     /** An authority name of the battery content provider. */
     public static final String AUTHORITY = "com.android.settings.battery.usage.provider";
@@ -74,8 +79,6 @@
     public static final String BATTERY_STATE_TABLE = "BatteryState";
     /** A path name for app usage latest timestamp query. */
     public static final String APP_USAGE_LATEST_TIMESTAMP_PATH = "appUsageLatestTimestamp";
-    /** A class name for battery usage data provider. */
-    public static final String SETTINGS_PACKAGE_PATH = "com.android.settings";
     /** Key for query parameter timestamp used in BATTERY_CONTENT_URI **/
     public static final String QUERY_KEY_TIMESTAMP = "timestamp";
     /** Key for query parameter userid used in APP_USAGE_EVENT_URI **/
@@ -114,6 +117,11 @@
     @VisibleForTesting
     static Supplier<Cursor> sFakeSupplier;
 
+    @VisibleForTesting
+    static IUsageStatsManager sUsageStatsManager =
+            IUsageStatsManager.Stub.asInterface(
+                    ServiceManager.getService(Context.USAGE_STATS_SERVICE));
+
     private DatabaseUtils() {
     }
 
@@ -468,6 +476,37 @@
                 SHARED_PREFS_FILE, Context.MODE_PRIVATE);
     }
 
+    static void removeUsageSource(Context context) {
+        final SharedPreferences sharedPreferences = getSharedPreferences(context);
+        if (sharedPreferences != null && sharedPreferences.contains(KEY_LAST_USAGE_SOURCE)) {
+            sharedPreferences.edit().remove(KEY_LAST_USAGE_SOURCE).apply();
+        }
+    }
+
+    /**
+     * Returns what App Usage Observers will consider the source of usage for an activity.
+     *
+     * @see UsageStatsManager#getUsageSource()
+     */
+    static int getUsageSource(Context context) {
+        final SharedPreferences sharedPreferences = getSharedPreferences(context);
+        if (sharedPreferences != null && sharedPreferences.contains(KEY_LAST_USAGE_SOURCE)) {
+            return sharedPreferences
+                    .getInt(KEY_LAST_USAGE_SOURCE, ConvertUtils.DEFAULT_USAGE_SOURCE);
+        }
+        int usageSource = ConvertUtils.DEFAULT_USAGE_SOURCE;
+
+        try {
+            usageSource = sUsageStatsManager.getUsageSource();
+        } catch (RemoteException e) {
+            Log.e(TAG, "Failed to getUsageSource", e);
+        }
+        if (sharedPreferences != null) {
+            sharedPreferences.edit().putInt(KEY_LAST_USAGE_SOURCE, usageSource).apply();
+        }
+        return usageSource;
+    }
+
     static void recordDateTime(Context context, String preferenceKey) {
         final SharedPreferences sharedPreferences = getSharedPreferences(context);
         if (sharedPreferences != null) {
@@ -564,7 +603,7 @@
 
     private static Map<Long, Map<String, BatteryHistEntry>> loadHistoryMapFromContentProvider(
             Context context, Uri batteryStateUri) {
-        context = DatabaseUtils.getParentContext(context);
+        context = getParentContext(context);
         if (context == null) {
             return null;
         }
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/BootBroadcastReceiverTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/BootBroadcastReceiverTest.java
index aa1ebd7..566df52 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/BootBroadcastReceiverTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/BootBroadcastReceiverTest.java
@@ -22,8 +22,10 @@
 
 import android.app.AlarmManager;
 import android.app.Application;
+import android.app.usage.UsageStatsManager;
 import android.content.Context;
 import android.content.Intent;
+import android.content.SharedPreferences;
 
 import androidx.test.core.app.ApplicationProvider;
 
@@ -40,7 +42,6 @@
 import org.robolectric.shadows.ShadowAlarmManager;
 
 import java.time.Clock;
-import java.time.Duration;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 
@@ -65,10 +66,12 @@
         BatteryTestUtils.insertDataToBatteryStateTable(
                 mContext, Clock.systemUTC().millis(), "com.android.systemui");
         mDao = database.batteryStateDao();
+        clearSharedPreferences();
     }
 
     @After
     public void tearDown() {
+        clearSharedPreferences();
         mPeriodicJobManager.reset();
     }
 
@@ -82,8 +85,21 @@
 
     @Test
     public void onReceive_withBootCompletedIntent_refreshesJob() {
+        final SharedPreferences sharedPreferences = DatabaseUtils.getSharedPreferences(mContext);
+        sharedPreferences
+                .edit()
+                .putInt(DatabaseUtils.KEY_LAST_USAGE_SOURCE,
+                        UsageStatsManager.USAGE_SOURCE_CURRENT_ACTIVITY)
+                .apply();
+
         mReceiver.onReceive(mContext, new Intent(Intent.ACTION_BOOT_COMPLETED));
+
         assertThat(mShadowAlarmManager.peekNextScheduledAlarm()).isNotNull();
+        assertThat(
+                DatabaseUtils
+                        .getSharedPreferences(mContext)
+                        .contains(DatabaseUtils.KEY_LAST_USAGE_SOURCE))
+                .isFalse();
     }
 
     @Test
@@ -133,15 +149,7 @@
                 BootBroadcastReceiver.ACTION_PERIODIC_JOB_RECHECK);
     }
 
-    private void insertExpiredData(int shiftDay) {
-        final long expiredTimeInMs =
-                Clock.systemUTC().millis() - Duration.ofDays(shiftDay).toMillis();
-        BatteryTestUtils.insertDataToBatteryStateTable(
-                mContext, expiredTimeInMs - 1, "com.android.systemui");
-        BatteryTestUtils.insertDataToBatteryStateTable(
-                mContext, expiredTimeInMs, "com.android.systemui");
-        // Ensures the testing environment is correct.
-        assertThat(mDao.getAllAfter(0)).hasSize(3);
+    private void clearSharedPreferences() {
+        DatabaseUtils.getSharedPreferences(mContext).edit().clear().apply();
     }
-
 }
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/ConvertUtilsTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/ConvertUtilsTest.java
index 6b8073b..3cbe8a4 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/ConvertUtilsTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/ConvertUtilsTest.java
@@ -25,7 +25,6 @@
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
 
-import android.app.usage.IUsageStatsManager;
 import android.app.usage.UsageEvents;
 import android.app.usage.UsageEvents.Event;
 import android.content.ContentValues;
@@ -35,7 +34,6 @@
 import android.os.BatteryManager;
 import android.os.BatteryUsageStats;
 import android.os.LocaleList;
-import android.os.RemoteException;
 import android.os.UserHandle;
 
 import com.android.settings.fuelgauge.batteryusage.db.AppUsageEventEntity;
@@ -62,14 +60,13 @@
     @Mock
     private BatteryUsageStats mBatteryUsageStats;
     @Mock
-    private IUsageStatsManager mUsageStatsManager;
-    @Mock
     private BatteryEntry mMockBatteryEntry;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         mContext = spy(RuntimeEnvironment.application);
+        ConvertUtils.sUsageSource = ConvertUtils.EMPTY_USAGE_SOURCE;
         when(mContext.getPackageManager()).thenReturn(mMockPackageManager);
     }
 
@@ -302,7 +299,7 @@
 
         final long userId = 2;
         final AppUsageEvent appUsageEvent = ConvertUtils.convertToAppUsageEvent(
-                mContext, mUsageStatsManager, event, userId);
+                mContext, event, userId);
         assertThat(appUsageEvent.getTimestamp()).isEqualTo(101L);
         assertThat(appUsageEvent.getType()).isEqualTo(AppUsageEventType.ACTIVITY_RESUMED);
         assertThat(appUsageEvent.getPackageName()).isEqualTo("com.android.settings1");
@@ -322,8 +319,8 @@
         when(mMockPackageManager.getPackageUidAsUser(any(), anyInt())).thenReturn(1001);
 
         final long userId = 1;
-        final AppUsageEvent appUsageEvent = ConvertUtils.convertToAppUsageEvent(
-                mContext, mUsageStatsManager, event, userId);
+        final AppUsageEvent appUsageEvent =
+                ConvertUtils.convertToAppUsageEvent(mContext, event, userId);
         assertThat(appUsageEvent.getTimestamp()).isEqualTo(101L);
         assertThat(appUsageEvent.getType()).isEqualTo(AppUsageEventType.DEVICE_SHUTDOWN);
         assertThat(appUsageEvent.getPackageName()).isEqualTo("com.android.settings1");
@@ -338,8 +335,8 @@
         final Event event = new Event();
         event.mPackage = null;
 
-        final AppUsageEvent appUsageEvent = ConvertUtils.convertToAppUsageEvent(
-                mContext, mUsageStatsManager, event, /*userId=*/ 0);
+        final AppUsageEvent appUsageEvent =
+                ConvertUtils.convertToAppUsageEvent(mContext, event, /*userId=*/ 0);
 
         assertThat(appUsageEvent).isNull();
     }
@@ -354,8 +351,8 @@
                 .thenThrow(new PackageManager.NameNotFoundException());
 
         final long userId = 1;
-        final AppUsageEvent appUsageEvent = ConvertUtils.convertToAppUsageEvent(
-                mContext, mUsageStatsManager, event, userId);
+        final AppUsageEvent appUsageEvent =
+                ConvertUtils.convertToAppUsageEvent(mContext, event, userId);
 
         assertThat(appUsageEvent).isNull();
     }
@@ -450,51 +447,47 @@
     }
 
     @Test
-    public void getEffectivePackageName_currentActivity_returnPackageName() throws RemoteException {
-        when(mUsageStatsManager.getUsageSource()).thenReturn(USAGE_SOURCE_CURRENT_ACTIVITY);
+    public void getEffectivePackageName_currentActivity_returnPackageName() {
+        ConvertUtils.sUsageSource = USAGE_SOURCE_CURRENT_ACTIVITY;
         final String packageName = "com.android.settings1";
         final String taskRootPackageName = "com.android.settings2";
 
         assertThat(ConvertUtils.getEffectivePackageName(
-                mUsageStatsManager, packageName, taskRootPackageName))
+                mContext, packageName, taskRootPackageName))
                 .isEqualTo(packageName);
     }
 
     @Test
-    public void getEffectivePackageName_usageSourceThrowException_returnPackageName()
-            throws RemoteException {
-        when(mUsageStatsManager.getUsageSource()).thenThrow(new RemoteException());
+    public void getEffectivePackageName_emptyUsageSource_returnPackageName() {
         final String packageName = "com.android.settings1";
         final String taskRootPackageName = "com.android.settings2";
 
         assertThat(ConvertUtils.getEffectivePackageName(
-                mUsageStatsManager, packageName, taskRootPackageName))
+                mContext, packageName, taskRootPackageName))
                 .isEqualTo(packageName);
     }
 
     @Test
-    public void getEffectivePackageName_rootActivity_returnTaskRootPackageName()
-            throws RemoteException {
-        when(mUsageStatsManager.getUsageSource()).thenReturn(USAGE_SOURCE_TASK_ROOT_ACTIVITY);
+    public void getEffectivePackageName_rootActivity_returnTaskRootPackageName() {
+        ConvertUtils.sUsageSource = USAGE_SOURCE_TASK_ROOT_ACTIVITY;
         final String packageName = "com.android.settings1";
         final String taskRootPackageName = "com.android.settings2";
 
         assertThat(ConvertUtils.getEffectivePackageName(
-                mUsageStatsManager, packageName, taskRootPackageName))
+                mContext, packageName, taskRootPackageName))
                 .isEqualTo(taskRootPackageName);
     }
 
     @Test
-    public void getEffectivePackageName_nullOrEmptyTaskRoot_returnPackageName()
-            throws RemoteException {
-        when(mUsageStatsManager.getUsageSource()).thenReturn(USAGE_SOURCE_TASK_ROOT_ACTIVITY);
+    public void getEffectivePackageName_nullOrEmptyTaskRoot_returnPackageName() {
+        ConvertUtils.sUsageSource = USAGE_SOURCE_TASK_ROOT_ACTIVITY;
         final String packageName = "com.android.settings1";
 
         assertThat(ConvertUtils.getEffectivePackageName(
-                mUsageStatsManager, packageName, /*taskRootPackageName=*/ null))
+                mContext, packageName, /*taskRootPackageName=*/ null))
                 .isEqualTo(packageName);
         assertThat(ConvertUtils.getEffectivePackageName(
-                mUsageStatsManager, packageName, /*taskRootPackageName=*/ ""))
+                mContext, packageName, /*taskRootPackageName=*/ ""))
                 .isEqualTo(packageName);
     }
 }
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessManagerTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessManagerTest.java
index b610cfb..7f7fe43 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessManagerTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessManagerTest.java
@@ -72,7 +72,7 @@
         MockitoAnnotations.initMocks(this);
 
         mContext = spy(RuntimeEnvironment.application);
-        DataProcessor.sUsageStatsManager = mUsageStatsManager;
+        DatabaseUtils.sUsageStatsManager = mUsageStatsManager;
         doReturn(mContext).when(mContext).getApplicationContext();
         doReturn(mUserManager)
                 .when(mContext)
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessorTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessorTest.java
index e2274e2..8bed054 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessorTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessorTest.java
@@ -93,7 +93,7 @@
         mPowerUsageFeatureProvider = mFeatureFactory.powerUsageFeatureProvider;
 
         DataProcessor.sTestSystemAppsPackageNames = Set.of();
-        DataProcessor.sUsageStatsManager = mUsageStatsManager;
+        DatabaseUtils.sUsageStatsManager = mUsageStatsManager;
         doReturn(mIntent).when(mContext).registerReceiver(
                 isA(BroadcastReceiver.class), isA(IntentFilter.class));
         doReturn(100).when(mIntent).getIntExtra(eq(BatteryManager.EXTRA_SCALE), anyInt());
@@ -249,7 +249,7 @@
 
         final Map<Integer, Map<Integer, Map<Long, Map<String, List<AppUsagePeriod>>>>> periodMap =
                 DataProcessor.generateAppUsagePeriodMap(
-                        14400000L, hourlyBatteryLevelsPerDay, appUsageEventList, new ArrayList<>());
+                        mContext, hourlyBatteryLevelsPerDay, appUsageEventList, new ArrayList<>());
 
         assertThat(periodMap).hasSize(3);
         // Day 1
@@ -287,7 +287,8 @@
         hourlyBatteryLevelsPerDay.add(
                 new BatteryLevelData.PeriodBatteryLevelData(new ArrayList<>(), new ArrayList<>()));
         assertThat(DataProcessor.generateAppUsagePeriodMap(
-                0L, hourlyBatteryLevelsPerDay, new ArrayList<>(), new ArrayList<>())).isNull();
+                mContext, hourlyBatteryLevelsPerDay, new ArrayList<>(), new ArrayList<>()))
+                .isNull();
     }
 
     @Test
@@ -1644,7 +1645,7 @@
 
         final Map<Long, Map<String, List<AppUsagePeriod>>> appUsagePeriodMap =
                 DataProcessor.buildAppUsagePeriodList(
-                        appUsageEvents, new ArrayList<>(), 0, 5);
+                        mContext, appUsageEvents, new ArrayList<>(), 0, 5);
 
         assertThat(appUsagePeriodMap).hasSize(2);
         final Map<String, List<AppUsagePeriod>> userMap1 = appUsagePeriodMap.get(1L);
@@ -1668,7 +1669,7 @@
     @Test
     public void buildAppUsagePeriodList_emptyEventList_returnNull() {
         assertThat(DataProcessor.buildAppUsagePeriodList(
-                new ArrayList<>(), new ArrayList<>(), 0, 1)).isNull();
+                mContext, new ArrayList<>(), new ArrayList<>(), 0, 1)).isNull();
     }
 
     @Test
@@ -1680,7 +1681,7 @@
                 AppUsageEventType.DEVICE_SHUTDOWN, /*timestamp=*/ 2));
 
         assertThat(DataProcessor.buildAppUsagePeriodList(
-                appUsageEvents, new ArrayList<>(), 0, 3)).isNull();
+                mContext, appUsageEvents, new ArrayList<>(), 0, 3)).isNull();
     }
 
     @Test
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 efce44e..24be769 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DatabaseUtilsTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DatabaseUtilsTest.java
@@ -16,6 +16,9 @@
 
 package com.android.settings.fuelgauge.batteryusage;
 
+import static android.app.usage.UsageStatsManager.USAGE_SOURCE_CURRENT_ACTIVITY;
+import static android.app.usage.UsageStatsManager.USAGE_SOURCE_TASK_ROOT_ACTIVITY;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
@@ -23,15 +26,19 @@
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
 
+import android.app.usage.IUsageStatsManager;
 import android.content.ContentResolver;
 import android.content.ContentValues;
 import android.content.Context;
 import android.content.Intent;
+import android.content.SharedPreferences;
 import android.content.pm.PackageManager;
 import android.database.MatrixCursor;
 import android.os.BatteryManager;
 import android.os.BatteryUsageStats;
+import android.os.RemoteException;
 import android.os.UserHandle;
 import android.os.UserManager;
 
@@ -67,6 +74,7 @@
     @Mock private BatteryEntry mMockBatteryEntry2;
     @Mock private BatteryEntry mMockBatteryEntry3;
     @Mock private Context mMockContext;
+    @Mock private IUsageStatsManager mUsageStatsManager;
 
     @Before
     public void setUp() {
@@ -77,6 +85,7 @@
         doReturn(mPackageManager).when(mMockContext).getPackageManager();
         doReturn(mPackageManager).when(mContext).getPackageManager();
         DatabaseUtils.getSharedPreferences(mContext).edit().clear().apply();
+        DatabaseUtils.sUsageStatsManager = mUsageStatsManager;
     }
 
     @Test
@@ -423,6 +432,71 @@
     }
 
     @Test
+    public void removeUsageSource_hasNoData() {
+        DatabaseUtils.removeUsageSource(mContext);
+        assertThat(
+                DatabaseUtils
+                        .getSharedPreferences(mContext)
+                        .contains(DatabaseUtils.KEY_LAST_USAGE_SOURCE))
+                .isFalse();
+    }
+
+    @Test
+    public void removeUsageSource_hasData_deleteUsageSource() {
+        final SharedPreferences sharedPreferences = DatabaseUtils.getSharedPreferences(mContext);
+        sharedPreferences
+                .edit()
+                .putInt(DatabaseUtils.KEY_LAST_USAGE_SOURCE, USAGE_SOURCE_TASK_ROOT_ACTIVITY)
+                .apply();
+
+        DatabaseUtils.removeUsageSource(mContext);
+
+        assertThat(
+                DatabaseUtils
+                        .getSharedPreferences(mContext)
+                        .contains(DatabaseUtils.KEY_LAST_USAGE_SOURCE))
+                .isFalse();
+    }
+
+    @Test
+    public void getUsageSource_hasData() {
+        final SharedPreferences sharedPreferences = DatabaseUtils.getSharedPreferences(mContext);
+        sharedPreferences
+                .edit()
+                .putInt(DatabaseUtils.KEY_LAST_USAGE_SOURCE, USAGE_SOURCE_TASK_ROOT_ACTIVITY)
+                .apply();
+
+        assertThat(DatabaseUtils.getUsageSource(mContext))
+                .isEqualTo(USAGE_SOURCE_TASK_ROOT_ACTIVITY);
+    }
+
+    @Test
+    public void getUsageSource_notHasData_writeLoadedData() throws RemoteException {
+        when(mUsageStatsManager.getUsageSource()).thenReturn(USAGE_SOURCE_TASK_ROOT_ACTIVITY);
+
+        assertThat(DatabaseUtils.getUsageSource(mContext))
+                .isEqualTo(USAGE_SOURCE_TASK_ROOT_ACTIVITY);
+        assertThat(
+                DatabaseUtils
+                        .getSharedPreferences(mContext)
+                        .getInt(DatabaseUtils.KEY_LAST_USAGE_SOURCE, USAGE_SOURCE_CURRENT_ACTIVITY))
+                .isEqualTo(USAGE_SOURCE_TASK_ROOT_ACTIVITY);
+    }
+
+    @Test
+    public void getUsageSource_throwException_writeDefaultData() throws RemoteException {
+        when(mUsageStatsManager.getUsageSource()).thenThrow(new RemoteException());
+
+        assertThat(DatabaseUtils.getUsageSource(mContext))
+                .isEqualTo(USAGE_SOURCE_CURRENT_ACTIVITY);
+        assertThat(
+                DatabaseUtils
+                        .getSharedPreferences(mContext)
+                        .getInt(DatabaseUtils.KEY_LAST_USAGE_SOURCE, USAGE_SOURCE_CURRENT_ACTIVITY))
+                .isEqualTo(USAGE_SOURCE_CURRENT_ACTIVITY);
+    }
+
+    @Test
     public void recordDateTime_writeDataIntoSharedPreferences() {
         final String preferenceKey = "test_preference_key";
         DatabaseUtils.recordDateTime(mContext, preferenceKey);