Use correct uid when querying usage data for specific user.

- the data usage list shows usage data for each app, and also the
aggregated data for work apps and other user. When user clicks on the
app item, we will shows the detail usage data for that item using the
app item key. However, if the app item is not for app but for user, the
key will not be the app uid, but a key built from the user id. So,
querying using the key will not get any data.
- for app item corresponding for user, need to sum up all usage data
for the uids associated with the user.

Change-Id: I1a185ee6b4d59e477f3a03cade97d85d6982e416
Fixes: 122200400
Test: make RunSettingsRoboTests
diff --git a/src/com/android/settings/applications/appinfo/AppDataUsagePreferenceController.java b/src/com/android/settings/applications/appinfo/AppDataUsagePreferenceController.java
index 1bfcd7d..7d9fdcf 100644
--- a/src/com/android/settings/applications/appinfo/AppDataUsagePreferenceController.java
+++ b/src/com/android/settings/applications/appinfo/AppDataUsagePreferenceController.java
@@ -89,7 +89,7 @@
     public Loader<List<NetworkCycleDataForUid>> onCreateLoader(int id, Bundle args) {
         final NetworkTemplate template = getTemplate(mContext);
         return NetworkCycleDataForUidLoader.builder(mContext)
-            .setUid(mParent.getAppEntry().info.uid)
+            .addUid(mParent.getAppEntry().info.uid)
             .setRetrieveDetail(false)
             .setNetworkTemplate(template)
             .setSubscriberId(template.getSubscriberId())
diff --git a/src/com/android/settings/datausage/AppDataUsage.java b/src/com/android/settings/datausage/AppDataUsage.java
index 329bf68..4f38843 100644
--- a/src/com/android/settings/datausage/AppDataUsage.java
+++ b/src/com/android/settings/datausage/AppDataUsage.java
@@ -94,10 +94,12 @@
     private SpinnerPreference mCycle;
     private RestrictedSwitchPreference mUnrestrictedData;
     private DataSaverBackend mDataSaverBackend;
+    private Context mContext;
 
     @Override
     public void onCreate(Bundle icicle) {
         super.onCreate(icicle);
+        mContext = getContext();
         mPackageManager = getPackageManager();
         final Bundle args = getArguments();
 
@@ -105,9 +107,8 @@
         mTemplate = (args != null) ? (NetworkTemplate) args.getParcelable(ARG_NETWORK_TEMPLATE)
                 : null;
         if (mTemplate == null) {
-            Context context = getContext();
-            mTemplate = DataUsageUtils.getDefaultTemplate(context,
-                    DataUsageUtils.getDefaultSubscriptionId(context));
+            mTemplate = DataUsageUtils.getDefaultTemplate(mContext,
+                    DataUsageUtils.getDefaultSubscriptionId(mContext));
         }
         if (mAppItem == null) {
             int uid = (args != null) ? args.getInt(AppInfoBase.ARG_PACKAGE_UID, -1)
@@ -131,7 +132,7 @@
         mBackgroundUsage = findPreference(KEY_BACKGROUND_USAGE);
 
         mCycle = (SpinnerPreference) findPreference(KEY_CYCLE);
-        mCycleAdapter = new CycleAdapter(getContext(), mCycle, mCycleListener);
+        mCycleAdapter = new CycleAdapter(mContext, mCycle, mCycleListener);
 
         if (mAppItem.key > 0) {
             if (mPackages.size() != 0) {
@@ -155,7 +156,7 @@
                         KEY_UNRESTRICTED_DATA);
                 mUnrestrictedData.setOnPreferenceChangeListener(this);
             }
-            mDataSaverBackend = new DataSaverBackend(getContext());
+            mDataSaverBackend = new DataSaverBackend(mContext);
             mAppSettings = findPreference(KEY_APP_SETTINGS);
 
             mAppSettingsIntent = new Intent(Intent.ACTION_MANAGE_NETWORK_USAGE);
@@ -256,7 +257,7 @@
 
     private void updatePrefs(boolean restrictBackground, boolean unrestrictData) {
         final EnforcedAdmin admin = RestrictedLockUtilsInternal.checkIfMeteredDataRestricted(
-                getContext(), mPackageName, UserHandle.getUserId(mAppItem.key));
+                mContext, mPackageName, UserHandle.getUserId(mAppItem.key));
         if (mRestrictBackground != null) {
             mRestrictBackground.setChecked(!restrictBackground);
             mRestrictBackground.setDisabledByAdmin(admin);
@@ -294,11 +295,10 @@
             foregroundBytes = data.getForegroudUsage();
         }
         final long totalBytes = backgroundBytes + foregroundBytes;
-        final Context context = getContext();
 
-        mTotalUsage.setSummary(DataUsageUtils.formatDataUsage(context, totalBytes));
-        mForegroundUsage.setSummary(DataUsageUtils.formatDataUsage(context, foregroundBytes));
-        mBackgroundUsage.setSummary(DataUsageUtils.formatDataUsage(context, backgroundBytes));
+        mTotalUsage.setSummary(DataUsageUtils.formatDataUsage(mContext, totalBytes));
+        mForegroundUsage.setSummary(DataUsageUtils.formatDataUsage(mContext, foregroundBytes));
+        mBackgroundUsage.setSummary(DataUsageUtils.formatDataUsage(mContext, backgroundBytes));
     }
 
     private boolean getAppRestrictBackground() {
@@ -364,16 +364,24 @@
         }
     };
 
-    private final LoaderManager.LoaderCallbacks<List<NetworkCycleDataForUid>> mUidDataCallbacks =
+    @VisibleForTesting
+    final LoaderManager.LoaderCallbacks<List<NetworkCycleDataForUid>> mUidDataCallbacks =
         new LoaderManager.LoaderCallbacks<List<NetworkCycleDataForUid>>() {
             @Override
             public Loader<List<NetworkCycleDataForUid>> onCreateLoader(int id, Bundle args) {
-                return NetworkCycleDataForUidLoader.builder(getContext())
-                    .setUid(mAppItem.key)
-                    .setRetrieveDetail(true)
+                final NetworkCycleDataForUidLoader.Builder builder
+                    = NetworkCycleDataForUidLoader.builder(mContext);
+                builder.setRetrieveDetail(true)
                     .setNetworkTemplate(mTemplate)
-                    .setSubscriberId(mTemplate.getSubscriberId())
-                    .build();
+                    .setSubscriberId(mTemplate.getSubscriberId());
+                if (mAppItem.category == AppItem.CATEGORY_USER) {
+                    for (int i = 0; i < mAppItem.uids.size(); i++) {
+                        builder.addUid(mAppItem.uids.keyAt(i));
+                    }
+                } else {
+                    builder.addUid(mAppItem.key);
+                }
+                return builder.build();
             }
 
             @Override
diff --git a/tests/robotests/src/com/android/settings/datausage/AppDataUsageTest.java b/tests/robotests/src/com/android/settings/datausage/AppDataUsageTest.java
index 37bb3cd..12559f4 100644
--- a/tests/robotests/src/com/android/settings/datausage/AppDataUsageTest.java
+++ b/tests/robotests/src/com/android/settings/datausage/AppDataUsageTest.java
@@ -16,6 +16,8 @@
 
 package com.android.settings.datausage;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
@@ -31,7 +33,9 @@
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.net.NetworkPolicyManager;
+import android.net.NetworkTemplate;
 import android.os.Bundle;
+import android.text.format.DateUtils;
 import android.util.ArraySet;
 import android.view.View;
 
@@ -47,6 +51,7 @@
 import com.android.settingslib.RestrictedLockUtils.EnforcedAdmin;
 import com.android.settingslib.RestrictedSwitchPreference;
 import com.android.settingslib.net.NetworkCycleDataForUid;
+import com.android.settingslib.net.NetworkCycleDataForUidLoader;
 
 import org.junit.After;
 import org.junit.Before;
@@ -188,7 +193,7 @@
         ReflectionHelpers.setField(mFragment, "mBackgroundUsage", preference);
         ReflectionHelpers.setField(mFragment, "mForegroundUsage", preference);
         ReflectionHelpers.setField(mFragment, "mTotalUsage", preference);
-        doReturn(RuntimeEnvironment.application).when(mFragment).getContext();
+        ReflectionHelpers.setField(mFragment, "mContext", RuntimeEnvironment.application);
 
         mFragment.bindData(0 /* position */);
 
@@ -199,7 +204,7 @@
     public void bindData_hasAppUsageData_shouldShowCycleSpinnerAndUpdateUsageSummary() {
         mFragment = spy(new AppDataUsage());
         final Context context = RuntimeEnvironment.application;
-        doReturn(context).when(mFragment).getContext();
+        ReflectionHelpers.setField(mFragment, "mContext", context);
         final long backgroundBytes = 1234L;
         final long foregroundBytes = 5678L;
         final List<NetworkCycleDataForUid> appUsage = new ArrayList<>();
@@ -223,4 +228,53 @@
         verify(backgroundPref).setSummary(DataUsageUtils.formatDataUsage(context, backgroundBytes));
         verify(foregroundPref).setSummary(DataUsageUtils.formatDataUsage(context, foregroundBytes));
     }
+
+    @Test
+    public void onCreateLoader_categoryApp_shouldQueryDataUsageUsingAppKey() {
+        mFragment = new AppDataUsage();
+        final Context context = RuntimeEnvironment.application;
+        final int testUid = 123123;
+        final AppItem appItem = new AppItem(testUid);
+        appItem.category = AppItem.CATEGORY_APP;
+        ReflectionHelpers.setField(mFragment, "mContext", context);
+        ReflectionHelpers.setField(mFragment, "mAppItem", appItem);
+        ReflectionHelpers.setField(mFragment, "mTemplate",
+            NetworkTemplate.buildTemplateWifiWildcard());
+        final long end = System.currentTimeMillis();
+        final long start = end - (DateUtils.WEEK_IN_MILLIS * 4);
+
+        final NetworkCycleDataForUidLoader loader = (NetworkCycleDataForUidLoader)
+            mFragment.mUidDataCallbacks.onCreateLoader(0, Bundle.EMPTY);
+
+        final List<Integer> uids = loader.getUids();
+        assertThat(uids).hasSize(1);
+        assertThat(uids.get(0)).isEqualTo(testUid);
+    }
+
+    @Test
+    public void onCreateLoader_categoryUser_shouldQueryDataUsageUsingAssociatedUids() {
+        mFragment = new AppDataUsage();
+        final Context context = RuntimeEnvironment.application;
+        final int testUserId = 11;
+        final AppItem appItem = new AppItem(testUserId);
+        appItem.category = AppItem.CATEGORY_USER;
+        appItem.addUid(123);
+        appItem.addUid(456);
+        appItem.addUid(789);
+        ReflectionHelpers.setField(mFragment, "mContext", context);
+        ReflectionHelpers.setField(mFragment, "mAppItem", appItem);
+        ReflectionHelpers.setField(mFragment, "mTemplate",
+            NetworkTemplate.buildTemplateWifiWildcard());
+        final long end = System.currentTimeMillis();
+        final long start = end - (DateUtils.WEEK_IN_MILLIS * 4);
+
+        final NetworkCycleDataForUidLoader loader = (NetworkCycleDataForUidLoader)
+            mFragment.mUidDataCallbacks.onCreateLoader(0, Bundle.EMPTY);
+
+        final List<Integer> uids = loader.getUids();
+        assertThat(uids).hasSize(3);
+        assertThat(uids.get(0)).isEqualTo(123);
+        assertThat(uids.get(1)).isEqualTo(456);
+        assertThat(uids.get(2)).isEqualTo(789);
+    }
 }