support multiple websocket clients in a scenario karatelabs#395

aruncsengr · Nov 2, 2018 · a626597 · a626597
1 parent 8cb881e
commit a626597
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 20 deletions.
diff --git a/karate-core/src/main/java/com/intuit/karate/core/ScenarioContext.java b/karate-core/src/main/java/com/intuit/karate/core/ScenarioContext.java
@@ -100,7 +100,7 @@ public class ScenarioContext {
 
     // websocket
     private final Object LOCK = new Object();
-    private WebSocketClient webSocketClient;
+    private List<WebSocketClient> webSocketClients;
     private Object signalResult;
 
     public void logLastPerfEvent(String failureMessage) {
@@ -150,11 +150,7 @@ public HttpRequest getPrevRequest() {
 
     public HttpClient getHttpClient() {
         return client;
-    }
-
-    public WebSocketClient getWebSocketClient() {
-        return webSocketClient;
-    }        
+    }       
 
     public int getCallDepth() {
         return callDepth;
@@ -288,7 +284,7 @@ private ScenarioContext(ScenarioContext sc, ScenarioInfo info) {
         prevRequest = sc.prevRequest;
         prevPerfEvent = sc.prevPerfEvent;
         callResults = sc.callResults;
-        webSocketClient = sc.webSocketClient;
+        webSocketClients = sc.webSocketClients;
         signalResult = sc.signalResult;
     }
 
@@ -825,16 +821,22 @@ public void embed(byte[] bytes, String contentType) {
     }
 
     // websocket / async =======================================================
+
+    private WebSocketClient createWebSocketClient(String url, Consumer<String> textHandler, Consumer<byte[]> binaryHandler) {
+        WebSocketClient webSocketClient = new WebSocketClient(url, textHandler, binaryHandler);
+        if (webSocketClients == null) {
+            webSocketClients = new ArrayList();
+        }
+        webSocketClients.add(webSocketClient);
+        return webSocketClient;
+    }
 
-    public void webSocket(String url, Consumer<String> textHandler) {
-        webSocket(url, textHandler, null);
+    public WebSocketClient webSocket(String url, Consumer<String> textHandler) {
+        return webSocket(url, textHandler, null);
     }
 
-    public void webSocket(String url, Consumer<String> textHandler, Consumer<byte[]> binaryHandler) {
-        if (webSocketClient != null) {
-            webSocketClient.close();
-        }
-        webSocketClient = new WebSocketClient(url, textHandler, binaryHandler);
+    public WebSocketClient webSocket(String url, Consumer<String> textHandler, Consumer<byte[]> binaryHandler) {
+        return createWebSocketClient(url, textHandler, binaryHandler);
     }
 
     public void signal(Object result) {
@@ -902,8 +904,8 @@ public void submit(String name) {
     }
 
     public void stop() {
-        if (webSocketClient != null) {
-            webSocketClient.close();
+        if (webSocketClients != null) {
+            webSocketClients.forEach(WebSocketClient::close);
         }
         if (driver != null) {
             driver.quit();

diff --git a/karate-core/src/main/java/com/intuit/karate/core/ScriptBridge.java b/karate-core/src/main/java/com/intuit/karate/core/ScriptBridge.java
@@ -361,8 +361,7 @@ public WebSocketClient webSocket(String url, Consumer<String> textHandler) {
     }
 
     public WebSocketClient webSocket(String url, Consumer<String> textHandler, Consumer<byte[]> binaryHandler) {
-        context.webSocket(url, textHandler, binaryHandler);
-        return context.getWebSocketClient();
+        return context.webSocket(url, textHandler, binaryHandler);
     }
 
     public void signal(Object result) {

diff --git a/karate-demo/src/test/java/demo/websocket/echo.feature b/karate-demo/src/test/java/demo/websocket/echo.feature
@@ -3,7 +3,7 @@ Feature: public test at
     http://www.websocket.org/echo.html
 
 Scenario: text messages
-    Given def handler = function(msg) { karate.signal(msg) }
+    Given def handler = function(msg){ karate.signal(msg) }
     And def socket = karate.webSocket('ws://echo.websocket.org', handler)
     When eval socket.send('hello world!')
     And def result = karate.listen(5000)
@@ -14,7 +14,7 @@ Scenario: text messages
     Then match result == 'another test'
 
 Scenario: binary message
-    Given def handler = function(msg) { karate.signal(msg) }
+    Given def handler = function(msg){ karate.signal(msg) }
     And def socket = karate.webSocket('ws://echo.websocket.org', null, handler)
     And bytes data = read('../upload/test.pdf')
     When eval socket.sendBytes(data)