1 files changed, 11 insertions, 10 deletions
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index f7befd140..443bc1af1 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -21,9 +21,9 @@ import static android.content.pm.PackageManager.DONT_KILL_APP;
 import static android.content.pm.PackageManager.MATCH_ALL;
 import static android.content.pm.PackageManager.MATCH_DISABLED_COMPONENTS;
 
+import static com.android.launcher3.tapl.TestHelpers.getOverviewPackageName;
 import static com.android.launcher3.testing.TestProtocol.BACKGROUND_APP_STATE_ORDINAL;
 import static com.android.launcher3.testing.TestProtocol.NORMAL_STATE_ORDINAL;
-import static com.android.launcher3.tapl.TestHelpers.getOverviewPackageName;
 
 import android.app.ActivityManager;
 import android.app.Instrumentation;
@@ -570,6 +570,16 @@ public final class LauncherInstrumentation {
         return object;
     }
 
+    @NonNull
+    UiObject2 waitForObjectInContainer(UiObject2 container, BySelector selector) {
+        final UiObject2 object = container.wait(
+                Until.findObject(selector),
+                WAIT_TIME_MS);
+        assertNotNull("Can't find a launcher object id: " + selector + " in container: " +
+                container.getResourceName(), object);
+        return object;
+    }
+
     @Nullable
     private boolean hasLauncherObject(String resId) {
         return mDevice.hasObject(getLauncherObjectSelector(resId));
@@ -581,11 +591,6 @@ public final class LauncherInstrumentation {
     }
 
     @NonNull
-    UiObject2 waitForLauncherObjectByClass(String clazz) {
-        return waitForObjectBySelector(getLauncherObjectSelectorByClass(clazz));
-    }
-
-    @NonNull
     UiObject2 waitForFallbackLauncherObject(String resName) {
         return waitForObjectBySelector(getFallbackLauncherObjectSelector(resName));
     }
@@ -600,10 +605,6 @@ public final class LauncherInstrumentation {
         return By.res(getLauncherPackageName(), resName);
     }
 
-    BySelector getLauncherObjectSelectorByClass(String clazz) {
-        return By.pkg(getLauncherPackageName()).clazz(clazz);
-    }
-
     BySelector getFallbackLauncherObjectSelector(String resName) {
         return By.res(getOverviewPackageName(), resName);
     }