diff --git a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
index 5744464..fd0243a 100644
--- a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
+++ b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
@@ -488,6 +488,15 @@
                 }
             });
         }
+
+        public void onEnterDesktopModeTransitionStarted(int transitionDuration) {
+
+        }
+
+        @Override
+        public void onExitDesktopModeTransitionStarted(int transitionDuration) {
+
+        }
     }
 
     /** A listener for Taskbar in Desktop Mode. */
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
index 2998892..dce377d 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
@@ -1009,7 +1009,12 @@
 
         @Override
         public void onRecentsAnimationFinished(RecentsAnimationController controller) {
-            endGestureStateOverride(!controller.getFinishTargetIsLauncher(), false /*canceled*/);
+            endGestureStateOverride(!controller.getFinishTargetIsLauncher(),
+                    controller.getLauncherIsVisibleAtFinish(), false /*canceled*/);
+        }
+
+        private void endGestureStateOverride(boolean finishedToApp, boolean canceled) {
+            endGestureStateOverride(finishedToApp, finishedToApp, canceled);
         }
 
         /**
@@ -1019,11 +1024,13 @@
          *
          * @param finishedToApp {@code true} if the recents animation finished to showing an app and
          *                      not workspace or overview
+         * @param launcherIsVisible {code true} if launcher is visible at finish
          * @param canceled      {@code true} if the recents animation was canceled instead of
          *                      finishing
          *                      to completion
          */
-        private void endGestureStateOverride(boolean finishedToApp, boolean canceled) {
+        private void endGestureStateOverride(boolean finishedToApp, boolean launcherIsVisible,
+                boolean canceled) {
             mCallbacks.removeListener(this);
             mTaskBarRecentsAnimationListener = null;
             ((RecentsView) mLauncher.getOverviewPanel()).setTaskLaunchListener(null);
@@ -1032,18 +1039,27 @@
                 mSkipNextRecentsAnimEnd = false;
                 return;
             }
-            updateStateForUserFinishedToApp(finishedToApp);
+            updateStateForUserFinishedToApp(finishedToApp, launcherIsVisible);
         }
     }
 
     /**
+     * @see #updateStateForUserFinishedToApp(boolean, boolean)
+     */
+    private void updateStateForUserFinishedToApp(boolean finishedToApp) {
+        updateStateForUserFinishedToApp(finishedToApp, !finishedToApp);
+    }
+
+    /**
      * Updates the visible state immediately to ensure a seamless handoff.
      *
      * @param finishedToApp True iff user is in an app.
+     * @param launcherIsVisible True iff launcher is still visible (ie. transparent app)
      */
-    private void updateStateForUserFinishedToApp(boolean finishedToApp) {
+    private void updateStateForUserFinishedToApp(boolean finishedToApp,
+            boolean launcherIsVisible) {
         // Update the visible state immediately to ensure a seamless handoff
-        boolean launcherVisible = !finishedToApp;
+        boolean launcherVisible = !finishedToApp || launcherIsVisible;
         updateStateForFlag(FLAG_TRANSITION_TO_VISIBLE, false);
         updateStateForFlag(FLAG_VISIBLE, launcherVisible);
         applyState();
@@ -1052,7 +1068,7 @@
         if (DEBUG) {
             Log.d(TAG, "endGestureStateOverride - FLAG_IN_APP: " + finishedToApp);
         }
-        controller.updateStateForFlag(FLAG_IN_APP, finishedToApp);
+        controller.updateStateForFlag(FLAG_IN_APP, finishedToApp && !launcherIsVisible);
         controller.applyState();
     }
 
diff --git a/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java b/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
index 1124aac..6719ab7 100644
--- a/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
+++ b/quickstep/src/com/android/quickstep/LauncherBackAnimationController.java
@@ -24,6 +24,7 @@
 import static com.android.launcher3.BaseActivity.INVISIBLE_ALL;
 import static com.android.launcher3.BaseActivity.INVISIBLE_BY_PENDING_FLAGS;
 import static com.android.launcher3.BaseActivity.PENDING_INVISIBLE_BY_WALLPAPER_ANIMATION;
+import static com.android.window.flags.Flags.predictiveBackThreeButtonNav;
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
@@ -60,6 +61,8 @@
 import com.android.launcher3.Utilities;
 import com.android.launcher3.taskbar.LauncherTaskbarUIController;
 import com.android.launcher3.uioverrides.QuickstepLauncher;
+import com.android.launcher3.util.DisplayController;
+import com.android.launcher3.util.NavigationMode;
 import com.android.launcher3.widget.LauncherAppWidgetHostView;
 import com.android.quickstep.util.BackAnimState;
 import com.android.systemui.shared.system.QuickStepContract;
@@ -295,8 +298,11 @@
 
         mStartRect.set(appTarget.windowConfiguration.getMaxBounds());
 
-        // inset bottom in case of pinned taskbar being present
-        mStartRect.inset(0, 0, 0, appTarget.contentInsets.bottom);
+        // inset bottom in case of taskbar being present
+        if (!predictiveBackThreeButtonNav() || mLauncher.getDeviceProfile().isTaskbarPresent
+                || DisplayController.getNavigationMode(mLauncher) == NavigationMode.NO_BUTTON) {
+            mStartRect.inset(0, 0, 0, appTarget.contentInsets.bottom);
+        }
 
         mLauncherTargetView = mQuickstepTransitionManager.findLauncherView(
                 new RemoteAnimationTarget[]{ mBackTarget });
diff --git a/quickstep/src/com/android/quickstep/RecentsAnimationController.java b/quickstep/src/com/android/quickstep/RecentsAnimationController.java
index dcb0108..60fcff8 100644
--- a/quickstep/src/com/android/quickstep/RecentsAnimationController.java
+++ b/quickstep/src/com/android/quickstep/RecentsAnimationController.java
@@ -53,6 +53,8 @@
     private boolean mFinishRequested = false;
     // Only valid when mFinishRequested == true.
     private boolean mFinishTargetIsLauncher;
+    // Only valid when mFinishRequested == true
+    private boolean mLauncherIsVisibleAtFinish;
     private RunnableList mPendingFinishCallbacks = new RunnableList();
 
     public RecentsAnimationController(RecentsAnimationControllerCompat controller,
@@ -117,13 +119,27 @@
     }
 
     @UiThread
+    public void finish(boolean toRecents, boolean launcherIsVisibleAtFinish,
+            Runnable onFinishComplete, boolean sendUserLeaveHint) {
+        Preconditions.assertUIThread();
+        finishController(toRecents, launcherIsVisibleAtFinish, onFinishComplete, sendUserLeaveHint,
+                false);
+    }
+
+    @UiThread
     public void finishController(boolean toRecents, Runnable callback, boolean sendUserLeaveHint) {
-        finishController(toRecents, callback, sendUserLeaveHint, false /* forceFinish */);
+        finishController(toRecents, false, callback, sendUserLeaveHint, false /* forceFinish */);
     }
 
     @UiThread
     public void finishController(boolean toRecents, Runnable callback, boolean sendUserLeaveHint,
             boolean forceFinish) {
+        finishController(toRecents, toRecents, callback, sendUserLeaveHint, forceFinish);
+    }
+
+    @UiThread
+    public void finishController(boolean toRecents, boolean launcherIsVisibleAtFinish,
+            Runnable callback, boolean sendUserLeaveHint, boolean forceFinish) {
         mPendingFinishCallbacks.add(callback);
         if (!forceFinish && mFinishRequested) {
             // If finish has already been requested, then add the callback to the pending list.
@@ -135,6 +151,7 @@
         // Finish not yet requested
         mFinishRequested = true;
         mFinishTargetIsLauncher = toRecents;
+        mLauncherIsVisibleAtFinish = launcherIsVisibleAtFinish;
         mOnFinishedListener.accept(this);
         Runnable finishCb = () -> {
             mController.finish(toRecents, sendUserLeaveHint, new IResultReceiver.Stub() {
@@ -211,6 +228,14 @@
         return mFinishTargetIsLauncher;
     }
 
+    /**
+     * RecentsAnimationListeners can check this in onRecentsAnimationFinished() to determine whether
+     * the animation was finished to launcher vs an app.
+     */
+    public boolean getLauncherIsVisibleAtFinish() {
+        return mLauncherIsVisibleAtFinish;
+    }
+
     public void dump(String prefix, PrintWriter pw) {
         pw.println(prefix + "RecentsAnimationController:");
 
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 8982850..6ab3e28 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -610,6 +610,8 @@
     private int mKeyboardTaskFocusSnapAnimationDuration;
     private int mKeyboardTaskFocusIndex = INVALID_PAGE;
 
+    private int[] mDismissPrimaryTranslations;
+
     /**
      * TODO: Call reloadIdNeeded in onTaskStackChanged.
      */
@@ -1418,7 +1420,7 @@
         if (showAsGrid()) {
             int screenStart = getPagedOrientationHandler().getPrimaryScroll(this);
             int screenEnd = screenStart + getPagedOrientationHandler().getMeasuredSize(this);
-            return isTaskViewWithinBounds(tv, screenStart, screenEnd);
+            return isTaskViewWithinBounds(tv, screenStart, screenEnd, /*taskViewTranslation=*/ 0);
         } else {
             // For now, just check if it's the active task or an adjacent task
             return Math.abs(indexOfChild(tv) - getNextPage()) <= 1;
@@ -1465,14 +1467,28 @@
         return clearAllScroll + (mIsRtl ? distance : -distance);
     }
 
-    private boolean isTaskViewWithinBounds(TaskView tv, int start, int end) {
-        int taskStart = getPagedOrientationHandler().getChildStart(tv)
-                + (int) tv.getOffsetAdjustment(showAsGrid());
-        int taskSize = (int) (getPagedOrientationHandler().getMeasuredSize(tv)
-                * tv.getSizeAdjustment(showAsFullscreen()));
+    /*
+     * Returns if TaskView is within screen bounds defined in [screenStart, screenEnd].
+     *
+     * @param taskViewTranslation taskView is considered within bounds if either translated or
+     * original position of taskView is within screen bounds.
+     */
+    private boolean isTaskViewWithinBounds(TaskView taskView, int screenStart, int screenEnd,
+            int taskViewTranslation) {
+        int taskStart = getPagedOrientationHandler().getChildStart(taskView)
+                + (int) taskView.getOffsetAdjustment(showAsGrid());
+        int taskSize = (int) (getPagedOrientationHandler().getMeasuredSize(taskView)
+                * taskView.getSizeAdjustment(showAsFullscreen()));
         int taskEnd = taskStart + taskSize;
-        return (taskStart >= start && taskStart <= end) || (taskEnd >= start
-                && taskEnd <= end);
+
+        int translatedTaskStart = taskStart + taskViewTranslation;
+        int translatedTaskEnd = taskEnd + taskViewTranslation;
+
+        taskStart = Math.min(taskStart, translatedTaskStart);
+        taskEnd = Math.max(taskEnd, translatedTaskEnd);
+
+        return (taskStart >= screenStart && taskStart <= screenEnd) || (taskEnd >= screenStart
+                && taskEnd <= screenEnd);
     }
 
     private boolean isTaskViewFullyWithinBounds(TaskView tv, int start, int end) {
@@ -2470,7 +2486,8 @@
             }
             boolean visible;
             if (showAsGrid()) {
-                visible = isTaskViewWithinBounds(taskView, visibleStart, visibleEnd);
+                visible = isTaskViewWithinBounds(taskView, visibleStart, visibleEnd,
+                        mDismissPrimaryTranslations != null ? mDismissPrimaryTranslations[i] : 0);
             } else {
                 visible = lower <= i && i <= upper;
             }
@@ -3837,6 +3854,7 @@
             stagingTranslation += mIsRtl ? newClearAllShortTotalWidthTranslation
                     : -newClearAllShortTotalWidthTranslation;
         }
+        mDismissPrimaryTranslations = new int[taskCount];
         for (int i = 0; i < count; i++) {
             View child = getChildAt(i);
             if (child == dismissedTaskView) {
@@ -3854,7 +3872,7 @@
                             Math.abs(i - dismissedIndex),
                             scrollDiff,
                             anim,
-                            splitTimings);
+                            splitTimings, i);
                     needsCurveUpdates = true;
                 }
             } else if (child instanceof TaskView taskView) {
@@ -3915,10 +3933,12 @@
                 primaryTranslation += mIsRtl ? stagingTranslation : -stagingTranslation;
 
                 if (primaryTranslation != 0) {
+                    float finalTranslation = mIsRtl ? primaryTranslation : -primaryTranslation;
                     anim.setFloat(taskView, taskView.getPrimaryDismissTranslationProperty(),
-                            mIsRtl ? primaryTranslation : -primaryTranslation,
+                            finalTranslation,
                             clampToProgress(dismissInterpolator, animationStartProgress,
                                     animationEndProgress));
+                    mDismissPrimaryTranslations[i] = (int) finalTranslation;
                     distanceFromDismissedTask++;
                 }
             }
@@ -3937,7 +3957,7 @@
         if (animateTaskView && dismissedTaskView != null) {
             dismissedTaskView.setTranslationZ(0.1f);
         }
-
+        loadVisibleTaskData(TaskView.FLAG_UPDATE_ALL);
         mPendingAnimation = anim;
         final TaskView finalNextFocusedTaskView = nextFocusedTaskView;
         final boolean finalCloseGapBetweenClearAll = closeGapBetweenClearAll;
@@ -4155,6 +4175,7 @@
                 updateCurrentTaskActionsVisibility();
                 onDismissAnimationEnds();
                 mPendingAnimation = null;
+                mDismissPrimaryTranslations = null;
             }
         });
     }
