Add more sane multi-profile app attribution.

Due to issues w.r.t. attribution across multiple users, we originally
duplicated the previous implementation which zeroed out the sizes. This,
however, caused system bloat due to the change in how we calculate the
system size.

By attributing apps which do not exist in the primary profile to the
first user that shows up with it installed, we can avoid accidentally
attributing it to the system size.

Bug: 62623731
Test: Settings unittest & Settings robotest
Change-Id: I9514c9ecef62ea6270723f62e6bf27c69b75f180
(cherry picked from commit b088010d12bdb709fc1aa4269a86924aab58be47)
diff --git a/src/com/android/settings/deviceinfo/StorageProfileFragment.java b/src/com/android/settings/deviceinfo/StorageProfileFragment.java
index dee6793..f5129ed 100644
--- a/src/com/android/settings/deviceinfo/StorageProfileFragment.java
+++ b/src/com/android/settings/deviceinfo/StorageProfileFragment.java
@@ -120,7 +120,6 @@
     @Override
     public void onLoadFinished(Loader<SparseArray<AppsStorageResult>> loader,
             SparseArray<AppsStorageResult> result) {
-        scrubAppsFromResult(result.get(mUserId));
         mPreferenceController.onLoadFinished(result, mUserId);
     }
 
@@ -132,17 +131,4 @@
     void setPreferenceController(StorageItemPreferenceController controller) {
         mPreferenceController = controller;
     }
-
-    private AppsStorageResult scrubAppsFromResult(AppsStorageResult result) {
-        if (result == null) {
-            return null;
-        }
-
-        // TODO(b/35927909): Attribute app sizes better than zeroing out for profiles.
-        result.gamesSize = 0;
-        result.musicAppsSize = 0;
-        result.videoAppsSize = 0;
-        result.otherAppsSize = 0;
-        return result;
-    }
 }
diff --git a/src/com/android/settings/deviceinfo/storage/StorageAsyncLoader.java b/src/com/android/settings/deviceinfo/storage/StorageAsyncLoader.java
index 74474b3..0b685d0 100644
--- a/src/com/android/settings/deviceinfo/storage/StorageAsyncLoader.java
+++ b/src/com/android/settings/deviceinfo/storage/StorageAsyncLoader.java
@@ -25,6 +25,7 @@
 import android.content.pm.PackageManager.NameNotFoundException;
 import android.content.pm.UserInfo;
 import android.os.UserHandle;
+import android.util.ArraySet;
 import android.util.Log;
 import android.util.SparseArray;
 
