Load preivews eagerly

Currently, content preview is not started until the metadata for all the
shared URIs is loaded.
This CL changes it to eagerly loaded previews i.e. previews are getting
loaded as soon as the corresponding metadata becomes available.

This is achieved by:
* Adding support for a Flow as a preview source into
  ScrollableImgePreviewView (patchset #1);
  Specifically, as a Flow may never complete, the flow collection is
  cancelled in `onDetachFromWindow` to avoid possible memory leaks.
  The view is used inside a RecyclerView (the single-profile case) and
  can be attached and detached multiple times per one `setPreviews`
  call. Thus BatchPreviewLoader is made relaunchable and captures a
  notion of a pending loading job; `#batchLoader` value gets updated
  accordingly. (patchset #8)
* Make PreviewDataProvider expose a Flow of shared URIs metadata,
  FileInfo (patchset #2, #3);
* Make content preview classes to pass the Flow from PreviewDataProvider
  to ScrollableImagePreviewView (patchset #4).

Bug: 292157413
Test: manual testing with ShareTest app: with and without artificial
image loading delays, orientation changes and on single- and
multi-profile cases.
Test: unit tests
Test: integration test

Change-Id: Ib663bab8917624493a9ba619e64e4cb81fa35a93

6 files changed, 300 insertions, 154 deletions

diff --git a/java/tests/src/com/android/intentresolver/TestContentProvider.kt b/java/tests/src/com/android/intentresolver/TestContentProvider.kt
index b3b53baa..426f9af2 100644
--- a/java/tests/src/com/android/intentresolver/TestContentProvider.kt
+++ b/java/tests/src/com/android/intentresolver/TestContentProvider.kt
@@ -30,15 +30,23 @@ class TestContentProvider : ContentProvider() {
         sortOrder: String?
     ): Cursor? = null
 
-    override fun getType(uri: Uri): String?
-        = runCatching {
-            uri.getQueryParameter("mimeType")
-        }.getOrNull()
+    override fun getType(uri: Uri): String? =
+        runCatching { uri.getQueryParameter(PARAM_MIME_TYPE) }.getOrNull()
 
-    override fun getStreamTypes(uri: Uri, mimeTypeFilter: String): Array<String>?
-        = runCatching {
-            uri.getQueryParameter("streamType")?.let { arrayOf(it) }
-        }.getOrNull()
+    override fun getStreamTypes(uri: Uri, mimeTypeFilter: String): Array<String>? {
+        val delay =
+            runCatching { uri.getQueryParameter(PARAM_STREAM_TYPE_TIMEOUT)?.toLong() ?: 0L }
+                .getOrDefault(0L)
+        if (delay > 0) {
+            try {
+                Thread.sleep(delay)
+            } catch (e: InterruptedException) {
+                Thread.currentThread().interrupt()
+            }
+        }
+        return runCatching { uri.getQueryParameter(PARAM_STREAM_TYPE)?.let { arrayOf(it) } }
+            .getOrNull()
+    }
 
     override fun insert(uri: Uri, values: ContentValues?): Uri? = null
 
@@ -52,4 +60,10 @@ class TestContentProvider : ContentProvider() {
     ): Int = 0
 
     override fun onCreate(): Boolean = true
-}
\ No newline at end of file
+
+    companion object {
+        const val PARAM_MIME_TYPE = "mimeType"
+        const val PARAM_STREAM_TYPE = "streamType"
+        const val PARAM_STREAM_TYPE_TIMEOUT = "streamTypeTo"
+    }
+}
diff --git a/java/tests/src/com/android/intentresolver/UnbundledChooserActivityTest.java b/java/tests/src/com/android/intentresolver/UnbundledChooserActivityTest.java
index 28a45051..5709c912 100644
--- a/java/tests/src/com/android/intentresolver/UnbundledChooserActivityTest.java
+++ b/java/tests/src/com/android/intentresolver/UnbundledChooserActivityTest.java
@@ -136,6 +136,10 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -1042,6 +1046,63 @@ public class UnbundledChooserActivityTest {
     }
 
     @Test
+    public void testPartiallyLoadedMetadata_previewIsShownForTheLoadedPart()
+            throws InterruptedException {
+        Uri imgOneUri = createTestContentProviderUri("image/png", null);
+        Uri imgTwoUri = createTestContentProviderUri("image/png", null)
+                .buildUpon()
+                .path("image-2.png")
+                .build();
+        Uri docUri = createTestContentProviderUri("application/pdf", "image/png", 3_000);
+        ArrayList<Uri> uris = new ArrayList<>(2);
+        // two large previews to fill the screen and be presented right away and one
+        // document that would be delayed by the URI metadata reading
+        uris.add(imgOneUri);
+        uris.add(imgTwoUri);
+        uris.add(docUri);
+
+        Intent sendIntent = createSendUriIntentWithPreview(uris);
+        Map<Uri, Bitmap> bitmaps = new HashMap<>();
+        bitmaps.put(imgOneUri, createWideBitmap(Color.RED));
+        bitmaps.put(imgTwoUri, createWideBitmap(Color.GREEN));
+        bitmaps.put(docUri, createWideBitmap(Color.BLUE));
+        ChooserActivityOverrideData.getInstance().imageLoader =
+                new TestPreviewImageLoader(bitmaps);
+
+        List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
+        setupResolverControllers(resolvedComponentInfos);
+
+        assertThat(launchActivityWithTimeout(Intent.createChooser(sendIntent, null), 1_000))
+                .isTrue();
+        waitForIdle();
+
+        onView(withId(R.id.scrollable_image_preview))
+                .check((view, exception) -> {
+                    if (exception != null) {
+                        throw exception;
+                    }
+                    RecyclerView recyclerView = (RecyclerView) view;
+                    assertThat(recyclerView.getChildCount()).isAtLeast(1);
+                    // the first view is a preview
+                    View imageView = recyclerView.getChildAt(0).findViewById(R.id.image);
+                    assertThat(imageView).isNotNull();
+                })
+                .perform(RecyclerViewActions.scrollToLastPosition())
+                .check((view, exception) -> {
+                    if (exception != null) {
+                        throw exception;
+                    }
+                    RecyclerView recyclerView = (RecyclerView) view;
+                    assertThat(recyclerView.getChildCount()).isAtLeast(1);
+                    // check that the last view is a loading indicator
+                    View loadingIndicator =
+                            recyclerView.getChildAt(recyclerView.getChildCount() - 1);
+                    assertThat(loadingIndicator).isNotNull();
+                });
+        waitForIdle();
+    }
+
+    @Test
     public void testImageAndTextPreview() {
         final Uri uri = createTestContentProviderUri("image/png", null);
         final String sharedText = "text-" + System.currentTimeMillis();
@@ -2641,15 +2702,25 @@ public class UnbundledChooserActivityTest {
 
     private Uri createTestContentProviderUri(
             @Nullable String mimeType, @Nullable String streamType) {
+        return createTestContentProviderUri(mimeType, streamType, 0);
+    }
+
+    private Uri createTestContentProviderUri(
+            @Nullable String mimeType, @Nullable String streamType, long streamTypeTimeout) {
         String packageName =
                 InstrumentationRegistry.getInstrumentation().getContext().getPackageName();
         Uri.Builder builder = Uri.parse("content://" + packageName + "/image.png")
                 .buildUpon();
         if (mimeType != null) {
-            builder.appendQueryParameter("mimeType", mimeType);
+            builder.appendQueryParameter(TestContentProvider.PARAM_MIME_TYPE, mimeType);
         }
         if (streamType != null) {
-            builder.appendQueryParameter("streamType", streamType);
+            builder.appendQueryParameter(TestContentProvider.PARAM_STREAM_TYPE, streamType);
+        }
+        if (streamTypeTimeout > 0) {
+            builder.appendQueryParameter(
+                    TestContentProvider.PARAM_STREAM_TYPE_TIMEOUT,
+                    Long.toString(streamTypeTimeout));
         }
         return builder.build();
     }
@@ -2779,11 +2850,44 @@ public class UnbundledChooserActivityTest {
         InstrumentationRegistry.getInstrumentation().waitForIdleSync();
     }
 
+    private boolean launchActivityWithTimeout(Intent intent, long timeout)
+            throws InterruptedException {
+        final int initialState = 0;
+        final int completedState = 1;
+        final int timeoutState = 2;
+        final AtomicInteger state = new AtomicInteger(initialState);
+        final CountDownLatch cdl = new CountDownLatch(1);
+
+        ScheduledExecutorService executor = Executors.newScheduledThreadPool(2);
+        try {
+            executor.execute(() -> {
+                mActivityRule.launchActivity(intent);
+                state.compareAndSet(initialState, completedState);
+                cdl.countDown();
+            });
+            executor.schedule(
+                    () -> {
+                        state.compareAndSet(initialState, timeoutState);
+                        cdl.countDown();
+                    },
+                    timeout,
+                    TimeUnit.MILLISECONDS);
+            cdl.await();
+            return state.get() == completedState;
+        } finally {
+            executor.shutdownNow();
+        }
+    }
+
     private Bitmap createBitmap() {
         return createBitmap(200, 200);
     }
 
     private Bitmap createWideBitmap() {
+        return createWideBitmap(Color.RED);
+    }
+
+    private Bitmap createWideBitmap(int bgColor) {
         WindowManager windowManager = InstrumentationRegistry.getInstrumentation()
                 .getTargetContext()
                 .getSystemService(WindowManager.class);
@@ -2792,15 +2896,19 @@ public class UnbundledChooserActivityTest {
             Rect bounds = windowManager.getMaximumWindowMetrics().getBounds();
             width = bounds.width() + 200;
         }
-        return createBitmap(width, 100);
+        return createBitmap(width, 100, bgColor);
     }
 
     private Bitmap createBitmap(int width, int height) {
+        return createBitmap(width, height, Color.RED);
+    }
+
+    private Bitmap createBitmap(int width, int height, int bgColor) {
         Bitmap bitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888);
         Canvas canvas = new Canvas(bitmap);
 
         Paint paint = new Paint();
-        paint.setColor(Color.RED);
+        paint.setColor(bgColor);
         paint.setStyle(Paint.Style.FILL);
         canvas.drawPaint(paint);
 
diff --git a/java/tests/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUiTest.kt b/java/tests/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUiTest.kt
index 9bfd2052..008cc162 100644
--- a/java/tests/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUiTest.kt
+++ b/java/tests/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUiTest.kt
@@ -20,7 +20,7 @@ import android.content.Intent
 import android.graphics.Bitmap
 import android.net.Uri
 import androidx.lifecycle.Lifecycle
-import com.android.intentresolver.any
+import com.android.intentresolver.TestLifecycleOwner
 import com.android.intentresolver.contentpreview.ChooserContentPreviewUi.ActionFactory
 import com.android.intentresolver.mock
 import com.android.intentresolver.whenever
@@ -28,13 +28,14 @@ import com.android.intentresolver.widget.ActionRow
 import com.android.intentresolver.widget.ImagePreviewView
 import com.google.common.truth.Truth.assertThat
 import java.util.function.Consumer
+import kotlinx.coroutines.flow.MutableSharedFlow
 import org.junit.Test
 import org.mockito.Mockito.never
 import org.mockito.Mockito.times
 import org.mockito.Mockito.verify
 
 class ChooserContentPreviewUiTest {
-    private val lifecycle = mock<Lifecycle>()
+    private val lifecycleOwner = TestLifecycleOwner()
     private val previewData = mock<PreviewDataProvider>()
     private val headlineGenerator = mock<HeadlineGenerator>()
     private val imageLoader =
@@ -64,7 +65,7 @@ class ChooserContentPreviewUiTest {
         whenever(previewData.previewType).thenReturn(ContentPreviewType.CONTENT_PREVIEW_TEXT)
         val testSubject =
             ChooserContentPreviewUi(
-                lifecycle,
+                lifecycleOwner.lifecycle,
                 previewData,
                 Intent(Intent.ACTION_VIEW),
                 imageLoader,
@@ -83,7 +84,7 @@ class ChooserContentPreviewUiTest {
         whenever(previewData.previewType).thenReturn(ContentPreviewType.CONTENT_PREVIEW_FILE)
         val testSubject =
             ChooserContentPreviewUi(
-                lifecycle,
+                lifecycleOwner.lifecycle,
                 previewData,
                 Intent(Intent.ACTION_SEND),
                 imageLoader,
@@ -104,9 +105,10 @@ class ChooserContentPreviewUiTest {
         whenever(previewData.uriCount).thenReturn(2)
         whenever(previewData.firstFileInfo)
             .thenReturn(FileInfo.Builder(uri).withPreviewUri(uri).withMimeType("image/png").build())
+        whenever(previewData.imagePreviewFileInfoFlow).thenReturn(MutableSharedFlow())
         val testSubject =
             ChooserContentPreviewUi(
-                lifecycle,
+                lifecycleOwner.lifecycle,
                 previewData,
                 Intent(Intent.ACTION_SEND).apply { putExtra(Intent.EXTRA_TEXT, "Shared text") },
                 imageLoader,
@@ -116,7 +118,7 @@ class ChooserContentPreviewUiTest {
             )
         assertThat(testSubject.mContentPreviewUi)
             .isInstanceOf(FilesPlusTextContentPreviewUi::class.java)
-        verify(previewData, times(1)).getFileMetadataForImagePreview(any(), any())
+        verify(previewData, times(1)).imagePreviewFileInfoFlow
         verify(transitionCallback, times(1)).onAllTransitionElementsReady()
     }
 
@@ -127,9 +129,10 @@ class ChooserContentPreviewUiTest {
         whenever(previewData.uriCount).thenReturn(2)
         whenever(previewData.firstFileInfo)
             .thenReturn(FileInfo.Builder(uri).withPreviewUri(uri).withMimeType("image/png").build())
+        whenever(previewData.imagePreviewFileInfoFlow).thenReturn(MutableSharedFlow())
         val testSubject =
             ChooserContentPreviewUi(
-                lifecycle,
+                lifecycleOwner.lifecycle,
                 previewData,
                 Intent(Intent.ACTION_SEND),
                 imageLoader,
@@ -140,7 +143,7 @@ class ChooserContentPreviewUiTest {
         assertThat(testSubject.preferredContentPreview)
             .isEqualTo(ContentPreviewType.CONTENT_PREVIEW_IMAGE)
         assertThat(testSubject.mContentPreviewUi).isInstanceOf(UnifiedContentPreviewUi::class.java)
-        verify(previewData, times(1)).getFileMetadataForImagePreview(any(), any())
+        verify(previewData, times(1)).imagePreviewFileInfoFlow
         verify(transitionCallback, never()).onAllTransitionElementsReady()
     }
 }
diff --git a/java/tests/src/com/android/intentresolver/contentpreview/PreviewDataProviderTest.kt b/java/tests/src/com/android/intentresolver/contentpreview/PreviewDataProviderTest.kt
index 145b89ad..6599baa9 100644
--- a/java/tests/src/com/android/intentresolver/contentpreview/PreviewDataProviderTest.kt
+++ b/java/tests/src/com/android/intentresolver/contentpreview/PreviewDataProviderTest.kt
@@ -22,18 +22,15 @@ import android.database.MatrixCursor
 import android.media.MediaMetadata
 import android.net.Uri
 import android.provider.DocumentsContract
-import androidx.lifecycle.Lifecycle
-import com.android.intentresolver.TestLifecycleOwner
 import com.android.intentresolver.mock
 import com.android.intentresolver.whenever
 import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.Dispatchers
+import kotlin.coroutines.EmptyCoroutineContext
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.toList
+import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
-import kotlinx.coroutines.test.resetMain
-import kotlinx.coroutines.test.setMain
-import org.junit.After
-import org.junit.Before
+import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.mockito.Mockito.any
 import org.mockito.Mockito.never
@@ -44,27 +41,13 @@ import org.mockito.Mockito.verify
 class PreviewDataProviderTest {
     private val contentResolver = mock<ContentInterface>()
     private val mimeTypeClassifier = DefaultMimeTypeClassifier
-
-    private val lifecycleOwner = TestLifecycleOwner()
-    private val dispatcher = UnconfinedTestDispatcher()
-
-    @Before
-    fun setup() {
-        Dispatchers.setMain(dispatcher)
-        lifecycleOwner.state = Lifecycle.State.CREATED
-    }
-
-    @After
-    fun cleanup() {
-        lifecycleOwner.state = Lifecycle.State.DESTROYED
-        Dispatchers.resetMain()
-    }
+    private val testScope = TestScope(EmptyCoroutineContext + UnconfinedTestDispatcher())
 
     @Test
     fun test_nonSendIntentAction_resolvesToTextPreviewUiSynchronously() {
         val targetIntent = Intent(Intent.ACTION_VIEW)
         val testSubject =
-            PreviewDataProvider(targetIntent, contentResolver, mimeTypeClassifier, dispatcher)
+            PreviewDataProvider(testScope, targetIntent, contentResolver, mimeTypeClassifier)
 
         assertThat(testSubject.previewType).isEqualTo(ContentPreviewType.CONTENT_PREVIEW_TEXT)
         verify(contentResolver, never()).getType(any())
@@ -73,14 +56,14 @@ class PreviewDataProviderTest {
     @Test
     fun test_sendSingleTextFileWithoutPreview_resolvesToFilePreviewUi() {
         val uri = Uri.parse("content://org.pkg.app/notes.txt")
-        val targetIntent = Intent(Intent.ACTION_SEND)
-            .apply {
+        val targetIntent =
+            Intent(Intent.ACTION_SEND).apply {
                 putExtra(Intent.EXTRA_STREAM, uri)
                 type = "text/plain"
             }
         whenever(contentResolver.getType(uri)).thenReturn("text/plain")
         val testSubject =
-            PreviewDataProvider(targetIntent, contentResolver, mimeTypeClassifier, dispatcher)
+            PreviewDataProvider(testScope, targetIntent, contentResolver, mimeTypeClassifier)
 
         assertThat(testSubject.previewType).isEqualTo(ContentPreviewType.CONTENT_PREVIEW_FILE)
         assertThat(testSubject.uriCount).isEqualTo(1)
@@ -90,12 +73,9 @@ class PreviewDataProviderTest {
 
     @Test
     fun test_sendIntentWithoutUris_resolvesToTextPreviewUiSynchronously() {
-        val targetIntent = Intent(Intent.ACTION_SEND)
-            .apply {
-                type = "image/png"
-            }
+        val targetIntent = Intent(Intent.ACTION_SEND).apply { type = "image/png" }
         val testSubject =
-            PreviewDataProvider(targetIntent, contentResolver, mimeTypeClassifier, dispatcher)
+            PreviewDataProvider(testScope, targetIntent, contentResolver, mimeTypeClassifier)
 
         assertThat(testSubject.previewType).isEqualTo(ContentPreviewType.CONTENT_PREVIEW_TEXT)
         verify(contentResolver, never()).getType(any())
@@ -104,13 +84,10 @@ class PreviewDataProviderTest {
     @Test
     fun test_sendSingleImage_resolvesToImagePreviewUi() {
         val uri = Uri.parse("content://org.pkg.app/image.png")
-        val targetIntent = Intent(Intent.ACTION_SEND)
-            .apply {
-                putExtra(Intent.EXTRA_STREAM, uri)
-            }
+        val targetIntent = Intent(Intent.ACTION_SEND).apply { putExtra(Intent.EXTRA_STREAM, uri) }
         whenever(contentResolver.getType(uri)).thenReturn("image/png")
         val testSubject =
-            PreviewDataProvider(targetIntent, contentResolver, mimeTypeClassifier, dispatcher)
+            PreviewDataProvider(testScope, targetIntent, contentResolver, mimeTypeClassifier)
 
         assertThat(testSubject.previewType).isEqualTo(ContentPreviewType.CONTENT_PREVIEW_IMAGE)
         assertThat(testSubject.uriCount).isEqualTo(1)
@@ -122,13 +99,10 @@ class PreviewDataProviderTest {
     @Test
     fun test_sendSingleNonImage_resolvesToFilePreviewUi() {
         val uri = Uri.parse("content://org.pkg.app/paper.pdf")
-        val targetIntent = Intent(Intent.ACTION_SEND)
-            .apply {
-                putExtra(Intent.EXTRA_STREAM, uri)
-            }
+        val targetIntent = Intent(Intent.ACTION_SEND).apply { putExtra(Intent.EXTRA_STREAM, uri) }
         whenever(contentResolver.getType(uri)).thenReturn("application/pdf")
         val testSubject =
-            PreviewDataProvider(targetIntent, contentResolver, mimeTypeClassifier, dispatcher)
+            PreviewDataProvider(testScope, targetIntent, contentResolver, mimeTypeClassifier)
 
         assertThat(testSubject.previewType).isEqualTo(ContentPreviewType.CONTENT_PREVIEW_FILE)
         assertThat(testSubject.uriCount).isEqualTo(1)
@@ -141,14 +115,13 @@ class PreviewDataProviderTest {
     fun test_sendSingleImageWithFailingGetType_resolvesToFilePreviewUi() {
         val uri = Uri.parse("content://org.pkg.app/image.png")
         val targetIntent =
-            Intent(Intent.ACTION_SEND)
-                .apply {
-                    type = "image/png"
-                    putExtra(Intent.EXTRA_STREAM, uri)
-                }
+            Intent(Intent.ACTION_SEND).apply {
+                type = "image/png"
+                putExtra(Intent.EXTRA_STREAM, uri)
+            }
         whenever(contentResolver.getType(uri)).thenThrow(SecurityException("test failure"))
         val testSubject =
-            PreviewDataProvider(targetIntent, contentResolver, mimeTypeClassifier, dispatcher)
+            PreviewDataProvider(testScope, targetIntent, contentResolver, mimeTypeClassifier)
 
         assertThat(testSubject.previewType).isEqualTo(ContentPreviewType.CONTENT_PREVIEW_FILE)
         assertThat(testSubject.uriCount).isEqualTo(1)
@@ -161,17 +134,16 @@ class PreviewDataProviderTest {
     fun test_sendSingleImageWithFailingMetadata_resolvesToFilePreviewUi() {
         val uri = Uri.parse("content://org.pkg.app/image.png")
         val targetIntent =
-            Intent(Intent.ACTION_SEND)
-                .apply {
-                    type = "image/png"
-                    putExtra(Intent.EXTRA_STREAM, uri)
-                }
+            Intent(Intent.ACTION_SEND).apply {
+                type = "image/png"
+                putExtra(Intent.EXTRA_STREAM, uri)
+            }
         whenever(contentResolver.getStreamTypes(uri, "*/*"))
             .thenThrow(SecurityException("test failure"))
         whenever(contentResolver.query(uri, METADATA_COLUMNS, null, null))
             .thenThrow(SecurityException("test failure"))
         val testSubject =
-            PreviewDataProvider(targetIntent, contentResolver, mimeTypeClassifier, dispatcher)
+            PreviewDataProvider(testScope, targetIntent, contentResolver, mimeTypeClassifier)
 
         assertThat(testSubject.previewType).isEqualTo(ContentPreviewType.CONTENT_PREVIEW_FILE)
         assertThat(testSubject.uriCount).isEqualTo(1)
@@ -183,14 +155,11 @@ class PreviewDataProviderTest {
     @Test
     fun test_SingleNonImageUriWithImageTypeInGetStreamTypes_useImagePreviewUi() {
         val uri = Uri.parse("content://org.pkg.app/paper.pdf")
-        val targetIntent = Intent(Intent.ACTION_SEND)
-            .apply {
-                putExtra(Intent.EXTRA_STREAM, uri)
-            }
+        val targetIntent = Intent(Intent.ACTION_SEND).apply { putExtra(Intent.EXTRA_STREAM, uri) }
         whenever(contentResolver.getStreamTypes(uri, "*/*"))
             .thenReturn(arrayOf("application/pdf", "image/png"))
         val testSubject =
-            PreviewDataProvider(targetIntent, contentResolver, mimeTypeClassifier, dispatcher)
+            PreviewDataProvider(testScope, targetIntent, contentResolver, mimeTypeClassifier)
 
         assertThat(testSubject.previewType).isEqualTo(ContentPreviewType.CONTENT_PREVIEW_IMAGE)
         assertThat(testSubject.uriCount).isEqualTo(1)
@@ -221,15 +190,12 @@ class PreviewDataProviderTest {
 
     private fun testMetadataToImagePreview(columns: Array<String>, values: Array<Any>) {
         val uri = Uri.parse("content://org.pkg.app/test.pdf")
-        val targetIntent = Intent(Intent.ACTION_SEND)
-            .apply {
-                putExtra(Intent.EXTRA_STREAM, uri)
-            }
+        val targetIntent = Intent(Intent.ACTION_SEND).apply { putExtra(Intent.EXTRA_STREAM, uri) }
         whenever(contentResolver.getType(uri)).thenReturn("application/pdf")
         whenever(contentResolver.query(uri, METADATA_COLUMNS, null, null))
             .thenReturn(MatrixCursor(columns).apply { addRow(values) })
         val testSubject =
-            PreviewDataProvider(targetIntent, contentResolver, mimeTypeClassifier, dispatcher)
+            PreviewDataProvider(testScope, targetIntent, contentResolver, mimeTypeClassifier)
 
         assertThat(testSubject.previewType).isEqualTo(ContentPreviewType.CONTENT_PREVIEW_IMAGE)
         assertThat(testSubject.uriCount).isEqualTo(1)
@@ -243,20 +209,19 @@ class PreviewDataProviderTest {
         val uri1 = Uri.parse("content://org.pkg.app/test.png")
         val uri2 = Uri.parse("content://org.pkg.app/test.jpg")
         val targetIntent =
-            Intent(Intent.ACTION_SEND_MULTIPLE)
-                .apply {
-                    putExtra(
-                        Intent.EXTRA_STREAM,
-                        ArrayList<Uri>().apply {
-                            add(uri1)
-                            add(uri2)
-                        }
-                    )
-                }
+            Intent(Intent.ACTION_SEND_MULTIPLE).apply {
+                putExtra(
+                    Intent.EXTRA_STREAM,
+                    ArrayList<Uri>().apply {
+                        add(uri1)
+                        add(uri2)
+                    }
+                )
+            }
         whenever(contentResolver.getType(uri1)).thenReturn("image/png")
         whenever(contentResolver.getType(uri2)).thenReturn("image/jpeg")
         val testSubject =
-            PreviewDataProvider(targetIntent, contentResolver, mimeTypeClassifier, dispatcher)
+            PreviewDataProvider(testScope, targetIntent, contentResolver, mimeTypeClassifier)
 
         assertThat(testSubject.previewType).isEqualTo(ContentPreviewType.CONTENT_PREVIEW_IMAGE)
         assertThat(testSubject.uriCount).isEqualTo(2)
@@ -273,18 +238,17 @@ class PreviewDataProviderTest {
         whenever(contentResolver.getType(uri1)).thenReturn("image/png")
         whenever(contentResolver.getType(uri2)).thenReturn("application/pdf")
         val targetIntent =
-            Intent(Intent.ACTION_SEND_MULTIPLE)
-                .apply {
-                    putExtra(
-                        Intent.EXTRA_STREAM,
-                        ArrayList<Uri>().apply {
-                            add(uri1)
-                            add(uri2)
-                        }
-                    )
-                }
+            Intent(Intent.ACTION_SEND_MULTIPLE).apply {
+                putExtra(
+                    Intent.EXTRA_STREAM,
+                    ArrayList<Uri>().apply {
+                        add(uri1)
+                        add(uri2)
+                    }
+                )
+            }
         val testSubject =
-            PreviewDataProvider(targetIntent, contentResolver, mimeTypeClassifier, dispatcher)
+            PreviewDataProvider(testScope, targetIntent, contentResolver, mimeTypeClassifier)
 
         assertThat(testSubject.previewType).isEqualTo(ContentPreviewType.CONTENT_PREVIEW_IMAGE)
         assertThat(testSubject.uriCount).isEqualTo(2)
@@ -299,21 +263,20 @@ class PreviewDataProviderTest {
         val uri1 = Uri.parse("content://org.pkg.app/test.mp4")
         val uri2 = Uri.parse("content://org.pkg.app/test.pdf")
         val targetIntent =
-            Intent(Intent.ACTION_SEND_MULTIPLE)
-                .apply {
-                    putExtra(
-                        Intent.EXTRA_STREAM,
-                        ArrayList<Uri>().apply {
-                            add(uri1)
-                            add(uri2)
-                        }
-                    )
-                }
+            Intent(Intent.ACTION_SEND_MULTIPLE).apply {
+                putExtra(
+                    Intent.EXTRA_STREAM,
+                    ArrayList<Uri>().apply {
+                        add(uri1)
+                        add(uri2)
+                    }
+                )
+            }
         whenever(contentResolver.getType(uri1)).thenReturn("video/mpeg4")
         whenever(contentResolver.getStreamTypes(uri1, "*/*")).thenReturn(arrayOf("image/png"))
         whenever(contentResolver.getType(uri2)).thenReturn("application/pdf")
         val testSubject =
-            PreviewDataProvider(targetIntent, contentResolver, mimeTypeClassifier, dispatcher)
+            PreviewDataProvider(testScope, targetIntent, contentResolver, mimeTypeClassifier)
 
         assertThat(testSubject.previewType).isEqualTo(ContentPreviewType.CONTENT_PREVIEW_IMAGE)
         assertThat(testSubject.uriCount).isEqualTo(2)
@@ -327,20 +290,19 @@ class PreviewDataProviderTest {
         val uri1 = Uri.parse("content://org.pkg.app/test.html")
         val uri2 = Uri.parse("content://org.pkg.app/test.pdf")
         val targetIntent =
-            Intent(Intent.ACTION_SEND_MULTIPLE)
-                .apply {
-                    putExtra(
-                        Intent.EXTRA_STREAM,
-                        ArrayList<Uri>().apply {
-                            add(uri1)
-                            add(uri2)
-                        }
-                    )
-                }
+            Intent(Intent.ACTION_SEND_MULTIPLE).apply {
+                putExtra(
+                    Intent.EXTRA_STREAM,
+                    ArrayList<Uri>().apply {
+                        add(uri1)
+                        add(uri2)
+                    }
+                )
+            }
         whenever(contentResolver.getType(uri1)).thenReturn("text/html")
         whenever(contentResolver.getType(uri2)).thenReturn("application/pdf")
         val testSubject =
-            PreviewDataProvider(targetIntent, contentResolver, mimeTypeClassifier, dispatcher)
+            PreviewDataProvider(testScope, targetIntent, contentResolver, mimeTypeClassifier)
 
         assertThat(testSubject.previewType).isEqualTo(ContentPreviewType.CONTENT_PREVIEW_FILE)
         assertThat(testSubject.uriCount).isEqualTo(2)
@@ -348,4 +310,40 @@ class PreviewDataProviderTest {
         assertThat(testSubject.firstFileInfo?.previewUri).isNull()
         verify(contentResolver, times(2)).getType(any())
     }
+
+    @Test
+    fun test_imagePreviewFileInfoFlow_dataLoadedOnce() =
+        testScope.runTest {
+            val uri1 = Uri.parse("content://org.pkg.app/test.html")
+            val uri2 = Uri.parse("content://org.pkg.app/test.pdf")
+            val targetIntent =
+                Intent(Intent.ACTION_SEND_MULTIPLE).apply {
+                    putExtra(
+                        Intent.EXTRA_STREAM,
+                        ArrayList<Uri>().apply {
+                            add(uri1)
+                            add(uri2)
+                        }
+                    )
+                }
+            whenever(contentResolver.getType(uri1)).thenReturn("text/html")
+            whenever(contentResolver.getType(uri2)).thenReturn("application/pdf")
+            whenever(contentResolver.getStreamTypes(uri1, "*/*"))
+                .thenReturn(arrayOf("text/html", "image/jpeg"))
+            whenever(contentResolver.getStreamTypes(uri2, "*/*"))
+                .thenReturn(arrayOf("application/pdf", "image/png"))
+            val testSubject =
+                PreviewDataProvider(testScope, targetIntent, contentResolver, mimeTypeClassifier)
+
+            val fileInfoListOne = testSubject.imagePreviewFileInfoFlow.toList()
+            val fileInfoListTwo = testSubject.imagePreviewFileInfoFlow.toList()
+
+            assertThat(fileInfoListOne).hasSize(2)
+            assertThat(fileInfoListOne).containsAtLeastElementsIn(fileInfoListTwo).inOrder()
+
+            verify(contentResolver, times(1)).getType(uri1)
+            verify(contentResolver, times(1)).getStreamTypes(uri1, "*/*")
+            verify(contentResolver, times(1)).getType(uri2)
+            verify(contentResolver, times(1)).getStreamTypes(uri2, "*/*")
+        }
 }
diff --git a/java/tests/src/com/android/intentresolver/contentpreview/UnifiedContentPreviewUiTest.kt b/java/tests/src/com/android/intentresolver/contentpreview/UnifiedContentPreviewUiTest.kt
index 08331209..e7de0b7b 100644
--- a/java/tests/src/com/android/intentresolver/contentpreview/UnifiedContentPreviewUiTest.kt
+++ b/java/tests/src/com/android/intentresolver/contentpreview/UnifiedContentPreviewUiTest.kt
@@ -25,6 +25,13 @@ import com.android.intentresolver.R.layout.chooser_grid
 import com.android.intentresolver.mock
 import com.android.intentresolver.whenever
 import com.android.intentresolver.widget.ImagePreviewView.TransitionElementStatusCallback
+import kotlin.coroutines.EmptyCoroutineContext
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.asFlow
+import kotlinx.coroutines.flow.takeWhile
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mockito.anyInt
@@ -33,6 +40,7 @@ import org.mockito.Mockito.verify
 
 @RunWith(AndroidJUnit4::class)
 class UnifiedContentPreviewUiTest {
+    private val testScope = TestScope(EmptyCoroutineContext + UnconfinedTestDispatcher())
     private val actionFactory =
         mock<ChooserContentPreviewUi.ActionFactory> {
             whenever(createCustomActions()).thenReturn(emptyList())
@@ -129,24 +137,30 @@ class UnifiedContentPreviewUiTest {
     }
 
     private fun testLoadingHeadline(intentMimeType: String, files: List<FileInfo>?) {
-        val testSubject =
-            UnifiedContentPreviewUi(
-                /*isSingleImage=*/ false,
-                intentMimeType,
-                actionFactory,
-                imageLoader,
-                DefaultMimeTypeClassifier,
-                object : TransitionElementStatusCallback {
-                    override fun onTransitionElementReady(name: String) = Unit
-                    override fun onAllTransitionElementsReady() = Unit
-                },
-                /*itemCount=*/ 2,
-                headlineGenerator
-            )
-        val layoutInflater = LayoutInflater.from(context)
-        val gridLayout = layoutInflater.inflate(chooser_grid, null, false) as ViewGroup
-
-        files?.let(testSubject::setFiles)
-        testSubject.display(context.resources, LayoutInflater.from(context), gridLayout)
+        testScope.runTest {
+            val endMarker = FileInfo.Builder(Uri.EMPTY).build()
+            val emptySourceFlow = MutableSharedFlow<FileInfo>(replay = 1)
+            val testSubject =
+                UnifiedContentPreviewUi(
+                    testScope,
+                    /*isSingleImage=*/ false,
+                    intentMimeType,
+                    actionFactory,
+                    imageLoader,
+                    DefaultMimeTypeClassifier,
+                    object : TransitionElementStatusCallback {
+                        override fun onTransitionElementReady(name: String) = Unit
+                        override fun onAllTransitionElementsReady() = Unit
+                    },
+                    files?.let { it.asFlow() } ?: emptySourceFlow.takeWhile { it !== endMarker },
+                    /*itemCount=*/ 2,
+                    headlineGenerator
+                )
+            val layoutInflater = LayoutInflater.from(context)
+            val gridLayout = layoutInflater.inflate(chooser_grid, null, false) as ViewGroup
+
+            testSubject.display(context.resources, LayoutInflater.from(context), gridLayout)
+            emptySourceFlow.tryEmit(endMarker)
+        }
     }
 }
diff --git a/java/tests/src/com/android/intentresolver/widget/BatchPreviewLoaderTest.kt b/java/tests/src/com/android/intentresolver/widget/BatchPreviewLoaderTest.kt
index a0211308..4f4223c0 100644
--- a/java/tests/src/com/android/intentresolver/widget/BatchPreviewLoaderTest.kt
+++ b/java/tests/src/com/android/intentresolver/widget/BatchPreviewLoaderTest.kt
@@ -31,6 +31,7 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.asFlow
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.resetMain
@@ -67,7 +68,13 @@ class BatchPreviewLoaderTest {
         val uriTwo = createUri(2)
         imageLoader.setUriLoadingOrder(succeed(uriTwo), succeed(uriOne))
         val testSubject =
-            BatchPreviewLoader(imageLoader, previews(uriOne, uriTwo), 0, onUpdate, onCompletion)
+            BatchPreviewLoader(
+                imageLoader,
+                previews(uriOne, uriTwo),
+                totalItemCount = 2,
+                onUpdate,
+                onCompletion
+            )
         testSubject.loadAspectRatios(200) { _, _, _ -> 100 }
         dispatcher.scheduler.advanceUntilIdle()
 
@@ -87,7 +94,7 @@ class BatchPreviewLoaderTest {
             BatchPreviewLoader(
                 imageLoader,
                 previews(uriOne, uriTwo, uriThree),
-                0,
+                totalItemCount = 3,
                 onUpdate,
                 onCompletion
             )
@@ -115,7 +122,7 @@ class BatchPreviewLoaderTest {
             }
         imageLoader.setUriLoadingOrder(*loadingOrder)
         val testSubject =
-            BatchPreviewLoader(imageLoader, previews(*uris), 0, onUpdate, onCompletion)
+            BatchPreviewLoader(imageLoader, previews(*uris), uris.size, onUpdate, onCompletion)
         testSubject.loadAspectRatios(200) { _, _, _ -> 100 }
         dispatcher.scheduler.advanceUntilIdle()
 
@@ -144,7 +151,7 @@ class BatchPreviewLoaderTest {
         val expectedUris = Array(uris.size / 2) { createUri(it * 2 + 1) }
         imageLoader.setUriLoadingOrder(*loadingOrder)
         val testSubject =
-            BatchPreviewLoader(imageLoader, previews(*uris), 0, onUpdate, onCompletion)
+            BatchPreviewLoader(imageLoader, previews(*uris), uris.size, onUpdate, onCompletion)
         testSubject.loadAspectRatios(200) { _, _, _ -> 100 }
         dispatcher.scheduler.advanceUntilIdle()
 
@@ -161,9 +168,11 @@ class BatchPreviewLoaderTest {
     private fun fail(uri: Uri) = uri to false
     private fun succeed(uri: Uri) = uri to true
     private fun previews(vararg uris: Uri) =
-        uris.fold(ArrayList<Preview>(uris.size)) { acc, uri ->
-            acc.apply { add(Preview(PreviewType.Image, uri, editAction = null)) }
-        }
+        uris
+            .fold(ArrayList<Preview>(uris.size)) { acc, uri ->
+                acc.apply { add(Preview(PreviewType.Image, uri, editAction = null)) }
+            }
+            .asFlow()
 }
 
 private class TestImageLoader(scope: CoroutineScope) : suspend (Uri, Boolean) -> Bitmap? {