Support switch to screenshot for TTV

- Moved live tile decision to TaskThumbnailViewModel, and added RecentsViewData for runningTask information
- Implemented a ThumbnailData override mechanism to override screenshot from TaskThunmbnailCache, the override will be added during switchToScreenshot and cleared when user exit recents (RecentsView.reset)
- During the thumbnail override, we wait until the new ThumbnailData propagates, before finsihing Recents animation to avoid a frame of flash

Bug: 343364498
Test: TaskThumbnailViewModelTest, TaskRepositoryTest, RecentsViewModelTest
Flag: com.android.launcher3.enable_refactor_task_thumbnail
Change-Id: I776943ecdfa0d65d94b054692297b42686f59f8a
diff --git a/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt
index 9c4248c..d5aaed5 100644
--- a/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt
@@ -41,4 +41,10 @@
      * populated e.g. icons/thumbnails etc.
      */
     fun setVisibleTasks(visibleTaskIdList: List<Int>)
+
+    /**
+     * Override [ThumbnailData] with a map of taskId to [ThumbnailData]. The override only applies
+     * if the tasks are already visible, and will be invalidated when tasks become invisible.
+     */
+    fun setThumbnailOverride(thumbnailOverride: Map<Int, ThumbnailData>)
 }
diff --git a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
index 71be75b..0714170 100644
--- a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
@@ -45,15 +45,18 @@
     private val _taskData =
         groupedTaskData.map { groupTaskList -> groupTaskList.flatMap { it.tasks } }
     private val visibleTaskIds = MutableStateFlow(emptySet<Int>())
+    private val thumbnailOverride = MutableStateFlow(mapOf<Int, ThumbnailData>())
 
     private val taskData: Flow<List<Task>> =
