From d81ec0f23bf4cc97060a75a0b12e299bd8dfdefc Mon Sep 17 00:00:00 2001
From: Yuri Schimke <yuri@schimke.ee>
Date: Wed, 26 Jun 2019 06:34:51 +0100
Subject: [PATCH] OkHttpClientTestRule check connectionCount instead of idle
 (#5226)

* OkHttpClientTestRule check connectionCount instead of idle

Clients should be clean after use, not just from idle connections.

* Abandon unclean clients

* Simplify logic
---
 .../main/java/okhttp3/OkHttpClientTestRule.kt | 25 +++++++++++++++----
 .../internal/ws/WebSocketHttpTest.java        |  1 +
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/okhttp-testing-support/src/main/java/okhttp3/OkHttpClientTestRule.kt b/okhttp-testing-support/src/main/java/okhttp3/OkHttpClientTestRule.kt
index 90e067b8a..d94f9f5c1 100644
--- a/okhttp-testing-support/src/main/java/okhttp3/OkHttpClientTestRule.kt
+++ b/okhttp-testing-support/src/main/java/okhttp3/OkHttpClientTestRule.kt
@@ -44,9 +44,11 @@ class OkHttpClientTestRule : TestRule {
   }
 
   fun ensureAllConnectionsReleased() {
-    val connectionPool = prototype!!.connectionPool
-    connectionPool.evictAll()
-    assertThat(connectionPool.idleConnectionCount()).isEqualTo(0)
+    prototype?.let {
+      val connectionPool = it.connectionPool
+      connectionPool.evictAll()
+      assertThat(connectionPool.connectionCount()).isEqualTo(0)
+    }
   }
 
   override fun apply(base: Statement, description: Description): Statement {
@@ -68,12 +70,25 @@ class OkHttpClientTestRule : TestRule {
       }
 
       private fun releaseClient() {
-        prototypes.push(prototype)
-        prototype = null
+        prototype?.let {
+          prototypes.push(it)
+          prototype = null
+        }
       }
     }
   }
 
+  /**
+   * Called if a test is known to be leaky.
+   */
+  fun abandonClient() {
+    prototype?.let {
+      prototype = null
+      it.dispatcher.executorService.shutdownNow()
+      it.connectionPool.evictAll()
+    }
+  }
+
   companion object {
     /**
      * Quick and dirty pool of OkHttpClient instances. Each has its own independent dispatcher and
diff --git a/okhttp/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.java b/okhttp/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.java
index ad966506b..4806f923d 100644
--- a/okhttp/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.java
+++ b/okhttp/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.java
@@ -84,6 +84,7 @@ public final class WebSocketHttpTest {
     clientListener.assertExhausted();
 
     // TODO: assert all connections are released once leaks are fixed
+    clientTestRule.abandonClient();
   }
 
   @Test public void textMessage() {