Reland "Introduce CriteriaHelper#pollUiThreadNested"

Introduce CriteriaHelper#pollUiThreadNested
This change adds an assertion that pollUiThread is not called from the
UI thread, which is equivalent to calling sleep(), and won't actually
run any UI thread tasks.

In the few cases where polling the UI Thread from the UI thread is
necessary, I introduce CriteriaHelper#pollUiThreadNested, which nests
the looper in order to continue running UI tasks.

Bug: 1115220
Change-Id: I9338a607590e8ee0a22bc32477f02867c6e37271
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2357289Reviewed-by: Yaron Friedman <yfriedman@chromium.org>
Commit-Queue: Michael Thiessen <mthiesse@chromium.org>
Cr-Commit-Position: refs/heads/master@{#798862}

chrome/android/BUILD.gn

@@ -3262,11 +3262,7 @@ android_library("native_java_unittests_java") {
     "//url/mojom:url_mojom_gurl_java",
   ]
 
-  sources = [
-    "native_java_unittests/src/org/chromium/chrome/browser/UnitTestUtils.java",
-  ]
-
-  sources += native_java_unittests_tests
+  sources = native_java_unittests_tests
 
   annotation_processor_deps = [ "//base/android/jni_generator:jni_processor" ]
 }

chrome/android/javatests/src/org/chromium/chrome/browser/bookmarks/BookmarkPersonalizedSigninPromoDismissTest.java

@@ -63,8 +63,8 @@ public class BookmarkPersonalizedSigninPromoDismissTest {
             BookmarkModel bookmarkModel = new BookmarkModel(Profile.fromWebContents(
                     mSyncTestRule.getActivity().getActivityTab().getWebContents()));
             bookmarkModel.loadFakePartnerBookmarkShimForTesting();
-            BookmarkTestUtil.waitForBookmarkModelLoaded();
         });
+        BookmarkTestUtil.waitForBookmarkModelLoaded();
     }
 
     @After

chrome/android/javatests/src/org/chromium/chrome/browser/compositor/layouts/LayoutManagerTest.java

