From 272d622bc4c572fbfc8e2b9362f810cf411def9d Mon Sep 17 00:00:00 2001
From: vadimt <vadimt@google.com>
Date: Tue, 17 Aug 2021 18:47:12 -0700
Subject: [PATCH] Improving diags for containers disappearing while getting
 children

Bug: 184609576
Test: presubmit
Change-Id: I7f527a9991a58fd148b17783078b5ac80979a5b3
---
 .../android/launcher3/tapl/LauncherInstrumentation.java  | 9 +++++++++
 tests/tapl/com/android/launcher3/tapl/Widgets.java       | 4 ++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 2ca40d8229..efd8a551d4 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -1016,6 +1016,15 @@ public final class LauncherInstrumentation {
         }
     }
 
+    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 99d988926a..6e7264ac0b 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -116,9 +116,9 @@ public final class Widgets extends LauncherInstrumentation.VisibleContainer {
                     "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);