added test for client side web socket re-connection logic

mockserver-client-java/src/test/java/org/mockserver/client/MockServerClientIntegrationTest.java

@@ -1,6 +1,9 @@
 package org.mockserver.client;
 
 import com.google.common.util.concurrent.SettableFuture;
+import io.netty.channel.Channel;
+import io.netty.channel.ChannelOutboundInvoker;
+import org.hamcrest.core.IsNot;
 import org.junit.*;
 import org.junit.rules.ExpectedException;
 import org.mockserver.echo.http.EchoServer;
@@ -17,6 +20,7 @@
 import org.mockserver.verify.VerificationSequence;
 import org.mockserver.verify.VerificationTimes;
 
+import java.net.SocketAddress;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -27,6 +31,7 @@
 import static org.hamcrest.core.Is.is;
 import static org.hamcrest.core.IsCollectionContaining.hasItems;
 import static org.junit.Assert.fail;
+import static org.mockito.Matchers.same;
 import static org.mockserver.character.Character.NEW_LINE;
 import static org.mockserver.matchers.Times.exactly;
 import static org.mockserver.matchers.Times.unlimited;
@@ -294,9 +299,7 @@ public void shouldSetupExpectationWithResponseObjectCallback() {
                     .withPath("/some_path")
                     .withBody(new StringBody("some_request_body"))
             )
-            .respond(httpRequest -> {
-                return response();
-            });
+            .respond(httpRequest -> response());
 
         // then
         assertThat(retrieveRequests(request()).size(), is(1));
@@ -334,6 +337,30 @@ public void shouldSetupExpectationWithResponseObjectCallback() {
         }
     }
 