@@ -686,12 +686,9 @@ public class LayoutManagerTest implements MockTabModelDelegate {
         CriteriaHelper.pollUiThread(
                 mActivityTestRule.getActivity().getTabModelSelector()::isTabStateInitialized);
 
-        TestThreadUtils.runOnUiThreadBlocking(() -> {
-            LayoutManagerChrome layoutManager = mActivityTestRule.getActivity().getLayoutManager();
-            layoutManager.showOverview(false);
-
-            CriteriaHelper.pollUiThread(layoutManager::overviewVisible);
-        });
+        LayoutManagerChrome layoutManager = mActivityTestRule.getActivity().getLayoutManager();
+        TestThreadUtils.runOnUiThreadBlocking(() -> layoutManager.showOverview(false));
+        CriteriaHelper.pollUiThread(layoutManager::overviewVisible);
     }
 
     private Layout getActiveLayout() {

chrome/android/javatests/src/org/chromium/chrome/browser/ntp/cards/promo/HomepagePromoTest.java

@@ -31,6 +31,7 @@ import androidx.test.espresso.contrib.RecyclerViewActions;
 import androidx.test.filters.MediumTest;
 import androidx.test.filters.SmallTest;
 
+import org.hamcrest.Matchers;
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
@@ -493,16 +494,23 @@ public class HomepagePromoTest {
     private void scrollToHomepagePromo() {
         onView(instanceOf(RecyclerView.class))
                 .perform(RecyclerViewActions.scrollToPosition(NTP_HEADER_POSITION + 1));
-        waitForView((ViewGroup) mActivityTestRule.getActivity().findViewById(R.id.homepage_promo),
-                allOf(withId(R.id.promo_primary_button), isDisplayed()));
+        TestThreadUtils.runOnUiThreadBlocking(() -> {
+            waitForView(
+                    (ViewGroup) mActivityTestRule.getActivity().findViewById(R.id.homepage_promo),
+                    allOf(withId(R.id.promo_primary_button), isDisplayed()));
+        });
 
-        // Verify impress tracking metrics is working.
-        Assert.assertEquals("Promo created should be seen.", 1,
-                RecordHistogram.getHistogramValueCountForTesting(
-                        METRICS_HOMEPAGE_PROMO, HomepagePromoAction.SEEN));
-        Assert.assertEquals("Impression should be tracked in shared preference.", 1,
-                SharedPreferencesManager.getInstance().readInt(
-                        HomepagePromoUtils.getTimesSeenKey()));
+        CriteriaHelper.pollUiThread(() -> {
+            // Verify impress tracking metrics is working.
+            Criteria.checkThat("Promo created should be seen.",
+                    RecordHistogram.getHistogramValueCountForTesting(
+                            METRICS_HOMEPAGE_PROMO, HomepagePromoAction.SEEN),
+                    Matchers.is(1));
+            Criteria.checkThat("Impression should be tracked in shared preference.",
+                    SharedPreferencesManager.getInstance().readInt(
+                            HomepagePromoUtils.getTimesSeenKey()),
+                    Matchers.is(1));
+        });
         Mockito.verify(mTracker).notifyEvent(EventConstants.HOMEPAGE_PROMO_SEEN);
     }
 
@@ -533,8 +541,10 @@ public class HomepagePromoTest {
         // screen.
         onView(instanceOf(RecyclerView.class))
                 .perform(RecyclerViewActions.scrollToPosition(feedHeaderPosition));
-        waitForView(rootView,
-                allOf(withId(R.id.header_status),
-                        withText(expanded ? R.string.hide : R.string.show)));
+        TestThreadUtils.runOnUiThreadBlocking(() -> {
+            waitForView(rootView,
+                    allOf(withId(R.id.header_status),
+                            withText(expanded ? R.string.hide : R.string.show)));
+        });
     }
 }

chrome/android/javatests/src/org/chromium/chrome/browser/signin/SigninSignoutIntegrationTest.java

@@ -216,7 +216,9 @@ public class SigninSignoutIntegrationTest {
             mBookmarkModel = new BookmarkModel(Profile.fromWebContents(
                     mActivityTestRule.getActivity().getActivityTab().getWebContents()));
             mBookmarkModel.loadFakePartnerBookmarkShimForTesting();
-            BookmarkTestUtil.waitForBookmarkModelLoaded();
+        });
+        BookmarkTestUtil.waitForBookmarkModelLoaded();
+        TestThreadUtils.runOnUiThreadBlocking(() -> {
             Assert.assertEquals(0, mBookmarkModel.getChildCount(mBookmarkModel.getDefaultFolder()));
             mBookmarkModel.addBookmark(
                     mBookmarkModel.getDefaultFolder(), 0, "Test Bookmark", "http://google.com");

chrome/android/native_java_unittests/src/org/chromium/chrome/browser/UnitTestUtils.java

-// Copyright 2020 The Chromium Authors. All rights reserved.
-// Use of this source code is governed by a BSD-style license that can be
-// found in the LICENSE file.
-
-package org.chromium.chrome.browser;
-
-import android.os.Handler;
-import android.os.Looper;
-
-import org.junit.Assert;
-
-import org.chromium.base.ThreadUtils;
-import org.chromium.base.test.util.TimeoutTimer;
-import org.chromium.content_public.browser.test.NestedSystemMessageHandler;
-import org.chromium.content_public.browser.test.util.CriteriaHelper;
-
-import java.util.concurrent.Callable;
-import java.util.concurrent.atomic.AtomicBoolean;
-
-/**
- * Utilities for use in Native Java Unit Tests.
- */
-public class UnitTestUtils {
-    /**
-     * Polls the UI thread waiting for the Callable |criteria| to return true.
-     *
-     * In practice, this nests the looper and checks the criteria after each task is run as these
-     * tests run on the UI thread and sleeping would block the thing we're waiting for.
-     */
-    public static void pollUiThread(final Callable<Boolean> criteria) throws Exception {
-        assert ThreadUtils.runningOnUiThread();
-        boolean isSatisfied = criteria.call();
-        TimeoutTimer timer = new TimeoutTimer(CriteriaHelper.DEFAULT_MAX_TIME_TO_POLL);
-        Handler handler = new Handler(Looper.myLooper());
-        AtomicBoolean called = new AtomicBoolean(true);
-
-        while (!isSatisfied && !timer.isTimedOut()) {
-            // Ensure we pump the message handler in case no new tasks arrive.
-            if (called.get()) {
-                called.set(false);
-                handler.postDelayed(
-                        () -> { called.set(true); }, CriteriaHelper.DEFAULT_POLLING_INTERVAL);
-            }
-
-            NestedSystemMessageHandler.runSingleNestedLooperTask(Looper.myQueue());
-            isSatisfied = criteria.call();
-        }
-        Assert.assertFalse("Timed out waiting for condition", timer.isTimedOut());
-        Assert.assertTrue(isSatisfied);
-    }
-}
\ No newline at end of file

chrome/android/native_java_unittests/src/org/chromium/chrome/browser/installedapp/InstalledAppProviderTest.java

@@ -18,9 +18,9 @@ import org.chromium.base.Callback;
 import org.chromium.base.annotations.CalledByNative;
 import org.chromium.base.annotations.CalledByNativeJavaTest;
 import org.chromium.base.annotations.DisabledCalledByNativeJavaTest;
-import org.chromium.chrome.browser.UnitTestUtils;
 import org.chromium.chrome.browser.instantapps.InstantAppsHandler;
 import org.chromium.chrome.browser.profiles.Profile;
+import org.chromium.content_public.browser.test.util.CriteriaHelper;
 import org.chromium.installedapp.mojom.InstalledAppProvider;
 import org.chromium.installedapp.mojom.RelatedApplication;
 import org.chromium.url.GURL;
@@ -372,7 +372,7 @@ public class InstalledAppProviderTest {
                         called.set(true);
                     }
                 });
