diff --git a/go/quickstep/src/com/android/launcher3/uioverrides/RecentsUiFactory.java b/go/quickstep/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
index 4582243..cbc77d2 100644
--- a/go/quickstep/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
+++ b/go/quickstep/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
@@ -89,4 +89,6 @@
     public static RotationMode getRotationMode(DeviceProfile dp) {
         return RotationMode.NORMAL;
     }
+
+    public static void clearSwipeSharedState(boolean finishAnimation) {}
 }
diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/RecentsUiFactory.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
index 8d5ac50..6ecf1c1 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
@@ -46,6 +46,7 @@
 import com.android.launcher3.util.UiThreadHelper.AsyncCommand;
 import com.android.quickstep.SysUINavigationMode;
 import com.android.quickstep.SysUINavigationMode.Mode;
+import com.android.quickstep.TouchInteractionService;
 import com.android.quickstep.views.RecentsView;
 import com.android.systemui.shared.system.WindowManagerWrapper;
 
@@ -184,6 +185,13 @@
     }
 
     /**
+     * Clears the swipe shared state for the current swipe gesture.
+     */
+    public static void clearSwipeSharedState(boolean finishAnimation) {
+        TouchInteractionService.getSwipeSharedState().clearAllState(finishAnimation);
+    }
+
+    /**
      * Recents logic that triggers when launcher state changes or launcher activity stops/resumes.
      *
      * @param launcher the launcher activity
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/SwipeSharedState.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/SwipeSharedState.java
index 6689ce3..c55f656 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/SwipeSharedState.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/SwipeSharedState.java
@@ -72,13 +72,15 @@
         mLastAnimationRunning = false;
     }
 
-    private void clearListenerState() {
+    private void clearListenerState(boolean finishAnimation) {
         if (mRecentsAnimationListener != null) {
             mRecentsAnimationListener.removeListener(this);
             mRecentsAnimationListener.cancelListener();
             if (mLastAnimationRunning && mLastAnimationTarget != null) {
                 Utilities.postAsyncCallback(MAIN_THREAD_EXECUTOR.getHandler(),
-                        mLastAnimationTarget::cancelAnimation);
+                        finishAnimation
+                                ? mLastAnimationTarget::finishAnimation
+                                : mLastAnimationTarget::cancelAnimation);
                 mLastAnimationTarget = null;
             }
         }
@@ -106,7 +108,7 @@
             }
         }
 
-        clearListenerState();
+        clearListenerState(false /* finishAnimation */);
         boolean shouldMinimiseSplitScreen = mOverviewComponentObserver == null ? false
                 : mOverviewComponentObserver.getActivityControlHelper().shouldMinimizeSplitScreen();
         mRecentsAnimationListener = new RecentsAnimationListenerSet(
@@ -138,8 +140,8 @@
         mLastAnimationTarget = mLastAnimationTarget.cloneWithoutTargets();
     }
 
