2 files changed, 68 insertions, 14 deletions

diff --git a/tests/testables/src/android/testing/TestableLooper.java b/tests/testables/src/android/testing/TestableLooper.java
index edd6dd3468ef..82e40b1eee6b 100644
--- a/tests/testables/src/android/testing/TestableLooper.java
+++ b/tests/testables/src/android/testing/TestableLooper.java
@@ -32,6 +32,7 @@ import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 import java.lang.reflect.Field;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * This is a wrapper around {@link TestLooperManager} to make it easier to manage
@@ -55,7 +56,6 @@ public class TestableLooper {
     private MessageHandler mMessageHandler;
 
     private Handler mHandler;
-    private Runnable mEmptyMessage;
     private TestLooperManager mQueueWrapper;
 
     static {
@@ -121,8 +121,12 @@ public class TestableLooper {
      * @param num Number of messages to parse
      */
     public int processMessages(int num) {
+        return processMessagesInternal(num, null);
+    }
+
+    private int processMessagesInternal(int num, Runnable barrierRunnable) {
         for (int i = 0; i < num; i++) {
-            if (!parseMessageInt()) {
+            if (!processSingleMessage(barrierRunnable)) {
                 return i + 1;
             }
         }
@@ -130,6 +134,27 @@ public class TestableLooper {
     }
 
     /**
+     * Process up to a certain number of messages, not blocking if the queue has less messages than
+     * that
+     * @param num the maximum number of messages to process
+     * @return the number of messages processed. This will be at most {@code num}.
+     */
+
+    public int processMessagesNonBlocking(int num) {
+        final AtomicBoolean reachedBarrier = new AtomicBoolean(false);
+        Runnable barrierRunnable = () -> {
+            reachedBarrier.set(true);
+        };
+        mHandler.post(barrierRunnable);
+        waitForMessage(mQueueWrapper, mHandler, barrierRunnable);
+        try {
+            return processMessagesInternal(num, barrierRunnable) + (reachedBarrier.get() ? -1 : 0);
+        } finally {
+            mHandler.removeCallbacks(barrierRunnable);
+        }
+    }
+
+    /**
      * Process messages in the queue until no more are found.
      */
     public void processAllMessages() {
@@ -165,19 +190,20 @@ public class TestableLooper {
 
     private int processQueuedMessages() {
         int count = 0;
-        mEmptyMessage = () -> { };
-        mHandler.post(mEmptyMessage);
-        waitForMessage(mQueueWrapper, mHandler, mEmptyMessage);
-        while (parseMessageInt()) count++;
+        Runnable barrierRunnable = () -> { };
+        mHandler.post(barrierRunnable);
+        waitForMessage(mQueueWrapper, mHandler, barrierRunnable);
+        while (processSingleMessage(barrierRunnable)) count++;
         return count;
     }
 
-    private boolean parseMessageInt() {
+    private boolean processSingleMessage(Runnable barrierRunnable) {
         try {
             Message result = mQueueWrapper.next();
             if (result != null) {
                 // This is a break message.
-                if (result.getCallback() == mEmptyMessage) {
+                if (result.getCallback() == barrierRunnable) {
+                    mQueueWrapper.execute(result);
                     mQueueWrapper.recycle(result);
                     return false;
                 }
diff --git a/tests/testables/tests/src/android/testing/TestableLooperTest.java b/tests/testables/tests/src/android/testing/TestableLooperTest.java
index 0f491b86626c..a02eb6b176dc 100644
--- a/tests/testables/tests/src/android/testing/TestableLooperTest.java
+++ b/tests/testables/tests/src/android/testing/TestableLooperTest.java
@@ -27,12 +27,6 @@ import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.InOrder;
-
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
@@ -40,6 +34,11 @@ import android.test.suitebuilder.annotation.SmallTest;
 import android.testing.TestableLooper.MessageHandler;
 import android.testing.TestableLooper.RunWithLooper;
 
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.InOrder;
+
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
 @RunWithLooper
@@ -240,4 +239,33 @@ public class TestableLooperTest {
         inOrder.verify(handler).dispatchMessage(messageC);
     }
 
+    @Test
+    public void testProcessMessagesNonBlocking_onlyArgNumber() {
+        Handler h = new Handler(mTestableLooper.getLooper());
+        Runnable r = mock(Runnable.class);
+
+        h.post(r);
+        h.post(r);
+        h.post(r);
+
+        int processed = mTestableLooper.processMessagesNonBlocking(2);
+
+        verify(r, times(2)).run();
+        assertEquals(2, processed);
+    }
+
+    @Test
+    public void testProcessMessagesNonBlocking_lessMessagesThanArg() {
+        Handler h = new Handler(mTestableLooper.getLooper());
+        Runnable r = mock(Runnable.class);
+
+        h.post(r);
+        h.post(r);
+        h.post(r);
+
+        int processed = mTestableLooper.processMessagesNonBlocking(5);
+
+        verify(r, times(3)).run();
+        assertEquals(3, processed);
+    }
 }