diff --git a/tests/tapl/com/android/launcher3/tapl/AllApps.java b/tests/tapl/com/android/launcher3/tapl/AllApps.java
index 9ff354a..f070280 100644
--- a/tests/tapl/com/android/launcher3/tapl/AllApps.java
+++ b/tests/tapl/com/android/launcher3/tapl/AllApps.java
@@ -120,7 +120,7 @@
             mLauncher.assertTrue("Unable to scroll to a clickable icon: " + appName,
                     hasClickableIcon(allAppsContainer, appListRecycler, appIconSelector));
 
-            final UiObject2 appIcon = mLauncher.getObjectInContainer(appListRecycler,
+            final UiObject2 appIcon = mLauncher.waitForObjectInContainer(appListRecycler,
                     appIconSelector);
             return new AppIcon(mLauncher, appIcon);
         }
diff --git a/tests/tapl/com/android/launcher3/tapl/BaseOverview.java b/tests/tapl/com/android/launcher3/tapl/BaseOverview.java
index bbd2c29..25e6e8c 100644
--- a/tests/tapl/com/android/launcher3/tapl/BaseOverview.java
+++ b/tests/tapl/com/android/launcher3/tapl/BaseOverview.java
@@ -74,7 +74,7 @@
                 flingForward();
             }
 
-            mLauncher.getObjectInContainer(verifyActiveContainer(), clearAllSelector).click();
+            mLauncher.waitForObjectInContainer(verifyActiveContainer(), clearAllSelector).click();
             try (LauncherInstrumentation.Closable c1 = mLauncher.addContextLayer(
                     "dismissed all tasks")) {
                 return new Workspace(mLauncher);
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 1371b21..6348c41 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -679,13 +679,6 @@
     }
 
     @NonNull
-    UiObject2 getObjectInContainer(UiObject2 container, BySelector selector) {
-        final UiObject2 object = container.findObject(selector);
-        assertNotNull("Can't find an object with selector: " + selector, object);
-        return object;
-    }
-
-    @NonNull
     List<UiObject2> getObjectsInContainer(UiObject2 container, String resName) {
         return container.findObjects(getLauncherObjectSelector(resName));
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index 07f8b64..639902f 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -108,10 +108,13 @@
      */
     @NonNull
     public AppIcon getWorkspaceAppIcon(String appName) {
-        return new AppIcon(mLauncher,
-                mLauncher.getObjectInContainer(
-                        verifyActiveContainer(),
-                        AppIcon.getAppIconSelector(appName, mLauncher)));
+        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                "want to get a workspace icon")) {
+            return new AppIcon(mLauncher,
+                    mLauncher.waitForObjectInContainer(
+                            verifyActiveContainer(),
+                            AppIcon.getAppIconSelector(appName, mLauncher)));
+        }
     }
 
     /**
@@ -142,13 +145,13 @@
 
     @NonNull
     public AppIcon getHotseatAppIcon(String appName) {
-        return new AppIcon(mLauncher, mLauncher.getObjectInContainer(
+        return new AppIcon(mLauncher, mLauncher.waitForObjectInContainer(
                 mHotseat, AppIcon.getAppIconSelector(appName, mLauncher)));
     }
 
     @NonNull
     public Folder getHotseatFolder(String appName) {
-        return new Folder(mLauncher, mLauncher.getObjectInContainer(
+        return new Folder(mLauncher, mLauncher.waitForObjectInContainer(
                 mHotseat, Folder.getSelector(appName, mLauncher)));
     }