Merge "Retain the image loader through configuration change" into udc-dev am: 6df7212d10

Original change: https://googleplex-android-review.googlesource.com/c/platform/packages/modules/IntentResolver/+/23192997

Change-Id: Ic4af9d17b4777d0449680a8ad11eb7d799cf73fa
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

13 files changed, 257 insertions, 137 deletions

diff --git a/java/src/com/android/intentresolver/ChooserActivity.java b/java/src/com/android/intentresolver/ChooserActivity.java
index 0786e088..014aa2a2 100644
--- a/java/src/com/android/intentresolver/ChooserActivity.java
+++ b/java/src/com/android/intentresolver/ChooserActivity.java
@@ -83,9 +83,9 @@ import com.android.intentresolver.NoCrossProfileEmptyStateProvider.DevicePolicyB
 import com.android.intentresolver.chooser.DisplayResolveInfo;
 import com.android.intentresolver.chooser.MultiDisplayResolveInfo;
 import com.android.intentresolver.chooser.TargetInfo;
+import com.android.intentresolver.contentpreview.BasePreviewViewModel;
 import com.android.intentresolver.contentpreview.ChooserContentPreviewUi;
 import com.android.intentresolver.contentpreview.HeadlineGeneratorImpl;
-import com.android.intentresolver.contentpreview.ImageLoader;
 import com.android.intentresolver.contentpreview.PreviewViewModel;
 import com.android.intentresolver.flags.FeatureFlagRepository;
 import com.android.intentresolver.flags.FeatureFlagRepositoryFactory;
@@ -272,13 +272,14 @@ public class ChooserActivity extends ResolverActivity implements
             }
         });
 
+        BasePreviewViewModel previewViewModel =
+                new ViewModelProvider(this, createPreviewViewModelFactory())
+                        .get(BasePreviewViewModel.class);
         mChooserContentPreviewUi = new ChooserContentPreviewUi(
                 getLifecycle(),
-                new ViewModelProvider(this, PreviewViewModel.Companion.getFactory())
-                        .get(PreviewViewModel.class)
-                        .createOrReuseProvider(mChooserRequest),
+                previewViewModel.createOrReuseProvider(mChooserRequest),
                 mChooserRequest.getTargetIntent(),
-                createPreviewImageLoader(),
+                previewViewModel.createOrReuseImageLoader(),
                 createChooserActionFactory(),
                 mEnterTransitionAnimationDelegate,
                 new HeadlineGeneratorImpl(this));
@@ -1314,14 +1315,8 @@ public class ChooserActivity extends ResolverActivity implements
     }
 
     @VisibleForTesting
