Merge "Make fetch thumbnails call on bg thread to reduce thread switching" into main
diff --git a/quickstep/src/com/android/quickstep/TaskIconCache.kt b/quickstep/src/com/android/quickstep/TaskIconCache.kt
index bb0a304..bf94d41 100644
--- a/quickstep/src/com/android/quickstep/TaskIconCache.kt
+++ b/quickstep/src/com/android/quickstep/TaskIconCache.kt
@@ -25,6 +25,7 @@
 import android.util.SparseArray
 import androidx.annotation.WorkerThread
 import com.android.launcher3.Flags.enableOverviewIconMenu
+import com.android.launcher3.Flags.enableRefactorTaskThumbnail
 import com.android.launcher3.R
 import com.android.launcher3.Utilities
 import com.android.launcher3.icons.BaseIconFactory
@@ -63,15 +64,8 @@
     @get:WorkerThread
     private val iconFactory: BaseIconFactory
         get() =
-            _iconFactory
-                ?: BaseIconFactory(
-                        context,
-                        DisplayController.INSTANCE[context].info.densityDpi,
-                        context.resources.getDimensionPixelSize(
-                            R.dimen.task_icon_cache_default_icon_size
-                        ),
-                    )
-                    .also { _iconFactory = it }
+            if (enableRefactorTaskThumbnail()) createIconFactory()
+            else _iconFactory ?: createIconFactory().also { _iconFactory = it }
 
     var taskVisualsChangeListener: TaskVisualsChangeListener? = null
 
@@ -85,6 +79,22 @@
         }
     }
 
+    // TODO(b/387496731): Add ensureActive() calls if they show performance benefit
+    override suspend fun getIcon(task: Task): TaskCacheEntry {
+        task.icon?.let {
+            // Nothing to load, the icon is already loaded
+            return TaskCacheEntry(it, task.titleDescription ?: "", task.title)
+        }
+
+        val entry = getCacheEntry(task)
+        task.icon = entry.icon
+        task.titleDescription = entry.contentDescription
+        task.title = entry.title
+
+        dispatchIconUpdate(task.key.id)
+        return entry
+    }
+
     /**
      * Asynchronously fetches the icon and other task data.
      *
@@ -92,14 +102,11 @@
      * @param callback The callback to receive the task after its data has been populated.
      * @return A cancelable handle to the request
      */
