From 70e3c22880a34b369e9404a207c76cb1ae6749c5 Mon Sep 17 00:00:00 2001
From: AndroidX Test Team <copybara-androidxtest@google.com>
Date: Mon, 6 Apr 2026 16:06:30 -0700
Subject: [PATCH] Ensure @Before and @Test run on the same thread in
 AndroidJUnit4ClassRunner.

Currently, AndroidJUnit4ClassRunner inherits standard JUnit 4 behavior where
timeouts are applied only to the test method itself. If a test timeout is
set (e.g., passed by TradeFed or runner arguments), FailOnTimeout executes
the test method in a separate thread.

Since setup (@Before) and teardown (@After) are applied outside this
timeout, they run on the caller thread, causing a thread mismatch. This
prevents ThreadLocal variables initialized in setup from being accessible in
the test.

This CL modifies AndroidJUnit4ClassRunner to apply the timeout to the
combined statement of setup, test, and teardown. This is achieved by:
  1. Overriding withAfters to wrap the combined statement in FailOnTimeout.
  2. Overriding withPotentialTimeout to be a no-op, preventing double
     timeout application.

This ensures all lifecycle methods run on the same thread spawned by the
timeout.

NOTE: Because the timeout now wraps setup and teardown, any time spent in
@Before or @After will now count towards the test timeout.

PiperOrigin-RevId: 895547043
---
 runner/android_junit_runner/CHANGELOG.md      |  2 ++
 .../junit4/AndroidJUnit4ClassRunner.java      | 34 +++++++++----------
 2 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/runner/android_junit_runner/CHANGELOG.md b/runner/android_junit_runner/CHANGELOG.md
index 43ae3fda7..d8d6e2ad4 100644
--- a/runner/android_junit_runner/CHANGELOG.md
+++ b/runner/android_junit_runner/CHANGELOG.md
@@ -12,6 +12,8 @@
 
 * Add logs at the start and end of RunBefore and RunAfters sections to help bug understanding. (b/445754263)
 
+* Ensure @Before and @Test run on the same thread in AndroidJUnit4ClassRunner.
+
 **Breaking Changes**
 
 **API Changes**
diff --git a/runner/android_junit_runner/java/androidx/test/internal/runner/junit4/AndroidJUnit4ClassRunner.java b/runner/android_junit_runner/java/androidx/test/internal/runner/junit4/AndroidJUnit4ClassRunner.java
index a4dfde2c7..ea7b0fba3 100644
--- a/runner/android_junit_runner/java/androidx/test/internal/runner/junit4/AndroidJUnit4ClassRunner.java
+++ b/runner/android_junit_runner/java/androidx/test/internal/runner/junit4/AndroidJUnit4ClassRunner.java
@@ -73,31 +73,31 @@ protected Statement withBefores(FrameworkMethod method, Object target, Statement
   @Override
   protected Statement withAfters(FrameworkMethod method, Object target, Statement statement) {
     List<FrameworkMethod> afters = getTestClass().getAnnotatedMethods(After.class);
-    return afters.isEmpty() ? statement : new RunAfters(method, statement, afters, target);
-  }
+    Statement combined =
+        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));
-
-    // use runner arg timeout if test level timeout is not present
     if (timeout <= 0 && perTestTimeout > 0) {
       timeout = perTestTimeout;
     }
 
-    if (timeout <= 0) {
-      // no timeout was set
-      return next;
+    if (timeout > 0) {
+      // Cannot switch to use builder as that is not supported in JUnit 4.10 which is what is
+      // available in AOSP.
+      @SuppressWarnings("deprecation")
+      var failOnTimeout = new FailOnTimeout(combined, timeout);
+      return failOnTimeout;
     }
+    return combined;
+  }
 
-    // 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);
+  /**
+   * 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) {