1 files changed, 16 insertions, 22 deletions

diff --git a/tests/PhotoPicker/src/android/photopicker/cts/util/PhotoPickerAssertionsUtils.java b/tests/PhotoPicker/src/android/photopicker/cts/util/PhotoPickerAssertionsUtils.java
index e0385ca89c9..d9ff84bed99 100644
--- a/tests/PhotoPicker/src/android/photopicker/cts/util/PhotoPickerAssertionsUtils.java
+++ b/tests/PhotoPicker/src/android/photopicker/cts/util/PhotoPickerAssertionsUtils.java
@@ -17,9 +17,7 @@
 package android.photopicker.cts.util;
 
 import static android.os.SystemProperties.getBoolean;
-import static android.provider.MediaStore.Files.FileColumns.MEDIA_TYPE_IMAGE;
-import static android.provider.MediaStore.Files.FileColumns.MEDIA_TYPE_VIDEO;
-import static android.provider.MediaStore.Files.FileColumns.MIME_TYPE;
+import static android.provider.MediaStore.Files.FileColumns;
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
@@ -65,16 +63,22 @@ public class PhotoPickerAssertionsUtils {
     public static void assertRedactedReadOnlyAccess(Uri uri) throws Exception {
         assertThat(uri).isNotNull();
         // TODO(b/205291616): Replace FileColumns.MIME_TYPE with PickerMediaColumns.MIME_TYPE
-        final String[] projection = new String[]{ MIME_TYPE };
+        final String[] projection = new String[]{ FileColumns.MIME_TYPE };
         final Context context = InstrumentationRegistry.getTargetContext();
         final ContentResolver resolver = context.getContentResolver();
-        final Cursor c = resolver.query(uri, projection, null, null);
-        assertThat(c).isNotNull();
-        assertThat(c.moveToFirst()).isTrue();
+        try (Cursor c = resolver.query(uri, projection, null, null)) {
+            assertThat(c).isNotNull();
+            assertThat(c.moveToFirst()).isTrue();
+
+            final String mimeType;
+            if (getBoolean("sys.photopicker.pickerdb.enabled", true)) {
+                // TODO(b/205291616): Replace FileColumns.MIME_TYPE with
+                // PickerMediaColumns.MIME_TYPE
+                mimeType = c.getString(c.getColumnIndex(FileColumns.MIME_TYPE));
+            } else {
+                mimeType = c.getString(c.getColumnIndex(FileColumns.MIME_TYPE));
+            }
 
-        if (getBoolean("sys.photopicker.pickerdb.enabled", true)) {
-            // TODO(b/205291616): Replace FileColumns.MIME_TYPE with PickerMediaColumns.MIME_TYPE
-            final String mimeType = c.getString(c.getColumnIndex(MIME_TYPE));
             if (mimeType.startsWith("image")) {
                 assertImageRedactedReadOnlyAccess(uri, resolver);
             } else if (mimeType.startsWith("video")) {
@@ -82,24 +86,14 @@ public class PhotoPickerAssertionsUtils {
             } else {
                 fail("The mime type is not as expected: " + mimeType);
             }
-        } else {
-            final int mediaType = c.getInt(1);
-            switch (mediaType) {
-                case MEDIA_TYPE_IMAGE:
-                    assertImageRedactedReadOnlyAccess(uri, resolver);
-                    break;
-                case MEDIA_TYPE_VIDEO:
-                    assertVideoRedactedReadOnlyAccess(uri, resolver);
-                    break;
-                default:
-                    fail("The media type is not as expected: " + mediaType);
-            }
         }
     }
 
     private static void assertVideoRedactedReadOnlyAccess(Uri uri, ContentResolver resolver)
             throws Exception {
         // The location is redacted
+        // TODO(b/201505595): Make this method work for test_video.mp4. Currently it works only for
+        //  test_video_dng.mp4
         try (InputStream in = resolver.openInputStream(uri);
                 ByteArrayOutputStream out = new ByteArrayOutputStream()) {
             FileUtils.copy(in, out);