[fenix] For https://github.com/mozilla-mobile/fenix/issues/21769 fix flaky tabMediaControlButtonTest Ui test

app/src/androidTest/java/org/mozilla/fenix/ui/SmokeTest.kt

@@ -1272,7 +1272,7 @@ class SmokeTest {
             assertPlaybackState(browserStore, MediaSession.PlaybackState.PLAYING)
         }.openTabDrawer {
             verifyTabMediaControlButtonState("Pause")
-            clickTabMediaControlButton()
+            clickTabMediaControlButton("Pause")
             verifyTabMediaControlButtonState("Play")
         }.openTab(audioTestPage.title) {
             assertPlaybackState(browserStore, MediaSession.PlaybackState.PAUSED)

app/src/androidTest/java/org/mozilla/fenix/ui/robots/TabDrawerRobot.kt

@@ -32,6 +32,7 @@ import androidx.test.uiautomator.UiSelector
 import androidx.test.uiautomator.Until
 import androidx.test.uiautomator.Until.findObject
 import com.google.android.material.bottomsheet.BottomSheetBehavior
+import junit.framework.AssertionFailedError
 import junit.framework.TestCase.assertTrue
 import org.hamcrest.CoreMatchers.allOf
 import org.hamcrest.CoreMatchers.anyOf
@@ -129,21 +130,73 @@ class TabDrawerRobot {
     }
 
     fun verifyTabMediaControlButtonState(action: String) {
-        mDevice.waitForIdle()
+        try {
+            mDevice.findObject(
+                UiSelector().resourceId("$packageName:id/tab_tray_empty_view")
+            ).waitUntilGone(waitingTime)
 
-        mDevice.findObject(
-            UiSelector()
-                .resourceId("$packageName:id/play_pause_button")
-        ).waitForExists(waitingTime)
+            mDevice.findObject(
+                UiSelector().resourceId("$packageName:id/tab_tray_grid_item")
+            ).waitForExists(waitingTime)
 
-        assertTrue(
             mDevice.findObject(
-                UiSelector().descriptionContains(action)
+                UiSelector()
+                    .resourceId("$packageName:id/play_pause_button")
+                    .descriptionContains(action)
             ).waitForExists(waitingTime)
-        )
+
+            assertTrue(
+                mDevice.findObject(UiSelector().descriptionContains(action)).waitForExists(waitingTime)
+            )
+        } catch (e: AssertionFailedError) {
+            // In some cases the tab media button isn't updated after performing an action on it
+            println("Failed to update the state of the tab media button")
+
+            // Let's dismiss the tabs tray and try again
+            mDevice.pressBack()
+            mDevice.findObject(
+                UiSelector()
+                    .resourceId("$packageName:id/toolbar")
+            ).waitForExists(waitingTime)
+
+            browserScreen {
+            }.openTabDrawer {
+                // Click again the tab media button
+                tabMediaControlButton().click()
+
+                mDevice.findObject(
+                    UiSelector().resourceId("$packageName:id/tab_tray_empty_view")
+                ).waitUntilGone(waitingTime)
+
+                mDevice.findObject(
+                    UiSelector().resourceId("$packageName:id/tab_tray_grid_item")
+                ).waitForExists(waitingTime)
+
+                mDevice.findObject(
+                    UiSelector()
+                        .resourceId("$packageName:id/play_pause_button")
+                        .descriptionContains(action)
+                ).waitForExists(waitingTime)
+
+                assertTrue(
+                    mDevice.findObject(UiSelector().descriptionContains(action)).waitForExists(waitingTime)
+                )
+            }
+        }
     }
 
-    fun clickTabMediaControlButton() = tabMediaControlButton().click()
+    fun clickTabMediaControlButton(action: String) {
+        mDevice.waitNotNull(
+            Until.findObjects(
+                By
+                    .res("$packageName:id/play_pause_button")
+                    .descContains(action)
+            ),
+            waitingTime
+        )
+
+        tabMediaControlButton().click()
+    }
 
     fun clickSelectTabs() {
         threeDotMenu().click()
@@ -355,7 +408,8 @@ fun tabDrawer(interact: TabDrawerRobot.() -> Unit): TabDrawerRobot.Transition {
     return TabDrawerRobot.Transition()
 }
 
-private fun tabMediaControlButton() = onView(withId(R.id.play_pause_button))
+private fun tabMediaControlButton() =
+    mDevice.findObject(UiSelector().resourceId("$packageName:id/play_pause_button"))
 
 private fun closeTabButton() =
     mDevice.findObject(UiSelector().resourceId("$packageName:id/mozac_browser_tabstray_close"))