diff --git a/tests/tapl/com/android/launcher3/tapl/AppIconMenu.java b/tests/tapl/com/android/launcher3/tapl/AppIconMenu.java
index 2a03f9a..7f28151 100644
--- a/tests/tapl/com/android/launcher3/tapl/AppIconMenu.java
+++ b/tests/tapl/com/android/launcher3/tapl/AppIconMenu.java
@@ -20,6 +20,8 @@
 
 import androidx.test.uiautomator.UiObject2;
 
+import java.util.List;
+
 /**
  * Context menu of an app icon.
  */
@@ -37,10 +39,9 @@
      * Returns a menu item with a given number. Fails if it doesn't exist.
      */
     public AppIconMenuItem getMenuItem(int itemNumber) {
-        assertTrue(mDeepShortcutsContainer.getChildCount() > itemNumber);
-
-        final UiObject2 shortcut = mLauncher.waitForObjectInContainer(
-                mDeepShortcutsContainer.getChildren().get(itemNumber), "bubble_text");
-        return new AppIconMenuItem(mLauncher, shortcut);
+        final List<UiObject2> menuItems = mLauncher.getObjectsInContainer(mDeepShortcutsContainer,
+                "bubble_text");
+        assertTrue(menuItems.size() > itemNumber);
+        return new AppIconMenuItem(mLauncher, menuItems.get(itemNumber));
     }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index bb399d5..49bd73a 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -41,6 +41,7 @@
 import org.junit.Assert;
 
 import java.lang.ref.WeakReference;
+import java.util.List;
 import java.util.concurrent.TimeoutException;
 
 /**
@@ -367,6 +368,11 @@
     }
 
     @NonNull
+    List<UiObject2> getObjectsInContainer(UiObject2 container, String resName) {
+        return container.findObjects(getLauncherObjectSelector(resName));
+    }
+
+    @NonNull
     UiObject2 waitForObjectInContainer(UiObject2 container, String resName) {
         final UiObject2 object = container.wait(
                 Until.findObject(getLauncherObjectSelector(resName)),