Add use case for split animation to retrieve thumbnail

Fix: 349120849
Test: TaskSplitAnimationHelperTest
Flag: com.android.launcher3.enable_refactor_task_thumbnail
Change-Id: I7d4bcbdbeacd56356c71eb18c73ceef8446011d1
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
index 170e018..a2278ec 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
@@ -270,7 +270,7 @@
                                     foundTask,
                                     taskContainer.getIconView().getDrawable(),
                                     taskContainer.getSnapshotView(),
-                                    taskContainer.getThumbnail(),
+                                    taskContainer.getSplitAnimationThumbnail(),
                                     null /* intent */,
                                     null /* user */,
                                     info);
diff --git a/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt
index c1eef0b..9c4248c 100644
--- a/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt
@@ -17,6 +17,7 @@
 package com.android.quickstep.recents.data
 
 import com.android.systemui.shared.recents.model.Task
+import com.android.systemui.shared.recents.model.ThumbnailData
 import kotlinx.coroutines.flow.Flow
 
 interface RecentTasksRepository {
@@ -25,11 +26,17 @@
 
     /**
      * Gets the data associated with a task that has id [taskId]. Flow will settle on null if the
-     * task was not found.
+     * task was not found. [Task.thumbnail] will settle on null if task is invisible.
      */
     fun getTaskDataById(taskId: Int): Flow<Task?>
 
     /**
+     * Gets the [ThumbnailData] associated with a task that has id [taskId]. Flow will settle on
+     * null if the task was not found or is invisible.
+     */
+    fun getThumbnailById(taskId: Int): Flow<ThumbnailData?>
+
+    /**
      * Sets the tasks that are visible, indicating that properties relating to visuals need to be
      * populated e.g. icons/thumbnails etc.
      */
diff --git a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
index 9f3ef4a..4d6dfc3 100644
--- a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
@@ -27,6 +27,7 @@
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.distinctUntilChangedBy
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.flow.flowOf
@@ -63,6 +64,9 @@
     override fun getTaskDataById(taskId: Int): Flow<Task?> =
         taskData.map { taskList -> taskList.firstOrNull { it.key.id == taskId } }
 
+    override fun getThumbnailById(taskId: Int): Flow<ThumbnailData?> =
+        getTaskDataById(taskId).map { it?.thumbnail }.distinctUntilChangedBy { it?.snapshotId }
+
     override fun setVisibleTasks(visibleTaskIdList: List<Int>) {
         this.visibleTaskIds.value = visibleTaskIdList.toSet()
     }
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
index 20a081b..22d49c1 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
@@ -52,10 +52,10 @@
             RecentsViewContainer.containerFromContext<RecentsViewContainer>(context)
                 .getOverviewPanel<RecentsView<*, *>>()
         TaskThumbnailViewModel(
-            recentsView.mRecentsViewData,
+            recentsView.mRecentsViewData!!,
             (parent as TaskView).taskViewData,
             (parent as TaskView).getTaskContainerForTaskThumbnailView(this)!!.taskContainerData,
-            recentsView.mTasksRepository,
+            recentsView.mTasksRepository!!,
         )
     }
 
diff --git a/quickstep/src/com/android/quickstep/task/util/GetThumbnailUseCase.kt b/quickstep/src/com/android/quickstep/task/util/GetThumbnailUseCase.kt
new file mode 100644
index 0000000..e8dd04c3
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/task/util/GetThumbnailUseCase.kt
@@ -0,0 +1,30 @@
+/*
+ * 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.task.util
+
+import android.graphics.Bitmap
+import com.android.quickstep.recents.data.RecentTasksRepository
+import kotlinx.coroutines.flow.firstOrNull
+import kotlinx.coroutines.runBlocking
+
+/** Use case for retrieving thumbnail. */
+class GetThumbnailUseCase(private val taskRepository: RecentTasksRepository) {
+    /** Returns the latest thumbnail associated with [taskId] if loaded, or null otherwise */
+    fun run(taskId: Int): Bitmap? = runBlocking {
+        taskRepository.getThumbnailById(taskId).firstOrNull()?.thumbnail
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt b/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
index de39584..5e55e2e 100644
--- a/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
+++ b/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
@@ -46,7 +46,7 @@
                     overlay.taskView.context
                 )
                 .getOverviewPanel<RecentsView<*, *>>()
-        TaskOverlayViewModel(task, recentsView.mRecentsViewData, recentsView.mTasksRepository)
+        TaskOverlayViewModel(task, recentsView.mRecentsViewData!!, recentsView.mTasksRepository!!)
     }
 
     // TODO(b/331753115): TaskOverlay should listen for state changes and react.
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
index 4682323..47f32fb 100644
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
@@ -24,7 +24,6 @@
 import com.android.systemui.shared.recents.model.Task
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
-import kotlinx.coroutines.flow.distinctUntilChangedBy
 import kotlinx.coroutines.flow.map
 
 /** View model for TaskOverlay */