@@ -4193,7 +4214,8 @@
             int indexDiff,
             int scrollDiffPerPage,
             PendingAnimation pendingAnimation,
-            SplitAnimationTimings splitTimings) {
+            SplitAnimationTimings splitTimings,
+            int index) {
         FloatProperty translationProperty = view instanceof TaskView
                 ? ((TaskView) view).getPrimaryDismissTranslationProperty()
                 : getPagedOrientationHandler().getPrimaryViewTranslate();
@@ -4227,6 +4249,9 @@
                 )
         );
 
+        if (view instanceof TaskView) {
+            mDismissPrimaryTranslations[index] = scrollDiffPerPage;
+        }
         if (mEnableDrawingLiveTile && view instanceof TaskView
                 && ((TaskView) view).isRunningTask()) {
             pendingAnimation.addOnFrameCallback(() -> {
@@ -5854,15 +5879,22 @@
      * Finish recents animation.
      */
     public void finishRecentsAnimation(boolean toRecents, @Nullable Runnable onFinishComplete) {
-        finishRecentsAnimation(toRecents, true /* shouldPip */, onFinishComplete);
+        finishRecentsAnimation(toRecents, false, true /* shouldPip */, onFinishComplete);
     }
 
     /**
+     * Finish recents animation.
+     */
+    public void finishRecentsAnimation(boolean toRecents, boolean shouldPip,
+            @Nullable Runnable onFinishComplete) {
+        finishRecentsAnimation(toRecents, shouldPip, false, onFinishComplete);
+    }
+    /**
      * NOTE: Whatever value gets passed through to the toRecents param may need to also be set on
      * {@link #mRecentsAnimationController#setWillFinishToHome}.
      */
     public void finishRecentsAnimation(boolean toRecents, boolean shouldPip,
-            @Nullable Runnable onFinishComplete) {
+            boolean allAppTargetsAreTranslucent, @Nullable Runnable onFinishComplete) {
         Log.d(TAG, "finishRecentsAnimation - mRecentsAnimationController: "
                 + mRecentsAnimationController);
         // TODO(b/197232424#comment#10) Move this back into onRecentsAnimationComplete(). Maybe?
@@ -5894,7 +5926,7 @@
                         tx, null /* overlay */);
             }
         }
-        mRecentsAnimationController.finish(toRecents, () -> {
+        mRecentsAnimationController.finish(toRecents, allAppTargetsAreTranslucent, () -> {
             if (onFinishComplete != null) {
                 onFinishComplete.run();
             }
diff --git a/res/layout/widgets_full_sheet_paged_view.xml b/res/layout/widgets_full_sheet_paged_view.xml
index 622f0d6..7c57726 100644
--- a/res/layout/widgets_full_sheet_paged_view.xml
+++ b/res/layout/widgets_full_sheet_paged_view.xml
@@ -81,6 +81,7 @@
             android:layout_marginTop="8dp"
             android:layout_marginBottom="8dp"
             android:background="@drawable/widgets_surface_background"
+            android:clipToOutline="true"
             android:orientation="vertical"
             android:layout_marginHorizontal="@dimen/widget_list_horizontal_margin"
             android:visibility="gone">
diff --git a/res/layout/widgets_full_sheet_recyclerview.xml b/res/layout/widgets_full_sheet_recyclerview.xml
index 5427732..1ce1c55 100644
--- a/res/layout/widgets_full_sheet_recyclerview.xml
+++ b/res/layout/widgets_full_sheet_recyclerview.xml
@@ -64,6 +64,7 @@
             android:layout_marginTop="8dp"
             android:layout_marginBottom="8dp"
             android:background="@drawable/widgets_surface_background"
+            android:clipToOutline="true"
             android:orientation="vertical"
             android:visibility="gone">
             <include layout="@layout/widget_recommendations" />
diff --git a/res/layout/widgets_two_pane_sheet.xml b/res/layout/widgets_two_pane_sheet.xml
index 5dc1b47..cf090ad 100644
--- a/res/layout/widgets_two_pane_sheet.xml
+++ b/res/layout/widgets_two_pane_sheet.xml
@@ -133,6 +133,7 @@
                                 android:layout_width="match_parent"
                                 android:layout_height="match_parent"
                                 android:background="@drawable/widgets_surface_background"
+                                android:clipToOutline="true"
                                 android:orientation="vertical"
                                 android:visibility="gone">
                                 <include layout="@layout/widget_recommendations" />
diff --git a/src/com/android/launcher3/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index ef5c88a..817cc40 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -127,6 +127,8 @@
 
     private static final int[] STATE_PRESSED = new int[]{android.R.attr.state_pressed};
 
+    private static final int APP_PILL_TITLE_PADDING = 8;
+
     private float mScaleForReorderBounce = 1f;
 
     private IntArray mBreakPointsIntArray;
@@ -730,16 +732,21 @@
         Paint.FontMetrics fm = getPaint().getFontMetrics();
         Rect tmpRect = new Rect();
         getDrawingRect(tmpRect);
+        CharSequence text = getText();
 
-        if (mIcon == null) {
-            appTitleBounds = new RectF(0, 0, tmpRect.right,
-                    (int) Math.ceil(fm.bottom - fm.top));
-        } else {
+        float titleLength = (getPaint().measureText(text, 0, text.length())
+                + APP_PILL_TITLE_PADDING * 2);
+        titleLength = Math.min(titleLength, tmpRect.width());
+        appTitleBounds = new RectF((tmpRect.width() - titleLength) / 2.f - getCompoundPaddingLeft(),
+                0, (tmpRect.width() + titleLength) / 2.f + getCompoundPaddingRight(),
+                (int) Math.ceil(fm.bottom - fm.top));
+
+
+        if (mIcon != null) {
             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));
+            appTitleBounds.offset(0, textStart);
         }
 
         canvas.drawRoundRect(appTitleBounds, appTitleBounds.height() / 2,
@@ -851,6 +858,11 @@
             setPadding(getPaddingLeft(), (height - cellHeightPx) / 2, getPaddingRight(),
                     getPaddingBottom());
         }
+        if (shouldDrawAppContrastTile()) {
+            setPadding(getPaddingLeft() + APP_PILL_TITLE_PADDING, getPaddingTop(),
+                    getPaddingRight() + APP_PILL_TITLE_PADDING,
+                    getPaddingBottom());
+        }
         // Only apply two line for all_apps and device search only if necessary.
         if (shouldUseTwoLine() && (mLastOriginalText != null)) {
             int allowedVerticalSpace = height - getPaddingTop() - getPaddingBottom()
diff --git a/src/com/android/launcher3/DeleteDropTarget.java b/src/com/android/launcher3/DeleteDropTarget.java
index 425f277..58789fd 100644
--- a/src/com/android/launcher3/DeleteDropTarget.java
+++ b/src/com/android/launcher3/DeleteDropTarget.java
@@ -130,6 +130,7 @@
     public void completeDrop(DragObject d) {
         ItemInfo item = d.dragInfo;
         if (canRemove(item)) {
+            onAccessibilityDrop(null, item);
             mDropTargetHandler.onDeleteComplete(item);
         }
     }
diff --git a/src/com/android/launcher3/SecondaryDropTarget.java b/src/com/android/launcher3/SecondaryDropTarget.java
index b3cb948..f4d3146 100644
--- a/src/com/android/launcher3/SecondaryDropTarget.java
+++ b/src/com/android/launcher3/SecondaryDropTarget.java
@@ -303,10 +303,11 @@
                     .setData(Uri.fromParts("package", cn.getPackageName(), cn.getClassName()))
                     .putExtra(Intent.EXTRA_USER, info.user);
             context.startActivity(i);
-            FileLog.d(TAG, "start uninstall activity " + cn.getPackageName());
+            FileLog.d(TAG, "start uninstall activity from drop target " + cn.getPackageName());
             return cn;
         } catch (URISyntaxException e) {
-            Log.e(TAG, "Failed to parse intent to start uninstall activity for item=" + info);
+            Log.e(TAG, "Failed to parse intent to start drop target uninstall activity for"
+                    + " item=" + info);
             return null;
         }
     }
diff --git a/src/com/android/launcher3/SessionCommitReceiver.java b/src/com/android/launcher3/SessionCommitReceiver.java
index 6168e41..ea5eb8f 100644
--- a/src/com/android/launcher3/SessionCommitReceiver.java
+++ b/src/com/android/launcher3/SessionCommitReceiver.java
@@ -73,8 +73,9 @@
                 || alreadyAddedPromiseIcon) {
             FileLog.d(LOG,
                     String.format(Locale.ENGLISH,
-                            "Removing PromiseIcon for package: %s, install reason: %d,"
-                            + " alreadyAddedPromiseIcon: %s",
+                            "Removing unneeded PromiseIcon for package: %s"
+                                    + ", install reason: %d,"
+                                    + " alreadyAddedPromiseIcon: %s",
                     info.getAppPackageName(),
                     info.getInstallReason(),
                     alreadyAddedPromiseIcon
diff --git a/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java b/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
index 1b58987..c938482 100644
--- a/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
@@ -28,6 +28,7 @@
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.ScrollableLayoutManager.PREDICTIVE_BACK_MIN_SCALE;
 import static com.android.launcher3.views.RecyclerViewFastScroller.FastScrollerLocation.ALL_APPS_SCROLLER;
+import static com.android.window.flags.Flags.predictiveBackThreeButtonNav;
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
@@ -1173,8 +1174,10 @@
         super.dispatchDraw(canvas);
 
         if (mNavBarScrimHeight > 0) {
-            canvas.drawRect(0, getHeight() - mNavBarScrimHeight, getWidth(), getHeight(),
-                    mNavBarScrimPaint);
+            float left = (getWidth() - getWidth() / getScaleX()) / 2;
+            float top = getHeight() / 2f + (getHeight() / 2f - mNavBarScrimHeight) / getScaleY();
+            canvas.drawRect(left, top, getWidth() / getScaleX(),
+                    top + mNavBarScrimHeight / getScaleY(), mNavBarScrimPaint);
         }
     }
 
@@ -1340,6 +1343,17 @@
         invalidateHeader();
     }
 
+    @Override
+    public void setScaleY(float scaleY) {
+        super.setScaleY(scaleY);
+        if (predictiveBackThreeButtonNav() && mNavBarScrimHeight > 0) {
+            // Call invalidate to prevent navbar scrim from scaling. The navbar scrim is drawn
+            // directly onto the canvas. To prevent it from being scaled with the canvas, there's a
+            // counter scale applied in dispatchDraw.
+            invalidate(20, getHeight() - mNavBarScrimHeight, getWidth(), getHeight());
+        }
+    }
+
     /**
      * Set {@link Animator.AnimatorListener} on {@link mAllAppsTransitionController} to observe
      * animation of backing out of all apps search view to all apps view.
diff --git a/src/com/android/launcher3/allapps/PrivateProfileManager.java b/src/com/android/launcher3/allapps/PrivateProfileManager.java
index 21dce14..609edd2 100644
--- a/src/com/android/launcher3/allapps/PrivateProfileManager.java
+++ b/src/com/android/launcher3/allapps/PrivateProfileManager.java
@@ -45,6 +45,7 @@
 import android.animation.ValueAnimator;
 import android.content.Context;
 import android.content.Intent;
+import android.os.Trace;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.util.Log;
@@ -220,6 +221,7 @@
      * when animation is not running.
      */
     public void reset() {
+        Trace.beginSection("PrivateProfileManager#reset");
         // Ensure the state of the header view is what it should be before animating.
         updateView();
         getMainRecyclerView().setChildAttachedConsumer(null);
@@ -239,6 +241,7 @@
             executeLock();
         }
         addPrivateSpaceDecorator(updatedState);
+        Trace.endSection();
     }
 
     /** Returns whether or not Private Space Settings Page is available. */
@@ -293,31 +296,12 @@
         }
     }
 
