diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index abaad20..4a0ca5c 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -20,6 +20,7 @@
 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;
 
@@ -37,7 +38,10 @@
 
 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;
@@ -46,6 +50,7 @@
 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;
@@ -81,6 +86,7 @@
     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";
 
@@ -175,6 +181,31 @@
     }
 
     /**
+     * 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 25c461b..06c56f2 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -216,8 +216,7 @@
             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 {
@@ -307,8 +306,11 @@
                 switchToAllApps();
         allApps.freeze();
         try {
-            allApps.getAppIcon(APP_NAME).dragToWorkspace();
-            mLauncher.getWorkspace().getWorkspaceAppIcon(APP_NAME).launch(getAppPackageName());
+            allApps.
+                    getAppIcon(APP_NAME).
+                    dragToWorkspace().
+                    getWorkspaceAppIcon(APP_NAME).
+                    launch(getAppPackageName());
         } finally {
             allApps.unfreeze();
         }
@@ -337,8 +339,13 @@
                     getMenuItem(0);
             final String shortcutName = menuItem.getText();
 
-            menuItem.dragToWorkspace();
-            mLauncher.getWorkspace().getWorkspaceAppIcon(shortcutName).launch(getAppPackageName());
+            // 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());
         } finally {
             allApps.unfreeze();
         }
diff --git a/tests/src/com/android/launcher3/ui/TestViewHelpers.java b/tests/src/com/android/launcher3/ui/TestViewHelpers.java
index d0df664..d13d319 100644
--- a/tests/src/com/android/launcher3/ui/TestViewHelpers.java
+++ b/tests/src/com/android/launcher3/ui/TestViewHelpers.java
@@ -18,15 +18,28 @@
 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;
@@ -37,6 +50,21 @@
 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.
      *
@@ -63,6 +91,92 @@
         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 dc72bda..5eb5f19 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -27,18 +27,20 @@
 
 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;
@@ -69,6 +71,7 @@
     }
 
     @Test
+    // Convert test to TAPL b/131116002
     public void testWidgetConfig() throws Throwable {
         runTest(false, true);
     }
@@ -80,6 +83,7 @@
     }
 
     @Test
+    // Convert test to TAPL b/131116002
     public void testConfigCancelled() throws Throwable {
         runTest(false, false);
     }
@@ -100,13 +104,15 @@
         clearHomescreen();
         mActivityMonitor.startLauncher();
 
-        final Widgets widgets = mLauncher.getWorkspace().openAllWidgets();
+        // Open widget tray and wait for load complete.
+        final UiObject2 widgetContainer = TestViewHelpers.openWidgetsTray();
+        Wait.atMost(null, Condition.minChildCount(widgetContainer, 2), DEFAULT_UI_TIMEOUT);
 
         // Drag widget to homescreen
         WidgetConfigStartupMonitor monitor = new WidgetConfigStartupMonitor();
-        widgets.
-                getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager())).
-                dragToWorkspace();
+        UiObject2 widget = scrollAndFind(widgetContainer, By.clazz(WidgetCell.class)
+                .hasDescendant(By.text(mWidgetInfo.getLabel(mTargetContext.getPackageManager()))));
+        TestViewHelpers.dragToWorkspace(widget, false);
         // 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 4529a80..0061568 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
@@ -19,12 +19,20 @@
 
 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;
@@ -53,6 +61,7 @@
         performTest();
     }
 
+    // Convert to TAPL b/131116002
     private void performTest() throws Throwable {
         clearHomescreen();
         mActivityMonitor.startLauncher();
@@ -60,15 +69,22 @@
         final LauncherAppWidgetProviderInfo widgetInfo =
                 TestViewHelpers.findWidgetProvider(this, false /* hasConfigureScreen */);
 
-        mLauncher.
-                getWorkspace().
-                openAllWidgets().
-                getWidget(widgetInfo.getLabel(mTargetContext.getPackageManager())).
-                dragToWorkspace();
+        // Open widget tray and wait for load complete.
+        final UiObject2 widgetContainer = TestViewHelpers.openWidgetsTray();
+        Wait.atMost(null, Condition.minChildCount(widgetContainer, 2), DEFAULT_UI_TIMEOUT);
 
-        assertTrue(mActivityMonitor.itemExists(
-                (info, view) -> info instanceof LauncherAppWidgetInfo &&
+        // 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 &&
                         ((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 04b8019..d4bdafa 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 void dragToWorkspace() {
+    public Workspace dragToWorkspace() {
         final Point launchableCenter = getObject().getVisibleCenter();
         final Point displaySize = mLauncher.getRealDisplaySize();
         final int width = displaySize.x / 2;
@@ -80,6 +80,10 @@
                                 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 1641d70..a7e6336 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(By.copy(selector).pkg(getLauncherPackageName()));
+        return waitForObjectBySelector(selector.pkg(getLauncherPackageName()));
     }
 
     @NonNull
     UiObject2 tryWaitForLauncherObject(BySelector selector, long timeout) {
-        return tryWaitForObjectBySelector(By.copy(selector).pkg(getLauncherPackageName()), timeout);
+        return tryWaitForObjectBySelector(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 1b6d8c4..128789d 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widget.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widget.java
@@ -18,16 +18,7 @@
 
 import androidx.test.uiautomator.UiObject2;
 
-/**
- * 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";
+public class Widget {
+    Widget(LauncherInstrumentation launcher, UiObject2 widget) {
     }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index cbc3a81..94003be 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -16,12 +16,6 @@
 
 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;
 
@@ -81,28 +75,4 @@
     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);
-        }
-    }
 }