-        combine(_taskData, getThumbnailQueryResults(), getIconQueryResults()) {
+        combine(_taskData, getThumbnailQueryResults(), getIconQueryResults(), thumbnailOverride) {
             tasks,
             thumbnailQueryResults,
-            iconQueryResults ->
+            iconQueryResults,
+            thumbnailOverride ->
             tasks.forEach { task ->
-                // Add retrieved thumbnails + remove unnecessary thumbnails
-                task.thumbnail = thumbnailQueryResults[task.key.id]
+                // Add retrieved thumbnails + remove unnecessary thumbnails (e.g. invisible)
+                task.thumbnail =
+                    thumbnailOverride[task.key.id] ?: thumbnailQueryResults[task.key.id]
 
                 // TODO(b/352331675) don't load icons for DesktopTaskView
                 // Add retrieved icons + remove unnecessary icons
@@ -79,6 +82,12 @@
 
     override fun setVisibleTasks(visibleTaskIdList: List<Int>) {
         this.visibleTaskIds.value = visibleTaskIdList.toSet()
+        setThumbnailOverride(thumbnailOverride.value)
+    }
+
+    override fun setThumbnailOverride(thumbnailOverride: Map<Int, ThumbnailData>) {
+        this.thumbnailOverride.value =
+            thumbnailOverride.filterKeys(this.visibleTaskIds.value::contains).toMap()
     }
 
     /** Flow wrapper for [TaskThumbnailDataSource.getThumbnailInBackground] api */
diff --git a/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewData.kt b/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewData.kt
index f5e0243..87446b0 100644
--- a/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewData.kt
+++ b/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewData.kt
@@ -36,4 +36,10 @@
     val tintAmount = MutableStateFlow(0f)
 
     val thumbnailSplashProgress = MutableStateFlow(0f)
+
+    // A list of taskIds that are associated with a RecentsAnimationController. */
+    val runningTaskIds = MutableStateFlow(emptySet<Int>())
+
+    // Whether we should use static screenshot instead of live tile for taskIds in [runningTaskIds]
+    val runningTaskShowScreenshot = MutableStateFlow(false)
 }
diff --git a/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewModel.kt b/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewModel.kt
index 6148d4b..7205fc8 100644
--- a/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewModel.kt
+++ b/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewModel.kt
@@ -17,6 +17,10 @@
 package com.android.quickstep.recents.viewmodel
 
 import com.android.quickstep.recents.data.RecentTasksRepository
+import com.android.systemui.shared.recents.model.ThumbnailData
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.filter
+import kotlinx.coroutines.flow.first
 
 class RecentsViewModel(
     private val recentsTasksRepository: RecentTasksRepository,
@@ -53,4 +57,35 @@
     fun updateThumbnailSplashProgress(taskThumbnailSplashAlpha: Float) {
         recentsViewData.thumbnailSplashProgress.value = taskThumbnailSplashAlpha
     }
+
+    fun setThumbnailOverride(thumbnailOverride: Map<Int, ThumbnailData>) {
+        recentsTasksRepository.setThumbnailOverride(thumbnailOverride)
+    }
+
+    suspend fun waitForThumbnailsToUpdate(updatedThumbnails: Map<Int, ThumbnailData>) {
+        combine(
+                updatedThumbnails.map {
+                    recentsTasksRepository.getThumbnailById(it.key).filter { thumbnailData ->
+                        thumbnailData == it.value
+                    }
+                }
+            ) {}
+            .first()
+    }
+
+    suspend fun waitForRunningTaskShowScreenshotToUpdate() {
+        recentsViewData.runningTaskShowScreenshot.filter { it }.first()
+    }
+
+    fun onReset() {
+        updateVisibleTasks(emptyList())
+    }
+
+    fun updateRunningTask(taskIds: Set<Int>) {
+        recentsViewData.runningTaskIds.value = taskIds
+    }
+
+    fun setRunningTaskShowScreenshot(showScreenshot: Boolean) {
+        recentsViewData.runningTaskShowScreenshot.value = showScreenshot
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailUiState.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailUiState.kt
index aa7d26c..a6be9f6 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailUiState.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailUiState.kt
@@ -45,5 +45,3 @@
         val size: Point,
     )
 }
-
-data class TaskThumbnail(val taskId: Int, val isRunning: Boolean)
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt
index de33919..4e29840 100644
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt
@@ -10,13 +10,14 @@
  * 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
+ * See the License for the specific language goveryning permissions and
  * limitations under the License.
  */
 
 package com.android.quickstep.task.viewmodel
 
 import android.annotation.ColorInt
+import android.app.ActivityTaskManager.INVALID_TASK_ID
 import android.graphics.Matrix
 import android.graphics.Point
 import androidx.core.graphics.ColorUtils
@@ -26,7 +27,6 @@
 import com.android.quickstep.recents.viewmodel.RecentsViewData
 import com.android.quickstep.task.thumbnail.GetSplashSizeUseCase
 import com.android.quickstep.task.thumbnail.SplashAlphaUseCase
-import com.android.quickstep.task.thumbnail.TaskThumbnail
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.BackgroundOnly
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
@@ -59,7 +59,7 @@
 ) {
     private val task = MutableStateFlow<Flow<Task?>>(flowOf(null))
     private val splashProgress = MutableStateFlow(flowOf(0f))
-    private lateinit var taskThumbnail: TaskThumbnail
+    private var taskId: Int = INVALID_TASK_ID
 
     /**
      * Progress for changes in corner radius. progress: 0 = overview corner radius; 1 = fullscreen
@@ -68,10 +68,12 @@
     val cornerRadiusProgress =
         if (taskViewData.isOutlineFormedByThumbnailView) recentsViewData.fullscreenProgress
         else MutableStateFlow(1f).asStateFlow()
+
     val inheritedScale =
         combine(recentsViewData.scale, taskViewData.scale) { recentsScale, taskScale ->
             recentsScale * taskScale
         }
+
     val dimProgress: Flow<Float> =
         combine(taskContainerData.taskMenuOpenProgress, recentsViewData.tintAmount) {
             taskMenuOpenProgress,
@@ -79,34 +81,42 @@
             max(taskMenuOpenProgress * MAX_SCRIM_ALPHA, tintAmount)
         }
     val splashAlpha = splashProgress.flatMapLatest { it }
+
+    private val isLiveTile =
+        combine(
+                task.flatMapLatest { it }.map { it?.key?.id }.distinctUntilChanged(),
+                recentsViewData.runningTaskIds,
+                recentsViewData.runningTaskShowScreenshot
+            ) { taskId, runningTaskIds, runningTaskShowScreenshot ->
+                runningTaskIds.contains(taskId) && !runningTaskShowScreenshot
+            }
+            .distinctUntilChanged()
+
     val uiState: Flow<TaskThumbnailUiState> =
-        task
-            .flatMapLatest { taskFlow ->
-                taskFlow.map { taskVal ->
-                    when {
-                        taskVal == null -> Uninitialized
-                        taskThumbnail.isRunning -> LiveTile
-                        isBackgroundOnly(taskVal) ->
-                            BackgroundOnly(taskVal.colorBackground.removeAlpha())
-                        isSnapshotSplashState(taskVal) ->
-                            SnapshotSplash(createSnapshotState(taskVal), createSplashState(taskVal))
-                        else -> Uninitialized
-                    }
+        combine(task.flatMapLatest { it }, isLiveTile) { taskVal, isRunning ->
+                when {
+                    taskVal == null -> Uninitialized
+                    isRunning -> LiveTile
+                    isBackgroundOnly(taskVal) ->
+                        BackgroundOnly(taskVal.colorBackground.removeAlpha())
+                    isSnapshotSplashState(taskVal) ->
+                        SnapshotSplash(createSnapshotState(taskVal), createSplashState(taskVal))
+                    else -> Uninitialized
                 }
             }
             .distinctUntilChanged()
 
-    fun bind(taskThumbnail: TaskThumbnail) {
-        this.taskThumbnail = taskThumbnail
-        task.value = tasksRepository.getTaskDataById(taskThumbnail.taskId)
-        splashProgress.value = splashAlphaUseCase.execute(taskThumbnail.taskId)
+    fun bind(taskId: Int) {
+        this.taskId = taskId
+        task.value = tasksRepository.getTaskDataById(taskId)
+        splashProgress.value = splashAlphaUseCase.execute(taskId)
     }
 
     fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): Matrix {
         return runBlocking {
             when (
                 val thumbnailPositionState =
-                    getThumbnailPositionUseCase.run(taskThumbnail.taskId, width, height, isRtl)
+                    getThumbnailPositionUseCase.run(taskId, width, height, isRtl)
             ) {
                 is ThumbnailPositionState.MatrixScaling -> thumbnailPositionState.matrix
                 is ThumbnailPositionState.MissingThumbnail -> Matrix.IDENTITY_MATRIX
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 8e232ee..8b6bc39 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -233,6 +233,7 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -807,6 +808,7 @@
     private boolean mAnyTaskHasBeenDismissed;
 
     private final RecentsViewModel mRecentsViewModel;
+    private final RecentsViewHelper mHelper;
 
     public RecentsView(Context context, @Nullable AttributeSet attrs, int defStyleAttr,
             BaseContainerInterface sizeStrategy) {
@@ -828,6 +830,7 @@
                     recentsDependencies.inject(RecentTasksRepository.class),
                     recentsDependencies.inject(RecentsViewData.class)
             );
+            mHelper = new RecentsViewHelper(mRecentsViewModel);
 
             recentsDependencies.provide(RecentsRotationStateRepository.class,
                     () -> new RecentsRotationStateRepositoryImpl(mOrientationState));
@@ -836,6 +839,7 @@
                     () -> new RecentsDeviceProfileRepositoryImpl(mContainer));
         } else {
             mRecentsViewModel = null;
+            mHelper = null;
         }
 
         mScrollHapticMinGapMillis = getResources()
@@ -1173,6 +1177,9 @@
         if (FeatureFlags.enableSplitContextually()) {
             mSplitSelectStateController.registerSplitListener(mSplitSelectionListener);
         }
+        if (enableRefactorTaskThumbnail()) {
+            mHelper.onAttachedToWindow();
+        }
     }
 
     @Override
@@ -1195,6 +1202,9 @@
             mSplitSelectStateController.unregisterSplitListener(mSplitSelectionListener);
         }
         reset();
+        if (enableRefactorTaskThumbnail()) {
+            mHelper.onDetachedFromWindow();
+        }
     }
 
     @Override
@@ -2565,14 +2575,19 @@
 
         // These are relatively expensive and don't need to be done this frame (RecentsView isn't
         // visible anyway), so defer by a frame to get off the critical path, e.g. app to home.
-        post(() -> {
-            unloadVisibleTaskData(TaskView.FLAG_UPDATE_ALL);
-            setCurrentPage(0);
-            LayoutUtils.setViewEnabled(mActionsView, true);
-            if (mOrientationState.setGestureActive(false)) {
-                updateOrientationHandler(/* forceRecreateDragLayerControllers = */ false);
-            }
-        });
+        post(this::onReset);
+    }
+
+    private void onReset() {
+        if (enableRefactorTaskThumbnail()) {
+            mRecentsViewModel.onReset();
+        }
+        unloadVisibleTaskData(TaskView.FLAG_UPDATE_ALL);
+        setCurrentPage(0);
+        LayoutUtils.setViewEnabled(mActionsView, true);
+        if (mOrientationState.setGestureActive(false)) {
+            updateOrientationHandler(/* forceRecreateDragLayerControllers = */ false);
+        }
     }
 
     public int getRunningTaskViewId() {
@@ -2948,18 +2963,13 @@
     }
 
     private void setRunningTaskViewId(int runningTaskViewId) {
-        int prevRunningTaskViewId = mRunningTaskViewId;
         mRunningTaskViewId = runningTaskViewId;
 
         if (enableRefactorTaskThumbnail()) {
-            TaskView previousRunningTaskView = getTaskViewFromTaskViewId(prevRunningTaskViewId);
-            if (previousRunningTaskView != null) {
-                previousRunningTaskView.notifyIsRunningTaskUpdated();
-            }
-            TaskView newRunningTaskView = getTaskViewFromTaskViewId(runningTaskViewId);
-            if (newRunningTaskView != null) {
-                newRunningTaskView.notifyIsRunningTaskUpdated();
-            }
+            TaskView runningTaskView = getTaskViewFromTaskViewId(runningTaskViewId);
+            mRecentsViewModel.updateRunningTask(
+                    runningTaskView != null ? runningTaskView.getTaskIdSet()
+                            : Collections.emptySet());
         }
     }
 
