From 2400498718ebf73f395c1331859afd35c5e1f0c1 Mon Sep 17 00:00:00 2001
From: AndiAJ <andiaj@users.noreply.github.com>
Date: Tue, 17 May 2022 15:45:44 +0300
Subject: [PATCH] For #25184 and #23417 fix flaky testDownloadPrompt and
 testCloseDownloadPrompt

---
 .../mozilla/fenix/ui/robots/DownloadRobot.kt  | 33 +++++++++++++------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git a/app/src/androidTest/java/org/mozilla/fenix/ui/robots/DownloadRobot.kt b/app/src/androidTest/java/org/mozilla/fenix/ui/robots/DownloadRobot.kt
index 5687dcb7cc14..61e82ddab308 100644
--- a/app/src/androidTest/java/org/mozilla/fenix/ui/robots/DownloadRobot.kt
+++ b/app/src/androidTest/java/org/mozilla/fenix/ui/robots/DownloadRobot.kt
@@ -133,16 +133,29 @@ fun downloadRobot(interact: DownloadRobot.() -> Unit): DownloadRobot.Transition
 }
 
 private fun assertDownloadPrompt(fileName: String) {
-    assertTrue(
-        "Download prompt button not visible",
-        mDevice.findObject(UiSelector().resourceId("$packageName:id/download_button"))
-            .waitForExists(waitingTime)
-    )
-    assertTrue(
-        "$fileName title doesn't match",
-        mDevice.findObject(UiSelector().text(fileName))
-            .waitForExists(waitingTime)
-    )
+    var currentTries = 0
+    while (currentTries++ < 3) {
+        try {
+            assertTrue(
+                "Download prompt button not visible",
+                mDevice.findObject(UiSelector().resourceId("$packageName:id/download_button"))
+                    .waitForExists(waitingTime)
+            )
+            assertTrue(
+                "$fileName title doesn't match",
+                mDevice.findObject(UiSelector().text(fileName))
+                    .waitForExists(waitingTime)
+            )
+
+            break
+        } catch (e: AssertionError) {
+            println("Failed to find locator: ${e.localizedMessage}")
+
+            browserScreen {
+            }.clickDownloadLink(fileName) {
+            }
+        }
+    }
 }
 
 private fun assertDownloadNotificationPopup() {