From 8f6452d5b8cdf2b10d9c2c723f907c42d9bf2c45 Mon Sep 17 00:00:00 2001
From: mlykotom <mlykotom@google.com>
Date: Thu, 20 Apr 2023 15:17:46 +0200
Subject: [PATCH] Fix forYou actions

Change-Id: I445cc9766b4035a40079ab3566d94706131d1db5
---
 .../test/uiautomator/UiAutomatorHelpers.kt    | 59 +++++++++++++++++++
 .../google/samples/apps/nowinandroid/Utils.kt |  1 +
 .../apps/nowinandroid/foryou/ForYouActions.kt |  5 +-
 3 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/benchmarks/src/main/java/androidx/test/uiautomator/UiAutomatorHelpers.kt b/benchmarks/src/main/java/androidx/test/uiautomator/UiAutomatorHelpers.kt
index 85867b982..b98c54daf 100644
--- a/benchmarks/src/main/java/androidx/test/uiautomator/UiAutomatorHelpers.kt
+++ b/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()
+}
diff --git a/benchmarks/src/main/java/com/google/samples/apps/nowinandroid/Utils.kt b/benchmarks/src/main/java/com/google/samples/apps/nowinandroid/Utils.kt
index 494e451d0..8ee7a983a 100644
--- a/benchmarks/src/main/java/com/google/samples/apps/nowinandroid/Utils.kt
+++ b/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)
 }
+
diff --git a/benchmarks/src/main/java/com/google/samples/apps/nowinandroid/foryou/ForYouActions.kt b/benchmarks/src/main/java/com/google/samples/apps/nowinandroid/foryou/ForYouActions.kt
index 5e4952fbb..918e8a505 100644
--- a/benchmarks/src/main/java/com/google/samples/apps/nowinandroid/foryou/ForYouActions.kt
+++ b/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)
 }
 
 /**