diff --git a/tests/tapl/com/android/launcher3/tapl/AddToHomeScreenPrompt.java b/tests/tapl/com/android/launcher3/tapl/AddToHomeScreenPrompt.java
index 5daac394b9..e1fde3ba8b 100644
--- a/tests/tapl/com/android/launcher3/tapl/AddToHomeScreenPrompt.java
+++ b/tests/tapl/com/android/launcher3/tapl/AddToHomeScreenPrompt.java
@@ -19,6 +19,7 @@ package com.android.launcher3.tapl;
 import static java.util.regex.Pattern.CASE_INSENSITIVE;
 
 import androidx.test.uiautomator.By;
+import androidx.test.uiautomator.BySelector;
 import androidx.test.uiautomator.UiObject2;
 
 import com.android.launcher3.testing.TestProtocol;
@@ -33,11 +34,14 @@ public class AddToHomeScreenPrompt {
 
     AddToHomeScreenPrompt(LauncherInstrumentation launcher) {
         mLauncher = launcher;
-        mWidgetCell = launcher.waitForLauncherObject(By.clazz(
-                "com.android.launcher3.widget.WidgetCell"));
+        mWidgetCell = launcher.waitForLauncherObject(getSelector());
         mLauncher.assertNotNull("Can't find widget cell object", mWidgetCell);
     }
 
+    private static BySelector getSelector() {
+        return By.clazz("com.android.launcher3.widget.WidgetCell");
+    }
+
     public void addAutomatically() {
         try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck()) {
             if (mLauncher.getNavigationModel()
@@ -53,6 +57,7 @@ public class AddToHomeScreenPrompt {
             mLauncher.waitForObjectInContainer(
                     mWidgetCell.getParent().getParent().getParent().getParent(),
                     By.text(ADD_AUTOMATICALLY)).click();
+            mLauncher.waitUntilLauncherObjectGone(getSelector());
         }
     }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index ce7d39d1de..240f5154a9 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -542,32 +542,32 @@ public final class LauncherInstrumentation {
                     if (mDevice.isNaturalOrientation()) {
                         waitForLauncherObject(APPS_RES_ID);
                     } else {
-                        waitUntilGone(APPS_RES_ID);
+                        waitUntilLauncherObjectGone(APPS_RES_ID);
                     }
-                    waitUntilGone(OVERVIEW_RES_ID);
-                    waitUntilGone(WIDGETS_RES_ID);
+                    waitUntilLauncherObjectGone(OVERVIEW_RES_ID);
+                    waitUntilLauncherObjectGone(WIDGETS_RES_ID);
                     return waitForLauncherObject(WORKSPACE_RES_ID);
                 }
                 case WIDGETS: {
-                    waitUntilGone(WORKSPACE_RES_ID);
-                    waitUntilGone(APPS_RES_ID);
-                    waitUntilGone(OVERVIEW_RES_ID);
+                    waitUntilLauncherObjectGone(WORKSPACE_RES_ID);
+                    waitUntilLauncherObjectGone(APPS_RES_ID);
+                    waitUntilLauncherObjectGone(OVERVIEW_RES_ID);
                     return waitForLauncherObject(WIDGETS_RES_ID);
                 }
                 case ALL_APPS: {
-                    waitUntilGone(WORKSPACE_RES_ID);
-                    waitUntilGone(OVERVIEW_RES_ID);
-                    waitUntilGone(WIDGETS_RES_ID);
+                    waitUntilLauncherObjectGone(WORKSPACE_RES_ID);
+                    waitUntilLauncherObjectGone(OVERVIEW_RES_ID);
+                    waitUntilLauncherObjectGone(WIDGETS_RES_ID);
                     return waitForLauncherObject(APPS_RES_ID);
                 }
                 case OVERVIEW: {
                     if (hasAllAppsInOverview()) {
                         waitForLauncherObject(APPS_RES_ID);
                     } else {
-                        waitUntilGone(APPS_RES_ID);
+                        waitUntilLauncherObjectGone(APPS_RES_ID);
                     }
-                    waitUntilGone(WORKSPACE_RES_ID);
-                    waitUntilGone(WIDGETS_RES_ID);
+                    waitUntilLauncherObjectGone(WORKSPACE_RES_ID);
+                    waitUntilLauncherObjectGone(WIDGETS_RES_ID);
 
                     return waitForLauncherObject(OVERVIEW_RES_ID);
                 }
@@ -575,10 +575,10 @@ public final class LauncherInstrumentation {
                     return waitForFallbackLauncherObject(OVERVIEW_RES_ID);
                 }
                 case BACKGROUND: {
-                    waitUntilGone(WORKSPACE_RES_ID);
-                    waitUntilGone(APPS_RES_ID);
-                    waitUntilGone(OVERVIEW_RES_ID);
-                    waitUntilGone(WIDGETS_RES_ID);
+                    waitUntilLauncherObjectGone(WORKSPACE_RES_ID);
+                    waitUntilLauncherObjectGone(APPS_RES_ID);
+                    waitUntilLauncherObjectGone(OVERVIEW_RES_ID);
+                    waitUntilLauncherObjectGone(WIDGETS_RES_ID);
                     return null;
                 }
                 default:
@@ -647,7 +647,7 @@ public final class LauncherInstrumentation {
                             false, GestureScope.INSIDE_TO_OUTSIDE);
                     try (LauncherInstrumentation.Closable c = addContextLayer(
                             "Swiped up from context menu to home")) {
-                        waitUntilGone(CONTEXT_MENU_RES_ID);
+                        waitUntilLauncherObjectGone(CONTEXT_MENU_RES_ID);
                     }
                 }
                 if (hasLauncherObject(WORKSPACE_RES_ID)) {
@@ -800,9 +800,17 @@ public final class LauncherInstrumentation {
         }
     }
 
-    void waitUntilGone(String resId) {
-        assertTrue("Unexpected launcher object visible: " + resId,
-                mDevice.wait(Until.gone(getLauncherObjectSelector(resId)),
+    void waitUntilLauncherObjectGone(String resId) {
+        waitUntilGoneBySelector(getLauncherObjectSelector(resId));
+    }
+
+    void waitUntilLauncherObjectGone(BySelector selector) {
+        waitUntilGoneBySelector(makeLauncherSelector(selector));
+    }
+
+    private void waitUntilGoneBySelector(BySelector launcherSelector) {
+        assertTrue("Unexpected launcher object visible: " + launcherSelector,
+                mDevice.wait(Until.gone(launcherSelector),
                         WAIT_TIME_MS));
     }
 
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index 9f809173b7..f0e686f77f 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -230,7 +230,7 @@ public final class Workspace extends Home {
             launcher.expectEvent(TestProtocol.SEQUENCE_MAIN, LauncherInstrumentation.EVENT_START);
         }
         LauncherInstrumentation.log("dragIconToWorkspace: end");
-        launcher.waitUntilGone("drop_target_bar");
+        launcher.waitUntilLauncherObjectGone("drop_target_bar");
     }
 
     /**