-    @Override
     public void setQuietMode(boolean enable) {
-        UI_HELPER_EXECUTOR.post(() ->
-                mUserCache.getUserProfiles()
-                        .stream()
-                        .filter(getUserMatcher())
-                        .findFirst()
-                        .ifPresent(userHandle -> setQuietModeSafely(enable, userHandle)));
+        setQuietMode(enable, mAllApps.mActivityContext);
         mReadyToAnimate = true;
     }
 
     /**
-     * Sets Quiet Mode for Private Profile.
-     * If {@link SecurityException} is thrown, prompts the user to set this launcher as HOME app.
-     */
-    private void setQuietModeSafely(boolean enable, UserHandle userHandle) {
-        try {
-            mUserManager.requestQuietModeEnabled(enable, userHandle);
-        } catch (SecurityException ex) {
-            ApiWrapper.INSTANCE.get(mAllApps.mActivityContext)
-                    .assignDefaultHomeRole(mAllApps.mActivityContext);
-        }
-    }
-
-    /**
      * Expand the private space after the app list has been added and updated from
      * {@link AlphabeticalAppsList#onAppsUpdated()}
      */
@@ -331,7 +315,9 @@
 
     /** Collapses the private space before the app list has been updated. */
     void executeLock() {
+        Trace.beginSection("PrivateProfileManager#executeLock");
         MAIN_EXECUTOR.execute(() -> updatePrivateStateAnimator(false));
+        Trace.endSection();
     }
 
     void setAnimationRunning(boolean isAnimationRunning) {
@@ -378,6 +364,7 @@
         if (mPSHeader == null) {
             return;
         }
+        Trace.beginSection("PrivateProfileManager#updateView");
         Log.d(TAG, "bindPrivateSpaceHeaderViewElements: " + "Updating view with state: "
                 + getCurrentState());
         mPSHeader.setAlpha(1);
@@ -436,6 +423,7 @@
             }
         }
         mPSHeader.invalidate();
+        Trace.endSection();
     }
 
     /** Sets the enablement of the profile when header or button is clicked. */
@@ -840,6 +828,7 @@
         ActivityAllAppsContainerView<?>.AdapterHolder mainAdapterHolder = mAllApps.mAH.get(MAIN);
         List<BaseAllAppsAdapter.AdapterItem> adapterItems =
                 mainAdapterHolder.mAppsList.getAdapterItems();
+        Trace.beginSection("PrivateProfileManager#expandPrivateSpace");
         if (Flags.enablePrivateSpace() && Flags.privateSpaceAnimation()
                 && mAllApps.isPersonalTab()) {
             // Animate the text and settings icon.
@@ -849,6 +838,7 @@
                     getPsHeaderHeight(), deviceProfile.allAppsCellHeightPx);
             updatePrivateStateAnimator(true);
         }
+        Trace.endSection();
     }
 
     private void exitSearchAndExpand() {
diff --git a/src/com/android/launcher3/allapps/UserProfileManager.java b/src/com/android/launcher3/allapps/UserProfileManager.java
index 93b6b29..765c29c 100644
--- a/src/com/android/launcher3/allapps/UserProfileManager.java
+++ b/src/com/android/launcher3/allapps/UserProfileManager.java
@@ -18,6 +18,7 @@
 
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 
+import android.content.Context;
 import android.os.UserHandle;
 import android.os.UserManager;
 
@@ -26,6 +27,7 @@
 import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.pm.UserCache;
+import com.android.launcher3.util.ApiWrapper;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -69,14 +71,26 @@
     }
 
     /** Sets quiet mode as enabled/disabled for the profile type. */
-    protected void setQuietMode(boolean enabled) {
+    protected void setQuietMode(boolean enabled, Context context) {
         UI_HELPER_EXECUTOR.post(() ->
                 mUserCache.getUserProfiles()
                         .stream()
                         .filter(getUserMatcher())
                         .findFirst()
                         .ifPresent(userHandle ->
-                                mUserManager.requestQuietModeEnabled(enabled, userHandle)));
+                                setQuietModeSafely(enabled, userHandle, context)));
+    }
+
+    /**
+     * Sets Quiet Mode for Private Profile.
+     * If {@link SecurityException} is thrown, prompts the user to set this launcher as HOME app.
+     */
+    private void setQuietModeSafely(boolean enable, UserHandle userHandle, Context context) {
+        try {
+            mUserManager.requestQuietModeEnabled(enable, userHandle);
+        } catch (SecurityException ex) {
+            ApiWrapper.INSTANCE.get(context).assignDefaultHomeRole(context);
+        }
     }
 
     /** Sets current state for the profile type. */
diff --git a/src/com/android/launcher3/allapps/WorkProfileManager.java b/src/com/android/launcher3/allapps/WorkProfileManager.java
index 3d0c1d0..6ebab5a 100644
--- a/src/com/android/launcher3/allapps/WorkProfileManager.java
+++ b/src/com/android/launcher3/allapps/WorkProfileManager.java
@@ -74,7 +74,7 @@
      */
     public void setWorkProfileEnabled(boolean enabled) {
         updateCurrentState(STATE_TRANSITION);
-        setQuietMode(!enabled);
+        setQuietMode(!enabled, mAllApps.mActivityContext);
     }
 
     @Override
diff --git a/src/com/android/launcher3/model/GridSizeMigrationDBController.java b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
index 617cac7..bfa00bd 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationDBController.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationDBController.java
@@ -17,12 +17,14 @@
 package com.android.launcher3.model;
 
 import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
+import static com.android.launcher3.Flags.oneGridSpecs;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.LauncherSettings.Favorites.TMP_TABLE;
 import static com.android.launcher3.Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET;
 import static com.android.launcher3.model.LoaderTask.SMARTSPACE_ON_HOME_SCREEN;
 import static com.android.launcher3.provider.LauncherDbUtils.copyTable;
 import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
+import static com.android.launcher3.provider.LauncherDbUtils.shiftTableByXCells;
 
 import android.content.ComponentName;
 import android.content.ContentValues;
@@ -130,6 +132,20 @@
             // Only use this strategy when comparing the previous grid to the new grid and the
             // columns are the same and the destination has more rows
             copyTable(source, TABLE_NAME, target.getWritableDatabase(), TABLE_NAME, context);
+
+            if (oneGridSpecs()) {
+                DbReader destReader = new DbReader(
+                        target.getWritableDatabase(), TABLE_NAME, context);
+                boolean shouldShiftCells = shouldShiftCells(destReader, srcDeviceState.getRows());
+                if (shouldShiftCells) {
+                    shiftTableByXCells(
+                            target.getWritableDatabase(),
+                            (destDeviceState.getRows() - srcDeviceState.getRows()),
+                            TABLE_NAME);
+                }
+            }
+
+            // Save current configuration, so that the migration does not run again.
             destDeviceState.writeToPrefs(context);
             return true;
         }
@@ -427,17 +443,22 @@
         }
     }
 
