Refactor timeout handling in AndroidJUnit4ClassRunner.

This change replaces the use of JUnit's FailOnTimeout with a custom implementation that ensures @before, @test, and @after methods run on the same thread. The previous FailOnTimeout would execute only the @test method in a separate thread, potentially causing issues with thread-local state. The new implementation wraps the entire test block (including Before/After) in a separate thread when a timeout is active, preserving thread locality within the test lifecycle.

PiperOrigin-RevId: 903547390

Author: copybara-androidxtest

runner/android_junit_runner/CHANGELOG.md

@@ -6,6 +6,8 @@
 
 **Bug Fixes**
 
+* Ensure @Before and @Test run on the same thread in AndroidJUnit4ClassRunner.
+
 **New Features**
 
 * Make perfetto trace sections for tests more identifiable by prefixing with "test:" and using fully qualified class name. (b/204992764)

runner/android_junit_runner/java/androidx/test/internal/runner/junit4/AndroidJUnit4ClassRunner.java

@@ -23,14 +23,17 @@
 import androidx.test.internal.runner.junit4.statement.UiThreadStatement;
 import androidx.test.internal.util.AndroidRunnerParams;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
-import org.junit.internal.runners.statements.FailOnTimeout;
 import org.junit.runners.BlockJUnit4ClassRunner;
 import org.junit.runners.model.FrameworkMethod;
 import org.junit.runners.model.InitializationError;
 import org.junit.runners.model.Statement;
