From b1ee43a7ed9e88c092946b7691af4603aea85db1 Mon Sep 17 00:00:00 2001
From: AndiAJ <andiaj@users.noreply.github.com>
Date: Mon, 13 Sep 2021 14:37:00 +0300
Subject: [PATCH] [fenix] Improve coverage for addPrivateBrowsingShortcut Ui
 test

---
 .../org/mozilla/fenix/ui/SettingsPrivacyTest.kt   |  1 +
 .../robots/SettingsSubMenuPrivateBrowsingRobot.kt | 15 +++++++++++++++
 2 files changed, 16 insertions(+)

diff --git a/app/src/androidTest/java/org/mozilla/fenix/ui/SettingsPrivacyTest.kt b/app/src/androidTest/java/org/mozilla/fenix/ui/SettingsPrivacyTest.kt
index ebe3333d2d..8594164362 100644
--- a/app/src/androidTest/java/org/mozilla/fenix/ui/SettingsPrivacyTest.kt
+++ b/app/src/androidTest/java/org/mozilla/fenix/ui/SettingsPrivacyTest.kt
@@ -364,6 +364,7 @@ class SettingsPrivacyTest {
         }.openThreeDotMenu {
         }.openSettings {
         }.openPrivateBrowsingSubMenu {
+            cancelPrivateShortcutAddition()
             addPrivateShortcutToHomescreen()
             verifyPrivateBrowsingShortcutIcon()
         }.openPrivateBrowsingShortcut {
diff --git a/app/src/androidTest/java/org/mozilla/fenix/ui/robots/SettingsSubMenuPrivateBrowsingRobot.kt b/app/src/androidTest/java/org/mozilla/fenix/ui/robots/SettingsSubMenuPrivateBrowsingRobot.kt
index 4a5a23438b..943b568b19 100644
--- a/app/src/androidTest/java/org/mozilla/fenix/ui/robots/SettingsSubMenuPrivateBrowsingRobot.kt
+++ b/app/src/androidTest/java/org/mozilla/fenix/ui/robots/SettingsSubMenuPrivateBrowsingRobot.kt
@@ -51,6 +51,18 @@ class SettingsSubMenuPrivateBrowsingRobot {
 
     fun clickOpenLinksInPrivateTabSwitch() = openLinksInPrivateTabSwitch().click()
 
+    fun cancelPrivateShortcutAddition() {
+        mDevice.wait(
+            Until.findObject(text("Add private browsing shortcut")),
+            waitingTime
+        )
+        addPrivateBrowsingShortcutButton().click()
+        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O) {
+            mDevice.wait(Until.findObject(By.textContains("CANCEL")), waitingTime)
+            cancelShortcutAdditionButton().click()
+        }
+    }
+
     fun addPrivateShortcutToHomescreen() {
         mDevice.wait(
             Until.findObject(text("Add private browsing shortcut")),
@@ -105,6 +117,9 @@ private fun goBackButton() = onView(withContentDescription("Navigate up"))
 private fun addAutomaticallyButton() =
     mDevice.findObject(UiSelector().textStartsWith("add automatically"))
 
+private fun cancelShortcutAdditionButton() =
+    mDevice.findObject(UiSelector().textContains("CANCEL"))
+
 private fun privateBrowsingShortcutIcon() = mDevice.findObject(text("Private $appName"))
 
 private fun assertAddPrivateBrowsingShortcutButton() {