-    static void copyCurrentGridToNewGrid(
-            @NonNull Context context,
-            @NonNull DeviceGridState destDeviceState,
-            @NonNull DatabaseHelper target,
-            @NonNull SQLiteDatabase source) {
-        // Only use this strategy when comparing the previous grid to the new grid and the
-        // columns are the same and the destination has more rows
-        copyTable(source, TABLE_NAME, target.getWritableDatabase(), TABLE_NAME, context);
-        destDeviceState.writeToPrefs(context);
+    private static boolean shouldShiftCells(final DbReader destReader, final int srcGridRowCount) {
+        List<DbEntry> workspaceItems = destReader.loadAllWorkspaceEntries();
+        int firstPageItemsRowPosSum = workspaceItems.stream()
+                .filter(entry -> entry.screenId == 0)
+                .mapToInt(entry -> entry.cellY).sum();
+        int firstPageWorkspaceItemsCount = (int) workspaceItems.stream()
+                .filter(entry -> entry.screenId == 0).count();
+        if (firstPageWorkspaceItemsCount == 0) {
+            return false;
+        }
+        float srcGridMidPoint = srcGridRowCount / 2f;
+        float firstPageItemPosAvg = (float) firstPageItemsRowPosSum / firstPageWorkspaceItemsCount;
+        return (firstPageItemPosAvg >= srcGridMidPoint);
     }
 
+
     @VisibleForTesting(otherwise = VisibleForTesting.PACKAGE_PRIVATE)
     public static class DbReader {
 
diff --git a/src/com/android/launcher3/model/GridSizeMigrationLogic.kt b/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
index c856d4b..3f52d8a 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
+++ b/src/com/android/launcher3/model/GridSizeMigrationLogic.kt
@@ -21,15 +21,20 @@
 import android.util.Log
 import androidx.annotation.VisibleForTesting
 import com.android.launcher3.Flags
+import com.android.launcher3.Flags.oneGridSpecs
 import com.android.launcher3.LauncherPrefs
 import com.android.launcher3.LauncherPrefs.Companion.get
 import com.android.launcher3.LauncherPrefs.Companion.getPrefs
 import com.android.launcher3.LauncherSettings
+import com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME
+import com.android.launcher3.LauncherSettings.Favorites.TMP_TABLE
 import com.android.launcher3.Utilities
 import com.android.launcher3.config.FeatureFlags
 import com.android.launcher3.model.GridSizeMigrationDBController.DbReader
-import com.android.launcher3.provider.LauncherDbUtils
 import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction
+import com.android.launcher3.provider.LauncherDbUtils.copyTable
+import com.android.launcher3.provider.LauncherDbUtils.dropTable
+import com.android.launcher3.provider.LauncherDbUtils.shiftTableByXCells
 import com.android.launcher3.util.CellAndSpan
 import com.android.launcher3.util.GridOccupancy
 import com.android.launcher3.util.IntArray
@@ -59,27 +64,30 @@
         // amount of rows we simply copy over the source grid to the destination grid, rather
         // than undergoing the general grid migration.
         if (shouldMigrateToStrictlyTallerGrid(isDestNewDb, srcDeviceState, destDeviceState)) {
-            GridSizeMigrationDBController.copyCurrentGridToNewGrid(
-                context,
-                destDeviceState,
-                target,
-                source,
-            )
+            copyTable(source, TABLE_NAME, target.writableDatabase, TABLE_NAME, context)
+            if (oneGridSpecs()) {
+                val destReader = DbReader(target.writableDatabase, TABLE_NAME, context)
+                val shouldShiftCells = shouldShiftCells(destReader, srcDeviceState.rows)
+                if (shouldShiftCells) {
+                    shiftTableByXCells(
+                        target.writableDatabase,
+                        (destDeviceState.rows - srcDeviceState.rows),
+                        TABLE_NAME,
+                    )
+                }
+            }
+            // Save current configuration, so that the migration does not run again.
+            destDeviceState.writeToPrefs(context)
             return
         }
-        LauncherDbUtils.copyTable(
-            source,
-            LauncherSettings.Favorites.TABLE_NAME,
-            target.writableDatabase,
-            LauncherSettings.Favorites.TMP_TABLE,
-            context,
-        )
+
+        copyTable(source, TABLE_NAME, target.writableDatabase, TMP_TABLE, context)
 
         val migrationStartTime = System.currentTimeMillis()
         try {
             SQLiteTransaction(target.writableDatabase).use { t ->
-                val srcReader = DbReader(t.db, LauncherSettings.Favorites.TMP_TABLE, context)
-                val destReader = DbReader(t.db, LauncherSettings.Favorites.TABLE_NAME, context)
+                val srcReader = DbReader(t.db, TMP_TABLE, context)
+                val destReader = DbReader(t.db, TABLE_NAME, context)
 
                 val targetSize = Point(destDeviceState.columns, destDeviceState.rows)
 
@@ -95,7 +103,7 @@
                 // Migrate workspace.
                 migrateWorkspace(srcReader, destReader, target, targetSize, idsInUse)
 
-                LauncherDbUtils.dropTable(t.db, LauncherSettings.Favorites.TMP_TABLE)
+                dropTable(t.db, TMP_TABLE)
                 t.commit()
             }
         } catch (e: Exception) {
@@ -112,6 +120,19 @@
         }
     }
 
+    private fun shouldShiftCells(destReader: DbReader, srcGridRowCount: Int): Boolean {
+        val workspaceItems = destReader.loadAllWorkspaceEntries()
+        val firstPageItemsRowPosSum =
+            workspaceItems.sumOf { entry -> if (entry.screenId == 0) entry.cellY else 0 }
+        val firstPageWorkspaceItemsCount = workspaceItems.count { entry -> entry.screenId == 0 }
+        if (firstPageWorkspaceItemsCount == 0) {
+            return false
+        }
+        val srcGridMidPoint = srcGridRowCount / 2f
+        val firstPageItemPosAvg = firstPageItemsRowPosSum / firstPageWorkspaceItemsCount.toFloat()
+        return (firstPageItemPosAvg >= srcGridMidPoint)
+    }
+
     /** Handles hotseat migration. */
     @VisibleForTesting
     fun migrateHotseat(
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index a830c96..83eace8 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -246,7 +246,7 @@
         TraceHelper.INSTANCE.beginSection(TAG);
         LoaderMemoryLogger memoryLogger = new LoaderMemoryLogger();
         mIsRestoreFromBackup =
-                (Boolean) LauncherPrefs.get(mApp.getContext()).get(IS_FIRST_LOAD_AFTER_RESTORE);
+                LauncherPrefs.get(mApp.getContext()).get(IS_FIRST_LOAD_AFTER_RESTORE);
         LauncherRestoreEventLogger restoreEventLogger = null;
         if (enableLauncherBrMetricsFixed()) {
             restoreEventLogger = LauncherRestoreEventLogger.Companion
@@ -266,21 +266,21 @@
                 sanitizeFolders(mItemsDeleted);
                 sanitizeAppPairs();
                 sanitizeWidgetsShortcutsAndPackages();
-                logASplit("sanitizeData");
+                logASplit("sanitizeData finished");
             }
 
             verifyNotStopped();
             mLauncherBinder.bindWorkspace(true /* incrementBindId */, /* isBindSync= */ false);
-            logASplit("bindWorkspace");
+            logASplit("bindWorkspace finished");
 
             mModelDelegate.workspaceLoadComplete();
             // Notify the installer packages of packages with active installs on the first screen.
             sendFirstScreenActiveInstallsBroadcast();
-            logASplit("sendFirstScreenBroadcast");
+            logASplit("sendFirstScreenBroadcast finished");
 
             // Take a break
             waitForIdle();
-            logASplit("step 1 complete");
+            logASplit("step 1 loading workspace complete");
             verifyNotStopped();
 
             // second step
@@ -291,11 +291,11 @@
             } finally {
                 Trace.endSection();
             }
-            logASplit("loadAllApps");
+            logASplit("loadAllApps finished");
 
             verifyNotStopped();
             mLauncherBinder.bindAllApps();
-            logASplit("bindAllApps");
+            logASplit("bindAllApps finished");
 
             verifyNotStopped();
             IconCacheUpdateHandler updateHandler = mIconCache.getUpdateHandler();
@@ -303,28 +303,28 @@
             updateHandler.updateIcons(allActivityList,
                     LauncherActivityCachingLogic.INSTANCE,
                     mApp.getModel()::onPackageIconsUpdated);
-            logASplit("update icon cache");
+            logASplit("update AllApps icon cache finished");
 
             verifyNotStopped();
-            logASplit("save shortcuts in icon cache");
+            logASplit("saving all shortcuts in icon cache");
             updateHandler.updateIcons(allShortcuts, CacheableShortcutCachingLogic.INSTANCE,
                     mApp.getModel()::onPackageIconsUpdated);
 
             // Take a break
             waitForIdle();
-            logASplit("step 2 complete");
+            logASplit("step 2 loading AllApps complete");
             verifyNotStopped();
 
             // third step
             List<ShortcutInfo> allDeepShortcuts = loadDeepShortcuts();
-            logASplit("loadDeepShortcuts");
+            logASplit("loadDeepShortcuts finished");
 
             verifyNotStopped();
             mLauncherBinder.bindDeepShortcuts();
-            logASplit("bindDeepShortcuts");
+            logASplit("bindDeepShortcuts finished");
 
             verifyNotStopped();
-            logASplit("save deep shortcuts in icon cache");
+            logASplit("saving deep shortcuts in icon cache");
             updateHandler.updateIcons(
                     convertShortcutsToCacheableShortcuts(allDeepShortcuts, allActivityList),
                     CacheableShortcutCachingLogic.INSTANCE,
@@ -332,7 +332,7 @@
 
             // Take a break
             waitForIdle();
-            logASplit("step 3 complete");
+            logASplit("step 3 loading all shortcuts complete");
             verifyNotStopped();
 
             // fourth step
@@ -345,11 +345,11 @@
                 widgetsModel.updateWidgetFilters(mWidgetsFilterDataProvider);
             }
             List<CachedObject> allWidgetsList = widgetsModel.update(mApp, /*packageUser=*/null);
-            logASplit("load widgets");
+            logASplit("load widgets finished");
 
             verifyNotStopped();
             mLauncherBinder.bindWidgets();
-            logASplit("bindWidgets");
+            logASplit("bindWidgets finished");
             verifyNotStopped();
             LauncherPrefs prefs = LauncherPrefs.get(mApp.getContext());
 
@@ -357,7 +357,7 @@
                 mLauncherBinder.bindSmartspaceWidget();
                 // Turn off pref.
                 prefs.putSync(SHOULD_SHOW_SMARTSPACE.to(false));
-                logASplit("bindSmartspaceWidget");
+                logASplit("bindSmartspaceWidget finished");
                 verifyNotStopped();
             } else if (!enableSmartspaceAsAWidget() && WIDGET_ON_FIRST_SCREEN
                     && !prefs.get(LauncherPrefs.SHOULD_SHOW_SMARTSPACE)) {
@@ -365,10 +365,10 @@
                 prefs.putSync(SHOULD_SHOW_SMARTSPACE.to(true));
             }
 
+            logASplit("saving all widgets in icon cache");
             updateHandler.updateIcons(allWidgetsList,
                     CachedObjectCachingLogic.INSTANCE,
                     mApp.getModel()::onWidgetLabelsUpdated);
-            logASplit("save widgets in icon cache");
 
             // fifth step
             loadFolderNames();
@@ -414,7 +414,7 @@
         } finally {
             Trace.endSection();
         }
