Log various TaskView properties in TaksContainer.itemInfo

- Extracted itemInfo creation in to TaskViewItemInfo and added a test
- Created a standalone atom for TaskView, logging data we want (please refer to the bug)
- Convert TaskView atom to StatsLog data in StatsLogCompatManager
- Simplified DesktopSystemShortcutTest and ExternalDisplaySystemShortcutTest to avoid extra testing on DesktopModeStatus

Fix: 341320349
Test: TaskViewItemInfoTest, DesktopSystemShortcutTest, ExternalDisplaySystemShortcutTest
Flag: EXEMPT logging
Change-Id: Icf04ff2b1670c3f3d5f6a66f99c0d28314efaf21
diff --git a/protos/launcher_atom.proto b/protos/launcher_atom.proto
index 823c821..ce99348 100644
--- a/protos/launcher_atom.proto
+++ b/protos/launcher_atom.proto
@@ -33,6 +33,7 @@
     FolderIcon folder_icon = 9;
     Slice slice = 10;
     SearchActionItem search_action_item = 11;
+    TaskView task_view = 15;
   }
   // When used for launch event, stores the global predictive rank
   optional int32 rank = 5;
@@ -262,6 +263,21 @@
   optional int32 index = 3;
 }
 
+// TaskView in RecentsView.
+message TaskView {
+  // TaskViewType.
+  optional int32 type = 1;
+
+  // Index of TaskView in RecentsView.
+  optional int32 index = 2;
+
+  // ComponentName of the Task.
+  optional string component_name = 3;
+
+  // Number of tasks in the TaskView.
+  optional int32 cardinality = 4;
+}
+
 // Represents folder in a closed state.
 message FolderIcon {
   // Number of items inside folder.
diff --git a/quickstep/src/com/android/launcher3/model/data/TaskViewItemInfo.kt b/quickstep/src/com/android/launcher3/model/data/TaskViewItemInfo.kt
new file mode 100644
index 0000000..ee28d7a
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/model/data/TaskViewItemInfo.kt
@@ -0,0 +1,80 @@
+/*
+ * 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.launcher3.model.data
+
+import android.content.Context
+import android.content.Intent
+import androidx.annotation.VisibleForTesting
+import androidx.annotation.VisibleForTesting.Companion.PRIVATE
+import com.android.launcher3.Flags.privateSpaceRestrictAccessibilityDrag
+import com.android.launcher3.LauncherSettings
+import com.android.launcher3.logger.LauncherAtom
+import com.android.launcher3.pm.UserCache
+import com.android.quickstep.TaskUtils
+import com.android.quickstep.views.TaskContainer
+
+class TaskViewItemInfo(taskContainer: TaskContainer) : WorkspaceItemInfo() {
+    @VisibleForTesting(otherwise = PRIVATE) val taskViewAtom: LauncherAtom.TaskView
+
+    init {
+        itemType = LauncherSettings.Favorites.ITEM_TYPE_TASK
+        container = LauncherSettings.Favorites.CONTAINER_TASKSWITCHER
+        val componentKey = TaskUtils.getLaunchComponentKeyForTask(taskContainer.task.key)
+        user = componentKey.user
+        intent = Intent().setComponent(componentKey.componentName)
+        title = taskContainer.task.title
+        if (privateSpaceRestrictAccessibilityDrag()) {
+            if (
+                UserCache.getInstance(taskContainer.taskView.context)
+                    .getUserInfo(componentKey.user)
+                    .isPrivate
+            ) {
+                runtimeStatusFlags = runtimeStatusFlags or ItemInfoWithIcon.FLAG_NOT_PINNABLE
+            }
+        }
+
+        taskViewAtom =
+            createTaskViewAtom(
+                type = taskContainer.taskView.type.ordinal,
+                index =
+                    taskContainer.taskView.recentsView?.indexOfChild(taskContainer.taskView) ?: -1,
+                componentName = componentKey.componentName.flattenToShortString(),
+                cardinality = taskContainer.taskView.taskContainers.size,
+            )
+    }
+
+    override fun buildProto(cInfo: CollectionInfo?, context: Context): LauncherAtom.ItemInfo =
+        super.buildProto(cInfo, context).toBuilder().setTaskView(taskViewAtom).build()
+
+    companion object {
+        @VisibleForTesting(otherwise = PRIVATE)
+        fun createTaskViewAtom(
+            type: Int,
+            index: Int,
+            componentName: String,
+            cardinality: Int,
+        ): LauncherAtom.TaskView =
+            LauncherAtom.TaskView.newBuilder()
+                .apply {
+                    this.type = type
+                    this.index = index
+                    this.componentName = componentName
+                    this.cardinality = cardinality
+                }
+                .build()
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index 2daaaf9..9bfe71f 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -733,48 +733,40 @@
                             .getQueryLength() : -1;
                 }
             default:
-                return info.getFolderIcon().getCardinality();
+                return switch (info.getItemCase()) {
+                    case FOLDER_ICON -> info.getFolderIcon().getCardinality();
+                    case TASK_VIEW -> info.getTaskView().getCardinality();
+                    default -> 0;
+                };
         }
     }
 
     private static String getPackageName(LauncherAtom.ItemInfo info) {
-        switch (info.getItemCase()) {
-            case APPLICATION:
-                return info.getApplication().getPackageName();
-            case SHORTCUT:
-                return info.getShortcut().getShortcutName();
-            case WIDGET:
-                return info.getWidget().getPackageName();
-            case TASK:
-                return info.getTask().getPackageName();
-            case SEARCH_ACTION_ITEM:
-                return info.getSearchActionItem().getPackageName();
-            default:
-                return null;
-        }
+        return switch (info.getItemCase()) {
+            case APPLICATION -> info.getApplication().getPackageName();
+            case SHORTCUT -> info.getShortcut().getShortcutName();
+            case WIDGET -> info.getWidget().getPackageName();
+            case TASK -> info.getTask().getPackageName();
+            case SEARCH_ACTION_ITEM -> info.getSearchActionItem().getPackageName();
+            default -> null;
+        };
     }
 
     private static String getComponentName(LauncherAtom.ItemInfo info) {
-        switch (info.getItemCase()) {
-            case APPLICATION:
-                return info.getApplication().getComponentName();
-            case SHORTCUT:
-                return info.getShortcut().getShortcutName();
-            case WIDGET:
-                return info.getWidget().getComponentName();
-            case TASK:
-                return info.getTask().getComponentName();
-            case SEARCH_ACTION_ITEM:
-                return info.getSearchActionItem().getTitle();
-            case SLICE:
-                return info.getSlice().getUri();
-            default:
-                return null;
-        }
+        return switch (info.getItemCase()) {
+            case APPLICATION -> info.getApplication().getComponentName();
+            case SHORTCUT -> info.getShortcut().getShortcutName();
+            case WIDGET -> info.getWidget().getComponentName();
+            case TASK -> info.getTask().getComponentName();
+            case TASK_VIEW -> info.getTaskView().getComponentName();
+            case SEARCH_ACTION_ITEM -> info.getSearchActionItem().getTitle();
+            case SLICE -> info.getSlice().getUri();
+            default -> null;
+        };
     }
 
     private static int getGridX(LauncherAtom.ItemInfo info, boolean parent) {
-        LauncherAtom.ContainerInfo containerInfo = info.getContainerInfo();
+        ContainerInfo containerInfo = info.getContainerInfo();
         if (containerInfo.getContainerCase() == FOLDER) {
             if (parent) {
                 return containerInfo.getFolder().getWorkspace().getGridX();
@@ -802,37 +794,38 @@
     }
 
     private static int getPageId(LauncherAtom.ItemInfo info) {
-        if (info.hasTask()) {
-            return info.getTask().getIndex();
-        }
-        switch (info.getContainerInfo().getContainerCase()) {
-            case FOLDER:
-                return info.getContainerInfo().getFolder().getPageIndex();
-            case HOTSEAT:
-                return info.getContainerInfo().getHotseat().getIndex();
-            case PREDICTED_HOTSEAT_CONTAINER:
-                return info.getContainerInfo().getPredictedHotseatContainer().getIndex();
-            case TASK_BAR_CONTAINER:
-                return info.getContainerInfo().getTaskBarContainer().getIndex();
-            default:
-                return info.getContainerInfo().getWorkspace().getPageIndex();
-        }
+        return switch (info.getItemCase()) {
+            case TASK -> info.getTask().getIndex();
+            case TASK_VIEW -> info.getTaskView().getIndex();
+            default -> getPageIdFromContainerInfo(info.getContainerInfo());
+        };
+    }
+
+    private static int getPageIdFromContainerInfo(LauncherAtom.ContainerInfo containerInfo) {
+        return switch (containerInfo.getContainerCase()) {
+            case FOLDER -> containerInfo.getFolder().getPageIndex();
+            case HOTSEAT -> containerInfo.getHotseat().getIndex();
+            case PREDICTED_HOTSEAT_CONTAINER ->
+                    containerInfo.getPredictedHotseatContainer().getIndex();
+            case TASK_BAR_CONTAINER -> containerInfo.getTaskBarContainer().getIndex();
+            default -> containerInfo.getWorkspace().getPageIndex();
+        };
     }
 
     private static int getParentPageId(LauncherAtom.ItemInfo info) {
-        switch (info.getContainerInfo().getContainerCase()) {
-            case FOLDER:
+        return switch (info.getContainerInfo().getContainerCase()) {
+            case FOLDER -> {
                 if (info.getContainerInfo().getFolder().getParentContainerCase()
                         == ParentContainerCase.HOTSEAT) {
-                    return info.getContainerInfo().getFolder().getHotseat().getIndex();
+                    yield info.getContainerInfo().getFolder().getHotseat().getIndex();
                 }
-                return info.getContainerInfo().getFolder().getWorkspace().getPageIndex();
-            case SEARCH_RESULT_CONTAINER:
-                return info.getContainerInfo().getSearchResultContainer().getWorkspace()
-                        .getPageIndex();
-            default:
-                return info.getContainerInfo().getWorkspace().getPageIndex();
-        }
+                yield info.getContainerInfo().getFolder().getWorkspace().getPageIndex();
+            }
+            case SEARCH_RESULT_CONTAINER ->
+                    info.getContainerInfo().getSearchResultContainer().getWorkspace()
+                            .getPageIndex();
+            default -> info.getContainerInfo().getWorkspace().getPageIndex();
+        };
     }
 
     private static int getHierarchy(LauncherAtom.ItemInfo info) {
@@ -857,25 +850,21 @@
     }
 
     private static String getStateString(int state) {
-        switch (state) {
-            case LAUNCHER_UICHANGED__DST_STATE__BACKGROUND:
-                return "BACKGROUND";
-            case LAUNCHER_UICHANGED__DST_STATE__HOME:
-                return "HOME";
-            case LAUNCHER_UICHANGED__DST_STATE__OVERVIEW:
-                return "OVERVIEW";
-            case LAUNCHER_UICHANGED__DST_STATE__ALLAPPS:
-                return "ALLAPPS";
-            default:
-                return "INVALID";
-        }
+        return switch (state) {
+            case LAUNCHER_UICHANGED__DST_STATE__BACKGROUND -> "BACKGROUND";
+            case LAUNCHER_UICHANGED__DST_STATE__HOME -> "HOME";
+            case LAUNCHER_UICHANGED__DST_STATE__OVERVIEW -> "OVERVIEW";
+            case LAUNCHER_UICHANGED__DST_STATE__ALLAPPS -> "ALLAPPS";
+            default -> "INVALID";
+        };
     }
 
     private static int getFeatures(LauncherAtom.ItemInfo info) {
-        if (info.getItemCase().equals(LauncherAtom.ItemInfo.ItemCase.WIDGET)) {
-            return info.getWidget().getWidgetFeatures();
-        }
-        return 0;
+        return switch (info.getItemCase()) {
+            case WIDGET -> info.getWidget().getWidgetFeatures();
+            case TASK_VIEW -> info.getTaskView().getType();
+            default -> 0;
+        };
     }
 
     private static int getSearchAttributes(LauncherAtom.ItemInfo info) {
diff --git a/quickstep/src/com/android/quickstep/views/TaskContainer.kt b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
index c940fb4..a7425e0 100644
--- a/quickstep/src/com/android/quickstep/views/TaskContainer.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
@@ -16,19 +16,13 @@
 
 package com.android.quickstep.views
 
-import android.content.Intent
 import android.graphics.Bitmap
 import android.view.View
 import com.android.launcher3.Flags.enableRefactorTaskThumbnail
-import com.android.launcher3.Flags.privateSpaceRestrictAccessibilityDrag
-import com.android.launcher3.LauncherSettings
-import com.android.launcher3.model.data.ItemInfoWithIcon
-import com.android.launcher3.model.data.WorkspaceItemInfo
-import com.android.launcher3.pm.UserCache
+import com.android.launcher3.model.data.TaskViewItemInfo
 import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.launcher3.util.TransformingTouchDelegate
 import com.android.quickstep.TaskOverlayFactory
-import com.android.quickstep.TaskUtils
 import com.android.quickstep.ViewUtils.addAccessibleChildToList
 import com.android.quickstep.recents.di.RecentsDependencies
 import com.android.quickstep.recents.di.get
@@ -123,27 +117,8 @@
             else thumbnailViewDeprecated.sysUiStatusNavFlags
 
     /** Builds proto for logging */
-    val itemInfo: WorkspaceItemInfo
-        get() =
-            WorkspaceItemInfo().apply {
-                itemType = LauncherSettings.Favorites.ITEM_TYPE_TASK
-                container = LauncherSettings.Favorites.CONTAINER_TASKSWITCHER
-                val componentKey = TaskUtils.getLaunchComponentKeyForTask(task.key)
-                user = componentKey.user
-                intent = Intent().setComponent(componentKey.componentName)
-                title = task.title
-                taskView.recentsView?.let { screenId = it.indexOfChild(taskView) }
-                if (privateSpaceRestrictAccessibilityDrag()) {
-                    if (
-                        UserCache.getInstance(taskView.context)
-                            .getUserInfo(componentKey.user)
-                            .isPrivate
-                    ) {
-                        runtimeStatusFlags =
-                            runtimeStatusFlags or ItemInfoWithIcon.FLAG_NOT_PINNABLE
-                    }
-                }
-            }
+    val itemInfo: TaskViewItemInfo
+        get() = TaskViewItemInfo(this)
 
     fun bind() {
         digitalWellBeingToast?.bind(task, taskView, snapshotView, stagePosition)
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index 7e489ea..3586dfb 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -101,7 +101,7 @@
     defStyleRes: Int = 0,
     focusBorderAnimator: BorderAnimator? = null,
     hoverBorderAnimator: BorderAnimator? = null,
-    private val type: TaskViewType = TaskViewType.SINGLE,
+    val type: TaskViewType = TaskViewType.SINGLE,
     protected val thumbnailFullscreenParams: FullscreenDrawParams = FullscreenDrawParams(context),
 ) : FrameLayout(context, attrs), ViewPool.Reusable {
     /**
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/model/data/TaskViewItemInfoTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/model/data/TaskViewItemInfoTest.kt
new file mode 100644
index 0000000..0103e7e
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/model/data/TaskViewItemInfoTest.kt
@@ -0,0 +1,154 @@
+/*
+ * 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.launcher3.model.data
+
+import android.content.ComponentName
+import android.content.Context
+import android.content.Intent
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.Flags.enableRefactorTaskThumbnail
+import com.android.launcher3.model.data.TaskViewItemInfo.Companion.createTaskViewAtom
+import com.android.launcher3.util.SplitConfigurationOptions
+import com.android.launcher3.util.TransformingTouchDelegate
+import com.android.quickstep.TaskOverlayFactory
+import com.android.quickstep.TaskOverlayFactory.TaskOverlay
+import com.android.quickstep.recents.di.RecentsDependencies
+import com.android.quickstep.task.thumbnail.TaskThumbnailView
+import com.android.quickstep.views.RecentsView
+import com.android.quickstep.views.TaskContainer
+import com.android.quickstep.views.TaskThumbnailViewDeprecated
+import com.android.quickstep.views.TaskView
+import com.android.quickstep.views.TaskViewIcon
+import com.android.quickstep.views.TaskViewType
+import com.android.systemui.shared.recents.model.Task
+import com.android.systemui.shared.recents.model.Task.TaskKey
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+/** Test for [TaskViewItemInfo] */
+@RunWith(AndroidJUnit4::class)
+class TaskViewItemInfoTest {
+    private val context = mock<Context>()
+    private val taskView = mock<TaskView>()
+    private val recentsView = mock<RecentsView<*, *>>()
+    private val overlayFactory = mock<TaskOverlayFactory>()
+
+    @Before
+    fun setUp() {
+        whenever(overlayFactory.createOverlay(any())).thenReturn(mock<TaskOverlay<*>>())
+        whenever(taskView.context).thenReturn(context)
+        whenever(taskView.recentsView).thenReturn(recentsView)
+        whenever(recentsView.indexOfChild(taskView)).thenReturn(TASK_VIEW_INDEX)
+        RecentsDependencies.initialize(InstrumentationRegistry.getInstrumentation().targetContext)
+    }
+
+    @Test
+    fun singleTask() {
+        val taskContainers = listOf(createTaskContainer(createTask(1)))
+        whenever(taskView.type).thenReturn(TaskViewType.SINGLE)
+        whenever(taskView.taskContainers).thenReturn(taskContainers)
+
+        val taskViewItemInfo = TaskViewItemInfo(taskContainers[0])
+        val taskViewAtom = taskViewItemInfo.taskViewAtom
+
+        assertThat(taskViewAtom)
+            .isEqualTo(
+                createTaskViewAtom(
+                    type = 0,
+                    index = TASK_VIEW_INDEX,
+                    componentName = "${PACKAGE}/${CLASS}",
+                    cardinality = 1,
+                )
+            )
+    }
+
+    @Test
+    fun splitTask() {
+        val taskContainers =
+            listOf(createTaskContainer(createTask(1)), createTaskContainer(createTask(2)))
+        whenever(taskView.type).thenReturn(TaskViewType.GROUPED)
+        whenever(taskView.taskContainers).thenReturn(taskContainers)
+
+        val taskViewItemInfo = TaskViewItemInfo(taskContainers[0])
+        val taskViewAtom = taskViewItemInfo.taskViewAtom
+
+        assertThat(taskViewAtom)
+            .isEqualTo(
+                createTaskViewAtom(
+                    type = 1,
+                    index = TASK_VIEW_INDEX,
+                    componentName = "${PACKAGE}/${CLASS}",
+                    cardinality = 2,
+                )
+            )
+    }
+
+    @Test
+    fun desktopTask() {
+        val taskContainers =
+            listOf(
+                createTaskContainer(createTask(1)),
+                createTaskContainer(createTask(2)),
+                createTaskContainer(createTask(3)),
+            )
+        whenever(taskView.type).thenReturn(TaskViewType.DESKTOP)
+        whenever(taskView.taskContainers).thenReturn(taskContainers)
+
+        val taskViewItemInfo = TaskViewItemInfo(taskContainers[0])
+        val taskViewAtom = taskViewItemInfo.taskViewAtom
+
+        assertThat(taskViewAtom)
+            .isEqualTo(
+                createTaskViewAtom(
+                    type = 2,
+                    index = TASK_VIEW_INDEX,
+                    componentName = "${PACKAGE}/${CLASS}",
+                    cardinality = 3,
+                )
+            )
+    }
+
+    private fun createTask(id: Int) =
+        Task(TaskKey(id, 0, Intent(), ComponentName(PACKAGE, CLASS), 0, 2000))
+
+    private fun createTaskContainer(task: Task): TaskContainer {
+        return TaskContainer(
+            taskView,
+            task,
+            if (enableRefactorTaskThumbnail()) mock<TaskThumbnailView>()
+            else mock<TaskThumbnailViewDeprecated>(),
+            mock<TaskViewIcon>(),
+            mock<TransformingTouchDelegate>(),
+            SplitConfigurationOptions.STAGE_POSITION_UNDEFINED,
+            digitalWellBeingToast = null,
+            showWindowsView = null,
+            overlayFactory,
+        )
+    }
+
+    companion object {
+        const val PACKAGE = "package"
+        const val CLASS = "class"
+        const val TASK_VIEW_INDEX = 4
+    }
+}
diff --git a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
index 231c113..94e7c2e 100644
--- a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
@@ -18,8 +18,6 @@
 
 import android.content.ComponentName
 import android.content.Intent
-import android.platform.test.flag.junit.SetFlagsRule
-import com.android.dx.mockito.inline.extended.ExtendedMockito
 import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
 import com.android.dx.mockito.inline.extended.StaticMockitoSession
 import com.android.launcher3.AbstractFloatingView
@@ -27,7 +25,7 @@
 import com.android.launcher3.Flags.enableRefactorTaskThumbnail
 import com.android.launcher3.logging.StatsLogManager
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent
-import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.model.data.TaskViewItemInfo
 import com.android.launcher3.uioverrides.QuickstepLauncher
 import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.launcher3.util.TransformingTouchDelegate
@@ -40,14 +38,13 @@
 import com.android.quickstep.views.TaskViewIcon
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.Task.TaskKey
-import com.android.window.flags.Flags
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.shared.desktopmode.DesktopModeTransitionSource
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
-import org.junit.Rule
 import org.junit.Test
+import org.mockito.Mockito.`when`
 import org.mockito.kotlin.any
 import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.eq
@@ -57,24 +54,18 @@
 import org.mockito.kotlin.whenever
 import org.mockito.quality.Strictness
 
-/** Test for DesktopSystemShortcut */
+/** Test for [DesktopSystemShortcut] */
 class DesktopSystemShortcutTest {
 
-    @get:Rule val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
-
     private val launcher: QuickstepLauncher = mock()
     private val statsLogManager: StatsLogManager = mock()
     private val statsLogger: StatsLogManager.StatsLogger = mock()
     private val recentsView: LauncherRecentsView = mock()
     private val taskView: TaskView = mock()
-    private val workspaceItemInfo: WorkspaceItemInfo = mock()
     private val abstractFloatingViewHelper: AbstractFloatingViewHelper = mock()
-    private val iconView: TaskViewIcon = mock()
-    private val transformingTouchDelegate: TransformingTouchDelegate = mock()
+    private val overlayFactory: TaskOverlayFactory = mock()
     private val factory: TaskShortcutFactory =
         DesktopSystemShortcut.createFactory(abstractFloatingViewHelper)
-    private val overlayFactory: TaskOverlayFactory = mock()
-    private val overlay: TaskOverlay<*> = mock()
 
     private lateinit var mockitoSession: StaticMockitoSession
 
@@ -83,11 +74,10 @@
         mockitoSession =
             mockitoSession()
                 .strictness(Strictness.LENIENT)
-                .spyStatic(DesktopModeStatus::class.java)
+                .mockStatic(DesktopModeStatus::class.java)
                 .startMocking()
-        ExtendedMockito.doReturn(true).`when` { DesktopModeStatus.enforceDeviceRestrictions() }
-        ExtendedMockito.doReturn(true).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
-        whenever(overlayFactory.createOverlay(any())).thenReturn(overlay)
+        whenever(DesktopModeStatus.canEnterDesktopMode(any())).thenReturn(true)
+        whenever(overlayFactory.createOverlay(any())).thenReturn(mock<TaskOverlay<*>>())
     }
 
     @After
@@ -97,22 +87,7 @@
 
     @Test
     fun createDesktopTaskShortcutFactory_desktopModeDisabled() {
-        setFlagsRule.disableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE)
-
-        val task =
-            Task(TaskKey(1, 0, Intent(), ComponentName("", ""), 0, 2000)).apply {
-                isDockable = true
-            }
-        val taskContainer = createTaskContainer(task)
-
-        val shortcuts = factory.getShortcuts(launcher, taskContainer)
-        assertThat(shortcuts).isNull()
-    }
-
-    @Test
-    fun createDesktopTaskShortcutFactory_desktopModeEnabled_DeviceNotSupported() {
-        setFlagsRule.enableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE)
-        ExtendedMockito.doReturn(false).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
+        `when`(DesktopModeStatus.canEnterDesktopMode(any())).thenReturn(false)
 
         val taskContainer = createTaskContainer(createTask())
 
@@ -121,22 +96,7 @@
     }
 
     @Test
