Extend the delay before exiting test to deflake

Make sure during screen orientation changes happens when the
disclaimer is still on the screen.

Change-Id: Icd6fa16de94b84c8acb8e4666ebfd27074629eca
Bug: 1157611
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2600042Reviewed-by: Theresa  <twellington@chromium.org>
Commit-Queue: Wenyu Fu <wenyufu@chromium.org>
Cr-Commit-Position: refs/heads/master@{#839001}

chrome/android/java/src/org/chromium/chrome/browser/firstrun/TosAndUmaFirstRunFragmentWithEnterpriseSupport.java

@@ -13,6 +13,7 @@ import android.view.accessibility.AccessibilityEvent;
 import android.widget.TextView;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.VisibleForTesting;
 
 import org.chromium.base.Log;
 import org.chromium.base.ThreadUtils;
@@ -31,6 +32,8 @@ public class TosAndUmaFirstRunFragmentWithEnterpriseSupport
         extends ToSAndUMAFirstRunFragment implements LoadingView.Observer {
     private static final String TAG = "TosAndUmaFragment";
 
+    private static Runnable sOverridenOnExitFreRunnableForTest;
+
     /** FRE page that instantiates this fragment. */
     public static class Page
             implements FirstRunPage<TosAndUmaFirstRunFragmentWithEnterpriseSupport> {
@@ -189,11 +192,21 @@ public class TosAndUmaFirstRunFragmentWithEnterpriseSupport
             mPrivacyDisclaimer.announceForAccessibility(mPrivacyDisclaimer.getText());
         }
 
-        mExitFreRunnable = () -> {
-            getPageDelegate().exitFirstRun();
-            mExitFreRunnable = null;
-        };
+        if (sOverridenOnExitFreRunnableForTest != null) {
+            mExitFreRunnable = sOverridenOnExitFreRunnableForTest;
+        } else {
+            mExitFreRunnable = () -> {
+                getPageDelegate().exitFirstRun();
+                mExitFreRunnable = null;
+            };
+        }
+
         mHandler = new Handler(ThreadUtils.getUiThreadLooper());
         mHandler.postDelayed(mExitFreRunnable, FirstRunUtils.getSkipTosExitDelayMs());
     }
+
+    @VisibleForTesting
+    static void setOverrideOnExitFreRunnableForTest(Runnable runnable) {
+        sOverridenOnExitFreRunnableForTest = runnable;
+    }
 }

chrome/android/javatests/src/org/chromium/chrome/browser/firstrun/TosAndUmaFirstRunFragmentWithEnterpriseSupportTest.java

@@ -41,7 +41,6 @@ import org.chromium.base.metrics.RecordHistogram;
 import org.chromium.base.test.util.CallbackHelper;
 import org.chromium.base.test.util.Criteria;
 import org.chromium.base.test.util.CriteriaHelper;
-import org.chromium.base.test.util.DisabledTest;
 import org.chromium.base.test.util.Feature;
 import org.chromium.chrome.R;
 import org.chromium.chrome.browser.customtabs.CustomTabsTestUtils;
@@ -166,6 +165,7 @@ public class TosAndUmaFirstRunFragmentWithEnterpriseSupportTest {
     public void tearDown() {
         FirstRunAppRestrictionInfo.setInitializedInstanceForTest(null);
         ToSAndUMAFirstRunFragment.setShowUmaCheckBoxForTesting(false);
+        TosAndUmaFirstRunFragmentWithEnterpriseSupport.setOverrideOnExitFreRunnableForTest(null);
         PolicyServiceFactory.setPolicyServiceForTest(null);
         FirstRunUtilsJni.TEST_HOOKS.setInstanceForTesting(mFirstRunUtils);
         EnterpriseInfo.setInstanceForTest(null);
@@ -430,8 +430,11 @@ public class TosAndUmaFirstRunFragmentWithEnterpriseSupportTest {
     @Test
     @SmallTest
     @Feature({"RenderTest", "FirstRun"})
-    @DisabledTest(message = "https://crbug.com/1157521")
     public void testRenderWithPolicy() throws Exception {
+        final CallbackHelper onExitFreCallback = new CallbackHelper();
+        TosAndUmaFirstRunFragmentWithEnterpriseSupport.setOverrideOnExitFreRunnableForTest(
+                onExitFreCallback::notifyCalled);
+
         setAppRestrictionsMockInitialized(true);
         setEnterpriseInfoInitializedWithDeviceOwner(true);
         launchFirstRunThroughCustomTab();
@@ -444,10 +447,10 @@ public class TosAndUmaFirstRunFragmentWithEnterpriseSupportTest {
         TestThreadUtils.runOnUiThreadBlocking(tosAndUmaFragment::clearFocus);
 
         setPolicyServiceMockInitializedWithDialogEnabled(false);
-        CriteriaHelper.pollUiThread(
-                ()
-                        -> Criteria.checkThat(
-                                mPrivacyDisclaimer.getVisibility(), Matchers.is(View.VISIBLE)));
+        onExitFreCallback.waitForFirst("OnExitFreCallback is never invoked.");
+        Assert.assertEquals("Privacy disclaimer is not visible", mPrivacyDisclaimer.getVisibility(),
+                View.VISIBLE);
+
         renderWithPortraitAndLandscape(tosAndUmaFragment, "fre_tosanduma_withpolicy");
     }