[Autofill Assistant] Add WaitForNavigation action to stabilize test

Bug: 1104900
Change-Id: Ied3096ac58166b2a254e60220302e47ac53f2621
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2299723Reviewed-by: Clemens Arbesser <arbesser@google.com>
Commit-Queue: Sandro Maggi <sandromaggi@google.com>
Cr-Commit-Position: refs/heads/master@{#788660}

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

@@ -31,7 +31,6 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 
 import org.chromium.base.test.util.CommandLineFlags;
-import org.chromium.base.test.util.DisabledTest;
 import org.chromium.chrome.autofill_assistant.R;
 import org.chromium.chrome.browser.autofill_assistant.proto.ActionProto;
 import org.chromium.chrome.browser.autofill_assistant.proto.ChipProto;
@@ -44,6 +43,7 @@ import org.chromium.chrome.browser.autofill_assistant.proto.StopProto;
 import org.chromium.chrome.browser.autofill_assistant.proto.SupportedScriptProto;
 import org.chromium.chrome.browser.autofill_assistant.proto.SupportedScriptProto.PresentationProto;
 import org.chromium.chrome.browser.autofill_assistant.proto.TellProto;
+import org.chromium.chrome.browser.autofill_assistant.proto.WaitForNavigationProto;
 import org.chromium.chrome.browser.flags.ChromeSwitches;
 import org.chromium.chrome.test.ChromeJUnit4ClassRunner;
 import org.chromium.chrome.test.ChromeTabbedActivityTestRule;
@@ -199,32 +199,45 @@ public class AutofillAssistantNavigationIntegrationTest {
 
     @Test
     @MediumTest
-    @DisabledTest(message = "Flaky - https://crbug.com/1104900")
     public void navigateActionDoesNotCauseError() {
         // Push something to navigation stack so we can use back and forth.
         ChromeTabUtils.loadUrlOnUiThread(
                 mTestRule.getActivity().getActivityTab(), getURL(TEST_PAGE_B));
 
         ArrayList<ActionProto> list = new ArrayList<>();
+        list.add((ActionProto) ActionProto.newBuilder()
+                         .setPrompt(PromptProto.newBuilder().setMessage("Page B").addChoices(
+                                 Choice.newBuilder().setChip(
+                                         ChipProto.newBuilder().setText("Navigate"))))
+                         .build());
         list.add((ActionProto) ActionProto.newBuilder()
                          .setNavigate(NavigateProto.newBuilder().setUrl(getURL(TEST_PAGE_A)))
                          .build());
+        list.add((ActionProto) ActionProto.newBuilder()
+                         .setWaitForNavigation(WaitForNavigationProto.newBuilder())
+                         .build());
         list.add((ActionProto) ActionProto.newBuilder()
                          .setPrompt(PromptProto.newBuilder().setMessage("Page A").addChoices(
                                  Choice.newBuilder().setChip(
-                                         ChipProto.newBuilder().setText("Next"))))
+                                         ChipProto.newBuilder().setText("Go back"))))
                          .build());
         list.add((ActionProto) ActionProto.newBuilder()
                          .setNavigate(NavigateProto.newBuilder().setGoBackward(true))
                          .build());
+        list.add((ActionProto) ActionProto.newBuilder()
+                         .setWaitForNavigation(WaitForNavigationProto.newBuilder())
+                         .build());
         list.add((ActionProto) ActionProto.newBuilder()
                          .setPrompt(PromptProto.newBuilder().setMessage("Page B").addChoices(
                                  Choice.newBuilder().setChip(
-                                         ChipProto.newBuilder().setText("Next"))))
+                                         ChipProto.newBuilder().setText("Go forward"))))
                          .build());
         list.add((ActionProto) ActionProto.newBuilder()
                          .setNavigate(NavigateProto.newBuilder().setGoForward(true))
                          .build());
+        list.add((ActionProto) ActionProto.newBuilder()
+                         .setWaitForNavigation(WaitForNavigationProto.newBuilder())
+                         .build());
         list.add((ActionProto) ActionProto.newBuilder()
                          .setPrompt(PromptProto.newBuilder().setMessage("Page A").addChoices(
                                  PromptProto.Choice.newBuilder()))
@@ -240,17 +253,20 @@ public class AutofillAssistantNavigationIntegrationTest {
         setupScripts(script);
         startAutofillAssistantOnTab(TEST_PAGE_B);
 
+        waitUntilViewMatchesCondition(withText("Page B"), isCompletelyDisplayed());
+        onView(withText("Navigate")).perform(click());
+
         waitUntilViewMatchesCondition(withText("Page A"), isCompletelyDisplayed());
         waitUntil(()
                           -> mTestRule.getActivity().getActivityTab().getUrl().getSpec().equals(
                                   getURL(TEST_PAGE_A)));
-        onView(withText("Next")).perform(click());
+        onView(withText("Go back")).perform(click());
 
         waitUntilViewMatchesCondition(withText("Page B"), isCompletelyDisplayed());
         waitUntil(()
                           -> mTestRule.getActivity().getActivityTab().getUrl().getSpec().equals(
                                   getURL(TEST_PAGE_B)));
-        onView(withText("Next")).perform(click());
+        onView(withText("Go forward")).perform(click());
 
         waitUntilViewMatchesCondition(withText("Page A"), isCompletelyDisplayed());
         waitUntil(()