3 files changed, 26 insertions, 3 deletions
diff --git a/juick-server/src/test/java/com/juick/server/tests/ServerTests.java b/juick-server/src/test/java/com/juick/server/tests/ServerTests.java
index 8e3e5dfd..9a78eac9 100644
--- a/juick-server/src/test/java/com/juick/server/tests/ServerTests.java
+++ b/juick-server/src/test/java/com/juick/server/tests/ServerTests.java
@@ -57,6 +57,9 @@ import org.springframework.util.LinkedMultiValueMap;
 import org.springframework.util.MultiValueMap;
 import org.springframework.web.client.RestTemplate;
 import org.springframework.web.context.WebApplicationContext;
+import org.springframework.web.socket.TextMessage;
+import org.springframework.web.socket.client.WebSocketConnectionManager;
+import org.springframework.web.socket.client.standard.StandardWebSocketClient;
 import org.springframework.web.util.UriComponents;
 import org.springframework.web.util.UriComponentsBuilder;
 import org.w3c.dom.Document;
@@ -841,6 +844,8 @@ public class ServerTests {
                 WebsocketClientConfig.class, PropertyPlaceholderAutoConfiguration.class)
                 .run("--spring.main.web_environment=false");
         long count = context.getBean(WebsocketClientConfig.class).getLatch().getCount();
+        WebsocketClientHandler handler = context.getBean(WebsocketClientHandler.class);
+        assertThat(handler.getLastMessage(), is("PONG"));
         context.close();
         Assert.assertThat(count, equalTo(0L));
         mockMvc.perform(get("/ws/_all")).andExpect(status().isBadRequest());
diff --git a/juick-server/src/test/java/com/juick/server/tests/WebsocketClientConfig.java b/juick-server/src/test/java/com/juick/server/tests/WebsocketClientConfig.java
index 6be8e3c6..074be9b0 100644
--- a/juick-server/src/test/java/com/juick/server/tests/WebsocketClientConfig.java
+++ b/juick-server/src/test/java/com/juick/server/tests/WebsocketClientConfig.java
@@ -28,7 +28,7 @@ public class WebsocketClientConfig implements CommandLineRunner {
     public WebSocketConnectionManager connectionManager() {
         WebSocketConnectionManager manager = new WebSocketConnectionManager(client(),
                 handler(),
-                "ws://localhost:8080/ws/");
+                "ws://localhost:8080/ws/_all");
         manager.setAutoStartup(true);
         return manager;
     }
diff --git a/juick-server/src/test/java/com/juick/server/tests/WebsocketClientHandler.java b/juick-server/src/test/java/com/juick/server/tests/WebsocketClientHandler.java
index 7d896a04..ae92ccf0 100644
--- a/juick-server/src/test/java/com/juick/server/tests/WebsocketClientHandler.java
+++ b/juick-server/src/test/java/com/juick/server/tests/WebsocketClientHandler.java
@@ -1,18 +1,36 @@
 package com.juick.server.tests;
 
+import org.springframework.web.socket.TextMessage;
 import org.springframework.web.socket.WebSocketSession;
 import org.springframework.web.socket.handler.TextWebSocketHandler;
 
+import java.io.IOException;
 import java.util.concurrent.CountDownLatch;
 
 class WebsocketClientHandler extends TextWebSocketHandler {
     private final CountDownLatch latch;
+    private WebSocketSession session;
+    private String lastMessage;
     public WebsocketClientHandler(CountDownLatch latch) {
         this.latch = latch;
     }
     @Override
-    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
-        super.afterConnectionEstablished(session);
+    public void afterConnectionEstablished(WebSocketSession session) throws IOException {
+        this.session = session;
+        session.sendMessage(new TextMessage("PING"));
+    }
+
+    @Override
+    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
+        this.lastMessage = message.getPayload();
         latch.countDown();
     }
+
+    public WebSocketSession getSession() {
+        return session;
+    }
+
+    public String getLastMessage() {
+        return lastMessage;
+    }
 }