-        UnitTestUtils.pollUiThread(() -> called.get());
+        CriteriaHelper.pollUiThreadNested(() -> called.get());
     }
 
     @CalledByNative

chrome/android/native_java_unittests/src/org/chromium/chrome/browser/payments/AndroidPaymentAppUnitTest.java

@@ -14,10 +14,10 @@ import org.mockito.MockitoAnnotations;
 
 import org.chromium.base.annotations.CalledByNative;
 import org.chromium.base.annotations.CalledByNativeJavaTest;
-import org.chromium.chrome.browser.UnitTestUtils;
 import org.chromium.components.payments.PayerData;
 import org.chromium.components.payments.PaymentApp;
 import org.chromium.components.payments.SupportedDelegations;
+import org.chromium.content_public.browser.test.util.CriteriaHelper;
 import org.chromium.payments.mojom.PaymentCurrencyAmount;
 import org.chromium.payments.mojom.PaymentDetailsModifier;
 import org.chromium.payments.mojom.PaymentItem;
@@ -93,7 +93,7 @@ public class AndroidPaymentAppUnitTest {
                         mReadyToPayResponse = isReadyToPay;
                     }
                 });
-        UnitTestUtils.pollUiThread(() -> mReadyToPayQueryFinished);
+        CriteriaHelper.pollUiThreadNested(() -> mReadyToPayQueryFinished);
         Assert.assertTrue("Payment app should be ready to pay", mReadyToPayResponse);
 
         PaymentItem total = new PaymentItem();
