Merge "Add more sane multi-profile app attribution." into oc-dev
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;
     }
+
 }