+import org.junit.runners.model.TestTimedOutException;
 
 /** A specialized {@link BlockJUnit4ClassRunner} that can handle timeouts */
 public class AndroidJUnit4ClassRunner extends BlockJUnit4ClassRunner {
@@ -55,13 +58,35 @@ public AndroidJUnit4ClassRunner(Class<?> klass) throws InitializationError {
     this(klass, RunnerArgs.parseTestTimeout(getArguments()));
   }
 
+  private static final ThreadLocal<CountDownLatch> currentTestStartedLatch = new ThreadLocal<>();
+  private static final ThreadLocal<CountDownLatch> currentTestFinishedLatch = new ThreadLocal<>();
+
   /** Returns a {@link Statement} that invokes {@code method} on {@code test} */
   @Override
   protected Statement methodInvoker(FrameworkMethod method, Object test) {
+    final Statement invoker;
     if (UiThreadStatement.shouldRunOnUiThread(method)) {
-      return new UiThreadStatement(super.methodInvoker(method, test), true);
+      invoker = new UiThreadStatement(super.methodInvoker(method, test), true);
+    } else {
+      invoker = super.methodInvoker(method, test);
     }
-    return super.methodInvoker(method, test);
+    return new Statement() {
+      @Override
+      public void evaluate() throws Throwable {
+        CountDownLatch startLatch = currentTestStartedLatch.get();
+        if (startLatch != null) {
+          startLatch.countDown();
+        }
+        try {
+          invoker.evaluate();
+        } finally {
+          CountDownLatch finishLatch = currentTestFinishedLatch.get();
+          if (finishLatch != null) {
+            finishLatch.countDown();
+          }
+        }
+      }
+    };
   }
 
   @Override
@@ -76,28 +101,76 @@ protected Statement withAfters(FrameworkMethod method, Object target, Statement
     return afters.isEmpty() ? statement : new RunAfters(method, statement, afters, target);
   }
 
-  /**
-   * Default to {@link org.junit.Test#timeout()} level timeout if set. Otherwise, set the timeout
-   * that was passed to the instrumentation via argument.
-   */
   @Override
-  protected Statement withPotentialTimeout(FrameworkMethod method, Object test, Statement next) {
-    // test level timeout i.e @Test(timeout = 123)
-    long timeout = getTimeout(method.getAnnotation(Test.class));
+  protected Statement methodBlock(FrameworkMethod method) {
+    final Statement statement = super.methodBlock(method);
 
-    // use runner arg timeout if test level timeout is not present
+    long timeout = getTimeout(method.getAnnotation(Test.class));
     if (timeout <= 0 && perTestTimeout > 0) {
       timeout = perTestTimeout;
     }
+    final long finalTimeout = timeout;
 
-    if (timeout <= 0) {
-      // no timeout was set
-      return next;
+    if (finalTimeout <= 0) {
+      return statement;
     }
 
-    // Cannot switch to use builder as that is not supported in JUnit 4.10 which is what is
-    // available in AOSP.
-    return new FailOnTimeout(next, timeout);
+    return new Statement() {
+      @Override
+      @SuppressWarnings("Interruption") // We want to interrupt the thread to stop the test.
+      public void evaluate() throws Throwable {
+        final AtomicReference<Throwable> failure = new AtomicReference<>();
+        final CountDownLatch testStartedLatch = new CountDownLatch(1);
+        final CountDownLatch testFinishedLatch = new CountDownLatch(1);
+        final CountDownLatch doneLatch = new CountDownLatch(1);
+
+        Thread thread =
+            new Thread(
+                new Runnable() {
+                  @Override
+                  public void run() {
+                    currentTestStartedLatch.set(testStartedLatch);
+                    currentTestFinishedLatch.set(testFinishedLatch);
+                    try {
+                      statement.evaluate();
+                    } catch (Throwable t) {
+                      failure.set(t);
+                    } finally {
+                      testStartedLatch.countDown();
+                      testFinishedLatch.countDown();
+                      doneLatch.countDown();
+                      currentTestStartedLatch.remove();
+                      currentTestFinishedLatch.remove();
+                    }
+                  }
+                },
+                "Time-limited test");
+        thread.setDaemon(true);
+        thread.start();
+
+        testStartedLatch.await();
+        boolean finishedInTime = testFinishedLatch.await(finalTimeout, TimeUnit.MILLISECONDS);
+
+        if (!finishedInTime) {
+          thread.interrupt();
+          throw new TestTimedOutException(finalTimeout, TimeUnit.MILLISECONDS);
+        }
+
+        doneLatch.await();
+        if (failure.get() != null) {
+          throw failure.get();
+        }
+      }
+    };
+  }
+
+  /**
+   * Default to {@link org.junit.Test#timeout()} level timeout if set. Otherwise, set the timeout
+   * that was passed to the instrumentation via argument.
+   */
+  @Override
+  protected Statement withPotentialTimeout(FrameworkMethod method, Object test, Statement next) {
+    return next;
   }
 
   private long getTimeout(Test annotation) {

runner/android_junit_runner/javatests/androidx/test/internal/runner/junit4/AndroidAnnotatedBuilderTest.java

@@ -26,9 +26,12 @@
 import java.lang.reflect.InvocationTargetException;
 import java.util.Arrays;
 import java.util.Collection;
+import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.JUnitCore;
+import org.junit.runner.Result;
 import org.junit.runner.RunWith;
 import org.junit.runner.Runner;
 import org.junit.runners.JUnit4;
@@ -133,4 +136,41 @@ public Runner buildAndroidRunner(Class<? extends Runner> runnerClass, Class<?> t
     // attempt to create a runner for a class with no @RunWith annotation
     ab.runnerForClass(NoRunWithClass.class);
   }
+
+  @SuppressWarnings("NonFinalStaticField") // Static fields are needed to check thread assignment.
+  public static class TimeoutTestClass {
+    static Thread beforeThread;
+    static Thread testThread;
+    static Thread afterThread;
+
+    @Before
+    public void before() {
+      beforeThread = Thread.currentThread();
+    }
+
+    @Test(timeout = 5000)
+    public void testWithTimeout() {
+      testThread = Thread.currentThread();
+    }
+
+    @After
+    public void after() {
+      afterThread = Thread.currentThread();
+    }
+  }
+
+  @Test
+  public void testThreadsSameWithTimeout() throws org.junit.runners.model.InitializationError {
+    TimeoutTestClass.beforeThread = null;
+    TimeoutTestClass.testThread = null;
+    TimeoutTestClass.afterThread = null;
+
+    AndroidJUnit4ClassRunner runner = new AndroidJUnit4ClassRunner(TimeoutTestClass.class, 0);
+    Result result = new JUnitCore().run(runner);
+
+    assertEquals(0, result.getFailureCount());
+    Assert.assertNotNull(TimeoutTestClass.beforeThread);
+    assertEquals(TimeoutTestClass.beforeThread, TimeoutTestClass.testThread);
+    assertEquals(TimeoutTestClass.testThread, TimeoutTestClass.afterThread);
+  }
 }