Merge "Use correct uid when querying usage data for specific user."
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);
+    }
 }