-        logASplit("loadWorkspace");
+        logASplit("loadWorkspace finished");
 
         mBgDataModel.isFirstPagePinnedItemEnabled = FeatureFlags.QSB_ON_FIRST_SCREEN
                 && (!enableSmartspaceRemovalToggle() || LauncherPrefs.getPrefs(
@@ -440,7 +440,7 @@
         } else {
             dbController.tryMigrateDB(restoreEventLogger);
         }
-        Log.d(TAG, "loadWorkspace: loading default favorites");
+        Log.d(TAG, "loadWorkspace: loading default favorites if necessary");
         dbController.loadDefaultFavoritesIfNecessary();
 
         synchronized (mBgDataModel) {
@@ -453,7 +453,7 @@
                 mInstallingPkgsCached = installingPkgs;
             }
             installingPkgs.forEach(mApp.getIconCache()::updateSessionCache);
-            FileLog.d(TAG, "loadWorkspace: Packages with active install sessions: "
+            FileLog.d(TAG, "loadWorkspace: Packages with active install/update sessions: "
                     + installingPkgs.keySet().stream().map(info -> info.mPackageName).toList());
 
             mFirstScreenBroadcast = new FirstScreenBroadcast(installingPkgs);
@@ -478,8 +478,12 @@
                         widgetInflater, mPmHelper, iconRequestInfos, unlockedUsers,
                         allDeepShortcuts);
 
-                while (!mStopped && c.moveToNext()) {
-                    itemProcessor.processItem();
+                if (mStopped) {
+                    Log.w(TAG, "loadWorkspaceImpl: Loader stopped, skipping item processing");
+                } else {
+                    while (!mStopped && c.moveToNext()) {
+                        itemProcessor.processItem();
+                    }
                 }
                 tryLoadWorkspaceIconsInBulk(iconRequestInfos);
             } finally {
diff --git a/src/com/android/launcher3/pm/InstallSessionTracker.java b/src/com/android/launcher3/pm/InstallSessionTracker.java
index 856c294..b9c928c 100644
--- a/src/com/android/launcher3/pm/InstallSessionTracker.java
+++ b/src/com/android/launcher3/pm/InstallSessionTracker.java
@@ -33,7 +33,6 @@
 import androidx.annotation.WorkerThread;
 
 import com.android.launcher3.Flags;
-import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.util.PackageUserKey;
 
 import java.lang.ref.WeakReference;
@@ -79,7 +78,7 @@
         }
         SessionInfo sessionInfo = pushSessionDisplayToLauncher(sessionId, helper, callback);
         if (sessionInfo != null) {
-            FileLog.d(TAG, "onCreated: Install session created for"
+            Log.d(TAG, "onCreated: Install session created for"
                     + " appPackageName=" + sessionInfo.getAppPackageName()
                     + ", sessionId=" + sessionInfo.getSessionId()
                     + ", appIcon=" + sessionInfo.getAppIcon()
@@ -111,7 +110,7 @@
         activeSessions.remove(sessionId);
 
         if (key != null && key.mPackageName != null) {
-            FileLog.d(TAG, "onFinished: active install session finished for"
+            Log.d(TAG, "onFinished: active install session finished for"
                     + " appPackageName=" + key.mPackageName
                     + ", sessionId=" + sessionId
                     + ", success=" + success);
diff --git a/src/com/android/launcher3/provider/LauncherDbUtils.kt b/src/com/android/launcher3/provider/LauncherDbUtils.kt
index 3c68e46..6f1d0dd 100644
--- a/src/com/android/launcher3/provider/LauncherDbUtils.kt
+++ b/src/com/android/launcher3/provider/LauncherDbUtils.kt
@@ -131,6 +131,11 @@
         }
     }
 
+    @JvmStatic
+    fun shiftTableByXCells(db: SQLiteDatabase, x: Int, toTable: String) {
+        db.run { execSQL("UPDATE $toTable SET cellY = cellY + $x") }
+    }
+
     /**
      * Migrates the legacy shortcuts to deep shortcuts pinned under Launcher. Removes any invalid
      * shortcut or any shortcut which requires some permission to launch
diff --git a/src/com/android/launcher3/widget/BaseWidgetSheet.java b/src/com/android/launcher3/widget/BaseWidgetSheet.java
index 1c0d94c..fda5175 100644
--- a/src/com/android/launcher3/widget/BaseWidgetSheet.java
+++ b/src/com/android/launcher3/widget/BaseWidgetSheet.java
@@ -20,6 +20,7 @@
 import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.anim.AnimatorListeners.forSuccessCallback;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_WIDGET_ADD_BUTTON_TAP;
+import static com.android.window.flags.Flags.predictiveBackThreeButtonNav;
 
 import android.content.Context;
 import android.graphics.Canvas;
@@ -128,6 +129,17 @@
     }
 
     @Override
+    public void setScaleY(float scaleY) {
+        super.setScaleY(scaleY);
+        if (predictiveBackThreeButtonNav() && mNavBarScrimHeight > 0) {
+            // Call invalidate to prevent navbar scrim from scaling. The navbar scrim is drawn
+            // directly onto the canvas. To prevent it from being scaled with the canvas, there's a
+            // counter scale applied in dispatchDraw.
+            invalidate();
+        }
+    }
+
+    @Override
     public final void onClick(View v) {
         WidgetCell wc;
         if (v instanceof WidgetCell view) {
@@ -318,8 +330,10 @@
         super.dispatchDraw(canvas);
 
         if (mNavBarScrimHeight > 0) {
-            canvas.drawRect(0, getHeight() - mNavBarScrimHeight, getWidth(), getHeight(),
-                    mNavBarScrimPaint);
+            float left = (getWidth() - getWidth() / getScaleX()) / 2;
+            float top = getHeight() / 2f + (getHeight() / 2f - mNavBarScrimHeight) / getScaleY();
+            canvas.drawRect(left, top, getWidth() / getScaleX(),
+                    top + mNavBarScrimHeight / getScaleY(), mNavBarScrimPaint);
         }
     }
 
diff --git a/tests/Android.bp b/tests/Android.bp
index b1d4ef6..e4fecc5 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -63,7 +63,6 @@
         "src/com/android/launcher3/dragging/TaplDragTest.java",
         "src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java",
         "src/com/android/launcher3/ui/TaplTestsLauncher3Test.java",
-        "src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java",
         "src/com/android/launcher3/ui/workspace/TaplWorkspaceTest.java",
     ],
 }
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index ed5762d..8e4db5c 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -21,11 +21,8 @@
 
 import static org.junit.Assert.assertTrue;
 
-import android.content.BroadcastReceiver;
 import android.content.ComponentName;
-import android.content.Context;
 import android.content.Intent;
-import android.content.IntentFilter;
 import android.os.Process;
 import android.system.OsConstants;
 import android.util.Log;
@@ -53,7 +50,6 @@
 
 import java.util.Objects;
 import java.util.concurrent.Callable;
-import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.function.Consumer;
@@ -216,38 +212,6 @@
         }, mLauncher, timeout);
     }
 
-    /**
-     * Broadcast receiver which blocks until the result is received.
-     */
-    public class BlockingBroadcastReceiver extends BroadcastReceiver {
-
-        private final CountDownLatch latch = new CountDownLatch(1);
-        private Intent mIntent;
-
-        public BlockingBroadcastReceiver(String action) {
-            mTargetContext.registerReceiver(this, new IntentFilter(action),
-                    Context.RECEIVER_EXPORTED/*UNAUDITED*/);
-        }
-
-        @Override
-        public void onReceive(Context context, Intent intent) {
-            mIntent = intent;
-            latch.countDown();
-        }
-
-        public Intent blockingGetIntent() throws InterruptedException {
-            assertTrue("Timed Out", latch.await(DEFAULT_BROADCAST_TIMEOUT_SECS, TimeUnit.SECONDS));
-            mTargetContext.unregisterReceiver(this);
-            return mIntent;
-        }
-
-        public Intent blockingGetExtraIntent() throws InterruptedException {
-            Intent intent = blockingGetIntent();
-            return intent == null ? null : (Intent) intent.getParcelableExtra(
-                    Intent.EXTRA_INTENT);
-        }
-    }
-
     public static void startAppFast(String packageName) {
         startIntent(
                 getInstrumentation().getContext().getPackageManager().getLaunchIntentForPackage(
diff --git a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
new file mode 100644
index 0000000..bb645d7
--- /dev/null
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -0,0 +1,188 @@
+/*
+ * Copyright (C) 2017 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.widget;
+
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
+import static com.android.launcher3.util.LauncherBindableItemsContainer.ItemOperator;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
+
+import android.appwidget.AppWidgetManager;
+import android.content.Context;
+import android.content.Intent;
+import android.os.Process;
+import android.view.View;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.LargeTest;
+
+import com.android.launcher3.AbstractFloatingView;
+import com.android.launcher3.Launcher;
+import com.android.launcher3.celllayout.FavoriteItemsTransaction;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.model.data.LauncherAppWidgetInfo;
+import com.android.launcher3.testcomponent.WidgetConfigActivity;
+import com.android.launcher3.ui.PortraitLandscapeRunner.PortraitLandscape;
+import com.android.launcher3.ui.TestViewHelpers;
+import com.android.launcher3.util.BaseLauncherActivityTest;
+import com.android.launcher3.util.BlockingBroadcastReceiver;
+import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.util.Wait;
+import com.android.launcher3.util.rule.ShellCommandRule;
+import com.android.launcher3.views.OptionsPopupView;
+import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.widget.PendingAddWidgetInfo;
+import com.android.launcher3.widget.WidgetCell;
+import com.android.launcher3.widget.picker.WidgetsFullSheet;
+import com.android.launcher3.widget.picker.WidgetsListAdapter;
+import com.android.launcher3.widget.picker.WidgetsRecyclerView;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/**
+ * Test to verify widget configuration is properly shown.
+ */
+@LargeTest
+@RunWith(AndroidJUnit4.class)
+public class AddConfigWidgetTest extends BaseLauncherActivityTest<Launcher> {
+
+    @Rule
+    public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grantWidgetBind();
+
+    private LauncherAppWidgetProviderInfo mWidgetInfo;
+    private AppWidgetManager mAppWidgetManager;
+
+    private int mWidgetId;
+
+    @Before
+    public void setUp() throws Exception {
+        mWidgetInfo = TestViewHelpers.findWidgetProvider(true /* hasConfigureScreen */);
+        mAppWidgetManager = AppWidgetManager.getInstance(targetContext());
+    }
+
+    @Test
+    @PortraitLandscape
+    public void testWidgetConfig() throws Throwable {
+        runTest(true);
+    }
+
+    @Test
+    @PortraitLandscape
+    public void testConfigCancelled() throws Throwable {
+        runTest(false);
+    }
+
+    /**
+     * @param acceptConfig accept the config activity
+     */
+    private void runTest(boolean acceptConfig) throws Throwable {
+        new FavoriteItemsTransaction(targetContext()).commit();
+        loadLauncherSync();
+
+        // Add widget to homescreen
+        WidgetConfigStartupMonitor monitor = new WidgetConfigStartupMonitor();
+        executeOnLauncher(OptionsPopupView::openWidgets);
+        uiDevice.waitForIdle();
+
+        // Select the widget header
+        Context testContext = getInstrumentation().getContext();
+        String packageName = testContext.getPackageName();
+        executeOnLauncher(l -> {
+            WidgetsRecyclerView wrv = WidgetsFullSheet.getWidgetsView(l);
+            WidgetsListAdapter adapter = (WidgetsListAdapter) wrv.getAdapter();
+            int pos = adapter.getItems().indexOf(
+                    adapter.getItems().stream()
+                            .filter(entry -> packageName.equals(entry.mPkgItem.packageName))
+                            .findFirst()
+                            .get());
+            wrv.getLayoutManager().scrollToPosition(pos);
+            adapter.onHeaderClicked(true, new PackageUserKey(packageName, Process.myUserHandle()));
+        });
+        uiDevice.waitForIdle();
+
+        View widgetView = getOnceNotNull("Widget not found", l -> searchView(l.getDragLayer(), v ->
+                v instanceof WidgetCell
+                        && v.getTag() instanceof PendingAddWidgetInfo pawi
+                        && mWidgetInfo.provider.equals(pawi.componentName)));
+        addToWorkspace(widgetView);
+
+        // Widget id for which the config activity was opened
+        mWidgetId = monitor.getWidgetId();
+
+        // Verify that the widget id is valid and bound
+        assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
+        setResult(acceptConfig);
+
+        if (acceptConfig) {
+            getOnceNotNull("Widget was not added", l -> {
+                // Close the resize frame before searching for widget
+                AbstractFloatingView.closeAllOpenViews(l);
+                return l.getWorkspace().getFirstMatch(new WidgetSearchCondition());
+            });
+            assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
+        } else {
+            // Verify that the widget id is deleted.
+            Wait.atMost("", () -> mAppWidgetManager.getAppWidgetInfo(mWidgetId) == null);
+        }
+    }
+
+    private void setResult(boolean success) {
+        getInstrumentation().getTargetContext().sendBroadcast(
+                WidgetConfigActivity.getCommandIntent(WidgetConfigActivity.class,
+                        success ? "clickOK" : "clickCancel"));
+        uiDevice.waitForIdle();
+    }
+
+    /**
+     * Condition for searching widget id
+     */
+    private class WidgetSearchCondition implements ItemOperator {
+
+        @Override
+        public boolean evaluate(ItemInfo info, View view) {
+            return info instanceof LauncherAppWidgetInfo lawi
+                    && lawi.providerName.equals(mWidgetInfo.provider)
+                    && lawi.appWidgetId == mWidgetId;
+        }
+    }
+
+    /**
+     * Broadcast receiver for receiving widget config activity status.
+     */
+    private static class WidgetConfigStartupMonitor extends BlockingBroadcastReceiver {
+
+        WidgetConfigStartupMonitor() {
+            super(WidgetConfigActivity.class.getName());
+        }
+
+        public int getWidgetId() throws InterruptedException {
+            Intent intent = blockingGetExtraIntent();
+            assertNotNull("Null EXTRA_INTENT", intent);
+            assertEquals("Intent action is not ACTION_APPWIDGET_CONFIGURE",
+                    AppWidgetManager.ACTION_APPWIDGET_CONFIGURE, intent.getAction());
+            int widgetId = intent.getIntExtra(AppWidgetManager.EXTRA_APPWIDGET_ID,
+                    LauncherAppWidgetInfo.NO_ID);
+            assertNotSame("Widget id is NO_ID", widgetId, LauncherAppWidgetInfo.NO_ID);
+            return widgetId;
+        }
+    }
+}
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
similarity index 77%
rename from tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java
rename to tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
index 4cdbd96..8846d65 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
@@ -1,17 +1,17 @@
 /*
  * Copyright (C) 2017 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
+ * 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
+ *      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.
+ * 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.widget;
 
@@ -22,12 +22,12 @@
 import static com.android.launcher3.model.data.LauncherAppWidgetInfo.FLAG_RESTORE_STARTED;
 import static com.android.launcher3.provider.LauncherDbUtils.itemIdMatch;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
+import static com.android.launcher3.util.TestUtil.getOnUiThread;
+import static com.android.launcher3.util.Wait.atMost;
 import static com.android.launcher3.util.WidgetUtils.createWidgetInfo;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertTrue;
 
 import android.appwidget.AppWidgetManager;
 import android.content.ComponentName;
@@ -36,6 +36,7 @@
 import android.content.pm.PackageManager;
 import android.database.Cursor;
 import android.os.Bundle;
+import android.text.TextUtils;
 import android.widget.RemoteViews;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
@@ -49,13 +50,12 @@
 import com.android.launcher3.celllayout.FavoriteItemsTransaction;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.pm.InstallSessionHelper;
-import com.android.launcher3.tapl.Widget;
-import com.android.launcher3.tapl.Workspace;
-import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.TestViewHelpers;
-import com.android.launcher3.util.TestUtil;
+import com.android.launcher3.util.BaseLauncherActivityTest;
 import com.android.launcher3.util.rule.ShellCommandRule;
+import com.android.launcher3.widget.LauncherAppWidgetHostView;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.widget.PendingAppWidgetHostView;
 import com.android.launcher3.widget.WidgetManagerHelper;
 
 import org.junit.After;
@@ -67,6 +67,7 @@
 import java.util.HashSet;
 import java.util.Set;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 /**
  * Tests for bind widget flow.
@@ -75,7 +76,7 @@
  */
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-public class TaplBindWidgetTest extends AbstractLauncherUiTest<Launcher> {
+public class BindWidgetTest extends BaseLauncherActivityTest<Launcher> {
 
     @Rule
     public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grantWidgetBind();
@@ -87,11 +88,9 @@
 
     private LauncherModel mModel;
 
-    @Override
     @Before
     public void setUp() throws Exception {
-        super.setUp();
-        mModel = LauncherAppState.getInstance(mTargetContext).getModel();
+        mModel = LauncherAppState.getInstance(targetContext()).getModel();
     }
 
     @After
@@ -101,7 +100,7 @@
         }
 
         if (mSessionId > -1) {
-            mTargetContext.getPackageManager().getPackageInstaller().abandonSession(mSessionId);
+            targetContext().getPackageManager().getPackageInstaller().abandonSession(mSessionId);
         }
     }
 
@@ -122,13 +121,12 @@
         LauncherAppWidgetProviderInfo info = addWidgetToScreen(false, false,
                 item -> item.appWidgetId = -33);
 
-        final Workspace workspace = mLauncher.getWorkspace();
         // Item deleted from db
         mCursor = queryItem();
         assertEquals(0, mCursor.getCount());
 
         // The view does not exist
-        assertTrue("Widget exists", workspace.tryGetWidget(info.label, 0) == null);
+        verifyItemEventuallyNull("Widget exists", widgetProvider(info));
     }
 
     @Test
@@ -154,18 +152,19 @@
         // Widget has a valid Id now.
         assertEquals(0, mCursor.getInt(mCursor.getColumnIndex(LauncherSettings.Favorites.RESTORED))
                 & FLAG_ID_NOT_VALID);
-        assertNotNull(AppWidgetManager.getInstance(mTargetContext)
+        assertNotNull(AppWidgetManager.getInstance(targetContext())
                 .getAppWidgetInfo(mCursor.getInt(mCursor.getColumnIndex(
                         LauncherSettings.Favorites.APPWIDGET_ID))));
 
         // send OPTION_APPWIDGET_RESTORE_COMPLETED
         int appWidgetId = mCursor.getInt(
                 mCursor.getColumnIndex(LauncherSettings.Favorites.APPWIDGET_ID));
-        AppWidgetManager appWidgetManager = AppWidgetManager.getInstance(mTargetContext);
+        AppWidgetManager appWidgetManager = AppWidgetManager.getInstance(targetContext());
 
         Bundle b = new Bundle();
         b.putBoolean(WidgetManagerHelper.WIDGET_OPTION_RESTORE_COMPLETED, true);
-        RemoteViews remoteViews = new RemoteViews(mTargetPackage, R.layout.appwidget_not_ready);
+        RemoteViews remoteViews = new RemoteViews(
+                targetContext().getPackageName(), R.layout.appwidget_not_ready);
         appWidgetManager.updateAppWidgetOptions(appWidgetId, b);
         appWidgetManager.updateAppWidget(appWidgetId, remoteViews);
 
@@ -175,15 +174,14 @@
                         WidgetManagerHelper.WIDGET_OPTION_RESTORE_COMPLETED));
         executeOnLauncher(l -> l.getAppWidgetHolder().startListening());
         verifyWidgetPresent(info);
-        assertNull(mLauncher.getWorkspace().tryGetPendingWidget(100));
+        verifyItemEventuallyNull("Pending widget exists", pendingWidgetProvider());
     }
 
     @Test
     public void testPendingWidget_notRestored_removed() {
         addPendingItemToScreen(getInvalidWidgetInfo(), FLAG_ID_NOT_VALID | FLAG_PROVIDER_NOT_READY);
 
-        assertTrue("Pending widget exists",
-                mLauncher.getWorkspace().tryGetPendingWidget(0) == null);
+        verifyItemEventuallyNull("Pending widget exists", pendingWidgetProvider());
         // Item deleted from db
         mCursor = queryItem();
         assertEquals(0, mCursor.getCount());
@@ -216,7 +214,7 @@
         // Create an active installer session
         SessionParams params = new SessionParams(SessionParams.MODE_FULL_INSTALL);
         params.setAppPackageName(item.providerName.getPackageName());
-        PackageInstaller installer = mTargetContext.getPackageManager().getPackageInstaller();
+        PackageInstaller installer = targetContext().getPackageManager().getPackageInstaller();
         mSessionId = installer.createSession(params);
 
         addPendingItemToScreen(item, FLAG_ID_NOT_VALID | FLAG_PROVIDER_NOT_READY);
@@ -234,36 +232,47 @@
     }
 
     private void verifyWidgetPresent(LauncherAppWidgetProviderInfo info) {
-        final Widget widget = mLauncher.getWorkspace().tryGetWidget(info.label,
-                TestUtil.DEFAULT_UI_TIMEOUT);
-        assertTrue("Widget is not present",
-                widget != null);
+        getOnceNotNull("Widget is not present", widgetProvider(info));
     }
 
     private void verifyPendingWidgetPresent() {
-        final Widget widget = mLauncher.getWorkspace().tryGetPendingWidget(
-                TestUtil.DEFAULT_UI_TIMEOUT);
-        assertTrue("Pending widget is not present",
-                widget != null);
+        getOnceNotNull("Widget is not present", pendingWidgetProvider());
+    }
+
+    private Function<Launcher, Object> pendingWidgetProvider() {
+        return l -> l.getWorkspace().getFirstMatch(
+                (item, view) -> view instanceof PendingAppWidgetHostView);
+    }
+
+    private Function<Launcher, Object> widgetProvider(LauncherAppWidgetProviderInfo info) {
+        return l -> l.getWorkspace().getFirstMatch((item, view) ->
+                view instanceof LauncherAppWidgetHostView
+                        && TextUtils.equals(info.label, view.getContentDescription()));
+    }
+
+    private void verifyItemEventuallyNull(String message, Function<Launcher, Object> provider) {
+        atMost(message, () -> getFromLauncher(provider) == null);
     }
 
     private void addPendingItemToScreen(LauncherAppWidgetInfo item, int restoreStatus) {
         item.restoreStatus = restoreStatus;
         item.screenId = FIRST_SCREEN_ID;
-        commitTransactionAndLoadHome(
-                new FavoriteItemsTransaction(mTargetContext).addItem(() -> item));
+        new FavoriteItemsTransaction(targetContext()).addItem(() -> item).commit();
+        loadLauncherSync();
     }
 
     private LauncherAppWidgetProviderInfo addWidgetToScreen(boolean hasConfigureScreen,
             boolean bindWidget, Consumer<LauncherAppWidgetInfo> itemOverride) {
         LauncherAppWidgetProviderInfo info = TestViewHelpers.findWidgetProvider(hasConfigureScreen);
-        commitTransactionAndLoadHome(new FavoriteItemsTransaction(mTargetContext)
+        new FavoriteItemsTransaction(targetContext())
                 .addItem(() -> {
-                    LauncherAppWidgetInfo item = createWidgetInfo(info, mTargetContext, bindWidget);
+                    LauncherAppWidgetInfo item =
+                            createWidgetInfo(info, targetContext(), bindWidget);
                     item.screenId = FIRST_SCREEN_ID;
                     itemOverride.accept(item);
                     return item;
-                }));
+                }).commit();
+        loadLauncherSync();
         return info;
     }
 
@@ -277,13 +286,13 @@
 
         Set<String> activePackage = getOnUiThread(() -> {
             Set<String> packages = new HashSet<>();
-            InstallSessionHelper.INSTANCE.get(mTargetContext).getActiveSessions()
+            InstallSessionHelper.INSTANCE.get(targetContext()).getActiveSessions()
                     .keySet().forEach(packageUserKey -> packages.add(packageUserKey.mPackageName));
             return packages;
         });
         while (true) {
             try {
-                mTargetContext.getPackageManager().getPackageInfo(
+                targetContext().getPackageManager().getPackageInfo(
                         pkg, PackageManager.GET_UNINSTALLED_PACKAGES);
             } catch (Exception e) {
                 if (!activePackage.contains(pkg)) {
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java b/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
similarity index 62%
rename from tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java
rename to tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
index fe3b2ee..2fb7987 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/RequestPinItemTest.java
@@ -1,34 +1,41 @@
 /*
  * Copyright (C) 2017 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
+ * 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
+ *      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.
+ * 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.widget;
 
 import static android.app.PendingIntent.FLAG_MUTABLE;
 import static android.app.PendingIntent.FLAG_ONE_SHOT;
 
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNotSame;
 
+import static java.util.regex.Pattern.CASE_INSENSITIVE;
+
 import android.app.PendingIntent;
 import android.appwidget.AppWidgetManager;
+import android.content.Context;
 import android.content.Intent;
 import android.graphics.Color;
 import android.view.View;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.uiautomator.By;
+import androidx.test.uiautomator.BySelector;
 
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherSettings.Favorites;
@@ -37,14 +44,13 @@
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.shortcuts.ShortcutKey;
-import com.android.launcher3.tapl.AddToHomeScreenPrompt;
 import com.android.launcher3.testcomponent.AppWidgetNoConfig;
 import com.android.launcher3.testcomponent.AppWidgetWithConfig;
 import com.android.launcher3.testcomponent.RequestPinItemActivity;
-import com.android.launcher3.ui.AbstractLauncherUiTest;
+import com.android.launcher3.util.BaseLauncherActivityTest;
+import com.android.launcher3.util.BlockingBroadcastReceiver;
 import com.android.launcher3.util.LauncherBindableItemsContainer.ItemOperator;
-import com.android.launcher3.util.Wait;
-import com.android.launcher3.util.Wait.Condition;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.rule.ShellCommandRule;
 
 import org.junit.Before;
@@ -53,25 +59,27 @@
 import org.junit.runner.RunWith;
 
 import java.util.UUID;
+import java.util.regex.Pattern;
 
 /**
  * Test to verify pin item request flow.
  */
 @LargeTest
 @RunWith(AndroidJUnit4.class)
-public class TaplRequestPinItemTest extends AbstractLauncherUiTest<Launcher> {
+public class RequestPinItemTest extends BaseLauncherActivityTest<Launcher> {
 
     @Rule
     public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grantWidgetBind();
 
+    @Rule
+    public ShellCommandRule mDefaultLauncherRule = ShellCommandRule.setDefaultLauncher();
+
     private String mCallbackAction;
     private String mShortcutId;
     private int mAppWidgetId;
 
-    @Override
     @Before
     public void setUp() throws Exception {
-        super.setUp();
         mCallbackAction = UUID.randomUUID().toString();
         mShortcutId = UUID.randomUUID().toString();
     }
@@ -81,9 +89,9 @@
 
     @Test
     public void testPinWidgetNoConfig() throws Throwable {
-        runTest("pinWidgetNoConfig", true, (info, view) -> info instanceof LauncherAppWidgetInfo &&
-                ((LauncherAppWidgetInfo) info).appWidgetId == mAppWidgetId &&
-                ((LauncherAppWidgetInfo) info).providerName.getClassName()
+        runTest("pinWidgetNoConfig", true, (info, view) -> info instanceof LauncherAppWidgetInfo
+                && ((LauncherAppWidgetInfo) info).appWidgetId == mAppWidgetId
+                && ((LauncherAppWidgetInfo) info).providerName.getClassName()
                         .equals(AppWidgetNoConfig.class.getName()));
     }
 
@@ -94,18 +102,18 @@
                 RequestPinItemActivity.class, "setRemoteViewColor").putExtra(
                 RequestPinItemActivity.EXTRA_PARAM + "0", Color.RED);
 
-        runTest("pinWidgetNoConfig", true, (info, view) -> info instanceof LauncherAppWidgetInfo &&
-                ((LauncherAppWidgetInfo) info).appWidgetId == mAppWidgetId &&
-                ((LauncherAppWidgetInfo) info).providerName.getClassName()
+        runTest("pinWidgetNoConfig", true, (info, view) -> info instanceof LauncherAppWidgetInfo
+                && ((LauncherAppWidgetInfo) info).appWidgetId == mAppWidgetId
+                && ((LauncherAppWidgetInfo) info).providerName.getClassName()
                         .equals(AppWidgetNoConfig.class.getName()), command);
     }
 
     @Test
     public void testPinWidgetWithConfig() throws Throwable {
         runTest("pinWidgetWithConfig", true,
-                (info, view) -> info instanceof LauncherAppWidgetInfo &&
-                        ((LauncherAppWidgetInfo) info).appWidgetId == mAppWidgetId &&
-                        ((LauncherAppWidgetInfo) info).providerName.getClassName()
+                (info, view) -> info instanceof LauncherAppWidgetInfo
+                        && ((LauncherAppWidgetInfo) info).appWidgetId == mAppWidgetId
+                        && ((LauncherAppWidgetInfo) info).providerName.getClassName()
                                 .equals(AppWidgetWithConfig.class.getName()));
     }
 
@@ -119,47 +127,48 @@
         runTest("pinShortcut", false, new ItemOperator() {
             @Override
             public boolean evaluate(ItemInfo info, View view) {
-                return info instanceof WorkspaceItemInfo &&
-                        info.itemType == Favorites.ITEM_TYPE_DEEP_SHORTCUT &&
-                        ShortcutKey.fromItemInfo(info).getId().equals(mShortcutId);
+                return info instanceof WorkspaceItemInfo
+                        && info.itemType == Favorites.ITEM_TYPE_DEEP_SHORTCUT
+                        && ShortcutKey.fromItemInfo(info).getId().equals(mShortcutId);
             }
         }, command);
     }
 
     private void runTest(String activityMethod, boolean isWidget, ItemOperator itemMatcher,
             Intent... commandIntents) throws Throwable {
-        commitTransactionAndLoadHome(new FavoriteItemsTransaction(mTargetContext));
+        new FavoriteItemsTransaction(targetContext()).commit();
+        loadLauncherSync();
 
         // Open Pin item activity
         BlockingBroadcastReceiver openMonitor = new BlockingBroadcastReceiver(
                 RequestPinItemActivity.class.getName());
-        mLauncher.
-                getWorkspace().
-                switchToAllApps().
-                getAppIcon("Test Pin Item").
-                launch(getAppPackageName());
+        Context testContext = getInstrumentation().getContext();
+        startAppFast(
+                testContext.getPackageName(),
+                new Intent(testContext, RequestPinItemActivity.class));
         assertNotNull(openMonitor.blockingGetExtraIntent());
 
         // Set callback
-        PendingIntent callback = PendingIntent.getBroadcast(mTargetContext, 0,
-                new Intent(mCallbackAction).setPackage(mTargetContext.getPackageName()),
+        PendingIntent callback = PendingIntent.getBroadcast(targetContext(), 0,
+                new Intent(mCallbackAction).setPackage(targetContext().getPackageName()),
                 FLAG_ONE_SHOT | FLAG_MUTABLE);
-        mTargetContext.sendBroadcast(RequestPinItemActivity.getCommandIntent(
+        targetContext().sendBroadcast(RequestPinItemActivity.getCommandIntent(
                 RequestPinItemActivity.class, "setCallback").putExtra(
                 RequestPinItemActivity.EXTRA_PARAM + "0", callback));
 
         for (Intent command : commandIntents) {
-            mTargetContext.sendBroadcast(command);
+            targetContext().sendBroadcast(command);
         }
 
         // call the requested method to start the flow
-        mTargetContext.sendBroadcast(RequestPinItemActivity.getCommandIntent(
+        targetContext().sendBroadcast(RequestPinItemActivity.getCommandIntent(
                 RequestPinItemActivity.class, activityMethod));
-        final AddToHomeScreenPrompt addToHomeScreenPrompt = mLauncher.getAddToHomeScreenPrompt();
 
         // Accept confirmation:
         BlockingBroadcastReceiver resultReceiver = new BlockingBroadcastReceiver(mCallbackAction);
-        addToHomeScreenPrompt.addAutomatically();
+        BySelector selector = By.text(Pattern.compile("^Add to home screen$", CASE_INSENSITIVE))
+                .pkg(targetContext().getPackageName());
+        uiDevice.wait(device -> device.findObject(selector), TestUtil.DEFAULT_UI_TIMEOUT).click();
         Intent result = resultReceiver.blockingGetIntent();
         assertNotNull(result);
         mAppWidgetId = result.getIntExtra(AppWidgetManager.EXTRA_APPWIDGET_ID, -1);
@@ -167,28 +176,9 @@
             assertNotSame(-1, mAppWidgetId);
         }
 
-        // Go back to home
-        mLauncher.goHome();
-        Wait.atMost("", new ItemSearchCondition(itemMatcher), mLauncher);
-    }
-
-    /**
-     * Condition for for an item
-     */
-    private class ItemSearchCondition implements Condition {
-
-        private final ItemOperator mOp;
-
-        ItemSearchCondition(ItemOperator op) {
-            mOp = op;
-        }
-
-        @Override
-        public boolean isTrue() throws Throwable {
-            return mMainThreadExecutor.submit(() -> {
-                Launcher l = Launcher.ACTIVITY_TRACKER.getCreatedContext();
-                return l != null && l.getWorkspace().getFirstMatch(mOp) != null;
-            }).get();
-        }
+        // Reload activity, so that the activity is focused
+        closeCurrentActivity();
+        loadLauncherSync();
+        getOnceNotNull("", l -> l.getWorkspace().getFirstMatch(itemMatcher));
     }
 }
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
deleted file mode 100644
index 7845222..0000000
--- a/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
+++ /dev/null
@@ -1,173 +0,0 @@
-/*
- * Copyright (C) 2017 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.widget;
-
-import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
-
-import static com.android.launcher3.util.LauncherBindableItemsContainer.ItemOperator;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertNotSame;
-
-import android.appwidget.AppWidgetManager;
-import android.content.Intent;
-import android.view.View;
-
-import androidx.test.filters.LargeTest;
-import androidx.test.runner.AndroidJUnit4;
-
-import com.android.launcher3.Launcher;
-import com.android.launcher3.celllayout.FavoriteItemsTransaction;
-import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.model.data.LauncherAppWidgetInfo;
-import com.android.launcher3.testcomponent.WidgetConfigActivity;
-import com.android.launcher3.ui.AbstractLauncherUiTest;
-import com.android.launcher3.ui.PortraitLandscapeRunner.PortraitLandscape;
-import com.android.launcher3.ui.TestViewHelpers;
-import com.android.launcher3.util.Wait;
-import com.android.launcher3.util.rule.ShellCommandRule;
-import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
-
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-
-/**
- * Test to verify widget configuration is properly shown.
- */
-@LargeTest
-@RunWith(AndroidJUnit4.class)
-public class TaplAddConfigWidgetTest extends AbstractLauncherUiTest<Launcher> {
-
-    @Rule
-    public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grantWidgetBind();
-
-    private LauncherAppWidgetProviderInfo mWidgetInfo;
-    private AppWidgetManager mAppWidgetManager;
-
-    private int mWidgetId;
-
-    @Override
-    @Before
-    public void setUp() throws Exception {
-        super.setUp();
-        mWidgetInfo = TestViewHelpers.findWidgetProvider(true /* hasConfigureScreen */);
-        mAppWidgetManager = AppWidgetManager.getInstance(mTargetContext);
-    }
-
-    @Test
-    @PortraitLandscape
-    public void testWidgetConfig() throws Throwable {
-        runTest(true);
-    }
-
-    @Test
-    @PortraitLandscape
-    public void testConfigCancelled() throws Throwable {
-        runTest(false);
-    }
-
-
-    /**
-     * @param acceptConfig accept the config activity
-     */
-    private void runTest(boolean acceptConfig) throws Throwable {
-        commitTransactionAndLoadHome(new FavoriteItemsTransaction(mTargetContext));
-
-        // Drag widget to homescreen
-        WidgetConfigStartupMonitor monitor = new WidgetConfigStartupMonitor();
-        mLauncher.getWorkspace()
-                .openAllWidgets()
-                .getWidget(mWidgetInfo.getLabel())
-                .dragToWorkspace(true, false);
-        // Widget id for which the config activity was opened
-        mWidgetId = monitor.getWidgetId();
-
-        // Verify that the widget id is valid and bound
-        assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
-
-        setResultAndWaitForAnimation(acceptConfig);
-        if (acceptConfig) {
-            Wait.atMost("", new WidgetSearchCondition(), mLauncher);
-            assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
-        } else {
-            // Verify that the widget id is deleted.
-            Wait.atMost("", () -> mAppWidgetManager.getAppWidgetInfo(mWidgetId) == null,
-                    mLauncher);
-        }
-    }
-
-    private static void setResult(boolean success) {
-        getInstrumentation().getTargetContext().sendBroadcast(
-                WidgetConfigActivity.getCommandIntent(WidgetConfigActivity.class,
-                        success ? "clickOK" : "clickCancel"));
-    }
-
-    private void setResultAndWaitForAnimation(boolean success) {
-        if (mLauncher.isLauncher3()) {
-            setResult(success);
-        } else {
-            mLauncher.executeAndWaitForWallpaperAnimation(
-                    () -> setResult(success),
-                    "setting widget coinfig result");
-        }
-    }
-
-    /**
-     * Condition for searching widget id
-     */
-    private class WidgetSearchCondition implements Wait.Condition, ItemOperator {
-
-        @Override
-        public boolean isTrue() throws Throwable {
-            return mMainThreadExecutor.submit(() -> {
-                Launcher l = Launcher.ACTIVITY_TRACKER.getCreatedContext();
-                return l != null && l.getWorkspace().getFirstMatch(this) != null;
-            }).get();
-        }
-
-        @Override
-        public boolean evaluate(ItemInfo info, View view) {
-            return info instanceof LauncherAppWidgetInfo
-                    && ((LauncherAppWidgetInfo) info).providerName.getClassName().equals(
-                            mWidgetInfo.provider.getClassName())
-                    && ((LauncherAppWidgetInfo) info).appWidgetId == mWidgetId;
-        }
-    }
-
-    /**
-     * Broadcast receiver for receiving widget config activity status.
-     */
-    private class WidgetConfigStartupMonitor extends BlockingBroadcastReceiver {
-
-        public WidgetConfigStartupMonitor() {
-            super(WidgetConfigActivity.class.getName());
-        }
-
-        public int getWidgetId() throws InterruptedException {
-            Intent intent = blockingGetExtraIntent();
-            assertNotNull("Null EXTRA_INTENT", intent);
-            assertEquals("Intent action is not ACTION_APPWIDGET_CONFIGURE",
-                    AppWidgetManager.ACTION_APPWIDGET_CONFIGURE, intent.getAction());
-            int widgetId = intent.getIntExtra(AppWidgetManager.EXTRA_APPWIDGET_ID,
-                    LauncherAppWidgetInfo.NO_ID);
-            assertNotSame("Widget id is NO_ID", widgetId, LauncherAppWidgetInfo.NO_ID);
-            return widgetId;
-        }
-    }
-}
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java b/tests/src/com/android/launcher3/ui/widget/WidgetPickerTest.java
similarity index 68%
rename from tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java
rename to tests/src/com/android/launcher3/ui/widget/WidgetPickerTest.java
index 19c5850..caad1d9 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplWidgetPickerTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/WidgetPickerTest.java
@@ -22,24 +22,30 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.MediumTest;
 
+import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.Launcher;
-import com.android.launcher3.tapl.Widgets;
-import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.PortraitLandscapeRunner.PortraitLandscape;
+import com.android.launcher3.util.BaseLauncherActivityTest;
+import com.android.launcher3.util.rule.ScreenRecordRule;
 import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
+import com.android.launcher3.views.OptionsPopupView;
 import com.android.launcher3.widget.picker.WidgetsFullSheet;
 import com.android.launcher3.widget.picker.WidgetsRecyclerView;
 
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.TestRule;
 import org.junit.runner.RunWith;
 
 /**
- * This test run in both Out of process (Oop) and in-process (Ipc).
  * Make sure the basic interactions with the WidgetPicker works.
  */
 @MediumTest
 @RunWith(AndroidJUnit4.class)
-public class TaplWidgetPickerTest extends AbstractLauncherUiTest<Launcher> {
+public class WidgetPickerTest extends BaseLauncherActivityTest<Launcher> {
+
+    @Rule
+    public TestRule screenRecordRule = new ScreenRecordRule();
 
     private WidgetsRecyclerView getWidgetsView(Launcher launcher) {
         return WidgetsFullSheet.getWidgetsView(launcher);
@@ -56,30 +62,21 @@
     @ScreenRecord
     @PortraitLandscape
     public void testWidgets() {
-        mLauncher.goHome();
+        loadLauncherSync();
         // Test opening widgets.
         executeOnLauncher(launcher ->
                 assertTrue("Widgets is initially opened", getWidgetsView(launcher) == null));
-        Widgets widgets = mLauncher.getWorkspace().openAllWidgets();
-        assertNotNull("openAllWidgets() returned null", widgets);
-        widgets = mLauncher.getAllWidgets();
+        assertNotNull("openAllWidgets() returned null",
+                getFromLauncher(OptionsPopupView::openWidgets));
+        WidgetsRecyclerView widgets = getFromLauncher(this::getWidgetsView);
         assertNotNull("getAllWidgets() returned null", widgets);
-        executeOnLauncher(launcher ->
-                assertTrue("Widgets is not shown", getWidgetsView(launcher).isShown()));
+        executeOnLauncher(launcher -> assertTrue("Widgets is not shown", widgets.isShown()));
         executeOnLauncher(launcher -> assertEquals("Widgets is scrolled upon opening",
                 0, getWidgetsScroll(launcher)));
 
-        // Test flinging widgets.
-        widgets.flingForward();
-        Integer flingForwardY = getFromLauncher(launcher -> getWidgetsScroll(launcher));
-        executeOnLauncher(launcher -> assertTrue("Flinging forward didn't scroll widgets",
-                flingForwardY > 0));
+        executeOnLauncher(AbstractFloatingView::closeAllOpenViews);
+        uiDevice.waitForIdle();
 
-        widgets.flingBackward();
-        executeOnLauncher(launcher -> assertTrue("Flinging backward didn't scroll widgets",
-                getWidgetsScroll(launcher) < flingForwardY));
-
-        mLauncher.goHome();
         waitForLauncherCondition("Widgets were not closed",
                 launcher -> getWidgetsView(launcher) == null);
     }
diff --git a/tests/src/com/android/launcher3/ui/workspace/ThemeIconsTest.java b/tests/src/com/android/launcher3/ui/workspace/ThemeIconsTest.java
index cfc0a6b..c623513 100644
--- a/tests/src/com/android/launcher3/ui/workspace/ThemeIconsTest.java
+++ b/tests/src/com/android/launcher3/ui/workspace/ThemeIconsTest.java
@@ -15,8 +15,6 @@
  */
 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;
@@ -29,11 +27,9 @@
 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;
@@ -49,9 +45,6 @@
 
 import org.junit.Test;
 
-import java.util.ArrayDeque;
-import java.util.Queue;
-
 /**
  * Tests for theme icon support in Launcher
  *
@@ -137,27 +130,10 @@
         } 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;
+        return (BubbleTextView) searchView(parent, v ->
+                v instanceof BubbleTextView btv
+                    && btv.getContentDescription() != null
+                        && title.equals(btv.getContentDescription().toString()));
     }
 
     private BubbleTextView verifyIconTheme(String title, ViewGroup parent, boolean isThemed) {
@@ -193,11 +169,4 @@
             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 476e497..6446592 100644
--- a/tests/src/com/android/launcher3/util/BaseLauncherActivityTest.kt
+++ b/tests/src/com/android/launcher3/util/BaseLauncherActivityTest.kt
@@ -22,6 +22,9 @@
 import android.view.KeyCharacterMap
 import android.view.KeyEvent
 import android.view.MotionEvent
+import android.view.View
+import android.view.ViewGroup
+import androidx.core.view.children
 import androidx.lifecycle.Lifecycle.State.RESUMED
 import androidx.test.core.app.ActivityScenario
 import androidx.test.core.app.ActivityScenario.ActivityAction
@@ -30,11 +33,13 @@
 import com.android.launcher3.Launcher
 import com.android.launcher3.LauncherAppState
 import com.android.launcher3.LauncherState
+import com.android.launcher3.R
 import com.android.launcher3.allapps.AllAppsStore.DEFER_UPDATES_TEST
 import com.android.launcher3.tapl.TestHelpers
 import com.android.launcher3.util.ModelTestExtensions.loadModelSync
 import com.android.launcher3.util.Wait.atMost
 import java.util.function.Function
+import java.util.function.Predicate
 import java.util.function.Supplier
 import org.junit.After
 
@@ -56,6 +61,8 @@
                     )
                     .also { currentScenario = it }
 
+    @JvmField val uiDevice = UiDevice.getInstance(getInstrumentation())
+
     @After
     fun closeCurrentActivity() {
         currentScenario?.close()
@@ -136,18 +143,43 @@
         event.recycle()
     }
 
-    fun startAppFast(packageName: String) {
-        val intent = targetContext().packageManager.getLaunchIntentForPackage(packageName)!!
+    @JvmOverloads
+    fun startAppFast(
+        packageName: String,
+        intent: Intent = targetContext().packageManager.getLaunchIntentForPackage(packageName)!!,
+    ) {
         intent.addCategory(Intent.CATEGORY_LAUNCHER)
         intent.addFlags(Intent.FLAG_ACTIVITY_NEW_TASK or Intent.FLAG_ACTIVITY_CLEAR_TASK)
         targetContext().startActivity(intent)
-        UiDevice.getInstance(getInstrumentation()).waitForIdle()
+        uiDevice.waitForIdle()
     }
 
     fun freezeAllApps() = executeOnLauncher {
         it.appsView.appsStore.enableDeferUpdates(DEFER_UPDATES_TEST)
     }
 
-    fun executeShellCommand(cmd: String) =
-        UiDevice.getInstance(getInstrumentation()).executeShellCommand(cmd)
+    fun executeShellCommand(cmd: String) = uiDevice.executeShellCommand(cmd)
+
+    fun addToWorkspace(view: View) {
+        TestUtil.runOnExecutorSync(Executors.MAIN_EXECUTOR) {
+            view.accessibilityDelegate.performAccessibilityAction(
+                view,
+                R.id.action_add_to_workspace,
+                null,
+            )
+        }
+        UiDevice.getInstance(getInstrumentation()).waitForIdle()
+    }
+
+    fun ViewGroup.searchView(filter: Predicate<View>): View? {
+        if (filter.test(this)) return this
+        for (child in children) {
+            if (filter.test(child)) return child
+            if (child is ViewGroup)
+                child.searchView(filter)?.let {
+                    return it
+                }
+        }
+        return null
+    }
 }
diff --git a/tests/src/com/android/launcher3/util/BlockingBroadcastReceiver.kt b/tests/src/com/android/launcher3/util/BlockingBroadcastReceiver.kt
new file mode 100644
index 0000000..20881d1
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/BlockingBroadcastReceiver.kt
@@ -0,0 +1,54 @@
+/*
+ * 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.util
+
+import android.content.BroadcastReceiver
+import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
+import android.os.Parcelable
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.TimeUnit.SECONDS
+
+private const val DEFAULT_BROADCAST_TIMEOUT_SECS: Long = 10
+
+/** Broadcast receiver which blocks until the result is received. */
+open class BlockingBroadcastReceiver(action: String) : BroadcastReceiver() {
+
+    val value = CompletableFuture<Intent>()
+
+    init {
+        getInstrumentation()
+            .targetContext
+            .registerReceiver(this, IntentFilter(action), Context.RECEIVER_EXPORTED)
+    }
+
+    override fun onReceive(context: Context, intent: Intent) {
+        value.complete(intent)
+    }
+
+    @Throws(InterruptedException::class)
+    fun blockingGetIntent(): Intent =
+        value.get(DEFAULT_BROADCAST_TIMEOUT_SECS, SECONDS).also {
+            getInstrumentation().targetContext.unregisterReceiver(this)
+        }
+
+    @Throws(InterruptedException::class)
+    fun blockingGetExtraIntent(): Intent? =
+        blockingGetIntent().getParcelableExtra<Parcelable>(Intent.EXTRA_INTENT) as Intent?
+}