Merge pull request square#1668 from square/jw/ws-response

Add optional Response to WS failure callback.
Rainmonth · May 21, 2015 · ce09d95 · ce09d95
2 parents aac8462 + 8e25853
commit ce09d95
Show file tree

Hide file tree

Showing 8 changed files with 37 additions and 30 deletions.
diff --git a/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockWebServer.java b/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockWebServer.java
@@ -674,7 +674,7 @@ private void handleWebSocketUpgrade(Socket socket, BufferedSource source, Buffer
     new Thread(new Runnable() {
       @Override public void run() {
         try {
-          listener.onOpen(webSocket, fancyRequest, fancyResponse);
+          listener.onOpen(webSocket, fancyResponse);
         } catch (IOException e) {
           // TODO try to write close frame?
           connectionClose.countDown();

diff --git a/okhttp-ws-tests/src/test/java/com/squareup/okhttp/ws/AutobahnTester.java b/okhttp-ws-tests/src/test/java/com/squareup/okhttp/ws/AutobahnTester.java
@@ -69,7 +69,7 @@ private void runTest(final long number, final long count) throws IOException {
           private final ExecutorService sendExecutor = Executors.newSingleThreadExecutor();
           private WebSocket webSocket;
 
-          @Override public void onOpen(WebSocket webSocket, Request request, Response response)
+          @Override public void onOpen(WebSocket webSocket, Response response)
               throws IOException {
             System.out.println("Executing test case " + number + "/" + count);
             this.webSocket = webSocket;
@@ -100,7 +100,7 @@ private void runTest(final long number, final long count) throws IOException {
             latch.countDown();
           }
 
-          @Override public void onFailure(IOException e) {
+          @Override public void onFailure(IOException e, Response response) {
             latch.countDown();
           }
         });
@@ -118,7 +118,7 @@ private long getTestCount() throws IOException {
     final AtomicLong countRef = new AtomicLong();
     final AtomicReference<IOException> failureRef = new AtomicReference<>();
     newWebSocket("/getCaseCount").enqueue(new WebSocketListener() {
-      @Override public void onOpen(WebSocket webSocket, Request request, Response response)
+      @Override public void onOpen(WebSocket webSocket, Response response)
           throws IOException {
       }
 
@@ -135,7 +135,7 @@ private long getTestCount() throws IOException {
         latch.countDown();
       }
 
-      @Override public void onFailure(IOException e) {
+      @Override public void onFailure(IOException e, Response response) {
         failureRef.set(e);
         latch.countDown();
       }
@@ -157,7 +157,7 @@ private long getTestCount() throws IOException {
   private void updateReports() {
     final CountDownLatch latch = new CountDownLatch(1);
     newWebSocket("/updateReports?agent=" + Version.userAgent()).enqueue(new WebSocketListener() {
-      @Override public void onOpen(WebSocket webSocket, Request request, Response response)
+      @Override public void onOpen(WebSocket webSocket, Response response)
           throws IOException {
       }
 
@@ -172,7 +172,7 @@ private void updateReports() {
         latch.countDown();
       }
 
-      @Override public void onFailure(IOException e) {
+      @Override public void onFailure(IOException e, Response response) {
         latch.countDown();
       }
     });

diff --git a/okhttp-ws-tests/src/test/java/com/squareup/okhttp/ws/WebSocketCallTest.java b/okhttp-ws-tests/src/test/java/com/squareup/okhttp/ws/WebSocketCallTest.java
@@ -70,7 +70,7 @@ public final class WebSocketCallTest {
 
   @Test public void serverMessage() throws IOException {
     WebSocketListener serverListener = new EmptyWebSocketListener() {
-      @Override public void onOpen(WebSocket webSocket, Request request, Response response)
+      @Override public void onOpen(WebSocket webSocket, Response response)
           throws IOException {
         webSocket.sendMessage(TEXT, new Buffer().writeUtf8("Hello, WebSockets!"));
       }
@@ -96,7 +96,7 @@ public final class WebSocketCallTest {
 
   @Test public void serverStreamingMessage() throws IOException {
     WebSocketListener serverListener = new EmptyWebSocketListener() {
-      @Override public void onOpen(WebSocket webSocket, Request request, Response response)
+      @Override public void onOpen(WebSocket webSocket, Response response)
           throws IOException {
         BufferedSink sink = webSocket.newMessageSink(TEXT);
         sink.writeUtf8("Hello, ").flush();
@@ -134,7 +134,8 @@ public final class WebSocketCallTest {
   }
 
   @Test public void wrongConnectionHeader() {
-    server.enqueue(new MockResponse().setResponseCode(101)
+    server.enqueue(new MockResponse()
+        .setResponseCode(101)
         .setHeader("Upgrade", "websocket")
         .setHeader("Connection", "Downgrade")
         .setHeader("Sec-WebSocket-Accept", "ujmZX4KXZqjwy6vi1aQFH5p4Ygk="));
@@ -234,7 +235,7 @@ private WebSocket awaitWebSocket(Request request) {
     final AtomicReference<IOException> failureRef = new AtomicReference<>();
     final CountDownLatch latch = new CountDownLatch(1);
     call.enqueue(new WebSocketListener() {
-      @Override public void onOpen(WebSocket webSocket, Request request, Response response)
+      @Override public void onOpen(WebSocket webSocket, Response response)
           throws IOException {
         webSocketRef.set(webSocket);
         responseRef.set(response);
@@ -254,8 +255,8 @@ private WebSocket awaitWebSocket(Request request) {
         listener.onClose(code, reason);
       }
 
-      @Override public void onFailure(IOException e) {
-        listener.onFailure(e);
+      @Override public void onFailure(IOException e, Response response) {
+        listener.onFailure(e, null);
         failureRef.set(e);
         latch.countDown();
       }
@@ -273,7 +274,7 @@ private WebSocket awaitWebSocket(Request request) {
   }
 
   private static class EmptyWebSocketListener implements WebSocketListener {
-    @Override public void onOpen(WebSocket webSocket, Request request, Response response)
+    @Override public void onOpen(WebSocket webSocket, Response response)
         throws IOException {
     }
 
@@ -287,7 +288,7 @@ private static class EmptyWebSocketListener implements WebSocketListener {
     @Override public void onClose(int code, String reason) {
     }
 
-    @Override public void onFailure(IOException e) {
+    @Override public void onFailure(IOException e, Response response) {
     }
   }
 }
diff --git a/okhttp-ws-tests/src/test/java/com/squareup/okhttp/ws/WebSocketRecorder.java b/okhttp-ws-tests/src/test/java/com/squareup/okhttp/ws/WebSocketRecorder.java
@@ -15,7 +15,6 @@
  */
 package com.squareup.okhttp.ws;
 
-import com.squareup.okhttp.Request;
 import com.squareup.okhttp.Response;
 import com.squareup.okhttp.internal.ws.WebSocketReader;
 import java.io.IOException;
@@ -44,7 +43,7 @@ public void setNextMessageDelegate(MessageDelegate delegate) {
     this.delegate = delegate;
   }
 
-  @Override public void onOpen(WebSocket webSocket, Request request, Response response) {
+  @Override public void onOpen(WebSocket webSocket, Response response) {
   }
 
   @Override public void onMessage(BufferedSource source, WebSocket.PayloadType type)
@@ -72,7 +71,7 @@ public void setNextMessageDelegate(MessageDelegate delegate) {
     events.add(new Close(code, reason));
   }
 
-  @Override public void onFailure(IOException e) {
+  @Override public void onFailure(IOException e, Response response) {
     events.add(e);
   }
 
@@ -109,7 +108,7 @@ public void assertPong(Buffer payload) {
   }
 
   public void assertClose(int code, String reason) {
-      assertEquals(new Close(code, reason), nextEvent());
+    assertEquals(new Close(code, reason), nextEvent());
   }
 
   public void assertFailure(Class<? extends IOException> cls, String message) {

diff --git a/okhttp-ws/src/main/java/com/squareup/okhttp/internal/ws/RealWebSocket.java b/okhttp-ws/src/main/java/com/squareup/okhttp/internal/ws/RealWebSocket.java
@@ -181,7 +181,7 @@ private void readerErrorClose(IOException e) {
     } catch (IOException ignored) {
     }
 
-    listener.onFailure(e);
+    listener.onFailure(e, null);
   }
 
   /** Perform any tear-down work on the connection (close the socket, recycle, etc.). */

diff --git a/okhttp-ws/src/main/java/com/squareup/okhttp/ws/WebSocketCall.java b/okhttp-ws/src/main/java/com/squareup/okhttp/ws/WebSocketCall.java
@@ -102,12 +102,12 @@ public void enqueue(final WebSocketListener listener) {
         try {
           createWebSocket(response, listener);
         } catch (IOException e) {
-          listener.onFailure(e);
+          listener.onFailure(e, response);
         }
       }
 
       @Override public void onFailure(Request request, IOException e) {
-        listener.onFailure(e);
+        listener.onFailure(e, null);
       }
     };
     // TODO call.enqueue(responseCallback, true);
@@ -167,7 +167,7 @@ private void createWebSocket(Response response, WebSocketListener listener)
     // TODO connection.setOwner(webSocket);
     Internal.instance.connectionSetOwner(connection, webSocket);
 
-    listener.onOpen(webSocket, request, response);
+    listener.onOpen(webSocket, response);
 
     // Start a dedicated thread for reading the web socket.
     new Thread(new NamedRunnable("OkHttp WebSocket reader %s", request.urlString()) {

diff --git a/okhttp-ws/src/main/java/com/squareup/okhttp/ws/WebSocketListener.java b/okhttp-ws/src/main/java/com/squareup/okhttp/ws/WebSocketListener.java
@@ -15,7 +15,6 @@
  */
 package com.squareup.okhttp.ws;
 
-import com.squareup.okhttp.Request;
 import com.squareup.okhttp.Response;
 import java.io.IOException;
 import okio.Buffer;
@@ -25,7 +24,18 @@
 
 /** Listener for server-initiated messages on a connected {@link WebSocket}. */
 public interface WebSocketListener {
-  void onOpen(WebSocket webSocket, Request request, Response response) throws IOException;
+  /**
+   * Called when the request has successfully been upgraded to a web socket.
+   */
+  void onOpen(WebSocket webSocket, Response response) throws IOException;
+
+  /**
+   * Called when the transport or protocol layer of this web socket errors during communication.
+   *
+   * @param response Present when the failure is a direct result of the response (e.g., failed
+   * upgrade, non-101 response code, etc.). {@code null} otherwise.
+   */
+  void onFailure(IOException e, Response response);
 
   /**
    * Called when a server message is received. The {@code type} indicates whether the
@@ -53,7 +63,4 @@ public interface WebSocketListener {
    * @param reason Reason for close or an empty string.
    */
   void onClose(int code, String reason);
-
-  /** Called when the transport or protocol layer of this web socket errors during communication. */
-  void onFailure(IOException e);
 }
diff --git a/samples/guide/src/main/java/com/squareup/okhttp/recipes/WebSocketEcho.java b/samples/guide/src/main/java/com/squareup/okhttp/recipes/WebSocketEcho.java
@@ -27,7 +27,7 @@ private void run() throws IOException {
     client.getDispatcher().getExecutorService().shutdown();
   }
 
-  @Override public void onOpen(WebSocket webSocket, Request request, Response response)
+  @Override public void onOpen(WebSocket webSocket, Response response)
       throws IOException {
     webSocket.sendMessage(TEXT, new Buffer().writeUtf8("Hello..."));
     webSocket.sendMessage(TEXT, new Buffer().writeUtf8("...World!"));
@@ -57,7 +57,7 @@ private void run() throws IOException {
     System.out.println("CLOSE: " + code + " " + reason);
   }
 
-  @Override public void onFailure(IOException e) {
+  @Override public void onFailure(IOException e, Response response) {
     e.printStackTrace();
   }