-    fun createDesktopTaskShortcutFactory_desktopModeEnabled_DeviceNotSupported_OverrideEnabled() {
-        setFlagsRule.enableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE)
-        ExtendedMockito.doReturn(false).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
-        ExtendedMockito.doReturn(false).`when` { DesktopModeStatus.enforceDeviceRestrictions() }
-
-        val taskContainer = spy(createTaskContainer(createTask()))
-        doReturn(workspaceItemInfo).whenever(taskContainer).itemInfo
-
-        val shortcuts = factory.getShortcuts(launcher, taskContainer)
-        assertThat(shortcuts).isNotNull()
-    }
-
-    @Test
     fun createDesktopTaskShortcutFactory_undockable() {
-        setFlagsRule.enableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE)
-
         val unDockableTask = createTask().apply { isDockable = false }
         val taskContainer = createTaskContainer(unDockableTask)
 
@@ -146,8 +106,6 @@
 
     @Test
     fun desktopSystemShortcutClicked() {
-        setFlagsRule.enableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE)
-
         val task = createTask()
         val taskContainer = spy(createTaskContainer(task))
 
@@ -159,13 +117,14 @@
             val successCallback = it.getArgument<Runnable>(2)
             successCallback.run()
         }
-        doReturn(workspaceItemInfo).whenever(taskContainer).itemInfo
+        val taskViewItemInfo = mock<TaskViewItemInfo>()
+        doReturn(taskViewItemInfo).whenever(taskContainer).itemInfo
 
         val shortcuts = factory.getShortcuts(launcher, taskContainer)
