Rethrow original failure in OkHttpClientTestRule (#6037)

okhttp-testing-support/src/main/kotlin/okhttp3/OkHttpClientTestRule.kt

@@ -19,6 +19,7 @@ import java.net.InetAddress
 import java.util.concurrent.TimeUnit
 import java.util.logging.Handler
 import java.util.logging.Level
+import java.util.logging.LogManager
 import java.util.logging.LogRecord
 import java.util.logging.Logger
 import okhttp3.internal.concurrent.TaskRunner
@@ -136,8 +137,13 @@ class OkHttpClientTestRule : TestRule {
   }
 
   private fun ensureAllTaskQueuesIdle() {
+    val entryTime = System.nanoTime()
+
     for (queue in TaskRunner.INSTANCE.activeQueues()) {
-      if (!queue.idleLatch().await(1_000L, TimeUnit.MILLISECONDS)) {
+      // We wait at most 1 second, so we don't ever turn multiple lost threads into
+      // a test timeout failure.
+      val waitTime = (entryTime + 1_000_000_000L - System.nanoTime())
+      if (!queue.idleLatch().await(waitTime, TimeUnit.NANOSECONDS)) {
         TaskRunner.INSTANCE.cancelAll()
         fail("Queue still active after 1000 ms")
       }
@@ -154,10 +160,13 @@ class OkHttpClientTestRule : TestRule {
         Thread.setDefaultUncaughtExceptionHandler { _, throwable ->
           initUncaughtException(throwable)
         }
+        val taskQueuesWereIdle = TaskRunner.INSTANCE.activeQueues().isEmpty()
+        var failure: Throwable? = null
         try {
           applyLogger {
             addHandler(testLogHandler)
             level = Level.FINEST
+            useParentHandlers = false
           }
 
           base.evaluate()
@@ -166,19 +175,42 @@ class OkHttpClientTestRule : TestRule {
           }
           logEventsIfFlaky(description)
         } catch (t: Throwable) {
+          failure = t
           logEvents()
           throw t
         } finally {
-          applyLogger {
+          LogManager.getLogManager().reset()
-            removeHandler(testLogHandler)
-            level = Level.INFO
-          }
 
           Thread.setDefaultUncaughtExceptionHandler(defaultUncaughtExceptionHandler)
+          try {
             ensureAllConnectionsReleased()
             releaseClient()
+          } catch (ae: AssertionError) {
+            // Prefer keeping the inflight failure, but don't release this in-use client.
+            if (failure != null) {
+              failure.addSuppressed(ae)
+            } else {
+              failure = ae
+            }
+          }
+
+          try {
+            if (taskQueuesWereIdle) {
               ensureAllTaskQueuesIdle()
             }
+          } catch (ae: AssertionError) {
+            // Prefer keeping the inflight failure, but don't release this in-use client.
+            if (failure != null) {
+              failure.addSuppressed(ae)
+            } else {
+              failure = ae
+            }
+          }
+
+          if (failure != null) {
+            throw failure
+          }
+        }
       }
 
       private fun releaseClient() {

okhttp/src/test/java/okhttp3/internal/http2/HttpOverHttp2Test.java

@@ -72,6 +72,7 @@ import okio.Buffer;
 import okio.BufferedSink;
 import okio.GzipSink;
 import okio.Okio;
+import org.jetbrains.annotations.NotNull;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Ignore;
@@ -112,7 +113,7 @@ public final class HttpOverHttp2Test {
   }
 
   private final PlatformRule platform = new PlatformRule();
-  private final OkHttpClientTestRule clientTestRule = new OkHttpClientTestRule();
+  private final OkHttpClientTestRule clientTestRule = configureClientTestRule();
   @Rule public final TestRule chain = RuleChain.outerRule(platform)
       .around(new Timeout(60, SECONDS))
       .around(clientTestRule);
@@ -129,6 +130,13 @@ public final class HttpOverHttp2Test {
     this.protocol = protocol;
   }
 
+
+  @NotNull private OkHttpClientTestRule configureClientTestRule() {
+    OkHttpClientTestRule clientTestRule = new OkHttpClientTestRule();
+    clientTestRule.setRecordTaskRunner(true);
+    return clientTestRule;
+  }
+
   @Before public void setUp() {
     platform.assumeNotOpenJSSE();
     platform.assumeNotBouncyCastle();

okhttp/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.java

@@ -53,6 +53,7 @@ import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.RuleChain;
 
 import static java.util.Arrays.asList;
 import static okhttp3.TestUtil.repeat;
@@ -66,8 +67,10 @@ public final class WebSocketHttpTest {
   // Flaky https://github.com/square/okhttp/issues/4515
   // Flaky https://github.com/square/okhttp/issues/4953
 
-  @Rule public final MockWebServer webServer = new MockWebServer();
+  final MockWebServer webServer = new MockWebServer();
-  @Rule public final OkHttpClientTestRule clientTestRule = new OkHttpClientTestRule();
+  final OkHttpClientTestRule clientTestRule = configureClientTestRule();
+
+  @Rule public final RuleChain orderedRules = RuleChain.outerRule(clientTestRule).around(webServer);
   @Rule public final PlatformRule platform = new PlatformRule();
   @Rule public final TestLogHandler testLogHandler = new TestLogHandler(OkHttpClient.class);
 
@@ -86,6 +89,12 @@ public final class WebSocketHttpTest {
       })
       .build();
 
+  private OkHttpClientTestRule configureClientTestRule() {
+    OkHttpClientTestRule clientTestRule = new OkHttpClientTestRule();
+    clientTestRule.setRecordTaskRunner(true);
+    return clientTestRule;
+  }
+
   @Before public void setUp() {
     platform.assumeNotOpenJSSE();
     platform.assumeNotBouncyCastle();