@@ -2989,6 +2999,9 @@
         if (runningTaskView != null) {
             runningTaskView.setShouldShowScreenshot(mRunningTaskShowScreenshot);
         }
+        if (enableRefactorTaskThumbnail()) {
+            mRecentsViewModel.setRunningTaskShowScreenshot(showScreenshot);
+        }
     }
 
     public void setTaskIconScaledDown(boolean isScaledDown) {
@@ -6024,38 +6037,28 @@
             return;
         }
 
-        switchToScreenshotInternal(onFinishRunnable);
-    }
-
-    private void switchToScreenshotInternal(Runnable onFinishRunnable) {
-        // TODO(b/342560598): Handle switchToScreenshot for new TTV.
-        if (enableRefactorTaskThumbnail()) {
-            onFinishRunnable.run();
-            return;
-        }
-
         TaskView taskView = getRunningTaskView();
         if (taskView == null) {
             onFinishRunnable.run();
             return;
         }
 
-        setRunningTaskViewShowScreenshot(true);
-        for (TaskContainer container : taskView.getTaskContainers()) {
-            if (container == null) {
-                continue;
+        if (enableRefactorTaskThumbnail()) {
+            mHelper.switchToScreenshot(taskView, mRecentsAnimationController, onFinishRunnable);
+        } else {
+            setRunningTaskViewShowScreenshot(true);
+            for (TaskContainer container : taskView.getTaskContainers()) {
+                ThumbnailData thumbnailData =
+                        mRecentsAnimationController.screenshotTask(container.getTask().key.id);
+                TaskThumbnailViewDeprecated thumbnailView = container.getThumbnailViewDeprecated();
+                if (thumbnailData != null) {
+                    thumbnailView.setThumbnail(container.getTask(), thumbnailData);
+                } else {
+                    thumbnailView.refresh();
+                }
             }
-
-            ThumbnailData td =
-                    mRecentsAnimationController.screenshotTask(container.getTask().key.id);
-            TaskThumbnailViewDeprecated thumbnailView = container.getThumbnailViewDeprecated();
-            if (td != null) {
-                thumbnailView.setThumbnail(container.getTask(), td);
-            } else {
-                thumbnailView.refresh();
-            }
+            ViewUtils.postFrameDrawn(taskView, onFinishRunnable);
         }
-        ViewUtils.postFrameDrawn(taskView, onFinishRunnable);
     }
 
     /**
@@ -6069,9 +6072,13 @@
             Runnable onFinishRunnable) {
         final TaskView taskView = getRunningTaskView();
         if (taskView != null) {
-            taskView.setShouldShowScreenshot(true);
-            taskView.refreshThumbnails(thumbnailDatas);
-            ViewUtils.postFrameDrawn(taskView, onFinishRunnable);
+            if (enableRefactorTaskThumbnail()) {
+                mHelper.switchToScreenshot(taskView, thumbnailDatas, onFinishRunnable);
+            } else {
+                taskView.setShouldShowScreenshot(true);
+                taskView.refreshThumbnails(thumbnailDatas);
+                ViewUtils.postFrameDrawn(taskView, onFinishRunnable);
+            }
         } else {
             onFinishRunnable.run();
         }
diff --git a/quickstep/src/com/android/quickstep/views/RecentsViewHelper.kt b/quickstep/src/com/android/quickstep/views/RecentsViewHelper.kt
new file mode 100644
index 0000000..05c2462
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/views/RecentsViewHelper.kt
@@ -0,0 +1,74 @@
+/*
+ * Copyright (C) 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.quickstep.views
+
+import com.android.quickstep.RecentsAnimationController
+import com.android.quickstep.ViewUtils
+import com.android.quickstep.recents.viewmodel.RecentsViewModel
+import com.android.systemui.shared.recents.model.ThumbnailData
+import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.launch
+
+/** Helper for [RecentsView] to interact with coroutine. */
+class RecentsViewHelper(private val recentsViewModel: RecentsViewModel) {
+    private lateinit var viewAttachedScope: CoroutineScope
+
+    fun onAttachedToWindow() {
+        viewAttachedScope =
+            CoroutineScope(SupervisorJob() + Dispatchers.Main + CoroutineName("RecentsView"))
+    }
+
+    fun onDetachedFromWindow() {
+        viewAttachedScope.cancel("RecentsView detaching from window")
+    }
+
+    fun switchToScreenshot(
+        taskView: TaskView,
+        recentsAnimationController: RecentsAnimationController,
+        onFinishRunnable: Runnable
+    ) {
+        val updatedThumbnails =
+            taskView.taskContainers.associate {
+                it.task.key.id to recentsAnimationController.screenshotTask(it.task.key.id)
+            }
+        switchToScreenshot(taskView, updatedThumbnails, onFinishRunnable)
+    }
+
+    fun switchToScreenshot(
+        taskView: TaskView,
+        updatedThumbnails: Map<Int, ThumbnailData>?,
+        onFinishRunnable: Runnable
+    ) {
+        // Update recentsViewModel and apply the thumbnailOverride ASAP, before waiting inside
+        // viewAttachedScope.
+        recentsViewModel.setRunningTaskShowScreenshot(true)
+        if (updatedThumbnails != null) {
+            recentsViewModel.setThumbnailOverride(updatedThumbnails)
+        }
+        viewAttachedScope.launch {
+            recentsViewModel.waitForRunningTaskShowScreenshotToUpdate()
+            if (updatedThumbnails != null) {
+                recentsViewModel.waitForThumbnailsToUpdate(updatedThumbnails)
+            }
+            ViewUtils.postFrameDrawn(taskView, onFinishRunnable)
+        }
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/views/TaskContainer.kt b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
index b1a25b5..a922e37 100644
--- a/quickstep/src/com/android/quickstep/views/TaskContainer.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
@@ -34,7 +34,6 @@
 import com.android.quickstep.recents.di.getScope
 import com.android.quickstep.recents.di.inject
 import com.android.quickstep.recents.viewmodel.TaskContainerViewModel
-import com.android.quickstep.task.thumbnail.TaskThumbnail
 import com.android.quickstep.task.thumbnail.TaskThumbnailView
 import com.android.quickstep.task.viewmodel.TaskContainerData
 import com.android.quickstep.task.viewmodel.TaskThumbnailViewModel
@@ -165,9 +164,7 @@
     }
 
     fun bindThumbnailView() {
-        // TODO(b/343364498): Existing view has shouldShowScreenshot as an override as well but
-        //  this should be decided inside TaskThumbnailViewModel.
-        taskThumbnailViewModel.bind(TaskThumbnail(task.key.id, taskView.isRunningTask))
+        taskThumbnailViewModel.bind(task.key.id)
     }
 
     fun setOverlayEnabled(enabled: Boolean) {
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index b2abe69..292595c 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -123,6 +123,10 @@
         /** Returns a copy of integer array containing taskIds of all tasks in the TaskView. */
         get() = taskContainers.map { it.task.key.id }.toIntArray()
 
+    val taskIdSet: Set<Int>
+        /** Returns a copy of integer array containing taskIds of all tasks in the TaskView. */
+        get() = taskContainers.map { it.task.key.id }.toSet()
+
     val snapshotViews: Array<View>
         get() = taskContainers.map { it.snapshotView }.toTypedArray()
 
@@ -569,9 +573,7 @@
         resetPersistentViewTransforms()
         // Clear any references to the thumbnail (it will be re-read either from the cache or the
         // system on next bind)
-        if (enableRefactorTaskThumbnail()) {
-            notifyIsRunningTaskUpdated()
-        } else {
+        if (!enableRefactorTaskThumbnail()) {
             taskContainers.forEach { it.thumbnailViewDeprecated.setThumbnail(it.task, null) }
         }
         setOverlayEnabled(false)
@@ -921,9 +923,8 @@
         iconView.setText(text)
     }
 
-    open fun refreshThumbnails(thumbnailDatas: HashMap<Int, ThumbnailData?>?) {
+    open fun refreshThumbnails(thumbnailDatas: Map<Int, ThumbnailData?>?) {
         if (enableRefactorTaskThumbnail()) {
-            // TODO(b/342560598) add thumbnail logic
             return
         }
 
@@ -1490,11 +1491,6 @@
         }
     }
 
-    /** Updates [TaskThumbnailView] to reflect the latest [Task] state (i.e., task isRunning). */
-    fun notifyIsRunningTaskUpdated() {
-        taskContainers.forEach { it.bindThumbnailView() }
-    }
-
     fun resetPersistentViewTransforms() {
         nonGridTranslationX = 0f
         gridTranslationX = 0f
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTasksRepository.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTasksRepository.kt
index ec1da5a..e488413 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTasksRepository.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTasksRepository.kt
@@ -20,6 +20,7 @@
 import com.android.systemui.shared.recents.model.ThumbnailData
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 
 class FakeTasksRepository : RecentTasksRepository {
@@ -27,11 +28,23 @@
     private var taskIconDataMap: Map<Int, TaskIconQueryResponse> = emptyMap()
     private var tasks: MutableStateFlow<List<Task>> = MutableStateFlow(emptyList())
     private var visibleTasks: MutableStateFlow<List<Int>> = MutableStateFlow(emptyList())
+    private var thumbnailOverrideMap: Map<Int, ThumbnailData> = emptyMap()
 
     override fun getAllTaskData(forceRefresh: Boolean): Flow<List<Task>> = tasks
 
     override fun getTaskDataById(taskId: Int): Flow<Task?> =
-        getAllTaskData().map { taskList -> taskList.firstOrNull { it.key.id == taskId } }
+        combine(getAllTaskData(), visibleTasks) { taskList, visibleTasks ->
+                taskList.filter { visibleTasks.contains(it.key.id) }
+            }
+            .map { taskList ->
+                val task = taskList.firstOrNull { it.key.id == taskId } ?: return@map null
+                Task(task).apply {
+                    thumbnail = thumbnailOverrideMap[taskId] ?: task.thumbnail
+                    icon = task.icon
+                    titleDescription = task.titleDescription
+                    title = task.title
+                }
+            }
 
     override fun getThumbnailById(taskId: Int): Flow<ThumbnailData?> =
         getTaskDataById(taskId).map { it?.thumbnail }
@@ -49,6 +62,16 @@
                     }
                 }
             }
+        setThumbnailOverrideInternal(thumbnailOverrideMap)
+    }
+
+    override fun setThumbnailOverride(thumbnailOverride: Map<Int, ThumbnailData>) {
+        setThumbnailOverrideInternal(thumbnailOverride)
+    }
+
+    private fun setThumbnailOverrideInternal(thumbnailOverride: Map<Int, ThumbnailData>) {
+        thumbnailOverrideMap =
+            thumbnailOverride.filterKeys(this.visibleTasks.value::contains).toMap()
     }
 
     fun seedTasks(tasks: List<Task>) {
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 88fa190..aee5d1e 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
@@ -18,9 +18,13 @@
 
 import android.content.ComponentName
 import android.content.Intent
+import android.graphics.Bitmap
+import android.view.Surface
+import com.android.quickstep.task.thumbnail.TaskThumbnailViewModelTest
 import com.android.quickstep.util.DesktopTask
 import com.android.quickstep.util.GroupTask
 import com.android.systemui.shared.recents.model.Task
+import com.android.systemui.shared.recents.model.ThumbnailData
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.drop
@@ -30,6 +34,8 @@
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 @OptIn(ExperimentalCoroutinesApi::class)
 class TasksRepositoryTest {
@@ -186,6 +192,68 @@
         assertThat(taskFlowValuesList[1]!!.thumbnail!!.thumbnail).isEqualTo(bitmap2)
     }
 
+    @Test
+    fun setThumbnailOverrideOverrideThumbnails() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+        val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+        val thumbnailOverride = createThumbnailData()
+        systemUnderTest.getAllTaskData(forceRefresh = true)
+
+        systemUnderTest.setVisibleTasks(listOf(1))
+        systemUnderTest.setThumbnailOverride(mapOf(2 to thumbnailOverride))
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+        // .drop(1) to ignore initial null content before from thumbnail was loaded.
+        assertThat(systemUnderTest.getThumbnailById(1).drop(1).first()!!.thumbnail)
+            .isEqualTo(bitmap1)
+        assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail).isEqualTo(bitmap2)
+    }
+
+    @Test
+    fun setThumbnailOverrideClearedWhenTaskBecomeInvisible() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+        val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+        val thumbnailOverride = createThumbnailData()
+        systemUnderTest.getAllTaskData(forceRefresh = true)
+
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
+        systemUnderTest.setThumbnailOverride(mapOf(2 to thumbnailOverride))
+        systemUnderTest.setVisibleTasks(listOf(1))
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+        // .drop(1) to ignore initial null content before from thumbnail was loaded.
+        assertThat(systemUnderTest.getThumbnailById(1).drop(1).first()!!.thumbnail)
+            .isEqualTo(bitmap1)
+        assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail).isEqualTo(bitmap2)
+    }
+
+    @Test
+    fun setThumbnailOverrideDoesNotOverrideInvisibleTasks() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+        val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+        val thumbnailOverride = createThumbnailData()
+        systemUnderTest.getAllTaskData(forceRefresh = true)
+
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
+        systemUnderTest.setThumbnailOverride(mapOf(2 to thumbnailOverride))
+
+        // .drop(1) to ignore initial null content before from thumbnail was loaded.
+        assertThat(systemUnderTest.getThumbnailById(1).drop(1).first()!!.thumbnail)
+            .isEqualTo(bitmap1)
+        assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail)
+            .isEqualTo(thumbnailOverride.thumbnail)
+    }
+
     private fun createTaskWithId(taskId: Int) =
         Task(Task.TaskKey(taskId, 0, Intent(), ComponentName("", ""), 0, 2000))