-        assertThat(shortcuts).hasSize(1)
-        assertThat(shortcuts!!.first()).isInstanceOf(DesktopSystemShortcut::class.java)
+        assertThat(shortcuts).isNotNull()
+        assertThat(shortcuts!!.single()).isInstanceOf(DesktopSystemShortcut::class.java)
 
-        val desktopShortcut = shortcuts.first() as DesktopSystemShortcut
+        val desktopShortcut = shortcuts.single() as DesktopSystemShortcut
 
         desktopShortcut.onClick(taskView)
 
@@ -178,30 +137,24 @@
                 eq(DesktopModeTransitionSource.APP_FROM_OVERVIEW),
                 any(),
             )
-        verify(statsLogger).withItemInfo(workspaceItemInfo)
+        verify(statsLogger).withItemInfo(taskViewItemInfo)
         verify(statsLogger).log(LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_DESKTOP_TAP)
     }
 
-    private fun createTask(): Task {
-        return Task(TaskKey(1, 0, Intent(), ComponentName("", ""), 0, 2000)).apply {
-            isDockable = true
-        }
-    }
+    private fun createTask() =
+        Task(TaskKey(1, 0, Intent(), ComponentName("", ""), 0, 2000)).apply { isDockable = true }
 
-    private fun createTaskContainer(task: Task): TaskContainer {
-        val snapshotView =
-            if (enableRefactorTaskThumbnail()) mock<TaskThumbnailView>()
-            else mock<TaskThumbnailViewDeprecated>()
-        return TaskContainer(
+    private fun createTaskContainer(task: Task) =
+        TaskContainer(
             taskView,
             task,
-            snapshotView,
-            iconView,
-            transformingTouchDelegate,
+            if (enableRefactorTaskThumbnail()) mock<TaskThumbnailView>()
+            else mock<TaskThumbnailViewDeprecated>(),
+            mock<TaskViewIcon>(),
+            mock<TransformingTouchDelegate>(),
             SplitConfigurationOptions.STAGE_POSITION_UNDEFINED,
             digitalWellBeingToast = null,
             showWindowsView = null,
             overlayFactory,
         )
-    }
 }
