1 /*
2  * Copyright (C) 2021 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package android.usage.cts_root;
17 
18 import static android.os.storage.StorageManager.UUID_DEFAULT;
19 
20 import static com.google.common.truth.Truth.assertThat;
21 
22 import android.app.usage.StorageStats;
23 import android.app.usage.StorageStatsManager;
24 import android.content.Context;
25 import android.content.pm.PackageInfo;
26 import android.content.pm.PackageStats;
27 import android.os.UserHandle;
28 import android.util.Log;
29 
30 import androidx.test.core.app.ApplicationProvider;
31 
32 import com.android.server.LocalManagerRegistry;
33 import com.android.server.usage.StorageStatsManagerLocal;
34 
35 import org.junit.After;
36 import org.junit.Before;
37 import org.junit.Test;
38 
39 import java.util.ArrayList;
40 import java.util.List;
41 import java.util.Objects;
42 
43 import org.junit.Test;
44 
45 public class StorageStatsManagerLocalTest {
46     private static final String TAG = "StorageStatsManagerLocalTest";
47 
48     private static final int AUGMENT_BYTES_FOR_PACKAGE = 7907;
49     private static final int AUGMENT_BYTES_FOR_UID = 6883;
50     private static final int AUGMENT_BYTES_FOR_USER = 5741;
51 
52     @Test
testStorageAugmenter()53     public void testStorageAugmenter() throws Exception {
54         Context context = ApplicationProvider.getApplicationContext();
55         StorageStatsManager storageStatsManager =
56                 context.getSystemService(StorageStatsManager.class);
57         String packageName = context.getPackageName();
58         UserHandle user = context.getUser();
59         int uid = android.os.Process.myUid();
60 
61         // Collect stats before augmentation
62         StorageStats beforeStatsForPkg =
63                 storageStatsManager.queryStatsForPackage(UUID_DEFAULT, packageName, user);
64         StorageStats beforeStatsForUid = storageStatsManager.queryStatsForUid(UUID_DEFAULT, uid);
65         StorageStats beforeStatsForUser = storageStatsManager.queryStatsForUser(UUID_DEFAULT, user);
66 
67         // Register our augmenter
68         LocalManagerRegistry.getManager(StorageStatsManagerLocal.class)
69                 .registerStorageStatsAugmenter(new TestStorageStatsAugmenter(), TAG);
70 
71         // Collect stats after augmentation
72         StorageStats afterStatsForPkg =
73                 storageStatsManager.queryStatsForPackage(UUID_DEFAULT, packageName, user);
74         StorageStats afterStatsForUid = storageStatsManager.queryStatsForUid(UUID_DEFAULT, uid);
75         StorageStats afterStatsForUser = storageStatsManager.queryStatsForUser(UUID_DEFAULT, user);
76 
77         // Verify the storage size increase.
78         assertThat(afterStatsForPkg.getDataBytes() - beforeStatsForPkg.getDataBytes())
79                 .isEqualTo(AUGMENT_BYTES_FOR_PACKAGE);
80         assertThat(afterStatsForUid.getDataBytes() - beforeStatsForUid.getDataBytes())
81                 .isEqualTo(AUGMENT_BYTES_FOR_UID);
82         assertThat(afterStatsForUser.getDataBytes() - beforeStatsForUser.getDataBytes())
83                 .isEqualTo(AUGMENT_BYTES_FOR_USER);
84     }
85 
86     private class TestStorageStatsAugmenter implements
87             StorageStatsManagerLocal.StorageStatsAugmenter {
88         @Override
augmentStatsForPackageForUser( PackageStats stats, String packageName, UserHandle userHandle, boolean canCallerAccessAllStats)89         public void augmentStatsForPackageForUser(
90                 PackageStats stats,
91                 String packageName,
92                 UserHandle userHandle,
93                 boolean canCallerAccessAllStats) {
94             stats.dataSize += AUGMENT_BYTES_FOR_PACKAGE;
95         }
96 
97         @Override
augmentStatsForUid( PackageStats stats, int uid, boolean canCallerAccessAllStats)98         public void augmentStatsForUid(
99                 PackageStats stats, int uid, boolean canCallerAccessAllStats) {
100             stats.dataSize += AUGMENT_BYTES_FOR_UID;
101         }
102 
103         @Override
augmentStatsForUser(PackageStats stats, UserHandle userHandle)104         public void augmentStatsForUser(PackageStats stats, UserHandle userHandle) {
105             stats.dataSize += AUGMENT_BYTES_FOR_USER;
106         }
107     }
108 }
109