diff --git a/docs/api.md b/docs/api.md index c6b496963e..9d5df84b99 100644 --- a/docs/api.md +++ b/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) #### 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 diff --git a/src/client/network.ts b/src/client/network.ts index 496fde703b..da9b366577 100644 --- a/src/client/network.ts +++ b/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 { + 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 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 { + 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) { diff --git a/src/server/frames.ts b/src/server/frames.ts index cf99a1bff5..69c430688f 100644 --- a/src/server/frames.ts +++ b/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}`); diff --git a/test/web-socket.spec.ts b/test/web-socket.spec.ts index ce36ad879b..b563ea718f 100644 --- a/test/web-socket.spec.ts +++ b/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: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'); +});