[Chooser/ResolverActivity] Fix flakiness in work profile tests

Sometimes work profile tests were failing
because work profile tab button click was not
handled. Most likely it is related to the
opening animation of the drawer layout.

Added retry mechanism to click on the button
again if the tab was not selected.

Test: repeated run of ResolverActivityWorkProfileTest
Test: atest com.android.internal.app.ChooserActivityWorkProfileTest
Test: atest com.android.internal.app.ResolverActivityWorkProfileTest
Test: atest com.android.intentresolve.UnbundledChooserActivityWorkProfileTest
Bug: 262018267
Change-Id: Icb6d0337d7f92f1ceb59ef3c03e5b9338dcf2890

2 files changed, 32 insertions, 4 deletions

diff --git a/core/tests/coretests/src/com/android/internal/app/ChooserActivityWorkProfileTest.java b/core/tests/coretests/src/com/android/internal/app/ChooserActivityWorkProfileTest.java
index c6537c0fbfb3..b6ea9dd3998d 100644
--- a/core/tests/coretests/src/com/android/internal/app/ChooserActivityWorkProfileTest.java
+++ b/core/tests/coretests/src/com/android/internal/app/ChooserActivityWorkProfileTest.java
@@ -16,11 +16,14 @@
 
 package com.android.internal.app;
 
+import static android.util.PollingCheck.waitFor;
+
 import static androidx.test.espresso.Espresso.onView;
 import static androidx.test.espresso.action.ViewActions.click;
 import static androidx.test.espresso.action.ViewActions.swipeUp;
 import static androidx.test.espresso.assertion.ViewAssertions.matches;
 import static androidx.test.espresso.matcher.ViewMatchers.isDisplayed;
+import static androidx.test.espresso.matcher.ViewMatchers.isSelected;
 import static androidx.test.espresso.matcher.ViewMatchers.withId;
 import static androidx.test.espresso.matcher.ViewMatchers.withText;
 
@@ -49,6 +52,8 @@ import com.android.internal.R;
 import com.android.internal.app.ResolverActivity.ResolvedComponentInfo;
 import com.android.internal.app.ChooserActivityWorkProfileTest.TestCase.Tab;
 
+import junit.framework.AssertionFailedError;
+
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -331,8 +336,17 @@ public class ChooserActivityWorkProfileTest {
         final int stringId = tab == Tab.WORK ? R.string.resolver_work_tab
                 : R.string.resolver_personal_tab;
 
-        onView(withText(stringId)).perform(click());
-        waitForIdle();
+        waitFor(() -> {
+            onView(withText(stringId)).perform(click());
+            waitForIdle();
+
+            try {
+                onView(withText(stringId)).check(matches(isSelected()));
+                return true;
+            } catch (AssertionFailedError e) {
+                return false;
+            }
+        });
 
         onView(withId(R.id.contentPanel))
                 .perform(swipeUp());
diff --git a/core/tests/coretests/src/com/android/internal/app/ResolverActivityWorkProfileTest.java b/core/tests/coretests/src/com/android/internal/app/ResolverActivityWorkProfileTest.java
index ce68906a5bff..4a79f5ee3f88 100644
--- a/core/tests/coretests/src/com/android/internal/app/ResolverActivityWorkProfileTest.java
+++ b/core/tests/coretests/src/com/android/internal/app/ResolverActivityWorkProfileTest.java
@@ -16,11 +16,14 @@
 
 package com.android.internal.app;
 
+import static android.util.PollingCheck.waitFor;
+
 import static androidx.test.espresso.Espresso.onView;
 import static androidx.test.espresso.action.ViewActions.click;
 import static androidx.test.espresso.action.ViewActions.swipeUp;
 import static androidx.test.espresso.assertion.ViewAssertions.matches;
 import static androidx.test.espresso.matcher.ViewMatchers.isDisplayed;
+import static androidx.test.espresso.matcher.ViewMatchers.isSelected;
 import static androidx.test.espresso.matcher.ViewMatchers.withId;
 import static androidx.test.espresso.matcher.ViewMatchers.withText;
 
@@ -48,6 +51,8 @@ import com.android.internal.R;
 import com.android.internal.app.ResolverActivity.ResolvedComponentInfo;
 import com.android.internal.app.ResolverActivityWorkProfileTest.TestCase.Tab;
 
+import junit.framework.AssertionFailedError;
+
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -309,8 +314,17 @@ public class ResolverActivityWorkProfileTest {
         final int stringId = tab == Tab.WORK ? R.string.resolver_work_tab
                 : R.string.resolver_personal_tab;
 
-        onView(withText(stringId)).perform(click());
-        waitForIdle();
+        waitFor(() -> {
+            onView(withText(stringId)).perform(click());
+            waitForIdle();
+
+            try {
+                onView(withText(stringId)).check(matches(isSelected()));
+                return true;
+            } catch (AssertionFailedError e) {
+                return false;
+            }
+        });
 
         onView(withId(R.id.contentPanel))
                 .perform(swipeUp());