diff --git a/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt b/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt
index 8968b9c..9c2c13c 100644
--- a/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt
@@ -18,10 +18,8 @@
 
 import android.content.ComponentName
 import android.content.Intent
-import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
-import com.android.dx.mockito.inline.extended.ExtendedMockito
 import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
 import com.android.dx.mockito.inline.extended.StaticMockitoSession
 import com.android.launcher3.AbstractFloatingView
@@ -29,7 +27,7 @@
 import com.android.launcher3.Flags.enableRefactorTaskThumbnail
 import com.android.launcher3.logging.StatsLogManager
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent
-import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.model.data.TaskViewItemInfo
 import com.android.launcher3.uioverrides.QuickstepLauncher
 import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.launcher3.util.TransformingTouchDelegate
@@ -49,6 +47,7 @@
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
+import org.mockito.Mockito.`when`
 import org.mockito.kotlin.any
 import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.eq
@@ -58,7 +57,7 @@
 import org.mockito.kotlin.whenever
 import org.mockito.quality.Strictness
 
-/** Test for ExternalDisplaySystemShortcut */
+/** Test for [ExternalDisplaySystemShortcut] */
 class ExternalDisplaySystemShortcutTest {
 
     @get:Rule val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
@@ -68,14 +67,10 @@
     private val statsLogger: StatsLogManager.StatsLogger = mock()
     private val recentsView: LauncherRecentsView = mock()
     private val taskView: TaskView = mock()
-    private val workspaceItemInfo: WorkspaceItemInfo = mock()
     private val abstractFloatingViewHelper: AbstractFloatingViewHelper = mock()
-    private val iconView: TaskViewIcon = mock()
-    private val transformingTouchDelegate: TransformingTouchDelegate = mock()
+    private val overlayFactory: TaskOverlayFactory = mock()
     private val factory: TaskShortcutFactory =
         ExternalDisplaySystemShortcut.createFactory(abstractFloatingViewHelper)
-    private val overlayFactory: TaskOverlayFactory = mock()
-    private val overlay: TaskOverlay<*> = mock()
 
     private lateinit var mockitoSession: StaticMockitoSession
 
@@ -84,11 +79,10 @@
         mockitoSession =
             mockitoSession()
                 .strictness(Strictness.LENIENT)
-                .spyStatic(DesktopModeStatus::class.java)
+                .mockStatic(DesktopModeStatus::class.java)
                 .startMocking()
-        ExtendedMockito.doReturn(true).`when` { DesktopModeStatus.enforceDeviceRestrictions() }
-        ExtendedMockito.doReturn(true).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
-        whenever(overlayFactory.createOverlay(any())).thenReturn(overlay)
+        whenever(DesktopModeStatus.canEnterDesktopMode(any())).thenReturn(true)
+        whenever(overlayFactory.createOverlay(any())).thenReturn(mock<TaskOverlay<*>>())
     }
 
     @After
