2 files changed, 135 insertions, 20 deletions
diff --git a/tests/BlobStoreTestUtils/src/com/android/utils/blob/DummyBlobData.java b/tests/BlobStoreTestUtils/src/com/android/utils/blob/DummyBlobData.java
index 504bd1727682..b805744a8387 100644
--- a/tests/BlobStoreTestUtils/src/com/android/utils/blob/DummyBlobData.java
+++ b/tests/BlobStoreTestUtils/src/com/android/utils/blob/DummyBlobData.java
@@ -38,38 +38,75 @@ import java.util.concurrent.TimeUnit;
 public class DummyBlobData {
     private static final long DEFAULT_SIZE_BYTES = 10 * 1024L * 1024L;
 
-    private final Context mContext;
     private final Random mRandom;
     private final File mFile;
     private final long mFileSize;
-    private final String mLabel;
+    private final CharSequence mLabel;
 
     byte[] mFileDigest;
     long mExpiryTimeMs;
 
-    public DummyBlobData(Context context) {
-        this(context, new Random(0), "blob_" + System.nanoTime());
+    public DummyBlobData(Builder builder) {
+        mRandom = new Random(builder.getRandomSeed());
+        mFile = new File(builder.getContext().getFilesDir(), builder.getFileName());
+        mFileSize = builder.getFileSize();
+        mLabel = builder.getLabel();
     }
 
-    public DummyBlobData(Context context, long fileSize) {
-        this(context, fileSize, new Random(0), "blob_" + System.nanoTime(), "Test label");
-    }
+    public static class Builder {
+        private final Context mContext;
+        private int mRandomSeed = 0;
+        private long mFileSize = DEFAULT_SIZE_BYTES;
+        private CharSequence mLabel = "Test label";
+        private String mFileName = "blob_" + System.nanoTime();
 
-    public DummyBlobData(Context context, Random random, String fileName) {
-        this(context, DEFAULT_SIZE_BYTES, random, fileName, "Test label");
-    }
+        public Builder(Context context) {
+            mContext = context;
+        }
 
-    public DummyBlobData(Context context, Random random, String fileName, String label) {
-        this(context, DEFAULT_SIZE_BYTES, random, fileName, label);
-    }
+        public Context getContext() {
+            return mContext;
+        }
+
+        public Builder setRandomSeed(int randomSeed) {
+            mRandomSeed = randomSeed;
+            return this;
+        }
+
+        public int getRandomSeed() {
+            return mRandomSeed;
+        }
+
+        public Builder setFileSize(int fileSize) {
+            mFileSize = fileSize;
+            return this;
+        }
 
-    public DummyBlobData(Context context, long fileSize, Random random, String fileName,
-            String label) {
-        mContext = context;
-        mRandom = random;
-        mFile = new File(mContext.getFilesDir(), fileName);
-        mFileSize = fileSize;
-        mLabel = label;
+        public long getFileSize() {
+            return mFileSize;
+        }
+
+        public Builder setLabel(CharSequence label) {
+            mLabel = label;
+            return this;
+        }
+
+        public CharSequence getLabel() {
+            return mLabel;
+        }
+
+        public Builder setFileName(String fileName) {
+            mFileName = fileName;
+            return this;
+        }
+
+        public String getFileName() {
+            return mFileName;
+        }
+
+        public DummyBlobData build() {
+            return new DummyBlobData(this);
+        }
     }
 
     public void prepare() throws Exception {
diff --git a/tests/BlobStoreTestUtils/src/com/android/utils/blob/Utils.java b/tests/BlobStoreTestUtils/src/com/android/utils/blob/Utils.java
index c35385cd0429..654c1e21999d 100644
--- a/tests/BlobStoreTestUtils/src/com/android/utils/blob/Utils.java
+++ b/tests/BlobStoreTestUtils/src/com/android/utils/blob/Utils.java
@@ -16,7 +16,13 @@
 
 package com.android.utils.blob;
 
+import static com.google.common.truth.Truth.assertThat;
+
+import android.app.blob.BlobHandle;
 import android.app.blob.BlobStoreManager;
+import android.app.blob.LeaseInfo;
+import android.content.Context;
+import android.content.res.Resources;
 import android.os.ParcelFileDescriptor;
 
 import java.io.FileInputStream;
@@ -56,4 +62,76 @@ public class Utils {
             copy(in, out, lengthBytes);
         }
     }
+
+    public static void assertLeasedBlobs(BlobStoreManager blobStoreManager,
+            BlobHandle... expectedBlobHandles) throws IOException {
+        assertThat(blobStoreManager.getLeasedBlobs()).containsExactly(expectedBlobHandles);
+    }
+
+    public static void assertNoLeasedBlobs(BlobStoreManager blobStoreManager)
+            throws IOException {
+        assertThat(blobStoreManager.getLeasedBlobs()).isEmpty();
+    }
+
+    public static void acquireLease(Context context,
+            BlobHandle blobHandle, CharSequence description) throws IOException {
+        final BlobStoreManager blobStoreManager = (BlobStoreManager) context.getSystemService(
+                Context.BLOB_STORE_SERVICE);
+        blobStoreManager.acquireLease(blobHandle, description);
+
+        final LeaseInfo leaseInfo = blobStoreManager.getLeaseInfo(blobHandle);
+        assertLeaseInfo(leaseInfo, context.getPackageName(), 0,
+                Resources.ID_NULL, description);
+    }
+
+    public static void acquireLease(Context context,
+            BlobHandle blobHandle, int descriptionResId) throws IOException {
+        final BlobStoreManager blobStoreManager = (BlobStoreManager) context.getSystemService(
+                Context.BLOB_STORE_SERVICE);
+        blobStoreManager.acquireLease(blobHandle, descriptionResId);
+
+        final LeaseInfo leaseInfo = blobStoreManager.getLeaseInfo(blobHandle);
+        assertLeaseInfo(leaseInfo, context.getPackageName(), 0,
+                descriptionResId, context.getString(descriptionResId));
+    }
+
+    public static void acquireLease(Context context,
+            BlobHandle blobHandle, CharSequence description,
+            long expiryTimeMs) throws IOException {
+        final BlobStoreManager blobStoreManager = (BlobStoreManager) context.getSystemService(
+                Context.BLOB_STORE_SERVICE);
+        blobStoreManager.acquireLease(blobHandle, description, expiryTimeMs);
+
+        final LeaseInfo leaseInfo = blobStoreManager.getLeaseInfo(blobHandle);
+        assertLeaseInfo(leaseInfo, context.getPackageName(), expiryTimeMs,
+                Resources.ID_NULL, description);
+    }
+
+    public static void acquireLease(Context context,
+            BlobHandle blobHandle, int descriptionResId,
+            long expiryTimeMs) throws IOException {
+        final BlobStoreManager blobStoreManager = (BlobStoreManager) context.getSystemService(
+                Context.BLOB_STORE_SERVICE);
+        blobStoreManager.acquireLease(blobHandle, descriptionResId, expiryTimeMs);
+
+        final LeaseInfo leaseInfo = blobStoreManager.getLeaseInfo(blobHandle);
+        assertLeaseInfo(leaseInfo, context.getPackageName(), expiryTimeMs,
+                descriptionResId, context.getString(descriptionResId));
+    }
+
+    public static void releaseLease(Context context,
+            BlobHandle blobHandle) throws IOException {
+        final BlobStoreManager blobStoreManager = (BlobStoreManager) context.getSystemService(
+                Context.BLOB_STORE_SERVICE);
+        blobStoreManager.releaseLease(blobHandle);
+        assertThat(blobStoreManager.getLeaseInfo(blobHandle)).isNull();
+    }
+
+    private static void assertLeaseInfo(LeaseInfo leaseInfo, String packageName,
+            long expiryTimeMs, int descriptionResId, CharSequence description) {
+        assertThat(leaseInfo.getPackageName()).isEqualTo(packageName);
+        assertThat(leaseInfo.getExpiryTimeMillis()).isEqualTo(expiryTimeMs);
+        assertThat(leaseInfo.getDescriptionResId()).isEqualTo(descriptionResId);
+        assertThat(leaseInfo.getDescription()).isEqualTo(description);
+    }
 }