1 package org.robolectric.shadows; 2 3 import static com.google.common.base.Preconditions.checkArgument; 4 5 import android.app.usage.StorageStats; 6 import android.app.usage.StorageStatsManager; 7 import android.content.pm.PackageManager; 8 import android.os.Build; 9 import android.os.Parcel; 10 import android.os.UserHandle; 11 import android.os.storage.StorageManager; 12 import com.google.auto.value.AutoValue; 13 import java.io.IOException; 14 import java.util.Map; 15 import java.util.UUID; 16 import java.util.concurrent.ConcurrentHashMap; 17 import org.robolectric.annotation.Implementation; 18 import org.robolectric.annotation.Implements; 19 20 /** 21 * Fake implementation of {@link android.app.usage.StorageStatsManager} that provides a fake 22 * implementation of query for {@link StorageStats} of a package. 23 */ 24 @Implements(value = StorageStatsManager.class, minSdk = Build.VERSION_CODES.O) 25 public class ShadowStorageStatsManager { 26 27 public static final long DEFAULT_STORAGE_FREE_BYTES = 4L * 1024L * 1024L * 1024L; // 4 GB 28 public static final long DEFAULT_STORAGE_TOTAL_BYTES = 8L * 1024L * 1024L * 1024L; // 8 GB 29 30 private final Map<UUID, FreeAndTotalBytesPair> freeAndTotalBytesMap = 31 createFreeAndTotalBytesMapWithSingleEntry( 32 StorageManager.UUID_DEFAULT, DEFAULT_STORAGE_FREE_BYTES, DEFAULT_STORAGE_TOTAL_BYTES); 33 private final Map<StorageStatsKey, StorageStats> storageStatsMapForPackage = 34 new ConcurrentHashMap<>(); 35 private final Map<StorageStatsKey, StorageStats> storageStatsMapForUser = 36 new ConcurrentHashMap<>(); 37 38 /** 39 * Sets the {@code storageUuid} to return the specified {@code freeBytes} and {@code totalBytes} 40 * when queried in {@link #getFreeBytes} and {@link #getTotalBytes} respectively. 41 * 42 * <p>Both {@code freeBytes} and {@code totalBytes} have to be non-negative, else this method will 43 * throw {@link IllegalArgumentException}. 44 */ setStorageDeviceFreeAndTotalBytes(UUID storageUuid, long freeBytes, long totalBytes)45 public void setStorageDeviceFreeAndTotalBytes(UUID storageUuid, long freeBytes, long totalBytes) { 46 checkArgument( 47 freeBytes >= 0 && totalBytes >= 0, "Both freeBytes and totalBytes must be non-negative!"); 48 freeAndTotalBytesMap.put(storageUuid, FreeAndTotalBytesPair.create(freeBytes, totalBytes)); 49 } 50 51 /** 52 * Removes a storage device identified by {@code storageUuid} if it's currently present. 53 * Otherwise, this method will be a no-op. 54 */ removeStorageDevice(UUID storageUuid)55 public void removeStorageDevice(UUID storageUuid) { 56 freeAndTotalBytesMap.remove(storageUuid); 57 } 58 59 /** 60 * Sets the {@link StorageStats} for given {@code storageUuid}, {@code packageName} and {@code 61 * userHandle}. If {@code queryStatsForPackage} is called with matching {@code storageUuid}, 62 * {@code packageName} and {@code userHandle}, the {@code storageStatsToReturn} will be returned 63 * directly. If {@code queryStatsForUser} is called with matching {@code storageUuid} and {@code 64 * userHandle}, then an accumulated {@link StorageStats} will be returned. 65 */ addStorageStats( UUID storageUuid, String packageName, UserHandle userHandle, StorageStats storageStatsToReturn)66 public void addStorageStats( 67 UUID storageUuid, 68 String packageName, 69 UserHandle userHandle, 70 StorageStats storageStatsToReturn) { 71 StorageStatsKey storageStatsKeyForPackage = 72 StorageStatsKey.create(storageUuid, packageName, userHandle); 73 StorageStats storageStatsForPackage = storageStatsMapForPackage.get(storageStatsKeyForPackage); 74 storageStatsMapForPackage.put(storageStatsKeyForPackage, storageStatsToReturn); 75 76 StorageStatsKey storageStatsKeyForUser = 77 StorageStatsKey.create(storageUuid, /* packageName= */ "", userHandle); 78 StorageStats storageStatsForUser = storageStatsMapForUser.get(storageStatsKeyForUser); 79 if (storageStatsForUser == null) { 80 storageStatsMapForUser.put(storageStatsKeyForUser, storageStatsToReturn); 81 } else { 82 long moreAppBytes = storageStatsToReturn.getAppBytes(); 83 long moreDataBytes = storageStatsToReturn.getDataBytes(); 84 long moreCacheBytes = storageStatsToReturn.getCacheBytes(); 85 if (storageStatsForPackage != null) { 86 moreAppBytes -= storageStatsForPackage.getAppBytes(); 87 moreDataBytes -= storageStatsForPackage.getDataBytes(); 88 moreCacheBytes -= storageStatsForPackage.getCacheBytes(); 89 } 90 Parcel parcel = Parcel.obtain(); 91 parcel.writeLong(storageStatsForUser.getAppBytes() + moreAppBytes); 92 parcel.writeLong(storageStatsForUser.getDataBytes() + moreDataBytes); 93 parcel.writeLong(storageStatsForUser.getCacheBytes() + moreCacheBytes); 94 parcel.setDataPosition(0); 95 storageStatsMapForUser.put( 96 storageStatsKeyForUser, StorageStats.CREATOR.createFromParcel(parcel)); 97 } 98 } 99 100 /** Clears all {@link StorageStats} set in {@link ShadowStorageStatsManager#addStorageStats}. */ clearStorageStats()101 public void clearStorageStats() { 102 storageStatsMapForPackage.clear(); 103 storageStatsMapForUser.clear(); 104 } 105 106 /** 107 * Fake implementation of {@link StorageStatsManager#getFreeBytes} that returns test setup values. 108 * This fake implementation does not check for access permission. It only checks for arguments 109 * matching those set in {@link ShadowStorageStatsManager#setStorageDeviceFreeAndTotalBytes}. 110 */ 111 @Implementation getFreeBytes(UUID storageUuid)112 protected long getFreeBytes(UUID storageUuid) throws IOException { 113 FreeAndTotalBytesPair freeAndTotalBytesPair = freeAndTotalBytesMap.get(storageUuid); 114 if (freeAndTotalBytesPair == null) { 115 throw new IOException( 116 "getFreeBytes with non-existent storageUuid! Did you forget to call" 117 + " setStorageDeviceFreeAndTotalBytes?"); 118 } 119 return freeAndTotalBytesPair.freeBytes(); 120 } 121 122 /** 123 * Fake implementation of {@link StorageStatsManager#getTotalBytes} that returns test setup 124 * values. This fake implementation does not check for access permission. It only checks for 125 * arguments matching those set in {@link 126 * ShadowStorageStatsManager#setStorageDeviceFreeAndTotalBytes}. 127 */ 128 @Implementation getTotalBytes(UUID storageUuid)129 protected long getTotalBytes(UUID storageUuid) throws IOException { 130 FreeAndTotalBytesPair freeAndTotalBytesPair = freeAndTotalBytesMap.get(storageUuid); 131 if (freeAndTotalBytesPair == null) { 132 throw new IOException( 133 "getTotalBytes with non-existent storageUuid! Did you forget to call" 134 + " setStorageDeviceFreeAndTotalBytes?"); 135 } 136 return freeAndTotalBytesPair.totalBytes(); 137 } 138 139 /** 140 * Fake implementation of {@link StorageStatsManager#queryStatsForPackage} that returns test setup 141 * values. This fake implementation does not check for access permission. It only checks for 142 * arguments matching those set in {@link ShadowStorageStatsManager#addStorageStats}. 143 */ 144 @Implementation queryStatsForPackage(UUID storageUuid, String packageName, UserHandle user)145 protected StorageStats queryStatsForPackage(UUID storageUuid, String packageName, UserHandle user) 146 throws PackageManager.NameNotFoundException, IOException { 147 StorageStats storageStat = 148 storageStatsMapForPackage.get(StorageStatsKey.create(storageUuid, packageName, user)); 149 if (storageStat == null) { 150 throw new PackageManager.NameNotFoundException( 151 "queryStatsForPackage with non matching arguments. Did you forget to call" 152 + " addStorageStats?"); 153 } 154 return storageStat; 155 } 156 157 /** 158 * Fake implementation of {@link StorageStatsManager#queryStatsForUser} that returns an 159 * accumulated {@link StorageStats} based on the setup values for the user. This fake 160 * implementation does not check for access permission. It only checks for arguments matching 161 * those set in {@link ShadowStorageStatsManager#addStorageStats}. 162 */ 163 @Implementation queryStatsForUser(UUID storageUuid, UserHandle user)164 protected StorageStats queryStatsForUser(UUID storageUuid, UserHandle user) 165 throws PackageManager.NameNotFoundException, IOException { 166 StorageStats storageStat = 167 storageStatsMapForUser.get( 168 StorageStatsKey.create(storageUuid, /* packageName= */ "", user)); 169 if (storageStat == null) { 170 throw new PackageManager.NameNotFoundException( 171 "queryStatsForUser with non matching arguments. Did you forget to call" 172 + " addStorageStats?"); 173 } 174 return storageStat; 175 } 176 createFreeAndTotalBytesMapWithSingleEntry( UUID storageUuid, long freeBytes, long totalBytes)177 private static Map<UUID, FreeAndTotalBytesPair> createFreeAndTotalBytesMapWithSingleEntry( 178 UUID storageUuid, long freeBytes, long totalBytes) { 179 Map<UUID, FreeAndTotalBytesPair> currMap = new ConcurrentHashMap<>(); 180 currMap.put(storageUuid, FreeAndTotalBytesPair.create(freeBytes, totalBytes)); 181 return currMap; 182 } 183 184 /** Simple wrapper to combine freeBytes and totalBytes in one object. */ 185 @AutoValue 186 abstract static class FreeAndTotalBytesPair { 187 FreeAndTotalBytesPair()188 FreeAndTotalBytesPair() {} 189 190 /** Returns the freeBytes. */ freeBytes()191 abstract long freeBytes(); 192 193 /** Returns the totalBytes. */ totalBytes()194 abstract long totalBytes(); 195 196 /** Creates {@link FreeAndTotalBytesPair}. */ create(long freeBytes, long totalBytes)197 static FreeAndTotalBytesPair create(long freeBytes, long totalBytes) { 198 return new AutoValue_ShadowStorageStatsManager_FreeAndTotalBytesPair(freeBytes, totalBytes); 199 } 200 } 201 202 /** Simple wrapper for parameters of {@link StorageStatsManager#queryStatsForPackage} method. */ 203 @AutoValue 204 abstract static class StorageStatsKey { 205 StorageStatsKey()206 StorageStatsKey() {} 207 208 /** Returns the storage UUID part of this key. */ storageUuid()209 abstract UUID storageUuid(); 210 211 /** Returns the package name part of this key. */ packageName()212 abstract String packageName(); 213 214 /** Returns the user handle part of this key. */ userHandle()215 abstract UserHandle userHandle(); 216 217 /** Creates StorageStatsKey. */ create(UUID storageUuid, String packageName, UserHandle userHandle)218 static StorageStatsKey create(UUID storageUuid, String packageName, UserHandle userHandle) { 219 return new AutoValue_ShadowStorageStatsManager_StorageStatsKey( 220 storageUuid, packageName, userHandle); 221 } 222 } 223 } 224