Fix forYou actions

Change-Id: I445cc9766b4035a40079ab3566d94706131d1db5

benchmarks/src/main/java/androidx/test/uiautomator/UiAutomatorHelpers.kt

@@ -16,9 +16,11 @@
 
 package androidx.test.uiautomator
 
+import android.util.Log
 import androidx.test.uiautomator.HasChildrenOp.AT_LEAST
 import androidx.test.uiautomator.HasChildrenOp.AT_MOST
 import androidx.test.uiautomator.HasChildrenOp.EXACTLY
+import java.io.ByteArrayOutputStream
 
 // These helpers need to be in the androidx.test.uiautomator package,
 // because the abstract class has package local method that needs to be implemented.
@@ -46,3 +48,60 @@ enum class HasChildrenOp {
     EXACTLY,
     AT_MOST,
 }
+
+/**
+ * Finds an element by [selector].
+ * If not found, fails with [AssertionError] and dumps the window hierarchy.
+ */
+fun UiDevice.findOrFail(
+    selector: BySelector,
+    message: String? = null,
+): UiObject2 {
+    val element = findObject(selector)
+    if (element == null) {
+        val hierarchy = getWindowHierarchy()
+        Log.d("Benchmark", hierarchy)
+        throw AssertionError((message ?: "Element not on screen ($selector)") + "\n$hierarchy")
+    }
+    return element
+}
+
+/**
+ * Waits until a [searchCondition] evaluates to true.
+ * If not found within [timeout], fails with [AssertionError] and dumps the window hierarchy.
+ * For example, wait [Until.hasObject] to wait until an element is present on screen.
+ */
+fun UiDevice.waitOrFail(
+    searchCondition: SearchCondition<Boolean>,
+    timeout: Long,
+    message: String? = null,
+) {
+    // If successfully found, just return
+    if (wait(searchCondition, timeout)) return
+
+    val hierarchy = getWindowHierarchy()
+    Log.d("Benchmark", hierarchy)
+    throw AssertionError((message ?: "Element not on screen") + "\n$hierarchy")
+}
+
+/**
+ * Combines waiting for an element and returning it.
+ * If an object is not present, it throws [AssertionError] and dumps the window hierarchy.
+ */
+fun UiDevice.waitAndFind(
+    selector: BySelector,
+    timeout: Long = 5_000,
+    message: String? = null,
+): UiObject2 {
+    waitOrFail(Until.hasObject(selector), timeout, message)
+    return findOrFail(selector, message)
+}
+
+/**
+ * Simplifies dumping window hierarchy
+ */
+fun UiDevice.getWindowHierarchy(): String {
+    val output = ByteArrayOutputStream()
+    dumpWindowHierarchy(output)
+    return output.toString()
+}

benchmarks/src/main/java/com/google/samples/apps/nowinandroid/Utils.kt

@@ -38,3 +38,4 @@ fun UiDevice.flingElementDownUp(element: UiObject2) {
     waitForIdle()
     element.fling(Direction.UP)
 }
+

benchmarks/src/main/java/com/google/samples/apps/nowinandroid/foryou/ForYouActions.kt

@@ -20,6 +20,7 @@ import androidx.benchmark.macro.MacrobenchmarkScope
 import androidx.test.uiautomator.By
 import androidx.test.uiautomator.Until
 import androidx.test.uiautomator.untilHasChildren
+import androidx.test.uiautomator.waitAndFind
 import com.google.samples.apps.nowinandroid.flingElementDownUp
 
 fun MacrobenchmarkScope.forYouWaitForContent() {
@@ -27,9 +28,9 @@ fun MacrobenchmarkScope.forYouWaitForContent() {
     device.wait(Until.gone(By.res("loadingWheel")), 5_000)
     // Sometimes, the loading wheel is gone, but the content is not loaded yet
     // So we'll wait here for topics to be sure
-    val obj = device.findObject(By.res("forYou:topicSelection"))
+    val topicSelection = device.waitAndFind(By.res("forYou:topicSelection"), 5_000)
     // Timeout here is quite big, because sometimes data loading takes a long time!
-    obj.wait(untilHasChildren(), 60_000)
+    topicSelection.wait(untilHasChildren(), 60_000)
 }
 
 /**