@@ -97,23 +91,9 @@
     }
 
     @Test
-    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE)
     @EnableFlags(Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT)
     fun createExternalDisplayTaskShortcut_desktopModeDisabled() {
-        val task = createTask()
-        val taskContainer = createTaskContainer(task)
-
-        val shortcuts = factory.getShortcuts(launcher, taskContainer)
-        assertThat(shortcuts).isNull()
-    }
-
-    @Test
-    @EnableFlags(
-        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE,
-        Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT,
-    )
-    fun createExternalDisplayTaskShortcut_desktopModeEnabled_deviceNotSupported() {
-        ExtendedMockito.doReturn(false).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
+        `when`(DesktopModeStatus.canEnterDesktopMode(any())).thenReturn(false)
 
         val taskContainer = createTaskContainer(createTask())
 
@@ -122,26 +102,7 @@
     }
 
     @Test
-    @EnableFlags(
-        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE,
-        Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT,
-    )
-    fun createExternalDisplayTaskShortcut_desktopModeEnabled_deviceNotSupported_overrideEnabled() {
-        ExtendedMockito.doReturn(false).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
-        ExtendedMockito.doReturn(false).`when` { DesktopModeStatus.enforceDeviceRestrictions() }
-
-        val taskContainer = spy(createTaskContainer(createTask()))
-        doReturn(workspaceItemInfo).whenever(taskContainer).itemInfo
-
-        val shortcuts = factory.getShortcuts(launcher, taskContainer)
-        assertThat(shortcuts).isNotNull()
-    }
-
-    @Test
-    @EnableFlags(
-        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE,
-        Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT,
-    )
+    @EnableFlags(Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT)
     fun externalDisplaySystemShortcutClicked() {
         val task = createTask()
         val taskContainer = spy(createTaskContainer(task))
@@ -154,7 +115,8 @@
             val successCallback = it.getArgument<Runnable>(1)
             successCallback.run()
         }