@@ -131,6 +131,6 @@ public class AndroidPaymentAppUnitTest {
         intentResult.data.putExtras(extras);
         app.onIntentCompletedForTesting(intentResult);
 
-        UnitTestUtils.pollUiThread(() -> mInvokePaymentAppFinished);
+        CriteriaHelper.pollUiThreadNested(() -> mInvokePaymentAppFinished);
     }
 }

chrome/android/native_java_unittests/src/org/chromium/chrome/browser/tabmodel/TabPersistentStoreUnitTest.java

@@ -27,7 +27,6 @@ import org.chromium.base.annotations.CalledByNative;
 import org.chromium.base.annotations.CalledByNativeJavaTest;
 import org.chromium.base.task.TaskRunner;
 import org.chromium.base.test.util.Feature;
-import org.chromium.chrome.browser.UnitTestUtils;
 import org.chromium.chrome.browser.tab.Tab;
 import org.chromium.chrome.browser.tab.TabImpl;
 import org.chromium.chrome.browser.tab.TabLaunchType;
@@ -37,6 +36,7 @@ import org.chromium.chrome.browser.tabmodel.TabPersistentStore.TabPersistentStor
 import org.chromium.chrome.browser.tabmodel.TabPersistentStore.TabRestoreDetails;
 import org.chromium.components.embedder_support.util.UrlConstants;
 import org.chromium.content_public.browser.LoadUrlParams;
+import org.chromium.content_public.browser.test.util.CriteriaHelper;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -107,7 +107,7 @@ public class TabPersistentStoreUnitTest {
         final AtomicBoolean flushed = new AtomicBoolean(false);
         if (mPersistentStore != null) {
             mPersistentStore.getTaskRunnerForTests().postTask(() -> { flushed.set(true); });
-            UnitTestUtils.pollUiThread(() -> flushed.get());
+            CriteriaHelper.pollUiThreadNested(() -> flushed.get());
         }
     }
 

chrome/test/android/javatests/src/org/chromium/chrome/test/util/ViewUtils.java