@@ -34,6 +35,8 @@
 import com.android.settingslib.applications.StorageStatsSource;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
 
 /**
@@ -48,6 +51,7 @@
     private String mUuid;
     private StorageStatsSource mStatsManager;
     private PackageManagerWrapper mPackageManager;
+    private ArraySet<String> mSeenPackages;
 
     public StorageAsyncLoader(Context context, UserManagerWrapper userManager,
             String uuid, StorageStatsSource source, PackageManagerWrapper pm) {
@@ -64,8 +68,18 @@
     }
 
     private SparseArray<AppsStorageResult> loadApps() {
+        mSeenPackages = new ArraySet<>();
         SparseArray<AppsStorageResult> result = new SparseArray<>();
         List<UserInfo> infos = mUserManager.getUsers();
+        // Sort the users by user id ascending.
+        Collections.sort(
+                infos,
+                new Comparator<UserInfo>() {
+                    @Override
+                    public int compare(UserInfo userInfo, UserInfo otherUser) {
+                        return Integer.compare(userInfo.id, otherUser.id);
+                    }
+                });
         for (int i = 0, userCount = infos.size(); i < userCount; i++) {
             UserInfo info = infos.get(i);
             result.put(info.id, getStorageResultForUser(info.id));
@@ -93,10 +107,11 @@
 
             long blamedSize = stats.getDataBytes() - stats.getCacheBytes();
 
-            // Only count app code against the current user; we don't want
-            // double-counting on multi-user devices.
-            if (userId == UserHandle.myUserId()) {
+            // This isn't quite right because it slams the first user by user id with the whole code
+            // size, but this ensures that we count all apps seen once.
+            if (!mSeenPackages.contains(app.packageName)) {
                 blamedSize += stats.getCodeBytes();
+                mSeenPackages.add(app.packageName);
             }
 
             switch (app.category) {
@@ -140,6 +155,7 @@
         public long musicAppsSize;
         public long videoAppsSize;
         public long otherAppsSize;
+        public long cacheSize;
         public StorageStatsSource.ExternalStorageStats externalStats;
     }
 
diff --git a/src/com/android/settings/deviceinfo/storage/UserProfileController.java b/src/com/android/settings/deviceinfo/storage/UserProfileController.java
index 18fa7b7..fc297ca 100644
--- a/src/com/android/settings/deviceinfo/storage/UserProfileController.java
+++ b/src/com/android/settings/deviceinfo/storage/UserProfileController.java
@@ -96,7 +96,13 @@
         int userId = mUser.id;
         StorageAsyncLoader.AppsStorageResult result = stats.get(userId);
         if (result != null) {
-            setSize(result.externalStats.totalBytes, mTotalSizeBytes);
+            setSize(
+                    result.externalStats.totalBytes
+                            + result.otherAppsSize
+                            + result.videoAppsSize
+                            + result.musicAppsSize
+                            + result.gamesSize,
+                    mTotalSizeBytes);
         }
     }
 
diff --git a/tests/robotests/src/com/android/settings/deviceinfo/StorageProfileFragmentTest.java b/tests/robotests/src/com/android/settings/deviceinfo/StorageProfileFragmentTest.java
index 03f15bb..3ea8016 100644
--- a/tests/robotests/src/com/android/settings/deviceinfo/StorageProfileFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/deviceinfo/StorageProfileFragmentTest.java
@@ -43,7 +43,7 @@
     private ArgumentCaptor<SparseArray<StorageAsyncLoader.AppsStorageResult>> mCaptor;
 
     @Test
-    public void verifyAppSizesAreZeroedOut() {
+    public void verifyAppSizesAreNotZeroedOut() {
         StorageItemPreferenceController controller = mock(StorageItemPreferenceController.class);
         StorageProfileFragment fragment = new StorageProfileFragment();
         StorageAsyncLoader.AppsStorageResult result = new StorageAsyncLoader.AppsStorageResult();
@@ -62,10 +62,10 @@
         verify(controller).onLoadFinished(mCaptor.capture(), anyInt());
 
         StorageAsyncLoader.AppsStorageResult extractedResult = mCaptor.getValue().get(0);
-        assertThat(extractedResult.musicAppsSize).isEqualTo(0);
-        assertThat(extractedResult.videoAppsSize).isEqualTo(0);
-        assertThat(extractedResult.otherAppsSize).isEqualTo(0);
-        assertThat(extractedResult.gamesSize).isEqualTo(0);
+        assertThat(extractedResult.musicAppsSize).isEqualTo(100);
+        assertThat(extractedResult.videoAppsSize).isEqualTo(400);
+        assertThat(extractedResult.otherAppsSize).isEqualTo(200);
+        assertThat(extractedResult.gamesSize).isEqualTo(300);
         assertThat(extractedResult.externalStats.audioBytes).isEqualTo(1);
         assertThat(extractedResult.externalStats.videoBytes).isEqualTo(2);
         assertThat(extractedResult.externalStats.imageBytes).isEqualTo(3);
diff --git a/tests/unit/src/com/android/settings/deviceinfo/storage/StorageAsyncLoaderTest.java b/tests/unit/src/com/android/settings/deviceinfo/storage/StorageAsyncLoaderTest.java
index 79a4595..28bd861 100644
--- a/tests/unit/src/com/android/settings/deviceinfo/storage/StorageAsyncLoaderTest.java
+++ b/tests/unit/src/com/android/settings/deviceinfo/storage/StorageAsyncLoaderTest.java
@@ -27,6 +27,7 @@
 
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager.NameNotFoundException;
 import android.content.pm.UserInfo;
 import android.os.UserHandle;
 import android.support.test.filters.SmallTest;
@@ -74,7 +75,8 @@
         MockitoAnnotations.initMocks(this);
         mInfo = new ArrayList<>();
         mLoader = new StorageAsyncLoader(mContext, mUserManager, "id", mSource, mPackageManager);
-        when(mPackageManager.getInstalledApplicationsAsUser(anyInt(), anyInt())).thenReturn(mInfo);
+        when(mPackageManager.getInstalledApplicationsAsUser(eq(PRIMARY_USER_ID), anyInt()))
+                .thenReturn(mInfo);
         UserInfo info = new UserInfo();
         mUsers = new ArrayList<>();
         mUsers.add(info);
@@ -174,13 +176,37 @@
         info.category = ApplicationInfo.CATEGORY_UNDEFINED;
         mInfo.add(info);
         when(mSource.getStatsForPackage(anyString(), anyString(), any(UserHandle.class)))
-                .thenThrow(new IllegalStateException());
+                .thenThrow(new NameNotFoundException());
 
         SparseArray<StorageAsyncLoader.AppsStorageResult> result = mLoader.loadInBackground();
 
         // Should not crash.
     }
 
+    @Test
+    public void testPackageIsNotDoubleCounted() throws Exception {
+        UserInfo info = new UserInfo();
+        info.id = SECONDARY_USER_ID;
+        mUsers.add(info);
+        when(mSource.getExternalStorageStats(anyString(), eq(UserHandle.SYSTEM)))
+                .thenReturn(new StorageStatsSource.ExternalStorageStats(9, 2, 3, 4, 0));
+        when(mSource.getExternalStorageStats(anyString(), eq(new UserHandle(SECONDARY_USER_ID))))
+                .thenReturn(new StorageStatsSource.ExternalStorageStats(10, 3, 3, 4, 0));
+        addPackage(PACKAGE_NAME_1, 0, 1, 10, ApplicationInfo.CATEGORY_VIDEO);
+        ArrayList<ApplicationInfo> secondaryUserApps = new ArrayList<>();
+        ApplicationInfo appInfo = new ApplicationInfo();
+        appInfo.packageName = PACKAGE_NAME_1;
+        appInfo.category = ApplicationInfo.CATEGORY_VIDEO;
+        secondaryUserApps.add(appInfo);
+
+        SparseArray<StorageAsyncLoader.AppsStorageResult> result = mLoader.loadInBackground();
+
+        assertThat(result.size()).isEqualTo(2);
+        assertThat(result.get(PRIMARY_USER_ID).videoAppsSize).isEqualTo(11L);
+        // No code size for the second user.
+        assertThat(result.get(SECONDARY_USER_ID).videoAppsSize).isEqualTo(10L);
+    }
+
     private ApplicationInfo addPackage(String packageName, long cacheSize, long codeSize,
             long dataSize, int category) throws Exception {
         StorageStatsSource.AppStorageStats storageStats =
@@ -196,4 +222,5 @@
         mInfo.add(info);
         return info;
     }
+
 }