-    public void clearAllState() {
-        clearListenerState();
+    public void clearAllState(boolean finishAnimation) {
+        clearListenerState(finishAnimation);
         canGestureBeContinued = false;
         recentsAnimationFinishInterrupted = false;
         nextRunningTaskId = -1;
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/TouchInteractionService.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/TouchInteractionService.java
index 8f08f0d..53da0f9 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/TouchInteractionService.java
@@ -225,14 +225,18 @@
     };
 
     private static boolean sConnected = false;
+    private static final SwipeSharedState sSwipeSharedState = new SwipeSharedState();
 
     public static boolean isConnected() {
         return sConnected;
     }
 
-    private final SwipeSharedState mSwipeSharedState = new SwipeSharedState();
+    public static SwipeSharedState getSwipeSharedState() {
+        return sSwipeSharedState;
+    }
+
     private final InputConsumer mResetGestureInputConsumer =
-            new ResetGestureInputConsumer(mSwipeSharedState);
+            new ResetGestureInputConsumer(sSwipeSharedState);
 
     private ActivityManagerWrapper mAM;
     private RecentsModel mRecentsModel;
@@ -436,7 +440,7 @@
         mInputConsumer = InputConsumerController.getRecentsAnimationInputConsumer();
         mIsUserUnlocked = true;
 
-        mSwipeSharedState.setOverviewComponentObserver(mOverviewComponentObserver);
+        sSwipeSharedState.setOverviewComponentObserver(mOverviewComponentObserver);
         mInputConsumer.registerInputConsumer();
         onSystemUiProxySet();
         onSystemUiFlagsChanged();
@@ -589,7 +593,7 @@
     private InputConsumer newBaseConsumer(boolean useSharedState, MotionEvent event) {
         final RunningTaskInfo runningTaskInfo = mAM.getRunningTask(0);
         if (!useSharedState) {
-            mSwipeSharedState.clearAllState();
+            sSwipeSharedState.clearAllState(false /* finishAnimation */);
         }
         if ((mSystemUiStateFlags & SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING_OCCLUDED) != 0) {
             // This handles apps showing over the lockscreen (e.g. camera)
@@ -599,16 +603,16 @@
         final ActivityControlHelper activityControl =
                 mOverviewComponentObserver.getActivityControlHelper();
 
-        if (runningTaskInfo == null && !mSwipeSharedState.goingToLauncher
-                && !mSwipeSharedState.recentsAnimationFinishInterrupted) {
+        if (runningTaskInfo == null && !sSwipeSharedState.goingToLauncher
+                && !sSwipeSharedState.recentsAnimationFinishInterrupted) {
             return mResetGestureInputConsumer;
-        } else if (mSwipeSharedState.recentsAnimationFinishInterrupted) {
+        } else if (sSwipeSharedState.recentsAnimationFinishInterrupted) {
             // If the finish animation was interrupted, then continue using the other activity input
             // consumer but with the next task as the running task
             RunningTaskInfo info = new ActivityManager.RunningTaskInfo();
-            info.id = mSwipeSharedState.nextRunningTaskId;
+            info.id = sSwipeSharedState.nextRunningTaskId;
             return createOtherActivityInputConsumer(event, info);
-        } else if (mSwipeSharedState.goingToLauncher || activityControl.isResumed()) {
+        } else if (sSwipeSharedState.goingToLauncher || activityControl.isResumed()) {
             return createOverviewInputConsumer(event);
         } else if (ENABLE_QUICKSTEP_LIVE_TILE.get() && activityControl.isInLiveTileMode()) {
             return createOverviewInputConsumer(event);
@@ -617,7 +621,7 @@
             return mResetGestureInputConsumer;
         } else if (mMode == Mode.NO_BUTTON && !mOverviewComponentObserver.isHomeAndOverviewSame()) {
             return new FallbackNoButtonInputConsumer(this, activityControl,
-                    mInputMonitorCompat, mSwipeSharedState, mSwipeTouchRegion,
+                    mInputMonitorCompat, sSwipeSharedState, mSwipeTouchRegion,
                     mOverviewComponentObserver, disableHorizontalSwipe(event), runningTaskInfo);
         } else {
             return createOtherActivityInputConsumer(event, runningTaskInfo);
@@ -640,13 +644,13 @@
         return new OtherActivityInputConsumer(this, runningTaskInfo, mRecentsModel,
                 mOverviewComponentObserver.getOverviewIntent(), activityControl,
                 shouldDefer, mOverviewCallbacks, mInputConsumer, this::onConsumerInactive,
-                mSwipeSharedState, mInputMonitorCompat, mSwipeTouchRegion,
+                sSwipeSharedState, mInputMonitorCompat, mSwipeTouchRegion,
                 disableHorizontalSwipe(event));
     }
 
     private InputConsumer createDeviceLockedInputConsumer(RunningTaskInfo taskInfo) {
         if (mMode == Mode.NO_BUTTON && taskInfo != null) {
-            return new DeviceLockedInputConsumer(this, mSwipeSharedState, mInputMonitorCompat,
+            return new DeviceLockedInputConsumer(this, sSwipeSharedState, mInputMonitorCompat,
                     mSwipeTouchRegion, taskInfo.taskId);
         } else {
             return mResetGestureInputConsumer;
@@ -661,7 +665,7 @@
             return mResetGestureInputConsumer;
         }
 
-        if (activity.getRootView().hasWindowFocus() || mSwipeSharedState.goingToLauncher) {
+        if (activity.getRootView().hasWindowFocus() || sSwipeSharedState.goingToLauncher) {
             return new OverviewInputConsumer(activity, mInputMonitorCompat,
                     false /* startingInActivityBounds */);
         } else {
@@ -708,7 +712,7 @@
                     + mOverviewComponentObserver.getActivityControlHelper().isResumed());
             pw.println("  useSharedState=" + mConsumer.useSharedSwipeState());
             if (mConsumer.useSharedSwipeState()) {
-                mSwipeSharedState.dump("    ", pw);
+                sSwipeSharedState.dump("    ", pw);
             }
             pw.println("  mConsumer=" + mConsumer.getName());
             pw.println("FeatureFlags:");
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/ResetGestureInputConsumer.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/ResetGestureInputConsumer.java
index 56cba21..8eede81 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/ResetGestureInputConsumer.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/ResetGestureInputConsumer.java
@@ -39,7 +39,7 @@
     public void onMotionEvent(MotionEvent ev) {
         if (ev.getAction() == MotionEvent.ACTION_DOWN
                 && mSwipeSharedState.getActiveListener() != null) {
-            mSwipeSharedState.clearAllState();
+            mSwipeSharedState.clearAllState(false /* finishAnimation */);
         }
     }
 }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/SwipeAnimationTargetSet.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/SwipeAnimationTargetSet.java
index df9efa2..381c27a 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/SwipeAnimationTargetSet.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/SwipeAnimationTargetSet.java
@@ -106,6 +106,10 @@
         finishController(false /* toRecents */, null, false /* sendUserLeaveHint */);
     }
 
+    public void finishAnimation() {
+        finishController(true /* toRecents */, null, false /* sendUserLeaveHint */);
+    }
+
     public interface SwipeAnimationListener {
 
         void onRecentsAnimationStart(SwipeAnimationTargetSet targetSet);
diff --git a/quickstep/src/com/android/launcher3/QuickstepAppTransitionManagerImpl.java b/quickstep/src/com/android/launcher3/QuickstepAppTransitionManagerImpl.java
index 8643160..44324cb 100644
--- a/quickstep/src/com/android/launcher3/QuickstepAppTransitionManagerImpl.java
+++ b/quickstep/src/com/android/launcher3/QuickstepAppTransitionManagerImpl.java
@@ -765,7 +765,7 @@
                 LauncherAnimationRunner.AnimationResult result) {
             if (!mLauncher.hasBeenResumed()) {
                 // If launcher is not resumed, wait until new async-frame after resume
-                mLauncher.setOnResumeCallback(() ->
+                mLauncher.addOnResumeCallback(() ->
                         postAsyncCallback(mHandler, () ->
                                 onCreateAnimation(targetCompats, result)));
                 return;
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index ad2783e..d9af4da 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -244,7 +244,7 @@
 
     @Thunk boolean mWorkspaceLoading = true;
 
-    private OnResumeCallback mOnResumeCallback;
+    private ArrayList<OnResumeCallback> mOnResumeCallbacks = new ArrayList<>();
 
     private ViewOnDrawExecutor mPendingExecutor;
 
@@ -869,6 +869,7 @@
     @Override
     protected void onStop() {
         super.onStop();
+
         if (mLauncherCallbacks != null) {
             mLauncherCallbacks.onStop();
         }
@@ -951,7 +952,10 @@
         mHandler.removeCallbacks(mHandleDeferredResume);
         Utilities.postAsyncCallback(mHandler, mHandleDeferredResume);
 
-        setOnResumeCallback(null);
+        for (OnResumeCallback cb : mOnResumeCallbacks) {
+            cb.onLauncherResume();
+        }
+        mOnResumeCallbacks.clear();
 
         if (mLauncherCallbacks != null) {
             mLauncherCallbacks.onResume();
@@ -1805,6 +1809,16 @@
             android.util.Log.d(TestProtocol.NO_START_TAG,
                     "startActivitySafely outer");
         }
+
+        if (!hasBeenResumed()) {
+            // Workaround an issue where the WM launch animation is clobbered when finishing the
+            // recents animation into launcher. Defer launching the activity until Launcher is
+            // next resumed.
+            addOnResumeCallback(() -> startActivitySafely(v, intent, item, sourceContainer));
+            UiFactory.clearSwipeSharedState(true /* finishAnimation */);
+            return true;
+        }
+
         boolean success = super.startActivitySafely(v, intent, item, sourceContainer);
         if (success && v instanceof BubbleTextView) {
             // This is set to the view that launched the activity that navigated the user away
@@ -1813,7 +1827,7 @@
             // state when we return to launcher.
             BubbleTextView btv = (BubbleTextView) v;
             btv.setStayPressed(true);
-            setOnResumeCallback(btv);
+            addOnResumeCallback(btv);
         }
         return success;
     }
@@ -1861,11 +1875,8 @@
         return result;
     }
 
-    public void setOnResumeCallback(OnResumeCallback callback) {
-        if (mOnResumeCallback != null) {
-            mOnResumeCallback.onLauncherResume();
-        }
-        mOnResumeCallback = callback;
+    public void addOnResumeCallback(OnResumeCallback callback) {
+        mOnResumeCallbacks.add(callback);
     }
 
     /**
diff --git a/src/com/android/launcher3/SecondaryDropTarget.java b/src/com/android/launcher3/SecondaryDropTarget.java
index 0cf6e44..55cb6f2 100644
--- a/src/com/android/launcher3/SecondaryDropTarget.java
+++ b/src/com/android/launcher3/SecondaryDropTarget.java
@@ -179,7 +179,7 @@
             DeferredOnComplete deferred = (DeferredOnComplete) d.dragSource;
             if (target != null) {
                 deferred.mPackageName = target.getPackageName();
-                mLauncher.setOnResumeCallback(deferred);
+                mLauncher.addOnResumeCallback(deferred);
             } else {
                 deferred.sendFailure();
             }
diff --git a/src_ui_overrides/com/android/launcher3/uioverrides/UiFactory.java b/src_ui_overrides/com/android/launcher3/uioverrides/UiFactory.java
index 17ff66e..5cc64dc 100644
--- a/src_ui_overrides/com/android/launcher3/uioverrides/UiFactory.java
+++ b/src_ui_overrides/com/android/launcher3/uioverrides/UiFactory.java
@@ -93,4 +93,6 @@
 
     public static void resetPendingActivityResults(Launcher launcher, int requestCode) { }
 
+    public static void clearSwipeSharedState(boolean finishAnimation) {}
+
 }
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 4a0ca5c..abaad20 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -20,7 +20,6 @@
 import static com.android.launcher3.ui.TaplTestsLauncher3.getAppPackageName;
 
 import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
 
 import static java.lang.System.exit;
 
@@ -38,10 +37,7 @@
 
 import androidx.test.InstrumentationRegistry;
 import androidx.test.uiautomator.By;
-import androidx.test.uiautomator.BySelector;
-import androidx.test.uiautomator.Direction;
 import androidx.test.uiautomator.UiDevice;
-import androidx.test.uiautomator.UiObject2;
 import androidx.test.uiautomator.Until;
 
 import com.android.launcher3.Launcher;
@@ -50,7 +46,6 @@
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherState;
 import com.android.launcher3.MainThreadExecutor;
-import com.android.launcher3.ResourceUtils;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.compat.LauncherAppsCompat;
 import com.android.launcher3.model.AppLaunchTracker;
@@ -86,7 +81,6 @@
     public static final long DEFAULT_ACTIVITY_TIMEOUT = TimeUnit.SECONDS.toMillis(10);
     public static final long DEFAULT_BROADCAST_TIMEOUT_SECS = 5;
 
-    public static final long SHORT_UI_TIMEOUT = 300;
     public static final long DEFAULT_UI_TIMEOUT = 10000;
     private static final String TAG = "AbstractLauncherUiTest";
 
@@ -181,31 +175,6 @@
     }
 
     /**
-     * Scrolls the {@param container} until it finds an object matching {@param condition}.
-     *
-     * @return the matching object.
-     */
-    protected UiObject2 scrollAndFind(UiObject2 container, BySelector condition) {
-        final int margin = ResourceUtils.getNavbarSize(
-                ResourceUtils.NAVBAR_BOTTOM_GESTURE_SIZE, mLauncher.getResources()) + 1;
-        container.setGestureMargins(0, 0, 0, margin);
-
-        int i = 0;
-        for (; ; ) {
-            // findObject can only execute after spring settles.
-            mDevice.wait(Until.findObject(condition), SHORT_UI_TIMEOUT);
-            UiObject2 widget = container.findObject(condition);
-            if (widget != null && widget.getVisibleBounds().intersects(
-                    0, 0, mDevice.getDisplayWidth(),
-                    mDevice.getDisplayHeight() - margin)) {
-                return widget;
-            }
-            if (++i > 40) fail("Too many attempts");
-            container.scroll(Direction.DOWN, 1f);
-        }
-    }
-
-    /**
      * Removes all icons from homescreen and hotseat.
      */
     public void clearHomescreen() throws Throwable {
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index 06c56f2..25c461b 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -216,7 +216,8 @@
             final AppIcon app = allApps.getAppIcon("TestActivity7");
             assertNotNull("AppIcon.launch returned null", app.launch(getAppPackageName()));
             test.executeOnLauncher(launcher -> assertTrue(
-                    "Launcher activity is the top activity; expecting another activity to be the top "
+                    "Launcher activity is the top activity; expecting another activity to be the "
+                            + "top "
                             + "one",
                     test.isInBackground(launcher)));
         } finally {
@@ -306,11 +307,8 @@
                 switchToAllApps();
         allApps.freeze();
         try {
-            allApps.
-                    getAppIcon(APP_NAME).
-                    dragToWorkspace().
-                    getWorkspaceAppIcon(APP_NAME).
-                    launch(getAppPackageName());
+            allApps.getAppIcon(APP_NAME).dragToWorkspace();
+            mLauncher.getWorkspace().getWorkspaceAppIcon(APP_NAME).launch(getAppPackageName());
         } finally {
             allApps.unfreeze();
         }
@@ -339,13 +337,8 @@
                     getMenuItem(0);
             final String shortcutName = menuItem.getText();
 
-            // 4. Drag the first shortcut to the home screen.
-            // 5. Verify that the shortcut works on home screen
-            //    (the app opens and has the same text as the shortcut).
-            menuItem.
-                    dragToWorkspace().
-                    getWorkspaceAppIcon(shortcutName).
-                    launch(getAppPackageName());
+            menuItem.dragToWorkspace();
+            mLauncher.getWorkspace().getWorkspaceAppIcon(shortcutName).launch(getAppPackageName());
         } finally {
             allApps.unfreeze();
         }
diff --git a/tests/src/com/android/launcher3/ui/TestViewHelpers.java b/tests/src/com/android/launcher3/ui/TestViewHelpers.java
index d13d319..d0df664 100644
--- a/tests/src/com/android/launcher3/ui/TestViewHelpers.java
+++ b/tests/src/com/android/launcher3/ui/TestViewHelpers.java
@@ -18,28 +18,15 @@
 import static androidx.test.InstrumentationRegistry.getInstrumentation;
 import static androidx.test.InstrumentationRegistry.getTargetContext;
 
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertNull;
-
 import android.content.ComponentName;
-import android.content.Context;
-import android.graphics.Point;
 import android.os.Process;
-import android.os.SystemClock;
 import android.util.Log;
-import android.view.KeyEvent;
-import android.view.MotionEvent;
 import android.view.View;
 import android.view.ViewGroup;
 
-import androidx.test.uiautomator.By;
-import androidx.test.uiautomator.BySelector;
 import androidx.test.uiautomator.UiDevice;
-import androidx.test.uiautomator.UiObject2;
-import androidx.test.uiautomator.Until;
 
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
-import com.android.launcher3.R;
 import com.android.launcher3.compat.AppWidgetManagerCompat;
 import com.android.launcher3.testcomponent.AppWidgetNoConfig;
 import com.android.launcher3.testcomponent.AppWidgetWithConfig;
@@ -50,21 +37,6 @@
 public class TestViewHelpers {
     private static final String TAG = "TestViewHelpers";
 
-    private static UiDevice getDevice() {
-        return UiDevice.getInstance(getInstrumentation());
-    }
-
-    public static UiObject2 findViewById(int id) {
-        return getDevice().wait(Until.findObject(getSelectorForId(id)),
-                AbstractLauncherUiTest.DEFAULT_UI_TIMEOUT);
-    }
-
-    public static BySelector getSelectorForId(int id) {
-        final Context targetContext = getTargetContext();
-        String name = targetContext.getResources().getResourceEntryName(id);
-        return By.res(targetContext.getPackageName(), name);
-    }
-
     /**
      * Finds a widget provider which can fit on the home screen.
      *
@@ -91,92 +63,6 @@
         return info;
     }
 
-    /**
-     * Drags an icon to the center of homescreen.
-     *
-     * @param icon object that is either app icon or shortcut icon
-     */
-    public static void dragToWorkspace(UiObject2 icon, boolean expectedToShowShortcuts) {
-        Point center = icon.getVisibleCenter();
-
-        // Action Down
-        final long downTime = SystemClock.uptimeMillis();
-        sendPointer(downTime, MotionEvent.ACTION_DOWN, center);
-
-        UiObject2 dragLayer = findViewById(R.id.drag_layer);
-
-        if (expectedToShowShortcuts) {
-            // Make sure shortcuts show up, and then move a bit to hide them.
-            assertNotNull(findViewById(R.id.deep_shortcuts_container));
-
-            Point moveLocation = new Point(center);
-            int distanceToMove =
-                    getTargetContext().getResources().getDimensionPixelSize(
-                            R.dimen.deep_shortcuts_start_drag_threshold) + 50;
-            if (moveLocation.y - distanceToMove >= dragLayer.getVisibleBounds().top) {
-                moveLocation.y -= distanceToMove;
-            } else {
-                moveLocation.y += distanceToMove;
-            }
-            movePointer(downTime, center, moveLocation);
-
-            assertNull(findViewById(R.id.deep_shortcuts_container));
-        }
-
-        // Wait until Remove/Delete target is visible
-        assertNotNull(findViewById(R.id.delete_target_text));
-
-        Point moveLocation = dragLayer.getVisibleCenter();
-
-        // Move to center
-        movePointer(downTime, center, moveLocation);
-        sendPointer(downTime, MotionEvent.ACTION_UP, moveLocation);
-
-        // Wait until remove target is gone.
-        getDevice().wait(Until.gone(getSelectorForId(R.id.delete_target_text)),
-                AbstractLauncherUiTest.DEFAULT_UI_TIMEOUT);
-    }
-
-    private static void movePointer(long downTime, Point from, Point to) {
-        while (!from.equals(to)) {
-            try {
-                Thread.sleep(20);
-            } catch (InterruptedException e) {
-                e.printStackTrace();
-            }
-            from.x = getNextMoveValue(to.x, from.x);
-            from.y = getNextMoveValue(to.y, from.y);
-            sendPointer(downTime, MotionEvent.ACTION_MOVE, from);
-        }
-    }
-
-    private static int getNextMoveValue(int targetValue, int oldValue) {
-        if (targetValue - oldValue > 10) {
-            return oldValue + 10;
-        } else if (targetValue - oldValue < -10) {
-            return oldValue - 10;
-        } else {
-            return targetValue;
-        }
-    }
-
-    public static void sendPointer(long downTime, int action, Point point) {
-        MotionEvent event = MotionEvent.obtain(downTime,
-                SystemClock.uptimeMillis(), action, point.x, point.y, 0);
-        getInstrumentation().sendPointerSync(event);
-        event.recycle();
-    }
-
-    /**
-     * Opens widget tray and returns the recycler view.
-     */
-    public static UiObject2 openWidgetsTray() {
-        final UiDevice device = getDevice();
-        device.pressKeyCode(KeyEvent.KEYCODE_W, KeyEvent.META_CTRL_ON);
-        device.waitForIdle();
-        return findViewById(R.id.widgets_list_view);
-    }
-
     public static View findChildView(ViewGroup parent, Function<View, Boolean> condition) {
         for (int i = 0; i < parent.getChildCount(); ++i) {
             final View child = parent.getChildAt(i);
diff --git a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
index 5eb5f19..dc72bda 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -27,20 +27,18 @@
 
 import androidx.test.filters.LargeTest;
 import androidx.test.runner.AndroidJUnit4;
-import androidx.test.uiautomator.By;
-import androidx.test.uiautomator.UiObject2;
 
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.LauncherAppWidgetInfo;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.Workspace;
+import com.android.launcher3.tapl.Widgets;
 import com.android.launcher3.testcomponent.WidgetConfigActivity;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.TestViewHelpers;
 import com.android.launcher3.util.Condition;
 import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.ShellCommandRule;
-import com.android.launcher3.widget.WidgetCell;
 
 import org.junit.Before;
 import org.junit.Ignore;
@@ -71,7 +69,6 @@
     }
 
     @Test
-    // Convert test to TAPL b/131116002
     public void testWidgetConfig() throws Throwable {
         runTest(false, true);
     }
@@ -83,7 +80,6 @@
     }
 
     @Test
-    // Convert test to TAPL b/131116002
     public void testConfigCancelled() throws Throwable {
         runTest(false, false);
     }
@@ -104,15 +100,13 @@
         clearHomescreen();
         mActivityMonitor.startLauncher();
 
-        // Open widget tray and wait for load complete.
-        final UiObject2 widgetContainer = TestViewHelpers.openWidgetsTray();
-        Wait.atMost(null, Condition.minChildCount(widgetContainer, 2), DEFAULT_UI_TIMEOUT);
+        final Widgets widgets = mLauncher.getWorkspace().openAllWidgets();
 
         // Drag widget to homescreen
         WidgetConfigStartupMonitor monitor = new WidgetConfigStartupMonitor();
-        UiObject2 widget = scrollAndFind(widgetContainer, By.clazz(WidgetCell.class)
-                .hasDescendant(By.text(mWidgetInfo.getLabel(mTargetContext.getPackageManager()))));
-        TestViewHelpers.dragToWorkspace(widget, false);
+        widgets.
+                getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager())).
+                dragToWorkspace();
         // Widget id for which the config activity was opened
         mWidgetId = monitor.getWidgetId();
 
diff --git a/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
index 0061568..4529a80 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
@@ -19,20 +19,12 @@
 
 import androidx.test.filters.LargeTest;
 import androidx.test.runner.AndroidJUnit4;
-import androidx.test.uiautomator.By;
-import androidx.test.uiautomator.UiObject2;
-import android.view.View;
 
-import com.android.launcher3.ItemInfo;
 import com.android.launcher3.LauncherAppWidgetInfo;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
-import com.android.launcher3.Workspace.ItemOperator;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.TestViewHelpers;
-import com.android.launcher3.util.Condition;
-import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.ShellCommandRule;
-import com.android.launcher3.widget.WidgetCell;
 
 import org.junit.Ignore;
 import org.junit.Rule;
@@ -61,7 +53,6 @@
         performTest();
     }
 
-    // Convert to TAPL b/131116002
     private void performTest() throws Throwable {
         clearHomescreen();
         mActivityMonitor.startLauncher();
@@ -69,22 +60,15 @@
         final LauncherAppWidgetProviderInfo widgetInfo =
                 TestViewHelpers.findWidgetProvider(this, false /* hasConfigureScreen */);
 
-        // Open widget tray and wait for load complete.
-        final UiObject2 widgetContainer = TestViewHelpers.openWidgetsTray();
-        Wait.atMost(null, Condition.minChildCount(widgetContainer, 2), DEFAULT_UI_TIMEOUT);
+        mLauncher.
+                getWorkspace().
+                openAllWidgets().
+                getWidget(widgetInfo.getLabel(mTargetContext.getPackageManager())).
+                dragToWorkspace();
 
-        // Drag widget to homescreen
-        UiObject2 widget = scrollAndFind(widgetContainer, By.clazz(WidgetCell.class)
-                .hasDescendant(By.text(widgetInfo.getLabel(mTargetContext.getPackageManager()))));
-        TestViewHelpers.dragToWorkspace(widget, false);
-
-        assertTrue(mActivityMonitor.itemExists(new ItemOperator() {
-            @Override
-            public boolean evaluate(ItemInfo info, View view) {
-                return info instanceof LauncherAppWidgetInfo &&
+        assertTrue(mActivityMonitor.itemExists(
+                (info, view) -> info instanceof LauncherAppWidgetInfo &&
                         ((LauncherAppWidgetInfo) info).providerName.getClassName().equals(
-                                widgetInfo.provider.getClassName());
-            }
-        }).call());
+                                widgetInfo.provider.getClassName())).call());
     }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/Launchable.java b/tests/tapl/com/android/launcher3/tapl/Launchable.java
index d4bdafa..04b8019 100644
--- a/tests/tapl/com/android/launcher3/tapl/Launchable.java
+++ b/tests/tapl/com/android/launcher3/tapl/Launchable.java
@@ -68,7 +68,7 @@
     /**
      * Drags an object to the center of homescreen.
      */
-    public Workspace dragToWorkspace() {
+    public void dragToWorkspace() {
         final Point launchableCenter = getObject().getVisibleCenter();
         final Point displaySize = mLauncher.getRealDisplaySize();
         final int width = displaySize.x / 2;
@@ -80,10 +80,6 @@
                                 launchableCenter.x - width / 2 : launchableCenter.x + width / 2,
                         displaySize.y / 2),
                 getLongPressIndicator());
-        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
-                "dragged launchable to workspace")) {
-            return new Workspace(mLauncher);
-        }
     }
 
     protected abstract String getLongPressIndicator();
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index a7e6336..1641d70 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -614,12 +614,12 @@
 
     @NonNull
     UiObject2 waitForLauncherObject(BySelector selector) {
-        return waitForObjectBySelector(selector.pkg(getLauncherPackageName()));
+        return waitForObjectBySelector(By.copy(selector).pkg(getLauncherPackageName()));
     }
 
     @NonNull
     UiObject2 tryWaitForLauncherObject(BySelector selector, long timeout) {
-        return tryWaitForObjectBySelector(selector.pkg(getLauncherPackageName()), timeout);
+        return tryWaitForObjectBySelector(By.copy(selector).pkg(getLauncherPackageName()), timeout);
     }
 
     @NonNull
diff --git a/tests/tapl/com/android/launcher3/tapl/Widget.java b/tests/tapl/com/android/launcher3/tapl/Widget.java
index 128789d..1b6d8c4 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widget.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widget.java
@@ -18,7 +18,16 @@
 
 import androidx.test.uiautomator.UiObject2;
 
-public class Widget {
-    Widget(LauncherInstrumentation launcher, UiObject2 widget) {
+/**
+ * Widget in workspace or a widget list.
+ */
+public final class Widget extends Launchable {
+    Widget(LauncherInstrumentation launcher, UiObject2 icon) {
+        super(launcher, icon);
+    }
+
+    @Override
+    protected String getLongPressIndicator() {
+        return "drop_target_bar";
     }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index 94003be..cbc3a81 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -16,6 +16,12 @@
 
 package com.android.launcher3.tapl;
 
+import static org.junit.Assert.fail;
+
+import android.graphics.Point;
+
+import androidx.test.uiautomator.By;
+import androidx.test.uiautomator.BySelector;
 import androidx.test.uiautomator.Direction;
 import androidx.test.uiautomator.UiObject2;
 
@@ -75,4 +81,28 @@
     protected LauncherInstrumentation.ContainerType getContainerType() {
         return LauncherInstrumentation.ContainerType.WIDGETS;
     }
+
+    public Widget getWidget(String label) {
+        final int margin = ResourceUtils.getNavbarSize(
+                ResourceUtils.NAVBAR_BOTTOM_GESTURE_SIZE, mLauncher.getResources()) + 1;
+        final UiObject2 widgetsContainer = verifyActiveContainer();
+        widgetsContainer.setGestureMargins(0, 0, 0, margin);
+
+        final Point displaySize = mLauncher.getRealDisplaySize();
+
+        int i = 0;
+        final BySelector selector = By.
+                clazz("com.android.launcher3.widget.WidgetCell").
+                hasDescendant(By.text(label));
+
+        for (; ; ) {
+            final UiObject2 widget = mLauncher.tryWaitForLauncherObject(selector, 300);
+            if (widget != null && widget.getVisibleBounds().intersects(
+                    0, 0, displaySize.x, displaySize.y - margin)) {
+                return new Widget(mLauncher, widget);
+            }
+            if (++i > 40) fail("Too many attempts");
+            widgetsContainer.scroll(Direction.DOWN, 1f);
+        }
+    }
 }