🔍 Don't display the search engine until First Run completes

chrome/android/sync_shell/javatests/src/org/chromium/chrome/browser/sync/FirstRunTest.java

Index: chrome/android/sync_shell/javatests/src/org/chromium/chrome/browser/sync/FirstRunTest.java
diff --git a/chrome/android/sync_shell/javatests/src/org/chromium/chrome/browser/sync/FirstRunTest.java b/chrome/android/sync_shell/javatests/src/org/chromium/chrome/browser/sync/FirstRunTest.java
index a6f9789da2f357166a858f6735fdd3b37c3013d3..bf4bb6fb8369fb532cfd76fc476946631a1ffec6 100644
--- a/chrome/android/sync_shell/javatests/src/org/chromium/chrome/browser/sync/FirstRunTest.java
+++ b/chrome/android/sync_shell/javatests/src/org/chromium/chrome/browser/sync/FirstRunTest.java
@@ -14,12 +14,14 @@ import android.os.Bundle;
 import android.support.test.filters.SmallTest;
 
 import org.chromium.base.ThreadUtils;
+import org.chromium.base.test.util.CallbackHelper;
 import org.chromium.base.test.util.CommandLineFlags;
 import org.chromium.base.test.util.Feature;
 import org.chromium.base.test.util.FlakyTest;
 import org.chromium.base.test.util.RetryOnFailure;
 import org.chromium.chrome.browser.ChromeSwitches;
 import org.chromium.chrome.browser.firstrun.FirstRunActivity;
+import org.chromium.chrome.browser.firstrun.FirstRunActivity.FirstRunActivityObserver;
 import org.chromium.chrome.browser.firstrun.FirstRunFlowSequencer;
 import org.chromium.chrome.browser.firstrun.FirstRunSignInProcessor;
 import org.chromium.chrome.browser.preferences.Preferences;
@@ -30,6 +32,8 @@ import org.chromium.chrome.test.util.browser.sync.SyncTestUtil;
 import org.chromium.content.browser.test.util.Criteria;
 import org.chromium.content.browser.test.util.CriteriaHelper;
 
+import java.util.concurrent.TimeoutException;
+
 /**
  * Tests for the first run experience.
  */
@@ -43,10 +47,34 @@ public class FirstRunTest extends SyncTestBase {
         NO;
     }
 
+    private static final class TestObserver implements FirstRunActivityObserver {
+        public final CallbackHelper flowIsKnownCallback = new CallbackHelper();
+
+        @Override
+        public void onFlowIsKnown() {
+            flowIsKnownCallback.notifyCalled();
+        }
+
+        @Override
+        public void onAcceptTermsOfService() {}
+
+        @Override
+        public void onJumpToPage(int position) {}
+
+        @Override
+        public void onUpdateCachedEngineName() {}
+
+        @Override
+        public void onAbortFirstRunExperience() {}
+    }
+
+    private final TestObserver mTestObserver = new TestObserver();
     private FirstRunActivity mActivity;
 
     @Override
     public void startMainActivity() throws InterruptedException {
+        FirstRunActivity.setObserverForTest(mTestObserver);
+
         // Starts up and waits for the FirstRunActivity to be ready.
         // This isn't exactly what startMainActivity is supposed to be doing, but short of a
         // refactoring of SyncTestBase to use something other than ChromeTabbedActivity, it's the
@@ -78,17 +106,21 @@ public class FirstRunTest extends SyncTestBase {
         assertTrue(activity instanceof FirstRunActivity);
         mActivity = (FirstRunActivity) activity;
 
-        CriteriaHelper.pollUiThread(new Criteria() {
-            @Override
-            public boolean isSatisfied() {
-                return mActivity.isPostNativePageSequenceCreated();
-            }
-        });
+        try {
+            mTestObserver.flowIsKnownCallback.waitForCallback(0);
+        } catch (TimeoutException e) {
+            fail();
+        }
 
         getInstrumentation().waitForIdleSync();
     }
 
     @Override
+    public void setUp() throws Exception {
+        super.setUp();
+    }
+
+    @Override
     public void tearDown() throws Exception {
         if (mActivity != null) mActivity.finish();
         super.tearDown();