feat(websocket): add WebSocket.waitForEvent and isClosed (#4301)

Author: pavelfeldman

docs/api.md

@@ -4158,7 +4158,9 @@ The [WebSocket] class represents websocket connections in the page.
 - [event: 'framereceived'](#event-framereceived)
 - [event: 'framesent'](#event-framesent)
 - [event: 'socketerror'](#event-socketerror)
+- [webSocket.isClosed()](#websocketisclosed)
 - [webSocket.url()](#websocketurl)
+- [webSocket.waitForEvent(event[, optionsOrPredicate])](#websocketwaitforeventevent-optionsorpredicate)
 <!-- GEN:stop -->
 
 #### event: 'close'
@@ -4182,11 +4184,25 @@ Fired when the websocket sends a frame.
 
 Fired when the websocket has an error.
 
+#### webSocket.isClosed()
+- returns: <[boolean]>
+
+Indicates that the web socket has been closed.
+
 #### webSocket.url()
 - returns: <[string]>
 
 Contains the URL of the WebSocket.
 
+#### webSocket.waitForEvent(event[, optionsOrPredicate])
+- `event` <[string]> Event name, same one would pass into `webSocket.on(event)`.
+- `optionsOrPredicate` <[Function]|[Object]> Either a predicate that receives an event or an options object.
+  - `predicate` <[Function]> receives the event data and resolves to truthy value when the waiting should resolve.
+  - `timeout` <[number]> maximum time to wait for in milliseconds. Defaults to `30000` (30 seconds). Pass `0` to disable timeout. The default value can be changed by using the [browserContext.setDefaultTimeout(timeout)](#browsercontextsetdefaulttimeouttimeout) or [page.setDefaultTimeout(timeout)](#pagesetdefaulttimeouttimeout) methods.
+- returns: <[Promise]<[Object]>> Promise which resolves to the event data value.
+
+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
+is fired.
 
 ### class: TimeoutError
 

src/client/network.ts

@@ -18,12 +18,14 @@ import { URLSearchParams } from 'url';
 import * as channels from '../protocol/channels';
 import { ChannelOwner } from './channelOwner';
 import { Frame } from './frame';
-import { Headers } from './types';
+import { Headers, WaitForEventOptions } from './types';
 import * as fs from 'fs';
 import * as mime from 'mime';
 import * as util from 'util';
 import { isString, headersObjectToArray, headersArrayToObject } from '../utils/utils';
 import { Events } from './events';
+import { Page } from './page';
+import { Waiter } from './waiter';
 
 export type NetworkCookie = {
   name: string,
@@ -314,12 +316,17 @@ export class Response extends ChannelOwner<channels.ResponseChannel, channels.Re
 }
 
 export class WebSocket extends ChannelOwner<channels.WebSocketChannel, channels.WebSocketInitializer> {
+  private _page: Page;
+  private _isClosed: boolean;
+
   static from(webSocket: channels.WebSocketChannel): WebSocket {
     return (webSocket as any)._object;
   }
 
   constructor(parent: ChannelOwner, type: string, guid: string, initializer: channels.WebSocketInitializer) {
     super(parent, type, guid, initializer);
+    this._isClosed = false;
+    this._page = parent as Page;
     this._channel.on('frameSent', (event: { opcode: number, data: string }) => {
       const payload = event.opcode === 2 ? Buffer.from(event.data, 'base64') : event.data;
       this.emit(Events.WebSocket.FrameSent, { payload });
@@ -329,12 +336,34 @@ export class WebSocket extends ChannelOwner<channels.WebSocketChannel, channels.
       this.emit(Events.WebSocket.FrameReceived, { payload });
     });
     this._channel.on('error', ({ error }) => this.emit(Events.WebSocket.Error, error));
-    this._channel.on('close', () => this.emit(Events.WebSocket.Close));
+    this._channel.on('close', () => {
+      this._isClosed = true;
+      this.emit(Events.WebSocket.Close);
+    });
   }
 
   url(): string {
     return this._initializer.url;
   }
+
+  isClosed(): boolean {
+    return this._isClosed;
+  }
+
+  async waitForEvent(event: string, optionsOrPredicate: WaitForEventOptions = {}): Promise<any> {
+    const timeout = this._page._timeoutSettings.timeout(typeof optionsOrPredicate === 'function' ? {} : optionsOrPredicate);
+    const predicate = typeof optionsOrPredicate === 'function' ? optionsOrPredicate : optionsOrPredicate.predicate;
+    const waiter = new Waiter();
+    waiter.rejectOnTimeout(timeout, `Timeout while waiting for event "${event}"`);
+    if (event !== Events.WebSocket.Error)
+      waiter.rejectOnEvent(this, Events.WebSocket.Error, new Error('Socket error'));
+    if (event !== Events.WebSocket.Close)
+      waiter.rejectOnEvent(this, Events.WebSocket.Close, new Error('Socket closed'));
+    waiter.rejectOnEvent(this._page, Events.Page.Close, new Error('Page closed'));
+    const result = await waiter.waitForEvent(this, event, predicate as any);
+    waiter.dispose();
+    return result;
+  }
 }
 
 export function validateHeaders(headers: Headers) {

src/server/frames.ts

@@ -352,7 +352,7 @@ export class FrameManager {
 
   onWebSocketResponse(requestId: string, status: number, statusText: string) {
     const ws = this._webSockets.get(requestId);
-    if (status >= 200 && status < 400)
+    if (status < 400)
       return;
     if (ws)
       ws.error(`${statusText}: ${status}`);

test/web-socket.spec.ts

@@ -32,8 +32,10 @@ it('should emit close events', async ({ page, server }) => {
   let socketClosed;
   const socketClosePromise = new Promise(f => socketClosed = f);
   const log = [];
+  let webSocket;
   page.on('websocket', ws => {
     log.push(`open<${ws.url()}>`);
+    webSocket = ws;
     ws.on('close', () => { log.push('close'); socketClosed(); });
   });
   await page.evaluate(port => {
@@ -42,6 +44,7 @@ it('should emit close events', async ({ page, server }) => {
   }, server.PORT);
   await socketClosePromise;
   expect(log.join(':')).toBe(`open<ws://localhost:${server.PORT}/ws>:close`);
+  expect(webSocket.isClosed()).toBeTruthy();
 });
 
 it('should emit frame events', async ({ page, server, isFirefox }) => {
@@ -104,3 +107,43 @@ it('should emit error', async ({page, server, isFirefox}) => {
   else
     expect(message).toContain(': 400');
 });
+
+it('should not have stray error events', async ({page, server, isFirefox}) => {
+  const [ws] = await Promise.all([
+    page.waitForEvent('websocket'),
+    page.evaluate(port => {
+      (window as any).ws = new WebSocket('ws://localhost:' + port + '/ws');
+    }, server.PORT)
+  ]);
+  let error;
+  ws.on('socketerror', e => error = e);
+  await ws.waitForEvent('framereceived');
+  await page.evaluate('window.ws.close()');
+  expect(error).toBeFalsy();
+});
+
+it('should reject waitForEvent on socket close', async ({page, server, isFirefox}) => {
+  const [ws] = await Promise.all([
+    page.waitForEvent('websocket'),
+    page.evaluate(port => {
+      (window as any).ws = new WebSocket('ws://localhost:' + port + '/ws');
+    }, server.PORT)
+  ]);
+  await ws.waitForEvent('framereceived');
+  const error = ws.waitForEvent('framesent').catch(e => e);
+  await page.evaluate('window.ws.close()');
+  expect((await error).message).toContain('Socket closed');
+});
+
+it('should reject waitForEvent on page close', async ({page, server, isFirefox}) => {
+  const [ws] = await Promise.all([
+    page.waitForEvent('websocket'),
+    page.evaluate(port => {
+      (window as any).ws = new WebSocket('ws://localhost:' + port + '/ws');
+    }, server.PORT)
+  ]);
+  await ws.waitForEvent('framereceived');
+  const error = ws.waitForEvent('framesent').catch(e => e);
+  await page.close();
+  expect((await error).message).toContain('Page closed');
+});