+
+    private fun createThumbnailData(rotation: Int = Surface.ROTATION_0): ThumbnailData {
+        val bitmap = mock<Bitmap>()
+        whenever(bitmap.width).thenReturn(TaskThumbnailViewModelTest.THUMBNAIL_WIDTH)
+        whenever(bitmap.height).thenReturn(TaskThumbnailViewModelTest.THUMBNAIL_HEIGHT)
+
+        return ThumbnailData(thumbnail = bitmap, rotation = rotation)
+    }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt
new file mode 100644
index 0000000..00dbcc1
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt
@@ -0,0 +1,123 @@
+/*
+ * Copyright (C) 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.quickstep.recents.viewmodel
+
+import android.content.ComponentName
+import android.content.Intent
+import android.graphics.Bitmap
+import android.graphics.Color
+import android.view.Surface
+import com.android.quickstep.recents.data.FakeTasksRepository
+import com.android.quickstep.task.thumbnail.TaskThumbnailViewModelTest
+import com.android.systemui.shared.recents.model.Task
+import com.android.systemui.shared.recents.model.ThumbnailData
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+class RecentsViewModelTest {
+    private val tasksRepository = FakeTasksRepository()
+    private val recentsViewData = RecentsViewData()
+    private val systemUnderTest = RecentsViewModel(tasksRepository, recentsViewData)
+
+    private val tasks = (0..5).map(::createTaskWithId)
+
+    @Test
+    fun taskVisibilityControlThumbnailsAvailability() = runTest {
+        val thumbnailData1 = createThumbnailData()
+        val thumbnailData2 = createThumbnailData()
+        tasksRepository.seedTasks(tasks)
+        tasksRepository.seedThumbnailData(mapOf(1 to thumbnailData1, 2 to thumbnailData2))
+
+        val thumbnailDataFlow1 = tasksRepository.getThumbnailById(1)
+        val thumbnailDataFlow2 = tasksRepository.getThumbnailById(2)
+
+        systemUnderTest.refreshAllTaskData()
+
+        assertThat(thumbnailDataFlow1.first()).isNull()
+        assertThat(thumbnailDataFlow2.first()).isNull()
+
+        systemUnderTest.updateVisibleTasks(listOf(1, 2))
+
+        assertThat(thumbnailDataFlow1.first()).isEqualTo(thumbnailData1)
+        assertThat(thumbnailDataFlow2.first()).isEqualTo(thumbnailData2)
+
+        systemUnderTest.updateVisibleTasks(listOf(1))
+
+        assertThat(thumbnailDataFlow1.first()).isEqualTo(thumbnailData1)
+        assertThat(thumbnailDataFlow2.first()).isNull()
+
+        systemUnderTest.onReset()
+
+        assertThat(thumbnailDataFlow1.first()).isNull()
+        assertThat(thumbnailDataFlow2.first()).isNull()
+    }
+
+    @Test
+    fun thumbnailOverrideWaitAndReset() = runTest {
+        val thumbnailData1 = createThumbnailData()
+        val thumbnailData2 = createThumbnailData()
+        val thumbnailDataOverride = createThumbnailData()
+        tasksRepository.seedTasks(tasks)
+        tasksRepository.seedThumbnailData(mapOf(1 to thumbnailData1, 2 to thumbnailData2))
+
+        val thumbnailDataFlow1 = tasksRepository.getThumbnailById(1)
+        val thumbnailDataFlow2 = tasksRepository.getThumbnailById(2)
+
+        systemUnderTest.refreshAllTaskData()
+        systemUnderTest.updateVisibleTasks(listOf(1, 2))
+
+        assertThat(thumbnailDataFlow1.first()).isEqualTo(thumbnailData1)
+        assertThat(thumbnailDataFlow2.first()).isEqualTo(thumbnailData2)
+
+        val thumbnailUpdate = mapOf(2 to thumbnailDataOverride)
+        systemUnderTest.setRunningTaskShowScreenshot(true)
+        systemUnderTest.setThumbnailOverride(thumbnailUpdate)
+
+        systemUnderTest.waitForRunningTaskShowScreenshotToUpdate()
+        systemUnderTest.waitForThumbnailsToUpdate(thumbnailUpdate)
+
+        assertThat(thumbnailDataFlow1.first()).isEqualTo(thumbnailData1)
+        assertThat(thumbnailDataFlow2.first()).isEqualTo(thumbnailDataOverride)
+
+        systemUnderTest.onReset()
+
+        assertThat(thumbnailDataFlow1.first()).isNull()
+        assertThat(thumbnailDataFlow2.first()).isNull()
+
+        systemUnderTest.updateVisibleTasks(listOf(1, 2))
+
+        assertThat(thumbnailDataFlow1.first()).isEqualTo(thumbnailData1)
+        assertThat(thumbnailDataFlow2.first()).isEqualTo(thumbnailData2)
+    }
+
+    private fun createTaskWithId(taskId: Int) =
+        Task(Task.TaskKey(taskId, 0, Intent(), ComponentName("", ""), 0, 2000)).apply {
+            colorBackground = Color.argb(taskId, taskId, taskId, taskId)
+        }
+
+    private fun createThumbnailData(rotation: Int = Surface.ROTATION_0): ThumbnailData {
+        val bitmap = mock<Bitmap>()
+        whenever(bitmap.width).thenReturn(TaskThumbnailViewModelTest.THUMBNAIL_WIDTH)
+        whenever(bitmap.height).thenReturn(TaskThumbnailViewModelTest.THUMBNAIL_HEIGHT)
+
+        return ThumbnailData(thumbnail = bitmap, rotation = rotation)
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
index 877528e..fe7d37a 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
@@ -53,6 +53,7 @@
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.whenever
 
+/** Test for [TaskThumbnailView] */
 @RunWith(AndroidJUnit4::class)
 class TaskThumbnailViewModelTest {
     private var taskViewType = TaskViewType.SINGLE
@@ -89,14 +90,42 @@
 
     @Test
     fun bindRunningTask_thenStateIs_LiveTile() = runTest {
+        val taskId = 1
         tasksRepository.seedTasks(tasks)
-        val taskThumbnail = TaskThumbnail(taskId = 1, isRunning = true)
-        systemUnderTest.bind(taskThumbnail)
+        tasksRepository.setVisibleTasks(listOf(taskId))
+        recentsViewData.runningTaskIds.value = setOf(taskId)
+        systemUnderTest.bind(taskId)
 
         assertThat(systemUnderTest.uiState.first()).isEqualTo(LiveTile)
     }
 
     @Test
+    fun bindRunningTaskShouldShowScreenshot_thenStateIs_SnapshotSplash() = runTest {
+        val taskId = 1
+        val expectedThumbnailData = createThumbnailData()
+        tasksRepository.seedThumbnailData(mapOf(taskId to expectedThumbnailData))
+        val expectedIconData = createIconData("Task 1")
+        tasksRepository.seedIconData(mapOf(taskId to expectedIconData))
+        tasksRepository.seedTasks(tasks)
+        tasksRepository.setVisibleTasks(listOf(taskId))
+        recentsViewData.runningTaskIds.value = setOf(taskId)
+        recentsViewData.runningTaskShowScreenshot.value = true
+        systemUnderTest.bind(taskId)
+
+        assertThat(systemUnderTest.uiState.first())
+            .isEqualTo(
+                SnapshotSplash(
+                    Snapshot(
+                        backgroundColor = Color.rgb(1, 1, 1),
+                        bitmap = expectedThumbnailData.thumbnail!!,
+                        thumbnailRotation = Surface.ROTATION_0,
+                    ),
+                    Splash(expectedIconData.icon, Point())
+                )
+            )
+    }
+
+    @Test
     fun setRecentsFullscreenProgress_thenCornerRadiusProgressIsPassedThrough() = runTest {
         recentsViewData.fullscreenProgress.value = 0.5f
 
@@ -130,50 +159,54 @@
     @Test
     fun bindRunningTaskThenStoppedTaskWithoutThumbnail_thenStateChangesToBackgroundOnly() =
         runTest {
+            val runningTaskId = 1
+            val stoppedTaskId = 2
             tasksRepository.seedTasks(tasks)
-            val runningTask = TaskThumbnail(taskId = 1, isRunning = true)
-            val stoppedTask = TaskThumbnail(taskId = 2, isRunning = false)
-            systemUnderTest.bind(runningTask)
+            tasksRepository.setVisibleTasks(listOf(runningTaskId, stoppedTaskId))
+            recentsViewData.runningTaskIds.value = setOf(runningTaskId)
+            systemUnderTest.bind(runningTaskId)
             assertThat(systemUnderTest.uiState.first()).isEqualTo(LiveTile)
 
-            systemUnderTest.bind(stoppedTask)
+            systemUnderTest.bind(stoppedTaskId)
             assertThat(systemUnderTest.uiState.first())
                 .isEqualTo(BackgroundOnly(backgroundColor = Color.rgb(2, 2, 2)))
         }
 
     @Test
     fun bindStoppedTaskWithoutThumbnail_thenStateIs_BackgroundOnly_withAlphaRemoved() = runTest {
+        val stoppedTaskId = 2
         tasksRepository.seedTasks(tasks)
-        val stoppedTask = TaskThumbnail(taskId = 2, isRunning = false)
+        tasksRepository.setVisibleTasks(listOf(stoppedTaskId))
 
-        systemUnderTest.bind(stoppedTask)
+        systemUnderTest.bind(stoppedTaskId)
         assertThat(systemUnderTest.uiState.first())
             .isEqualTo(BackgroundOnly(backgroundColor = Color.rgb(2, 2, 2)))
     }
 
     @Test
     fun bindLockedTaskWithThumbnail_thenStateIs_BackgroundOnly() = runTest {
-        tasksRepository.seedThumbnailData(mapOf(2 to createThumbnailData()))
-        tasks[2].isLocked = true
+        val taskId = 2
+        tasksRepository.seedThumbnailData(mapOf(taskId to createThumbnailData()))
+        tasks[taskId].isLocked = true
         tasksRepository.seedTasks(tasks)
-        val recentTask = TaskThumbnail(taskId = 2, isRunning = false)
+        tasksRepository.setVisibleTasks(listOf(taskId))
 
-        systemUnderTest.bind(recentTask)
+        systemUnderTest.bind(taskId)
         assertThat(systemUnderTest.uiState.first())
             .isEqualTo(BackgroundOnly(backgroundColor = Color.rgb(2, 2, 2)))
     }
 
     @Test
     fun bindStoppedTaskWithThumbnail_thenStateIs_SnapshotSplash_withAlphaRemoved() = runTest {
+        val taskId = 2
         val expectedThumbnailData = createThumbnailData(rotation = Surface.ROTATION_270)
-        tasksRepository.seedThumbnailData(mapOf(2 to expectedThumbnailData))
+        tasksRepository.seedThumbnailData(mapOf(taskId to expectedThumbnailData))
         val expectedIconData = createIconData("Task 2")
-        tasksRepository.seedIconData(mapOf(2 to expectedIconData))
+        tasksRepository.seedIconData(mapOf(taskId to expectedIconData))
         tasksRepository.seedTasks(tasks)
-        tasksRepository.setVisibleTasks(listOf(2))
-        val recentTask = TaskThumbnail(taskId = 2, isRunning = false)
+        tasksRepository.setVisibleTasks(listOf(taskId))
 
-        systemUnderTest.bind(recentTask)
+        systemUnderTest.bind(taskId)
         assertThat(systemUnderTest.uiState.first())
             .isEqualTo(
                 SnapshotSplash(
@@ -189,17 +222,17 @@
 
     @Test
     fun bindNonVisibleStoppedTask_whenMadeVisible_thenStateIsSnapshotSplash() = runTest {
+        val taskId = 2
         val expectedThumbnailData = createThumbnailData()
-        tasksRepository.seedThumbnailData(mapOf(2 to expectedThumbnailData))
+        tasksRepository.seedThumbnailData(mapOf(taskId to expectedThumbnailData))
         val expectedIconData = createIconData("Task 2")
-        tasksRepository.seedIconData(mapOf(2 to expectedIconData))
+        tasksRepository.seedIconData(mapOf(taskId to expectedIconData))
         tasksRepository.seedTasks(tasks)
-        val recentTask = TaskThumbnail(taskId = 2, isRunning = false)
 
-        systemUnderTest.bind(recentTask)
-        assertThat(systemUnderTest.uiState.first())
-            .isEqualTo(BackgroundOnly(backgroundColor = Color.rgb(2, 2, 2)))
-        tasksRepository.setVisibleTasks(listOf(2))
+        systemUnderTest.bind(taskId)
+        assertThat(systemUnderTest.uiState.first()).isEqualTo(Uninitialized)
+
+        tasksRepository.setVisibleTasks(listOf(taskId))
         assertThat(systemUnderTest.uiState.first())
             .isEqualTo(
                 SnapshotSplash(
@@ -215,17 +248,17 @@
 
     @Test
     fun bindStoppedTask_thenStateContainsSplashSizeFromUseCase() = runTest {
+        val taskId = 2
         val expectedSplashSize = Point(100, 150)
         whenever(getSplashSizeUseCase.execute(any())).thenReturn(expectedSplashSize)
         val expectedThumbnailData = createThumbnailData(rotation = Surface.ROTATION_270)
-        tasksRepository.seedThumbnailData(mapOf(2 to expectedThumbnailData))
+        tasksRepository.seedThumbnailData(mapOf(taskId to expectedThumbnailData))
         val expectedIconData = createIconData("Task 2")
-        tasksRepository.seedIconData(mapOf(2 to expectedIconData))
+        tasksRepository.seedIconData(mapOf(taskId to expectedIconData))
         tasksRepository.seedTasks(tasks)
-        tasksRepository.setVisibleTasks(listOf(2))
-        val recentTask = TaskThumbnail(taskId = 2, isRunning = false)
+        tasksRepository.setVisibleTasks(listOf(taskId))
 
-        systemUnderTest.bind(recentTask)
+        systemUnderTest.bind(taskId)
         val uiState = systemUnderTest.uiState.first() as SnapshotSplash
         assertThat(uiState.splash.size).isEqualTo(expectedSplashSize)
     }
@@ -233,13 +266,12 @@
     @Test
     fun getSnapshotMatrix_MissingThumbnail() = runTest {
         val taskId = 2
-        val recentTask = TaskThumbnail(taskId = taskId, isRunning = false)
         val isRtl = true
 
         whenever(mGetThumbnailPositionUseCase.run(taskId, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
             .thenReturn(MissingThumbnail)
 
-        systemUnderTest.bind(recentTask)
+        systemUnderTest.bind(taskId)
         assertThat(systemUnderTest.getThumbnailPositionState(CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
             .isEqualTo(Matrix.IDENTITY_MATRIX)
     }
@@ -247,13 +279,12 @@
     @Test
     fun getSnapshotMatrix_MatrixScaling() = runTest {
         val taskId = 2
-        val recentTask = TaskThumbnail(taskId = taskId, isRunning = false)
         val isRtl = true
 
         whenever(mGetThumbnailPositionUseCase.run(taskId, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
             .thenReturn(MatrixScaling(MATRIX, isRotated = false))
 
-        systemUnderTest.bind(recentTask)
+        systemUnderTest.bind(taskId)
         assertThat(systemUnderTest.getThumbnailPositionState(CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
             .isEqualTo(MATRIX)
     }