-        doReturn(workspaceItemInfo).whenever(taskContainer).itemInfo
+        val taskViewItemInfo = mock<TaskViewItemInfo>()
+        doReturn(taskViewItemInfo).whenever(taskContainer).itemInfo
 
         val shortcuts = factory.getShortcuts(launcher, taskContainer)
         assertThat(shortcuts).hasSize(1)
@@ -168,26 +130,23 @@
             AbstractFloatingView.TYPE_ALL and AbstractFloatingView.TYPE_REBIND_SAFE.inv()
         verify(abstractFloatingViewHelper).closeOpenViews(launcher, true, allTypesExceptRebindSafe)
         verify(recentsView).moveTaskToExternalDisplay(eq(taskContainer), any())
-        verify(statsLogger).withItemInfo(workspaceItemInfo)
+        verify(statsLogger).withItemInfo(taskViewItemInfo)
         verify(statsLogger).log(LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_EXTERNAL_DISPLAY_TAP)
     }
 
-    private fun createTask(): Task = Task(TaskKey(1, 0, Intent(), ComponentName("", ""), 0, 2000))
+    private fun createTask() = Task(TaskKey(1, 0, Intent(), ComponentName("", ""), 0, 2000))
 
-    private fun createTaskContainer(task: Task): TaskContainer {
-        val snapshotView =
-            if (enableRefactorTaskThumbnail()) mock<TaskThumbnailView>()
-            else mock<TaskThumbnailViewDeprecated>()
-        return TaskContainer(
+    private fun createTaskContainer(task: Task) =
+        TaskContainer(
             taskView,
             task,
-            snapshotView,
-            iconView,
-            transformingTouchDelegate,
+            if (enableRefactorTaskThumbnail()) mock<TaskThumbnailView>()
+            else mock<TaskThumbnailViewDeprecated>(),
+            mock<TaskViewIcon>(),
+            mock<TransformingTouchDelegate>(),
             SplitConfigurationOptions.STAGE_POSITION_UNDEFINED,
             digitalWellBeingToast = null,
             showWindowsView = null,
             overlayFactory,
         )
-    }
 }