-    protected ImageLoader createPreviewImageLoader() {
-        final int cacheSize;
-        float chooserWidth = getResources().getDimension(R.dimen.chooser_width);
-        // imageWidth = imagePreviewHeight * minAspectRatio (see ScrollableImagePreviewView)
-        float imageWidth =
-                getResources().getDimension(R.dimen.chooser_preview_image_height_tall) * 2 / 5;
-        cacheSize = (int) (Math.ceil(chooserWidth / imageWidth) + 2);
-        return new ImagePreviewImageLoader(this, getLifecycle(), cacheSize);
+    protected ViewModelProvider.Factory createPreviewViewModelFactory() {
+        return PreviewViewModel.Companion.getFactory();
     }
 
     private ChooserActionFactory createChooserActionFactory() {
diff --git a/java/src/com/android/intentresolver/contentpreview/BasePreviewViewModel.kt b/java/src/com/android/intentresolver/contentpreview/BasePreviewViewModel.kt
new file mode 100644
index 00000000..103e8bf4
--- /dev/null
+++ b/java/src/com/android/intentresolver/contentpreview/BasePreviewViewModel.kt
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.intentresolver.contentpreview
+
+import androidx.annotation.MainThread
+import androidx.lifecycle.ViewModel
+import com.android.intentresolver.ChooserRequestParameters
+
+/** A contract for the preview view model. Added for testing. */
+abstract class BasePreviewViewModel : ViewModel() {
+    @MainThread
+    abstract fun createOrReuseProvider(
+        chooserRequest: ChooserRequestParameters
+    ): PreviewDataProvider
+
+    @MainThread abstract fun createOrReuseImageLoader(): ImageLoader
+}
diff --git a/java/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUi.java b/java/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUi.java
index 787af95f..9100b392 100644
--- a/java/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUi.java
+++ b/java/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUi.java
@@ -118,6 +118,7 @@ public final class ChooserContentPreviewUi {
         int previewType = previewData.getPreviewType();
         if (previewType == CONTENT_PREVIEW_TEXT) {
             return createTextPreview(
+                    mLifecycle,
                     targetIntent,
                     actionFactory,
                     imageLoader,
@@ -140,6 +141,7 @@ public final class ChooserContentPreviewUi {
         if (!TextUtils.isEmpty(text)) {
             FilesPlusTextContentPreviewUi previewUi =
                     new FilesPlusTextContentPreviewUi(
+                            mLifecycle,
                             isSingleImageShare,
                             previewData.getUriCount(),
                             targetIntent.getCharSequenceExtra(Intent.EXTRA_TEXT),
@@ -180,6 +182,7 @@ public final class ChooserContentPreviewUi {
     }
 
     private static TextContentPreviewUi createTextPreview(
+            Lifecycle lifecycle,
             Intent targetIntent,
             ChooserContentPreviewUi.ActionFactory actionFactory,
             ImageLoader imageLoader,
@@ -195,6 +198,7 @@ public final class ChooserContentPreviewUi {
             }
         }
         return new TextContentPreviewUi(
+                lifecycle,
                 sharingText,
                 previewTitle,
                 previewThumbnail,
diff --git a/java/src/com/android/intentresolver/contentpreview/FilesPlusTextContentPreviewUi.java b/java/src/com/android/intentresolver/contentpreview/FilesPlusTextContentPreviewUi.java
index 860423c4..e4e33839 100644
--- a/java/src/com/android/intentresolver/contentpreview/FilesPlusTextContentPreviewUi.java
+++ b/java/src/com/android/intentresolver/contentpreview/FilesPlusTextContentPreviewUi.java
@@ -31,6 +31,7 @@ import android.widget.ImageView;
 import android.widget.TextView;
 
 import androidx.annotation.Nullable;
+import androidx.lifecycle.Lifecycle;
 
 import com.android.intentresolver.R;
 import com.android.intentresolver.widget.ActionRow;
@@ -48,6 +49,7 @@ import java.util.function.Consumer;
  * file content).
  */
 class FilesPlusTextContentPreviewUi extends ContentPreviewUi {
+    private final Lifecycle mLifecycle;
     private final CharSequence mText;
     private final ChooserContentPreviewUi.ActionFactory mActionFactory;
     private final ImageLoader mImageLoader;
@@ -63,6 +65,7 @@ class FilesPlusTextContentPreviewUi extends ContentPreviewUi {
     private boolean mAllVideos;
 
     FilesPlusTextContentPreviewUi(
+            Lifecycle lifecycle,
             boolean isSingleImage,
             int fileCount,
             CharSequence text,
@@ -70,6 +73,7 @@ class FilesPlusTextContentPreviewUi extends ContentPreviewUi {
             ImageLoader imageLoader,
             MimeTypeClassifier typeClassifier,
             HeadlineGenerator headlineGenerator) {
+        mLifecycle = lifecycle;
         if (isSingleImage && fileCount != 1) {
             throw new IllegalArgumentException(
                     "fileCount = " + fileCount + " and isSingleImage = true");
@@ -155,13 +159,16 @@ class FilesPlusTextContentPreviewUi extends ContentPreviewUi {
 
         ImageView imagePreview = mContentPreviewView.requireViewById(R.id.image_view);
         if (mIsSingleImage && mFirstFilePreviewUri != null) {
-            mImageLoader.loadImage(mFirstFilePreviewUri, bitmap -> {
-                if (bitmap == null) {
-                    imagePreview.setVisibility(View.GONE);
-                } else {
-                    imagePreview.setImageBitmap(bitmap);
-                }
-            });
+            mImageLoader.loadImage(
+                    mLifecycle,
+                    mFirstFilePreviewUri,
+                    bitmap -> {
+                        if (bitmap == null) {
+                            imagePreview.setVisibility(View.GONE);
+                        } else {
+                            imagePreview.setImageBitmap(bitmap);
+                        }
+                    });
         } else {
             imagePreview.setVisibility(View.GONE);
         }
diff --git a/java/src/com/android/intentresolver/contentpreview/ImageLoader.kt b/java/src/com/android/intentresolver/contentpreview/ImageLoader.kt
index 225807ee..8d0fb84b 100644
--- a/java/src/com/android/intentresolver/contentpreview/ImageLoader.kt
+++ b/java/src/com/android/intentresolver/contentpreview/ImageLoader.kt
@@ -18,32 +18,29 @@ package com.android.intentresolver.contentpreview
 
 import android.graphics.Bitmap
 import android.net.Uri
+import androidx.lifecycle.Lifecycle
 import java.util.function.Consumer
 
-/**
- * A content preview image loader.
- */
+/** A content preview image loader. */
 interface ImageLoader : suspend (Uri) -> Bitmap?, suspend (Uri, Boolean) -> Bitmap? {
     /**
      * Load preview image asynchronously; caching is allowed.
+     *
      * @param uri content URI
      * @param callback a callback that will be invoked with the loaded image or null if loading has
-     * failed.
+     *   failed.
      */
-    fun loadImage(uri: Uri, callback: Consumer<Bitmap?>)
+    fun loadImage(callerLifecycle: Lifecycle, uri: Uri, callback: Consumer<Bitmap?>)
 
-    /**
-     * Prepopulate the image loader cache.
-     */
+    /** Prepopulate the image loader cache. */
     fun prePopulate(uris: List<Uri>)
 
-    /**
-     * Load preview image; caching is allowed.
-     */
+    /** Load preview image; caching is allowed. */
     override suspend fun invoke(uri: Uri) = invoke(uri, true)
 
     /**
      * Load preview image.
+     *
      * @param uri content URI
      * @param caching indicates if the loaded image could be cached.
      */
diff --git a/java/src/com/android/intentresolver/ImagePreviewImageLoader.kt b/java/src/com/android/intentresolver/contentpreview/ImagePreviewImageLoader.kt
index c97efdd1..89b79a0a 100644
--- a/java/src/com/android/intentresolver/ImagePreviewImageLoader.kt
+++ b/java/src/com/android/intentresolver/contentpreview/ImagePreviewImageLoader.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2022 The Android Open Source Project
+ * Copyright (C) 2023 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,9 +14,9 @@
  * limitations under the License.
  */
 
-package com.android.intentresolver
+package com.android.intentresolver.contentpreview
 
-import android.content.Context
+import android.content.ContentResolver
 import android.graphics.Bitmap
 import android.net.Uri
 import android.util.Log
@@ -26,12 +26,10 @@ import androidx.annotation.VisibleForTesting
 import androidx.collection.LruCache
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.coroutineScope
-import com.android.intentresolver.contentpreview.ImageLoader
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CompletableDeferred
-import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Deferred
-import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
 import java.util.function.Consumer
@@ -42,29 +40,24 @@ private const val TAG = "ImagePreviewImageLoader"
  * Implements preview image loading for the content preview UI. Provides requests deduplication and
  * image caching.
  */
-@VisibleForTesting
-class ImagePreviewImageLoader @JvmOverloads constructor(
-    private val context: Context,
-    private val lifecycle: Lifecycle,
+@VisibleForTesting(otherwise = VisibleForTesting.PACKAGE_PRIVATE)
+class ImagePreviewImageLoader(
+    private val scope: CoroutineScope,
+    thumbnailSize: Int,
+    private val contentResolver: ContentResolver,
     cacheSize: Int,
-    private val dispatcher: CoroutineDispatcher = Dispatchers.IO
 ) : ImageLoader {
 
-    private val thumbnailSize: Size =
-        context.resources.getDimensionPixelSize(R.dimen.chooser_preview_image_max_dimen).let {
-            Size(it, it)
-        }
+    private val thumbnailSize: Size = Size(thumbnailSize, thumbnailSize)
 
     private val lock = Any()
-    @GuardedBy("lock")
-    private val cache = LruCache<Uri, RequestRecord>(cacheSize)
-    @GuardedBy("lock")
-    private val runningRequests = HashMap<Uri, RequestRecord>()
+    @GuardedBy("lock") private val cache = LruCache<Uri, RequestRecord>(cacheSize)
+    @GuardedBy("lock") private val runningRequests = HashMap<Uri, RequestRecord>()
 
     override suspend fun invoke(uri: Uri, caching: Boolean): Bitmap? = loadImageAsync(uri, caching)
 
-    override fun loadImage(uri: Uri, callback: Consumer<Bitmap?>) {
-        lifecycle.coroutineScope.launch {
+    override fun loadImage(callerLifecycle: Lifecycle, uri: Uri, callback: Consumer<Bitmap?>) {
+        callerLifecycle.coroutineScope.launch {
             val image = loadImageAsync(uri, caching = true)
             if (isActive) {
                 callback.accept(image)
@@ -74,28 +67,26 @@ class ImagePreviewImageLoader @JvmOverloads constructor(
 
     override fun prePopulate(uris: List<Uri>) {
         uris.asSequence().take(cache.maxSize()).forEach { uri ->
-            lifecycle.coroutineScope.launch {
-                loadImageAsync(uri, caching = true)
-            }
+            scope.launch { loadImageAsync(uri, caching = true) }
         }
     }
 
     private suspend fun loadImageAsync(uri: Uri, caching: Boolean): Bitmap? {
-        return getRequestDeferred(uri, caching)
-            .await()
+        return getRequestDeferred(uri, caching).await()
     }
 
     private fun getRequestDeferred(uri: Uri, caching: Boolean): Deferred<Bitmap?> {
         var shouldLaunchImageLoading = false
-        val request = synchronized(lock) {
-            cache[uri]
-                ?: runningRequests.getOrPut(uri) {
-                    shouldLaunchImageLoading = true
-                    RequestRecord(uri, CompletableDeferred(), caching)
-                }.apply {
-                    this.caching = this.caching || caching
-                }
-        }
+        val request =
+            synchronized(lock) {
+                cache[uri]
+                    ?: runningRequests
+                        .getOrPut(uri) {
+                            shouldLaunchImageLoading = true
+                            RequestRecord(uri, CompletableDeferred(), caching)
+                        }
+                        .apply { this.caching = this.caching || caching }
+            }
         if (shouldLaunchImageLoading) {
             request.loadBitmapAsync()
         }
@@ -103,22 +94,23 @@ class ImagePreviewImageLoader @JvmOverloads constructor(
     }
 
     private fun RequestRecord.loadBitmapAsync() {
-        lifecycle.coroutineScope.launch(dispatcher) {
-            loadBitmap()
-        }.invokeOnCompletion { cause ->
-            if (cause is CancellationException) {
-                cancel()
+        scope
+            .launch { loadBitmap() }
+            .invokeOnCompletion { cause ->
+                if (cause is CancellationException) {
+                    cancel()
+                }
             }
-        }
     }
 
     private fun RequestRecord.loadBitmap() {
-        val bitmap = try {
-            context.contentResolver.loadThumbnail(uri,  thumbnailSize, null)
-        } catch (t: Throwable) {
-            Log.d(TAG, "failed to load $uri preview", t)
-            null
-        }
+        val bitmap =
+            try {
+                contentResolver.loadThumbnail(uri, thumbnailSize, null)
+            } catch (t: Throwable) {
+                Log.d(TAG, "failed to load $uri preview", t)
+                null
+            }
         complete(bitmap)
     }
 
@@ -144,4 +136,4 @@ class ImagePreviewImageLoader @JvmOverloads constructor(
         val deferred: CompletableDeferred<Bitmap?>,
         @GuardedBy("lock") var caching: Boolean
     )
-}
+}
\ No newline at end of file
diff --git a/java/src/com/android/intentresolver/contentpreview/PreviewViewModel.kt b/java/src/com/android/intentresolver/contentpreview/PreviewViewModel.kt
index 2f4b0211..331b0cb6 100644
--- a/java/src/com/android/intentresolver/contentpreview/PreviewViewModel.kt
+++ b/java/src/com/android/intentresolver/contentpreview/PreviewViewModel.kt
@@ -16,24 +16,45 @@
 
 package com.android.intentresolver.contentpreview
 
-import android.content.ContentResolver
-import android.content.Context
+import android.app.Application
+import androidx.annotation.MainThread
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.ViewModelProvider.AndroidViewModelFactory.Companion.APPLICATION_KEY
+import androidx.lifecycle.viewModelScope
 import androidx.lifecycle.viewmodel.CreationExtras
 import com.android.intentresolver.ChooserRequestParameters
+import com.android.intentresolver.R
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.plus
 
 /** A trivial view model to keep a [PreviewDataProvider] instance over a configuration change */
-class PreviewViewModel(private val contentResolver: ContentResolver) : ViewModel() {
+class PreviewViewModel(private val application: Application) : BasePreviewViewModel() {
     private var previewDataProvider: PreviewDataProvider? = null
+    private var imageLoader: ImagePreviewImageLoader? = null
 
-    fun createOrReuseProvider(chooserRequest: ChooserRequestParameters): PreviewDataProvider {
-        return previewDataProvider
-            ?: PreviewDataProvider(chooserRequest.targetIntent, contentResolver).also {
+    @MainThread
+    override fun createOrReuseProvider(
+        chooserRequest: ChooserRequestParameters
+    ): PreviewDataProvider =
+        previewDataProvider
+            ?: PreviewDataProvider(chooserRequest.targetIntent, application.contentResolver).also {
                 previewDataProvider = it
             }
-    }
+
+    @MainThread
+    override fun createOrReuseImageLoader(): ImageLoader =
+        imageLoader
+            ?: ImagePreviewImageLoader(
+                    viewModelScope + Dispatchers.IO,
+                    thumbnailSize =
+                        application.resources.getDimensionPixelSize(
+                            R.dimen.chooser_preview_image_max_dimen
+                        ),
+                    application.contentResolver,
+                    cacheSize = 16
+                )
+                .also { imageLoader = it }
 
     companion object {
         val Factory: ViewModelProvider.Factory =
@@ -42,10 +63,7 @@ class PreviewViewModel(private val contentResolver: ContentResolver) : ViewModel
                 override fun <T : ViewModel> create(
                     modelClass: Class<T>,
                     extras: CreationExtras
-                ): T =
-                    PreviewViewModel(
-                        (checkNotNull(extras[APPLICATION_KEY]) as Context).contentResolver
-                    ) as T
+                ): T = PreviewViewModel(checkNotNull(extras[APPLICATION_KEY])) as T
             }
     }
 }
diff --git a/java/src/com/android/intentresolver/contentpreview/TextContentPreviewUi.java b/java/src/com/android/intentresolver/contentpreview/TextContentPreviewUi.java
index 3c8a6e48..19fd3bb4 100644
--- a/java/src/com/android/intentresolver/contentpreview/TextContentPreviewUi.java
+++ b/java/src/com/android/intentresolver/contentpreview/TextContentPreviewUi.java
@@ -28,6 +28,7 @@ import android.widget.ImageView;
 import android.widget.TextView;
 
 import androidx.annotation.Nullable;
+import androidx.lifecycle.Lifecycle;
 
 import com.android.intentresolver.R;
 import com.android.intentresolver.widget.ActionRow;
@@ -36,6 +37,7 @@ import java.util.ArrayList;
 import java.util.List;
 
 class TextContentPreviewUi extends ContentPreviewUi {
+    private final Lifecycle mLifecycle;
     @Nullable
     private final CharSequence mSharingText;
     @Nullable
@@ -47,12 +49,14 @@ class TextContentPreviewUi extends ContentPreviewUi {
     private final HeadlineGenerator mHeadlineGenerator;
 
     TextContentPreviewUi(
+            Lifecycle lifecycle,
             @Nullable CharSequence sharingText,
             @Nullable CharSequence previewTitle,
             @Nullable Uri previewThumbnail,
             ChooserContentPreviewUi.ActionFactory actionFactory,
             ImageLoader imageLoader,
             HeadlineGenerator headlineGenerator) {
+        mLifecycle = lifecycle;
         mSharingText = sharingText;
         mPreviewTitle = previewTitle;
         mPreviewThumbnail = previewThumbnail;
@@ -117,6 +121,7 @@ class TextContentPreviewUi extends ContentPreviewUi {
             previewThumbnailView.setVisibility(View.GONE);
         } else {
             mImageLoader.loadImage(
+                    mLifecycle,
                     mPreviewThumbnail,
                     (bitmap) -> updateViewWithImage(
                             contentPreviewLayout.findViewById(
diff --git a/java/tests/src/com/android/intentresolver/ChooserWrapperActivity.java b/java/tests/src/com/android/intentresolver/ChooserWrapperActivity.java
index 82ba8d4d..fa934f87 100644
--- a/java/tests/src/com/android/intentresolver/ChooserWrapperActivity.java
+++ b/java/tests/src/com/android/intentresolver/ChooserWrapperActivity.java
@@ -32,10 +32,11 @@ import android.net.Uri;
 import android.os.Bundle;
 import android.os.UserHandle;
 
+import androidx.lifecycle.ViewModelProvider;
+
 import com.android.intentresolver.AbstractMultiProfilePagerAdapter.CrossProfileIntentsChecker;
 import com.android.intentresolver.chooser.DisplayResolveInfo;
 import com.android.intentresolver.chooser.TargetInfo;
-import com.android.intentresolver.contentpreview.ImageLoader;
 import com.android.intentresolver.flags.FeatureFlagRepository;
 import com.android.intentresolver.grid.ChooserGridAdapter;
 import com.android.intentresolver.shortcuts.ShortcutLoader;
@@ -194,10 +195,10 @@ public class ChooserWrapperActivity
     }
 
     @Override
-    protected ImageLoader createPreviewImageLoader() {
-        return sOverrides.imageLoader == null
-                ? super.createPreviewImageLoader()
-                : sOverrides.imageLoader;
+    protected ViewModelProvider.Factory createPreviewViewModelFactory() {
+        return TestContentPreviewViewModel.Companion.wrap(
+                super.createPreviewViewModelFactory(),
+                sOverrides.imageLoader);
     }
 
     @Override
diff --git a/java/tests/src/com/android/intentresolver/TestContentPreviewViewModel.kt b/java/tests/src/com/android/intentresolver/TestContentPreviewViewModel.kt
new file mode 100644
index 00000000..d239f612
--- /dev/null
+++ b/java/tests/src/com/android/intentresolver/TestContentPreviewViewModel.kt
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.intentresolver
+
+import androidx.lifecycle.ViewModel
+import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.viewmodel.CreationExtras
+import com.android.intentresolver.contentpreview.BasePreviewViewModel
+import com.android.intentresolver.contentpreview.ImageLoader
+import com.android.intentresolver.contentpreview.PreviewDataProvider
+
+/** A test content preview model that supports image loader override. */
+class TestContentPreviewViewModel(
+    private val viewModel: BasePreviewViewModel,
+    private val imageLoader: ImageLoader? = null,
+) : BasePreviewViewModel() {
+    override fun createOrReuseProvider(
+        chooserRequest: ChooserRequestParameters
+    ): PreviewDataProvider = viewModel.createOrReuseProvider(chooserRequest)
+
+    override fun createOrReuseImageLoader(): ImageLoader =
+        imageLoader ?: viewModel.createOrReuseImageLoader()
+
+    companion object {
+        fun wrap(
+            factory: ViewModelProvider.Factory,
+            imageLoader: ImageLoader?,
+        ): ViewModelProvider.Factory =
+            object : ViewModelProvider.Factory {
+                @Suppress("UNCHECKED_CAST")
+                override fun <T : ViewModel> create(
+                    modelClass: Class<T>,
+                    extras: CreationExtras
+                ): T {
+                    return TestContentPreviewViewModel(
+                        factory.create(modelClass, extras) as BasePreviewViewModel,
+                        imageLoader,
+                    ) as T
+                }
+            }
+    }
+}
diff --git a/java/tests/src/com/android/intentresolver/TestPreviewImageLoader.kt b/java/tests/src/com/android/intentresolver/TestPreviewImageLoader.kt
index 74a253b8..bf87ed8a 100644
--- a/java/tests/src/com/android/intentresolver/TestPreviewImageLoader.kt
+++ b/java/tests/src/com/android/intentresolver/TestPreviewImageLoader.kt
@@ -18,13 +18,12 @@ package com.android.intentresolver
 
 import android.graphics.Bitmap
 import android.net.Uri
+import androidx.lifecycle.Lifecycle
 import com.android.intentresolver.contentpreview.ImageLoader
 import java.util.function.Consumer
 
-internal class TestPreviewImageLoader(
-    private val bitmaps: Map<Uri, Bitmap>
-) : ImageLoader {
-    override fun loadImage(uri: Uri, callback: Consumer<Bitmap?>) {
+internal class TestPreviewImageLoader(private val bitmaps: Map<Uri, Bitmap>) : ImageLoader {
+    override fun loadImage(callerLifecycle: Lifecycle, uri: Uri, callback: Consumer<Bitmap?>) {
         callback.accept(bitmaps[uri])
     }
 
diff --git a/java/tests/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUiTest.kt b/java/tests/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUiTest.kt
index 0bcd8423..c62f36ce 100644
--- a/java/tests/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUiTest.kt
+++ b/java/tests/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUiTest.kt
@@ -39,7 +39,11 @@ class ChooserContentPreviewUiTest {
     private val headlineGenerator = mock<HeadlineGenerator>()
     private val imageLoader =
         object : ImageLoader {
-            override fun loadImage(uri: Uri, callback: Consumer<Bitmap?>) {
+            override fun loadImage(
+                callerLifecycle: Lifecycle,
+                uri: Uri,
+                callback: Consumer<Bitmap?>,
+            ) {
                 callback.accept(null)
             }
             override fun prePopulate(uris: List<Uri>) = Unit
diff --git a/java/tests/src/com/android/intentresolver/ImagePreviewImageLoaderTest.kt b/java/tests/src/com/android/intentresolver/contentpreview/ImagePreviewImageLoaderTest.kt
index 3c399cc4..184401a0 100644
--- a/java/tests/src/com/android/intentresolver/ImagePreviewImageLoaderTest.kt
+++ b/java/tests/src/com/android/intentresolver/contentpreview/ImagePreviewImageLoaderTest.kt
@@ -14,15 +14,19 @@
  * limitations under the License.
  */
 
-package com.android.intentresolver
+package com.android.intentresolver.contentpreview
 
 import android.content.ContentResolver
-import android.content.Context
-import android.content.res.Resources
 import android.graphics.Bitmap
 import android.net.Uri
 import android.util.Size
 import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.coroutineScope
+import com.android.intentresolver.TestLifecycleOwner
+import com.android.intentresolver.any
+import com.android.intentresolver.anyOrNull
+import com.android.intentresolver.mock
+import com.android.intentresolver.whenever
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineStart.UNDISPATCHED
 import kotlinx.coroutines.Dispatchers
@@ -30,6 +34,7 @@ import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.async
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.plus
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestCoroutineScheduler
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
@@ -49,28 +54,26 @@ class ImagePreviewImageLoaderTest {
     private val uriOne = Uri.parse("content://org.package.app/image-1.png")
     private val uriTwo = Uri.parse("content://org.package.app/image-2.png")
     private val bitmap = Bitmap.createBitmap(1, 1, Bitmap.Config.ARGB_8888)
-    private val contentResolver = mock<ContentResolver> {
-        whenever(loadThumbnail(any(), any(), anyOrNull())).thenReturn(bitmap)
-    }
-    private val resources = mock<Resources> {
-        whenever(getDimensionPixelSize(R.dimen.chooser_preview_image_max_dimen))
-            .thenReturn(imageSize.width)
-    }
-    private val context = mock<Context> {
-        whenever(this.resources).thenReturn(this@ImagePreviewImageLoaderTest.resources)
-        whenever(this.contentResolver).thenReturn(this@ImagePreviewImageLoaderTest.contentResolver)
-    }
-    private val scheduler = TestCoroutineScheduler()
+    private val contentResolver =
+        mock<ContentResolver> {
+            whenever(loadThumbnail(any(), any(), anyOrNull())).thenReturn(bitmap)
+        }
     private val lifecycleOwner = TestLifecycleOwner()
-    private val dispatcher = UnconfinedTestDispatcher(scheduler)
-    private val testSubject = ImagePreviewImageLoader(
-        context, lifecycleOwner.lifecycle, 1, dispatcher
-    )
+    private val dispatcher = UnconfinedTestDispatcher()
+    private lateinit var testSubject: ImagePreviewImageLoader
 
     @Before
     fun setup() {
         Dispatchers.setMain(dispatcher)
         lifecycleOwner.state = Lifecycle.State.CREATED
+        // create test subject after we've updated the lifecycle dispatcher
+        testSubject =
+            ImagePreviewImageLoader(
+                lifecycleOwner.lifecycle.coroutineScope + dispatcher,
+                imageSize.width,
+                contentResolver,
+                1,
+            )
     }
 
     @After
@@ -110,14 +113,16 @@ class ImagePreviewImageLoaderTest {
     fun invoke_overlappedRequests_Deduplicate() = runTest {
         val scheduler = TestCoroutineScheduler()
         val dispatcher = StandardTestDispatcher(scheduler)
-        val testSubject = ImagePreviewImageLoader(context, lifecycleOwner.lifecycle, 1, dispatcher)
+        val testSubject =
+            ImagePreviewImageLoader(
+                lifecycleOwner.lifecycle.coroutineScope + dispatcher,
+                imageSize.width,
+                contentResolver,
+                1,
+            )
         coroutineScope {
-            launch(start = UNDISPATCHED) {
-                testSubject(uriOne, false)
-            }
-            launch(start = UNDISPATCHED) {
-                testSubject(uriOne, false)
-            }
+            launch(start = UNDISPATCHED) { testSubject(uriOne, false) }
+            launch(start = UNDISPATCHED) { testSubject(uriOne, false) }
             scheduler.advanceUntilIdle()
         }
 
@@ -154,11 +159,15 @@ class ImagePreviewImageLoaderTest {
     fun invoke_imageLoaderScopeClosedMidflight_throwsCancellationException() = runTest {
         val scheduler = TestCoroutineScheduler()
         val dispatcher = StandardTestDispatcher(scheduler)
-        val testSubject = ImagePreviewImageLoader(context, lifecycleOwner.lifecycle, 1, dispatcher)
+        val testSubject =
+            ImagePreviewImageLoader(
+                lifecycleOwner.lifecycle.coroutineScope + dispatcher,
+                imageSize.width,
+                contentResolver,
+                1
+            )
         coroutineScope {
-            val deferred = async(start = UNDISPATCHED) {
-                testSubject(uriOne, false)
-            }
+            val deferred = async(start = UNDISPATCHED) { testSubject(uriOne, false) }
             lifecycleOwner.state = Lifecycle.State.DESTROYED
             scheduler.advanceUntilIdle()
             deferred.await()
@@ -169,14 +178,16 @@ class ImagePreviewImageLoaderTest {
     fun invoke_multipleCallsWithDifferentCacheInstructions_cachingPrevails() = runTest {
         val scheduler = TestCoroutineScheduler()
         val dispatcher = StandardTestDispatcher(scheduler)
-        val testSubject = ImagePreviewImageLoader(context, lifecycleOwner.lifecycle, 1, dispatcher)
+        val testSubject =
+            ImagePreviewImageLoader(
+                lifecycleOwner.lifecycle.coroutineScope + dispatcher,
+                imageSize.width,
+                contentResolver,
+                1
+            )
         coroutineScope {
-            launch(start = UNDISPATCHED) {
-                testSubject(uriOne, false)
-            }
-            launch(start = UNDISPATCHED) {
-                testSubject(uriOne, true)
-            }
+            launch(start = UNDISPATCHED) { testSubject(uriOne, false) }
+            launch(start = UNDISPATCHED) { testSubject(uriOne, true) }
             scheduler.advanceUntilIdle()
         }
         testSubject(uriOne, true)