[Autofill Assistant] Fix for failing tests.

Bug: 1123958
Change-Id: Ie0ba79cc3ace8d7a115dac30888b2430972239d4
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2390205
Commit-Queue: Michele Mancina <micantox@google.com>
Reviewed-by: Marian Fechete <marianfe@google.com>
Cr-Commit-Position: refs/heads/master@{#803874}

chrome/android/features/autofill_assistant/javatests/src/org/chromium/chrome/browser/autofill_assistant/AutofillAssistantChromeTabIntegrationTest.java

@@ -22,6 +22,7 @@ import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.not;
 
+import static org.chromium.chrome.browser.autofill_assistant.AutofillAssistantUiTestUtil.waitUntil;
 import static org.chromium.chrome.browser.autofill_assistant.AutofillAssistantUiTestUtil.waitUntilKeyboardMatchesCondition;
 import static org.chromium.chrome.browser.autofill_assistant.AutofillAssistantUiTestUtil.waitUntilViewAssertionTrue;
 import static org.chromium.chrome.browser.autofill_assistant.AutofillAssistantUiTestUtil.waitUntilViewMatchesCondition;
@@ -39,8 +40,8 @@ import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import org.chromium.base.task.PostTask;
 import org.chromium.base.test.util.CommandLineFlags;
-import org.chromium.base.test.util.DisabledTest;
 import org.chromium.base.test.util.Restriction;
 import org.chromium.chrome.autofill_assistant.R;
 import org.chromium.chrome.browser.autofill_assistant.proto.ActionProto;
@@ -56,7 +57,10 @@ import org.chromium.chrome.browser.tabmodel.TabModelUtils;
 import org.chromium.chrome.test.ChromeJUnit4ClassRunner;
 import org.chromium.chrome.test.ChromeTabbedActivityTestRule;
 import org.chromium.chrome.test.util.ChromeTabUtils;
+import org.chromium.components.browser_ui.bottomsheet.BottomSheetController;
+import org.chromium.components.browser_ui.bottomsheet.BottomSheetController.SheetState;
 import org.chromium.components.browser_ui.widget.scrim.ScrimCoordinator;
+import org.chromium.content_public.browser.UiThreadTaskTraits;
 import org.chromium.content_public.browser.test.util.TestThreadUtils;
 import org.chromium.net.test.EmbeddedTestServer;
 import org.chromium.ui.test.util.UiRestriction;
@@ -342,7 +346,6 @@ public class AutofillAssistantChromeTabIntegrationTest {
 
     @Test
     @MediumTest
-    @DisabledTest(message = "Flaky - https://crbug.com/1123958")
     public void switchTabBetweenDifferentPeekModes() {
         ArrayList<ActionProto> listA = new ArrayList<>();
         listA.add((ActionProto) ActionProto.newBuilder()
@@ -367,8 +370,7 @@ public class AutofillAssistantChromeTabIntegrationTest {
                                                                               .proto.ChipType
                                                                               .HIGHLIGHTED_ACTION)))
                                              .setAllowInterrupt(false)
-                                             .setDisableForceExpandSheet(true)
-                                  /*.setBrowseMode(true)*/)
+                                             .setDisableForceExpandSheet(true))
                           .build());
 
         AutofillAssistantTestScript scriptA = new AutofillAssistantTestScript(
@@ -397,7 +399,7 @@ public class AutofillAssistantChromeTabIntegrationTest {
         int initialTabId =
                 TabModelUtils.getCurrentTabId(mTestRule.getActivity().getCurrentTabModel());
 
-        AutofillAssistantTestService autofillAssistantTestService = setupScripts(scriptA, scriptB);
+        setupScripts(scriptA, scriptB);
         startAutofillAssistantOnTab(TEST_PAGE_A);
         waitUntilViewMatchesCondition(
                 allOf(withText("Sticky"), isDescendantOfA(withId(R.id.header))),
@@ -410,11 +412,13 @@ public class AutofillAssistantChromeTabIntegrationTest {
 
         startAutofillAssistantOnTab(TEST_PAGE_B);
         waitUntilViewMatchesCondition(withText("Prompt B"), isCompletelyDisplayed());
+        waitForBottomSheetAnimationToEnd();
 
         Espresso.pressBack();
         waitUntilViewMatchesCondition(
                 withId(R.id.status_message), allOf(withText("Prompt B"), not(isDisplayed())));
         onView(withId(R.id.autofill_assistant)).check(matches(isDisplayed()));
+        waitForBottomSheetAnimationToEnd();
 
         int secondTabId =
                 TabModelUtils.getCurrentTabId(mTestRule.getActivity().getCurrentTabModel());
@@ -426,18 +430,18 @@ public class AutofillAssistantChromeTabIntegrationTest {
         waitUntilViewMatchesCondition(
                 allOf(withText("Sticky"), isDescendantOfA(withId(R.id.header))),
                 isCompletelyDisplayed());
-
-        ChromeTabUtils.switchTabInCurrentTabModel(mTestRule.getActivity(),
-                TabModelUtils.getTabIndexById(
-                        mTestRule.getActivity().getCurrentTabModel(), secondTabId));
-        waitUntilViewMatchesCondition(withId(R.id.autofill_assistant), isDisplayed());
-        onView(allOf(withId(R.id.status_message), withText("Prompt B")))
-                .check(matches(not(isDisplayed())));
+        PostTask.postTask(UiThreadTaskTraits.DEFAULT, () -> {
+            ChromeTabUtils.switchTabInCurrentTabModel(mTestRule.getActivity(),
+                    TabModelUtils.getTabIndexById(
+                            mTestRule.getActivity().getCurrentTabModel(), secondTabId));
+            waitUntilViewMatchesCondition(withId(R.id.autofill_assistant), isDisplayed());
+            waitUntilViewMatchesCondition(
+                    allOf(withId(R.id.status_message), withText("Prompt B")), not(isDisplayed()));
+        });
     }
 
     @Test
     @MediumTest
-    @DisabledTest(message = "Flaky - https://crbug.com/1115681")
     public void startingNewAutofillAssistantCloseTabResumesRunOnPreviousTab() {
         ArrayList<ActionProto> listA = new ArrayList<>();
         listA.add((ActionProto) ActionProto.newBuilder()
@@ -472,6 +476,7 @@ public class AutofillAssistantChromeTabIntegrationTest {
         startAutofillAssistantOnTab(TEST_PAGE_A);
 
         waitUntilViewMatchesCondition(withText("Prompt A"), isCompletelyDisplayed());
+        waitForBottomSheetAnimationToEnd();
 
         ChromeTabUtils.fullyLoadUrlInNewTab(InstrumentationRegistry.getInstrumentation(),
                 mTestRule.getActivity(), getURL(TEST_PAGE_B), false);
@@ -479,6 +484,7 @@ public class AutofillAssistantChromeTabIntegrationTest {
 
         startAutofillAssistantOnTab(TEST_PAGE_B);
         waitUntilViewMatchesCondition(withText("Prompt B"), isCompletelyDisplayed());
+        waitForBottomSheetAnimationToEnd();
 
         ChromeTabUtils.closeCurrentTab(
                 InstrumentationRegistry.getInstrumentation(), mTestRule.getActivity());
@@ -486,6 +492,10 @@ public class AutofillAssistantChromeTabIntegrationTest {
         waitUntilViewMatchesCondition(withText("Prompt A"), isCompletelyDisplayed());
     }
 
+    private void waitForBottomSheetAnimationToEnd() {
+        waitUntil(() -> getBottomSheetController().getSheetState() != SheetState.SCROLLING);
+    }
+
     @Test
     @MediumTest
     public void interactingWithLocationBarHidesAutofillAssistant() {
@@ -552,4 +562,8 @@ public class AutofillAssistantChromeTabIntegrationTest {
                 InstrumentationRegistry.getInstrumentation(), mTestRule.getActivity());
         waitUntilViewMatchesCondition(withText("Shutdown"), isCompletelyDisplayed());
     }
+
+    private BottomSheetController getBottomSheetController() {
+        return mTestRule.getActivity().getRootUiCoordinatorForTesting().getBottomSheetController();
+    }
 }