Use taskRootPackageName instead of packageName to load uid when usage
resource is USAGE_SOURCE_TASK_ROOT_ACTIVITY and taskRootPackageName is
not empty.

This logic is consistent with digital wellbeing: assign the screen-on
time onto task root activity when usage resource is
USAGE_SOURCE_TASK_ROOT_ACTIVITY.

Bug: 260964679
Test: make RunSettingsRoboTests + manual
Change-Id: I4c7ed342d8c00951879f5826bf79575f330ce86e
diff --git a/src/com/android/settings/fuelgauge/batteryusage/ConvertUtils.java b/src/com/android/settings/fuelgauge/batteryusage/ConvertUtils.java
index 2e977fd..c1b7818 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/ConvertUtils.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/ConvertUtils.java
@@ -17,7 +17,9 @@
 
 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;
 import android.content.Context;
 import android.content.pm.PackageManager;
@@ -25,7 +27,9 @@
 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;
 import android.util.Base64;
 import android.util.Log;
@@ -167,8 +171,10 @@
     /** Converts to {@link AppUsageEvent} from {@link Event} */
     @Nullable
     public static AppUsageEvent convertToAppUsageEvent(
-            Context context, final Event event, final long userId) {
-        if (event.getPackageName() == null) {
+            Context context, final IUsageStatsManager usageStatsManager, 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.
             // Note that system events like device shutting down should still come with the android
             // package name.
@@ -182,13 +188,20 @@
         appUsageEventBuilder
                 .setTimestamp(event.getTimeStamp())
                 .setType(getAppUsageEventType(event.getEventType()))
-                .setPackageName(event.getPackageName())
+                .setPackageName(packageName)
                 .setUserId(userId);
 
+        final String taskRootPackageName = getTaskRootPackageName(event);
+        if (taskRootPackageName != null) {
+            appUsageEventBuilder.setTaskRootPackageName(taskRootPackageName);
+        }
+
+        final String effectivePackageName =
+                getEffectivePackageName(usageStatsManager, packageName, taskRootPackageName);
         try {
             final long uid = context
                     .getPackageManager()
-                    .getPackageUidAsUser(event.getPackageName(), (int) userId);
+                    .getPackageUidAsUser(effectivePackageName, (int) userId);
             appUsageEventBuilder.setUid(uid);
         } catch (PackageManager.NameNotFoundException e) {
             Log.w(TAG, String.format(
@@ -201,10 +214,6 @@
         } catch (NoClassDefFoundError | NoSuchMethodError e) {
             Log.w(TAG, "UsageEvent instance ID API error");
         }
-        String taskRootPackageName = getTaskRootPackageName(event);
-        if (taskRootPackageName != null) {
-            appUsageEventBuilder.setTaskRootPackageName(taskRootPackageName);
-        }
 
         return appUsageEventBuilder.build();
     }
@@ -268,6 +277,35 @@
     }
 
     /**
+     * Returns the package name the app usage should be attributed to.
+     *
+     * <ul>
+     *   <li>If {@link UsageStatsManager#getUsageSource()} returns {@link
+     *       UsageStatsManager#USAGE_SOURCE_CURRENT_ACTIVITY}, this method will return packageName.
+     *   <li>If {@link UsageStatsManager#getUsageSource()} returns {@link
+     *       UsageStatsManager#USAGE_SOURCE_TASK_ROOT_ACTIVITY}, this method will return
+     *       taskRootPackageName if it exists, or packageName otherwise.
+     * </ul>
+     */
+    @VisibleForTesting
+    static String getEffectivePackageName(
+            final IUsageStatsManager usageStatsManager, final String packageName,
+            final String taskRootPackageName) {
+        int usageSource = getUsageSource(usageStatsManager);
+        switch (usageSource) {
+            case UsageStatsManager.USAGE_SOURCE_TASK_ROOT_ACTIVITY:
+                return !TextUtils.isEmpty(taskRootPackageName)
+                        ? taskRootPackageName
+                        : packageName;
+            case UsageStatsManager.USAGE_SOURCE_CURRENT_ACTIVITY:
+                return packageName;
+            default:
+                Log.e(TAG, "Unexpected usage source: " + usageSource);
+                return packageName;
+        }
+    }
+
+    /**
      * Returns the package name of the task root when this event was reported when {@code event} is
      * one of:
      *
@@ -298,6 +336,20 @@
         }
     }
 
+    /**
+     * 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 AppUsageEventType getAppUsageEventType(final int eventType) {
         switch (eventType) {
             case Event.ACTIVITY_RESUMED:
diff --git a/src/com/android/settings/fuelgauge/batteryusage/DataProcessor.java b/src/com/android/settings/fuelgauge/batteryusage/DataProcessor.java
index 3b3a135..c33e0a3 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/DataProcessor.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/DataProcessor.java
@@ -256,7 +256,8 @@
                             break;
                         }
                         final AppUsageEvent appUsageEvent =
-                                ConvertUtils.convertToAppUsageEvent(context, event, userId);
+                                ConvertUtils.convertToAppUsageEvent(
+                                        context, sUsageStatsManager, event, userId);
                         if (appUsageEvent != null) {
                             numEventsFetched++;
                             appUsageEventList.add(appUsageEvent);
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 aee8398..fe1bff6 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/ConvertUtilsTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/ConvertUtilsTest.java
@@ -15,6 +15,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;
@@ -22,6 +25,7 @@
 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;
@@ -31,6 +35,7 @@
 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;
@@ -56,6 +61,8 @@
     @Mock
     private BatteryUsageStats mBatteryUsageStats;
     @Mock
+    private IUsageStatsManager mUsageStatsManager;
+    @Mock
     private BatteryEntry mMockBatteryEntry;
 
     @Before
@@ -278,7 +285,7 @@
 
         final long userId = 2;
         final AppUsageEvent appUsageEvent = ConvertUtils.convertToAppUsageEvent(
-                mContext, event, userId);
+                mContext, mUsageStatsManager, event, userId);
         assertThat(appUsageEvent.getTimestamp()).isEqualTo(101L);
         assertThat(appUsageEvent.getType()).isEqualTo(AppUsageEventType.ACTIVITY_RESUMED);
         assertThat(appUsageEvent.getPackageName()).isEqualTo("com.android.settings1");
@@ -299,7 +306,7 @@
 
         final long userId = 1;
         final AppUsageEvent appUsageEvent = ConvertUtils.convertToAppUsageEvent(
-                mContext, event, userId);
+                mContext, mUsageStatsManager, event, userId);
         assertThat(appUsageEvent.getTimestamp()).isEqualTo(101L);
         assertThat(appUsageEvent.getType()).isEqualTo(AppUsageEventType.DEVICE_SHUTDOWN);
         assertThat(appUsageEvent.getPackageName()).isEqualTo("com.android.settings1");
@@ -315,7 +322,7 @@
         event.mPackage = null;
 
         final AppUsageEvent appUsageEvent = ConvertUtils.convertToAppUsageEvent(
-                mContext, event, /*userId=*/ 0);
+                mContext, mUsageStatsManager, event, /*userId=*/ 0);
 
         assertThat(appUsageEvent).isNull();
     }
