feat(api): implement WebSocket API (#86)

api-generator/src/main/java/com/microsoft/playwright/tools/ApiGenerator.java

@@ -364,6 +364,7 @@ class Method extends Element {
     };
     customSignature.put("Page.waitForEvent", waitForEvent);
     customSignature.put("BrowserContext.waitForEvent", waitForEvent);
+    customSignature.put("WebSocket.waitForEvent", waitForEvent);
 
     String[] selectOption = {
       "default List<String> selectOption(String selector, String value) {",
@@ -423,8 +424,9 @@ class Method extends Element {
 
   private static Set<String> skipJavadoc = new HashSet<>(asList(
     "Page.waitForEvent.optionsOrPredicate",
-    "Page.frame.options"
+    "Page.frame.options",
-  ));
+    "WebSocket.waitForEvent.optionsOrPredicate"
+    ));
 
   Method(TypeDefinition parent, JsonObject jsonElement) {
     super(parent, jsonElement);
@@ -737,8 +739,10 @@ class Interface extends TypeDefinition {
     if (asList("Page", "BrowserContext").contains(jsonName)) {
       output.add("import java.util.function.Consumer;");
     }
-    if (asList("Page", "Frame", "BrowserContext").contains(jsonName)) {
+    if (asList("Page", "Frame", "BrowserContext", "WebSocket").contains(jsonName)) {
       output.add("import java.util.function.Predicate;");
+    }
+    if (asList("Page", "Frame", "BrowserContext").contains(jsonName)) {
       output.add("import java.util.regex.Pattern;");
     }
     output.add("");
@@ -909,8 +913,16 @@ class Interface extends TypeDefinition {
         output.add("");
         break;
       }
+      case "WebSocket": {
+        output.add(offset + "interface FrameData {");
+        output.add(offset + "  byte[] body();");
+        output.add(offset + "  String text();");
+        output.add(offset + "}");
+        output.add("");
+        break;
+      }
     }
-    if (asList("Page", "BrowserContext").contains(jsonName)){
+    if (asList("Page", "BrowserContext", "WebSocket").contains(jsonName)){
       output.add(offset + "class WaitForEventOptions {");
       output.add(offset + "  public Integer timeout;");
       output.add(offset + "  public Predicate<Event<EventType>> predicate;");

api-generator/src/main/java/com/microsoft/playwright/tools/Types.java

@@ -196,6 +196,8 @@ class Types {
     add("ConsoleMessage.location", "Object", "Location");
     add("ElementHandle.boundingBox", "Promise<null|Object>", "BoundingBox", new Empty());
     add("Accessibility.snapshot", "Promise<null|Object>", "AccessibilityNode", new Empty());
+    add("WebSocket.framereceived", "Object", "FrameData", new Empty());
+    add("WebSocket.framesent", "Object", "FrameData", new Empty());
 
     add("Page.waitForRequest", "Promise<Request>", "Deferred<Request>");
     add("Page.waitForResponse", "Promise<Response>", "Deferred<Response>");

playwright/pom.xml

@@ -44,5 +44,11 @@
       <groupId>org.junit.jupiter</groupId>
       <artifactId>junit-jupiter-engine</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.java-websocket</groupId>
+      <artifactId>Java-WebSocket</artifactId>
+      <version>1.5.1</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 </project>

playwright/src/main/java/com/microsoft/playwright/WebSocket.java

@@ -17,11 +17,30 @@
 package com.microsoft.playwright;
 
 import java.util.*;
+import java.util.function.Predicate;
 
 /**
  * The WebSocket class represents websocket connections in the page.
  */
 public interface WebSocket {
+  interface FrameData {
+    byte[] body();
+    String text();
+  }
+
+  class WaitForEventOptions {
+    public Integer timeout;
+    public Predicate<Event<EventType>> predicate;
+    public WaitForEventOptions withTimeout(int millis) {
+      timeout = millis;
+      return this;
+    }
+    public WaitForEventOptions withPredicate(Predicate<Event<EventType>> predicate) {
+      this.predicate = predicate;
+      return this;
+    }
+  }
+
   enum EventType {
     CLOSE,
     FRAMERECEIVED,
@@ -31,28 +50,6 @@ public interface WebSocket {
 
   void addListener(EventType type, Listener<EventType> listener);
   void removeListener(EventType type, Listener<EventType> listener);
-  class WebSocketFramereceived {
-    /**
-     * frame payload
-     */
-    public byte[] payload;
-
-    public WebSocketFramereceived withPayload(byte[] payload) {
-      this.payload = payload;
-      return this;
-    }
-  }
-  class WebSocketFramesent {
-    /**
-     * frame payload
-     */
-    public byte[] payload;
-
-    public WebSocketFramesent withPayload(byte[] payload) {
-      this.payload = payload;
-      return this;
-    }
-  }
   /**
    * Indicates that the web socket has been closed.
    */
@@ -61,17 +58,21 @@ public interface WebSocket {
    * Contains the URL of the WebSocket.
    */
   String url();
-  default Object waitForEvent(String event) {
+  default Deferred<Event<EventType>> waitForEvent(EventType event) {
-    return waitForEvent(event, null);
+    return waitForEvent(event, (WaitForEventOptions) null);
+  }
+  default Deferred<Event<EventType>> waitForEvent(EventType event, Predicate<Event<EventType>> predicate) {
+    WaitForEventOptions options = new WaitForEventOptions();
+    options.predicate = predicate;
+    return waitForEvent(event, options);
   }
   /**
    * Waits for event to fire and passes its value into the predicate function. Resolves when the predicate returns truthy value. Will throw an error if the webSocket is closed before the event
    * <p>
    * is fired.
    * @param event Event name, same one would pass into {@code webSocket.on(event)}.
-   * @param optionsOrPredicate Either a predicate that receives an event or an options object.
    * @return Promise which resolves to the event data value.
    */
-  Object waitForEvent(String event, String optionsOrPredicate);
+  Deferred<Event<EventType>> waitForEvent(EventType event, WaitForEventOptions options);
 }
 

playwright/src/main/java/com/microsoft/playwright/impl/Connection.java

@@ -239,6 +239,9 @@ public class Connection {
       case "Selectors":
         result = new SelectorsImpl(parent, type, guid, initializer);
         break;
+      case "WebSocket":
+        result = new WebSocketImpl(parent, type, guid, initializer);
+        break;
       case "Worker":
         result = new WorkerImpl(parent, type, guid, initializer);
         break;

playwright/src/main/java/com/microsoft/playwright/impl/PageImpl.java

@@ -81,6 +81,10 @@ public class PageImpl extends ChannelOwner implements Page {
       worker.page = this;
       workers.add(worker);
       listeners.notify(EventType.WORKER, worker);
+    } else if ("webSocket".equals(event)) {
+      String guid = params.getAsJsonObject("webSocket").get("guid").getAsString();
+      WebSocketImpl webSocket = connection.getExistingObject(guid);
+      listeners.notify(EventType.WEBSOCKET, webSocket);
     } else if ("console".equals(event)) {
       String guid = params.getAsJsonObject("message").get("guid").getAsString();
       ConsoleMessageImpl message = connection.getExistingObject(guid);

playwright/src/main/java/com/microsoft/playwright/impl/WebSocketImpl.java

@@ -0,0 +1,164 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.microsoft.playwright.impl;
+
+import com.google.gson.JsonObject;
+import com.microsoft.playwright.*;
+
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.List;
+
+class WebSocketImpl extends ChannelOwner implements WebSocket {
+  private final ListenerCollection<EventType> listeners = new ListenerCollection<>();
+  private final PageImpl page;
+  private boolean isClosed;
+
+  public WebSocketImpl(ChannelOwner parent, String type, String guid, JsonObject initializer) {
+    super(parent, type, guid, initializer);
+    page = (PageImpl) parent;
+  }
+
+  @Override
+  public void addListener(EventType type, Listener<EventType> listener) {
+    listeners.add(type, listener);
+  }
+
+  @Override
+  public void removeListener(EventType type, Listener<EventType> listener) {
+    listeners.remove(type, listener);
+  }
+
+  @Override
+  public boolean isClosed() {
+    return isClosed;
+  }
+
+  @Override
+  public String url() {
+    return initializer.get("url").getAsString();
+  }
+
+  private class WaitableWebSocketError<R> implements Waitable<R>, Listener<EventType> {
+    private final List<EventType> subscribedEvents;
+    private String errorMessage;
+
+    WaitableWebSocketError() {
+      subscribedEvents = Arrays.asList(EventType.CLOSE, EventType.SOCKETERROR);
+      for (EventType e : subscribedEvents) {
+        addListener(e, this);
+      }
+    }
+
+    @Override
+    public void handle(Event<EventType> event) {
+      if (EventType.SOCKETERROR == event.type()) {
+        errorMessage = "Socket error";
+      } else if (EventType.CLOSE == event.type()) {
+        errorMessage = "Socket closed";
+      } else {
+        return;
+      }
+      dispose();
+    }
+
+    @Override
+    public boolean isDone() {
+      return errorMessage != null;
+    }
+
+    @Override
+    public R get() {
+      throw new PlaywrightException(errorMessage);
+    }
+
+    @Override
+    public void dispose() {
+      for (EventType e : subscribedEvents) {
+        removeListener(e, this);
+      }
+    }
+  }
+
+  @Override
+  public Deferred<Event<EventType>> waitForEvent(EventType event, WaitForEventOptions options) {
+    if (options == null) {
+      options = new WaitForEventOptions();
+    }
+    List<Waitable<Event<EventType>>> waitables = new ArrayList<>();
+    waitables.add(new WaitableEvent<>(listeners, event, options.predicate));
+    waitables.add(new WaitableWebSocketError<>());
+    waitables.add(page.createWaitForCloseHelper());
+    waitables.add(page.createWaitableTimeout(options.timeout));
+    return toDeferred(new WaitableRace<>(waitables));
+  }
+
+  private static class FrameDataImpl implements FrameData {
+    private final byte[] bytes;
+
+    FrameDataImpl(String payload, boolean isBase64) {
+      if (isBase64) {
+        bytes = Base64.getDecoder().decode(payload);
+      } else {
+        bytes = payload.getBytes();
+      }
+    }
+
+    @Override
+    public byte[] body() {
+      return bytes;
+    }
+
+    @Override
+    public String text() {
+      return new String(bytes, StandardCharsets.UTF_8);
+    }
+  }
+
+  @Override
+  void handleEvent(String event, JsonObject parameters) {
+    switch (event) {
+      case "frameSent": {
+        FrameDataImpl frameData = new FrameDataImpl(
+          parameters.get("data").getAsString(), parameters.get("opcode").getAsInt() == 2);
+        listeners.notify(EventType.FRAMESENT, frameData);
+        break;
+      }
+      case "frameReceived": {
+        FrameDataImpl frameData = new FrameDataImpl(
+          parameters.get("data").getAsString(), parameters.get("opcode").getAsInt() == 2);
+        listeners.notify(EventType.FRAMERECEIVED, frameData);
+        break;
+      }
+      case "socketError": {
+        String error = parameters.get("error").getAsString();
+        listeners.notify(EventType.SOCKETERROR, error);
+        break;
+      }
+      case "close": {
+        isClosed = true;
+        listeners.notify(EventType.CLOSE, null);
+        break;
+      }
+      default: {
+        throw new PlaywrightException("Unknown event: " + event);
+      }
+    }
+  }
+}

playwright/src/test/java/com/microsoft/playwright/TestWebSocket.java

@@ -0,0 +1,252 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.microsoft.playwright;
+
+import org.java_websocket.WebSocket;
+import org.java_websocket.drafts.Draft;
+import org.java_websocket.exceptions.InvalidDataException;
+import org.java_websocket.framing.Framedata;
+import org.java_websocket.handshake.ClientHandshake;
+import org.java_websocket.handshake.ServerHandshakeBuilder;
+import org.java_websocket.server.WebSocketServer;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.nio.channels.SelectionKey;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.List;
+
+import static com.microsoft.playwright.Page.EventType.WEBSOCKET;
+import static com.microsoft.playwright.WebSocket.EventType.*;
+import static java.util.Arrays.asList;
+import static org.junit.jupiter.api.Assertions.*;
+
+public class TestWebSocket extends TestBase {
+  private static WebSocketServerImpl webSocketServer;
+  private static int WS_SERVER_PORT = 8910;
+
+  private static class WebSocketServerImpl extends WebSocketServer {
+    WebSocketServerImpl(InetSocketAddress address) {
+      super(address, 1);
+    }
+
+    @Override
+    public void onOpen(org.java_websocket.WebSocket webSocket, ClientHandshake clientHandshake) {
+      webSocket.send("incoming");
+    }
+
+    @Override
+    public void onClose(org.java_websocket.WebSocket webSocket, int i, String s, boolean b) { }
+
+    @Override
+    public void onMessage(org.java_websocket.WebSocket webSocket, String s) { }
+
+    @Override
+    public void onError(WebSocket webSocket, Exception e) { }
+
+    @Override
+    public void onStart() { }
+  }
+
+  @BeforeAll
+  static void startWebSockerServer() {
+    webSocketServer = new WebSocketServerImpl(new InetSocketAddress("localhost", WS_SERVER_PORT));
+    new Thread(webSocketServer).start();
+  }
+
+  @AfterAll
+  static void stopWebSockerServer() throws IOException, InterruptedException {
+    webSocketServer.stop();
+  }
+
+
+  private void waitForCondition(boolean[] condition) {
+    assertEquals(1, condition.length);
+    Instant start = Instant.now();
+    while (!condition[0]) {
+      page.waitForTimeout(100).get();
+      assertTrue(Duration.between(start, Instant.now()).getSeconds() < 30, "Timed out");
+    }
+  }
+
+  @Test
+  void shouldWork() {
+    Object value = page.evaluate("port => {\n" +
+      "  let cb;\n" +
+      "  const result = new Promise(f => cb = f);\n" +
+      "  const ws = new WebSocket('ws://localhost:' + port + '/ws');\n" +
+      "  ws.addEventListener('message', data => { ws.close(); cb(data.data); });\n" +
+      "  return result;\n" +
+      "}", webSocketServer.getPort());
+    assertEquals("incoming", value);
+  }
+
+  @Test
+  void shouldEmitCloseEvents() {
+    boolean[] socketClosed = {false};
+    List<String> log = new ArrayList<>();
+    com.microsoft.playwright.WebSocket[] webSocket = {null};
+    page.addListener(WEBSOCKET, event -> {
+      com.microsoft.playwright.WebSocket ws = (com.microsoft.playwright.WebSocket) event.data();
+      log.add("open<" + ws.url() + ">");
+      webSocket[0] = ws;
+      ws.addListener(com.microsoft.playwright.WebSocket.EventType.CLOSE, closeEvent -> {
+        log.add("close");
+        socketClosed[0] = true;
+      });
+    });
+    page.evaluate("port => {\n" +
+      "  const ws = new WebSocket('ws://localhost:' + port + '/ws');\n" +
+      "  ws.addEventListener('open', () => ws.close());\n" +
+      "}", webSocketServer.getPort());
+    waitForCondition(socketClosed);
+    assertEquals(asList("open<ws://localhost:" + webSocketServer.getPort() + "/ws>", "close"), log);
+    assertTrue(webSocket[0].isClosed());
+  }
+
+  @Test
+  void shouldEmitFrameEvents() {
+    boolean[] socketClosed = {false};
+    List<String> log = new ArrayList<>();
+    page.addListener(WEBSOCKET, event -> {
+      com.microsoft.playwright.WebSocket ws = (com.microsoft.playwright.WebSocket) event.data();
+      log.add("open");
+      ws.addListener(FRAMESENT, e -> log.add("sent<" + ((com.microsoft.playwright.WebSocket.FrameData) e.data()).text() + ">"));
+      ws.addListener(FRAMERECEIVED, e -> log.add("received<" + ((com.microsoft.playwright.WebSocket.FrameData) e.data()).text()  + ">"));
+      ws.addListener(CLOSE, e -> { log.add("close"); socketClosed[0] = true; });
+    });
+    page.evaluate("port => {\n" +
+      "    const ws = new WebSocket('ws://localhost:' + port + '/ws');\n" +
+      "    ws.addEventListener('open', () => ws.send('outgoing'));\n" +
+      "    ws.addEventListener('message', () => { ws.close(); });\n" +
+      "  }", webSocketServer.getPort());
+    waitForCondition(socketClosed);
+    assertEquals("open", log.get(0));
+    assertEquals("close", log.get(3));
+    log.sort(String::compareTo);
+    assertEquals(asList("close", "open", "received<incoming>", "sent<outgoing>"), log);
+  }
+
+  @Test
+  void shouldEmitBinaryFrameEvents() {
+    boolean[] socketClosed = {false};
+    List<com.microsoft.playwright.WebSocket.FrameData> sent = new ArrayList<>();
+    page.addListener(WEBSOCKET, event -> {
+      com.microsoft.playwright.WebSocket ws = (com.microsoft.playwright.WebSocket) event.data();
+      ws.addListener(CLOSE, e -> { socketClosed[0] = true; });
+      ws.addListener(FRAMESENT, e -> sent.add((com.microsoft.playwright.WebSocket.FrameData) e.data()));
+    });
+    page.evaluate("port => {\n" +
+      "  const ws = new WebSocket('ws://localhost:' + port + '/ws');\n" +
+      "  ws.addEventListener('open', () => {\n" +
+      "    const binary = new Uint8Array(5);\n" +
+      "    for (let i = 0; i < 5; ++i)\n" +
+      "      binary[i] = i;\n" +
+      "    ws.send('text');\n" +
+      "    ws.send(binary);\n" +
+      "    ws.close();\n" +
+      "  });\n" +
+      "}", webSocketServer.getPort());
+    waitForCondition(socketClosed);
+    assertEquals("text", sent.get(0).text());
+    for (int i = 0; i < 5; ++i) {
+      assertEquals(i, sent.get(1).body()[i]);
+    }
+  }
+
+  @Test
+  void shouldEmitError() {
+    boolean[] socketError = {false};
+    String[] error = {null};
+    page.addListener(WEBSOCKET, event -> {
+      com.microsoft.playwright.WebSocket ws = (com.microsoft.playwright.WebSocket) event.data();
+      ws.addListener(SOCKETERROR, e -> {
+        error[0] = (String) e.data();
+        socketError[0] = true;
+      });
+    });
+    page.evaluate("port => {\n" +
+      "  new WebSocket('ws://localhost:' + port + '/bogus-ws');\n" +
+      "}", server.PORT);
+    waitForCondition(socketError);
+    if (isFirefox()) {
+      assertEquals("CLOSE_ABNORMAL", error[0]);
+    } else {
+      assertTrue(error[0].contains("404"), error[0]);
+    }
+  }
+
+  @Test
+  void shouldNotHaveStrayErrorEvents() {
+    Deferred<Event<Page.EventType>> wsEvent = page.waitForEvent(WEBSOCKET);
+    page.evaluate("port => {\n" +
+      "  window.ws = new WebSocket('ws://localhost:' + port + '/ws');\n" +
+      "}", webSocketServer.getPort());
+
+    com.microsoft.playwright.WebSocket ws = (com.microsoft.playwright.WebSocket) wsEvent.get().data();
+    boolean[] error = {false};
+    ws.addListener(SOCKETERROR, e -> error[0] = true);
+    Deferred<Event<com.microsoft.playwright.WebSocket.EventType>> frameReceivedEvent = ws.waitForEvent(FRAMERECEIVED);
+    frameReceivedEvent.get();
+    System.out.println("will close");
+    page.evaluate("window.ws.close()");
+    assertFalse(error[0]);
+  }
+
+  @Test
+  void shouldRejectWaitForEventOnSocketClose() {
+    Deferred<Event<Page.EventType>> wsEvent = page.waitForEvent(WEBSOCKET);
+    page.evaluate("port => {\n" +
+      "  window.ws = new WebSocket('ws://localhost:' + port + '/ws');\n" +
+      "}", webSocketServer.getPort());
+
+    com.microsoft.playwright.WebSocket ws = (com.microsoft.playwright.WebSocket) wsEvent.get().data();
+    ws.waitForEvent(FRAMERECEIVED).get();
+    Deferred<Event<com.microsoft.playwright.WebSocket.EventType>> frameSentEvent = ws.waitForEvent(FRAMESENT);
+    page.evaluate("window.ws.close()");
+    try {
+      frameSentEvent.get();
+      fail("did not throw");
+    } catch (PlaywrightException exception) {
+      assertTrue(exception.getMessage().contains("Socket closed"));
+    }
+  }
+
+  @Test
+  void shouldRejectWaitForEventOnPageClose() {
+    Deferred<Event<Page.EventType>> wsEvent = page.waitForEvent(WEBSOCKET);
+    page.evaluate("port => {\n" +
+      "  window.ws = new WebSocket('ws://localhost:' + port + '/ws');\n" +
+      "}", webSocketServer.getPort());
+
+    com.microsoft.playwright.WebSocket ws = (com.microsoft.playwright.WebSocket) wsEvent.get().data();
+    ws.waitForEvent(FRAMERECEIVED).get();
+    Deferred<Event<com.microsoft.playwright.WebSocket.EventType>> frameSentEvent = ws.waitForEvent(FRAMESENT);
+    page.close();
+    try {
+      frameSentEvent.get();
+      fail("did not throw");
+    } catch (PlaywrightException exception) {
+      assertTrue(exception.getMessage().contains("Page closed"));
+    }
+  }
+}

scripts/download_driver.sh

@@ -3,7 +3,7 @@
 set -e
 set +x
 
-FILE_PREFIX=playwright-cli-0.170.0-next.1605573954344
+FILE_PREFIX=playwright-cli-0.170.0-next.1607022026758
 
 trap "cd $(pwd -P)" EXIT
 cd "$(dirname $0)"