Merge pull request #3033 from square/jwilson.1211.flaky_ws_test

Fix flakiness in RealWebSocketTest.clientAndServerCloseClosesConnection

okhttp-tests/src/test/java/okhttp3/internal/ws/RealWebSocketTest.java

@@ -162,7 +162,7 @@ public final class RealWebSocketTest {
     server.listener.assertClosed(1000, "Hello!");
   }
 
-  @Test public void clientAndServerCloseClosesConnection() throws IOException {
+  @Test public void clientAndServerCloseClosesConnection() throws Exception {
     // Send close from both sides at the same time.
     server.webSocket.close(1000, "Hello!");
     client.processNextFrame(); // Read close, close connection close.
@@ -175,6 +175,7 @@ public final class RealWebSocketTest {
     server.listener.assertClosing(1000, "Hi!");
     client.listener.assertClosed(1000, "Hello!");
     server.listener.assertClosed(1000, "Hi!");
+    client.webSocket.awaitTermination(5, TimeUnit.SECONDS);
     assertTrue(client.closed);
 
     server.listener.assertExhausted(); // Client should not have sent second close.

okhttp/src/main/java/okhttp3/internal/ws/RealWebSocket.java

@@ -25,6 +25,7 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
 import okhttp3.Call;
 import okhttp3.Callback;
 import okhttp3.OkHttpClient;
@@ -276,6 +277,13 @@ boolean processNextFrame() throws IOException {
     }
   }
 
+  /**
+   * For testing: wait until the web socket's executor has terminated.
+   */
+  void awaitTermination(int timeout, TimeUnit timeUnit) throws InterruptedException {
+    executor.awaitTermination(timeout, timeUnit);
+  }
+
   synchronized int pingCount() {
     return pingCount;
   }