Websocket acceptance tests now can use websockets (#851)

PegaSysEng · Feb 13, 2019 · 9e872a5 · 9e872a5
1 parent e603d3e
commit 9e872a5
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 24 deletions.
diff --git a/...nce-tests/src/test/java/tech/pegasys/pantheon/tests/acceptance/dsl/node/PantheonNode.java b/...nce-tests/src/test/java/tech/pegasys/pantheon/tests/acceptance/dsl/node/PantheonNode.java
@@ -43,7 +43,9 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.concurrent.TimeUnit;
@@ -187,23 +189,34 @@ public String hostName() {
   }
 
   private JsonRequestFactories jsonRequestFactories() {
+    Optional<WebSocketService> websocketService = Optional.empty();
     if (jsonRequestFactories == null) {
-      final Optional<String> baseUrl;
-      final String port;
+      final Web3jService web3jService;
+
       if (useWsForJsonRpc) {
-        baseUrl = wsRpcBaseUrl();
-        port = "8546";
+        final String url = wsRpcBaseUrl().orElse("ws://" + LOCALHOST + ":" + 8546);
+        final Map<String, String> headers = new HashMap<>();
+        if (token != null) {
+          headers.put("Bearer", token);
+        }
+        final WebSocketClient wsClient = new WebSocketClient(URI.create(url), headers);
+
+        web3jService = new WebSocketService(wsClient, false);
+        try {
+          ((WebSocketService) web3jService).connect();
+        } catch (ConnectException e) {
+          throw new RuntimeException(e);
+        }
+
+        websocketService = Optional.of((WebSocketService) web3jService);
       } else {
-        baseUrl = jsonRpcBaseUrl();
-        port = "8545";
-      }
-      final Web3jService web3jService =
-          baseUrl
-              .map(url -> new HttpService(url))
-              .orElse(new HttpService("http://" + LOCALHOST + ":" + port));
-
-      if (token != null) {
-        ((HttpService) web3jService).addHeader("Bearer", token);
+        web3jService =
+            jsonRpcBaseUrl()
+                .map(HttpService::new)
+                .orElse(new HttpService("http://" + LOCALHOST + ":" + 8545));
+        if (token != null) {
+          ((HttpService) web3jService).addHeader("Bearer", token);
+        }
       }
 
       jsonRequestFactories =
@@ -212,7 +225,8 @@ private JsonRequestFactories jsonRequestFactories() {
               new CliqueJsonRpcRequestFactory(web3jService),
               new IbftJsonRpcRequestFactory(web3jService),
               new PermissioningJsonRpcRequestFactory(web3jService),
-              new AdminJsonRpcRequestFactory(web3jService));
+              new AdminJsonRpcRequestFactory(web3jService),
+              websocketService);
     }
 
     return jsonRequestFactories;
@@ -242,22 +256,16 @@ public void useWebSocketsForJsonRpc() {
 
     checkIfWebSocketEndpointIsAvailable(url);
 
-    final WebSocketService webSocketService = new WebSocketService(url, true);
-    try {
-      webSocketService.connect();
-    } catch (final ConnectException e) {
-      throw new RuntimeException("Error connection to WebSocket endpoint", e);
-    }
+    useWsForJsonRpc = true;
 
     if (jsonRequestFactories != null) {
       jsonRequestFactories.shutdown();
+      jsonRequestFactories = null;
     }
 
     if (httpRequestFactory != null) {
       httpRequestFactory = null;
     }
-
-    useWsForJsonRpc = true;
   }
 
   /** All future JSON-RPC calls will include the authentication token. */
@@ -266,6 +274,7 @@ public void useAuthenticationTokenInHeaderForJsonRpc(final String token) {
 
     if (jsonRequestFactories != null) {
       jsonRequestFactories.shutdown();
+      jsonRequestFactories = null;
     }
 
     if (httpRequestFactory != null) {

diff --git a/...est/java/tech/pegasys/pantheon/tests/acceptance/dsl/transaction/JsonRequestFactories.java b/...est/java/tech/pegasys/pantheon/tests/acceptance/dsl/transaction/JsonRequestFactories.java
@@ -12,7 +12,10 @@
  */
 package tech.pegasys.pantheon.tests.acceptance.dsl.transaction;
 
+import java.util.Optional;
+
 import org.web3j.protocol.core.JsonRpc2_0Web3j;
+import org.web3j.protocol.websocket.WebSocketService;
 
 public class JsonRequestFactories {
 
@@ -21,18 +24,21 @@ public class JsonRequestFactories {
   private final IbftJsonRpcRequestFactory ibft;
   private final PermissioningJsonRpcRequestFactory perm;
   private final AdminJsonRpcRequestFactory admin;
+  private final Optional<WebSocketService> websocketService;
 
   public JsonRequestFactories(
       final JsonRpc2_0Web3j netEth,
       final CliqueJsonRpcRequestFactory clique,
       final IbftJsonRpcRequestFactory ibft,
       final PermissioningJsonRpcRequestFactory perm,
-      final AdminJsonRpcRequestFactory admin) {
+      final AdminJsonRpcRequestFactory admin,
+      final Optional<WebSocketService> websocketService) {
     this.netEth = netEth;
     this.clique = clique;
     this.ibft = ibft;
     this.perm = perm;
     this.admin = admin;
+    this.websocketService = websocketService;
   }
 
   public JsonRpc2_0Web3j eth() {
@@ -61,5 +67,6 @@ public AdminJsonRpcRequestFactory admin() {
 
   public void shutdown() {
     netEth.shutdown();
+    websocketService.ifPresent(WebSocketService::close);
   }
 }