@@ -113,7 +113,7 @@ public class ViewUtils {
      */
     public static void waitForView(
             ViewGroup root, Matcher<View> viewMatcher, @ExpectedViewState int viewState) {
-        CriteriaHelper.pollUiThread(new ExpectedViewCriteria(viewMatcher, viewState, root));
+        CriteriaHelper.pollUiThreadNested(new ExpectedViewCriteria(viewMatcher, viewState, root));
     }
 
     /**
@@ -133,7 +133,7 @@ public class ViewUtils {
             Matcher<View> viewMatcher, @ExpectedViewState int viewState) {
         return (View view, NoMatchingViewException noMatchException) -> {
             if (noMatchException != null) throw noMatchException;
-            CriteriaHelper.pollUiThread(
+            CriteriaHelper.pollUiThreadNested(
                     new ExpectedViewCriteria(viewMatcher, viewState, (ViewGroup) view));
         };
     }
@@ -167,12 +167,13 @@ public class ViewUtils {
      * @return An interaction on the matching view.
      */
     public static ViewInteraction onViewWaiting(Matcher<View> viewMatcher) {
-        Espresso.onView(ViewMatchers.isRoot())
-                .check((View view, NoMatchingViewException noMatchException) -> {
-                    if (noMatchException != null) throw noMatchException;
-                    CriteriaHelper.pollUiThread(
-                            new ExpectedViewCriteria(viewMatcher, VIEW_VISIBLE, (ViewGroup) view));
-                });
+        CriteriaHelper.pollInstrumentationThread(() -> {
+            Espresso.onView(ViewMatchers.isRoot())
+                    .check((View view, NoMatchingViewException noMatchException) -> {
+                        if (noMatchException != null) throw noMatchException;
+                        new ExpectedViewCriteria(viewMatcher, VIEW_VISIBLE, (ViewGroup) view).run();
+                    });
+        });
         return Espresso.onView(viewMatcher);
     }
 

content/public/android/BUILD.gn

@@ -481,6 +481,7 @@ android_library("content_javatests") {
     "javatests/src/org/chromium/content/browser/ContentViewLocationTest.java",
     "javatests/src/org/chromium/content/browser/ContentViewPointerTypeTest.java",
     "javatests/src/org/chromium/content/browser/ContentViewScrollingTest.java",
+    "javatests/src/org/chromium/content/browser/CriteriaHelperTest.java",
     "javatests/src/org/chromium/content/browser/EncodeHtmlDataUriTest.java",
     "javatests/src/org/chromium/content/browser/GestureDetectorResetTest.java",
     "javatests/src/org/chromium/content/browser/GestureListenerManagerTest.java",
@@ -542,7 +543,6 @@ junit_binary("content_junit_tests") {
   sources = [
     "junit/src/org/chromium/content/browser/BindingManagerTest.java",
     "junit/src/org/chromium/content/browser/ChildProcessRankingTest.java",
-    "junit/src/org/chromium/content/browser/CriteriaHelperTest.java",
     "junit/src/org/chromium/content/browser/ScreenOrientationProviderImplTest.java",
     "junit/src/org/chromium/content/browser/SpareChildConnectionTest.java",
     "junit/src/org/chromium/content/browser/UiThreadTaskTraitsImplTest.java",

content/public/android/junit/src/org/chromium/content/browser/CriteriaHelperTest.java → content/public/android/javatests/src/org/chromium/content/browser/CriteriaHelperTest.java

@@ -5,11 +5,8 @@
 package org.chromium.content.browser;
 
 import static org.hamcrest.CoreMatchers.containsString;
-import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThat;
-import static org.junit.Assert.assertTrue;
 
-import static org.chromium.base.task.TaskTraits.THREAD_POOL;
 import static org.chromium.content_public.browser.test.util.CriteriaHelper.DEFAULT_POLLING_INTERVAL;
 
 import androidx.test.filters.MediumTest;
@@ -22,8 +19,9 @@ import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 
 import org.chromium.base.ThreadUtils;
-import org.chromium.base.task.PostTask;
-import org.chromium.base.test.BaseRobolectricTestRunner;
+import org.chromium.base.test.BaseJUnit4ClassRunner;
+import org.chromium.base.test.UiThreadTest;
+import org.chromium.base.test.util.Batch;
 import org.chromium.content_public.browser.test.util.Criteria;
 import org.chromium.content_public.browser.test.util.CriteriaHelper;
 import org.chromium.content_public.browser.test.util.CriteriaNotSatisfiedException;
@@ -34,7 +32,8 @@ import java.io.StringWriter;
 /**
  * Tests for {@link CriteriaHelper}.
  */
-@RunWith(BaseRobolectricTestRunner.class)
+@RunWith(BaseJUnit4ClassRunner.class)
+@Batch(Batch.UNIT_TESTS)
 public class CriteriaHelperTest {
     private static final String ERROR_MESSAGE = "my special error message";
 
@@ -44,29 +43,23 @@ public class CriteriaHelperTest {
     @Test
     @MediumTest
     public void testUiThread() {
-        // Robolectric runs the test on UI thread.
-        assertTrue(ThreadUtils.runningOnUiThread());
-        // In Instrumented tests, the tests would be on instrumentation thread instead.
-        // Emulate that behavior by posting the body of the test.
-        PostTask.postTask(THREAD_POOL, () -> {
-            assertFalse(ThreadUtils.runningOnUiThread());
-            CriteriaHelper.pollUiThread(
-                    () -> Criteria.checkThat(ThreadUtils.runningOnUiThread(), Matchers.is(true)));
-        });
+        CriteriaHelper.pollUiThread(
+                () -> Criteria.checkThat(ThreadUtils.runningOnUiThread(), Matchers.is(true)));
+    }
+
+    @Test
+    @MediumTest
+    @UiThreadTest
+    public void testUiThreadNested() {
+        CriteriaHelper.pollUiThreadNested(
+                () -> Criteria.checkThat(ThreadUtils.runningOnUiThread(), Matchers.is(true)));
     }
 
     @Test
     @MediumTest
     public void testInstrumentationThread() {
-        // Robolectric runs the test on UI thread.
-        assertTrue(ThreadUtils.runningOnUiThread());
-        // In Instrumented tests, the tests would be on instrumentation thread instead.
-        // Emulate that behavior by posting the body of the test.
-        PostTask.postTask(THREAD_POOL, () -> {
-            assertFalse(ThreadUtils.runningOnUiThread());
-            CriteriaHelper.pollInstrumentationThread(
-                    () -> Criteria.checkThat(ThreadUtils.runningOnUiThread(), Matchers.is(false)));
-        });
+        CriteriaHelper.pollInstrumentationThread(
+                () -> Criteria.checkThat(ThreadUtils.runningOnUiThread(), Matchers.is(false)));
     }
 
     @Test
@@ -75,6 +68,13 @@ public class CriteriaHelperTest {
         CriteriaHelper.pollUiThread(() -> {});
     }
 
+    @Test
+    @MediumTest
+    @UiThreadTest
+    public void testPass_Runnable_UiThreadNested() {
+        CriteriaHelper.pollUiThreadNested(() -> {});
+    }
+
     @Test
     @MediumTest
     public void testPass_Runnable_InstrumentationThread() {
@@ -87,6 +87,13 @@ public class CriteriaHelperTest {
         CriteriaHelper.pollUiThread(() -> true);
     }
 
+    @Test
+    @MediumTest
+    @UiThreadTest
+    public void testPass_Callable_UiThreadNested() {
+        CriteriaHelper.pollUiThreadNested(() -> true);
+    }
+
     @Test
     @MediumTest
     public void testPass_Callable_InstrumentationThread() {
@@ -102,6 +109,16 @@ public class CriteriaHelperTest {
         }, 0, DEFAULT_POLLING_INTERVAL);
     }
 
+    @Test
+    @MediumTest
+    @UiThreadTest
+    public void testThrow_Runnable_UiThreadNested() {
+        thrown.expect(AssertionError.class);
+        CriteriaHelper.pollUiThreadNested(() -> {
+            throw new CriteriaNotSatisfiedException("");
+        }, 0, DEFAULT_POLLING_INTERVAL);
+    }
+
     @Test
     @MediumTest
     public void testThrow_Runnable_InstrumentationThread() {
@@ -118,6 +135,14 @@ public class CriteriaHelperTest {
         CriteriaHelper.pollUiThread(() -> false, 0, DEFAULT_POLLING_INTERVAL);
     }
 
+    @Test
+    @MediumTest
+    @UiThreadTest
+    public void testThrow_Callable_UiThreadNested() {
+        thrown.expect(AssertionError.class);
+        CriteriaHelper.pollUiThreadNested(() -> false, 0, DEFAULT_POLLING_INTERVAL);
+    }
+
     @Test
     @MediumTest
     public void testThrow_Callable_InstrumentationThread() {
@@ -134,6 +159,16 @@ public class CriteriaHelperTest {
         }, 0, DEFAULT_POLLING_INTERVAL);
     }
 
+    @Test
+    @MediumTest
+    @UiThreadTest
+    public void testMessage_Runnable_UiThreadNested() {
+        thrown.expectMessage(ERROR_MESSAGE);
+        CriteriaHelper.pollUiThreadNested(() -> {
+            throw new CriteriaNotSatisfiedException(ERROR_MESSAGE);
+        }, 0, DEFAULT_POLLING_INTERVAL);
+    }
+
     @Test
     @MediumTest
     public void testMessage_Runnable_InstrumentationThread() {
@@ -179,6 +214,22 @@ public class CriteriaHelperTest {
         Assert.fail();
     }
 
+    @Test
+    @MediumTest
+    @UiThreadTest
+    public void testStack_Runnable_UiThreadNested() {
+        try {
+            CriteriaHelper.pollUiThreadNested(() -> {
+                throw new CriteriaNotSatisfiedException("test");
+            }, 0, DEFAULT_POLLING_INTERVAL);
+        } catch (AssertionError e) {
+            assertThat(getStackTrace(e),
+                    containsString("CriteriaHelperTest.testStack_Runnable_UiThreadNested("));
+            return;
+        }
+        Assert.fail();
+    }
+
     @Test
     @MediumTest
     public void testStack_Runnable_InstrumentationThread() {
@@ -207,6 +258,20 @@ public class CriteriaHelperTest {
         Assert.fail();
     }
 
+    @Test
+    @MediumTest
+    @UiThreadTest
+    public void testStack_Callable_UiThreadNested() {
+        try {
+            CriteriaHelper.pollUiThreadNested(() -> false, 0, DEFAULT_POLLING_INTERVAL);
+        } catch (AssertionError e) {
+            assertThat(getStackTrace(e),
+                    containsString("CriteriaHelperTest.testStack_Callable_UiThreadNested("));
+            return;
+        }
+        Assert.fail();
+    }
+
     @Test
     @MediumTest
     public void testStack_Callable_InstrumentationThread() {

content/public/test/android/javatests/src/org/chromium/content_public/browser/test/util/CriteriaHelper.java

@@ -4,12 +4,18 @@
 
 package org.chromium.content_public.browser.test.util;
 
+import android.os.Handler;
+import android.os.Looper;
+
 import org.hamcrest.Matchers;
 
 import org.chromium.base.ThreadUtils;
 import org.chromium.base.test.util.TimeoutTimer;
+import org.chromium.content_public.browser.test.NestedSystemMessageHandler;
 
+import java.lang.reflect.InvocationTargetException;
 import java.util.concurrent.Callable;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
 /**
@@ -73,6 +79,12 @@ public class CriteriaHelper {
      */
     public static void pollInstrumentationThread(
             Runnable criteria, long maxTimeoutMs, long checkIntervalMs) {
+        assert !ThreadUtils.runningOnUiThread();
+        pollThreadInternal(criteria, maxTimeoutMs, checkIntervalMs, false);
+    }
+
+    private static void pollThreadInternal(
+            Runnable criteria, long maxTimeoutMs, long checkIntervalMs, boolean isUiThread) {
         CriteriaNotSatisfiedException throwable;
         try {
             criteria.run();
@@ -82,11 +94,10 @@ public class CriteriaHelper {
         }
         TimeoutTimer timer = new TimeoutTimer(maxTimeoutMs);
         while (!timer.isTimedOut()) {
-            try {
-                Thread.sleep(checkIntervalMs);
-            } catch (InterruptedException e) {
-                // Catch the InterruptedException. If the exception occurs before maxTimeoutMs
-                // and the criteria is not satisfied, the while loop will run again.
+            if (isUiThread) {
+                loopUiThread(checkIntervalMs);
+            } else {
+                sleepInstrumentationThread(checkIntervalMs);
             }
             try {
                 criteria.run();
@@ -98,6 +109,34 @@ public class CriteriaHelper {
         throw new AssertionError(throwable);
     }
 
+    private static void sleepInstrumentationThread(long checkIntervalMs) {
+        assert !ThreadUtils.runningOnUiThread();
+        try {
+            Thread.sleep(checkIntervalMs);
+        } catch (InterruptedException e) {
+            // Catch the InterruptedException. If the exception occurs before maxTimeoutMs
+            // and the criteria is not satisfied, the while loop will run again.
+        }
+    }
+
+    private static void loopUiThread(long checkIntervalMs) {
+        assert ThreadUtils.runningOnUiThread();
+        AtomicBoolean called = new AtomicBoolean(false);
+
+        // Ensure we pump the message handler in case no new tasks arrive.
+        new Handler(Looper.myLooper()).postDelayed(() -> { called.set(true); }, checkIntervalMs);
+
+        TimeoutTimer timer = new TimeoutTimer(checkIntervalMs);
+        while (!timer.isTimedOut() && !called.get()) {
+            try {
+                NestedSystemMessageHandler.runSingleNestedLooperTask(Looper.myQueue());
+            } catch (IllegalArgumentException | IllegalAccessException | SecurityException
+                    | InvocationTargetException e) {
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
     /**
      * Checks whether the given Runnable completes without exception at the default interval.
      *
@@ -193,6 +232,7 @@ public class CriteriaHelper {
      */
     public static void pollUiThread(
             final Runnable criteria, long maxTimeoutMs, long checkIntervalMs) {
+        assert !ThreadUtils.runningOnUiThread();
         pollInstrumentationThread(() -> {
             AtomicReference<Throwable> throwableRef = new AtomicReference<>();
             ThreadUtils.runOnUiThreadBlocking(() -> {
@@ -283,6 +323,63 @@ public class CriteriaHelper {
         pollUiThread(criteria, null);
     }
 
+    /**
+     * Checks whether the given Runnable completes without exception at a given interval on the UI
+     * thread, until either the Runnable successfully completes, or the maxTimeoutMs number of ms
+     * has elapsed.
+     * This call will nest the Looper in order to wait for the Runnable to complete.
+     *
+     * @param criteria The Runnable that will be attempted.
+     * @param maxTimeoutMs The maximum number of ms that this check will be performed for
+     *                     before timeout.
+     * @param checkIntervalMs The number of ms between checks.
+     *
+     * @see #pollInstrumentationThread(Runnable)
+     */
+    public static void pollUiThreadNested(
+            Runnable criteria, long maxTimeoutMs, long checkIntervalMs) {
+        pollThreadInternal(criteria, maxTimeoutMs, checkIntervalMs, true);
+    }
+
+    /**
+     * Checks whether the given Runnable is satisfied polling at a given interval on the UI
+     * thread, until either the criteria is satisfied, or the maxTimeoutMs number of ms has elapsed.
+     * This call will nest the Looper in order to wait for the Criteria to be satisfied.
+     *
+     * @param criteria The Callable<Boolean> that will be checked.
+     * @param maxTimeoutMs The maximum number of ms that this check will be performed for
+     *                     before timeout.
+     * @param checkIntervalMs The number of ms between checks.
+     *
+     * @see #pollInstrumentationThread(Criteria)
+     */
+    public static void pollUiThreadNested(
+            final Callable<Boolean> criteria, long maxTimeoutMs, long checkIntervalMs) {
+        pollUiThreadNested(toNotSatisfiedRunnable(criteria, null), maxTimeoutMs, checkIntervalMs);
+    }
+
+    /**
+     * Checks whether the given Runnable completes without exception at the default interval on
+     * the UI thread. This call will nest the Looper in order to wait for the Runnable to complete.
+     * @param criteria The Runnable that will be attempted.
+     *
+     * @see #pollInstrumentationThread(Runnable)
+     */
+    public static void pollUiThreadNested(final Runnable criteria) {
+        pollUiThreadNested(criteria, DEFAULT_MAX_TIME_TO_POLL, DEFAULT_POLLING_INTERVAL);
+    }
+
+    /**
+     * Checks whether the given Callable<Boolean> is satisfied polling at a default interval on the
+     * UI thread. This call will nest the Looper in order to wait for the Criteria to be satisfied.
+     * @param criteria The Callable<Boolean> that will be checked.
+     *
+     * @see #pollInstrumentationThread(Criteria)
+     */
+    public static void pollUiThreadNested(final Callable<Boolean> criteria) {
+        pollUiThreadNested(toNotSatisfiedRunnable(criteria, null));
+    }
+
     private static Runnable toNotSatisfiedRunnable(
             Callable<Boolean> criteria, String failureReason) {
         return () -> {