@@ -331,7 +338,7 @@
 
         final long userId = 1;
         final AppUsageEvent appUsageEvent = ConvertUtils.convertToAppUsageEvent(
-                mContext, event, userId);
+                mContext, mUsageStatsManager, event, userId);
 
         assertThat(appUsageEvent).isNull();
     }
@@ -415,4 +422,53 @@
         mContext.getResources().getConfiguration().setLocales(new LocaleList());
         assertThat(ConvertUtils.getLocale(mContext)).isEqualTo(Locale.getDefault());
     }
+
+    @Test
+    public void getEffectivePackageName_currentActivity_returnPackageName() throws RemoteException {
+        when(mUsageStatsManager.getUsageSource()).thenReturn(USAGE_SOURCE_CURRENT_ACTIVITY);
+        final String packageName = "com.android.settings1";
+        final String taskRootPackageName = "com.android.settings2";
+
+        assertThat(ConvertUtils.getEffectivePackageName(
+                mUsageStatsManager, packageName, taskRootPackageName))
+                .isEqualTo(packageName);
+    }
+
+    @Test
+    public void getEffectivePackageName_usageSourceThrowException_returnPackageName()
+            throws RemoteException {
+        when(mUsageStatsManager.getUsageSource()).thenThrow(new RemoteException());
+        final String packageName = "com.android.settings1";
+        final String taskRootPackageName = "com.android.settings2";
+
+        assertThat(ConvertUtils.getEffectivePackageName(
+                mUsageStatsManager, packageName, taskRootPackageName))
+                .isEqualTo(packageName);
+    }
+
+    @Test
+    public void getEffectivePackageName_rootActivity_returnTaskRootPackageName()
+            throws RemoteException {
+        when(mUsageStatsManager.getUsageSource()).thenReturn(USAGE_SOURCE_TASK_ROOT_ACTIVITY);
+        final String packageName = "com.android.settings1";
+        final String taskRootPackageName = "com.android.settings2";
+
+        assertThat(ConvertUtils.getEffectivePackageName(
+                mUsageStatsManager, packageName, taskRootPackageName))
+                .isEqualTo(taskRootPackageName);
+    }
+
+    @Test
+    public void getEffectivePackageName_nullOrEmptyTaskRoot_returnPackageName()
+            throws RemoteException {
+        when(mUsageStatsManager.getUsageSource()).thenReturn(USAGE_SOURCE_TASK_ROOT_ACTIVITY);
+        final String packageName = "com.android.settings1";
+
+        assertThat(ConvertUtils.getEffectivePackageName(
+                mUsageStatsManager, packageName, /*taskRootPackageName=*/ null))
+                .isEqualTo(packageName);
+        assertThat(ConvertUtils.getEffectivePackageName(
+                mUsageStatsManager, packageName, /*taskRootPackageName=*/ ""))
+                .isEqualTo(packageName);
+    }
 }