-    override fun getIconInBackground(
-        task: Task,
-        callback: GetTaskIconCallback,
-    ): CancellableTask<*>? {
+    fun getIconInBackground(task: Task, callback: GetTaskIconCallback): CancellableTask<*>? {
         Preconditions.assertUIThread()
-        if (task.icon != null) {
+        task.icon?.let {
             // Nothing to load, the icon is already loaded
-            callback.onTaskIconReceived(task.icon, task.titleDescription ?: "", task.title ?: "")
+            callback.onTaskIconReceived(it, task.titleDescription ?: "", task.title ?: "")
             return null
         }
         val request =
@@ -141,10 +148,17 @@
     }
 
     @WorkerThread
+    private fun createIconFactory() =
+        BaseIconFactory(
+            context,
+            DisplayController.INSTANCE.get(context).info.densityDpi,
+            context.resources.getDimensionPixelSize(R.dimen.task_icon_cache_default_icon_size),
+        )
+
+    @WorkerThread
     private fun getCacheEntry(task: Task): TaskCacheEntry {
-        var entry = iconCache.getAndInvalidateIfModified(task.key)
-        if (entry != null) {
-            return entry
+        iconCache.getAndInvalidateIfModified(task.key)?.let {
+            return it
         }
 
         val desc = task.taskDescription
@@ -152,11 +166,10 @@
         var activityInfo: ActivityInfo? = null
 
         // Create new cache entry
-        entry = TaskCacheEntry()
 
         // Load icon
         val icon = getIcon(desc, key.userId)
-        entry.icon =
+        val entryIcon =
             if (icon != null) {
                 getBitmapInfo(
                         BitmapDrawable(context.resources, icon),
@@ -182,21 +195,29 @@
                 }
             }
 
-        // Skip loading the content description if the activity no longer exists
         activityInfo =
             activityInfo
                 ?: PackageManagerWrapper.getInstance().getActivityInfo(key.component, key.userId)
 
-        if (activityInfo != null) {
-            entry.contentDescription =
-                getBadgedContentDescription(activityInfo, task.key.userId, task.taskDescription)
-            if (enableOverviewIconMenu()) {
-                entry.title = Utilities.trim(activityInfo.loadLabel(context.packageManager))
-            }
-        }
-
-        iconCache.put(task.key, entry)
-        return entry
+        return when {
+            // Skip loading the content description if the activity no longer exists
+            activityInfo == null -> TaskCacheEntry(entryIcon)
+            enableOverviewIconMenu() ->
+                TaskCacheEntry(
+                    entryIcon,
+                    getBadgedContentDescription(
+                        activityInfo,
+                        task.key.userId,
+                        task.taskDescription,
+                    ),
+                    Utilities.trim(activityInfo.loadLabel(context.packageManager)),
+                )
+            else ->
+                TaskCacheEntry(
+                    entryIcon,
+                    getBadgedContentDescription(activityInfo, task.key.userId, task.taskDescription),
+                )
+        }.also { iconCache.put(task.key, it) }
     }
 
     private fun getIcon(desc: ActivityManager.TaskDescription, userId: Int): Bitmap? =
@@ -272,16 +293,16 @@
         iconCache.evictAll()
     }
 
-    private data class TaskCacheEntry(
-        var icon: Drawable? = null,
-        var contentDescription: String = "",
-        var title: String = "",
+    data class TaskCacheEntry(
+        val icon: Drawable,
+        val contentDescription: String = "",
+        val title: String = "",
     )
 
     /** Callback used when retrieving app icons from cache. */
     fun interface GetTaskIconCallback {
         /** Called when task icon is retrieved. */
-        fun onTaskIconReceived(icon: Drawable?, contentDescription: String, title: String)
+        fun onTaskIconReceived(icon: Drawable, contentDescription: String, title: String)
     }
 
     fun registerTaskVisualsChangeListener(newListener: TaskVisualsChangeListener?) {
diff --git a/quickstep/src/com/android/quickstep/TaskThumbnailCache.kt b/quickstep/src/com/android/quickstep/TaskThumbnailCache.kt
index 7de4481..7b56213 100644
--- a/quickstep/src/com/android/quickstep/TaskThumbnailCache.kt
+++ b/quickstep/src/com/android/quickstep/TaskThumbnailCache.kt
@@ -17,6 +17,7 @@
 
 import android.content.Context
 import androidx.annotation.VisibleForTesting
+import androidx.annotation.WorkerThread
 import com.android.launcher3.Flags.enableGridOnlyOverview
 import com.android.launcher3.R
 import com.android.launcher3.util.CancellableTask
@@ -78,14 +79,61 @@
         cache.updateIfAlreadyInCache(taskId, thumbnail)
     }
 
+    // TODO(b/387496731): Add ensureActive() calls if they show performance benefit
+    /**
+     * Retrieves a thumbnail for the provided `task` on the current thread. This should not be
+     * called from the main thread.
+     */
+    @WorkerThread
+    override suspend fun getThumbnail(task: Task): ThumbnailData? {
+        val lowResolution: Boolean = !highResLoadingState.isEnabled
+        // Check task for thumbnail
+        val taskThumbnail: ThumbnailData? = task.thumbnail
+        if (
+            taskThumbnail?.thumbnail != null && (!taskThumbnail.reducedResolution || lowResolution)
+        ) {
+            return taskThumbnail
+        }
+
+        // Check cache for thumbnail
+        val cachedThumbnail: ThumbnailData? = cache.getAndInvalidateIfModified(task.key)
+        if (
+            cachedThumbnail?.thumbnail != null &&
+                (!cachedThumbnail.reducedResolution || lowResolution)
+        ) {
+            return cachedThumbnail
+        }
+
+        // Get thumbnail from system
+        var thumbnailData =
+            ActivityManagerWrapper.getInstance().getTaskThumbnail(task.key.id, lowResolution)
+        if (thumbnailData.thumbnail == null) {
+            thumbnailData = ActivityManagerWrapper.getInstance().takeTaskThumbnail(task.key.id)
+        }
+
+        // Avoid an async timing issue that a low res entry replaces an existing high
+        // res entry in high res enabled state, so we check before putting it to cache
+        if (
+            enableGridOnlyOverview() &&
+                thumbnailData.reducedResolution &&
+                highResLoadingState.isEnabled
+        ) {
+            val newCachedThumbnail = cache.getAndInvalidateIfModified(task.key)
+            if (newCachedThumbnail.thumbnail != null && !newCachedThumbnail.reducedResolution) {
+                return newCachedThumbnail
+            }
+        }
+        cache.put(task.key, thumbnailData)
+        return thumbnailData
+    }
+
     /**
      * Asynchronously fetches the thumbnail for the given `task`.
      *
      * @param callback The callback to receive the task after its data has been populated.
-     *
      * @return a cancelable handle to the request
      */
-    override fun getThumbnailInBackground(
+    fun getThumbnailInBackground(
         task: Task,
         callback: Consumer<ThumbnailData>,
     ): CancellableTask<ThumbnailData>? {
diff --git a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
index f950f47..703d631 100644
--- a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
@@ -17,6 +17,7 @@
 package com.android.quickstep.recents.data
 
 import android.graphics.drawable.Drawable
+import android.graphics.drawable.ShapeDrawable
 import android.util.Log
 import com.android.launcher3.util.coroutines.DispatcherProvider
 import com.android.quickstep.recents.data.TaskVisualsChangedDelegate.TaskIconChangedCallback
@@ -25,16 +26,16 @@
 import com.android.quickstep.task.thumbnail.data.TaskThumbnailDataSource
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
-import kotlin.coroutines.resume
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
+import kotlinx.coroutines.async
+import kotlinx.coroutines.awaitAll
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.distinctUntilChangedBy
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.suspendCancellableCoroutine
 import kotlinx.coroutines.withContext
 
 class TasksRepository(
@@ -112,10 +113,11 @@
         taskRequests[taskId] =
             Pair(
                 task.key,
-                recentsCoroutineScope.launch(dispatcherProvider.main) {
+                recentsCoroutineScope.launch(dispatcherProvider.background) {
                     Log.i(TAG, "requestTaskData: $taskId")
-                    fetchIcon(task)
-                    fetchThumbnail(task)
+                    val thumbnailFetchDeferred = async { fetchThumbnail(task) }
+                    val iconFetchDeferred = async { fetchIcon(task) }
+                    awaitAll(thumbnailFetchDeferred, iconFetchDeferred)
                 },
             )
     }
@@ -150,7 +152,7 @@
             task.key,
             object : TaskIconChangedCallback {
                 override fun onTaskIconChanged() {
-                    recentsCoroutineScope.launch(dispatcherProvider.main) {
+                    recentsCoroutineScope.launch(dispatcherProvider.background) {
                         updateIcon(task.key.id, getIconFromDataSource(task))
                     }
                 }
@@ -168,7 +170,7 @@
                 }
 
                 override fun onHighResLoadingStateChanged() {
-                    recentsCoroutineScope.launch(dispatcherProvider.main) {
+                    recentsCoroutineScope.launch(dispatcherProvider.background) {
                         updateThumbnail(task.key.id, getThumbnailFromDataSource(task))
                     }
                 }
@@ -191,34 +193,18 @@
     }
 
     private suspend fun getThumbnailFromDataSource(task: Task) =
-        withContext(dispatcherProvider.main) {
-            suspendCancellableCoroutine { continuation ->
-                val cancellableTask =
-                    taskThumbnailDataSource.getThumbnailInBackground(task) {
-                        continuation.resume(it)
-                    }
-                continuation.invokeOnCancellation { cancellableTask?.cancel() }
-            }
-        }
+        withContext(dispatcherProvider.background) { taskThumbnailDataSource.getThumbnail(task) }
 
     private suspend fun getIconFromDataSource(task: Task) =
-        withContext(dispatcherProvider.main) {
-            suspendCancellableCoroutine { continuation ->
-                val cancellableTask =
-                    taskIconDataSource.getIconInBackground(task) { icon, contentDescription, title
-                        ->
-                        icon?.constantState?.let {
-                            continuation.resume(
-                                IconData(it.newDrawable().mutate(), contentDescription, title)
-                            )
-                        }
-                    }
-                continuation.invokeOnCancellation { cancellableTask?.cancel() }
-            }
+        withContext(dispatcherProvider.background) {
+            val iconCacheEntry = taskIconDataSource.getIcon(task)
+            val icon = iconCacheEntry.icon.constantState?.newDrawable()?.mutate() ?: EMPTY_DRAWABLE
+            IconData(icon, iconCacheEntry.contentDescription, iconCacheEntry.title)
         }
 
     companion object {
         private const val TAG = "TasksRepository"
+        private val EMPTY_DRAWABLE = ShapeDrawable()
     }
 
     /** Helper class to support StateFlow emissions when using a Map with a MutableStateFlow. */
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskIconDataSource.kt b/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskIconDataSource.kt
index ab699c6..c45458c 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskIconDataSource.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskIconDataSource.kt
@@ -16,10 +16,9 @@
 
 package com.android.quickstep.task.thumbnail.data
 
-import com.android.launcher3.util.CancellableTask
-import com.android.quickstep.TaskIconCache.GetTaskIconCallback
+import com.android.quickstep.TaskIconCache
 import com.android.systemui.shared.recents.model.Task
 
 interface TaskIconDataSource {
-    fun getIconInBackground(task: Task, callback: GetTaskIconCallback): CancellableTask<*>?
+    suspend fun getIcon(task: Task): TaskIconCache.TaskCacheEntry
 }
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskThumbnailDataSource.kt b/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskThumbnailDataSource.kt
index 986acbe..6e63ea9 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskThumbnailDataSource.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskThumbnailDataSource.kt
@@ -16,14 +16,9 @@
 
 package com.android.quickstep.task.thumbnail.data
 
-import com.android.launcher3.util.CancellableTask
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
-import java.util.function.Consumer
 
 interface TaskThumbnailDataSource {
-    fun getThumbnailInBackground(
-        task: Task,
-        callback: Consumer<ThumbnailData>
-    ): CancellableTask<ThumbnailData>?
+    suspend fun getThumbnail(task: Task): ThumbnailData?
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskIconDataSource.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskIconDataSource.kt
index 5de876a..f6f158f 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskIconDataSource.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskIconDataSource.kt
@@ -17,11 +17,11 @@
 package com.android.quickstep.recents.data
 
 import android.graphics.drawable.Drawable
-import com.android.launcher3.util.CancellableTask
-import com.android.quickstep.TaskIconCache
+import com.android.quickstep.TaskIconCache.TaskCacheEntry
 import com.android.quickstep.task.thumbnail.data.TaskIconDataSource
 import com.android.systemui.shared.recents.model.Task
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.yield
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.whenever
 
@@ -29,28 +29,30 @@
 
     val taskIdToDrawable: MutableMap<Int, Drawable> =
         (0..10).associateWith { mockCopyableDrawable() }.toMutableMap()
-
-    val taskIdToUpdatingTask: MutableMap<Int, () -> Unit> = mutableMapOf()
-    var shouldLoadSynchronously: Boolean = true
+    private val completionPrevented: MutableSet<Int> = mutableSetOf()
 
     /** Retrieves and sets an icon on [task] from [taskIdToDrawable]. */
-    override fun getIconInBackground(
-        task: Task,
-        callback: TaskIconCache.GetTaskIconCallback
-    ): CancellableTask<*>? {
-        val wrappedCallback = {
-            callback.onTaskIconReceived(
-                taskIdToDrawable.getValue(task.key.id),
-                "content desc ${task.key.id}",
-                "title ${task.key.id}"
-            )
+    override suspend fun getIcon(task: Task): TaskCacheEntry {
+        while (task.key.id in completionPrevented) {
+            yield()
         }
-        if (shouldLoadSynchronously) {
-            wrappedCallback()
-        } else {
-            taskIdToUpdatingTask[task.key.id] = wrappedCallback
-        }
-        return null
+        return TaskCacheEntry(
+            taskIdToDrawable.getValue(task.key.id),
+            "content desc ${task.key.id}",
+            "title ${task.key.id}",
+        )
+    }
+
+    fun preventIconLoad(taskId: Int) {
+        completionPrevented.add(taskId)
+    }
+
+    fun completeLoadingForTask(taskId: Int) {
+        completionPrevented.remove(taskId)
+    }
+
+    fun completeLoading() {
+        completionPrevented.clear()
     }
 
     companion object {
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskThumbnailDataSource.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskThumbnailDataSource.kt
index d12c0b0..e10afc4 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskThumbnailDataSource.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskThumbnailDataSource.kt
@@ -17,11 +17,10 @@
 package com.android.quickstep.recents.data
 
 import android.graphics.Bitmap
-import com.android.launcher3.util.CancellableTask
 import com.android.quickstep.task.thumbnail.data.TaskThumbnailDataSource
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
-import java.util.function.Consumer
+import kotlinx.coroutines.yield
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.whenever
 
@@ -29,25 +28,28 @@
 
     val taskIdToBitmap: MutableMap<Int, Bitmap> =
         (0..10).associateWith { mock<Bitmap>() }.toMutableMap()
-    val taskIdToUpdatingTask: MutableMap<Int, () -> Unit> = mutableMapOf()
-    var shouldLoadSynchronously: Boolean = true
+    private val completionPrevented: MutableSet<Int> = mutableSetOf()
+    private val getThumbnailCalls = mutableMapOf<Int, Int>()
 
     /** Retrieves and sets a thumbnail on [task] from [taskIdToBitmap]. */
-    override fun getThumbnailInBackground(
-        task: Task,
-        callback: Consumer<ThumbnailData>
-    ): CancellableTask<ThumbnailData>? {
-        val thumbnailData = mock<ThumbnailData>()
-        whenever(thumbnailData.thumbnail).thenReturn(taskIdToBitmap[task.key.id])
-        val wrappedCallback = {
-            task.thumbnail = thumbnailData
-            callback.accept(thumbnailData)
+    override suspend fun getThumbnail(task: Task): ThumbnailData {
+        getThumbnailCalls[task.key.id] = (getThumbnailCalls[task.key.id] ?: 0) + 1
+
+        while (task.key.id in completionPrevented) {
+            yield()
         }
-        if (shouldLoadSynchronously) {
-            wrappedCallback()
-        } else {
-            taskIdToUpdatingTask[task.key.id] = wrappedCallback
+        return mock<ThumbnailData>().also {
+            whenever(it.thumbnail).thenReturn(taskIdToBitmap[task.key.id])
         }
-        return null
+    }
+
+    fun getNumberOfGetThumbnailCalls(taskId: Int): Int = getThumbnailCalls[taskId] ?: 0
+
+    fun preventThumbnailLoad(taskId: Int) {
+        completionPrevented.add(taskId)
+    }
+
+    fun completeLoadingForTask(taskId: Int) {
+        completionPrevented.remove(taskId)
     }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
index ee1ec6e..b6cf5bd 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
@@ -214,7 +214,7 @@
                 .isEqualTo(bitmap2)
 
             // Prevent new loading of Bitmaps
-            taskThumbnailDataSource.shouldLoadSynchronously = false
+            taskThumbnailDataSource.preventThumbnailLoad(2)
             systemUnderTest.setVisibleTasks(setOf(2, 3))
 
             assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
@@ -235,7 +235,7 @@
                 .assertHasIconDataFromSource(taskIconDataSource)
 
             // Prevent new loading of Drawables
-            taskThumbnailDataSource.shouldLoadSynchronously = false
+            taskIconDataSource.preventIconLoad(2)
             systemUnderTest.setVisibleTasks(setOf(2, 3))
 
             systemUnderTest
@@ -257,9 +257,6 @@
             assertThat(task2.thumbnail!!.thumbnail).isEqualTo(bitmap2)
             task2.assertHasIconDataFromSource(taskIconDataSource)
 
-            // Prevent new loading of Bitmaps
-            taskThumbnailDataSource.shouldLoadSynchronously = false
-            taskIconDataSource.shouldLoadSynchronously = false
             systemUnderTest.setVisibleTasks(setOf(0, 1))
 
             val task2AfterVisibleTasksChanged = systemUnderTest.getTaskDataById(2).first()!!
@@ -275,21 +272,22 @@
             // Setup fakes
             recentsModel.seedTasks(defaultTaskList)
             val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
-            taskThumbnailDataSource.shouldLoadSynchronously = false
 
             // Setup TasksRepository
             systemUnderTest.getAllTaskData(forceRefresh = true)
-            systemUnderTest.setVisibleTasks(setOf(1, 2))
 
-            // Assert there is no bitmap in first emission
-            assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail).isNull()
+            val task2DataFlow = systemUnderTest.getTaskDataById(2)
+            val task2BitmapValues = mutableListOf<Bitmap?>()
+            testScope.backgroundScope.launch {
+                task2DataFlow.map { it?.thumbnail?.thumbnail }.toList(task2BitmapValues)
+            }
 
-            // Simulate bitmap loading after first emission
-            taskThumbnailDataSource.taskIdToUpdatingTask.getValue(2).invoke()
+            // Check for first emission
+            assertThat(task2BitmapValues.single()).isNull()
 
+            systemUnderTest.setVisibleTasks(setOf(2))
             // Check for second emission
-            assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
-                .isEqualTo(bitmap2)
+            assertThat(task2BitmapValues).isEqualTo(listOf(null, bitmap2))
         }
 
     @Test
@@ -365,7 +363,6 @@
         testScope.runTest {
             recentsModel.seedTasks(defaultTaskList)
             systemUnderTest.getAllTaskData(forceRefresh = true)
-            taskThumbnailDataSource.shouldLoadSynchronously = false
 
             val taskDataFlow = systemUnderTest.getTaskDataById(1)
             val task1IconValues = mutableListOf<Drawable?>()
@@ -374,14 +371,10 @@
             }
 
             systemUnderTest.setVisibleTasks(setOf(1))
-            val task1UpdatingTaskOld = taskThumbnailDataSource.taskIdToUpdatingTask[1]
-            println(task1UpdatingTaskOld)
+            assertThat(taskThumbnailDataSource.getNumberOfGetThumbnailCalls(1)).isEqualTo(1)
 
             systemUnderTest.setVisibleTasks(setOf(1, 2))
-            val task1UpdatingTaskNew = taskThumbnailDataSource.taskIdToUpdatingTask[1]
-            println(task1UpdatingTaskNew)
-
-            assertThat(task1UpdatingTaskNew).isEqualTo(task1UpdatingTaskOld)
+            assertThat(taskThumbnailDataSource.getNumberOfGetThumbnailCalls(1)).isEqualTo(1)
         }
 
     private fun createTaskWithId(taskId: Int) =