From 15c52d91884988dc3c714739f901b510eb6ffc56 Mon Sep 17 00:00:00 2001
From: Vitaly Takmazov
Date: Sun, 27 May 2018 16:45:59 +0300
Subject: ws: again, accept commands

---
 .../juick/server/tests/WebsocketClientHandler.java | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

(limited to 'juick-server/src/test/java/com/juick/server/tests/WebsocketClientHandler.java')

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;
+    }
 }
-- 
cgit v1.2.3