Fix Espresso.onIdle for non-main Loopers on Baklava.

espresso/core/java/androidx/test/espresso/base/IdlingResourceRegistry.java

@@ -216,6 +216,9 @@ public Boolean call() {
         boolean found = false;
         for (int i = 0; i < idlingStates.size(); i++) {
           if (idlingStates.get(i).resource.getName().equals(resource.getName())) {
+            if (idlingStates.get(i).resource instanceof LooperIdlingResourceInterrogationHandler) {
+              ((LooperIdlingResourceInterrogationHandler) idlingStates.get(i).resource).release();
+            }
             idlingStates.get(i).closeSpan();
             idlingStates.remove(i);
             found = true;

espresso/core/java/androidx/test/espresso/base/Interrogator.java

@@ -92,18 +92,30 @@ interface InterrogationHandler<R> extends QueueInterrogationHandler<R> {
     public String getMessage();
   }
 
+  private final TestLooperManagerCompat testLooperManager;
+
+  static Interrogator acquire(Looper looper) {
+    return new Interrogator(TestLooperManagerCompat.acquire(looper));
+  }
+
+  private Interrogator(TestLooperManagerCompat testLooperManager) {
+    this.testLooperManager = testLooperManager;
+  }
+
+  void release() {
+    testLooperManager.release();
+  }
 
   /**
    * Loops the main thread and informs the interrogation handler at interesting points in the exec
    * state.
    *
    * @param handler an interrogation handler that controls whether to continue looping or not.
    */
-  static <R> R loopAndInterrogate(InterrogationHandler<R> handler) {
+  <T> T loopAndInterrogate(InterrogationHandler<T> handler) {
     checkSanity();
     interrogating.set(Boolean.TRUE);
     boolean stillInterested = true;
-    TestLooperManagerCompat testLooperManager = TestLooperManagerCompat.acquire(Looper.myLooper());
 
     // We may have an identity when we're called - we want to restore it at the end of the fn.
     final long entryIdentity = Binder.clearCallingIdentity();
@@ -112,7 +124,7 @@ static <R> R loopAndInterrogate(InterrogationHandler<R> handler) {
       final long threadIdentity = Binder.clearCallingIdentity();
       while (stillInterested) {
         // run until the observer is no longer interested.
-        stillInterested = interrogateQueueState(testLooperManager, handler);
+        stillInterested = interrogateQueueState(handler);
         if (stillInterested) {
           Message m = testLooperManager.next();
 
@@ -149,7 +161,6 @@ static <R> R loopAndInterrogate(InterrogationHandler<R> handler) {
     } finally {
       Binder.restoreCallingIdentity(entryIdentity);
       interrogating.set(Boolean.FALSE);
-      testLooperManager.release();
     }
     return handler.get();
   }
@@ -168,37 +179,25 @@ static <R> R loopAndInterrogate(InterrogationHandler<R> handler) {
    *     queueEmpty(), taskDueSoon(), taskDueLong() or barrierUp(). once and only once.
    * @return the result of handler.get()
    */
-  static <R> R peekAtQueueState(
-      TestLooperManagerCompat testLooperManager, QueueInterrogationHandler<R> handler) {
-    checkNotNull(testLooperManager);
+  <T> T peekAtQueueState(QueueInterrogationHandler<T> handler) {
     checkNotNull(handler);
     checkState(
-        !interrogateQueueState(testLooperManager, handler),
+        !interrogateQueueState(handler),
         "It is expected that %s would stop interrogation after a single peak at the queue.",
         handler);
     return handler.get();
   }
 
-  static <R> R peekAtQueueState(Looper looper, QueueInterrogationHandler<R> handler) {
-    TestLooperManagerCompat testLooperManager = TestLooperManagerCompat.acquire(looper);
-    try {
-      return peekAtQueueState(testLooperManager, handler);
-    } finally {
-      testLooperManager.release();
-    }
-  }
-
-  private static boolean interrogateQueueState(
-      TestLooperManagerCompat testLooperManager, QueueInterrogationHandler<?> handler) {
+  private boolean interrogateQueueState(QueueInterrogationHandler<?> handler) {
     synchronized (testLooperManager.getQueue()) {
+      if (testLooperManager.isBlockedOnSyncBarrier()) {
+        if (Log.isLoggable(TAG, Log.DEBUG)) {
+          Log.d(TAG, "barrier is up");
+        }
+        return handler.barrierUp();
+      }
       Long headWhen = testLooperManager.peekWhen();
       if (headWhen == null) {
-        if (testLooperManager.isBlockedOnSyncBarrier()) {
-          if (Log.isLoggable(TAG, Log.DEBUG)) {
-            Log.d(TAG, "barrier is up");
-          }
-          return handler.barrierUp();
-        }
         return handler.queueEmpty();
       }
 
@@ -215,7 +214,7 @@ private static boolean interrogateQueueState(
     }
   }
 
-  private static void checkSanity() {
+  private void checkSanity() {
     checkState(Looper.myLooper() != null, "Calling non-looper thread!");
     checkState(Boolean.FALSE.equals(interrogating.get()), "Already interrogating!");
   }

espresso/core/java/androidx/test/espresso/base/LooperIdlingResourceInterrogationHandler.java

@@ -71,6 +71,8 @@ public boolean barrierUp() {
   private volatile Looper looper = null;
   private volatile boolean idle = true;
 
+  private volatile Interrogator interrogator = null;
+
   // written on main - read on looper
   private volatile IdlingResource.ResourceCallback cb = null;
 
@@ -97,8 +99,9 @@ static LooperIdlingResourceInterrogationHandler forLooper(Looper l) {
               @Override
               public void run() {
                 ir.looper = Looper.myLooper();
+                ir.interrogator = Interrogator.acquire(ir.looper);
                 ir.started = true;
-                Interrogator.loopAndInterrogate(ir);
+                ir.interrogator.loopAndInterrogate(ir);
               }
             });
 
@@ -115,6 +118,7 @@ public String getMessage() {
 
   @Override
   public void quitting() {
+    interrogator.release();
     transitionToIdle();
   }
 
@@ -162,7 +166,7 @@ public boolean isIdleNow() {
       // make sure nothing has arrived in the queue while the looper thread is waiting to pull a
       // new task out of it. There can be some delay between a new message entering the queue and
       // the looper thread pulling it out and processing it.
-      return Boolean.FALSE.equals(Interrogator.peekAtQueueState(looper, queueHasNewTasks));
+      return Boolean.FALSE.equals(interrogator.peekAtQueueState(queueHasNewTasks));
     }
     return false;
   }
@@ -183,4 +187,8 @@ private void transitionToIdle() {
       cb.onTransitionToIdle();
     }
   }
+
+  public void release() {
+    interrogator.release();
+  }
 }

espresso/core/java/androidx/test/espresso/base/TestLooperManagerCompat.java

@@ -9,6 +9,7 @@
 import android.os.MessageQueue;
 import android.os.TestLooperManager;
 import androidx.annotation.Nullable;
+import androidx.test.internal.util.Checks;
 import androidx.test.platform.app.InstrumentationRegistry;
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
@@ -77,6 +78,7 @@ private TestLooperManagerCompat(TestLooperManager testLooperManager) {
   static TestLooperManagerCompat acquire(Looper looper) {
     if (peekWhenMethod != null) {
       // running on a newer Android version that has the supported TestLooperManagerCompat changes
+      Checks.checkState(looper.isCurrentThread());
       TestLooperManager testLooperManager =
           InstrumentationRegistry.getInstrumentation().acquireLooperManager(looper);
       return new TestLooperManagerCompat(testLooperManager);

espresso/core/java/androidx/test/espresso/base/UiControllerImpl.java

@@ -156,6 +156,7 @@ private enum InterrogationStatus {
   private IdleNotifier<Runnable> asyncIdle;
   private IdleNotifier<Runnable> compatIdle;
   private Provider<IdleNotifier<IdleNotificationCallback>> dynamicIdleProvider;
+  private Interrogator interrogator;
 
   @VisibleForTesting
   @Inject
@@ -507,7 +508,7 @@ private IdleNotifier<IdleNotificationCallback> loopUntil(
           start + masterIdlePolicy.getIdleTimeoutUnit().toMillis(masterIdlePolicy.getIdleTimeout());
       interrogation = new MainThreadInterrogation(conditions, conditionSet, end);
 
-      InterrogationStatus result = Interrogator.loopAndInterrogate(interrogation);
+      InterrogationStatus result = getInterrogator().loopAndInterrogate(interrogation);
       if (InterrogationStatus.COMPLETED == result) {
         // did not time out, all conditions happy.
         return dynamicIdle;
@@ -585,6 +586,13 @@ private IdleNotifier<IdleNotificationCallback> loopUntil(
     return dynamicIdle;
   }
 
+  private Interrogator getInterrogator() {
+    if (interrogator == null) {
+      interrogator = Interrogator.acquire(mainLooper);
+    }
+    return interrogator;
+  }
+
   @Override
   public void interruptEspressoTasks() {
     controllerHandler.post(

espresso/core/javatests/androidx/test/espresso/base/EspressoIdleTest.java

@@ -4,7 +4,10 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import android.os.Handler;
+import android.os.HandlerThread;
+import android.os.Looper;
 import androidx.test.espresso.Espresso;
+import androidx.test.espresso.IdlingRegistry;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.junit.Test;
@@ -43,4 +46,22 @@ public void onIdle_afterPostLongDelay() {
     Espresso.onIdle();
     assertThat(wasRun.get()).isFalse();
   }
+
+  @Test
+  public void onIdle_afterPost_backgroundLooper() {
+    HandlerThread ht = new HandlerThread("onIdle_afterPost_backgroundLooper");
+    ht.start();
+    Looper looper = ht.getLooper();
+
+    try {
+      IdlingRegistry.getInstance().registerLooperAsIdlingResource(looper);
+      AtomicBoolean wasRun = new AtomicBoolean(false);
+      new Handler(looper).post(() -> wasRun.set(true));
+      Espresso.onIdle();
+      assertThat(wasRun.get()).isTrue();
+    } finally {
+      IdlingRegistry.getInstance().unregisterLooperAsIdlingResource(looper);
+      ht.quit();
+    }
+  }
 }