diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 2ca40d8..efd8a55 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -1016,6 +1016,15 @@
         }
     }
 
+    List<UiObject2> getChildren(UiObject2 container) {
+        try {
+            return container.getChildren();
+        } catch (StaleObjectException e) {
+            fail("The container disappeared from screen");
+            return null;
+        }
+    }
+
     private boolean hasLauncherObject(String resId) {
         return mDevice.hasObject(getLauncherObjectSelector(resId));
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index 99d9889..6e7264a 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -116,9 +116,9 @@
                     "widget_preview");
             int i = 0;
             for (; ; ) {
-                final Collection<UiObject2> tableRows = widgetsContainer.getChildren();
+                final Collection<UiObject2> tableRows = mLauncher.getChildren(widgetsContainer);
                 for (UiObject2 row : tableRows) {
-                    final Collection<UiObject2> widgetCells = row.getChildren();
+                    final Collection<UiObject2> widgetCells = mLauncher.getChildren(row);
                     for (UiObject2 widget : widgetCells) {
                         final UiObject2 label = mLauncher.findObjectInContainer(widget,
                                 labelSelector);