+    @Test
+    public void shouldReconnectWebSocketIfClosed() {
+        // given
+        echoServerOne.withNextResponse(response().withStatusCode(201));
+        echoServerOne.getRegisteredClients().clear();
+        mockServerClientOne
+            .when(
+                request()
+                    .withPath("/some_path")
+                    .withBody(new StringBody("some_request_body"))
+            )
+            .respond(httpRequest -> response());
+
+        // when
+        assertThat(echoServerOne.getWebsocketChannels().size(), is(1));
+        Channel initialChannel = echoServerOne.getWebsocketChannels().get(0);
+        echoServerOne.getWebsocketChannels().forEach(ChannelOutboundInvoker::close);
+
+        // then
+        assertThat(retrieveRequests(request()).size(), is(1));
+        assertThat(echoServerOne.getWebsocketChannels().size(), is(1));
+        assertThat(echoServerOne.getWebsocketChannels().get(0), IsNot.not(same(initialChannel)));
+    }
+
     @Test
     public void shouldCloseWebsocketAfterStop() {
         // given

mockserver-core/src/main/java/org/mockserver/echo/http/EchoServer.java

@@ -2,6 +2,7 @@
 
 import com.google.common.util.concurrent.SettableFuture;
 import io.netty.bootstrap.ServerBootstrap;
+import io.netty.channel.Channel;
 import io.netty.channel.ChannelFutureListener;
 import io.netty.channel.ChannelOption;
 import io.netty.channel.EventLoopGroup;
@@ -19,10 +20,7 @@
 import org.slf4j.event.Level;
 
 import java.net.InetSocketAddress;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.LinkedList;
-import java.util.Queue;
+import java.util.*;
 import java.util.concurrent.TimeUnit;
 
 
@@ -39,8 +37,9 @@ public class EchoServer implements Stoppable {
     private final OnlyResponse onlyResponse = new OnlyResponse();
     private final NioEventLoopGroup eventLoopGroup = new NioEventLoopGroup();
     private final SettableFuture<Integer> boundPort = SettableFuture.create();
-    private final ArrayList<String> registeredClients;
-    private final ArrayList<TextWebSocketFrame> textWebSocketFrames;
+    private final List<String> registeredClients;
+    private final List<Channel> websocketChannels;
+    private final List<TextWebSocketFrame> textWebSocketFrames;
     private EventLoopGroup bossGroup;
     private EventLoopGroup workerGroup;
 
@@ -50,6 +49,7 @@ public EchoServer(final boolean secure) {
 
     public EchoServer(final boolean secure, final Error error) {
         registeredClients = new ArrayList<>();
+        websocketChannels = new ArrayList<>();
         textWebSocketFrames = new ArrayList<>();
         new Thread(() -> {
             bossGroup = new NioEventLoopGroup(1);
@@ -58,7 +58,7 @@ public EchoServer(final boolean secure, final Error error) {
                 .channel(NioServerSocketChannel.class)
                 .option(ChannelOption.SO_BACKLOG, 100)
                 .handler(new LoggingHandler(EchoServer.class))
-                .childHandler(new EchoServerInitializer(mockServerLogger, secure, error, registeredClients, textWebSocketFrames))
+                .childHandler(new EchoServerInitializer(mockServerLogger, secure, error, registeredClients, websocketChannels, textWebSocketFrames))
                 .childAttr(LOG_FILTER, logFilter)
                 .childAttr(NEXT_RESPONSE, nextResponse)
                 .childAttr(ONLY_RESPONSE, onlyResponse)
@@ -124,11 +124,15 @@ public EchoServer withOnlyResponse(HttpResponse httpResponse) {
         return this;
     }
 
-    public ArrayList<String> getRegisteredClients() {
+    public List<String> getRegisteredClients() {
         return registeredClients;
     }
 
-    public ArrayList<TextWebSocketFrame> getTextWebSocketFrames() {
+    public List<Channel> getWebsocketChannels() {
+        return websocketChannels;
+    }
+
+    public List<TextWebSocketFrame> getTextWebSocketFrames() {
         return textWebSocketFrames;
     }
 

mockserver-core/src/main/java/org/mockserver/echo/http/EchoServerInitializer.java

@@ -1,5 +1,6 @@
 package org.mockserver.echo.http;
 
+import io.netty.channel.Channel;
 import io.netty.channel.ChannelInitializer;
 import io.netty.channel.ChannelPipeline;
 import io.netty.channel.socket.SocketChannel;
@@ -13,6 +14,7 @@
 import org.mockserver.socket.tls.NettySslContextFactory;
 
 import java.util.ArrayList;
+import java.util.List;
 
 import static org.mockserver.echo.http.EchoServer.*;
 import static org.slf4j.event.Level.TRACE;
@@ -25,17 +27,19 @@ public class EchoServerInitializer extends ChannelInitializer<SocketChannel> {
     private final MockServerLogger mockServerLogger;
     private final boolean secure;
     private final EchoServer.Error error;
-    private final ArrayList<TextWebSocketFrame> textWebSocketFrames;
-    private final ArrayList<String> registeredClients;
+    private final List<TextWebSocketFrame> textWebSocketFrames;
+    private final List<Channel> websocketChannels;
+    private final List<String> registeredClients;
 
-    EchoServerInitializer(MockServerLogger mockServerLogger, boolean secure, EchoServer.Error error, ArrayList<String> registeredClients, ArrayList<TextWebSocketFrame> textWebSocketFrames) {
+    EchoServerInitializer(MockServerLogger mockServerLogger, boolean secure, EchoServer.Error error, List<String> registeredClients, List<Channel> websocketChannels, List<TextWebSocketFrame> textWebSocketFrames) {
         if (!secure && error == EchoServer.Error.CLOSE_CONNECTION) {
             throw new IllegalArgumentException("Error type CLOSE_CONNECTION is not supported in non-secure mode");
         }
         this.mockServerLogger = mockServerLogger;
         this.secure = secure;
         this.error = error;
         this.registeredClients = registeredClients;
+        this.websocketChannels = websocketChannels;
         this.textWebSocketFrames = textWebSocketFrames;
     }
 
@@ -60,7 +64,7 @@ public void initChannel(SocketChannel channel) {
 
         pipeline.addLast(new HttpObjectAggregator(Integer.MAX_VALUE));
 
-        pipeline.addLast(new WebSocketServerHandler(mockServerLogger, registeredClients, textWebSocketFrames));
+        pipeline.addLast(new WebSocketServerHandler(mockServerLogger, registeredClients, websocketChannels, textWebSocketFrames));
 
         pipeline.addLast(new MockServerServerCodec(mockServerLogger, secure));
 

mockserver-core/src/main/java/org/mockserver/echo/http/WebSocketServerHandler.java

@@ -1,10 +1,7 @@
 
 package org.mockserver.echo.http;
 
-import io.netty.channel.ChannelFutureListener;
-import io.netty.channel.ChannelHandler;
-import io.netty.channel.ChannelHandlerContext;
-import io.netty.channel.ChannelInboundHandlerAdapter;
+import io.netty.channel.*;
 import io.netty.handler.codec.http.DefaultHttpHeaders;
 import io.netty.handler.codec.http.FullHttpRequest;
 import io.netty.handler.codec.http.websocketx.*;
@@ -30,12 +27,14 @@ public class WebSocketServerHandler extends ChannelInboundHandlerAdapter {
     private static final String UPGRADE_CHANNEL_FOR_CALLBACK_WEB_SOCKET_URI = "/_mockserver_callback_websocket";
     private final MockServerLogger mockServerLogger;
     private final List<String> registeredClients;
+    private final List<Channel> websocketChannels;
     private final List<TextWebSocketFrame> textWebSocketFrames;
     private WebSocketServerHandshaker handshaker;
 
-    WebSocketServerHandler(MockServerLogger mockServerLogger, List<String> registeredClients, List<TextWebSocketFrame> textWebSocketFrames) {
+    WebSocketServerHandler(MockServerLogger mockServerLogger, List<String> registeredClients, List<Channel> websocketChannels, List<TextWebSocketFrame> textWebSocketFrames) {
         this.mockServerLogger = mockServerLogger;
         this.registeredClients = registeredClients;
+        this.websocketChannels = websocketChannels;
         this.textWebSocketFrames = textWebSocketFrames;
     }
 
@@ -93,6 +92,7 @@ private void upgradeChannel(final ChannelHandlerContext ctx, FullHttpRequest htt
                             .setMessageFormat("Registering client " + clientId)
                     );
                     registeredClients.add(clientId);
+                    websocketChannels.add(future.channel());
                     future.channel().closeFuture().addListener((ChannelFutureListener) future1 -> {
                         mockServerLogger.logEvent(
                             new LogEntry()
@@ -101,6 +101,7 @@ private void upgradeChannel(final ChannelHandlerContext ctx, FullHttpRequest htt
                                 .setMessageFormat("Unregistering callback for client " + clientId)
                         );
                         registeredClients.remove(clientId);
+                        websocketChannels.remove(future.channel());
                     });
                 });
         }