From fea8fbba5fd9eadf3f88b91c1479290a60e6d462 Mon Sep 17 00:00:00 2001
From: Yuri Schimke <yuri@schimke.ee>
Date: Wed, 13 May 2020 07:46:19 +0100
Subject: [PATCH] Attempt to minimise WebSocket test flakiness (#6045)

---
 .../main/kotlin/okhttp3/ClientRuleEventListener.kt | 11 +++++++++--
 .../main/kotlin/okhttp3/OkHttpClientTestRule.kt    | 14 +++++++++++++-
 .../src/test/java/okhttp3/EventListenerTest.java   |  1 -
 .../okhttp3/internal/ws/WebSocketHttpTest.java     |  2 +-
 4 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/okhttp-testing-support/src/main/kotlin/okhttp3/ClientRuleEventListener.kt b/okhttp-testing-support/src/main/kotlin/okhttp3/ClientRuleEventListener.kt
index cf5a1038f..104c45e45 100644
--- a/okhttp-testing-support/src/main/kotlin/okhttp3/ClientRuleEventListener.kt
+++ b/okhttp-testing-support/src/main/kotlin/okhttp3/ClientRuleEventListener.kt
@@ -26,7 +26,7 @@ class ClientRuleEventListener(
   var logger: (String) -> Unit
 ) : EventListener(),
     EventListener.Factory {
-  private var startNs: Long = 0
+  private var startNs: Long? = null
 
   override fun create(call: Call): EventListener = this
 
@@ -266,7 +266,14 @@ class ClientRuleEventListener(
   }
 
   private fun logWithTime(message: String) {
-    val timeMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNs)
+    val startNs = startNs
+    val timeMs = if (startNs == null) {
+      // Event occurred before start, for an example an early cancel.
+      0L
+    } else {
+      TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNs)
+    }
+
     logger.invoke("[$timeMs ms] $message")
   }
 }
diff --git a/okhttp-testing-support/src/main/kotlin/okhttp3/OkHttpClientTestRule.kt b/okhttp-testing-support/src/main/kotlin/okhttp3/OkHttpClientTestRule.kt
index 6d3028245..26ff08918 100644
--- a/okhttp-testing-support/src/main/kotlin/okhttp3/OkHttpClientTestRule.kt
+++ b/okhttp-testing-support/src/main/kotlin/okhttp3/OkHttpClientTestRule.kt
@@ -43,6 +43,7 @@ class OkHttpClientTestRule : TestRule {
   private var testClient: OkHttpClient? = null
   private var uncaughtException: Throwable? = null
   var logger: Logger? = null
+  lateinit var testName: String
 
   var recordEvents = true
   var recordTaskRunner = false
@@ -131,7 +132,16 @@ class OkHttpClientTestRule : TestRule {
   fun ensureAllConnectionsReleased() {
     testClient?.let {
       val connectionPool = it.connectionPool
+
       connectionPool.evictAll()
+      if (connectionPool.connectionCount() > 0) {
+        // Minimise test flakiness due to possible race conditions with connections closing.
+        // Some number of tests will report here, but not fail due to this delay.
+        println("Delaying to avoid flakes")
+        Thread.sleep(500L)
+        println("After delay: " + connectionPool.connectionCount())
+      }
+
       assertEquals(0, connectionPool.connectionCount())
     }
   }
@@ -156,6 +166,8 @@ class OkHttpClientTestRule : TestRule {
   ): Statement {
     return object : Statement() {
       override fun evaluate() {
+        testName = description.methodName
+
         val defaultUncaughtExceptionHandler = Thread.getDefaultUncaughtExceptionHandler()
         Thread.setDefaultUncaughtExceptionHandler { _, throwable ->
           initUncaughtException(throwable)
@@ -233,7 +245,7 @@ class OkHttpClientTestRule : TestRule {
   @Synchronized private fun logEvents() {
     // Will be ineffective if test overrides the listener
     synchronized(clientEventsList) {
-      println("Events (${clientEventsList.size})")
+      println("$testName Events (${clientEventsList.size})")
 
       for (e in clientEventsList) {
         println(e)
diff --git a/okhttp/src/test/java/okhttp3/EventListenerTest.java b/okhttp/src/test/java/okhttp3/EventListenerTest.java
index 4e893f184..afa1e9bdb 100644
--- a/okhttp/src/test/java/okhttp3/EventListenerTest.java
+++ b/okhttp/src/test/java/okhttp3/EventListenerTest.java
@@ -78,7 +78,6 @@ import static okhttp3.tls.internal.TlsUtil.localhost;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.hamcrest.CoreMatchers.any;
 import static org.hamcrest.CoreMatchers.equalTo;
-import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeThat;
 
diff --git a/okhttp/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.java b/okhttp/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.java
index 16bd86f62..166a1263f 100644
--- a/okhttp/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.java
+++ b/okhttp/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.java
@@ -100,7 +100,7 @@ public final class WebSocketHttpTest {
     platform.assumeNotBouncyCastle();
   }
 
-  @After public void tearDown() {
+  @After public void tearDown() throws InterruptedException {
     clientListener.assertExhausted();
   }