diff --git a/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
index 68558fa..0fb9718 100644
--- a/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
+++ b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
@@ -31,6 +31,7 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.SharedPreferences;
+import android.graphics.Bitmap;
 import android.graphics.Color;
 import android.graphics.Matrix;
 import android.graphics.drawable.ColorDrawable;
@@ -47,6 +48,7 @@
 import android.widget.TextView;
 
 import androidx.annotation.IntDef;
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.BaseActivity;
@@ -58,7 +60,6 @@
 import com.android.quickstep.views.GoOverviewActionsView;
 import com.android.quickstep.views.TaskContainer;
 import com.android.systemui.shared.recents.model.Task;
-import com.android.systemui.shared.recents.model.ThumbnailData;
 
 import java.lang.annotation.Retention;
 
@@ -131,7 +132,7 @@
          * Called when the current task is interactive for the user
          */
         @Override
-        public void initOverlay(Task task, ThumbnailData thumbnail, Matrix matrix,
+        public void initOverlay(Task task, @Nullable Bitmap thumbnail, Matrix matrix,
                 boolean rotated) {
             if (mDialog != null && mDialog.isShowing()) {
                 // Redraw the dialog in case the layout changed
diff --git a/quickstep/res/layout/task_desktop.xml b/quickstep/res/layout/task_desktop.xml
index 89e9b3d..453057c 100644
--- a/quickstep/res/layout/task_desktop.xml
+++ b/quickstep/res/layout/task_desktop.xml
@@ -36,15 +36,6 @@
         android:layout_width="match_parent"
         android:layout_height="match_parent" />
 
-    <!--
-         TODO(b249371338): DesktopTaskView extends from TaskView. TaskView expects TaskThumbnailView
-         and IconView with these ids to be present. Need to refactor RecentsView to accept child
-         views that do not inherint from TaskView only or create a generic TaskView that have
-         N number of tasks.
-     -->
-    <include layout="@layout/task_thumbnail"
-        android:visibility="gone" />
-
     <ViewStub
         android:id="@+id/icon"
         android:inflatedId="@id/icon"
diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
index 358d703..46501c4 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
@@ -245,11 +245,20 @@
         }
 
         void updateThumbnailInBackground(Task task, Consumer<ThumbnailData> callback) {
-            mModel.getThumbnailCache().updateThumbnailInBackground(task, callback);
+            mModel.getThumbnailCache().getThumbnailInBackground(task,
+                    thumbnailData -> {
+                        task.thumbnail = thumbnailData;
+                        callback.accept(thumbnailData);
+                    });
         }
 
         void updateIconInBackground(Task task, Consumer<Task> callback) {
-            mModel.getIconCache().updateIconInBackground(task, callback);
+            mModel.getIconCache().getIconInBackground(task, (icon, contentDescription, title) -> {
+                task.icon = icon;
+                task.titleDescription = contentDescription;
+                task.title = title;
+                callback.accept(task);
+            });
         }
 
         void onCloseComplete() {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 5a8a1a3..6c3b4ad 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -43,6 +43,8 @@
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_VOICE_INTERACTION_WINDOW_SHOWING;
 import static com.android.wm.shell.Flags.enableTinyTaskbar;
 
+import static java.lang.invoke.MethodHandles.Lookup.PROTECTED;
+
 import android.animation.AnimatorSet;
 import android.animation.ValueAnimator;
 import android.app.ActivityOptions;
@@ -1515,7 +1517,8 @@
         return mIsNavBarKidsMode && isThreeButtonNav();
     }
 
-    protected boolean isNavBarForceVisible() {
+    @VisibleForTesting(otherwise = PROTECTED)
+    public boolean isNavBarForceVisible() {
         return mIsNavBarForceVisible;
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index ee79fbf..b90e5fd 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -115,6 +115,7 @@
     private WindowManager mWindowManager;
     private FrameLayout mTaskbarRootLayout;
     private boolean mAddedWindow;
+    private boolean mIsSuspended;
     private final TaskbarNavButtonController mNavButtonController;
     private final ComponentCallbacks mComponentCallbacks;
 
@@ -443,6 +444,8 @@
      */
     @VisibleForTesting
     public synchronized void recreateTaskbar() {
+        if (mIsSuspended) return;
+
         Trace.beginSection("recreateTaskbar");
         try {
             DeviceProfile dp = mUserUnlocked ?
@@ -648,8 +651,22 @@
         }
     }
 
+    /**
+     * Removes Taskbar from the window manager and prevents recreation if {@code true}.
+     * <p>
+     * Suspending is for testing purposes only; avoid calling this method in production.
+     */
     @VisibleForTesting
-    public void addTaskbarRootViewToWindow() {
+    public void setSuspended(boolean isSuspended) {
+        mIsSuspended = isSuspended;
+        if (mIsSuspended) {
+            removeTaskbarRootViewFromWindow();
+        } else {
+            addTaskbarRootViewToWindow();
+        }
+    }
+
+    private void addTaskbarRootViewToWindow() {
         if (enableTaskbarNoRecreate() && !mAddedWindow && mTaskbarActivityContext != null) {
             mWindowManager.addView(mTaskbarRootLayout,
                     mTaskbarActivityContext.getWindowLayoutParams());
@@ -657,8 +674,7 @@
         }
     }
 
-    @VisibleForTesting
-    public void removeTaskbarRootViewFromWindow() {
+    private void removeTaskbarRootViewFromWindow() {
         if (enableTaskbarNoRecreate() && mAddedWindow) {
             mWindowManager.removeViewImmediate(mTaskbarRootLayout);
             mAddedWindow = false;
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
index 0a81f78..36828a8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
@@ -24,11 +24,9 @@
 import com.android.quickstep.RecentsModel
 import com.android.quickstep.util.DesktopTask
 import com.android.quickstep.util.GroupTask
-import com.android.systemui.shared.recents.model.Task
 import com.android.window.flags.Flags.enableDesktopWindowingMode
 import com.android.window.flags.Flags.enableDesktopWindowingTaskbarRunningApps
 import java.io.PrintWriter
-import java.util.function.Consumer
 
 /**
  * Provides recent apps functionality, when the Taskbar Recent Apps section is enabled. Behavior:
@@ -185,9 +183,16 @@
 
         for (groupTask in shownTasks) {
             for (task in groupTask.tasks) {
-                val callback =
-                    Consumer<Task> { controllers.taskbarViewController.onTaskUpdated(it) }
-                val cancellableTask = recentsModel.iconCache.updateIconInBackground(task, callback)
+                val cancellableTask =
+                    recentsModel.iconCache.getIconInBackground(task) {
+                        icon,
+                        contentDescription,
+                        title ->
+                        task.icon = icon
+                        task.titleDescription = contentDescription
+                        task.title = title
+                        controllers.taskbarViewController.onTaskUpdated(task)
+                    }
                 if (cancellableTask != null) {
                     iconLoadRequests.add(cancellableTask)
                 }
diff --git a/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt b/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt
index 45e5554..358f644 100644
--- a/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt
+++ b/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt
@@ -25,7 +25,7 @@
 import com.android.quickstep.views.RecentsViewContainer
 import com.android.quickstep.views.TaskContainer
 import com.android.wm.shell.common.desktopmode.DesktopModeTransitionSource
-import com.android.wm.shell.shared.DesktopModeStatus
+import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 
 /** A menu item, "Desktop", that allows the user to bring the current app into Desktop Windowing. */
 class DesktopSystemShortcut(
diff --git a/quickstep/src/com/android/quickstep/TaskIconCache.java b/quickstep/src/com/android/quickstep/TaskIconCache.java
index e6febff..b3a9199 100644
--- a/quickstep/src/com/android/quickstep/TaskIconCache.java
+++ b/quickstep/src/com/android/quickstep/TaskIconCache.java
@@ -55,7 +55,6 @@
 import com.android.systemui.shared.system.PackageManagerWrapper;
 
 import java.util.concurrent.Executor;
-import java.util.function.Consumer;
 
 /**
  * Manages the caching of task icons and related data.
@@ -103,21 +102,21 @@
      * @param callback The callback to receive the task after its data has been populated.
      * @return A cancelable handle to the request
      */
-    public CancellableTask updateIconInBackground(Task task, Consumer<Task> callback) {
+    public CancellableTask getIconInBackground(Task task, GetTaskIconCallback callback) {
         Preconditions.assertUIThread();
         if (task.icon != null) {
             // Nothing to load, the icon is already loaded
-            callback.accept(task);
+            callback.onTaskIconReceived(task.icon, task.titleDescription, task.title);
             return null;
         }
         CancellableTask<TaskCacheEntry> request = new CancellableTask<>(
                 () -> getCacheEntry(task),
                 MAIN_EXECUTOR,
                 result -> {
-                    task.icon = result.icon;
-                    task.titleDescription = result.contentDescription;
-                    task.title = result.title;
-                    callback.accept(task);
+                    callback.onTaskIconReceived(
+                            result.icon,
+                            result.contentDescription,
+                            result.title);
                     dispatchIconUpdate(task.key.id);
                 }
         );
@@ -280,6 +279,12 @@
         public String title = "";
     }
 
+    /** Callback used when retrieving app icons from cache. */
+    public interface GetTaskIconCallback {
+        /** Called when task icon is retrieved. */
+        void onTaskIconReceived(Drawable icon, String contentDescription, String title);
+    }
+
     void registerTaskVisualsChangeListener(TaskVisualsChangeListener newListener) {
         mTaskVisualsChangeListener = newListener;
     }
diff --git a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
index c243a24..80902e3 100644
--- a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
@@ -52,7 +52,6 @@
 import com.android.quickstep.views.TaskContainer;
 import com.android.quickstep.views.TaskView;
 import com.android.systemui.shared.recents.model.Task;
-import com.android.systemui.shared.recents.model.ThumbnailData;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -187,17 +186,6 @@
 
         /**
          * Called when the current task is interactive for the user
-         *
-         * @deprecated TODO(b/350931107): Remove this interface once TaskOverlayFactoryGo is updated
-         */
-        @Deprecated
-        public void initOverlay(Task task, ThumbnailData thumbnail, Matrix matrix,
-                boolean rotated) {
-            initOverlay(task, thumbnail.getThumbnail(), matrix, rotated);
-        }
-
-        /**
-         * Called when the current task is interactive for the user
          */
         public void initOverlay(Task task, @Nullable Bitmap thumbnail, Matrix matrix,
                 boolean rotated) {
diff --git a/quickstep/src/com/android/quickstep/TaskThumbnailCache.java b/quickstep/src/com/android/quickstep/TaskThumbnailCache.java
index 38e927f..3c6c3e4 100644
--- a/quickstep/src/com/android/quickstep/TaskThumbnailCache.java
+++ b/quickstep/src/com/android/quickstep/TaskThumbnailCache.java
@@ -131,8 +131,7 @@
         Preconditions.assertUIThread();
         // Fetch the thumbnail for this task and put it in the cache
         if (task.thumbnail == null) {
-            updateThumbnailInBackground(task.key, lowResolution,
-                    t -> task.thumbnail = t);
+            getThumbnailInBackground(task.key, lowResolution, t -> task.thumbnail = t);
         }
     }
 
@@ -145,13 +144,13 @@
     }
 
     /**
-     * Asynchronously fetches the icon and other task data for the given {@param task}.
+     * Asynchronously fetches the thumbnail for the given {@code task}.
      *
      * @param callback The callback to receive the task after its data has been populated.
      * @return A cancelable handle to the request
      */
     @Override
-    public CancellableTask<ThumbnailData> updateThumbnailInBackground(
+    public CancellableTask<ThumbnailData> getThumbnailInBackground(
             Task task, @NonNull Consumer<ThumbnailData> callback) {
         Preconditions.assertUIThread();
 
@@ -164,10 +163,7 @@
             return null;
         }
 
-        return updateThumbnailInBackground(task.key, !mHighResLoadingState.isEnabled(), t -> {
-            task.thumbnail = t;
-            callback.accept(t);
-        });
+        return getThumbnailInBackground(task.key, !mHighResLoadingState.isEnabled(), callback);
     }
 
     /**
@@ -187,7 +183,7 @@
         return newSize > oldSize;
     }
 
-    private CancellableTask<ThumbnailData> updateThumbnailInBackground(TaskKey key,
+    private CancellableTask<ThumbnailData> getThumbnailInBackground(TaskKey key,
             boolean lowResolution, Consumer<ThumbnailData> callback) {
         Preconditions.assertUIThread();
 
diff --git a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
index b21a1b4..9f3ef4a 100644
--- a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
@@ -67,14 +67,15 @@
         this.visibleTaskIds.value = visibleTaskIdList.toSet()
     }
 
-    /** Flow wrapper for [TaskThumbnailDataSource.updateThumbnailInBackground] api */
+    /** Flow wrapper for [TaskThumbnailDataSource.getThumbnailInBackground] api */
     private fun getThumbnailDataRequest(task: Task): ThumbnailDataRequest =
         flow {
                 emit(task.key.id to task.thumbnail)
                 val thumbnailDataResult: ThumbnailData? =
                     suspendCancellableCoroutine { continuation ->
                         val cancellableTask =
-                            taskThumbnailDataSource.updateThumbnailInBackground(task) {
+                            taskThumbnailDataSource.getThumbnailInBackground(task) {
+                                task.thumbnail = it
                                 continuation.resume(it)
                             }
                         continuation.invokeOnCancellation { cancellableTask?.cancel() }
@@ -94,12 +95,7 @@
                 tasks.filter { it.key.id in visibleIds }
             }
         val visibleThumbnailDataRequests: Flow<List<ThumbnailDataRequest>> =
-            visibleTasks.map {
-                it.map { visibleTask ->
-                    val taskCopy = Task(visibleTask).apply { thumbnail = visibleTask.thumbnail }
-                    getThumbnailDataRequest(taskCopy)
-                }
-            }
+            visibleTasks.map { visibleTasksList -> visibleTasksList.map(::getThumbnailDataRequest) }
         return visibleThumbnailDataRequests.flatMapLatest {
             thumbnailRequestFlows: List<ThumbnailDataRequest> ->
             if (thumbnailRequestFlows.isEmpty()) {
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
index dbe2b19..20a081b 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
@@ -30,6 +30,7 @@
 import android.view.ViewOutlineProvider
 import androidx.annotation.ColorInt
 import com.android.launcher3.Utilities
+import com.android.launcher3.util.ViewPool
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.BackgroundOnly
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Snapshot
@@ -42,7 +43,7 @@
 import kotlinx.coroutines.MainScope
 import kotlinx.coroutines.launch
 
-class TaskThumbnailView : View {
+class TaskThumbnailView : View, ViewPool.Reusable {
     // TODO(b/335649589): Ideally create and obtain this from DI. This ViewModel should be scoped
     //  to [TaskView], and also shared between [TaskView] and [TaskThumbnailView]
     //  This is using a lazy for now because the dependencies cannot be obtained without DI.
@@ -71,7 +72,7 @@
             return _measuredBounds
         }
 
-    private var cornerRadius: Float = TaskCornerRadius.get(context)
+    private var overviewCornerRadius: Float = TaskCornerRadius.get(context)
     private var fullscreenCornerRadius: Float = QuickStepContract.getWindowCornerRadius(context)
 
     constructor(context: Context?) : super(context)
@@ -100,7 +101,7 @@
                 invalidate()
             }
         }
-        MainScope().launch { viewModel.recentsFullscreenProgress.collect { invalidateOutline() } }
+        MainScope().launch { viewModel.cornerRadiusProgress.collect { invalidateOutline() } }
         MainScope().launch {
             viewModel.inheritedScale.collect { viewModelInheritedScale ->
                 inheritedScale = viewModelInheritedScale
@@ -117,6 +118,11 @@
             }
     }
 
+    override fun onRecycle() {
+        // Do nothing
+        uiState = Uninitialized
+    }
+
     override fun onDraw(canvas: Canvas) {
         when (val uiStateVal = uiState) {
             is Uninitialized -> drawBackgroundOnly(canvas, Color.BLACK)
@@ -138,7 +144,7 @@
     override fun onConfigurationChanged(newConfig: Configuration?) {
         super.onConfigurationChanged(newConfig)
 
-        cornerRadius = TaskCornerRadius.get(context)
+        overviewCornerRadius = TaskCornerRadius.get(context)
         fullscreenCornerRadius = QuickStepContract.getWindowCornerRadius(context)
         invalidateOutline()
     }
@@ -159,8 +165,8 @@
 
     private fun getCurrentCornerRadius() =
         Utilities.mapRange(
-            viewModel.recentsFullscreenProgress.value,
-            cornerRadius,
+            viewModel.cornerRadiusProgress.value,
+            overviewCornerRadius,
             fullscreenCornerRadius
         ) / inheritedScale
 
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModel.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModel.kt
index fe21174..d8729a6 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModel.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModel.kt
@@ -31,6 +31,7 @@
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flatMapLatest
@@ -47,7 +48,13 @@
     private val task = MutableStateFlow<Flow<Task?>>(flowOf(null))
     private var boundTaskIsRunning = false
 
-    val recentsFullscreenProgress = recentsViewData.fullscreenProgress
+    /**
+     * Progress for changes in corner radius. progress: 0 = overview corner radius; 1 = fullscreen
+     * corner radius.
+     */
+    val cornerRadiusProgress =
+        if (taskViewData.isOutlineFormedByThumbnailView) recentsViewData.fullscreenProgress
+        else MutableStateFlow(1f).asStateFlow()
     val inheritedScale =
         combine(recentsViewData.scale, taskViewData.scale) { recentsScale, taskScale ->
             recentsScale * taskScale
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskThumbnailDataSource.kt b/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskThumbnailDataSource.kt
index 55598f0..986acbe 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskThumbnailDataSource.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskThumbnailDataSource.kt
@@ -22,7 +22,7 @@
 import java.util.function.Consumer
 
 interface TaskThumbnailDataSource {
-    fun updateThumbnailInBackground(
+    fun getThumbnailInBackground(
         task: Task,
         callback: Consumer<ThumbnailData>
     ): CancellableTask<ThumbnailData>?
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskViewData.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskViewData.kt
index a8b5112..7a9ecf2 100644
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskViewData.kt
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskViewData.kt
@@ -16,9 +16,14 @@
 
 package com.android.quickstep.task.viewmodel
 
+import com.android.quickstep.views.TaskViewType
 import kotlinx.coroutines.flow.MutableStateFlow
 
-class TaskViewData {
+class TaskViewData(taskViewType: TaskViewType) {
     // This is typically a View concern but it is used to invalidate rendering in other Views
     val scale = MutableStateFlow(1f)
+
+    // TODO(b/331753115): This property should not be in TaskViewData once TaskView is MVVM.
+    /** Whether outline of TaskView is formed by outline thumbnail view(s). */
+    val isOutlineFormedByThumbnailView: Boolean = taskViewType != TaskViewType.DESKTOP
 }
diff --git a/quickstep/src/com/android/quickstep/util/DesktopTask.java b/quickstep/src/com/android/quickstep/util/DesktopTask.java
index 307b2fa..a727aa2 100644
--- a/quickstep/src/com/android/quickstep/util/DesktopTask.java
+++ b/quickstep/src/com/android/quickstep/util/DesktopTask.java
@@ -18,7 +18,7 @@
 
 import androidx.annotation.NonNull;
 
-import com.android.quickstep.views.TaskView;
+import com.android.quickstep.views.TaskViewType;
 import com.android.systemui.shared.recents.model.Task;
 
 import java.util.List;
@@ -34,7 +34,7 @@
     public final List<Task> tasks;
 
     public DesktopTask(@NonNull List<Task> tasks) {
-        super(tasks.get(0), null, null, TaskView.Type.DESKTOP);
+        super(tasks.get(0), null, null, TaskViewType.DESKTOP);
         this.tasks = tasks;
     }
 
diff --git a/quickstep/src/com/android/quickstep/util/GroupTask.java b/quickstep/src/com/android/quickstep/util/GroupTask.java
index e8b611c..fba08a9 100644
--- a/quickstep/src/com/android/quickstep/util/GroupTask.java
+++ b/quickstep/src/com/android/quickstep/util/GroupTask.java
@@ -20,7 +20,7 @@
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.util.SplitConfigurationOptions.SplitBounds;
-import com.android.quickstep.views.TaskView;
+import com.android.quickstep.views.TaskViewType;
 import com.android.systemui.shared.recents.model.Task;
 
 import java.util.Arrays;
@@ -39,19 +39,18 @@
     public final Task task2;
     @Nullable
     public final SplitBounds mSplitBounds;
-    @TaskView.Type
-    public final int taskViewType;
+    public final TaskViewType taskViewType;
 
     public GroupTask(@NonNull Task task) {
         this(task, null, null);
     }
 
     public GroupTask(@NonNull Task t1, @Nullable Task t2, @Nullable SplitBounds splitBounds) {
-        this(t1, t2, splitBounds, t2 != null ? TaskView.Type.GROUPED : TaskView.Type.SINGLE);
+        this(t1, t2, splitBounds, t2 != null ? TaskViewType.GROUPED : TaskViewType.SINGLE);
     }
 
     protected GroupTask(@NonNull Task t1, @Nullable Task t2, @Nullable SplitBounds splitBounds,
-            @TaskView.Type int taskViewType) {
+            TaskViewType taskViewType) {
         task1 = t1;
         task2 = t2;
         mSplitBounds = splitBounds;
diff --git a/quickstep/src/com/android/quickstep/util/SplitWithKeyboardShortcutController.java b/quickstep/src/com/android/quickstep/util/SplitWithKeyboardShortcutController.java
index 5e42b90..27fb31d 100644
--- a/quickstep/src/com/android/quickstep/util/SplitWithKeyboardShortcutController.java
+++ b/quickstep/src/com/android/quickstep/util/SplitWithKeyboardShortcutController.java
@@ -138,12 +138,13 @@
                     mLauncher, mLauncher.getDragLayer(),
                     controller.screenshotTask(runningTaskInfo.taskId).getThumbnail(),
                     null /* icon */, startingTaskRect);
+            Task task = Task.from(new Task.TaskKey(runningTaskInfo), runningTaskInfo,
+                    false /* isLocked */);
             RecentsModel.INSTANCE.get(mLauncher.getApplicationContext())
                     .getIconCache()
-                    .updateIconInBackground(
-                            Task.from(new Task.TaskKey(runningTaskInfo), runningTaskInfo,
-                                    false /* isLocked */),
-                            (task) -> floatingTaskView.setIcon(task.icon));
+                    .getIconInBackground(
+                            task,
+                            (icon, contentDescription, title) -> floatingTaskView.setIcon(icon));
             floatingTaskView.setAlpha(1);
             floatingTaskView.addStagingAnimation(anim, startingTaskRect, mTempRect,
                     false /* fadeWithThumbnail */, true /* isStagedTask */);
diff --git a/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt b/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
index 55bbd50..4333c8b 100644
--- a/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
@@ -27,6 +27,7 @@
 import android.view.View
 import android.view.ViewGroup
 import androidx.core.view.updateLayoutParams
+import com.android.launcher3.Flags.enableRefactorTaskThumbnail
 import com.android.launcher3.R
 import com.android.launcher3.util.RunnableList
 import com.android.launcher3.util.SplitConfigurationOptions
@@ -35,12 +36,13 @@
 import com.android.launcher3.util.rects.set
 import com.android.quickstep.BaseContainerInterface
 import com.android.quickstep.TaskOverlayFactory
+import com.android.quickstep.task.thumbnail.TaskThumbnailView
 import com.android.quickstep.util.RecentsOrientedState
 import com.android.systemui.shared.recents.model.Task
 
 /** TaskView that contains all tasks that are part of the desktop. */
 class DesktopTaskView @JvmOverloads constructor(context: Context, attrs: AttributeSet? = null) :
-    TaskView(context, attrs) {
+    TaskView(context, attrs, type = TaskViewType.DESKTOP) {
 
     private val snapshotDrawParams =
         object : FullscreenDrawParams(context) {
@@ -48,7 +50,7 @@
             override fun computeTaskCornerRadius(context: Context) =
                 computeWindowCornerRadius(context)
         }
-    private val taskThumbnailViewPool =
+    private val taskThumbnailViewDeprecatedPool =
         ViewPool<TaskThumbnailViewDeprecated>(
             context,
             this,
@@ -89,6 +91,66 @@
         childCountAtInflation = childCount
     }
 
+    /** Updates this desktop task to the gives task list defined in `tasks` */
+    fun bind(
+        tasks: List<Task>,
+        orientedState: RecentsOrientedState,
+        taskOverlayFactory: TaskOverlayFactory
+    ) {
+        if (DEBUG) {
+            val sb = StringBuilder()
+            sb.append("bind tasks=").append(tasks.size).append("\n")
+            tasks.forEach { sb.append(" key=${it.key}\n") }
+            Log.d(TAG, sb.toString())
+        }
+        cancelPendingLoadTasks()
+        taskContainers =
+            tasks.map { task ->
+                val snapshotView =
+                    if (enableRefactorTaskThumbnail()) {
+                            TaskThumbnailView(context)
+                        } else {
+                            taskThumbnailViewDeprecatedPool.view
+                        }
+                        .also { snapshotView ->
+                            addView(
+                                snapshotView,
+                                // Add snapshotView to the front after initial views e.g. icon and
+                                // background.
+                                childCountAtInflation,
+                                LayoutParams(
+                                    ViewGroup.LayoutParams.WRAP_CONTENT,
+                                    ViewGroup.LayoutParams.WRAP_CONTENT
+                                )
+                            )
+                        }
+                TaskContainer(
+                    this,
+                    task,
+                    snapshotView,
+                    iconView,
+                    TransformingTouchDelegate(iconView.asView()),
+                    SplitConfigurationOptions.STAGE_POSITION_UNDEFINED,
+                    digitalWellBeingToast = null,
+                    showWindowsView = null,
+                    taskOverlayFactory
+                )
+            }
+        taskContainers.forEach { it.bind() }
+        setOrientationState(orientedState)
+    }
+
+    override fun onRecycle() {
+        super.onRecycle()
+        visibility = VISIBLE
+        taskContainers.forEach {
+            if (!enableRefactorTaskThumbnail()) {
+                removeView(it.thumbnailViewDeprecated)
+                taskThumbnailViewDeprecatedPool.recycle(it.thumbnailViewDeprecated)
+            }
+        }
+    }
+
     override fun onMeasure(widthMeasureSpec: Int, heightMeasureSpec: Int) {
         super.onMeasure(widthMeasureSpec, heightMeasureSpec)
         val containerWidth = MeasureSpec.getSize(widthMeasureSpec)
@@ -151,77 +213,6 @@
         }
     }
 
-    override fun onRecycle() {
-        super.onRecycle()
-        visibility = VISIBLE
-    }
-
-    /** Updates this desktop task to the gives task list defined in `tasks` */
-    fun bind(
-        tasks: List<Task>,
-        orientedState: RecentsOrientedState,
-        taskOverlayFactory: TaskOverlayFactory
-    ) {
-        if (DEBUG) {
-            val sb = StringBuilder()
-            sb.append("bind tasks=").append(tasks.size).append("\n")
-            tasks.forEach { sb.append(" key=${it.key}\n") }
-            Log.d(TAG, sb.toString())
-        }
-        cancelPendingLoadTasks()
-
-        if (!isTaskContainersInitialized()) {
-            taskContainers = arrayListOf()
-        }
-        val taskContainers = taskContainers as ArrayList
-        taskContainers.ensureCapacity(tasks.size)
-        tasks.forEachIndexed { index, task ->
-            val thumbnailViewDeprecated: TaskThumbnailViewDeprecated
-            if (index >= taskContainers.size) {
-                thumbnailViewDeprecated = taskThumbnailViewPool.view
-                // Add thumbnailView from to position after the initial child views.
-                addView(
-                    thumbnailViewDeprecated,
-                    childCountAtInflation,
-                    LayoutParams(
-                        ViewGroup.LayoutParams.WRAP_CONTENT,
-                        ViewGroup.LayoutParams.WRAP_CONTENT
-                    )
-                )
-            } else {
-                thumbnailViewDeprecated = taskContainers[index].thumbnailViewDeprecated
-            }
-            val taskContainer =
-                TaskContainer(
-                    this,
-                    task,
-                    // TODO(b/338360089): Support new TTV for DesktopTaskView
-                    thumbnailView = null,
-                    thumbnailViewDeprecated,
-                    iconView,
-                    TransformingTouchDelegate(iconView.asView()),
-                    SplitConfigurationOptions.STAGE_POSITION_UNDEFINED,
-                    digitalWellBeingToast = null,
-                    showWindowsView = null,
-                    taskOverlayFactory
-                )
-            if (index >= taskContainers.size) {
-                taskContainers.add(taskContainer)
-            } else {
-                taskContainers[index] = taskContainer
-            }
-            taskContainer.bind()
-        }
-        repeat(taskContainers.size - tasks.size) {
-            with(taskContainers.removeLast()) {
-                removeView(thumbnailViewDeprecated)
-                taskThumbnailViewPool.recycle(thumbnailViewDeprecated)
-            }
-        }
-
-        setOrientationState(orientedState)
-    }
-
     override fun needsUpdate(dataChange: Int, flag: Int) =
         if (flag == FLAG_UPDATE_THUMBNAIL) super.needsUpdate(dataChange, flag) else false
 
diff --git a/quickstep/src/com/android/quickstep/views/GroupedTaskView.kt b/quickstep/src/com/android/quickstep/views/GroupedTaskView.kt
index b070244..6523ba7 100644
--- a/quickstep/src/com/android/quickstep/views/GroupedTaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/GroupedTaskView.kt
@@ -49,7 +49,7 @@
  * (Icon loading sold separately, fees may apply. Shipping & Handling for Overlays not included).
  */
 class GroupedTaskView @JvmOverloads constructor(context: Context, attrs: AttributeSet? = null) :
-    TaskView(context, attrs) {
+    TaskView(context, attrs, type = TaskViewType.GROUPED) {
     // TODO(b/336612373): Support new TTV for GroupedTaskView
     var splitBoundsConfig: SplitConfigurationOptions.SplitBounds? = null
         private set
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index ee1b3e7..7b6d383 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -221,7 +221,7 @@
 import com.android.systemui.shared.system.TaskStackChangeListeners;
 import com.android.wm.shell.common.desktopmode.DesktopModeTransitionSource;
 import com.android.wm.shell.common.pip.IPipAnimationListener;
-import com.android.wm.shell.shared.DesktopModeStatus;
+import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
 
 import kotlin.Unit;
 
@@ -1824,7 +1824,7 @@
             // If we need to remove half of a pair of tasks, force a TaskView with Type.SINGLE
             // to be a temporary container for the remaining task.
             TaskView taskView = getTaskViewFromPool(
-                    isRemovalNeeded ? TaskView.Type.SINGLE : groupTask.taskViewType);
+                    isRemovalNeeded ? TaskViewType.SINGLE : groupTask.taskViewType);
             if (taskView instanceof GroupedTaskView) {
                 boolean firstTaskIsLeftTopTask =
                         groupTask.mSplitBounds.leftTopTaskId == groupTask.task1.key.id;
@@ -2600,16 +2600,16 @@
      * Handle the edge case where Recents could increment task count very high over long
      * period of device usage. Probably will never happen, but meh.
      */
-    private TaskView getTaskViewFromPool(@TaskView.Type int type) {
+    private TaskView getTaskViewFromPool(TaskViewType type) {
         TaskView taskView;
         switch (type) {
-            case TaskView.Type.GROUPED:
+            case GROUPED:
                 taskView = mGroupedTaskViewPool.getView();
                 break;
-            case TaskView.Type.DESKTOP:
+            case DESKTOP:
                 taskView = mDesktopTaskViewPool.getView();
                 break;
-            case TaskView.Type.SINGLE:
+            case SINGLE:
             default:
                 taskView = mTaskViewPool.getView();
         }
@@ -2840,12 +2840,12 @@
             // Add an empty view for now until the task plan is loaded and applied
             final TaskView taskView;
             if (needDesktopTask) {
-                taskView = getTaskViewFromPool(TaskView.Type.DESKTOP);
+                taskView = getTaskViewFromPool(TaskViewType.DESKTOP);
                 mTmpRunningTasks = Arrays.copyOf(runningTasks, runningTasks.length);
                 ((DesktopTaskView) taskView).bind(Arrays.asList(mTmpRunningTasks),
                         mOrientationState, mTaskOverlayFactory);
             } else if (needGroupTaskView) {
-                taskView = getTaskViewFromPool(TaskView.Type.GROUPED);
+                taskView = getTaskViewFromPool(TaskViewType.GROUPED);
                 mTmpRunningTasks = new Task[]{runningTasks[0], runningTasks[1]};
                 // When we create a placeholder task view mSplitBoundsConfig will be null, but with
                 // the actual app running we won't need to show the thumbnail until all the tasks
@@ -2853,7 +2853,7 @@
                 ((GroupedTaskView)taskView).bind(mTmpRunningTasks[0], mTmpRunningTasks[1],
                         mOrientationState, mTaskOverlayFactory, mSplitBoundsConfig);
             } else {
-                taskView = getTaskViewFromPool(TaskView.Type.SINGLE);
+                taskView = getTaskViewFromPool(TaskViewType.SINGLE);
                 // The temporary running task is only used for the duration between the start of the
                 // gesture and the task list is loaded and applied
                 mTmpRunningTasks = new Task[]{runningTasks[0]};
diff --git a/quickstep/src/com/android/quickstep/views/TaskContainer.kt b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
index 2e01e7e..0648986 100644
--- a/quickstep/src/com/android/quickstep/views/TaskContainer.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
@@ -38,8 +38,7 @@
 class TaskContainer(
     val taskView: TaskView,
     val task: Task,
-    val thumbnailView: TaskThumbnailView?,
-    val thumbnailViewDeprecated: TaskThumbnailViewDeprecated,
+    val snapshotView: View,
     val iconView: TaskViewIcon,
     /**
      * This technically can be a vanilla [android.view.TouchDelegate] class, however that class
@@ -57,12 +56,29 @@
     val overlay: TaskOverlayFactory.TaskOverlay<*> = taskOverlayFactory.createOverlay(this)
     val taskContainerData = TaskContainerData()
 
-    val snapshotView: View
-        get() = thumbnailView ?: thumbnailViewDeprecated
+    init {
+        if (enableRefactorTaskThumbnail()) {
+            require(snapshotView is TaskThumbnailView)
+        } else {
+            require(snapshotView is TaskThumbnailViewDeprecated)
+        }
+    }
+
+    val thumbnailView: TaskThumbnailView
+        get() {
+            require(enableRefactorTaskThumbnail())
+            return snapshotView as TaskThumbnailView
+        }
+
+    val thumbnailViewDeprecated: TaskThumbnailViewDeprecated
+        get() {
+            require(!enableRefactorTaskThumbnail())
+            return snapshotView as TaskThumbnailViewDeprecated
+        }
 
     // TODO(b/349120849): Extract ThumbnailData from TaskContainerData/TaskThumbnailViewModel
     val thumbnail: Bitmap?
-        get() = thumbnailViewDeprecated.thumbnail
+        get() = if (enableRefactorTaskThumbnail()) null else thumbnailViewDeprecated.thumbnail
 
     // TODO(b/334826842): Support shouldShowSplashView for new TTV.
     val shouldShowSplashView: Boolean
@@ -100,13 +116,14 @@
 
     fun destroy() {
         digitalWellBeingToast?.destroy()
-        thumbnailView?.let { taskView.removeView(it) }
+        if (enableRefactorTaskThumbnail()) {
+            taskView.removeView(thumbnailView)
+        }
         overlay.destroy()
     }
 
     fun bind() {
-        if (enableRefactorTaskThumbnail() && thumbnailView != null) {
-            thumbnailViewDeprecated.setTaskOverlay(overlay)
+        if (enableRefactorTaskThumbnail()) {
             bindThumbnailView()
             overlay.init()
         } else {
@@ -119,7 +136,7 @@
     fun bindThumbnailView() {
         // TODO(b/343364498): Existing view has shouldShowScreenshot as an override as well but
         //  this should be decided inside TaskThumbnailViewModel.
-        thumbnailView?.viewModel?.bind(TaskThumbnail(task.key.id, taskView.isRunningTask))
+        thumbnailView.viewModel.bind(TaskThumbnail(task.key.id, taskView.isRunningTask))
     }
 
     fun setOverlayEnabled(enabled: Boolean) {
diff --git a/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java b/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
index 5b7e6c7..56ca043 100644
--- a/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
+++ b/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
@@ -165,17 +165,6 @@
     }
 
     /**
-     * Sets TaskOverlay without binding a task.
-     *
-     * @deprecated Should only be used when using new
-     * {@link com.android.quickstep.task.thumbnail.TaskThumbnailView}.
-     */
-    @Deprecated
-    public void setTaskOverlay(TaskOverlay<?> overlay) {
-        mOverlay = overlay;
-    }
-
-    /**
      * Updates the thumbnail.
      *
      * @param refreshNow whether the {@code thumbnailData} will be used to redraw immediately.
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index 888b24a..2e07e36 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -102,7 +102,8 @@
     defStyleAttr: Int = 0,
     defStyleRes: Int = 0,
     focusBorderAnimator: BorderAnimator? = null,
-    hoverBorderAnimator: BorderAnimator? = null
+    hoverBorderAnimator: BorderAnimator? = null,
+    type: TaskViewType = TaskViewType.SINGLE
 ) : FrameLayout(context, attrs), ViewPool.Reusable {
     /**
      * Used in conjunction with [onTaskListVisibilityChanged], providing more granularity on which
@@ -112,18 +113,7 @@
     @IntDef(FLAG_UPDATE_ALL, FLAG_UPDATE_ICON, FLAG_UPDATE_THUMBNAIL, FLAG_UPDATE_CORNER_RADIUS)
     annotation class TaskDataChanges
 
-    /** Type of task view */
-    @Retention(AnnotationRetention.SOURCE)
-    @IntDef(Type.SINGLE, Type.GROUPED, Type.DESKTOP)
-    annotation class Type {
-        companion object {
-            const val SINGLE = 1
-            const val GROUPED = 2
-            const val DESKTOP = 3
-        }
-    }
-
-    val taskViewData = TaskViewData()
+    val taskViewData = TaskViewData(type)
     val taskIds: IntArray
         /** Returns a copy of integer array containing taskIds of all tasks in the TaskView. */
         get() = taskContainers.map { it.task.key.id }.toIntArray()
@@ -671,24 +661,22 @@
         taskOverlayFactory: TaskOverlayFactory
     ): TaskContainer {
         val thumbnailViewDeprecated: TaskThumbnailViewDeprecated = findViewById(thumbnailViewId)!!
-        val thumbnailView: TaskThumbnailView?
-        if (enableRefactorTaskThumbnail()) {
-            val indexOfSnapshotView = indexOfChild(thumbnailViewDeprecated)
-            thumbnailView =
+        val snapshotView =
+            if (enableRefactorTaskThumbnail()) {
+                thumbnailViewDeprecated.visibility = GONE
+                val indexOfSnapshotView = indexOfChild(thumbnailViewDeprecated)
                 TaskThumbnailView(context).apply {
                     layoutParams = thumbnailViewDeprecated.layoutParams
                     addView(this, indexOfSnapshotView)
                 }
-            thumbnailViewDeprecated.visibility = GONE
-        } else {
-            thumbnailView = null
-        }
+            } else {
+                thumbnailViewDeprecated
+            }
         val iconView = getOrInflateIconView(iconViewId)
         return TaskContainer(
             this,
             task,
-            thumbnailView,
-            thumbnailViewDeprecated,
+            snapshotView,
             iconView,
             TransformingTouchDelegate(iconView.asView()),
             stagePosition,
@@ -710,8 +698,6 @@
                 .inflate() as TaskViewIcon
     }
 
-    protected fun isTaskContainersInitialized() = this::taskContainers.isInitialized
-
     fun containsMultipleTasks() = taskContainers.size > 1
 
     /**
@@ -846,7 +832,8 @@
             taskContainers.forEach {
                 if (visible) {
                     recentsModel.thumbnailCache
-                        .updateThumbnailInBackground(it.task) { thumbnailData ->
+                        .getThumbnailInBackground(it.task) { thumbnailData ->
+                            it.task.thumbnail = thumbnailData
                             it.thumbnailViewDeprecated.setThumbnail(it.task, thumbnailData)
                         }
                         ?.also { request -> pendingThumbnailLoadRequests.add(request) }
@@ -862,12 +849,15 @@
             taskContainers.forEach {
                 if (visible) {
                     recentsModel.iconCache
-                        .updateIconInBackground(it.task) { task ->
-                            setIcon(it.iconView, task.icon)
+                        .getIconInBackground(it.task) { icon, contentDescription, title ->
+                            it.task.icon = icon
+                            it.task.titleDescription = contentDescription
+                            it.task.title = title
+                            setIcon(it.iconView, icon)
                             if (enableOverviewIconMenu()) {
-                                setText(it.iconView, task.title)
+                                setText(it.iconView, title)
                             }
-                            it.digitalWellBeingToast?.initialize(task)
+                            it.digitalWellBeingToast?.initialize(it.task)
                         }
                         ?.also { request -> pendingIconLoadRequests.add(request) }
                 } else {
diff --git a/quickstep/src/com/android/quickstep/views/TaskViewType.kt b/quickstep/src/com/android/quickstep/views/TaskViewType.kt
new file mode 100644
index 0000000..b2a32a9
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/views/TaskViewType.kt
@@ -0,0 +1,24 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.views
+
+/** Type of the [TaskView] */
+enum class TaskViewType {
+    SINGLE,
+    GROUPED,
+    DESKTOP
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
index 9ecd935..2f0b446 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
@@ -20,7 +20,6 @@
 import android.content.ComponentName
 import android.content.Intent
 import android.os.Process
-import androidx.test.annotation.UiThreadTest
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.BubbleTextView
 import com.android.launcher3.appprediction.PredictionRowView
@@ -34,6 +33,7 @@
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
 import com.android.launcher3.util.PackageUserKey
+import com.android.launcher3.util.TestUtil
 import com.google.common.truth.Truth.assertThat
 import org.junit.Rule
 import org.junit.Test
@@ -55,17 +55,17 @@
     @InjectController lateinit var overlayController: TaskbarOverlayController
 
     @Test
-    @UiThreadTest
     fun testToggle_once_showsAllApps() {
-        allAppsController.toggle()
+        getInstrumentation().runOnMainSync { allAppsController.toggle() }
         assertThat(allAppsController.isOpen).isTrue()
     }
 
     @Test
-    @UiThreadTest
     fun testToggle_twice_closesAllApps() {
-        allAppsController.toggle()
-        allAppsController.toggle()
+        getInstrumentation().runOnMainSync {
+            allAppsController.toggle()
+            allAppsController.toggle()
+        }
         assertThat(allAppsController.isOpen).isFalse()
     }
 
@@ -77,54 +77,62 @@
     }
 
     @Test
-    @UiThreadTest
     fun testSetApps_beforeOpened_cachesInfo() {
-        allAppsController.setApps(TEST_APPS, 0, emptyMap())
-        allAppsController.toggle()
+        val overlayContext =
+            TestUtil.getOnUiThread {
+                allAppsController.setApps(TEST_APPS, 0, emptyMap())
+                allAppsController.toggle()
+                overlayController.requestWindow()
+            }
 
-        val overlayContext = overlayController.requestWindow()
         assertThat(overlayContext.appsView.appsStore.apps).isEqualTo(TEST_APPS)
     }
 
     @Test
-    @UiThreadTest
     fun testSetApps_afterOpened_updatesStore() {
-        allAppsController.toggle()
-        allAppsController.setApps(TEST_APPS, 0, emptyMap())
+        val overlayContext =
+            TestUtil.getOnUiThread {
+                allAppsController.toggle()
+                allAppsController.setApps(TEST_APPS, 0, emptyMap())
+                overlayController.requestWindow()
+            }
 
-        val overlayContext = overlayController.requestWindow()
         assertThat(overlayContext.appsView.appsStore.apps).isEqualTo(TEST_APPS)
     }
 
     @Test
-    @UiThreadTest
     fun testSetPredictedApps_beforeOpened_cachesInfo() {
-        allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
-        allAppsController.toggle()
-
         val predictedApps =
-            overlayController
-                .requestWindow()
-                .appsView
-                .floatingHeaderView
-                .findFixedRowByType(PredictionRowView::class.java)
-                .predictedApps
+            TestUtil.getOnUiThread {
+                allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
+                allAppsController.toggle()
+
+                overlayController
+                    .requestWindow()
+                    .appsView
+                    .floatingHeaderView
+                    .findFixedRowByType(PredictionRowView::class.java)
+                    .predictedApps
+            }
+
         assertThat(predictedApps).isEqualTo(TEST_PREDICTED_APPS)
     }
 
     @Test
-    @UiThreadTest
     fun testSetPredictedApps_afterOpened_cachesInfo() {
-        allAppsController.toggle()
-        allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
-
         val predictedApps =
-            overlayController
-                .requestWindow()
-                .appsView
-                .floatingHeaderView
-                .findFixedRowByType(PredictionRowView::class.java)
-                .predictedApps
+            TestUtil.getOnUiThread {
+                allAppsController.toggle()
+                allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
+
+                overlayController
+                    .requestWindow()
+                    .appsView
+                    .floatingHeaderView
+                    .findFixedRowByType(PredictionRowView::class.java)
+                    .predictedApps
+            }
+
         assertThat(predictedApps).isEqualTo(TEST_PREDICTED_APPS)
     }
 
@@ -140,36 +148,38 @@
         }
 
         // Ensure the recycler view fully inflates before trying to grab an icon.
-        getInstrumentation().runOnMainSync {
-            val btv =
+        val btv =
+            TestUtil.getOnUiThread {
                 overlayController
                     .requestWindow()
                     .appsView
                     .activeRecyclerView
                     .findViewHolderForAdapterPosition(0)
                     ?.itemView as? BubbleTextView
-            assertThat(btv?.hasDot()).isTrue()
-        }
+            }
+        assertThat(btv?.hasDot()).isTrue()
     }
 
     @Test
-    @UiThreadTest
     fun testUpdateNotificationDots_predictedApp_hasDot() {
-        allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
-        allAppsController.toggle()
+        getInstrumentation().runOnMainSync {
+            allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
+            allAppsController.toggle()
+            taskbarUnitTestRule.activityContext.popupDataProvider.onNotificationPosted(
+                PackageUserKey.fromItemInfo(TEST_PREDICTED_APPS[0]),
+                NotificationKeyData("key"),
+            )
+        }
 
-        taskbarUnitTestRule.activityContext.popupDataProvider.onNotificationPosted(
-            PackageUserKey.fromItemInfo(TEST_PREDICTED_APPS[0]),
-            NotificationKeyData("key"),
-        )
-
-        val predictionRowView =
-            overlayController
-                .requestWindow()
-                .appsView
-                .floatingHeaderView
-                .findFixedRowByType(PredictionRowView::class.java)
-        val btv = predictionRowView.getChildAt(0) as BubbleTextView
+        val btv =
+            TestUtil.getOnUiThread {
+                overlayController
+                    .requestWindow()
+                    .appsView
+                    .floatingHeaderView
+                    .findFixedRowByType(PredictionRowView::class.java)
+                    .getChildAt(0) as BubbleTextView
+            }
         assertThat(btv.hasDot()).isTrue()
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
index fae5562..f946d4d 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
@@ -18,7 +18,6 @@
 
 import android.app.ActivityManager.RunningTaskInfo
 import android.view.MotionEvent
-import androidx.test.annotation.UiThreadTest
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.AbstractFloatingView
 import com.android.launcher3.AbstractFloatingView.TYPE_OPTIONS_POPUP
@@ -31,7 +30,7 @@
 import com.android.launcher3.taskbar.rules.TaskbarWindowSandboxContext
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
-import com.android.launcher3.views.BaseDragLayer
+import com.android.launcher3.util.TestUtil.getOnUiThread
 import com.android.systemui.shared.system.TaskStackChangeListeners
 import com.google.common.truth.Truth.assertThat
 import org.junit.Rule
@@ -54,74 +53,69 @@
         get() = taskbarUnitTestRule.activityContext
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_twice_reusesWindow() {
-        val context1 = overlayController.requestWindow()
-        val context2 = overlayController.requestWindow()
+        val (context1, context2) =
+            getOnUiThread {
+                Pair(overlayController.requestWindow(), overlayController.requestWindow())
+            }
         assertThat(context1).isSameInstanceAs(context2)
     }
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_afterHidingExistingWindow_createsNewWindow() {
-        val context1 = overlayController.requestWindow()
-        overlayController.hideWindow()
+        val context1 = getOnUiThread { overlayController.requestWindow() }
+        getInstrumentation().runOnMainSync { overlayController.hideWindow() }
 
-        val context2 = overlayController.requestWindow()
+        val context2 = getOnUiThread { overlayController.requestWindow() }
         assertThat(context1).isNotSameInstanceAs(context2)
     }
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_afterHidingOverlay_createsNewWindow() {
-        val context1 = overlayController.requestWindow()
-        TestOverlayView.show(context1)
-        overlayController.hideWindow()
+        val context1 = getOnUiThread { overlayController.requestWindow() }
+        getInstrumentation().runOnMainSync {
+            TestOverlayView.show(context1)
+            overlayController.hideWindow()
+        }
 
-        val context2 = overlayController.requestWindow()
+        val context2 = getOnUiThread { overlayController.requestWindow() }
         assertThat(context1).isNotSameInstanceAs(context2)
     }
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_addsProxyView() {
-        TestOverlayView.show(overlayController.requestWindow())
+        getInstrumentation().runOnMainSync {
+            TestOverlayView.show(overlayController.requestWindow())
+        }
         assertThat(hasOpenView(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)).isTrue()
     }
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_closeProxyView_closesOverlay() {
-        val overlay = TestOverlayView.show(overlayController.requestWindow())
-        AbstractFloatingView.closeOpenContainer(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
+        getInstrumentation().runOnMainSync {
+            AbstractFloatingView.closeOpenContainer(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)
+        }
         assertThat(overlay.isOpen).isFalse()
     }
 
     @Test
     fun testRequestWindow_attachesDragLayer() {
-        lateinit var dragLayer: BaseDragLayer<*>
-        getInstrumentation().runOnMainSync {
-            dragLayer = overlayController.requestWindow().dragLayer
-        }
-
+        val dragLayer = getOnUiThread { overlayController.requestWindow().dragLayer }
         // Allow drag layer to attach before checking.
         getInstrumentation().runOnMainSync { assertThat(dragLayer.isAttachedToWindow).isTrue() }
     }
 
     @Test
-    @UiThreadTest
     fun testHideWindow_closesOverlay() {
-        val overlay = TestOverlayView.show(overlayController.requestWindow())
-        overlayController.hideWindow()
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
+        getInstrumentation().runOnMainSync { overlayController.hideWindow() }
         assertThat(overlay.isOpen).isFalse()
     }
 
     @Test
     fun testHideWindow_detachesDragLayer() {
-        lateinit var dragLayer: BaseDragLayer<*>
-        getInstrumentation().runOnMainSync {
-            dragLayer = overlayController.requestWindow().dragLayer
-        }
+        val dragLayer = getOnUiThread { overlayController.requestWindow().dragLayer }
 
         // Wait for drag layer to be attached to window before hiding.
         getInstrumentation().runOnMainSync {
@@ -131,26 +125,30 @@
     }
 
     @Test
-    @UiThreadTest
     fun testTwoOverlays_closeOne_windowStaysOpen() {
-        val context = overlayController.requestWindow()
-        val overlay1 = TestOverlayView.show(context)
-        val overlay2 = TestOverlayView.show(context)
+        val (overlay1, overlay2) =
+            getOnUiThread {
+                val context = overlayController.requestWindow()
+                Pair(TestOverlayView.show(context), TestOverlayView.show(context))
+            }
 
-        overlay1.close(false)
+        getInstrumentation().runOnMainSync { overlay1.close(false) }
         assertThat(overlay2.isOpen).isTrue()
         assertThat(hasOpenView(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)).isTrue()
     }
 
     @Test
-    @UiThreadTest
     fun testTwoOverlays_closeAll_closesWindow() {
-        val context = overlayController.requestWindow()
-        val overlay1 = TestOverlayView.show(context)
-        val overlay2 = TestOverlayView.show(context)
+        val (overlay1, overlay2) =
+            getOnUiThread {
+                val context = overlayController.requestWindow()
+                Pair(TestOverlayView.show(context), TestOverlayView.show(context))
+            }
 
-        overlay1.close(false)
-        overlay2.close(false)
+        getInstrumentation().runOnMainSync {
+            overlay1.close(false)
+            overlay2.close(false)
+        }
         assertThat(hasOpenView(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)).isFalse()
     }
 
@@ -165,11 +163,7 @@
 
     @Test
     fun testTaskMovedToFront_closesOverlay() {
-        lateinit var overlay: TestOverlayView
-        getInstrumentation().runOnMainSync {
-            overlay = TestOverlayView.show(overlayController.requestWindow())
-        }
-
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
         TaskStackChangeListeners.getInstance().listenerImpl.onTaskMovedToFront(RunningTaskInfo())
         // Make sure TaskStackChangeListeners' Handler posts the callback before checking state.
         getInstrumentation().runOnMainSync { assertThat(overlay.isOpen).isFalse() }
@@ -177,9 +171,8 @@
 
     @Test
     fun testTaskStackChanged_allAppsClosed_overlayStaysOpen() {
-        lateinit var overlay: TestOverlayView
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
         getInstrumentation().runOnMainSync {
-            overlay = TestOverlayView.show(overlayController.requestWindow())
             taskbarContext.controllers.sharedState?.allAppsVisible = false
         }
 
@@ -189,9 +182,8 @@
 
     @Test
     fun testTaskStackChanged_allAppsOpen_closesOverlay() {
-        lateinit var overlay: TestOverlayView
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
         getInstrumentation().runOnMainSync {
-            overlay = TestOverlayView.show(overlayController.requestWindow())
             taskbarContext.controllers.sharedState?.allAppsVisible = true
         }
 
@@ -200,33 +192,39 @@
     }
 
     @Test
-    @UiThreadTest
     fun testUpdateLauncherDeviceProfile_overlayNotRebindSafe_closesOverlay() {
-        val overlayContext = overlayController.requestWindow()
-        val overlay = TestOverlayView.show(overlayContext).apply { type = TYPE_OPTIONS_POPUP }
+        val context = getOnUiThread { overlayController.requestWindow() }
+        val overlay = getOnUiThread {
+            TestOverlayView.show(context).apply { type = TYPE_OPTIONS_POPUP }
+        }
 
-        overlayController.updateLauncherDeviceProfile(
-            overlayController.launcherDeviceProfile
-                .toBuilder(overlayContext)
-                .setGestureMode(false)
-                .build()
-        )
+        getInstrumentation().runOnMainSync {
+            overlayController.updateLauncherDeviceProfile(
+                overlayController.launcherDeviceProfile
+                    .toBuilder(context)
+                    .setGestureMode(false)
+                    .build()
+            )
+        }
 
         assertThat(overlay.isOpen).isFalse()
     }
 
     @Test
-    @UiThreadTest
     fun testUpdateLauncherDeviceProfile_overlayRebindSafe_overlayStaysOpen() {
-        val overlayContext = overlayController.requestWindow()
-        val overlay = TestOverlayView.show(overlayContext).apply { type = TYPE_TASKBAR_ALL_APPS }
+        val context = getOnUiThread { overlayController.requestWindow() }
+        val overlay = getOnUiThread {
+            TestOverlayView.show(context).apply { type = TYPE_TASKBAR_ALL_APPS }
+        }
 
-        overlayController.updateLauncherDeviceProfile(
-            overlayController.launcherDeviceProfile
-                .toBuilder(overlayContext)
-                .setGestureMode(false)
-                .build()
-        )
+        getInstrumentation().runOnMainSync {
+            overlayController.updateLauncherDeviceProfile(
+                overlayController.launcherDeviceProfile
+                    .toBuilder(context)
+                    .setGestureMode(false)
+                    .build()
+            )
+        }
 
         assertThat(overlay.isOpen).isTrue()
     }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
index 6638736..c48947e 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
@@ -16,6 +16,7 @@
 
 package com.android.launcher3.taskbar.rules
 
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode
 import com.android.launcher3.taskbar.rules.TaskbarModeRule.TaskbarMode
 import com.android.launcher3.util.DisplayController
@@ -59,23 +60,25 @@
             override fun evaluate() {
                 val mode = taskbarMode.mode
 
-                context.applicationContext.putObject(
-                    DisplayController.INSTANCE,
-                    object : DisplayController(context) {
-                        override fun getInfo(): Info {
-                            return spy(super.getInfo()) {
-                                on { isTransientTaskbar } doReturn (mode == Mode.TRANSIENT)
-                                on { isPinnedTaskbar } doReturn (mode == Mode.PINNED)
-                                on { navigationMode } doReturn
-                                    when (mode) {
-                                        Mode.TRANSIENT,
-                                        Mode.PINNED -> NavigationMode.NO_BUTTON
-                                        Mode.THREE_BUTTONS -> NavigationMode.THREE_BUTTONS
-                                    }
+                getInstrumentation().runOnMainSync {
+                    context.applicationContext.putObject(
+                        DisplayController.INSTANCE,
+                        object : DisplayController(context) {
+                            override fun getInfo(): Info {
+                                return spy(super.getInfo()) {
+                                    on { isTransientTaskbar } doReturn (mode == Mode.TRANSIENT)
+                                    on { isPinnedTaskbar } doReturn (mode == Mode.PINNED)
+                                    on { navigationMode } doReturn
+                                        when (mode) {
+                                            Mode.TRANSIENT,
+                                            Mode.PINNED -> NavigationMode.NO_BUTTON
+                                            Mode.THREE_BUTTONS -> NavigationMode.THREE_BUTTONS
+                                        }
+                                }
                             }
-                        }
-                    },
-                )
+                        },
+                    )
+                }
 
                 base.evaluate()
             }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
index 12f946e..8a64949 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
@@ -20,18 +20,25 @@
 import android.app.PendingIntent
 import android.content.IIntentSender
 import android.content.Intent
+import android.provider.Settings
+import android.provider.Settings.Secure.NAV_BAR_KIDS_MODE
+import android.provider.Settings.Secure.USER_SETUP_COMPLETE
 import androidx.test.platform.app.InstrumentationRegistry
 import androidx.test.rule.ServiceTestRule
 import com.android.launcher3.LauncherAppState
 import com.android.launcher3.taskbar.TaskbarActivityContext
 import com.android.launcher3.taskbar.TaskbarManager
 import com.android.launcher3.taskbar.TaskbarNavButtonController.TaskbarNavButtonCallbacks
+import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR
 import com.android.launcher3.util.LauncherMultivalentJUnit.Companion.isRunningInRobolectric
+import com.android.launcher3.util.ModelTestExtensions.loadModelSync
+import com.android.launcher3.util.TestUtil
 import com.android.quickstep.AllAppsActionManager
 import com.android.quickstep.TouchInteractionService
 import com.android.quickstep.TouchInteractionService.TISBinder
 import org.junit.Assume.assumeTrue
+import org.junit.rules.RuleChain
 import org.junit.rules.TestRule
 import org.junit.runner.Description
 import org.junit.runners.model.Statement
@@ -48,12 +55,11 @@
  * that code that is executed on the main thread in production should also happen on that thread
  * when tested.
  *
- * `@UiThreadTest` is a simple way to run an entire test body on the main thread. But if a test
- * executes code that appends message(s) to the main thread's `MessageQueue`, the annotation will
- * prevent those messages from being processed until after the test body finishes.
+ * `@UiThreadTest` is incompatible with this rule. The annotation causes this rule to run on the
+ * main thread, but it needs to be run on the test thread for it to work properly. Instead, only run
+ * code that requires the main thread using something like [Instrumentation.runOnMainSync] or
+ * [TestUtil.getOnUiThread].
  *
- * To test pending messages, instead use something like [Instrumentation.runOnMainSync] to perform
- * only sections of the test body on the main thread synchronously:
  * ```
  * @Test
  * fun example() {
@@ -71,6 +77,10 @@
     private val instrumentation = InstrumentationRegistry.getInstrumentation()
     private val serviceTestRule = ServiceTestRule()
 
+    private val userSetupCompleteRule = TaskbarSecureSettingRule(USER_SETUP_COMPLETE)
+    private val kidsModeRule = TaskbarSecureSettingRule(NAV_BAR_KIDS_MODE)
+    private val settingRules = RuleChain.outerRule(userSetupCompleteRule).around(kidsModeRule)
+
     private lateinit var taskbarManager: TaskbarManager
 
     val activityContext: TaskbarActivityContext
@@ -80,15 +90,34 @@
         }
 
     override fun apply(base: Statement, description: Description): Statement {
+        return settingRules.apply(createStatement(base, description), description)
+    }
+
+    private fun createStatement(base: Statement, description: Description): Statement {
         return object : Statement() {
             override fun evaluate() {
 
+                // Only run test when Taskbar is enabled.
                 instrumentation.runOnMainSync {
                     assumeTrue(
                         LauncherAppState.getIDP(context).getDeviceProfile(context).isTaskbarPresent
                     )
                 }
 
+                // Process secure setting annotations.
+                instrumentation.runOnMainSync {
+                    userSetupCompleteRule.putInt(
+                        if (description.getAnnotation(UserSetupMode::class.java) != null) {
+                            0
+                        } else {
+                            1
+                        }
+                    )
+                    kidsModeRule.putInt(
+                        if (description.getAnnotation(NavBarKidsMode::class.java) != null) 1 else 0
+                    )
+                }
+
                 // Check for existing Taskbar instance from Launcher process.
                 val launcherTaskbarManager: TaskbarManager? =
                     if (!isRunningInRobolectric) {
@@ -105,8 +134,8 @@
                         null
                     }
 
-                instrumentation.runOnMainSync {
-                    taskbarManager =
+                taskbarManager =
+                    TestUtil.getOnUiThread {
                         TaskbarManager(
                             context,
                             AllAppsActionManager(context, UI_HELPER_EXECUTOR) {
@@ -114,12 +143,14 @@
                             },
                             object : TaskbarNavButtonCallbacks {},
                         )
-                }
+                    }
 
                 try {
+                    LauncherAppState.getInstance(context).model.loadModelSync()
+
                     // Replace Launcher Taskbar window with test instance.
                     instrumentation.runOnMainSync {
-                        launcherTaskbarManager?.removeTaskbarRootViewFromWindow()
+                        launcherTaskbarManager?.setSuspended(true)
                         taskbarManager.onUserUnlocked() // Required to complete initialization.
                     }
 
@@ -129,7 +160,7 @@
                     // Revert Taskbar window.
                     instrumentation.runOnMainSync {
                         taskbarManager.destroy()
-                        launcherTaskbarManager?.addTaskbarRootViewToWindow()
+                        launcherTaskbarManager?.setSuspended(false)
                     }
                 }
             }
@@ -167,4 +198,35 @@
     @Retention(AnnotationRetention.RUNTIME)
     @Target(AnnotationTarget.FIELD)
     annotation class InjectController
+
+    /** Overrides [USER_SETUP_COMPLETE] to be `false` for tests. */
+    @Retention(AnnotationRetention.RUNTIME)
+    @Target(AnnotationTarget.CLASS, AnnotationTarget.FUNCTION)
+    annotation class UserSetupMode
+
+    /** Overrides [NAV_BAR_KIDS_MODE] to be `true` for tests. */
+    @Retention(AnnotationRetention.RUNTIME)
+    @Target(AnnotationTarget.CLASS, AnnotationTarget.FUNCTION)
+    annotation class NavBarKidsMode
+
+    /** Rule for Taskbar integer-based secure settings. */
+    private inner class TaskbarSecureSettingRule(private val settingName: String) : TestRule {
+
+        override fun apply(base: Statement, description: Description): Statement {
+            return object : Statement() {
+                override fun evaluate() {
+                    val originalValue =
+                        Settings.Secure.getInt(context.contentResolver, settingName, /* def= */ 0)
+                    try {
+                        base.evaluate()
+                    } finally {
+                        instrumentation.runOnMainSync { putInt(originalValue) }
+                    }
+                }
+            }
+        }
+
+        /** Puts [value] into secure settings under [settingName]. */
+        fun putInt(value: Int) = Settings.Secure.putInt(context.contentResolver, settingName, value)
+    }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
index 8262e0f..234e499 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
@@ -22,6 +22,8 @@
 import com.android.launcher3.taskbar.TaskbarManager
 import com.android.launcher3.taskbar.TaskbarStashController
 import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
+import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.NavBarKidsMode
+import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.UserSetupMode
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
 import com.google.common.truth.Truth.assertThat
@@ -125,9 +127,40 @@
         }
     }
 
-    /** Executes [runTest] after the [testRule] setup phase completes. */
+    @Test
+    fun testUserSetupMode_default_isComplete() {
+        onSetup { assertThat(activityContext.isUserSetupComplete).isTrue() }
+    }
+
+    @Test
+    fun testUserSetupMode_withAnnotation_isIncomplete() {
+        @UserSetupMode class Mode
+        onSetup(description = Description.createSuiteDescription(Mode::class.java)) {
+            assertThat(activityContext.isUserSetupComplete).isFalse()
+        }
+    }
+
+    @Test
+    fun testNavBarKidsMode_default_navBarNotForcedVisible() {
+        onSetup { assertThat(activityContext.isNavBarForceVisible).isFalse() }
+    }
+
+    @Test
+    fun testNavBarKidsMode_withAnnotation_navBarForcedVisible() {
+        @NavBarKidsMode class Mode
+        onSetup(description = Description.createSuiteDescription(Mode::class.java)) {
+            assertThat(activityContext.isNavBarForceVisible).isTrue()
+        }
+    }
+
+    /**
+     * Executes [runTest] after the [testRule] setup phase completes.
+     *
+     * A [description] can also be provided to mimic annotating a test or test class.
+     */
     private fun onSetup(
         testRule: TaskbarUnitTestRule = TaskbarUnitTestRule(this, context),
+        description: Description = DESCRIPTION,
         runTest: TaskbarUnitTestRule.() -> Unit,
     ) {
         testRule
@@ -135,7 +168,7 @@
                 object : Statement() {
                     override fun evaluate() = runTest(testRule)
                 },
-                DESCRIPTION,
+                description,
             )
             .evaluate()
     }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskThumbnailDataSource.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskThumbnailDataSource.kt
index b66b735..30fc491 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskThumbnailDataSource.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskThumbnailDataSource.kt
@@ -32,7 +32,7 @@
     var shouldLoadSynchronously: Boolean = true
 
     /** Retrieves and sets a thumbnail on [task] from [taskIdToBitmap]. */
-    override fun updateThumbnailInBackground(
+    override fun getThumbnailInBackground(
         task: Task,
         callback: Consumer<ThumbnailData>
     ): CancellableTask<ThumbnailData>? {
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
index a394b65..b78f871 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
@@ -30,6 +30,7 @@
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
 import com.android.quickstep.task.viewmodel.TaskContainerData
 import com.android.quickstep.task.viewmodel.TaskViewData
+import com.android.quickstep.views.TaskViewType
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
 import com.google.common.truth.Truth.assertThat
@@ -42,12 +43,14 @@
 
 @RunWith(AndroidJUnit4::class)
 class TaskThumbnailViewModelTest {
+    private var taskViewType = TaskViewType.SINGLE
     private val recentsViewData = RecentsViewData()
-    private val taskViewData = TaskViewData()
+    private val taskViewData by lazy { TaskViewData(taskViewType) }
     private val taskContainerData = TaskContainerData()
     private val tasksRepository = FakeTasksRepository()
-    private val systemUnderTest =
+    private val systemUnderTest by lazy {
         TaskThumbnailViewModel(recentsViewData, taskViewData, taskContainerData, tasksRepository)
+    }
 
     private val tasks = (0..5).map(::createTaskWithId)
 
@@ -66,14 +69,26 @@
     }
 
     @Test
-    fun setRecentsFullscreenProgress_thenProgressIsPassedThrough() = runTest {
+    fun setRecentsFullscreenProgress_thenCornerRadiusProgressIsPassedThrough() = runTest {
         recentsViewData.fullscreenProgress.value = 0.5f
 
-        assertThat(systemUnderTest.recentsFullscreenProgress.first()).isEqualTo(0.5f)
+        assertThat(systemUnderTest.cornerRadiusProgress.first()).isEqualTo(0.5f)
 
         recentsViewData.fullscreenProgress.value = 0.6f
 
-        assertThat(systemUnderTest.recentsFullscreenProgress.first()).isEqualTo(0.6f)
+        assertThat(systemUnderTest.cornerRadiusProgress.first()).isEqualTo(0.6f)
+    }
+
+    @Test
+    fun setRecentsFullscreenProgress_thenCornerRadiusProgressIsConstantForDesktop() = runTest {
+        taskViewType = TaskViewType.DESKTOP
+        recentsViewData.fullscreenProgress.value = 0.5f
+
+        assertThat(systemUnderTest.cornerRadiusProgress.first()).isEqualTo(1f)
+
+        recentsViewData.fullscreenProgress.value = 0.6f
+
+        assertThat(systemUnderTest.cornerRadiusProgress.first()).isEqualTo(1f)
     }
 
     @Test
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/GroupTaskTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/GroupTaskTest.kt
index a6d3887..f11cd0b 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/GroupTaskTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/GroupTaskTest.kt
@@ -21,7 +21,7 @@
 import android.graphics.Rect
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.SplitConfigurationOptions
-import com.android.quickstep.views.TaskView
+import com.android.quickstep.views.TaskViewType
 import com.android.systemui.shared.recents.model.Task
 import com.android.wm.shell.common.split.SplitScreenConstants
 import com.google.common.truth.Truth.assertThat
@@ -68,8 +68,8 @@
                 2,
                 SplitScreenConstants.SNAP_TO_50_50
             )
-        val task1 = GroupTask(createTask(1), createTask(2), splitBounds, TaskView.Type.GROUPED)
-        val task2 = GroupTask(createTask(1), createTask(2), splitBounds, TaskView.Type.GROUPED)
+        val task1 = GroupTask(createTask(1), createTask(2), splitBounds, TaskViewType.GROUPED)
+        val task2 = GroupTask(createTask(1), createTask(2), splitBounds, TaskViewType.GROUPED)
         assertThat(task1).isEqualTo(task2)
     }
 
@@ -91,15 +91,15 @@
                 2,
                 SplitScreenConstants.SNAP_TO_30_70
             )
-        val task1 = GroupTask(createTask(1), createTask(2), splitBounds1, TaskView.Type.GROUPED)
-        val task2 = GroupTask(createTask(1), createTask(2), splitBounds2, TaskView.Type.GROUPED)
+        val task1 = GroupTask(createTask(1), createTask(2), splitBounds1, TaskViewType.GROUPED)
+        val task2 = GroupTask(createTask(1), createTask(2), splitBounds2, TaskViewType.GROUPED)
         assertThat(task1).isNotEqualTo(task2)
     }
 
     @Test
     fun testGroupTask_differentType_isNotEqual() {
-        val task1 = GroupTask(createTask(1), null, null, TaskView.Type.SINGLE)
-        val task2 = GroupTask(createTask(1), null, null, TaskView.Type.DESKTOP)
+        val task1 = GroupTask(createTask(1), null, null, TaskViewType.SINGLE)
+        val task2 = GroupTask(createTask(1), null, null, TaskViewType.DESKTOP)
         assertThat(task1).isNotEqualTo(task2)
     }
 
diff --git a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
index f160ce2..d9d5585 100644
--- a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
@@ -40,7 +40,7 @@
 import com.android.systemui.shared.recents.model.Task.TaskKey
 import com.android.window.flags.Flags
 import com.android.wm.shell.common.desktopmode.DesktopModeTransitionSource
-import com.android.wm.shell.shared.DesktopModeStatus
+import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
@@ -191,7 +191,6 @@
         return TaskContainer(
             taskView,
             task,
-            thumbnailView = null,
             thumbnailViewDeprecated,
             iconView,
             transformingTouchDelegate,
diff --git a/quickstep/tests/src/com/android/quickstep/RecentTasksListTest.java b/quickstep/tests/src/com/android/quickstep/RecentTasksListTest.java
index ce16b70..5d00255 100644
--- a/quickstep/tests/src/com/android/quickstep/RecentTasksListTest.java
+++ b/quickstep/tests/src/com/android/quickstep/RecentTasksListTest.java
@@ -32,7 +32,7 @@
 
 import com.android.launcher3.util.LooperExecutor;
 import com.android.quickstep.util.GroupTask;
-import com.android.quickstep.views.TaskView;
+import com.android.quickstep.views.TaskViewType;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.wm.shell.util.GroupedRecentTaskInfo;
 
@@ -125,7 +125,7 @@
                 Integer.MAX_VALUE /* numTasks */, -1 /* requestId */, false /* loadKeysOnly */);
 
         assertEquals(1, taskList.size());
-        assertEquals(TaskView.Type.DESKTOP, taskList.get(0).taskViewType);
+        assertEquals(TaskViewType.DESKTOP, taskList.get(0).taskViewType);
         List<Task> actualFreeformTasks = taskList.get(0).getTasks();
         assertEquals(3, actualFreeformTasks.size());
         assertEquals(1, actualFreeformTasks.get(0).key.id);