Merge "Make ImageLoader injectable" into main

7 files changed, 359 insertions, 291 deletions

diff --git a/java/src/com/android/intentresolver/contentpreview/ImageLoaderModule.kt b/java/src/com/android/intentresolver/contentpreview/ImageLoaderModule.kt
new file mode 100644
index 00000000..b861a24a
--- /dev/null
+++ b/java/src/com/android/intentresolver/contentpreview/ImageLoaderModule.kt
@@ -0,0 +1,44 @@
+/*
+ * Copyright 2024 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 android.content.res.Resources
+import com.android.intentresolver.R
+import com.android.intentresolver.inject.ApplicationOwned
+import dagger.Binds
+import dagger.Module
+import dagger.Provides
+import dagger.hilt.InstallIn
+import dagger.hilt.android.components.ActivityRetainedComponent
+import dagger.hilt.android.scopes.ActivityRetainedScoped
+
+@Module
+@InstallIn(ActivityRetainedComponent::class)
+interface ImageLoaderModule {
+    @Binds
+    @ActivityRetainedScoped
+    fun imageLoader(previewImageLoader: ImagePreviewImageLoader): ImageLoader
+
+    companion object {
+        @Provides
+        @ThumbnailSize
+        fun thumbnailSize(@ApplicationOwned resources: Resources): Int =
+            resources.getDimensionPixelSize(R.dimen.chooser_preview_image_max_dimen)
+
+        @Provides @PreviewCacheSize fun cacheSize() = 16
+    }
+}
diff --git a/java/src/com/android/intentresolver/contentpreview/ImagePreviewImageLoader.kt b/java/src/com/android/intentresolver/contentpreview/ImagePreviewImageLoader.kt
index 572ccf0b..fab7203e 100644
--- a/java/src/com/android/intentresolver/contentpreview/ImagePreviewImageLoader.kt
+++ b/java/src/com/android/intentresolver/contentpreview/ImagePreviewImageLoader.kt
@@ -24,17 +24,31 @@ import android.util.Size
 import androidx.annotation.GuardedBy
 import androidx.annotation.VisibleForTesting
 import androidx.collection.LruCache
+import com.android.intentresolver.inject.Background
 import java.util.function.Consumer
+import javax.inject.Inject
+import javax.inject.Qualifier
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineExceptionHandler
+import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Deferred
+import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.sync.Semaphore
 
 private const val TAG = "ImagePreviewImageLoader"
 
+@Qualifier @MustBeDocumented @Retention(AnnotationRetention.BINARY) annotation class ThumbnailSize
+
+@Qualifier
+@MustBeDocumented
+@Retention(AnnotationRetention.BINARY)
+annotation class PreviewCacheSize
+
 /**
  * Implements preview image loading for the content preview UI. Provides requests deduplication,
  * image caching, and a limit on the number of parallel loadings.
@@ -52,6 +66,26 @@ constructor(
     private val contentResolverSemaphore: Semaphore,
 ) : ImageLoader {
 
+    @Inject
+    constructor(
+        @Background dispatcher: CoroutineDispatcher,
+        @ThumbnailSize thumbnailSize: Int,
+        contentResolver: ContentResolver,
+        @PreviewCacheSize cacheSize: Int,
+    ) : this(
+        CoroutineScope(
+            SupervisorJob() +
+                dispatcher +
+                CoroutineExceptionHandler { _, exception ->
+                    Log.w(TAG, "Uncaught exception in ImageLoader", exception)
+                } +
+                CoroutineName("ImageLoader")
+        ),
+        thumbnailSize,
+        contentResolver,
+        cacheSize,
+    )
+
     constructor(
         scope: CoroutineScope,
         thumbnailSize: Int,
diff --git a/tests/activity/src/com/android/intentresolver/UnbundledChooserActivityTest.java b/tests/activity/src/com/android/intentresolver/UnbundledChooserActivityTest.java
index f597d7f2..c7b41ce0 100644
--- a/tests/activity/src/com/android/intentresolver/UnbundledChooserActivityTest.java
+++ b/tests/activity/src/com/android/intentresolver/UnbundledChooserActivityTest.java
@@ -800,7 +800,7 @@ public class UnbundledChooserActivityTest {
         Uri uri = createTestContentProviderUri("image/png", null);
         Intent sendIntent = createSendImageIntent(uri);
         ChooserActivityOverrideData.getInstance().imageLoader =
-                new TestPreviewImageLoader(Collections.emptyMap());
+                new FakeImageLoader(Collections.emptyMap());
         sendIntent.putExtra(Intent.EXTRA_TEXT, "https://google.com/search?q=google");
 
         List<ResolvedComponentInfo> resolvedComponentInfos = Arrays.asList(
@@ -958,7 +958,7 @@ public class UnbundledChooserActivityTest {
 
         Intent sendIntent = createSendUriIntentWithPreview(uris);
         ChooserActivityOverrideData.getInstance().imageLoader =
-                new TestPreviewImageLoader(Collections.emptyMap());
+                new FakeImageLoader(Collections.emptyMap());
 
         List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
 
@@ -1076,7 +1076,7 @@ public class UnbundledChooserActivityTest {
         bitmaps.put(imgTwoUri, createWideBitmap(Color.GREEN));
         bitmaps.put(docUri, createWideBitmap(Color.BLUE));
         ChooserActivityOverrideData.getInstance().imageLoader =
-                new TestPreviewImageLoader(bitmaps);
+                new FakeImageLoader(bitmaps);
 
         List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
         setupResolverControllers(resolvedComponentInfos);
@@ -3122,6 +3122,6 @@ public class UnbundledChooserActivityTest {
     }
 
     private static ImageLoader createImageLoader(Uri uri, Bitmap bitmap) {
-        return new TestPreviewImageLoader(Collections.singletonMap(uri, bitmap));
+        return new FakeImageLoader(Collections.singletonMap(uri, bitmap));
     }
 }
diff --git a/tests/activity/src/com/android/intentresolver/v2/UnbundledChooserActivityTest.java b/tests/activity/src/com/android/intentresolver/v2/UnbundledChooserActivityTest.java
index b8113422..a7221c10 100644
--- a/tests/activity/src/com/android/intentresolver/v2/UnbundledChooserActivityTest.java
+++ b/tests/activity/src/com/android/intentresolver/v2/UnbundledChooserActivityTest.java
@@ -119,15 +119,16 @@ import androidx.test.rule.ActivityTestRule;
 
 import com.android.intentresolver.AnnotatedUserHandles;
 import com.android.intentresolver.ChooserListAdapter;
+import com.android.intentresolver.FakeImageLoader;
 import com.android.intentresolver.Flags;
 import com.android.intentresolver.IChooserWrapper;
 import com.android.intentresolver.R;
 import com.android.intentresolver.ResolvedComponentInfo;
 import com.android.intentresolver.ResolverDataProvider;
 import com.android.intentresolver.TestContentProvider;
-import com.android.intentresolver.TestPreviewImageLoader;
 import com.android.intentresolver.chooser.DisplayResolveInfo;
 import com.android.intentresolver.contentpreview.ImageLoader;
+import com.android.intentresolver.contentpreview.ImageLoaderModule;
 import com.android.intentresolver.inject.PackageManagerModule;
 import com.android.intentresolver.logging.EventLog;
 import com.android.intentresolver.logging.FakeEventLog;
@@ -160,7 +161,6 @@ import org.mockito.Mockito;
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -183,7 +183,8 @@ import javax.inject.Inject;
 @UninstallModules({
         AppPredictionModule.class,
         ImageEditorModule.class,
-        PackageManagerModule.class
+        PackageManagerModule.class,
+        ImageLoaderModule.class,
 })
 public class UnbundledChooserActivityTest {
 
@@ -239,6 +240,11 @@ public class UnbundledChooserActivityTest {
     @BindValue
     PackageManager mPackageManager;
 
+    private final FakeImageLoader mFakeImageLoader = new FakeImageLoader();
+
+    @BindValue
+    final ImageLoader mImageLoader = mFakeImageLoader;
+
     @Before
     public void setUp() {
         // TODO: use the other form of `adoptShellPermissionIdentity()` where we explicitly list the
@@ -257,6 +263,9 @@ public class UnbundledChooserActivityTest {
         // values to the dependency graph at activity launch time. This allows replacing
         // arbitrary bindings per-test case if needed.
         mPackageManager = mContext.getPackageManager();
+
+        // TODO: inject image loader in the prod code and remove this override
+        ChooserActivityOverrideData.getInstance().imageLoader = mFakeImageLoader;
     }
 
     public UnbundledChooserActivityTest(boolean appPredictionAvailable) {
@@ -434,14 +443,13 @@ public class UnbundledChooserActivityTest {
     }
 
     @Test
-    public void visiblePreviewTitleAndThumbnail() throws InterruptedException {
+    public void visiblePreviewTitleAndThumbnail() {
         String previewTitle = "My Content Preview Title";
         Uri uri = Uri.parse(
                 "android.resource://com.android.frameworks.coretests/"
                 + com.android.intentresolver.tests.R.drawable.test320x240);
         Intent sendIntent = createSendTextIntentWithPreview(previewTitle, uri);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                createImageLoader(uri, createBitmap());
+        mFakeImageLoader.setBitmap(uri, createBitmap());
         List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
 
         setupResolverControllers(resolvedComponentInfos);
@@ -707,8 +715,7 @@ public class UnbundledChooserActivityTest {
     public void testFilePlusTextSharing_ExcludeText() {
         Uri uri = createTestContentProviderUri(null, "image/png");
         Intent sendIntent = createSendImageIntent(uri);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                createImageLoader(uri, createBitmap());
+        mFakeImageLoader.setBitmap(uri, createBitmap());
         sendIntent.putExtra(Intent.EXTRA_TEXT, "https://google.com/search?q=google");
 
         List<ResolvedComponentInfo> resolvedComponentInfos = Arrays.asList(
@@ -749,8 +756,7 @@ public class UnbundledChooserActivityTest {
     public void testFilePlusTextSharing_RemoveAndAddBackText() {
         Uri uri = createTestContentProviderUri("application/pdf", "image/png");
         Intent sendIntent = createSendImageIntent(uri);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                createImageLoader(uri, createBitmap());
+        mFakeImageLoader.setBitmap(uri, createBitmap());
         final String text = "https://google.com/search?q=google";
         sendIntent.putExtra(Intent.EXTRA_TEXT, text);
 
@@ -797,8 +803,7 @@ public class UnbundledChooserActivityTest {
     public void testFilePlusTextSharing_TextExclusionDoesNotAffectAlternativeIntent() {
         Uri uri = createTestContentProviderUri("image/png", null);
         Intent sendIntent = createSendImageIntent(uri);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                createImageLoader(uri, createBitmap());
+        mFakeImageLoader.setBitmap(uri, createBitmap());
         sendIntent.putExtra(Intent.EXTRA_TEXT, "https://google.com/search?q=google");
 
         Intent alternativeIntent = createSendTextIntent();
@@ -841,8 +846,6 @@ public class UnbundledChooserActivityTest {
     public void testImagePlusTextSharing_failedThumbnailAndExcludedText_textChanges() {
         Uri uri = createTestContentProviderUri("image/png", null);
         Intent sendIntent = createSendImageIntent(uri);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                new TestPreviewImageLoader(Collections.emptyMap());
         sendIntent.putExtra(Intent.EXTRA_TEXT, "https://google.com/search?q=google");
 
         List<ResolvedComponentInfo> resolvedComponentInfos = Arrays.asList(
@@ -937,8 +940,7 @@ public class UnbundledChooserActivityTest {
 
         Uri uri = createTestContentProviderUri("image/png", null);
         Intent sendIntent = createSendImageIntent(uri);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                createImageLoader(uri, createBitmap());
+        mFakeImageLoader.setBitmap(uri, createBitmap());
 
         List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
 
@@ -962,8 +964,7 @@ public class UnbundledChooserActivityTest {
         uris.add(uri);
 
         Intent sendIntent = createSendUriIntentWithPreview(uris);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                createImageLoader(uri, createWideBitmap());
+        mFakeImageLoader.setBitmap(uri, createWideBitmap());
 
         List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
 
@@ -1000,8 +1001,6 @@ public class UnbundledChooserActivityTest {
         uris.add(uri);
 
         Intent sendIntent = createSendUriIntentWithPreview(uris);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                new TestPreviewImageLoader(Collections.emptyMap());
 
         List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
 
@@ -1019,8 +1018,7 @@ public class UnbundledChooserActivityTest {
         ArrayList<Uri> uris = new ArrayList<>(1);
         uris.add(uri);
         Intent sendIntent = createSendUriIntentWithPreview(uris);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                createImageLoader(uri, createBitmap());
+        mFakeImageLoader.setBitmap(uri, createBitmap());
 
         List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
 
@@ -1046,8 +1044,7 @@ public class UnbundledChooserActivityTest {
         }
         uris.add(imageUri);
         Intent sendIntent = createSendUriIntentWithPreview(uris);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                createImageLoader(imageUri, createBitmap());
+        mFakeImageLoader.setBitmap(imageUri, createBitmap());
 
         List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
         setupResolverControllers(resolvedComponentInfos);
@@ -1079,8 +1076,7 @@ public class UnbundledChooserActivityTest {
         uris.add(uri);
 
         Intent sendIntent = createSendUriIntentWithPreview(uris);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                createImageLoader(uri, createBitmap());
+        mFakeImageLoader.setBitmap(uri, createBitmap());
 
         List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
 
@@ -1114,12 +1110,9 @@ public class UnbundledChooserActivityTest {
         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);
+        mFakeImageLoader.setBitmap(imgOneUri, createWideBitmap(Color.RED));
+        mFakeImageLoader.setBitmap(imgTwoUri, createWideBitmap(Color.GREEN));
+        mFakeImageLoader.setBitmap(docUri, createWideBitmap(Color.BLUE));
 
         List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
         setupResolverControllers(resolvedComponentInfos);
@@ -1167,8 +1160,7 @@ public class UnbundledChooserActivityTest {
 
         Intent sendIntent = createSendUriIntentWithPreview(uris);
         sendIntent.putExtra(Intent.EXTRA_TEXT, sharedText);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                createImageLoader(uri, createBitmap());
+        mFakeImageLoader.setBitmap(uri, createBitmap());
 
         List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
 
@@ -1197,8 +1189,7 @@ public class UnbundledChooserActivityTest {
 
         Intent sendIntent = createSendUriIntentWithPreview(uris);
         sendIntent.putExtra(Intent.EXTRA_TEXT, sharedText);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                createImageLoader(uri, createBitmap());
+        mFakeImageLoader.setBitmap(uri, createBitmap());
 
         List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
 
@@ -1234,8 +1225,7 @@ public class UnbundledChooserActivityTest {
 
         Intent sendIntent = createSendUriIntentWithPreview(uris);
         sendIntent.putExtra(Intent.EXTRA_TEXT, sharedText);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                createImageLoader(uri, createBitmap());
+        mFakeImageLoader.setBitmap(uri, createBitmap());
 
         List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
 
@@ -1331,8 +1321,7 @@ public class UnbundledChooserActivityTest {
         uris.add(uri);
 
         Intent sendIntent = createSendUriIntentWithPreview(uris);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                createImageLoader(uri, createBitmap());
+        mFakeImageLoader.setBitmap(uri, createBitmap());
 
         List<ResolvedComponentInfo> resolvedComponentInfos = createResolvedComponentsForTest(2);
 
@@ -2228,8 +2217,7 @@ public class UnbundledChooserActivityTest {
         uris.add(uri);
 
         Intent sendIntent = createSendUriIntentWithPreview(uris);
-        ChooserActivityOverrideData.getInstance().imageLoader =
-                createImageLoader(uri, createWideBitmap());
+        mFakeImageLoader.setBitmap(uri, createWideBitmap());
 
         mActivityRule.launchActivity(Intent.createChooser(sendIntent, "Scrollable preview test"));
         waitForIdle();
@@ -3134,8 +3122,4 @@ public class UnbundledChooserActivityTest {
                 };
         return shortcutLoaders;
     }
-
-    private static ImageLoader createImageLoader(Uri uri, Bitmap bitmap) {
-        return new TestPreviewImageLoader(Collections.singletonMap(uri, bitmap));
-    }
 }
diff --git a/tests/shared/src/com/android/intentresolver/TestPreviewImageLoader.kt b/tests/shared/src/com/android/intentresolver/FakeImageLoader.kt
index f0203bb6..c57ea78b 100644
--- a/tests/shared/src/com/android/intentresolver/TestPreviewImageLoader.kt
+++ b/tests/shared/src/com/android/intentresolver/FakeImageLoader.kt
@@ -22,7 +22,9 @@ import com.android.intentresolver.contentpreview.ImageLoader
 import java.util.function.Consumer
 import kotlinx.coroutines.CoroutineScope
 
-class TestPreviewImageLoader(private val bitmaps: Map<Uri, Bitmap>) : ImageLoader {
+class FakeImageLoader(initialBitmaps: Map<Uri, Bitmap> = emptyMap()) : ImageLoader {
+    private val bitmaps = HashMap<Uri, Bitmap>().apply { putAll(initialBitmaps) }
+
     override fun loadImage(callerScope: CoroutineScope, uri: Uri, callback: Consumer<Bitmap?>) {
         callback.accept(bitmaps[uri])
     }
@@ -30,4 +32,8 @@ class TestPreviewImageLoader(private val bitmaps: Map<Uri, Bitmap>) : ImageLoade
     override suspend fun invoke(uri: Uri, caching: Boolean): Bitmap? = bitmaps[uri]
 
     override fun prePopulate(uris: List<Uri>) = Unit
+
+    fun setBitmap(uri: Uri, bitmap: Bitmap) {
+        bitmaps[uri] = bitmap
+    }
 }
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUiTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUiTest.kt
index c7c3c516..68b277e7 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUiTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/ChooserContentPreviewUiTest.kt
@@ -21,7 +21,7 @@ import android.net.Uri
 import android.platform.test.flag.junit.CheckFlagsRule
 import android.platform.test.flag.junit.DeviceFlagsValueProvider
 import com.android.intentresolver.ContentTypeHint
-import com.android.intentresolver.TestPreviewImageLoader
+import com.android.intentresolver.FakeImageLoader
 import com.android.intentresolver.contentpreview.ChooserContentPreviewUi.ActionFactory
 import com.android.intentresolver.mock
 import com.android.intentresolver.whenever
@@ -43,7 +43,7 @@ class ChooserContentPreviewUiTest {
     private val testScope = TestScope(EmptyCoroutineContext + UnconfinedTestDispatcher())
     private val previewData = mock<PreviewDataProvider>()
     private val headlineGenerator = mock<HeadlineGenerator>()
-    private val imageLoader = TestPreviewImageLoader(emptyMap())
+    private val imageLoader = FakeImageLoader(emptyMap())
     private val testMetadataText: CharSequence = "Test metadata text"
     private val actionFactory =
         object : ActionFactory {
diff --git a/tests/unit/src/com/android/intentresolver/contentpreview/ImagePreviewImageLoaderTest.kt b/tests/unit/src/com/android/intentresolver/contentpreview/ImagePreviewImageLoaderTest.kt
index 89978707..41989bda 100644
--- a/tests/unit/src/com/android/intentresolver/contentpreview/ImagePreviewImageLoaderTest.kt
+++ b/tests/unit/src/com/android/intentresolver/contentpreview/ImagePreviewImageLoaderTest.kt
@@ -20,9 +20,6 @@ import android.content.ContentResolver
 import android.graphics.Bitmap
 import android.net.Uri
 import android.util.Size
-import androidx.lifecycle.Lifecycle
-import androidx.lifecycle.coroutineScope
-import androidx.lifecycle.testing.TestLifecycleOwner
 import com.android.intentresolver.any
 import com.android.intentresolver.anyOrNull
 import com.android.intentresolver.mock
@@ -38,25 +35,22 @@ import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.CoroutineStart.UNDISPATCHED
-import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.Runnable
 import kotlinx.coroutines.async
+import kotlinx.coroutines.cancel
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.plus
 import kotlinx.coroutines.sync.Semaphore
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestCoroutineScheduler
+import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
-import kotlinx.coroutines.test.resetMain
 import kotlinx.coroutines.test.runTest
-import kotlinx.coroutines.test.setMain
 import kotlinx.coroutines.yield
-import org.junit.After
 import org.junit.Assert.assertTrue
-import org.junit.Before
 import org.junit.Test
 import org.mockito.Mockito.never
 import org.mockito.Mockito.times
@@ -72,281 +66,287 @@ class ImagePreviewImageLoaderTest {
         mock<ContentResolver> {
             whenever(loadThumbnail(any(), any(), anyOrNull())).thenReturn(bitmap)
         }
-    private val lifecycleOwner = TestLifecycleOwner()
-    private val dispatcher = UnconfinedTestDispatcher()
-    private lateinit var testSubject: ImagePreviewImageLoader
-
-    @Before
-    fun setup() {
-        Dispatchers.setMain(dispatcher)
-        lifecycleOwner.handleLifecycleEvent(Lifecycle.Event.ON_CREATE)
-        // create test subject after we've updated the lifecycle dispatcher
-        testSubject =
-            ImagePreviewImageLoader(
-                lifecycleOwner.lifecycle.coroutineScope + dispatcher,
-                imageSize.width,
-                contentResolver,
-                cacheSize = 1,
-            )
-    }
-
-    @After
-    fun cleanup() {
-        lifecycleOwner.handleLifecycleEvent(Lifecycle.Event.ON_DESTROY)
-        Dispatchers.resetMain()
-    }
+    private val scheduler = TestCoroutineScheduler()
+    private val dispatcher = UnconfinedTestDispatcher(scheduler)
+    private val scope = TestScope(dispatcher)
+    private val testSubject =
+        ImagePreviewImageLoader(
+            dispatcher,
+            imageSize.width,
+            contentResolver,
+            cacheSize = 1,
+        )
 
     @Test
-    fun prePopulate_cachesImagesUpToTheCacheSize() = runTest {
-        testSubject.prePopulate(listOf(uriOne, uriTwo))
+    fun prePopulate_cachesImagesUpToTheCacheSize() =
+        scope.runTest {
+            testSubject.prePopulate(listOf(uriOne, uriTwo))
 
-        verify(contentResolver, times(1)).loadThumbnail(uriOne, imageSize, null)
-        verify(contentResolver, never()).loadThumbnail(uriTwo, imageSize, null)
+            verify(contentResolver, times(1)).loadThumbnail(uriOne, imageSize, null)
+            verify(contentResolver, never()).loadThumbnail(uriTwo, imageSize, null)
 
-        testSubject(uriOne)
-        verify(contentResolver, times(1)).loadThumbnail(uriOne, imageSize, null)
-    }
+            testSubject(uriOne)
+            verify(contentResolver, times(1)).loadThumbnail(uriOne, imageSize, null)
+        }
 
     @Test
-    fun invoke_returnCachedImageWhenCalledTwice() = runTest {
-        testSubject(uriOne)
-        testSubject(uriOne)
+    fun invoke_returnCachedImageWhenCalledTwice() =
+        scope.runTest {
+            testSubject(uriOne)
+            testSubject(uriOne)
 
-        verify(contentResolver, times(1)).loadThumbnail(any(), any(), anyOrNull())
-    }
+            verify(contentResolver, times(1)).loadThumbnail(any(), any(), anyOrNull())
+        }
 
     @Test
-    fun invoke_whenInstructed_doesNotCache() = runTest {
-        testSubject(uriOne, false)
-        testSubject(uriOne, false)
+    fun invoke_whenInstructed_doesNotCache() =
+        scope.runTest {
+            testSubject(uriOne, false)
+            testSubject(uriOne, false)
 
-        verify(contentResolver, times(2)).loadThumbnail(any(), any(), anyOrNull())
-    }
+            verify(contentResolver, times(2)).loadThumbnail(any(), any(), anyOrNull())
+        }
 
     @Test
-    fun invoke_overlappedRequests_Deduplicate() = runTest {
-        val scheduler = TestCoroutineScheduler()
-        val dispatcher = StandardTestDispatcher(scheduler)
-        val testSubject =
-            ImagePreviewImageLoader(
-                lifecycleOwner.lifecycle.coroutineScope + dispatcher,
-                imageSize.width,
-                contentResolver,
-                cacheSize = 1,
-            )
-        coroutineScope {
-            launch(start = UNDISPATCHED) { testSubject(uriOne, false) }
-            launch(start = UNDISPATCHED) { testSubject(uriOne, false) }
-            scheduler.advanceUntilIdle()
-        }
+    fun invoke_overlappedRequests_Deduplicate() =
+        scope.runTest {
+            val dispatcher = StandardTestDispatcher(scheduler)
+            val testSubject =
+                ImagePreviewImageLoader(
+                    dispatcher,
+                    imageSize.width,
+                    contentResolver,
+                    cacheSize = 1,
+                )
+            coroutineScope {
+                launch(start = UNDISPATCHED) { testSubject(uriOne, false) }
+                launch(start = UNDISPATCHED) { testSubject(uriOne, false) }
+                scheduler.advanceUntilIdle()
+            }
 
-        verify(contentResolver, times(1)).loadThumbnail(any(), any(), anyOrNull())
-    }
+            verify(contentResolver, times(1)).loadThumbnail(any(), any(), anyOrNull())
+        }
 
     @Test
-    fun invoke_oldRecordsEvictedFromTheCache() = runTest {
-        testSubject(uriOne)
-        testSubject(uriTwo)
-        testSubject(uriTwo)
-        testSubject(uriOne)
-
-        verify(contentResolver, times(2)).loadThumbnail(uriOne, imageSize, null)
-        verify(contentResolver, times(1)).loadThumbnail(uriTwo, imageSize, null)
-    }
+    fun invoke_oldRecordsEvictedFromTheCache() =
+        scope.runTest {
+            testSubject(uriOne)
+            testSubject(uriTwo)
+            testSubject(uriTwo)
+            testSubject(uriOne)
+
+            verify(contentResolver, times(2)).loadThumbnail(uriOne, imageSize, null)
+            verify(contentResolver, times(1)).loadThumbnail(uriTwo, imageSize, null)
+        }
 
     @Test
-    fun invoke_doNotCacheNulls() = runTest {
-        whenever(contentResolver.loadThumbnail(any(), any(), anyOrNull())).thenReturn(null)
-        testSubject(uriOne)
-        testSubject(uriOne)
+    fun invoke_doNotCacheNulls() =
+        scope.runTest {
+            whenever(contentResolver.loadThumbnail(any(), any(), anyOrNull())).thenReturn(null)
+            testSubject(uriOne)
+            testSubject(uriOne)
 
-        verify(contentResolver, times(2)).loadThumbnail(uriOne, imageSize, null)
-    }
+            verify(contentResolver, times(2)).loadThumbnail(uriOne, imageSize, null)
+        }
 
     @Test(expected = CancellationException::class)
-    fun invoke_onClosedImageLoaderScope_throwsCancellationException() = runTest {
-        lifecycleOwner.handleLifecycleEvent(Lifecycle.Event.ON_DESTROY)
-        testSubject(uriOne)
-    }
+    fun invoke_onClosedImageLoaderScope_throwsCancellationException() =
+        scope.runTest {
+            val imageLoaderScope = CoroutineScope(coroutineContext)
+            val testSubject =
+                ImagePreviewImageLoader(
+                    imageLoaderScope,
+                    imageSize.width,
+                    contentResolver,
+                    cacheSize = 1,
+                )
+            imageLoaderScope.cancel()
+            testSubject(uriOne)
+        }
 
     @Test(expected = CancellationException::class)
-    fun invoke_imageLoaderScopeClosedMidflight_throwsCancellationException() = runTest {
-        val scheduler = TestCoroutineScheduler()
-        val dispatcher = StandardTestDispatcher(scheduler)
-        val testSubject =
-            ImagePreviewImageLoader(
-                lifecycleOwner.lifecycle.coroutineScope + dispatcher,
-                imageSize.width,
-                contentResolver,
-                cacheSize = 1,
-            )
-        coroutineScope {
-            val deferred = async(start = UNDISPATCHED) { testSubject(uriOne, false) }
-            lifecycleOwner.handleLifecycleEvent(Lifecycle.Event.ON_DESTROY)
-            scheduler.advanceUntilIdle()
-            deferred.await()
+    fun invoke_imageLoaderScopeClosedMidflight_throwsCancellationException() =
+        scope.runTest {
+            val dispatcher = StandardTestDispatcher(scheduler)
+            val imageLoaderScope = CoroutineScope(coroutineContext + dispatcher)
+            val testSubject =
+                ImagePreviewImageLoader(
+                    imageLoaderScope,
+                    imageSize.width,
+                    contentResolver,
+                    cacheSize = 1,
+                )
+            coroutineScope {
+                val deferred = async(start = UNDISPATCHED) { testSubject(uriOne, false) }
+                imageLoaderScope.cancel()
+                scheduler.advanceUntilIdle()
+                deferred.await()
+            }
         }
-    }
 
     @Test
-    fun invoke_multipleCallsWithDifferentCacheInstructions_cachingPrevails() = runTest {
-        val scheduler = TestCoroutineScheduler()
-        val dispatcher = StandardTestDispatcher(scheduler)
-        val testSubject =
-            ImagePreviewImageLoader(
-                lifecycleOwner.lifecycle.coroutineScope + dispatcher,
-                imageSize.width,
-                contentResolver,
-                cacheSize = 1,
-            )
-        coroutineScope {
-            launch(start = UNDISPATCHED) { testSubject(uriOne, false) }
-            launch(start = UNDISPATCHED) { testSubject(uriOne, true) }
-            scheduler.advanceUntilIdle()
-        }
-        testSubject(uriOne, true)
+    fun invoke_multipleCallsWithDifferentCacheInstructions_cachingPrevails() =
+        scope.runTest {
+            val dispatcher = StandardTestDispatcher(scheduler)
+            val imageLoaderScope = CoroutineScope(coroutineContext + dispatcher)
+            val testSubject =
+                ImagePreviewImageLoader(
+                    imageLoaderScope,
+                    imageSize.width,
+                    contentResolver,
+                    cacheSize = 1,
+                )
+            coroutineScope {
+                launch(start = UNDISPATCHED) { testSubject(uriOne, false) }
+                launch(start = UNDISPATCHED) { testSubject(uriOne, true) }
+                scheduler.advanceUntilIdle()
+            }
+            testSubject(uriOne, true)
 
-        verify(contentResolver, times(1)).loadThumbnail(uriOne, imageSize, null)
-    }
+            verify(contentResolver, times(1)).loadThumbnail(uriOne, imageSize, null)
+        }
 
     @Test
-    fun invoke_semaphoreGuardsContentResolverCalls() = runTest {
-        val contentResolver =
-            mock<ContentResolver> {
-                whenever(loadThumbnail(any(), any(), anyOrNull()))
-                    .thenThrow(SecurityException("test"))
-            }
-        val acquireCount = AtomicInteger()
-        val releaseCount = AtomicInteger()
-        val testSemaphore =
-            object : Semaphore {
-                override val availablePermits: Int
-                    get() = error("Unexpected invocation")
-
-                override suspend fun acquire() {
-                    acquireCount.getAndIncrement()
+    fun invoke_semaphoreGuardsContentResolverCalls() =
+        scope.runTest {
+            val contentResolver =
+                mock<ContentResolver> {
+                    whenever(loadThumbnail(any(), any(), anyOrNull()))
+                        .thenThrow(SecurityException("test"))
                 }
-
-                override fun tryAcquire(): Boolean {
-                    error("Unexpected invocation")
+            val acquireCount = AtomicInteger()
+            val releaseCount = AtomicInteger()
+            val testSemaphore =
+                object : Semaphore {
+                    override val availablePermits: Int
+                        get() = error("Unexpected invocation")
+
+                    override suspend fun acquire() {
+                        acquireCount.getAndIncrement()
+                    }
+
+                    override fun tryAcquire(): Boolean {
+                        error("Unexpected invocation")
+                    }
+
+                    override fun release() {
+                        releaseCount.getAndIncrement()
+                    }
                 }
 
-                override fun release() {
-                    releaseCount.getAndIncrement()
-                }
-            }
-
-        val testSubject =
-            ImagePreviewImageLoader(
-                lifecycleOwner.lifecycle.coroutineScope + dispatcher,
-                imageSize.width,
-                contentResolver,
-                cacheSize = 1,
-                testSemaphore,
-            )
-        testSubject(uriOne, false)
-
-        verify(contentResolver, times(1)).loadThumbnail(uriOne, imageSize, null)
-        assertThat(acquireCount.get()).isEqualTo(1)
-        assertThat(releaseCount.get()).isEqualTo(1)
-    }
+            val testSubject =
+                ImagePreviewImageLoader(
+                    CoroutineScope(coroutineContext + dispatcher),
+                    imageSize.width,
+                    contentResolver,
+                    cacheSize = 1,
+                    testSemaphore,
+                )
+            testSubject(uriOne, false)
+
+            verify(contentResolver, times(1)).loadThumbnail(uriOne, imageSize, null)
+            assertThat(acquireCount.get()).isEqualTo(1)
+            assertThat(releaseCount.get()).isEqualTo(1)
+        }
 
     @Test
-    fun invoke_semaphoreIsReleasedAfterContentResolverFailure() = runTest {
-        val semaphoreDeferred = CompletableDeferred<Unit>()
-        val releaseCount = AtomicInteger()
-        val testSemaphore =
-            object : Semaphore {
-                override val availablePermits: Int
-                    get() = error("Unexpected invocation")
-
-                override suspend fun acquire() {
-                    semaphoreDeferred.await()
-                }
-
-                override fun tryAcquire(): Boolean {
-                    error("Unexpected invocation")
+    fun invoke_semaphoreIsReleasedAfterContentResolverFailure() =
+        scope.runTest {
+            val semaphoreDeferred = CompletableDeferred<Unit>()
+            val releaseCount = AtomicInteger()
+            val testSemaphore =
+                object : Semaphore {
+                    override val availablePermits: Int
+                        get() = error("Unexpected invocation")
+
+                    override suspend fun acquire() {
+                        semaphoreDeferred.await()
+                    }
+
+                    override fun tryAcquire(): Boolean {
+                        error("Unexpected invocation")
+                    }
+
+                    override fun release() {
+                        releaseCount.getAndIncrement()
+                    }
                 }
 
-                override fun release() {
-                    releaseCount.getAndIncrement()
-                }
-            }
-
-        val testSubject =
-            ImagePreviewImageLoader(
-                lifecycleOwner.lifecycle.coroutineScope + dispatcher,
-                imageSize.width,
-                contentResolver,
-                cacheSize = 1,
-                testSemaphore,
-            )
-        launch(start = UNDISPATCHED) { testSubject(uriOne, false) }
+            val testSubject =
+                ImagePreviewImageLoader(
+                    CoroutineScope(coroutineContext + dispatcher),
+                    imageSize.width,
+                    contentResolver,
+                    cacheSize = 1,
+                    testSemaphore,
+                )
+            launch(start = UNDISPATCHED) { testSubject(uriOne, false) }
 
-        verify(contentResolver, never()).loadThumbnail(any(), any(), anyOrNull())
+            verify(contentResolver, never()).loadThumbnail(any(), any(), anyOrNull())
 
-        semaphoreDeferred.complete(Unit)
+            semaphoreDeferred.complete(Unit)
 
-        verify(contentResolver, times(1)).loadThumbnail(uriOne, imageSize, null)
-        assertThat(releaseCount.get()).isEqualTo(1)
-    }
+            verify(contentResolver, times(1)).loadThumbnail(uriOne, imageSize, null)
+            assertThat(releaseCount.get()).isEqualTo(1)
+        }
 
     @Test
-    fun invoke_multipleSimultaneousCalls_limitOnNumberOfSimultaneousOutgoingCallsIsRespected() {
-        val requestCount = 4
-        val thumbnailCallsCdl = CountDownLatch(requestCount)
-        val pendingThumbnailCalls = ArrayDeque<CountDownLatch>()
-        val contentResolver =
-            mock<ContentResolver> {
-                whenever(loadThumbnail(any(), any(), anyOrNull())).thenAnswer {
-                    val latch = CountDownLatch(1)
-                    synchronized(pendingThumbnailCalls) { pendingThumbnailCalls.offer(latch) }
-                    thumbnailCallsCdl.countDown()
-                    assertTrue("Timeout waiting thumbnail calls", latch.await(1, SECONDS))
-                    bitmap
+    fun invoke_multipleSimultaneousCalls_limitOnNumberOfSimultaneousOutgoingCallsIsRespected() =
+        scope.runTest {
+            val requestCount = 4
+            val thumbnailCallsCdl = CountDownLatch(requestCount)
+            val pendingThumbnailCalls = ArrayDeque<CountDownLatch>()
+            val contentResolver =
+                mock<ContentResolver> {
+                    whenever(loadThumbnail(any(), any(), anyOrNull())).thenAnswer {
+                        val latch = CountDownLatch(1)
+                        synchronized(pendingThumbnailCalls) { pendingThumbnailCalls.offer(latch) }
+                        thumbnailCallsCdl.countDown()
+                        assertTrue("Timeout waiting thumbnail calls", latch.await(1, SECONDS))
+                        bitmap
+                    }
                 }
-            }
-        val name = "LoadImage"
-        val maxSimultaneousRequests = 2
-        val threadsStartedCdl = CountDownLatch(requestCount)
-        val dispatcher = NewThreadDispatcher(name) { threadsStartedCdl.countDown() }
-        val testSubject =
-            ImagePreviewImageLoader(
-                lifecycleOwner.lifecycle.coroutineScope + dispatcher + CoroutineName(name),
-                imageSize.width,
-                contentResolver,
-                cacheSize = 1,
-                maxSimultaneousRequests,
-            )
-        runTest {
-            repeat(requestCount) {
-                launch { testSubject(Uri.parse("content://org.pkg.app/image-$it.png")) }
-            }
-            yield()
-            // wait for all requests to be dispatched
-            assertThat(threadsStartedCdl.await(5, SECONDS)).isTrue()
+            val name = "LoadImage"
+            val maxSimultaneousRequests = 2
+            val threadsStartedCdl = CountDownLatch(requestCount)
+            val dispatcher = NewThreadDispatcher(name) { threadsStartedCdl.countDown() }
+            val testSubject =
+                ImagePreviewImageLoader(
+                    CoroutineScope(coroutineContext + dispatcher + CoroutineName(name)),
+                    imageSize.width,
+                    contentResolver,
+                    cacheSize = 1,
+                    maxSimultaneousRequests,
+                )
+            coroutineScope {
+                repeat(requestCount) {
+                    launch { testSubject(Uri.parse("content://org.pkg.app/image-$it.png")) }
+                }
+                yield()
+                // wait for all requests to be dispatched
+                assertThat(threadsStartedCdl.await(5, SECONDS)).isTrue()
 
-            assertThat(thumbnailCallsCdl.await(100, MILLISECONDS)).isFalse()
-            synchronized(pendingThumbnailCalls) {
-                assertThat(pendingThumbnailCalls.size).isEqualTo(maxSimultaneousRequests)
-            }
+                assertThat(thumbnailCallsCdl.await(100, MILLISECONDS)).isFalse()
+                synchronized(pendingThumbnailCalls) {
+                    assertThat(pendingThumbnailCalls.size).isEqualTo(maxSimultaneousRequests)
+                }
 
-            pendingThumbnailCalls.poll()?.countDown()
-            assertThat(thumbnailCallsCdl.await(100, MILLISECONDS)).isFalse()
-            synchronized(pendingThumbnailCalls) {
-                assertThat(pendingThumbnailCalls.size).isEqualTo(maxSimultaneousRequests)
-            }
+                pendingThumbnailCalls.poll()?.countDown()
+                assertThat(thumbnailCallsCdl.await(100, MILLISECONDS)).isFalse()
+                synchronized(pendingThumbnailCalls) {
+                    assertThat(pendingThumbnailCalls.size).isEqualTo(maxSimultaneousRequests)
+                }
 
-            pendingThumbnailCalls.poll()?.countDown()
-            assertThat(thumbnailCallsCdl.await(100, MILLISECONDS)).isTrue()
-            synchronized(pendingThumbnailCalls) {
-                assertThat(pendingThumbnailCalls.size).isEqualTo(maxSimultaneousRequests)
-            }
-            for (cdl in pendingThumbnailCalls) {
-                cdl.countDown()
+                pendingThumbnailCalls.poll()?.countDown()
+                assertThat(thumbnailCallsCdl.await(100, MILLISECONDS)).isTrue()
+                synchronized(pendingThumbnailCalls) {
+                    assertThat(pendingThumbnailCalls.size).isEqualTo(maxSimultaneousRequests)
+                }
+                for (cdl in pendingThumbnailCalls) {
+                    cdl.countDown()
+                }
             }
         }
-    }
 }
 
 private class NewThreadDispatcher(