Retain the image loader through configuration change

Make PreviewViewModel own an image loader instance.
This required ImagePreviewImageLoader to:
 * receive a CoroutineScope instance instead of a Lifecycle as ViewModel
   does not offer an associated Lifecycle;
 * make ImagePreviewImageLoader#loadImage to receive a caller's
   Lifecycle to avoid potential memory (Activity) leak.
With ImagePreviewImageLoader now receiving a scope, the dispatcher
argument is dropped as a scope can be configured with a proper
dispatcher.
As ChooserActivity is no longer responsible for the image loader
instantiation, to facilitate testing, a method hook to create an image
loader is replaced with a method hook to create a preview view model and
all the associated plumbing is done for the tests.

Plus some kfmt complience changes.
ImagePreviewImageLoader (and ImagePreviewImageLoaderTest) is moved into
contentpreview package as it's accessed only from within this package.

Fix: 282029067
Test: manual testing with an injected debug logging
Change-Id: I5dcdee2599714a2c51c3e1b63f5c727e43b26f6a
diff --git a/java/src/com/android/intentresolver/ChooserActivity.java b/java/src/com/android/intentresolver/ChooserActivity.java
index 0786e08..014aa2a 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.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 @@
             }
         });
 
+        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 @@
     }
 
     @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 0000000..103e8bf
--- /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 787af95..9100b39 100644
--- a/java/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUi.java
+++ b/java/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUi.java
@@ -118,6 +118,7 @@
         int previewType = previewData.getPreviewType();
         if (previewType == CONTENT_PREVIEW_TEXT) {
             return createTextPreview(
+                    mLifecycle,
                     targetIntent,
                     actionFactory,
                     imageLoader,
@@ -140,6 +141,7 @@
         if (!TextUtils.isEmpty(text)) {
             FilesPlusTextContentPreviewUi previewUi =
                     new FilesPlusTextContentPreviewUi(
+                            mLifecycle,
                             isSingleImageShare,
                             previewData.getUriCount(),
                             targetIntent.getCharSequenceExtra(Intent.EXTRA_TEXT),
@@ -180,6 +182,7 @@
     }
 
     private static TextContentPreviewUi createTextPreview(
+            Lifecycle lifecycle,
             Intent targetIntent,
             ChooserContentPreviewUi.ActionFactory actionFactory,
             ImageLoader imageLoader,
@@ -195,6 +198,7 @@
             }
         }
         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 860423c..e4e3383 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.TextView;
 
 import androidx.annotation.Nullable;
+import androidx.lifecycle.Lifecycle;
 
 import com.android.intentresolver.R;
 import com.android.intentresolver.widget.ActionRow;
@@ -48,6 +49,7 @@
  * 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 @@
     private boolean mAllVideos;
 
     FilesPlusTextContentPreviewUi(
+            Lifecycle lifecycle,
             boolean isSingleImage,
             int fileCount,
             CharSequence text,
@@ -70,6 +73,7 @@
             ImageLoader imageLoader,
             MimeTypeClassifier typeClassifier,
             HeadlineGenerator headlineGenerator) {
+        mLifecycle = lifecycle;
         if (isSingleImage && fileCount != 1) {
             throw new IllegalArgumentException(
                     "fileCount = " + fileCount + " and isSingleImage = true");
@@ -155,13 +159,16 @@
 
         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 225807e..8d0fb84 100644
--- a/java/src/com/android/intentresolver/contentpreview/ImageLoader.kt
+++ b/java/src/com/android/intentresolver/contentpreview/ImageLoader.kt
@@ -18,32 +18,29 @@
 
 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
similarity index 60%
rename from java/src/com/android/intentresolver/ImagePreviewImageLoader.kt
rename to java/src/com/android/intentresolver/contentpreview/ImagePreviewImageLoader.kt
index c97efdd..89b79a0 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.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 @@
  * 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 @@
 
     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 @@
     }
 
     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 @@
         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 2f4b021..331b0cb 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 @@
                 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 3c8a6e4..19fd3bb 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.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.List;
 
 class TextContentPreviewUi extends ContentPreviewUi {
+    private final Lifecycle mLifecycle;
     @Nullable
     private final CharSequence mSharingText;
     @Nullable
@@ -47,12 +49,14 @@
     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 @@
             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 82ba8d4..fa934f8 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.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 @@
     }
 
     @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 0000000..d239f61
--- /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 74a253b..bf87ed8 100644
--- a/java/tests/src/com/android/intentresolver/TestPreviewImageLoader.kt
+++ b/java/tests/src/com/android/intentresolver/TestPreviewImageLoader.kt
@@ -18,13 +18,12 @@
 
 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 0bcd842..c62f36c 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 @@
     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
similarity index 71%
rename from java/tests/src/com/android/intentresolver/ImagePreviewImageLoaderTest.kt
rename to java/tests/src/com/android/intentresolver/contentpreview/ImagePreviewImageLoaderTest.kt
index 3c399cc..184401a 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.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 @@
     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 @@
     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 @@
     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 @@
     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)