WebSocketUpgradeHandler.onCompleted should throw Exception when status != 101, close AsyncHttpClient#836

Author: slandelle

src/main/java/com/ning/http/client/ws/WebSocketUpgradeHandler.java

@@ -77,11 +77,11 @@ public WebSocket onCompleted() throws Exception {
             for (WebSocketListener listener : listeners) {
                 listener.onError(e);
             }
-            return null;
+            throw e;
         }
 
         if (webSocket == null) {
-            throw new IllegalStateException("WebSocket is null");
+            throw new NullPointerException("webSocket");
         }
         return webSocket;
     }

src/test/java/com/ning/http/client/ws/CloseCodeReasonMessageTest.java

@@ -12,22 +12,18 @@
  */
 package com.ning.http.client.ws;
 
-import com.ning.http.client.AsyncHttpClient;
-import com.ning.http.client.ws.WebSocket;
-import com.ning.http.client.ws.WebSocketCloseCodeReasonListener;
-import com.ning.http.client.ws.WebSocketListener;
-import com.ning.http.client.ws.WebSocketTextListener;
-import com.ning.http.client.ws.WebSocketUpgradeHandler;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
 
 import org.testng.annotations.Test;
 
+import com.ning.http.client.AsyncHttpClient;
+
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicReference;
 
-import static org.testng.Assert.assertEquals;
-import static org.testng.Assert.assertNotNull;
-import static org.testng.Assert.assertTrue;
-
 public abstract class CloseCodeReasonMessageTest extends TextMessageTest {
 
     @Test(timeOut = 60000)
@@ -94,13 +90,41 @@ public void onError(Throwable t) {
         }
     }
 
+    @Test(timeOut = 60000, expectedExceptions = { ExecutionException.class })
+    public void getWebSocketThrowsException() throws Throwable {
+        final CountDownLatch latch = new CountDownLatch(1);
+        try (AsyncHttpClient client = getAsyncHttpClient(null)) {
+            client.prepareGet("http://apache.org").execute(new WebSocketUpgradeHandler.Builder().addWebSocketListener(new WebSocketTextListener() {
+
+                @Override
+                public void onMessage(String message) {
+                }
+
+                @Override
+                public void onOpen(com.ning.http.client.ws.WebSocket websocket) {
+                }
+
+                @Override
+                public void onClose(com.ning.http.client.ws.WebSocket websocket) {
+                }
+
+                @Override
+                public void onError(Throwable t) {
+                    latch.countDown();
+                }
+            }).build()).get();
+        }
+        
+        latch.await();
+    }
+    
     @Test(timeOut = 60000)
     public void wrongStatusCode() throws Throwable {
         try (AsyncHttpClient client = getAsyncHttpClient(null)) {
             final CountDownLatch latch = new CountDownLatch(1);
             final AtomicReference<Throwable> throwable = new AtomicReference<>();
 
-            WebSocket websocket = client.prepareGet("http://apache.org/").execute(new WebSocketUpgradeHandler.Builder().addWebSocketListener(new WebSocketTextListener() {
+            client.prepareGet("http://apache.org/").execute(new WebSocketUpgradeHandler.Builder().addWebSocketListener(new WebSocketTextListener() {
 
                 @Override
                 public void onMessage(String message) {
@@ -119,7 +143,7 @@ public void onError(Throwable t) {
                     throwable.set(t);
                     latch.countDown();
                 }
-            }).build()).get();
+            }).build());
 
             latch.await();
             assertNotNull(throwable.get());
@@ -133,7 +157,7 @@ public void wrongProtocolCode() throws Throwable {
             final CountDownLatch latch = new CountDownLatch(1);
             final AtomicReference<Throwable> throwable = new AtomicReference<>();
 
-            WebSocket websocket = client.prepareGet("ws://www.google.com/").execute(new WebSocketUpgradeHandler.Builder().addWebSocketListener(new WebSocketTextListener() {
+            client.prepareGet("ws://www.google.com/").execute(new WebSocketUpgradeHandler.Builder().addWebSocketListener(new WebSocketTextListener() {
 
                 @Override
                 public void onMessage(String message) {
@@ -152,7 +176,7 @@ public void onError(Throwable t) {
                     throwable.set(t);
                     latch.countDown();
                 }
-            }).build()).get();
+            }).build());
 
             latch.await();
             assertNotNull(throwable.get());