From 8a8255916f37fe32e232eabfcac0fe02d464795c Mon Sep 17 00:00:00 2001
From: vadimt <vadimt@google.com>
Date: Wed, 20 May 2020 15:05:49 -0700
Subject: [PATCH] Failing with diags when UiObject2 container disappears

Change-Id: If6af49087a5afd6374776e55067b0502dabb7585
---
 .../tapl/LauncherInstrumentation.java         | 41 +++++++++++++------
 1 file changed, 28 insertions(+), 13 deletions(-)

diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 79008a76f3..681865e775 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -833,27 +833,42 @@ public final class LauncherInstrumentation {
 
     @NonNull
     List<UiObject2> getObjectsInContainer(UiObject2 container, String resName) {
-        return container.findObjects(getLauncherObjectSelector(resName));
+        try {
+            return container.findObjects(getLauncherObjectSelector(resName));
+        } catch (StaleObjectException e) {
+            fail("The container disappeared from screen");
+            return null;
+        }
     }
 
     @NonNull
     UiObject2 waitForObjectInContainer(UiObject2 container, String resName) {
-        final UiObject2 object = container.wait(
-                Until.findObject(getLauncherObjectSelector(resName)),
-                WAIT_TIME_MS);
-        assertNotNull("Can't find a view in Launcher, id: " + resName + " in container: "
-                + container.getResourceName(), object);
-        return object;
+        try {
+            final UiObject2 object = container.wait(
+                    Until.findObject(getLauncherObjectSelector(resName)),
+                    WAIT_TIME_MS);
+            assertNotNull("Can't find a view in Launcher, id: " + resName + " in container: "
+                    + container.getResourceName(), object);
+            return object;
+        } catch (StaleObjectException e) {
+            fail("The container disappeared from screen");
+            return null;
+        }
     }
 
     @NonNull
     UiObject2 waitForObjectInContainer(UiObject2 container, BySelector selector) {
-        final UiObject2 object = container.wait(
-                Until.findObject(selector),
-                WAIT_TIME_MS);
-        assertNotNull("Can't find a view in Launcher, id: " + selector + " in container: "
-                + container.getResourceName(), object);
-        return object;
+        try {
+            final UiObject2 object = container.wait(
+                    Until.findObject(selector),
+                    WAIT_TIME_MS);
+            assertNotNull("Can't find a view in Launcher, id: " + selector + " in container: "
+                    + container.getResourceName(), object);
+            return object;
+        } catch (StaleObjectException e) {
+            fail("The container disappeared from screen");
+            return null;
+        }
     }
 
     private boolean hasLauncherObject(String resId) {