chore(agnostification): agnostify web socket connections (#6520)

This PR updates the socket transport code to swap between a Node web
socket transport or a web one based on the `isNode` environment. It also
adds unit tests to the browser tests that show we can connect in a
browser.

src/common/BrowserConnector.ts

@@ -19,9 +19,9 @@ import { Browser } from './Browser.js';
 import { assert } from './assert.js';
 import { debugError } from '../common/helper.js';
 import { Connection } from './Connection.js';
-import { WebSocketTransport } from './WebSocketTransport.js';
 import { getFetch } from './fetch.js';
 import { Viewport } from './PuppeteerViewport.js';
+import { isNode } from '../environment.js';
 
 /**
  * Generic browser options that can be passed when launching any browser.
@@ -33,6 +33,13 @@ export interface BrowserOptions {
   slowMo?: number;
 }
 
+const getWebSocketTransportClass = async () => {
+  return isNode
+    ? (await import('../node/NodeWebSocketTransport.js')).NodeWebSocketTransport
+    : (await import('./BrowserWebSocketTransport.js'))
+        .BrowserWebSocketTransport;
+};
+
 /**
  * Users should never call this directly; it's called when calling
  * `puppeteer.connect`.
@@ -44,7 +51,7 @@ export const connectToBrowser = async (
     browserURL?: string;
     transport?: ConnectionTransport;
   }
-) => {
+): Promise<Browser> => {
   const {
     browserWSEndpoint,
     browserURL,
@@ -64,13 +71,17 @@ export const connectToBrowser = async (
   if (transport) {
     connection = new Connection('', transport, slowMo);
   } else if (browserWSEndpoint) {
-    const connectionTransport = await WebSocketTransport.create(
+    const WebSocketClass = await getWebSocketTransportClass();
+    const connectionTransport: ConnectionTransport = await WebSocketClass.create(
       browserWSEndpoint
     );
     connection = new Connection(browserWSEndpoint, connectionTransport, slowMo);
   } else if (browserURL) {
     const connectionURL = await getWSEndpoint(browserURL);
-    const connectionTransport = await WebSocketTransport.create(connectionURL);
+    const WebSocketClass = await getWebSocketTransportClass();
+    const connectionTransport: ConnectionTransport = await WebSocketClass.create(
+      connectionURL
+    );
     connection = new Connection(connectionURL, connectionTransport, slowMo);
   }
 

src/common/BrowserWebSocketTransport.ts

@@ -0,0 +1,55 @@
+/**
+ * Copyright 2020 Google Inc. All rights reserved.
+ *
+ * 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.
+ */
+import { ConnectionTransport } from './ConnectionTransport.js';
+
+export class BrowserWebSocketTransport implements ConnectionTransport {
+  static create(url: string): Promise<BrowserWebSocketTransport> {
+    return new Promise((resolve, reject) => {
+      const ws = new WebSocket(url);
+
+      ws.addEventListener('open', () =>
+        resolve(new BrowserWebSocketTransport(ws))
+      );
+      ws.addEventListener('error', reject);
+    });
+  }
+
+  private _ws: WebSocket;
+  onmessage?: (message: string) => void;
+  onclose?: () => void;
+
+  constructor(ws: WebSocket) {
+    this._ws = ws;
+    this._ws.addEventListener('message', (event) => {
+      if (this.onmessage) this.onmessage.call(null, event.data);
+    });
+    this._ws.addEventListener('close', () => {
+      if (this.onclose) this.onclose.call(null);
+    });
+    // Silently ignore all errors - we don't know what to do with them.
+    this._ws.addEventListener('error', () => {});
+    this.onmessage = null;
+    this.onclose = null;
+  }
+
+  send(message: string): void {
+    this._ws.send(message);
+  }
+
+  close(): void {
+    this._ws.close();
+  }
+}

src/node/BrowserRunner.ts

@@ -22,7 +22,7 @@ import { assert } from '../common/assert.js';
 import { helper, debugError } from '../common/helper.js';
 import { LaunchOptions } from './LaunchOptions.js';
 import { Connection } from '../common/Connection.js';
-import { WebSocketTransport } from '../common/WebSocketTransport.js';
+import { NodeWebSocketTransport as WebSocketTransport } from '../node/NodeWebSocketTransport.js';
 import { PipeTransport } from './PipeTransport.js';
 import * as readline from 'readline';
 import { TimeoutError } from '../common/Errors.js';

src/node/NodeWebSocketTransport.ts

@@ -13,23 +13,25 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import { ConnectionTransport } from './ConnectionTransport.js';
+import { ConnectionTransport } from '../common/ConnectionTransport.js';
 import NodeWebSocket from 'ws';
 
-export class WebSocketTransport implements ConnectionTransport {
-  static create(url: string): Promise<WebSocketTransport> {
+export class NodeWebSocketTransport implements ConnectionTransport {
+  static create(url: string): Promise<NodeWebSocketTransport> {
     return new Promise((resolve, reject) => {
       const ws = new NodeWebSocket(url, [], {
         perMessageDeflate: false,
         maxPayload: 256 * 1024 * 1024, // 256Mb
       });
 
-      ws.addEventListener('open', () => resolve(new WebSocketTransport(ws)));
+      ws.addEventListener('open', () =>
+        resolve(new NodeWebSocketTransport(ws))
+      );
       ws.addEventListener('error', reject);
     });
   }
 
-  _ws: NodeWebSocket;
+  private _ws: NodeWebSocket;
   onmessage?: (message: string) => void;
   onclose?: () => void;
 
@@ -47,7 +49,7 @@ export class WebSocketTransport implements ConnectionTransport {
     this.onclose = null;
   }
 
-  send(message): void {
+  send(message: string): void {
     this._ws.send(message);
   }
 

test-browser/connection.spec.js

@@ -14,48 +14,42 @@
  * limitations under the License.
  */
 import { Connection } from '../lib/esm/puppeteer/common/Connection.js';
+import { BrowserWebSocketTransport } from '../lib/esm/puppeteer/common/BrowserWebSocketTransport.js';
+import puppeteer from '../lib/esm/puppeteer/web.js';
 import expect from '../node_modules/expect/build-es5/index.js';
+import { getWebSocketEndpoint } from './helper.js';
 
-/**
- * A fake transport that echoes the message it got back and pretends to have got a result.
- *
- * In actual pptr code we expect that `result` is returned from the message
- * being sent with some data, so we fake that in the `send` method.
- *
- * We don't define `onmessage` here because Puppeteer's Connection class will
- * define an `onmessage` for us.
- */
-class EchoTransport {
-  send(message) {
-    const object = JSON.parse(message);
-    const fakeMessageResult = {
-      result: `fake-test-result-${object.method}`,
-    };
-    this.onmessage(
-      JSON.stringify({
-        ...object,
-        ...fakeMessageResult,
-      })
-    );
-  }
+describe('creating a Connection', () => {
+  it('can create a real connection to the backend and send messages', async () => {
+    const wsUrl = getWebSocketEndpoint();
+    const transport = await BrowserWebSocketTransport.create(wsUrl);
 
-  close() {}
-}
-
-describe('Connection', () => {
-  it('can be created in the browser and send/receive messages', async () => {
-    let receivedOutput = '';
-    const connection = new Connection('fake-url', new EchoTransport());
-
-    /**
-     * Puppeteer increments a counter from 0 for each
-     * message it sends So we have to register a callback for the object with
-     * the ID of `1` as the message we send will be the first.
+    const connection = new Connection(wsUrl, transport);
+    const result = await connection.send('Browser.getVersion');
+    /* We can't expect exact results as the version of Chrome/CDP might change
+     * and we don't want flakey tests, so let's assert the structure, which is
+     * enough to confirm the result was recieved successfully.
      */
-    connection._callbacks.set(1, {
-      resolve: (data) => (receivedOutput = data),
+    expect(result).toEqual({
+      protocolVersion: expect.any(String),
+      jsVersion: expect.any(String),
+      revision: expect.any(String),
+      userAgent: expect.any(String),
+      product: expect.any(String),
     });
-    connection.send('Browser.getVersion');
-    expect(receivedOutput).toEqual('fake-test-result-Browser.getVersion');
+  });
+});
+
+describe('puppeteer.connect', () => {
+  it('can connect over websocket and make requests to the backend', async () => {
+    const wsUrl = getWebSocketEndpoint();
+    const browser = await puppeteer.connect({
+      browserWSEndpoint: wsUrl,
+    });
+
+    const version = await browser.version();
+    const versionLooksCorrect = /.+Chrome\/\d{2}/.test(version);
+    expect(version).toEqual(expect.any(String));
+    expect(versionLooksCorrect).toEqual(true);
   });
 });

test-browser/helper.js

@@ -0,0 +1,27 @@
+/**
+ * Copyright 2020 Google Inc. All rights reserved.
+ *
+ * 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.
+ */
+
+/**
+ * Returns the web socket endpoint for the backend of the browser the tests run
+ * in. Used to create connections to that browser in Puppeteer for unit tests.
+ *
+ * It's available on window.__ENV__ because setup code in
+ * web-test-runner.config.js puts it there. If you're changing this code (or
+ * that code), make sure the other is updated accordingly.
+ */
+export function getWebSocketEndpoint() {
+  return window.__ENV__.wsEndpoint;
+}

web-test-runner.config.js

@@ -13,9 +13,22 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+const { chromeLauncher } = require('@web/test-runner-chrome');
 
 module.exports = {
   files: ['test-browser/**/*.spec.js'],
+  browsers: [
+    chromeLauncher({
+      async createPage({ browser }) {
+        const page = await browser.newPage();
+        page.evaluateOnNewDocument((wsEndpoint) => {
+          window.__ENV__ = { wsEndpoint };
+        }, browser.wsEndpoint());
+
+        return page;
+      },
+    }),
+  ],
   plugins: [
     {
       // turn expect UMD into an es module