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

LooperIdlingResourceInterrogationHandler, which is used for non main Loopers,
requires calling peekAtQueueState from the main Looper thread. The previous Interrogator implementation,
when used on Baklava, didn't effectively support this because TestLooperManager.acquire
needs to be called from the Looper's thread.

This commit refactors Interrogator so it be acquired and stored on the Looper's current thread.

PiperOrigin-RevId: 714149090

Author: brettchabot

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();
+    }
+  }
 }