@@ -37,10 +36,7 @@
         combine(
                 recentsViewData.overlayEnabled,
                 recentsViewData.settledFullyVisibleTaskIds.map { it.contains(task.key.id) },
-                tasksRepository
-                    .getTaskDataById(task.key.id)
-                    .map { it?.thumbnail }
-                    .distinctUntilChangedBy { it?.snapshotId }
+                tasksRepository.getThumbnailById(task.key.id)
             ) { isOverlayEnabled, isFullyVisible, thumbnailData ->
                 if (isOverlayEnabled && isFullyVisible) {
                     Enabled(
diff --git a/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt b/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
index 0cd36f4..e31a828 100644
--- a/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
+++ b/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
@@ -122,7 +122,7 @@
                     val drawable = getDrawable(container.iconView, splitSelectSource)
                     return SplitAnimInitProps(
                         container.snapshotView,
-                        container.thumbnail,
+                        container.splitAnimationThumbnail,
                         drawable,
                         fadeWithThumbnail = true,
                         isStagedTask = true,
@@ -141,7 +141,7 @@
                 val drawable = getDrawable(it.iconView, splitSelectSource)
                 return SplitAnimInitProps(
                     it.snapshotView,
-                    it.thumbnail,
+                    it.splitAnimationThumbnail,
                     drawable,
                     fadeWithThumbnail = true,
                     isStagedTask = true,
@@ -536,8 +536,13 @@
             val appPairLaunchingAppIndex = hasChangesForBothAppPairs(launchingIconView, info)
             if (appPairLaunchingAppIndex == -1) {
                 // Launch split app pair animation
-                composeIconSplitLaunchAnimator(launchingIconView, info, t, finishCallback,
-                        cornerRadius)
+                composeIconSplitLaunchAnimator(
+                    launchingIconView,
+                    info,
+                    t,
+                    finishCallback,
+                    cornerRadius
+                )
             } else {
                 composeFullscreenIconSplitLaunchAnimator(
                     launchingIconView,
@@ -554,8 +559,14 @@
                     "unexpected null"
             }
 
-            composeFadeInSplitLaunchAnimator(initialTaskId, secondTaskId, info, t, finishCallback,
-                    cornerRadius)
+            composeFadeInSplitLaunchAnimator(
+                initialTaskId,
+                secondTaskId,
+                info,
+                t,
+                finishCallback,
+                cornerRadius
+            )
         }
     }
 
@@ -701,7 +712,7 @@
         val launchAnimation = AnimatorSet()
 
         val splitRoots: Pair<Change, List<Change>>? =
-                SplitScreenUtils.extractTopParentAndChildren(transitionInfo)
+            SplitScreenUtils.extractTopParentAndChildren(transitionInfo)
         check(splitRoots != null) { "Could not find split roots" }
 
         // Will point to change (0) in diagram above
@@ -711,10 +722,11 @@
 
         // Find the place where our left/top app window meets the divider (used for the
         // launcher side animation)
-        val leftTopApp = leafRoots.single { change ->
-            (dp.isLeftRightSplit && change.endAbsBounds.left == 0) ||
+        val leftTopApp =
+            leafRoots.single { change ->
+                (dp.isLeftRightSplit && change.endAbsBounds.left == 0) ||
                     (!dp.isLeftRightSplit && change.endAbsBounds.top == 0)
-        }
+            }
         val dividerPos =
             if (dp.isLeftRightSplit) leftTopApp.endAbsBounds.right
             else leftTopApp.endAbsBounds.bottom
@@ -736,17 +748,24 @@
             )
         floatingView.bringToFront()
 
-        val iconLaunchValueAnimator = getIconLaunchValueAnimator(t, dp, finishCallback, launcher,
-                floatingView, mainRootCandidate)
+        val iconLaunchValueAnimator =
+            getIconLaunchValueAnimator(
+                t,
+                dp,
+                finishCallback,
+                launcher,
+                floatingView,
+                mainRootCandidate
+            )
         iconLaunchValueAnimator.addListener(
-                object : AnimatorListenerAdapter() {
-                    override fun onAnimationStart(animation: Animator, isReverse: Boolean) {
-                        for (c in leafRoots) {
-                            t.setCornerRadius(c.leash, windowRadius)
-                            t.apply()
-                        }
+            object : AnimatorListenerAdapter() {
+                override fun onAnimationStart(animation: Animator, isReverse: Boolean) {
+                    for (c in leafRoots) {
+                        t.setCornerRadius(c.leash, windowRadius)
+                        t.apply()
                     }
                 }
+            }
         )
         launchAnimation.play(iconLaunchValueAnimator)
         launchAnimation.start()
@@ -1017,12 +1036,12 @@
      */
     @VisibleForTesting
     fun composeFadeInSplitLaunchAnimator(
-            initialTaskId: Int,
-            secondTaskId: Int,
-            transitionInfo: TransitionInfo,
-            t: Transaction,
-            finishCallback: Runnable,
-            cornerRadius: Float
+        initialTaskId: Int,
+        secondTaskId: Int,
+        transitionInfo: TransitionInfo,
+        t: Transaction,
+        finishCallback: Runnable,
+        cornerRadius: Float
     ) {
         var splitRoot1: Change? = null
         var splitRoot2: Change? = null
@@ -1100,7 +1119,7 @@
                 override fun onAnimationStart(animation: Animator) {
                     for (leash in openingTargets) {
                         animTransaction.show(leash).setAlpha(leash, 0.0f)
-                        animTransaction.setCornerRadius(leash, cornerRadius);
+                        animTransaction.setCornerRadius(leash, cornerRadius)
                     }
                     animTransaction.apply()
                 }
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 7b6d383..d10bc50 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -461,7 +461,9 @@
 
     private static final float FOREGROUND_SCRIM_TINT = 0.32f;
 
+    @Nullable
     public final RecentsViewData mRecentsViewData = new RecentsViewData();
+    @Nullable
     public final TasksRepository mTasksRepository;
 
     protected final RecentsOrientedState mOrientationState;
diff --git a/quickstep/src/com/android/quickstep/views/TaskContainer.kt b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
index 0648986..74d120f 100644
--- a/quickstep/src/com/android/quickstep/views/TaskContainer.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
@@ -31,6 +31,7 @@
 import com.android.quickstep.TaskUtils
 import com.android.quickstep.task.thumbnail.TaskThumbnail
 import com.android.quickstep.task.thumbnail.TaskThumbnailView
+import com.android.quickstep.task.util.GetThumbnailUseCase
 import com.android.quickstep.task.viewmodel.TaskContainerData
 import com.android.systemui.shared.recents.model.Task
 
@@ -56,6 +57,16 @@
     val overlay: TaskOverlayFactory.TaskOverlay<*> = taskOverlayFactory.createOverlay(this)
     val taskContainerData = TaskContainerData()
 
+    private val getThumbnailUseCase by lazy {
+        // TODO(b/335649589): Ideally create and obtain this from DI.
+        val recentsView =
+            RecentsViewContainer.containerFromContext<RecentsViewContainer>(
+                    overlay.taskView.context
+                )
+                .getOverviewPanel<RecentsView<*, *>>()
+        GetThumbnailUseCase(recentsView.mTasksRepository!!)
+    }
+
     init {
         if (enableRefactorTaskThumbnail()) {
             require(snapshotView is TaskThumbnailView)
@@ -64,6 +75,14 @@
         }
     }
 
+    val splitAnimationThumbnail: Bitmap?
+        get() =
+            if (enableRefactorTaskThumbnail()) {
+                getThumbnailUseCase.run(task.key.id)
+            } else {
+                thumbnailViewDeprecated.thumbnail
+            }
+
     val thumbnailView: TaskThumbnailView
         get() {
             require(enableRefactorTaskThumbnail())
@@ -76,10 +95,6 @@
             return snapshotView as TaskThumbnailViewDeprecated
         }
 
-    // TODO(b/349120849): Extract ThumbnailData from TaskContainerData/TaskThumbnailViewModel
-    val thumbnail: Bitmap?
-        get() = if (enableRefactorTaskThumbnail()) null else thumbnailViewDeprecated.thumbnail
-
     // TODO(b/334826842): Support shouldShowSplashView for new TTV.
     val shouldShowSplashView: Boolean
         get() =
@@ -114,6 +129,15 @@
                 }
             }
 
+    fun bind() {
+        if (enableRefactorTaskThumbnail()) {
+            bindThumbnailView()
+        } else {
+            thumbnailViewDeprecated.bind(task, overlay)
+        }
+        overlay.init()
+    }
+
     fun destroy() {
         digitalWellBeingToast?.destroy()
         if (enableRefactorTaskThumbnail()) {
@@ -122,15 +146,6 @@
         overlay.destroy()
     }
 
-    fun bind() {
-        if (enableRefactorTaskThumbnail()) {
-            bindThumbnailView()
-            overlay.init()
-        } else {
-            thumbnailViewDeprecated.bind(task, overlay)
-        }
-    }
-
     // TODO(b/335649589): TaskView's VM will already have access to TaskThumbnailView's VM
     //  so there will be no need to access TaskThumbnailView's VM through the TaskThumbnailView
     fun bindThumbnailView() {
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index 2e07e36..5c95aaa 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -1179,7 +1179,7 @@
             container.task,
             container.iconView.drawable,
             container.snapshotView,
-            container.thumbnail,
+            container.splitAnimationThumbnail,
             /* intent */ null,
             /* user */ null,
             container.itemInfo
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 e160627..19990a8 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
@@ -32,6 +32,9 @@
     override fun getTaskDataById(taskId: Int): Flow<Task?> =
         getAllTaskData().map { taskList -> taskList.firstOrNull { it.key.id == taskId } }
 
+    override fun getThumbnailById(taskId: Int): Flow<ThumbnailData?> =
+        getTaskDataById(taskId).map { it?.thumbnail }
+
     override fun setVisibleTasks(visibleTaskIdList: List<Int>) {
         visibleTasks.value = visibleTaskIdList
         tasks.value = tasks.value.map { it.apply { thumbnail = thumbnailDataMap[it.key.id] } }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/util/GetThumbnailUseCaseTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/util/GetThumbnailUseCaseTest.kt
new file mode 100644
index 0000000..414f8ca
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/util/GetThumbnailUseCaseTest.kt
@@ -0,0 +1,84 @@
+/*
+ * 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.task.util
+
+import android.content.ComponentName
+import android.content.Intent
+import android.graphics.Bitmap
+import android.graphics.Color
+import com.android.quickstep.recents.data.FakeTasksRepository
+import com.android.quickstep.task.viewmodel.TaskOverlayViewModelTest
+import com.android.systemui.shared.recents.model.Task
+import com.android.systemui.shared.recents.model.ThumbnailData
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+/** Test for [GetThumbnailUseCase] */
+class GetThumbnailUseCaseTest {
+    private val task =
+        Task(Task.TaskKey(TASK_ID, 0, Intent(), ComponentName("", ""), 0, 2000)).apply {
+            colorBackground = Color.BLACK
+        }
+    private val thumbnailData =
+        ThumbnailData(
+            thumbnail =
+                mock<Bitmap>().apply {
+                    whenever(width).thenReturn(THUMBNAIL_WIDTH)
+                    whenever(height).thenReturn(THUMBNAIL_HEIGHT)
+                }
+        )
+
+    private val tasksRepository = FakeTasksRepository()
+    private val systemUnderTest = GetThumbnailUseCase(tasksRepository)
+
+    @Test
+    fun taskNotSeeded_returnsNull() {
+        assertThat(systemUnderTest.run(TASK_ID)).isNull()
+    }
+
+    @Test
+    fun taskNotLoaded_returnsNull() {
+        tasksRepository.seedTasks(listOf(task))
+
+        assertThat(systemUnderTest.run(TASK_ID)).isNull()
+    }
+
+    @Test
+    fun taskNotVisible_returnsNull() {
+        tasksRepository.seedTasks(listOf(task))
+        tasksRepository.seedThumbnailData(mapOf(TaskOverlayViewModelTest.TASK_ID to thumbnailData))
+
+        assertThat(systemUnderTest.run(TASK_ID)).isNull()
+    }
+
+    @Test
+    fun taskVisible_returnsThumbnail() {
+        tasksRepository.seedTasks(listOf(task))
+        tasksRepository.seedThumbnailData(mapOf(TaskOverlayViewModelTest.TASK_ID to thumbnailData))
+        tasksRepository.setVisibleTasks(listOf(TaskOverlayViewModelTest.TASK_ID))
+
+        assertThat(systemUnderTest.run(TASK_ID)).isEqualTo(thumbnailData.thumbnail)
+    }
+
+    companion object {
+        const val TASK_ID = 0
+        const val THUMBNAIL_WIDTH = 100
+        const val THUMBNAIL_HEIGHT = 200
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
index a9f5dcd..f3cde52 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
@@ -87,7 +87,7 @@
     @Before
     fun setup() {
         whenever(mockTaskContainer.snapshotView).thenReturn(mockSnapshotView)
-        whenever(mockTaskContainer.thumbnail).thenReturn(mockBitmap)
+        whenever(mockTaskContainer.splitAnimationThumbnail).thenReturn(mockBitmap)
         whenever(mockTaskContainer.iconView).thenReturn(mockIconView)
         whenever(mockIconView.drawable).thenReturn(mockTaskViewDrawable)
         whenever(mockTaskView.taskContainers).thenReturn(List(1) { mockTaskContainer })