Merge "Fixed bubble bar size for the 3 buttons navigation mode." into main
diff --git a/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransition.kt b/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransition.kt
index e160f82..87a82f0 100644
--- a/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransition.kt
+++ b/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransition.kt
@@ -143,6 +143,7 @@
     }
 
     companion object {
+        /** Change modes that represent a task becoming visible / launching in Desktop mode. */
         val LAUNCH_CHANGE_MODES = intArrayOf(TRANSIT_OPEN, TRANSIT_TO_FRONT)
 
         private val launchBoundsAnimationDef =
diff --git a/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransitionManager.kt b/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransitionManager.kt
new file mode 100644
index 0000000..e32bcd1
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransitionManager.kt
@@ -0,0 +1,88 @@
+/*
+ * 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.desktop
+
+import android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.content.Context
+import android.window.DesktopModeFlags
+import android.window.RemoteTransition
+import android.window.TransitionFilter
+import android.window.TransitionFilter.CONTAINER_ORDER_TOP
+import com.android.launcher3.desktop.DesktopAppLaunchTransition.AppLaunchType
+import com.android.launcher3.util.Executors.MAIN_EXECUTOR
+import com.android.quickstep.SystemUiProxy
+import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
+
+/** Manages transitions related to app launches in Desktop Mode. */
+class DesktopAppLaunchTransitionManager(
+    private val context: Context,
+    private val systemUiProxy: SystemUiProxy,
+) {
+    private var remoteWindowLimitUnminimizeTransition: RemoteTransition? = null
+
+    /**
+     * Register a [RemoteTransition] supporting Desktop app launches, and window limit
+     * minimizations.
+     */
+    fun registerTransitions() {
+        if (!shouldRegisterTransitions()) {
+            return
+        }
+        remoteWindowLimitUnminimizeTransition =
+            RemoteTransition(
+                DesktopAppLaunchTransition(context, MAIN_EXECUTOR, AppLaunchType.UNMINIMIZE)
+            )
+        systemUiProxy.registerRemoteTransition(
+            remoteWindowLimitUnminimizeTransition,
+            buildAppLaunchFilter(),
+        )
+    }
+
+    /**
+     * Unregister the [RemoteTransition] supporting Desktop app launches and window limit
+     * minimizations.
+     */
+    fun unregisterTransitions() {
+        if (!shouldRegisterTransitions()) {
+            return
+        }
+        systemUiProxy.unregisterRemoteTransition(remoteWindowLimitUnminimizeTransition)
+        remoteWindowLimitUnminimizeTransition = null
+    }
+
+    private fun shouldRegisterTransitions(): Boolean =
+        DesktopModeStatus.canEnterDesktopMode(context) &&
+            DesktopModeFlags.ENABLE_DESKTOP_APP_LAUNCH_TRANSITIONS.isTrue
+
+    companion object {
+        private fun buildAppLaunchFilter(): TransitionFilter {
+            val openRequirement =
+                TransitionFilter.Requirement().apply {
+                    mActivityType = ACTIVITY_TYPE_STANDARD
+                    mWindowingMode = WINDOWING_MODE_FREEFORM
+                    mModes = DesktopAppLaunchTransition.LAUNCH_CHANGE_MODES
+                    mMustBeTask = true
+                    mOrder = CONTAINER_ORDER_TOP
+                }
+            return TransitionFilter().apply {
+                mTypeSet = DesktopAppLaunchTransition.LAUNCH_CHANGE_MODES
+                mRequirements = arrayOf(openRequirement)
+            }
+        }
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
index b1cb2c6..4a94be7 100644
--- a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
@@ -362,20 +362,21 @@
             // This method can be called before init() is called.
             return;
         }
-        if (mControllers.uiController.isIconAlignedWithHotseat()
-                && !mTaskbarLauncherStateController.isAnimatingToLauncher()) {
-            // Only animate the nav buttons while home and not animating home, otherwise let
-            // the TaskbarViewController handle it.
-            mControllers.navbarButtonsViewController
-                    .getTaskbarNavButtonTranslationYForInAppDisplay()
-                    .updateValue(mLauncher.getDeviceProfile().getTaskbarOffsetY()
-                            * mTaskbarInAppDisplayProgress.value);
-            mControllers.navbarButtonsViewController
-                    .getOnTaskbarBackgroundNavButtonColorOverride().updateValue(progress);
-
+        if (mControllers.uiController.isIconAlignedWithHotseat()) {
+            if (!mTaskbarLauncherStateController.isAnimatingToLauncher()) {
+                // Only animate the nav buttons while home and not animating home, otherwise let
+                // the TaskbarViewController handle it.
+                mControllers.navbarButtonsViewController
+                        .getTaskbarNavButtonTranslationYForInAppDisplay()
+                        .updateValue(mLauncher.getDeviceProfile().getTaskbarOffsetY()
+                                * mTaskbarInAppDisplayProgress.value);
+                mControllers.navbarButtonsViewController
+                        .getOnTaskbarBackgroundNavButtonColorOverride().updateValue(progress);
+            }
             if (isBubbleBarEnabled()) {
                 mControllers.bubbleControllers.ifPresent(
-                        c -> c.bubbleStashController.setInAppDisplayOverrideProgress(progress));
+                        c -> c.bubbleStashController.setInAppDisplayOverrideProgress(
+                                mTaskbarInAppDisplayProgress.value));
             }
         }
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashController.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashController.kt
index 69da13d..45f5568 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashController.kt
@@ -120,7 +120,10 @@
             if (field == value) return
             field = value
             if (launcherState == BubbleLauncherState.HOME) {
-                bubbleBarViewController.bubbleBarTranslationY.updateValue(bubbleBarTranslationY)
+                if (bubbleBarTranslationYAnimator.isAnimating) {
+                    bubbleBarTranslationYAnimator.cancelAnimation()
+                }
+                bubbleBarTranslationYAnimator.updateValue(bubbleBarTranslationY)
                 if (value == 0f || value == 1f) {
                     // Update insets only when we reach the end values
                     taskbarInsetsController.onTaskbarOrBubblebarWindowHeightOrInsetsChanged()
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index fcc5121..ad5720f 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -85,6 +85,7 @@
 import com.android.launcher3.Flags;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.anim.AnimatedFloat;
+import com.android.launcher3.desktop.DesktopAppLaunchTransitionManager;
 import com.android.launcher3.provider.RestoreDbTask;
 import com.android.launcher3.statehandlers.DesktopVisibilityController;
 import com.android.launcher3.statemanager.StatefulActivity;
@@ -662,6 +663,7 @@
     private NavigationMode mGestureStartNavMode = null;
 
     private DesktopVisibilityController mDesktopVisibilityController;
+    private DesktopAppLaunchTransitionManager mDesktopAppLaunchTransitionManager;
 
     @Override
     public void onCreate() {
@@ -686,6 +688,9 @@
         mDesktopVisibilityController = new DesktopVisibilityController(this);
         mTaskbarManager = new TaskbarManager(
                 this, mAllAppsActionManager, mNavCallbacks, mDesktopVisibilityController);
+        mDesktopAppLaunchTransitionManager =
+                new DesktopAppLaunchTransitionManager(this, SystemUiProxy.INSTANCE.get(this));
+        mDesktopAppLaunchTransitionManager.registerTransitions();
         if (Flags.enableLauncherOverviewInWindow() || Flags.enableFallbackOverviewInWindow()) {
             mRecentsWindowManager = new RecentsWindowManager(this);
         }
@@ -851,6 +856,10 @@
         if (mRecentsWindowManager != null) {
             mRecentsWindowManager.destroy();
         }
+        if (mDesktopAppLaunchTransitionManager != null) {
+            mDesktopAppLaunchTransitionManager.unregisterTransitions();
+        }
+        mDesktopAppLaunchTransitionManager = null;
         mDesktopVisibilityController.onDestroy();
         sConnected = false;
 
diff --git a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
index 275af00..6c627ef 100644
--- a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
@@ -73,26 +73,29 @@
         getTaskDataById(taskId).map { it?.thumbnail }.distinctUntilChangedBy { it?.snapshotId }
 
     override fun setVisibleTasks(visibleTaskIdList: Set<Int>) {
-        // Remove tasks are no longer visible
         val tasksNoLongerVisible = taskRequests.keys.subtract(visibleTaskIdList)
+        val newlyVisibleTasks = visibleTaskIdList.subtract(taskRequests.keys)
+        if (tasksNoLongerVisible.isNotEmpty() || newlyVisibleTasks.isNotEmpty()) {
+            Log.d(
+                TAG,
+                "setVisibleTasks to: $visibleTaskIdList, " +
+                    "removed: $tasksNoLongerVisible, added: $newlyVisibleTasks",
+            )
+        }
+
+        // Remove tasks are no longer visible
         removeTasks(tasksNoLongerVisible)
         // Add new tasks to be requested
-        val newlyVisibleTasks = visibleTaskIdList.subtract(taskRequests.keys)
         newlyVisibleTasks.forEach { taskId -> requestTaskData(taskId) }
-
-        if (tasksNoLongerVisible.isNotEmpty() || newlyVisibleTasks.isNotEmpty()) {
-            Log.d(TAG, "setVisibleTasks to: $visibleTaskIdList, " +
-                    "removed: $tasksNoLongerVisible, added: $newlyVisibleTasks")
-        }
     }
 
     private fun requestTaskData(taskId: Int) {
-        Log.i(TAG, "requestTaskData: $taskId")
         val task = tasks.value[taskId] ?: return
         taskRequests[taskId] =
             Pair(
                 task.key,
                 recentsCoroutineScope.launch {
+                    Log.i(TAG, "requestTaskData: $taskId")
                     fetchIcon(task)
                     fetchThumbnail(task)
                 },
@@ -102,8 +105,8 @@
     private fun removeTasks(tasksToRemove: Set<Int>) {
         if (tasksToRemove.isEmpty()) return
 
+        Log.i(TAG, "removeTasks: $tasksToRemove")
         tasksToRemove.forEach { taskId ->
-            Log.i(TAG, "removeTask: $taskId")
             val request = taskRequests.remove(taskId) ?: return
             val (taskKey, job) = request
             job.cancel()
diff --git a/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt b/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
index f2b9976..dd11d48 100644
--- a/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
+++ b/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
@@ -19,6 +19,7 @@
 import android.content.Context
 import android.util.Log
 import android.view.View
+import com.android.launcher3.util.coroutines.DispatcherProvider
 import com.android.launcher3.util.coroutines.ProductionDispatchers
 import com.android.quickstep.RecentsModel
 import com.android.quickstep.recents.data.RecentTasksRepository
@@ -63,6 +64,7 @@
             val recentsCoroutineScope =
                 CoroutineScope(SupervisorJob() + Dispatchers.Main + CoroutineName("RecentsView"))
             set(CoroutineScope::class.java.simpleName, recentsCoroutineScope)
+            set(DispatcherProvider::class.java.simpleName, ProductionDispatchers)
             val recentsModel = RecentsModel.INSTANCE.get(appContext)
             val taskVisualsChangedDelegate =
                 TaskVisualsChangedDelegateImpl(
@@ -196,6 +198,7 @@
                         recentsViewData = inject(),
                         taskViewData = inject(scopeId, extras),
                         taskContainerData = inject(scopeId),
+                        dispatcherProvider = inject(),
                         getThumbnailPositionUseCase = inject(),
                         tasksRepository = inject(),
                         splashAlphaUseCase = inject(scopeId),
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
index eb9c047..a8c8659 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
@@ -33,7 +33,6 @@
 import com.android.launcher3.util.ViewPool
 import com.android.quickstep.recents.di.RecentsDependencies
 import com.android.quickstep.recents.di.get
-import com.android.quickstep.recents.di.inject
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.BackgroundOnly
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Snapshot
@@ -43,7 +42,6 @@
 import com.android.quickstep.util.TaskCornerRadius
 import com.android.quickstep.views.FixedSizeImageView
 import com.android.systemui.shared.system.QuickStepContract
-import kotlin.math.abs
 import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -53,8 +51,7 @@
 import kotlinx.coroutines.flow.onEach
 
 class TaskThumbnailView : ConstraintLayout, ViewPool.Reusable {
-
-    private val viewData: TaskThumbnailViewData by RecentsDependencies.inject(this)
+    private lateinit var viewData: TaskThumbnailViewData
     private lateinit var viewModel: TaskThumbnailViewModel
 
     private lateinit var viewAttachedScope: CoroutineScope
@@ -92,10 +89,12 @@
         super.onAttachedToWindow()
         viewAttachedScope =
             CoroutineScope(SupervisorJob() + Dispatchers.Main + CoroutineName("TaskThumbnailView"))
+        viewData = RecentsDependencies.get(this)
+        updateViewDataValues()
         viewModel = RecentsDependencies.get(this)
         viewModel.uiState
             .onEach { viewModelUiState ->
-                Log.d(TAG, "viewModelUiState changed from $uiState to: $viewModelUiState")
+                Log.d(TAG, "viewModelUiState changed from: $uiState to: $viewModelUiState")
                 uiState = viewModelUiState
                 resetViews()
                 when (viewModelUiState) {
@@ -144,11 +143,15 @@
     override fun onLayout(changed: Boolean, left: Int, top: Int, right: Int, bottom: Int) {
         super.onLayout(changed, left, top, right, bottom)
         if (changed) {
-            viewData.width.value = abs(right - left)
-            viewData.height.value = abs(bottom - top)
+            updateViewDataValues()
         }
     }
 
+    private fun updateViewDataValues() {
+        viewData.width.value = width
+        viewData.height.value = height
+    }
+
     override fun onSizeChanged(w: Int, h: Int, oldw: Int, oldh: Int) {
         super.onSizeChanged(w, h, oldw, oldh)
         if (uiState is SnapshotSplash) {
diff --git a/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt b/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
index c82ed9a..203177a 100644
--- a/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
+++ b/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
@@ -85,14 +85,18 @@
     }
 
     private fun initOverlay(enabledState: Enabled) {
-        Log.d(TAG, "initOverlay - taskId: ${task.key.id}, thumbnail: ${enabledState.thumbnail}")
+        if (DEBUG) {
+            Log.d(TAG, "initOverlay - taskId: ${task.key.id}, thumbnail: ${enabledState.thumbnail}")
+        }
         with(getThumbnailPositionState()) {
             overlay.initOverlay(task, enabledState.thumbnail, matrix, isRotated)
         }
     }
 
     private fun reset() {
-        Log.d(TAG, "reset - taskId: ${task.key.id}")
+        if (DEBUG) {
+            Log.d(TAG, "reset - taskId: ${task.key.id}")
+        }
         overlay.reset()
     }
 
@@ -105,5 +109,6 @@
 
     companion object {
         private const val TAG = "TaskOverlayHelper"
+        private const val DEBUG = false
     }
 }
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
index bd47cec..8b15a82 100644
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
@@ -21,6 +21,7 @@
 import android.graphics.Matrix
 import android.util.Log
 import androidx.core.graphics.ColorUtils
+import com.android.launcher3.util.coroutines.DispatcherProvider
 import com.android.quickstep.recents.data.RecentTasksRepository
 import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
 import com.android.quickstep.recents.usecase.ThumbnailPositionState
@@ -42,6 +43,7 @@
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.runBlocking
 
@@ -50,6 +52,7 @@
     recentsViewData: RecentsViewData,
     taskViewData: TaskViewData,
     taskContainerData: TaskContainerData,
+    dispatcherProvider: DispatcherProvider,
     private val tasksRepository: RecentTasksRepository,
     private val getThumbnailPositionUseCase: GetThumbnailPositionUseCase,
     private val splashAlphaUseCase: SplashAlphaUseCase,
@@ -73,6 +76,7 @@
             tintAmount ->
             max(taskMenuOpenProgress * MAX_SCRIM_ALPHA, tintAmount)
         }
+
     override val splashAlpha = splashProgress.flatMapLatest { it }
 
     private val isLiveTile =
@@ -84,6 +88,7 @@
                 runningTaskIds.contains(taskId) && !runningTaskShowScreenshot
             }
             .distinctUntilChanged()
+            .flowOn(dispatcherProvider.default)
 
     override val uiState: Flow<TaskThumbnailUiState> =
         combine(task.flatMapLatest { it }, isLiveTile) { taskVal, isRunning ->
@@ -105,6 +110,7 @@
                 }
             }
             .distinctUntilChanged()
+            .flowOn(dispatcherProvider.default)
 
     override fun bind(taskId: Int) {
         Log.d(TAG, "bind taskId: $taskId")
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 839c42e..3a4e328 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -3694,31 +3694,43 @@
         int currentPageScroll = getScrollForPage(mCurrentPage);
         int lastGridTaskScroll = getScrollForPage(indexOfChild(lastGridTaskView));
         boolean currentPageSnapsToEndOfGrid = currentPageScroll == lastGridTaskScroll;
+
+        int topGridRowSize = mTopRowIdSet.size();
+        int numLargeTiles = mUtils.getLargeTileCount(getTaskViews());
+        int bottomGridRowSize = taskCount - mTopRowIdSet.size() - numLargeTiles;
+        boolean topRowLonger = topGridRowSize > bottomGridRowSize;
+        boolean bottomRowLonger = bottomGridRowSize > topGridRowSize;
+        boolean dismissedTaskFromTop = mTopRowIdSet.contains(dismissedTaskViewId);
+        boolean dismissedTaskFromBottom = !dismissedTaskFromTop && !isFocusedTaskDismissed;
+        if (dismissedTaskFromTop || (isFocusedTaskDismissed && nextFocusedTaskFromTop)) {
+            topGridRowSize--;
+        }
+        if (dismissedTaskFromBottom || (isFocusedTaskDismissed && !nextFocusedTaskFromTop)) {
+            bottomGridRowSize--;
+        }
+        int longRowWidth = Math.max(topGridRowSize, bottomGridRowSize)
+                * (mLastComputedGridTaskSize.width() + mPageSpacing);
+        if (!enableGridOnlyOverview() && !isStagingFocusedTask) {
+            longRowWidth += mLastComputedTaskSize.width() + mPageSpacing;
+        }
+        // Compensate the removed gap if we don't already have shortTotalCompensation,
+        // and adjust accordingly to the new shortTotalCompensation after dismiss.
+        int newClearAllShortTotalWidthTranslation = 0;
+        if (mClearAllShortTotalWidthTranslation == 0) {
+            // If first task is not in the expected position (mLastComputedTaskSize) and being too
+            // close  to ClearAllButton, then apply extra translation to ClearAllButton.
+            int firstTaskStart = mLastComputedGridSize.left + longRowWidth;
+            int expectedFirstTaskStart = mLastComputedTaskSize.right;
+            if (firstTaskStart < expectedFirstTaskStart) {
+                newClearAllShortTotalWidthTranslation = expectedFirstTaskStart - firstTaskStart;
+            }
+        }
         if (lastGridTaskView != null && lastGridTaskView.isVisibleToUser()) {
             // After dismissal, animate translation of the remaining tasks to fill any gap left
             // between the end of the grid and the clear all button. Only animate if the clear
             // all button is visible or would become visible after dismissal.
             float longGridRowWidthDiff = 0;
 
-            int topGridRowSize = mTopRowIdSet.size();
-            int numLargeTiles = mUtils.getLargeTileCount(getTaskViews());
-            int bottomGridRowSize = taskCount - mTopRowIdSet.size() - numLargeTiles;
-            boolean topRowLonger = topGridRowSize > bottomGridRowSize;
-            boolean bottomRowLonger = bottomGridRowSize > topGridRowSize;
-            boolean dismissedTaskFromTop = mTopRowIdSet.contains(dismissedTaskViewId);
-            boolean dismissedTaskFromBottom = !dismissedTaskFromTop && !isFocusedTaskDismissed;
-            if (dismissedTaskFromTop || (isFocusedTaskDismissed && nextFocusedTaskFromTop)) {
-                topGridRowSize--;
-            }
-            if (dismissedTaskFromBottom || (isFocusedTaskDismissed && !nextFocusedTaskFromTop)) {
-                bottomGridRowSize--;
-            }
-            int longRowWidth = Math.max(topGridRowSize, bottomGridRowSize)
-                    * (mLastComputedGridTaskSize.width() + mPageSpacing);
-            if (!enableGridOnlyOverview() && !isStagingFocusedTask) {
-                longRowWidth += mLastComputedTaskSize.width() + mPageSpacing;
-            }
-
             float gapWidth = 0;
             if ((topRowLonger && dismissedTaskFromTop)
                     || (bottomRowLonger && dismissedTaskFromBottom)) {
@@ -3730,17 +3742,6 @@
             }
             if (gapWidth > 0) {
                 if (mClearAllShortTotalWidthTranslation == 0) {
-                    // Compensate the removed gap if we don't already have shortTotalCompensation,
-                    // and adjust accordingly to the new shortTotalCompensation after dismiss.
-                    int newClearAllShortTotalWidthTranslation = 0;
-                    if (longRowWidth < mLastComputedGridSize.width()) {
-                        DeviceProfile deviceProfile = mContainer.getDeviceProfile();
-                        newClearAllShortTotalWidthTranslation =
-                                (mIsRtl
-                                        ? mLastComputedTaskSize.right
-                                        : deviceProfile.widthPx - mLastComputedTaskSize.left)
-                                        - longRowWidth - deviceProfile.overviewGridSideMargin;
-                    }
                     float gapCompensation = gapWidth - newClearAllShortTotalWidthTranslation;
                     longGridRowWidthDiff += mIsRtl ? -gapCompensation : gapCompensation;
                 }
@@ -3830,6 +3831,8 @@
                     : mUtils.getDesktopTaskViewCount(getTaskViews());
             stagingTranslation = getPagedOrientationHandler().getPrimaryScroll(this)
                     - getScrollForPage(nextSnappedPage);
+            stagingTranslation += mIsRtl ? newClearAllShortTotalWidthTranslation
+                    : -newClearAllShortTotalWidthTranslation;
         }
         for (int i = 0; i < count; i++) {
             View child = getChildAt(i);
@@ -3888,7 +3891,7 @@
                     anim.setFloat(taskView, TaskView.DISMISS_SCALE, scale,
                             clampToProgress(LINEAR, animationStartProgress,
                                     dismissTranslationInterpolationEnd));
-                    primaryTranslation += mIsRtl ? dismissedTaskWidth : -dismissedTaskWidth;
+                    primaryTranslation += dismissedTaskWidth;
                     animationEndProgress = dismissTranslationInterpolationEnd;
                     float secondaryTranslation = -mTaskGridVerticalDiff;
                     if (!nextFocusedTaskFromTop) {
diff --git a/quickstep/src/com/android/quickstep/views/RecentsViewModelHelper.kt b/quickstep/src/com/android/quickstep/views/RecentsViewModelHelper.kt
index f22c672..3616fbb 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsViewModelHelper.kt
+++ b/quickstep/src/com/android/quickstep/views/RecentsViewModelHelper.kt
@@ -25,6 +25,7 @@
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 
 /** Helper for [RecentsView] to interact with the [RecentsViewModel]. */
 class RecentsViewModelHelper(private val recentsViewModel: RecentsViewModel) {
@@ -32,7 +33,7 @@
 
     fun onAttachedToWindow() {
         viewAttachedScope =
-            CoroutineScope(SupervisorJob() + Dispatchers.Main + CoroutineName("RecentsView"))
+            CoroutineScope(SupervisorJob() + Dispatchers.Default + CoroutineName("RecentsView"))
     }
 
     fun onDetachedFromWindow() {
@@ -50,7 +51,7 @@
         viewAttachedScope.launch {
             recentsViewModel.waitForRunningTaskShowScreenshotToUpdate()
             recentsViewModel.waitForThumbnailsToUpdate(updatedThumbnails)
-            ViewUtils.postFrameDrawn(taskView, onFinishRunnable)
+            withContext(Dispatchers.Main) { 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 25aba39..c940fb4 100644
--- a/quickstep/src/com/android/quickstep/views/TaskContainer.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
@@ -110,7 +110,6 @@
             return snapshotView as TaskThumbnailViewDeprecated
         }
 
-    // TODO(b/334826842): Support shouldShowSplashView for new TTV.
     val shouldShowSplashView: Boolean
         get() =
             if (enableRefactorTaskThumbnail())
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index 0760618..b1cb407 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -1221,8 +1221,6 @@
                     if (isQuickSwitch) {
                         setFreezeRecentTasksReordering()
                     }
-                    // TODO(b/334826842) no work required - add splash functionality to new TTV -
-                    // cold start e.g. restart device. Small splash moving to bigger splash
                     disableStartingWindow = firstContainer.shouldShowSplashView
                 }
         Executors.UI_HELPER_EXECUTOR.execute {
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashControllerTest.kt
index 9580d50..f795ab1 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/stashing/PersistentBubbleStashControllerTest.kt
@@ -298,6 +298,21 @@
     }
 
     @Test
+    fun inAppDisplayOverrideProgress_onHome_cancelExistingAnimation() {
+        whenever(bubbleBarViewController.hasBubbles()).thenReturn(false)
+        persistentTaskBarStashController.launcherState = BubbleLauncherState.HOME
+
+        bubbleBarViewController.bubbleBarTranslationY.animateToValue(100f)
+        advanceTimeBy(10)
+        assertThat(bubbleBarViewController.bubbleBarTranslationY.isAnimating).isTrue()
+
+        getInstrumentation().runOnMainSync {
+            persistentTaskBarStashController.inAppDisplayOverrideProgress = 0.5f
+        }
+        assertThat(bubbleBarViewController.bubbleBarTranslationY.isAnimating).isFalse()
+    }
+
+    @Test
     fun inAppDisplayProgressUpdate_inApp_noTranslationUpdate() {
         whenever(bubbleBarViewController.hasBubbles()).thenReturn(false)
         persistentTaskBarStashController.launcherState = BubbleLauncherState.IN_APP
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt
index c541d3d..e3a6adf 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt
@@ -24,6 +24,7 @@
 import android.graphics.drawable.Drawable
 import android.view.Surface
 import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.launcher3.util.TestDispatcherProvider
 import com.android.quickstep.recents.data.FakeTasksRepository
 import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
 import com.android.quickstep.recents.usecase.ThumbnailPositionState.MatrixScaling
@@ -42,6 +43,8 @@
 import com.android.systemui.shared.recents.model.ThumbnailData
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -51,18 +54,24 @@
 /** Test for [TaskThumbnailView] */
 @RunWith(AndroidJUnit4::class)
 class TaskThumbnailViewModelImplTest {
+    private val dispatcher = StandardTestDispatcher()
+    private val testScope = TestScope(dispatcher)
+
     private var taskViewType = TaskViewType.SINGLE
     private val recentsViewData = RecentsViewData()
     private val taskViewData by lazy { TaskViewData(taskViewType) }
     private val taskContainerData = TaskContainerData()
+    private val dispatcherProvider = TestDispatcherProvider(dispatcher)
     private val tasksRepository = FakeTasksRepository()
     private val mGetThumbnailPositionUseCase = mock<GetThumbnailPositionUseCase>()
     private val splashAlphaUseCase: SplashAlphaUseCase = mock()
+
     private val systemUnderTest by lazy {
         TaskThumbnailViewModelImpl(
             recentsViewData,
             taskViewData,
             taskContainerData,
+            dispatcherProvider,
             tasksRepository,
             mGetThumbnailPositionUseCase,
             splashAlphaUseCase,
@@ -72,81 +81,85 @@
     private val tasks = (0..5).map(::createTaskWithId)
 
     @Test
-    fun initialStateIsUninitialized() = runTest {
-        assertThat(systemUnderTest.uiState.first()).isEqualTo(Uninitialized)
-    }
+    fun initialStateIsUninitialized() =
+        testScope.runTest { assertThat(systemUnderTest.uiState.first()).isEqualTo(Uninitialized) }
 
     @Test
-    fun bindRunningTask_thenStateIs_LiveTile() = runTest {
-        val taskId = 1
-        tasksRepository.seedTasks(tasks)
-        tasksRepository.setVisibleTasks(setOf(taskId))
-        recentsViewData.runningTaskIds.value = setOf(taskId)
-        systemUnderTest.bind(taskId)
+    fun bindRunningTask_thenStateIs_LiveTile() =
+        testScope.runTest {
+            val taskId = 1
+            tasksRepository.seedTasks(tasks)
+            tasksRepository.setVisibleTasks(setOf(taskId))
+            recentsViewData.runningTaskIds.value = setOf(taskId)
+            systemUnderTest.bind(taskId)
 
-        assertThat(systemUnderTest.uiState.first()).isEqualTo(LiveTile)
-    }
+            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 = mock<Drawable>()
-        tasksRepository.seedIconData(taskId, "Task $taskId", "", expectedIconData)
-        tasksRepository.seedTasks(tasks)
-        tasksRepository.setVisibleTasks(setOf(taskId))
-        recentsViewData.runningTaskIds.value = setOf(taskId)
-        recentsViewData.runningTaskShowScreenshot.value = true
-        systemUnderTest.bind(taskId)
+    fun bindRunningTaskShouldShowScreenshot_thenStateIs_SnapshotSplash() =
+        testScope.runTest {
+            val taskId = 1
+            val expectedThumbnailData = createThumbnailData()
+            tasksRepository.seedThumbnailData(mapOf(taskId to expectedThumbnailData))
+            val expectedIconData = mock<Drawable>()
+            tasksRepository.seedIconData(taskId, "Task $taskId", "", expectedIconData)
+            tasksRepository.seedTasks(tasks)
+            tasksRepository.setVisibleTasks(setOf(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,
-                    ),
-                    expectedIconData,
+            assertThat(systemUnderTest.uiState.first())
+                .isEqualTo(
+                    SnapshotSplash(
+                        Snapshot(
+                            backgroundColor = Color.rgb(1, 1, 1),
+                            bitmap = expectedThumbnailData.thumbnail!!,
+                            thumbnailRotation = Surface.ROTATION_0,
+                        ),
+                        expectedIconData,
+                    )
                 )
-            )
-    }
+        }
 
     @Test
-    fun setRecentsFullscreenProgress_thenCornerRadiusProgressIsPassedThrough() = runTest {
-        recentsViewData.fullscreenProgress.value = 0.5f
+    fun setRecentsFullscreenProgress_thenCornerRadiusProgressIsPassedThrough() =
+        testScope.runTest {
+            recentsViewData.fullscreenProgress.value = 0.5f
 
-        assertThat(systemUnderTest.cornerRadiusProgress.first()).isEqualTo(0.5f)
+            assertThat(systemUnderTest.cornerRadiusProgress.first()).isEqualTo(0.5f)
 
-        recentsViewData.fullscreenProgress.value = 0.6f
+            recentsViewData.fullscreenProgress.value = 0.6f
 
-        assertThat(systemUnderTest.cornerRadiusProgress.first()).isEqualTo(0.6f)
-    }
+            assertThat(systemUnderTest.cornerRadiusProgress.first()).isEqualTo(0.6f)
+        }
 
     @Test
-    fun setRecentsFullscreenProgress_thenCornerRadiusProgressIsConstantForDesktop() = runTest {
-        taskViewType = TaskViewType.DESKTOP
-        recentsViewData.fullscreenProgress.value = 0.5f
+    fun setRecentsFullscreenProgress_thenCornerRadiusProgressIsConstantForDesktop() =
+        testScope.runTest {
+            taskViewType = TaskViewType.DESKTOP
+            recentsViewData.fullscreenProgress.value = 0.5f
 
-        assertThat(systemUnderTest.cornerRadiusProgress.first()).isEqualTo(1f)
+            assertThat(systemUnderTest.cornerRadiusProgress.first()).isEqualTo(1f)
 
-        recentsViewData.fullscreenProgress.value = 0.6f
+            recentsViewData.fullscreenProgress.value = 0.6f
 
-        assertThat(systemUnderTest.cornerRadiusProgress.first()).isEqualTo(1f)
-    }
+            assertThat(systemUnderTest.cornerRadiusProgress.first()).isEqualTo(1f)
+        }
 
     @Test
-    fun setAncestorScales_thenScaleIsCalculated() = runTest {
-        recentsViewData.scale.value = 0.5f
-        taskViewData.scale.value = 0.6f
+    fun setAncestorScales_thenScaleIsCalculated() =
+        testScope.runTest {
+            recentsViewData.scale.value = 0.5f
+            taskViewData.scale.value = 0.6f
 
-        assertThat(systemUnderTest.inheritedScale.first()).isEqualTo(0.3f)
-    }
+            assertThat(systemUnderTest.inheritedScale.first()).isEqualTo(0.3f)
+        }
 
     @Test
     fun bindRunningTaskThenStoppedTaskWithoutThumbnail_thenStateChangesToBackgroundOnly() =
-        runTest {
+        testScope.runTest {
             val runningTaskId = 1
             val stoppedTaskId = 2
             tasksRepository.seedTasks(tasks)
@@ -161,125 +174,138 @@
         }
 
     @Test
-    fun bindStoppedTaskWithoutThumbnail_thenStateIs_BackgroundOnly_withAlphaRemoved() = runTest {
-        val stoppedTaskId = 2
-        tasksRepository.seedTasks(tasks)
-        tasksRepository.setVisibleTasks(setOf(stoppedTaskId))
+    fun bindStoppedTaskWithoutThumbnail_thenStateIs_BackgroundOnly_withAlphaRemoved() =
+        testScope.runTest {
+            val stoppedTaskId = 2
+            tasksRepository.seedTasks(tasks)
+            tasksRepository.setVisibleTasks(setOf(stoppedTaskId))
 
-        systemUnderTest.bind(stoppedTaskId)
-        assertThat(systemUnderTest.uiState.first())
-            .isEqualTo(BackgroundOnly(backgroundColor = Color.rgb(2, 2, 2)))
-    }
+            systemUnderTest.bind(stoppedTaskId)
+            assertThat(systemUnderTest.uiState.first())
+                .isEqualTo(BackgroundOnly(backgroundColor = Color.rgb(2, 2, 2)))
+        }
 
     @Test
-    fun bindLockedTaskWithThumbnail_thenStateIs_BackgroundOnly() = runTest {
-        val taskId = 2
-        tasksRepository.seedThumbnailData(mapOf(taskId to createThumbnailData()))
-        tasks[taskId].isLocked = true
-        tasksRepository.seedTasks(tasks)
-        tasksRepository.setVisibleTasks(setOf(taskId))
+    fun bindLockedTaskWithThumbnail_thenStateIs_BackgroundOnly() =
+        testScope.runTest {
+            val taskId = 2
+            tasksRepository.seedThumbnailData(mapOf(taskId to createThumbnailData()))
+            tasks[taskId].isLocked = true
+            tasksRepository.seedTasks(tasks)
+            tasksRepository.setVisibleTasks(setOf(taskId))
 
-        systemUnderTest.bind(taskId)
-        assertThat(systemUnderTest.uiState.first())
-            .isEqualTo(BackgroundOnly(backgroundColor = Color.rgb(2, 2, 2)))
-    }
+            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(taskId to expectedThumbnailData))
-        val expectedIconData = mock<Drawable>()
-        tasksRepository.seedIconData(taskId, "Task $taskId", "", expectedIconData)
-        tasksRepository.seedTasks(tasks)
-        tasksRepository.setVisibleTasks(setOf(taskId))
+    fun bindStoppedTaskWithThumbnail_thenStateIs_SnapshotSplash_withAlphaRemoved() =
+        testScope.runTest {
+            val taskId = 2
+            val expectedThumbnailData = createThumbnailData(rotation = Surface.ROTATION_270)
+            tasksRepository.seedThumbnailData(mapOf(taskId to expectedThumbnailData))
+            val expectedIconData = mock<Drawable>()
+            tasksRepository.seedIconData(taskId, "Task $taskId", "", expectedIconData)
+            tasksRepository.seedTasks(tasks)
+            tasksRepository.setVisibleTasks(setOf(taskId))
 
-        systemUnderTest.bind(taskId)
-        assertThat(systemUnderTest.uiState.first())
-            .isEqualTo(
-                SnapshotSplash(
-                    Snapshot(
-                        backgroundColor = Color.rgb(2, 2, 2),
-                        bitmap = expectedThumbnailData.thumbnail!!,
-                        thumbnailRotation = Surface.ROTATION_270,
-                    ),
-                    expectedIconData,
+            systemUnderTest.bind(taskId)
+            assertThat(systemUnderTest.uiState.first())
+                .isEqualTo(
+                    SnapshotSplash(
+                        Snapshot(
+                            backgroundColor = Color.rgb(2, 2, 2),
+                            bitmap = expectedThumbnailData.thumbnail!!,
+                            thumbnailRotation = Surface.ROTATION_270,
+                        ),
+                        expectedIconData,
+                    )
                 )
-            )
-    }
+        }
 
     @Test
-    fun bindNonVisibleStoppedTask_whenMadeVisible_thenStateIsSnapshotSplash() = runTest {
-        val taskId = 2
-        val expectedThumbnailData = createThumbnailData()
-        tasksRepository.seedThumbnailData(mapOf(taskId to expectedThumbnailData))
-        val expectedIconData = mock<Drawable>()
-        tasksRepository.seedIconData(taskId, "Task $taskId", "", expectedIconData)
-        tasksRepository.seedTasks(tasks)
+    fun bindNonVisibleStoppedTask_whenMadeVisible_thenStateIsSnapshotSplash() =
+        testScope.runTest {
+            val taskId = 2
+            val expectedThumbnailData = createThumbnailData()
+            tasksRepository.seedThumbnailData(mapOf(taskId to expectedThumbnailData))
+            val expectedIconData = mock<Drawable>()
+            tasksRepository.seedIconData(taskId, "Task $taskId", "", expectedIconData)
+            tasksRepository.seedTasks(tasks)
 
-        systemUnderTest.bind(taskId)
-        assertThat(systemUnderTest.uiState.first()).isEqualTo(Uninitialized)
+            systemUnderTest.bind(taskId)
+            assertThat(systemUnderTest.uiState.first()).isEqualTo(Uninitialized)
 
-        tasksRepository.setVisibleTasks(setOf(taskId))
-        assertThat(systemUnderTest.uiState.first())
-            .isEqualTo(
-                SnapshotSplash(
-                    Snapshot(
-                        backgroundColor = Color.rgb(2, 2, 2),
-                        bitmap = expectedThumbnailData.thumbnail!!,
-                        thumbnailRotation = Surface.ROTATION_0,
-                    ),
-                    expectedIconData,
+            tasksRepository.setVisibleTasks(setOf(taskId))
+            assertThat(systemUnderTest.uiState.first())
+                .isEqualTo(
+                    SnapshotSplash(
+                        Snapshot(
+                            backgroundColor = Color.rgb(2, 2, 2),
+                            bitmap = expectedThumbnailData.thumbnail!!,
+                            thumbnailRotation = Surface.ROTATION_0,
+                        ),
+                        expectedIconData,
+                    )
                 )
-            )
-    }
+        }
 
     @Test
-    fun getSnapshotMatrix_MissingThumbnail() = runTest {
-        val taskId = 2
-        val isRtl = true
+    fun getSnapshotMatrix_MissingThumbnail() =
+        testScope.runTest {
+            val taskId = 2
+            val isRtl = true
 
-        whenever(mGetThumbnailPositionUseCase.run(taskId, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
-            .thenReturn(MissingThumbnail)
+            whenever(mGetThumbnailPositionUseCase.run(taskId, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
+                .thenReturn(MissingThumbnail)
 
-        systemUnderTest.bind(taskId)
-        assertThat(systemUnderTest.getThumbnailPositionState(CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
-            .isEqualTo(Matrix.IDENTITY_MATRIX)
-    }
+            systemUnderTest.bind(taskId)
+            assertThat(
+                    systemUnderTest.getThumbnailPositionState(CANVAS_WIDTH, CANVAS_HEIGHT, isRtl)
+                )
+                .isEqualTo(Matrix.IDENTITY_MATRIX)
+        }
 
     @Test
-    fun getSnapshotMatrix_MatrixScaling() = runTest {
-        val taskId = 2
-        val isRtl = true
+    fun getSnapshotMatrix_MatrixScaling() =
+        testScope.runTest {
+            val taskId = 2
+            val isRtl = true
 
-        whenever(mGetThumbnailPositionUseCase.run(taskId, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
-            .thenReturn(MatrixScaling(MATRIX, isRotated = false))
+            whenever(mGetThumbnailPositionUseCase.run(taskId, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
+                .thenReturn(MatrixScaling(MATRIX, isRotated = false))
 
-        systemUnderTest.bind(taskId)
-        assertThat(systemUnderTest.getThumbnailPositionState(CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
-            .isEqualTo(MATRIX)
-    }
+            systemUnderTest.bind(taskId)
+            assertThat(
+                    systemUnderTest.getThumbnailPositionState(CANVAS_WIDTH, CANVAS_HEIGHT, isRtl)
+                )
+                .isEqualTo(MATRIX)
+        }
 
     @Test
-    fun getForegroundScrimDimProgress_returnsForegroundMaxScrim() = runTest {
-        recentsViewData.tintAmount.value = 0.32f
-        taskContainerData.taskMenuOpenProgress.value = 0f
-        assertThat(systemUnderTest.dimProgress.first()).isEqualTo(0.32f)
-    }
+    fun getForegroundScrimDimProgress_returnsForegroundMaxScrim() =
+        testScope.runTest {
+            recentsViewData.tintAmount.value = 0.32f
+            taskContainerData.taskMenuOpenProgress.value = 0f
+            assertThat(systemUnderTest.dimProgress.first()).isEqualTo(0.32f)
+        }
 
     @Test
-    fun getTaskMenuScrimDimProgress_returnsTaskMenuScrim() = runTest {
-        recentsViewData.tintAmount.value = 0f
-        taskContainerData.taskMenuOpenProgress.value = 1f
-        assertThat(systemUnderTest.dimProgress.first()).isEqualTo(0.4f)
-    }
+    fun getTaskMenuScrimDimProgress_returnsTaskMenuScrim() =
+        testScope.runTest {
+            recentsViewData.tintAmount.value = 0f
+            taskContainerData.taskMenuOpenProgress.value = 1f
+            assertThat(systemUnderTest.dimProgress.first()).isEqualTo(0.4f)
+        }
 
     @Test
-    fun getForegroundScrimDimProgress_returnsNoScrim() = runTest {
-        recentsViewData.tintAmount.value = 0f
-        taskContainerData.taskMenuOpenProgress.value = 0f
-        assertThat(systemUnderTest.dimProgress.first()).isEqualTo(0f)
-    }
+    fun getForegroundScrimDimProgress_returnsNoScrim() =
+        testScope.runTest {
+            recentsViewData.tintAmount.value = 0f
+            taskContainerData.taskMenuOpenProgress.value = 0f
+            assertThat(systemUnderTest.dimProgress.first()).isEqualTo(0f)
+        }
 
     private fun createTaskWithId(taskId: Int) =
         Task(Task.TaskKey(taskId, 0, Intent(), ComponentName("", ""), 0, 2000)).apply {
diff --git a/quickstep/tests/src/com/android/quickstep/TaplStartLauncherViaGestureTests.java b/quickstep/tests/src/com/android/quickstep/TaplStartLauncherViaGestureTests.java
index a8f39af..2fb08dd 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplStartLauncherViaGestureTests.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplStartLauncherViaGestureTests.java
@@ -16,6 +16,8 @@
 
 package com.android.quickstep;
 
+import android.util.Log;
+
 import androidx.test.filters.LargeTest;
 import androidx.test.runner.AndroidJUnit4;
 
@@ -29,6 +31,8 @@
 @RunWith(AndroidJUnit4.class)
 public class TaplStartLauncherViaGestureTests extends AbstractQuickStepTest {
 
+    public static final String TAG = "TaplStartLauncherViaGestureTests";
+
     static final int STRESS_REPEAT_COUNT = 10;
 
     private enum TestCase {
@@ -69,7 +73,9 @@
     }
 
     private void runTest(TestCase testCase) {
+        long testStartTime = System.currentTimeMillis();
         for (int i = 0; i < STRESS_REPEAT_COUNT; ++i) {
+            long loopStartTime = System.currentTimeMillis();
             // Destroy Launcher activity.
             closeLauncherActivity();
 
@@ -84,7 +90,10 @@
                 default:
                     throw new IllegalStateException("Cannot run test case: " + testCase);
             }
+            Log.d(TAG, "Loop " + (i + 1) + " runtime="
+                    + (System.currentTimeMillis() - loopStartTime) + "ms");
         }
+        Log.d(TAG, "Test runtime=" + (System.currentTimeMillis() - testStartTime) + "ms");
         switch (testCase) {
             case TO_OVERVIEW:
                 closeLauncherActivity();
diff --git a/quickstep/tests/src/com/android/quickstep/desktop/DesktopAppLaunchTransitionManagerTest.kt b/quickstep/tests/src/com/android/quickstep/desktop/DesktopAppLaunchTransitionManagerTest.kt
new file mode 100644
index 0000000..26189df
--- /dev/null
+++ b/quickstep/tests/src/com/android/quickstep/desktop/DesktopAppLaunchTransitionManagerTest.kt
@@ -0,0 +1,110 @@
+/*
+ * 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.desktop
+
+import android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.content.Context
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
+import android.view.WindowManager.TRANSIT_OPEN
+import android.view.WindowManager.TRANSIT_TO_FRONT
+import android.window.TransitionFilter
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
+import com.android.quickstep.SystemUiProxy
+import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_APP_LAUNCH_TRANSITIONS
+import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
+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.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.times
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
+import org.mockito.quality.Strictness
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class DesktopAppLaunchTransitionManagerTest {
+
+    @get:Rule val mSetFlagsRule = SetFlagsRule()
+
+    private val mockitoSession =
+        mockitoSession()
+            .strictness(Strictness.LENIENT)
+            .mockStatic(DesktopModeStatus::class.java)
+            .startMocking()
+
+    private val context = mock<Context>()
+    private val systemUiProxy = mock<SystemUiProxy>()
+    private lateinit var transitionManager: DesktopAppLaunchTransitionManager
+
+    @Before
+    fun setUp() {
+        whenever(context.resources).thenReturn(mock())
+        whenever(DesktopModeStatus.canEnterDesktopMode(context)).thenReturn(true)
+        transitionManager = DesktopAppLaunchTransitionManager(context, systemUiProxy)
+    }
+
+    @After
+    fun tearDown() {
+        mockitoSession.finishMocking()
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_APP_LAUNCH_TRANSITIONS)
+    fun registerTransitions_appLaunchFlagEnabled_registersTransition() {
+        transitionManager.registerTransitions()
+
+        verify(systemUiProxy, times(1)).registerRemoteTransition(any(), any())
+    }
+
+    @Test
+    @DisableFlags(FLAG_ENABLE_DESKTOP_APP_LAUNCH_TRANSITIONS)
+    fun registerTransitions_appLaunchFlagDisabled_doesntRegisterTransition() {
+        transitionManager.registerTransitions()
+
+        verify(systemUiProxy, times(0)).registerRemoteTransition(any(), any())
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_APP_LAUNCH_TRANSITIONS)
+    fun registerTransitions_usesCorrectFilter() {
+        transitionManager.registerTransitions()
+        val filterArgumentCaptor = argumentCaptor<TransitionFilter>()
+
+        verify(systemUiProxy, times(1))
+            .registerRemoteTransition(any(), filterArgumentCaptor.capture())
+
+        assertThat(filterArgumentCaptor.lastValue).isNotNull()
+        assertThat(filterArgumentCaptor.lastValue.mTypeSet)
+            .isEqualTo(intArrayOf(TRANSIT_OPEN, TRANSIT_TO_FRONT))
+        assertThat(filterArgumentCaptor.lastValue.mRequirements).hasLength(1)
+        val launchRequirement = filterArgumentCaptor.lastValue.mRequirements!![0]
+        assertThat(launchRequirement.mModes).isEqualTo(intArrayOf(TRANSIT_OPEN, TRANSIT_TO_FRONT))
+        assertThat(launchRequirement.mActivityType).isEqualTo(ACTIVITY_TYPE_STANDARD)
+        assertThat(launchRequirement.mWindowingMode).isEqualTo(WINDOWING_MODE_FREEFORM)
+    }
+}
diff --git a/res/layout/widgets_list_expand_button.xml b/res/layout/widgets_list_expand_button.xml
index 17c19ac..ff2d777 100644
--- a/res/layout/widgets_list_expand_button.xml
+++ b/res/layout/widgets_list_expand_button.xml
@@ -15,6 +15,7 @@
   -->
 
 <Button xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/widget_list_expand_button"
     style="@style/Button.Rounded.Colored"
     android:layout_width="wrap_content"
     android:layout_height="wrap_content"
diff --git a/src/com/android/launcher3/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index 34cf56b..ef5c88a 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -20,6 +20,7 @@
 import static android.graphics.fonts.FontStyle.FONT_WEIGHT_NORMAL;
 import static android.text.Layout.Alignment.ALIGN_NORMAL;
 
+import static com.android.launcher3.Flags.enableContrastTiles;
 import static com.android.launcher3.Flags.enableCursorHoverStates;
 import static com.android.launcher3.graphics.PreloadIconDrawable.newPendingIcon;
 import static com.android.launcher3.icons.BitmapInfo.FLAG_NO_BADGE;
@@ -39,6 +40,7 @@
 import android.graphics.Color;
 import android.graphics.Paint;
 import android.graphics.Rect;
+import android.graphics.RectF;
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
 import android.icu.text.MessageFormat;
@@ -722,6 +724,29 @@
         }
     }
 
+    /** Draws a background behind the App Title label when required. **/
+    public void drawAppContrastTile(Canvas canvas) {
+        RectF appTitleBounds;
+        Paint.FontMetrics fm = getPaint().getFontMetrics();
+        Rect tmpRect = new Rect();
+        getDrawingRect(tmpRect);
+
+        if (mIcon == null) {
+            appTitleBounds = new RectF(0, 0, tmpRect.right,
+                    (int) Math.ceil(fm.bottom - fm.top));
+        } else {
+            Rect iconBounds = new Rect();
+            getIconBounds(iconBounds);
+            int textStart = iconBounds.bottom + getCompoundDrawablePadding();
+            appTitleBounds = new RectF(tmpRect.left, textStart, tmpRect.right,
+                    textStart + (int) Math.ceil(fm.bottom - fm.top));
+        }
+
+        canvas.drawRoundRect(appTitleBounds, appTitleBounds.height() / 2,
+                appTitleBounds.height() / 2,
+                PillColorProvider.getInstance(getContext()).getAppTitlePillPaint());
+    }
+
     /** Draws a line under the app icon if this is representing a running app in Desktop Mode. */
     protected void drawRunningAppIndicatorIfNecessary(Canvas canvas) {
         if (mRunningAppState == RunningAppState.NOT_RUNNING || mDisplay != DISPLAY_TASKBAR) {
@@ -909,7 +934,9 @@
 
     @Override
     public void setTextColor(ColorStateList colors) {
-        mTextColor = colors.getDefaultColor();
+        mTextColor = shouldDrawAppContrastTile() ? PillColorProvider.getInstance(
+                getContext()).getAppTitleTextPaint().getColor()
+                : colors.getDefaultColor();
         mTextColorStateList = colors;
         if (Float.compare(mTextAlpha, 1) == 0) {
             super.setTextColor(colors);
@@ -926,6 +953,15 @@
                 && info.container != LauncherSettings.Favorites.CONTAINER_HOTSEAT_PREDICTION);
     }
 
+    /**
+     * Whether or not an App title contrast tile should be drawn for this element.
+     **/
+    public boolean shouldDrawAppContrastTile() {
+        return mDisplay == DISPLAY_WORKSPACE && shouldTextBeVisible()
+                && PillColorProvider.getInstance(getContext()).isMatchaEnabled()
+                && enableContrastTiles();
+    }
+
     public void setTextVisibility(boolean visible) {
         setTextAlpha(visible ? 1 : 0);
     }
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 6145077..305941e 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -537,6 +537,7 @@
 
         mPopupDataProvider = new PopupDataProvider(this::updateNotificationDots);
         mWidgetPickerDataProvider = new WidgetPickerDataProvider();
+        PillColorProvider.getInstance(mWorkspace.getContext()).registerObserver();
 
         boolean internalStateHandled = ACTIVITY_TRACKER.handleCreate(this);
         if (internalStateHandled) {
@@ -1813,6 +1814,7 @@
         // changes while launcher is still loading.
         getRootView().getViewTreeObserver().removeOnPreDrawListener(mOnInitialBindListener);
         mOverlayManager.onActivityDestroyed();
+        PillColorProvider.getInstance(mWorkspace.getContext()).unregisterObserver();
     }
 
     public LauncherAccessibilityDelegate getAccessibilityDelegate() {
diff --git a/src/com/android/launcher3/PillColorPorovider.kt b/src/com/android/launcher3/PillColorPorovider.kt
new file mode 100644
index 0000000..347c5d6
--- /dev/null
+++ b/src/com/android/launcher3/PillColorPorovider.kt
@@ -0,0 +1,85 @@
+/*
+ * 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
+
+import android.content.Context
+import android.database.ContentObserver
+import android.graphics.Paint
+import android.net.Uri
+import android.provider.Settings
+import com.android.launcher3.util.Executors.ORDERED_BG_EXECUTOR
+
+class PillColorProvider private constructor(c: Context) {
+    private val context = c.applicationContext
+
+    private val matchaUri by lazy { Settings.Secure.getUriFor(MATCHA_SETTING) }
+    var appTitlePillPaint = Paint()
+        private set
+
+    var appTitleTextPaint = Paint()
+        private set
+
+    private var isMatchaEnabledInternal = 0
+
+    var isMatchaEnabled = isMatchaEnabledInternal != 0
+
+    private val pillColorObserver =
+        object : ContentObserver(ORDERED_BG_EXECUTOR.handler) {
+            override fun onChange(selfChange: Boolean, uri: Uri?) {
+                if (uri == matchaUri) {
+                    isMatchaEnabledInternal =
+                        Settings.Secure.getInt(context.contentResolver, MATCHA_SETTING, 0)
+                    isMatchaEnabled = isMatchaEnabledInternal != 0
+                }
+            }
+        }
+
+    fun registerObserver() {
+        context.contentResolver.registerContentObserver(matchaUri, false, pillColorObserver)
+        setup()
+    }
+
+    fun unregisterObserver() {
+        context.contentResolver.unregisterContentObserver(pillColorObserver)
+    }
+
+    fun setup() {
+        appTitlePillPaint.color =
+            context.resources.getColor(
+                R.color.material_color_surface_container_lowest,
+                context.theme,
+            )
+        appTitleTextPaint.color =
+            context.resources.getColor(R.color.material_color_on_surface, context.theme)
+        isMatchaEnabledInternal = Settings.Secure.getInt(context.contentResolver, MATCHA_SETTING, 0)
+        isMatchaEnabled = isMatchaEnabledInternal != 0
+    }
+
+    companion object {
+        private var INSTANCE: PillColorProvider? = null
+        private const val MATCHA_SETTING = "matcha_enable"
+
+        // TODO: Replace with a Dagger injection that is a singleton.
+        @JvmStatic
+        fun getInstance(context: Context): PillColorProvider {
+            if (INSTANCE == null) {
+                INSTANCE = PillColorProvider(context)
+            }
+            return INSTANCE!!
+        }
+    }
+}
diff --git a/src/com/android/launcher3/views/DoubleShadowBubbleTextView.java b/src/com/android/launcher3/views/DoubleShadowBubbleTextView.java
index ef66ffe..392d9a7 100644
--- a/src/com/android/launcher3/views/DoubleShadowBubbleTextView.java
+++ b/src/com/android/launcher3/views/DoubleShadowBubbleTextView.java
@@ -102,6 +102,9 @@
 
     @Override
     public void onDraw(Canvas canvas) {
+        if (shouldDrawAppContrastTile()) {
+            drawAppContrastTile(canvas);
+        }
         // If text is transparent or shadow alpha is 0, don't draw any shadow
         if (skipDoubleShadow()) {
             super.onDraw(canvas);
diff --git a/src/com/android/launcher3/views/StickyHeaderLayout.java b/src/com/android/launcher3/views/StickyHeaderLayout.java
index 090251f..4142e1f 100644
--- a/src/com/android/launcher3/views/StickyHeaderLayout.java
+++ b/src/com/android/launcher3/views/StickyHeaderLayout.java
@@ -120,7 +120,19 @@
     }
 
     private float getCurrentScroll() {
-        return mScrollOffset + (mCurrentEmptySpaceView == null ? 0 : mCurrentEmptySpaceView.getY());
+        float scroll;
+        if (mCurrentRecyclerView.getVisibility() != VISIBLE) {
+            // When no list is displayed, assume no scroll.
+            scroll = 0f;
+        } else if (mCurrentEmptySpaceView != null) {
+            // Otherwise use empty space view as reference to position.
+            scroll = mCurrentEmptySpaceView.getY();
+        } else {
+            // If there is no empty space view, but the list is visible, we are scrolled away
+            // completely, so assume all non-sticky children should also be scrolled away.
+            scroll = -mHeaderHeight;
+        }
+        return mScrollOffset + scroll;
     }
 
     @Override
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 2f64ab1..8bebfb2 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -605,9 +605,12 @@
             mAdapters.get(AdapterHolder.SEARCH).mWidgetsRecyclerView.setVisibility(GONE);
             mAdapters.get(getCurrentAdapterHolderType()).mWidgetsRecyclerView.setVisibility(
                     VISIBLE);
-            // Visibility of recommended widgets, recycler views and headers are handled in methods
-            // below.
-            post(this::onRecommendedWidgetsBound);
+            if (mRecommendedWidgetsCount > 0) {
+                // Display recommendations immediately, if present, so that other parts of sticky
+                // header (e.g. personal / work tabs) don't flash in interim.
+                mWidgetRecommendationsContainer.setVisibility(VISIBLE);
+            }
+            // Visibility of recycler views and headers are handled in methods below.
             onWidgetsBound();
         }
     }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
index 3c67538..74a9a5c 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
@@ -103,7 +103,7 @@
                             .equals(mWidgetsContentVisiblePackageUserKey);
     @Nullable private Predicate<WidgetsListBaseEntry> mFilter = null;
     @Nullable private RecyclerView mRecyclerView;
-    @Nullable private PackageUserKey mPendingClickHeader;
+    @Nullable private PackageUserKey mHeaderPositionToMaintain;
     @Px private int mMaxHorizontalSpan;
 
     private boolean mShowOnlyDefaultList = true;
@@ -215,7 +215,7 @@
         // Get the current top of the header with the matching key before adjusting the visible
         // entries.
         OptionalInt previousPositionForPackageUserKey =
-                getPositionForPackageUserKey(mPendingClickHeader);
+                getPositionForPackageUserKey(mHeaderPositionToMaintain);
         OptionalInt topForPackageUserKey =
                 getOffsetForPosition(previousPositionForPackageUserKey);
 
@@ -247,13 +247,15 @@
         mVisibleEntries.addAll(newVisibleEntries);
         diffResult.dispatchUpdatesTo(this);
 
-        if (mPendingClickHeader != null) {
+        if (mHeaderPositionToMaintain != null && mRecyclerView != null) {
             // Get the position for the clicked header after adjusting the visible entries. The
             // position may have changed if another header had previously been expanded.
             OptionalInt positionForPackageUserKey =
-                    getPositionForPackageUserKey(mPendingClickHeader);
-            scrollToPositionAndMaintainOffset(positionForPackageUserKey, topForPackageUserKey);
-            mPendingClickHeader = null;
+                    getPositionForPackageUserKey(mHeaderPositionToMaintain);
+            // Post scroll updates to be applied after diff updates.
+            mRecyclerView.post(() -> scrollToPositionAndMaintainOffset(positionForPackageUserKey,
+                    topForPackageUserKey));
+            mHeaderPositionToMaintain = null;
         }
     }
 
@@ -384,7 +386,7 @@
 
         // Store the header that was clicked so that its position will be maintained the next time
         // we update the entries.
-        mPendingClickHeader = packageUserKey;
+        mHeaderPositionToMaintain = packageUserKey;
 
         updateVisibleEntries();
 
@@ -470,6 +472,16 @@
      */
     public void useExpandedList() {
         mShowOnlyDefaultList = false;
+        if (mWidgetsContentVisiblePackageUserKey != null) {
+            // Maintain selected header for the next update that expands the list.
+            mHeaderPositionToMaintain = mWidgetsContentVisiblePackageUserKey;
+        } else if (mVisibleEntries.size() > 2) {
+            // Maintain last visible header shown above expand button since there was no selected
+            // header.
+            mHeaderPositionToMaintain = PackageUserKey.fromPackageItemInfo(
+                    mVisibleEntries.get(mVisibleEntries.size() - 2).mPkgItem);
+        }
+
     }
 
     /** Comparator for sorting WidgetListRowEntry based on package title. */
diff --git a/tests/AndroidManifest-common.xml b/tests/AndroidManifest-common.xml
index 68e493d..1ddd453 100644
--- a/tests/AndroidManifest-common.xml
+++ b/tests/AndroidManifest-common.xml
@@ -183,7 +183,6 @@
         </activity>
         <activity-alias android:name="Activity2"
                         android:label="TestActivity2"
-                        android:icon="@drawable/test_icon"
                         android:exported="true"
                         android:targetActivity="com.android.launcher3.testcomponent.BaseTestingActivity">
             <intent-filter>
@@ -193,7 +192,6 @@
         </activity-alias>
         <activity-alias android:name="Activity3"
                         android:label="TestActivity3"
-                        android:icon="@drawable/test_icon"
                         android:exported="true"
                         android:targetActivity="com.android.launcher3.testcomponent.BaseTestingActivity">
             <intent-filter>
@@ -203,7 +201,6 @@
         </activity-alias>
         <activity-alias android:name="Activity4"
                         android:label="TestActivity4"
-                        android:icon="@drawable/test_icon"
                         android:exported="true"
                         android:targetActivity="com.android.launcher3.testcomponent.BaseTestingActivity">
             <intent-filter>
@@ -213,7 +210,6 @@
         </activity-alias>
         <activity-alias android:name="Activity5"
                         android:label="TestActivity5"
-                        android:icon="@drawable/test_icon"
                         android:exported="true"
                         android:targetActivity="com.android.launcher3.testcomponent.BaseTestingActivity">
             <intent-filter>
@@ -223,7 +219,6 @@
         </activity-alias>
         <activity-alias android:name="Activity6"
                         android:label="TestActivity6"
-                        android:icon="@drawable/test_icon"
                         android:exported="true"
                         android:targetActivity="com.android.launcher3.testcomponent.BaseTestingActivity">
             <intent-filter>
@@ -233,7 +228,6 @@
         </activity-alias>
         <activity-alias android:name="Activity7"
                         android:label="TestActivity7"
-                        android:icon="@drawable/test_icon"
                         android:exported="true"
                         android:targetActivity="com.android.launcher3.testcomponent.BaseTestingActivity">
             <intent-filter>
@@ -243,7 +237,6 @@
         </activity-alias>
         <activity-alias android:name="Activity8"
                         android:label="TestActivity8"
-                        android:icon="@drawable/test_icon"
                         android:exported="true"
                         android:targetActivity="com.android.launcher3.testcomponent.BaseTestingActivity">
             <intent-filter>
@@ -253,7 +246,6 @@
         </activity-alias>
         <activity-alias android:name="Activity9" android:exported="true"
             android:label="TestActivity9"
-            android:icon="@drawable/test_icon"
             android:targetActivity="com.android.launcher3.testcomponent.OtherBaseTestingActivity">
             <intent-filter>
                 <action android:name="android.intent.action.MAIN" />
@@ -262,7 +254,6 @@
         </activity-alias>
         <activity-alias android:name="Activity10" android:exported="true"
             android:label="TestActivity10"
-            android:icon="@drawable/test_icon"
             android:targetActivity="com.android.launcher3.testcomponent.OtherBaseTestingActivity">
             <intent-filter>
                 <action android:name="android.intent.action.MAIN" />
@@ -271,7 +262,6 @@
         </activity-alias>
         <activity-alias android:name="Activity11" android:exported="true"
             android:label="TestActivity11"
-            android:icon="@drawable/test_icon"
             android:targetActivity="com.android.launcher3.testcomponent.OtherBaseTestingActivity">
             <intent-filter>
                 <action android:name="android.intent.action.MAIN" />
@@ -280,7 +270,6 @@
         </activity-alias>
         <activity-alias android:name="Activity12" android:exported="true"
             android:label="TestActivity12"
-            android:icon="@drawable/test_icon"
             android:targetActivity="com.android.launcher3.testcomponent.OtherBaseTestingActivity">
             <intent-filter>
                 <action android:name="android.intent.action.MAIN" />
@@ -289,7 +278,6 @@
         </activity-alias>
         <activity-alias android:name="Activity13" android:exported="true"
             android:label="TestActivity13"
-            android:icon="@drawable/test_icon"
             android:targetActivity="com.android.launcher3.testcomponent.OtherBaseTestingActivity">
             <intent-filter>
                 <action android:name="android.intent.action.MAIN" />
@@ -298,7 +286,6 @@
         </activity-alias>
         <activity-alias android:name="Activity14" android:exported="true"
             android:label="TestActivity14"
-            android:icon="@drawable/test_icon"
             android:targetActivity="com.android.launcher3.testcomponent.OtherBaseTestingActivity">
             <intent-filter>
                 <action android:name="android.intent.action.MAIN" />
@@ -376,7 +363,7 @@
         </activity>
         <activity android:name="com.android.launcher3.testcomponent.ImeTestActivity"
             android:label="ImeTestActivity"
-            android:icon="@drawable/test_icon"
+            android:icon="@drawable/test_theme_icon"
             android:exported="true">
             <intent-filter>
                 <action android:name="android.intent.action.MAIN" />
diff --git a/tests/multivalentTests/src/com/android/launcher3/AbstractDeviceProfileTest.kt b/tests/multivalentTests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
index 62c8426..f8794f8 100644
--- a/tests/multivalentTests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
@@ -126,6 +126,7 @@
         deviceSpec: DeviceSpec,
         isGestureMode: Boolean = true,
         isVerticalBar: Boolean = false,
+        isFixedLandscape: Boolean = false,
     ) {
         val (naturalX, naturalY) = deviceSpec.naturalSize
         val windowsBounds = phoneWindowsBounds(deviceSpec, isGestureMode, naturalX, naturalY)
@@ -138,6 +139,7 @@
             rotation = if (isVerticalBar) Surface.ROTATION_90 else Surface.ROTATION_0,
             isGestureMode,
             densityDpi = deviceSpec.densityDpi,
+            isFixedLandscape = isFixedLandscape,
         )
     }
 
@@ -166,6 +168,7 @@
         isLandscape: Boolean = false,
         isGestureMode: Boolean = true,
         isFolded: Boolean = false,
+        isFixedLandscape: Boolean = false,
     ) {
         val (unfoldedNaturalX, unfoldedNaturalY) = deviceSpecUnfolded.naturalSize
         val unfoldedWindowsBounds =
@@ -192,6 +195,7 @@
                 rotation = if (isLandscape) Surface.ROTATION_90 else Surface.ROTATION_0,
                 isGestureMode = isGestureMode,
                 densityDpi = deviceSpecFolded.densityDpi,
+                isFixedLandscape = isFixedLandscape,
             )
         } else {
             initializeCommonVars(
@@ -200,6 +204,7 @@
                 rotation = if (isLandscape) Surface.ROTATION_0 else Surface.ROTATION_90,
                 isGestureMode = isGestureMode,
                 densityDpi = deviceSpecUnfolded.densityDpi,
+                isFixedLandscape = isFixedLandscape,
             )
         }
     }
@@ -274,6 +279,7 @@
         rotation: Int,
         isGestureMode: Boolean = true,
         densityDpi: Int,
+        isFixedLandscape: Boolean = false,
     ) {
         setFlagsRule.setFlags(true, Flags.FLAG_ENABLE_TWOLINE_TOGGLE)
         LauncherPrefs.get(testContext).put(LauncherPrefs.ENABLE_TWOLINE_ALLAPPS_TOGGLE, true)
@@ -307,7 +313,7 @@
 
         whenever(launcherPrefs.get(LauncherPrefs.TASKBAR_PINNING)).thenReturn(false)
         whenever(launcherPrefs.get(LauncherPrefs.TASKBAR_PINNING_IN_DESKTOP_MODE)).thenReturn(true)
-        whenever(launcherPrefs.get(LauncherPrefs.FIXED_LANDSCAPE_MODE)).thenReturn(false)
+        whenever(launcherPrefs.get(LauncherPrefs.FIXED_LANDSCAPE_MODE)).thenReturn(isFixedLandscape)
         whenever(launcherPrefs.get(LauncherPrefs.HOTSEAT_COUNT)).thenReturn(-1)
         whenever(launcherPrefs.get(LauncherPrefs.DEVICE_TYPE)).thenReturn(-1)
         whenever(launcherPrefs.get(LauncherPrefs.WORKSPACE_SIZE)).thenReturn("")
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt
index e8f778f..7cd5da4 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt
@@ -18,7 +18,6 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.launcher3.LauncherAppState
-import com.android.launcher3.icons.BitmapInfo
 import com.android.launcher3.icons.waitForUpdateHandlerToFinish
 import com.android.launcher3.model.data.WorkspaceItemInfo
 import com.android.launcher3.util.Executors
@@ -161,9 +160,6 @@
             assertWithMessage("Index $index was not highRes")
                 .that(items[index].bitmap.isNullOrLowRes)
                 .isFalse()
-            assertWithMessage("Index $index was the default icon")
-                .that(isDefaultIcon(items[index].bitmap))
-                .isFalse()
         }
     }
 
@@ -172,17 +168,9 @@
             assertWithMessage("Index $index was not lowRes")
                 .that(items[index].bitmap.isNullOrLowRes)
                 .isTrue()
-            assertWithMessage("Index $index was the default icon")
-                .that(isDefaultIcon(items[index].bitmap))
-                .isFalse()
         }
     }
 
-    private fun isDefaultIcon(bitmap: BitmapInfo) =
-        LauncherAppState.getInstance(modelHelper.sandboxContext)
-            .iconCache
-            .isDefaultIcon(bitmap, modelHelper.sandboxContext.user)
-
     /** Recreate DeviceProfiles after changing InvariantDeviceProfile */
     private fun recreateSupportedDeviceProfiles() {
         LauncherAppState.getIDP(modelHelper.sandboxContext).supportedProfiles =
diff --git a/tests/res/drawable/test_icon.xml b/tests/res/drawable/test_icon.xml
deleted file mode 100644
index 72ebfeb..0000000
--- a/tests/res/drawable/test_icon.xml
+++ /dev/null
@@ -1,29 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-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.
--->
-<adaptive-icon xmlns:android="http://schemas.android.com/apk/res/android">
-    <background android:drawable="@android:color/white"/>
-    <foreground>
-        <color android:color="#FFFF0000" />
-    </foreground>
-    <monochrome>
-        <vector android:width="48dp" android:height="48dp" android:viewportWidth="48.0" android:viewportHeight="48.0">
-            <path
-                android:fillColor="#FF000000"
-                android:pathData="M0,24L48,24 48,48, 0,48 Z"/>
-        </vector>
-    </monochrome>
-</adaptive-icon>
diff --git a/tests/src/com/android/launcher3/nonquickstep/DeviceProfileDumpTest.kt b/tests/src/com/android/launcher3/nonquickstep/DeviceProfileDumpTest.kt
index 60385a7..2e2b6cd 100644
--- a/tests/src/com/android/launcher3/nonquickstep/DeviceProfileDumpTest.kt
+++ b/tests/src/com/android/launcher3/nonquickstep/DeviceProfileDumpTest.kt
@@ -20,6 +20,7 @@
 import com.android.launcher3.DeviceProfile
 import com.android.launcher3.Flags
 import com.android.launcher3.InvariantDeviceProfile
+import com.android.launcher3.util.rule.setFlags
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -35,11 +36,11 @@
 
     @Before
     fun setUp() {
-        if (instance.decoupleDepth) {
-            setFlagsRule.enableFlags(Flags.FLAG_ENABLE_SCALING_REVEAL_HOME_ANIMATION)
-        } else {
-            setFlagsRule.disableFlags(Flags.FLAG_ENABLE_SCALING_REVEAL_HOME_ANIMATION)
-        }
+        setFlagsRule.setFlags(
+            instance.decoupleDepth,
+            Flags.FLAG_ENABLE_SCALING_REVEAL_HOME_ANIMATION,
+        )
+        setFlagsRule.setFlags(false, Flags.FLAG_ONE_GRID_SPECS)
     }
 
     @Test
@@ -105,13 +106,13 @@
                 initializeVarsForTablet(
                     deviceSpec = deviceSpec,
                     isLandscape = isLandscape,
-                    isGestureMode = isGestureMode
+                    isGestureMode = isGestureMode,
                 )
             else ->
                 initializeVarsForPhone(
                     deviceSpec = deviceSpec,
                     isVerticalBar = isLandscape,
-                    isGestureMode = isGestureMode
+                    isGestureMode = isGestureMode,
                 )
         }
     }
@@ -136,7 +137,7 @@
                     "twopanel-tablet",
                     gridName = "4_by_4",
                     isTaskbarPresentInApps = true,
-                    decoupleDepth = true
+                    decoupleDepth = true,
                 ),
             )
         }
@@ -145,7 +146,7 @@
             val deviceName: String,
             val gridName: String,
             val isTaskbarPresentInApps: Boolean = false,
-            val decoupleDepth: Boolean = false
+            val decoupleDepth: Boolean = false,
         ) {
             fun filename(testName: String = ""): String {
                 val device =
diff --git a/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
similarity index 78%
rename from tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
rename to tests/src/com/android/launcher3/ui/WorkProfileTest.java
index 33ffd1d..d866a9f 100644
--- a/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
+++ b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
@@ -17,9 +17,7 @@
 
 import static com.android.launcher3.LauncherPrefs.WORK_EDU_STEP;
 import static com.android.launcher3.LauncherState.ALL_APPS;
-import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.allapps.AllAppsStore.DEFER_UPDATES_TEST;
-import static com.android.launcher3.testing.shared.TestProtocol.NORMAL_STATE_ORDINAL;
 import static com.android.launcher3.util.TestUtil.installDummyAppForUser;
 import static com.android.launcher3.util.rule.TestStabilityRule.LOCAL;
 import static com.android.launcher3.util.rule.TestStabilityRule.PLATFORM_POSTSUBMIT;
@@ -43,36 +41,39 @@
 import com.android.launcher3.allapps.WorkEduCard;
 import com.android.launcher3.allapps.WorkPausedCard;
 import com.android.launcher3.allapps.WorkProfileManager;
-import com.android.launcher3.tapl.LauncherInstrumentation;
+import com.android.launcher3.util.BaseLauncherActivityTest;
 import com.android.launcher3.util.TestUtil;
+import com.android.launcher3.util.rule.ScreenRecordRule;
+import com.android.launcher3.util.rule.TestStabilityRule;
 import com.android.launcher3.util.rule.TestStabilityRule.Stability;
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import java.io.IOException;
-import java.util.Objects;
 import java.util.function.Predicate;
 
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-public class TaplWorkProfileTest extends AbstractLauncherUiTest<Launcher> {
+public class WorkProfileTest extends BaseLauncherActivityTest<Launcher> {
 
     private static final int WORK_PAGE = ActivityAllAppsContainerView.AdapterHolder.WORK;
+    public static final int WAIT_TIME_MS = 30000;
+
+    @Rule
+    public ScreenRecordRule mScreenRecordRule = new ScreenRecordRule();
+    @Rule
+    public TestStabilityRule mTestStabilityRule = new TestStabilityRule();
 
     private int mProfileUserId;
     private boolean mWorkProfileSetupSuccessful;
-    private final String TAG = "WorkProfileTest";
+    private static final String TAG = "WorkProfileTest";
 
     @Before
-    @Override
     public void setUp() throws Exception {
-        super.setUp();
-        String output =
-                mDevice.executeShellCommand(
-                        "pm create-user --profileOf 0 --managed TestProfile");
+        String output = executeShellCommand("pm create-user --profileOf 0 --managed TestProfile");
         updateWorkProfileSetupSuccessful("pm create-user", output);
 
         String[] tokens = output.split("\\s+");
@@ -88,36 +89,15 @@
             return; // no need to setup launcher since all tests will skip.
         }
 
-        mDevice.pressHome();
-        waitForLauncherCondition("Launcher didn't start", Objects::nonNull);
-        waitForStateTransitionToEnd("Launcher internal state didn't switch to Normal",
-                () -> NORMAL);
-        waitForResumed("Launcher internal state is still Background");
-        mLauncher.getWorkspace().switchToAllApps();
-        waitForStateTransitionToEnd("Launcher internal state didn't switch to All Apps",
-                () -> ALL_APPS);
+        loadLauncherSync();
+        goToState(ALL_APPS);
+        waitForState("Launcher internal state didn't switch to All Apps", () -> ALL_APPS);
     }
 
     @After
     public void removeWorkProfile() throws Exception {
-        executeOnLauncherInTearDown(launcher -> {
-            if (launcher.getAppsView() == null) {
-                return;
-            }
-            launcher.getAppsView().getAppsStore().disableDeferUpdates(DEFER_UPDATES_TEST);
-        });
         TestUtil.uninstallDummyApp();
-
-        mLauncher.runToState(
-                () -> {
-                    try {
-                        mDevice.executeShellCommand("pm remove-user --wait " + mProfileUserId);
-                    } catch (IOException e) {
-                        throw new RuntimeException(e);
-                    }
-                },
-                NORMAL_STATE_ORDINAL,
-                "executing pm 'remove-user' command");
+        executeShellCommand("pm remove-user --wait " + mProfileUserId);
     }
 
     private void waitForWorkTabSetup() {
@@ -127,7 +107,7 @@
                 return true;
             }
             return false;
-        }, LauncherInstrumentation.WAIT_TIME_MS);
+        }, WAIT_TIME_MS);
     }
 
     @Test
@@ -137,10 +117,10 @@
         waitForWorkTabSetup();
         waitForLauncherCondition("Personal tab is missing",
                 launcher -> launcher.getAppsView().isPersonalTabVisible(),
-                LauncherInstrumentation.WAIT_TIME_MS);
+                WAIT_TIME_MS);
         waitForLauncherCondition("Work tab is missing",
                 launcher -> launcher.getAppsView().isWorkTabVisible(),
-                LauncherInstrumentation.WAIT_TIME_MS);
+                WAIT_TIME_MS);
     }
 
     // Staging; will be promoted to presubmit if stable
@@ -156,12 +136,11 @@
 
         WorkProfileManager manager = getFromLauncher(l -> l.getAppsView().getWorkManager());
 
-
         waitForLauncherCondition("work profile initial state check failed", launcher ->
                         manager.getWorkUtilityView() != null
                                 && manager.getCurrentState() == WorkProfileManager.STATE_ENABLED
                                 && manager.getWorkUtilityView().isEnabled(),
-                LauncherInstrumentation.WAIT_TIME_MS);
+                WAIT_TIME_MS);
 
         //start work profile toggle OFF test
         executeOnLauncher(l -> {
@@ -173,7 +152,7 @@
         waitForLauncherCondition("Work profile toggle OFF failed", launcher -> {
             manager.reset(); // pulls current state from system
             return manager.getCurrentState() == WorkProfileManager.STATE_DISABLED;
-        }, LauncherInstrumentation.WAIT_TIME_MS);
+        }, WAIT_TIME_MS);
 
         waitForWorkCard("Work paused card not shown", view -> view instanceof WorkPausedCard);
 
@@ -188,7 +167,7 @@
         waitForLauncherCondition("Work profile toggle ON failed", launcher -> {
             manager.reset(); // pulls current state from system
             return manager.getCurrentState() == WorkProfileManager.STATE_ENABLED;
-        }, LauncherInstrumentation.WAIT_TIME_MS);
+        }, WAIT_TIME_MS);
 
     }
 
@@ -215,7 +194,7 @@
             } finally {
                 l.getAppsView().getAppsStore().enableDeferUpdates(DEFER_UPDATES_TEST);
             }
-        }, LauncherInstrumentation.WAIT_TIME_MS);
+        }, WAIT_TIME_MS);
     }
 
     private void updateWorkProfileSetupSuccessful(String cli, String output) {
diff --git a/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java b/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
deleted file mode 100644
index d653317..0000000
--- a/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
+++ /dev/null
@@ -1,177 +0,0 @@
-/*
- * Copyright (C) 2022 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.ui.workspace;
-
-import static com.android.launcher3.util.TestConstants.AppNames.TEST_APP_NAME;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertTrue;
-
-import android.content.ContentProviderClient;
-import android.content.ContentResolver;
-import android.content.ContentValues;
-import android.net.Uri;
-import android.view.View;
-import android.view.ViewGroup;
-
-import androidx.test.filters.LargeTest;
-
-import com.android.launcher3.BubbleTextView;
-import com.android.launcher3.Launcher;
-import com.android.launcher3.icons.ThemedIconDrawable;
-import com.android.launcher3.tapl.HomeAllApps;
-import com.android.launcher3.tapl.HomeAppIcon;
-import com.android.launcher3.tapl.HomeAppIconMenuItem;
-import com.android.launcher3.ui.AbstractLauncherUiTest;
-import com.android.launcher3.util.Executors;
-
-import org.junit.Test;
-
-import java.util.ArrayDeque;
-import java.util.Queue;
-
-/**
- * Tests for theme icon support in Launcher
- *
- * Note running these tests will clear the workspace on the device.
- */
-@LargeTest
-public class TaplThemeIconsTest extends AbstractLauncherUiTest<Launcher> {
-
-    private static final String APP_NAME = "IconThemedActivity";
-    private static final String SHORTCUT_NAME = "Shortcut 1";
-
-    @Test
-    public void testIconWithoutTheme() throws Exception {
-        setThemeEnabled(false);
-        initialize(this);
-
-        HomeAllApps allApps = mLauncher.getWorkspace().switchToAllApps();
-        allApps.freeze();
-
-        try {
-            HomeAppIcon icon = allApps.getAppIcon(APP_NAME);
-            executeOnLauncher(l -> verifyIconTheme(APP_NAME, l.getAppsView(), false));
-            icon.dragToWorkspace(false, false);
-            executeOnLauncher(l -> verifyIconTheme(APP_NAME, l.getWorkspace(), false));
-        } finally {
-            allApps.unfreeze();
-        }
-    }
-
-    @Test
-    public void testShortcutIconWithoutTheme() throws Exception {
-        setThemeEnabled(false);
-        initialize(this);
-
-        HomeAllApps allApps = mLauncher.getWorkspace().switchToAllApps();
-        allApps.freeze();
-
-        try {
-            HomeAppIcon icon = allApps.getAppIcon(TEST_APP_NAME);
-            HomeAppIconMenuItem shortcutItem =
-                    (HomeAppIconMenuItem) icon.openDeepShortcutMenu().getMenuItem(SHORTCUT_NAME);
-            shortcutItem.dragToWorkspace(false, false);
-            executeOnLauncher(l -> verifyIconTheme(SHORTCUT_NAME, l.getWorkspace(), false));
-        } finally {
-            allApps.unfreeze();
-        }
-    }
-
-    @Test
-    public void testIconWithTheme() throws Exception {
-        setThemeEnabled(true);
-        initialize(this);
-
-        HomeAllApps allApps = mLauncher.getWorkspace().switchToAllApps();
-        allApps.freeze();
-
-        try {
-            HomeAppIcon icon = allApps.getAppIcon(APP_NAME);
-            executeOnLauncher(l -> verifyIconTheme(APP_NAME, l.getAppsView(), false));
-            icon.dragToWorkspace(false, false);
-            executeOnLauncher(l -> verifyIconTheme(APP_NAME, l.getWorkspace(), true));
-        } finally {
-            allApps.unfreeze();
-        }
-    }
-
-    @Test
-    public void testShortcutIconWithTheme() throws Exception {
-        setThemeEnabled(true);
-        initialize(this);
-
-        HomeAllApps allApps = mLauncher.getWorkspace().switchToAllApps();
-        allApps.freeze();
-
-        try {
-            HomeAppIcon icon = allApps.getAppIcon(TEST_APP_NAME);
-            HomeAppIconMenuItem shortcutItem =
-                    (HomeAppIconMenuItem) icon.openDeepShortcutMenu().getMenuItem(SHORTCUT_NAME);
-            shortcutItem.dragToWorkspace(false, false);
-            executeOnLauncher(l -> verifyIconTheme(SHORTCUT_NAME, l.getWorkspace(), true));
-        } finally {
-            allApps.unfreeze();
-        }
-    }
-
-    private void verifyIconTheme(String title, ViewGroup parent, boolean isThemed) {
-        // Wait for Launcher model to be completed
-        try {
-            Executors.MODEL_EXECUTOR.submit(() -> { }).get();
-        } catch (Exception e) {
-            throw new RuntimeException(e);
-        }
-
-        // Find the app icon
-        Queue<View> viewQueue = new ArrayDeque<>();
-        viewQueue.add(parent);
-        BubbleTextView icon = null;
-        while (!viewQueue.isEmpty()) {
-            View view = viewQueue.poll();
-            if (view instanceof ViewGroup) {
-                parent = (ViewGroup) view;
-                for (int i = parent.getChildCount() - 1; i >= 0; i--) {
-                    viewQueue.add(parent.getChildAt(i));
-                }
-            } else if (view instanceof BubbleTextView btv) {
-                if (title.equals(btv.getContentDescription().toString())) {
-                    icon = btv;
-                    break;
-                }
-            }
-        }
-
-        assertNotNull(icon.getIcon());
-        assertEquals(isThemed, icon.getIcon() instanceof ThemedIconDrawable);
-    }
-
-    private void setThemeEnabled(boolean isEnabled) throws Exception {
-        Uri uri = new Uri.Builder()
-                .scheme(ContentResolver.SCHEME_CONTENT)
-                .authority(mTargetPackage + ".grid_control")
-                .appendPath("set_icon_themed")
-                .build();
-        ContentValues values = new ContentValues();
-        values.put("boolean_value", isEnabled);
-        try (ContentProviderClient client = mTargetContext.getContentResolver()
-                .acquireContentProviderClient(uri)) {
-            int result = client.update(uri, values, null);
-            assertTrue(result > 0);
-        }
-    }
-}
diff --git a/tests/src/com/android/launcher3/ui/workspace/ThemeIconsTest.java b/tests/src/com/android/launcher3/ui/workspace/ThemeIconsTest.java
new file mode 100644
index 0000000..cfc0a6b
--- /dev/null
+++ b/tests/src/com/android/launcher3/ui/workspace/ThemeIconsTest.java
@@ -0,0 +1,203 @@
+/*
+ * Copyright (C) 2022 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.ui.workspace;
+
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
+import static com.android.launcher3.AbstractFloatingView.TYPE_ACTION_POPUP;
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
+import static com.android.launcher3.util.TestConstants.AppNames.TEST_APP_NAME;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import android.content.ContentProviderClient;
+import android.content.ContentResolver;
+import android.content.ContentValues;
+import android.net.Uri;
+import android.view.View;
+import android.view.ViewGroup;
+
+import androidx.test.filters.LargeTest;
+import androidx.test.uiautomator.UiDevice;
+
+import com.android.launcher3.AbstractFloatingView;
+import com.android.launcher3.BubbleTextView;
+import com.android.launcher3.Launcher;
+import com.android.launcher3.LauncherState;
+import com.android.launcher3.allapps.AllAppsRecyclerView;
+import com.android.launcher3.celllayout.FavoriteItemsTransaction;
+import com.android.launcher3.icons.ThemedIconDrawable;
+import com.android.launcher3.popup.ArrowPopup;
+import com.android.launcher3.util.BaseLauncherActivityTest;
+import com.android.launcher3.util.Executors;
+import com.android.launcher3.util.TestUtil;
+
+import org.junit.Test;
+
+import java.util.ArrayDeque;
+import java.util.Queue;
+
+/**
+ * Tests for theme icon support in Launcher
+ *
+ * Note running these tests will clear the workspace on the device.
+ */
+@LargeTest
+public class ThemeIconsTest extends BaseLauncherActivityTest<Launcher> {
+
+    private static final String APP_NAME = "IconThemedActivity";
+    private static final String SHORTCUT_NAME = "Shortcut 1";
+
+    @Test
+    public void testIconWithoutTheme() throws Exception {
+        setThemeEnabled(false);
+        new FavoriteItemsTransaction(targetContext()).commit();
+        loadLauncherSync();
+        goToState(LauncherState.ALL_APPS);
+        freezeAllApps();
+
+        scrollToAppIcon(APP_NAME);
+        BubbleTextView btv = getFromLauncher(
+                l -> verifyIconTheme(APP_NAME, l.getAppsView(), false));
+        addToWorkspace(btv);
+        executeOnLauncher(l -> verifyIconTheme(APP_NAME, l.getWorkspace(), false));
+    }
+
+    @Test
+    public void testShortcutIconWithoutTheme() throws Exception {
+        setThemeEnabled(false);
+        new FavoriteItemsTransaction(targetContext()).commit();
+        loadLauncherSync();
+        goToState(LauncherState.ALL_APPS);
+        freezeAllApps();
+
+        scrollToAppIcon(TEST_APP_NAME);
+        BubbleTextView btv = getFromLauncher(l -> findBtv(TEST_APP_NAME, l.getAppsView()));
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR, btv::performLongClick);
+
+        BubbleTextView menuItem = getOnceNotNull("Popup menu not open", l ->
+                (AbstractFloatingView.getOpenView(l, TYPE_ACTION_POPUP) instanceof ArrowPopup ap)
+                        ? findBtv(SHORTCUT_NAME, ap) : null);
+        addToWorkspace(menuItem);
+        executeOnLauncher(l -> verifyIconTheme(SHORTCUT_NAME, l.getWorkspace(), false));
+    }
+
+    @Test
+    public void testIconWithTheme() throws Exception {
+        setThemeEnabled(true);
+        new FavoriteItemsTransaction(targetContext()).commit();
+        loadLauncherSync();
+        goToState(LauncherState.ALL_APPS);
+        freezeAllApps();
+
+        scrollToAppIcon(APP_NAME);
+        BubbleTextView btv = getFromLauncher(l ->
+                verifyIconTheme(APP_NAME, l.getAppsView(), false));
+        addToWorkspace(btv);
+        executeOnLauncher(l -> verifyIconTheme(APP_NAME, l.getWorkspace(), true));
+    }
+
+    @Test
+    public void testShortcutIconWithTheme() throws Exception {
+        setThemeEnabled(true);
+        loadLauncherSync();
+        goToState(LauncherState.ALL_APPS);
+        freezeAllApps();
+
+        scrollToAppIcon(TEST_APP_NAME);
+        BubbleTextView btv = getFromLauncher(l -> findBtv(TEST_APP_NAME, l.getAppsView()));
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR, btv::performLongClick);
+
+        BubbleTextView menuItem = getOnceNotNull("Popup menu not open", l ->
+                (AbstractFloatingView.getOpenView(l, TYPE_ACTION_POPUP) instanceof ArrowPopup ap)
+                        ? findBtv(SHORTCUT_NAME, ap) : null);
+        addToWorkspace(menuItem);
+        executeOnLauncher(l -> verifyIconTheme(SHORTCUT_NAME, l.getWorkspace(), true));
+    }
+
+    private BubbleTextView findBtv(String title, ViewGroup parent) {
+        // Wait for Launcher model to be completed
+        try {
+            Executors.MODEL_EXECUTOR.submit(() -> { }).get();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+
+        // Find the app icon
+        Queue<View> viewQueue = new ArrayDeque<>();
+        viewQueue.add(parent);
+        BubbleTextView icon = null;
+        while (!viewQueue.isEmpty()) {
+            View view = viewQueue.poll();
+            if (view instanceof ViewGroup) {
+                parent = (ViewGroup) view;
+                for (int i = parent.getChildCount() - 1; i >= 0; i--) {
+                    viewQueue.add(parent.getChildAt(i));
+                }
+            } else if (view instanceof BubbleTextView btv) {
+                if (btv.getContentDescription() != null
+                        && title.equals(btv.getContentDescription().toString())) {
+                    icon = btv;
+                    break;
+                }
+            }
+        }
+        return icon;
+    }
+
+    private BubbleTextView verifyIconTheme(String title, ViewGroup parent, boolean isThemed) {
+        BubbleTextView icon = findBtv(title, parent);
+        assertNotNull(icon.getIcon());
+        assertEquals(isThemed, icon.getIcon() instanceof ThemedIconDrawable);
+        return icon;
+    }
+
+    private void setThemeEnabled(boolean isEnabled) throws Exception {
+        Uri uri = new Uri.Builder()
+                .scheme(ContentResolver.SCHEME_CONTENT)
+                .authority(targetContext().getPackageName() + ".grid_control")
+                .appendPath("set_icon_themed")
+                .build();
+        ContentValues values = new ContentValues();
+        values.put("boolean_value", isEnabled);
+        try (ContentProviderClient client = targetContext().getContentResolver()
+                .acquireContentProviderClient(uri)) {
+            int result = client.update(uri, values, null);
+            assertTrue(result > 0);
+        }
+    }
+
+    private void scrollToAppIcon(String appName) {
+        executeOnLauncher(l -> {
+            l.hideKeyboard();
+            AllAppsRecyclerView rv = l.getAppsView().getActiveRecyclerView();
+            int pos = rv.getApps().getAdapterItems().indexOf(rv.getApps().getAdapterItems().stream()
+                    .filter(i -> i.itemInfo != null && appName.equals(i.itemInfo.title.toString()))
+                    .findFirst()
+                    .get());
+            rv.getLayoutManager().scrollToPosition(pos);
+        });
+    }
+
+    private void addToWorkspace(View btv) {
+        TestUtil.runOnExecutorSync(MAIN_EXECUTOR, () ->
+                btv.getAccessibilityDelegate().performAccessibilityAction(
+                        btv, com.android.launcher3.R.id.action_add_to_workspace, null));
+        UiDevice.getInstance(getInstrumentation()).waitForIdle();
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/BaseLauncherActivityTest.kt b/tests/src/com/android/launcher3/util/BaseLauncherActivityTest.kt
index bacce40..476e497 100644
--- a/tests/src/com/android/launcher3/util/BaseLauncherActivityTest.kt
+++ b/tests/src/com/android/launcher3/util/BaseLauncherActivityTest.kt
@@ -69,8 +69,9 @@
 
     protected fun targetContext(): Context = getInstrumentation().targetContext
 
-    protected fun goToState(state: LauncherState) = executeOnLauncher {
-        it.stateManager.goToState(state, 0)
+    protected fun goToState(state: LauncherState) {
+        executeOnLauncher { it.stateManager.goToState(state, 0) }
+        UiDevice.getInstance(getInstrumentation()).waitForIdle()
     }
 
     protected fun executeOnLauncher(f: ActivityAction<LAUNCHER_TYPE>) = scenario.onActivity(f)
@@ -92,6 +93,12 @@
         condition: Function<LAUNCHER_TYPE, Boolean>,
     ) = atMost(message, { getFromLauncher(condition)!! })
 
+    protected fun waitForLauncherCondition(
+        message: String,
+        condition: Function<LAUNCHER_TYPE, Boolean>,
+        timeout: Long,
+    ) = atMost(message, { getFromLauncher(condition)!! }, null, timeout)
+
     protected fun <T> getOnceNotNull(message: String, f: Function<LAUNCHER_TYPE, T?>): T? {
         var output: T? = null
         atMost(
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 08c5552..fac73d3 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -1728,6 +1728,27 @@
         scrollDownByDistance(container, distance, appsListBottomPadding);
     }
 
+    /** Scrolls up by given distance within the container. */
+    void scrollUpByDistance(UiObject2 container, int distance) {
+        scrollUpByDistance(container, distance, 0);
+    }
+
+    /** Scrolls up by given distance within the container considering the given bottom padding. */
+    void scrollUpByDistance(UiObject2 container, int distance, int bottomPadding) {
+        final Rect containerRect = getVisibleBounds(container);
+        final int bottomGestureMarginInContainer = getBottomGestureMarginInContainer(container);
+        scroll(
+                container,
+                Direction.UP,
+                new Rect(
+                        0,
+                        containerRect.height() - bottomGestureMarginInContainer - distance,
+                        0,
+                        bottomGestureMarginInContainer + bottomPadding),
+                /* steps= */ 10,
+                /* slowDown= */ true);
+    }
+
     void scrollDownByDistance(UiObject2 container, int distance) {
         scrollDownByDistance(container, distance, 0);
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index 3097d9c..ac2748e 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -19,6 +19,7 @@
 import static com.android.launcher3.tapl.LauncherInstrumentation.WAIT_TIME_MS;
 import static com.android.launcher3.tapl.LauncherInstrumentation.log;
 
+import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.graphics.Rect;
 
@@ -31,6 +32,7 @@
 import com.android.launcher3.testing.shared.TestProtocol;
 
 import java.util.Collection;
+import java.util.List;
 
 /**
  * All widgets container.
@@ -128,8 +130,10 @@
             final UiObject2 searchBar = findSearchBar();
             final int searchBarHeight = searchBar.getVisibleBounds().height();
             final UiObject2 fullWidgetsPicker = verifyActiveContainer();
-            mLauncher.assertTrue("Widgets container didn't become scrollable",
-                    fullWidgetsPicker.wait(Until.scrollable(true), WAIT_TIME_MS));
+
+            // Widget picker may not be scrollable if there are few items. Instead of waiting on
+            // picker being scrollable, we wait on widget headers to be available.
+            waitForWidgetListItems(fullWidgetsPicker);
 
             final UiObject2 widgetsContainer =
                     findTestAppWidgetsTableContainer(testAppWidgetPackage);
@@ -176,6 +180,13 @@
         }
     }
 
+    private void waitForWidgetListItems(UiObject2 fullWidgetsPicker) {
+        List<UiObject2> headers = fullWidgetsPicker.wait(Until.findObjects(
+                By.res(mLauncher.getLauncherPackageName(), "widgets_list_header")), WAIT_TIME_MS);
+        mLauncher.assertTrue("Widgets list is not available",
+                headers != null && !headers.isEmpty());
+    }
+
     private UiObject2 findSearchBar() {
         final BySelector searchBarContainerSelector = By.res(mLauncher.getLauncherPackageName(),
                 "search_and_recommendations_container");
@@ -199,19 +210,38 @@
                 "container");
 
         String packageName =  mLauncher.getContext().getPackageName();
+        String packageNameToFind =
+                (testAppWidgetPackage == null || testAppWidgetPackage.isEmpty()) ? packageName
+                        : testAppWidgetPackage;
+
         final BySelector targetAppSelector = By
                 .clazz("android.widget.TextView")
-                .text((testAppWidgetPackage == null || testAppWidgetPackage.isEmpty())
-                                ? packageName
-                                : testAppWidgetPackage);
+                .text(packageNameToFind);
+        final BySelector expandListButtonSelector =
+                By.res(mLauncher.getLauncherPackageName(), "widget_list_expand_button");
         final BySelector widgetsContainerSelector = By.res(mLauncher.getLauncherPackageName(),
                 "widgets_table");
 
         boolean hasHeaderExpanded = false;
+        // List was expanded by clicking "Show all" button.
+        boolean hasListExpanded = false;
+
         int scrollDistance = 0;
         for (int i = 0; i < SCROLL_ATTEMPTS; i++) {
             UiObject2 widgetPicker = mLauncher.waitForLauncherObject(widgetPickerSelector);
             UiObject2 widgetListView = verifyActiveContainer();
+
+            // Press "Show all" button if it exists. Otherwise, keep scrolling to
+            // find the header or show all button.
+            UiObject2 expandListButton =
+                    mLauncher.findObjectInContainer(widgetListView, expandListButtonSelector);
+            if (expandListButton != null) {
+                expandListButton.click();
+                hasListExpanded = true;
+                i = -1;
+                continue;
+            }
+
             UiObject2 header = mLauncher.waitForObjectInContainer(widgetListView,
                     headerSelector);
             // If a header is barely visible in the bottom edge of the screen, its height could be
@@ -222,6 +252,17 @@
             // Look for a header that has the test app name.
             UiObject2 headerTitle = mLauncher.findObjectInContainer(widgetListView,
                     targetAppSelector);
+
+            final UiObject2 searchBar = findSearchBar();
+            // If header's title is under or above search bar, let's not process the header yet,
+            // scroll a bit more to bring the header into visible area.
+            if (headerTitle != null
+                    && headerTitle.getVisibleCenter().y <= searchBar.getVisibleCenter().y) {
+                log("Test app's header is behind the searchbar, scrolling up");
+                mLauncher.scrollUpByDistance(widgetListView, scrollDistance);
+                continue;
+            }
+
             if (headerTitle != null) {
                 // If we find the header and it has not been expanded, let's click it to see the
                 // widgets list. Note that we wait until the header is out of the gesture region at
@@ -258,11 +299,24 @@
                     widgetPicker,
                     widgetsContainerSelector);
 
-            mLauncher.scrollDownByDistance(hasHeaderExpanded && rightPane != null
-                    ? rightPane
-                    : widgetListView, scrollDistance);
+            if (hasListExpanded && packageNameToFind.compareToIgnoreCase(
+                    getFirstHeaderTitle(widgetListView)) < 0) {
+                mLauncher.scrollUpByDistance(hasHeaderExpanded && rightPane != null
+                        ? rightPane
+                        : widgetListView, scrollDistance);
+            } else {
+                mLauncher.scrollDownByDistance(hasHeaderExpanded && rightPane != null
+                        ? rightPane
+                        : widgetListView, scrollDistance);
+            }
         }
 
         return null;
     }
+
+    @NonNull
+    private String getFirstHeaderTitle(UiObject2 widgetListView) {
+        UiObject2 firstHeader = mLauncher.getObjectsInContainer(widgetListView, "app_title").get(0);
+        return firstHeader != null ? firstHeader.getText() : "";
+    }
 }