feat(routeWebSocket): address api review feedback (#32850)

This commit is contained in:
Dmitry Gozman 2024-09-27 04:01:31 -07:00 committed by GitHub
parent d6f584c2d4
commit a395fb22c4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 632 additions and 333 deletions

View File

@ -3686,65 +3686,54 @@ Note that only `WebSocket`s created after this method was called will be routed.
**Usage**
Below is an example of a simple handler that blocks some websocket messages.
See [WebSocketRoute] for more details and examples.
Below is an example of a simple mock that responds to a single message. See [WebSocketRoute] for more details and examples.
```js
await page.routeWebSocket('/ws', async ws => {
ws.routeSend(message => {
if (message === 'to-be-blocked')
return;
ws.send(message);
await page.routeWebSocket('/ws', ws => {
ws.onMessage(message => {
if (message === 'request')
ws.send('response');
});
await ws.connect();
});
```
```java
page.routeWebSocket("/ws", ws -> {
ws.routeSend(message -> {
if ("to-be-blocked".equals(message))
return;
ws.send(message);
ws.onMessage(message -> {
if ("request".equals(message))
ws.send("response");
});
ws.connect();
});
```
```python async
def message_handler(ws: WebSocketRoute, message: Union[str, bytes]):
if message == "to-be-blocked":
return
ws.send(message)
if message == "request":
ws.send("response")
async def handler(ws: WebSocketRoute):
ws.route_send(lambda message: message_handler(ws, message))
await ws.connect()
def handler(ws: WebSocketRoute):
ws.on_message(lambda message: message_handler(ws, message))
await page.route_web_socket("/ws", handler)
```
```python sync
def message_handler(ws: WebSocketRoute, message: Union[str, bytes]):
if message == "to-be-blocked":
return
ws.send(message)
if message == "request":
ws.send("response")
def handler(ws: WebSocketRoute):
ws.route_send(lambda message: message_handler(ws, message))
ws.connect()
ws.on_message(lambda message: message_handler(ws, message))
page.route_web_socket("/ws", handler)
```
```csharp
await page.RouteWebSocketAsync("/ws", async ws => {
ws.RouteSend(message => {
if (message == "to-be-blocked")
return;
ws.Send(message);
await page.RouteWebSocketAsync("/ws", ws => {
ws.OnMessage(message => {
if (message == "request")
ws.Send("response");
});
await ws.ConnectAsync();
});
```

View File

@ -1,69 +1,222 @@
# class: WebSocketRoute
* since: v1.48
Whenever a [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) route is set up with [`method: Page.routeWebSocket`] or [`method: BrowserContext.routeWebSocket`], the `WebSocketRoute` object allows to handle the WebSocket.
Whenever a [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) route is set up with [`method: Page.routeWebSocket`] or [`method: BrowserContext.routeWebSocket`], the `WebSocketRoute` object allows to handle the WebSocket, like an actual server would do.
By default, the routed WebSocket will not actually connect to the server. This way, you can mock entire communcation over the WebSocket. Here is an example that responds to a `"query"` with a `"result"`.
**Mocking**
By default, the routed WebSocket will not connect to the server. This way, you can mock entire communcation over the WebSocket. Here is an example that responds to a `"request"` with a `"response"`.
```js
await page.routeWebSocket('/ws', async ws => {
ws.routeSend(message => {
if (message === 'query')
ws.receive('result');
await page.routeWebSocket('/ws', ws => {
ws.onMessage(message => {
if (message === 'request')
ws.send('response');
});
});
```
```java
page.routeWebSocket("/ws", ws -> {
ws.routeSend(message -> {
if ("query".equals(message))
ws.receive("result");
ws.onMessage(message -> {
if ("request".equals(message))
ws.send("response");
});
});
```
```python async
def message_handler(ws, message):
if message == "query":
ws.receive("result")
if message == "request":
ws.send("response")
await page.route_web_socket("/ws", lambda ws: ws.route_send(
await page.route_web_socket("/ws", lambda ws: ws.on_message(
lambda message: message_handler(ws, message)
))
```
```python sync
def message_handler(ws, message):
if message == "query":
ws.receive("result")
if message == "request":
ws.send("response")
page.route_web_socket("/ws", lambda ws: ws.route_send(
page.route_web_socket("/ws", lambda ws: ws.on_message(
lambda message: message_handler(ws, message)
))
```
```csharp
await page.RouteWebSocketAsync("/ws", async ws => {
ws.RouteSend(message => {
if (message == "query")
ws.receive("result");
await page.RouteWebSocketAsync("/ws", ws => {
ws.OnMessage(message => {
if (message == "request")
ws.Send("response");
});
});
```
Since we do not call [`method: WebSocketRoute.connectToServer`] inside the WebSocket route handler, Playwright assumes that WebSocket will be mocked, and opens the WebSocket inside the page automatically.
## event: WebSocketRoute.close
* since: v1.48
**Intercepting**
Emitted when the [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) closes.
Alternatively, you may want to connect to the actual server, but intercept messages in-between and modify or block them. Calling [`method: WebSocketRoute.connectToServer`] returns a server-side `WebSocketRoute` instance that you can send messages to, or handle incoming messages.
Below is an example that modifies some messages sent by the page to the server. Messages sent from the server to the page are left intact, relying on the default forwarding.
```js
await page.routeWebSocket('/ws', ws => {
const server = ws.connectToServer();
ws.onMessage(message => {
if (message === 'request')
server.send('request2');
else
server.send(message);
});
});
```
```java
page.routeWebSocket("/ws", ws -> {
WebSocketRoute server = ws.connectToServer();
ws.onMessage(message -> {
if ("request".equals(message))
server.send("request2");
else
server.send(message);
});
});
```
```python async
def message_handler(server: WebSocketRoute, message: Union[str, bytes]):
if message == "request":
server.send("request2")
else:
server.send(message)
def handler(ws: WebSocketRoute):
server = ws.connect_to_server()
ws.on_message(lambda message: message_handler(server, message))
await page.route_web_socket("/ws", handler)
```
```python sync
def message_handler(server: WebSocketRoute, message: Union[str, bytes]):
if message == "request":
server.send("request2")
else:
server.send(message)
def handler(ws: WebSocketRoute):
server = ws.connect_to_server()
ws.on_message(lambda message: message_handler(server, message))
page.route_web_socket("/ws", handler)
```
```csharp
await page.RouteWebSocketAsync("/ws", ws => {
var server = ws.ConnectToServer();
ws.OnMessage(message => {
if (message == "request")
server.Send("request2");
else
server.Send(message);
});
});
```
After connecting to the server, all **messages are forwarded** between the page and the server by default.
However, if you call [`method: WebSocketRoute.onMessage`] on the original route, messages from the page to the server **will not be forwarded** anymore, but should instead be handled by the [`param: WebSocketRoute.onMessage.handler`].
Similarly, calling [`method: WebSocketRoute.onMessage`] on the server-side WebSocket will **stop forwarding messages** from the server to the page, and [`param: WebSocketRoute.onMessage.handler`] should take care of them.
The following example blocks some messages in both directions. Since it calls [`method: WebSocketRoute.onMessage`] in both directions, there is no automatic forwarding at all.
```js
await page.routeWebSocket('/ws', ws => {
const server = ws.connectToServer();
ws.onMessage(message => {
if (message !== 'blocked-from-the-page')
server.send(message);
});
server.onMessage(message => {
if (message !== 'blocked-from-the-server')
ws.send(message);
});
});
```
```java
page.routeWebSocket("/ws", ws -> {
WebSocketRoute server = ws.connectToServer();
ws.onMessage(message -> {
if (!"blocked-from-the-page".equals(message))
server.send(message);
});
server.onMessage(message -> {
if (!"blocked-from-the-server".equals(message))
ws.send(message);
});
});
```
```python async
def ws_message_handler(server: WebSocketRoute, message: Union[str, bytes]):
if message != "blocked-from-the-page":
server.send(message)
def server_message_handler(ws: WebSocketRoute, message: Union[str, bytes]):
if message != "blocked-from-the-server":
ws.send(message)
def handler(ws: WebSocketRoute):
server = ws.connect_to_server()
ws.on_message(lambda message: ws_message_handler(server, message))
server.on_message(lambda message: server_message_handler(ws, message))
await page.route_web_socket("/ws", handler)
```
```python sync
def ws_message_handler(server: WebSocketRoute, message: Union[str, bytes]):
if message != "blocked-from-the-page":
server.send(message)
def server_message_handler(ws: WebSocketRoute, message: Union[str, bytes]):
if message != "blocked-from-the-server":
ws.send(message)
def handler(ws: WebSocketRoute):
server = ws.connect_to_server()
ws.on_message(lambda message: ws_message_handler(server, message))
server.on_message(lambda message: server_message_handler(ws, message))
page.route_web_socket("/ws", handler)
```
```csharp
await page.RouteWebSocketAsync("/ws", ws => {
var server = ws.ConnectToServer();
ws.OnMessage(message => {
if (message != "blocked-from-the-page")
server.Send(message);
});
server.OnMessage(message => {
if (message != "blocked-from-the-server")
ws.Send(message);
});
});
```
## async method: WebSocketRoute.close
* since: v1.48
Closes the server connection and the [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) object in the page.
Closes one side of the WebSocket connection.
### option: WebSocketRoute.close.code
* since: v1.48
@ -78,79 +231,67 @@ Optional [close code](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket
Optional [close reason](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close#reason).
## async method: WebSocketRoute.connect
## method: WebSocketRoute.connectToServer
* since: v1.48
- returns: <[WebSocketRoute]>
By default, routed WebSocket does not connect to the server, so you can mock entire WebSocket communication. This method connects to the actual WebSocket server, and returns the server-side [WebSocketRoute] instance, giving the ability to send and receive messages from the server.
Once connected to the server:
* Messages received from the server will be **automatically forwarded** to the WebSocket in the page, unless [`method: WebSocketRoute.onMessage`] is called on the server-side `WebSocketRoute`.
* Messages sent by the [`WebSocket.send()`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/send) call in the page will be **automatically forwarded** to the server, unless [`method: WebSocketRoute.onMessage`] is called on the original `WebSocketRoute`.
See examples at the top for more details.
## method: WebSocketRoute.onClose
* since: v1.48
By default, routed WebSocket does not connect to the server, so you can mock entire WebSocket communication. This method connects to the actual WebSocket server, giving the ability to send and receive messages from the server.
Allows to handle [`WebSocket.close`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close).
Once connected:
* Messages received from the server will be automatically dispatched to the [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) object in the page, unless [`method: WebSocketRoute.routeReceive`] is called.
* Messages sent by the `WebSocket.send()` call in the page will be automatically sent to the server, unless [`method: WebSocketRoute.routeSend`] is called.
By default, closing one side of the connection, either in the page or on the server, will close the other side. However, when [`method: WebSocketRoute.onClose`] handler is set up, the default forwarding of closure is disabled, and handler should take care of it.
### param: WebSocketRoute.onClose.handler
* since: v1.48
- `handler` <[function]\([number]|[undefined], [string]|[undefined]\): [Promise<any>|any]>
Function that will handle WebSocket closure. Received an optional [close code](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close#code) and an optional [close reason](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close#reason).
## method: WebSocketRoute.receive
## async method: WebSocketRoute.onMessage
* since: v1.48
Dispatches a message to the [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) object in the page, like it was received from the server.
This method allows to handle messages that are sent by the WebSocket, either from the page or from the server.
### param: WebSocketRoute.receive.message
* since: v1.48
- `message` <[string]|[Buffer]>
When called on the original WebSocket route, this method handles messages sent from the page. You can handle this messages by responding to them with [`method: WebSocketRoute.send`], forwarding them to the server-side connection returned by [`method: WebSocketRoute.connectToServer`] or do something else.
Message to receive.
Once this method is called, messages are not automatically forwarded to the server or to the page - you should do that manually by calling [`method: WebSocketRoute.send`]. See examples at the top for more details.
Calling this method again will override the handler with a new one.
## async method: WebSocketRoute.routeReceive
* since: v1.48
This method allows to route messages that are received by the [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) object in the page from the server. This method only makes sense if you are also calling [`method: WebSocketRoute.connect`].
Once this method is called, received messages are not automatically dispatched to the [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) object in the page - you should do that manually by calling [`method: WebSocketRoute.receive`].
Calling this method again times will override the handler with a new one.
### param: WebSocketRoute.routeReceive.handler
### param: WebSocketRoute.onMessage.handler
* since: v1.48
* langs: js, python
- `handler` <[function]\([string]\): [Promise<any>|any]>
Handler function to route received messages.
Function that will handle messages.
### param: WebSocketRoute.routeReceive.handler
### param: WebSocketRoute.onMessage.handler
* since: v1.48
* langs: csharp, java
- `handler` <[function]\([WebSocketFrame]\)>
Handler function to route received messages.
Function that will handle messages.
## async method: WebSocketRoute.routeSend
* since: v1.48
This method allows to route messages that are sent by `WebSocket.send()` call in the page, instead of actually sending them to the server. Once this method is called, sent messages **are not** automatically forwarded to the server - you should do that manually by calling [`method: WebSocketRoute.send`].
Calling this method again times will override the handler with a new one.
### param: WebSocketRoute.routeSend.handler
* since: v1.48
* langs: js, python
- `handler` <[function]\([string]|[Buffer]\): [Promise<any>|any]>
Handler function to route sent messages.
### param: WebSocketRoute.routeSend.handler
* since: v1.48
* langs: csharp, java
- `handler` <[function]\([WebSocketFrame]\)>
Handler function to route sent messages.
## method: WebSocketRoute.send
* since: v1.48
Sends a message to the server, like it was sent in the page with `WebSocket.send()`.
Sends a message to the WebSocket. When called on the original WebSocket, sends the message to the page. When called on the result of [`method: WebSocketRoute.connectToServer`], sends the message to the server. See examples at the top for more details.
### param: WebSocketRoute.send.message
* since: v1.48
@ -159,6 +300,7 @@ Sends a message to the server, like it was sent in the page with `WebSocket.send
Message to send.
## method: WebSocketRoute.url
* since: v1.48
- returns: <[string]>

View File

@ -228,7 +228,7 @@ export class BrowserContext extends ChannelOwner<channels.BrowserContextChannel>
if (routeHandler)
await routeHandler.handle(webSocketRoute);
else
await webSocketRoute.connect();
webSocketRoute.connectToServer();
}
async _onBinding(bindingCall: BindingCall) {

View File

@ -85,10 +85,6 @@ export const Events = {
FrameSent: 'framesent',
},
WebSocketRoute: {
Close: 'close',
},
Worker: {
Close: 'close',
},

View File

@ -453,28 +453,76 @@ export class WebSocketRoute extends ChannelOwner<channels.WebSocketRouteChannel>
return (route as any)._object;
}
private _routeSendHandler?: (message: string | Buffer) => any;
private _routeReceiveHandler?: (message: string | Buffer) => any;
private _onPageMessage?: (message: string | Buffer) => any;
private _onPageClose?: (code: number | undefined, reason: string | undefined) => any;
private _onServerMessage?: (message: string | Buffer) => any;
private _onServerClose?: (code: number | undefined, reason: string | undefined) => any;
private _server: api.WebSocketRoute;
private _connected = false;
constructor(parent: ChannelOwner, type: string, guid: string, initializer: channels.WebSocketRouteInitializer) {
super(parent, type, guid, initializer);
this._channel.on('messageFromPage', ({ message, isBase64 }) => {
if (this._routeSendHandler)
this._routeSendHandler(isBase64 ? Buffer.from(message, 'base64') : message);
this._server = {
onMessage: (handler: (message: string | Buffer) => any) => {
this._onServerMessage = handler;
},
onClose: (handler: (code: number | undefined, reason: string | undefined) => any) => {
this._onServerClose = handler;
},
connectToServer: () => {
throw new Error(`connectToServer must be called on the page-side WebSocketRoute`);
},
url: () => {
return this._initializer.url;
},
close: async (options: { code?: number, reason?: string } = {}) => {
await this._channel.closeServer({ ...options, wasClean: true }).catch(() => {});
},
send: (message: string | Buffer) => {
if (isString(message))
this._channel.sendToServer({ message, isBase64: false }).catch(() => {});
else
this._channel.sendToServer({ message: message.toString('base64'), isBase64: true }).catch(() => {});
},
async [Symbol.asyncDispose]() {
await this.close();
},
};
this._channel.on('messageFromPage', ({ message, isBase64 }) => {
if (this._onPageMessage)
this._onPageMessage(isBase64 ? Buffer.from(message, 'base64') : message);
else if (this._connected)
this._channel.sendToServer({ message, isBase64 }).catch(() => {});
});
this._channel.on('messageFromServer', ({ message, isBase64 }) => {
if (this._routeReceiveHandler)
this._routeReceiveHandler(isBase64 ? Buffer.from(message, 'base64') : message);
if (this._onServerMessage)
this._onServerMessage(isBase64 ? Buffer.from(message, 'base64') : message);
else
this._channel.sendToPage({ message, isBase64 }).catch(() => {});
});
this._channel.on('close', () => this.emit(Events.WebSocketRoute.Close));
this._channel.on('closePage', ({ code, reason, wasClean }) => {
if (this._onPageClose)
this._onPageClose(code, reason);
else
this._channel.closeServer({ code, reason, wasClean }).catch(() => {});
});
this._channel.on('closeServer', ({ code, reason, wasClean }) => {
if (this._onServerClose)
this._onServerClose(code, reason);
else
this._channel.closePage({ code, reason, wasClean }).catch(() => {});
});
}
url() {
@ -482,40 +530,30 @@ export class WebSocketRoute extends ChannelOwner<channels.WebSocketRouteChannel>
}
async close(options: { code?: number, reason?: string } = {}) {
try {
await this._channel.close(options);
} catch (e) {
if (isTargetClosedError(e))
return;
throw e;
}
await this._channel.closePage({ ...options, wasClean: true }).catch(() => {});
}
async connect() {
connectToServer() {
if (this._connected)
throw new Error('Already connected to the server');
this._connected = true;
await this._channel.connect();
this._channel.connect().catch(() => {});
return this._server;
}
send(message: string | Buffer) {
if (isString(message))
this._channel.sendToServer({ message, isBase64: false }).catch(() => {});
else
this._channel.sendToServer({ message: message.toString('base64'), isBase64: true }).catch(() => {});
}
receive(message: string | Buffer) {
if (isString(message))
this._channel.sendToPage({ message, isBase64: false }).catch(() => {});
else
this._channel.sendToPage({ message: message.toString('base64'), isBase64: true }).catch(() => {});
}
routeSend(handler: (message: string | Buffer) => any) {
this._routeSendHandler = handler;
onMessage(handler: (message: string | Buffer) => any) {
this._onPageMessage = handler;
}
routeReceive(handler: (message: string | Buffer) => any) {
this._routeReceiveHandler = handler;
onClose(handler: (code: number | undefined, reason: string | undefined) => any) {
this._onPageClose = handler;
}
async [Symbol.asyncDispose]() {
@ -525,10 +563,7 @@ export class WebSocketRoute extends ChannelOwner<channels.WebSocketRouteChannel>
async _afterHandle() {
if (this._connected)
return;
if (this._routeReceiveHandler)
throw new Error(`WebSocketRoute.routeReceive() call had no effect. Make sure to call WebSocketRoute.connect() as well.`);
// Ensure that websocket is "open", so that test can send messages to it
// without an actual server connection.
// Ensure that websocket is "open" and can send messages without an actual server connection.
await this._channel.ensureOpened();
}
}

View File

@ -2139,7 +2139,16 @@ scheme.WebSocketRouteMessageFromServerEvent = tObject({
message: tString,
isBase64: tBoolean,
});
scheme.WebSocketRouteCloseEvent = tOptional(tObject({}));
scheme.WebSocketRouteClosePageEvent = tObject({
code: tOptional(tNumber),
reason: tOptional(tString),
wasClean: tBoolean,
});
scheme.WebSocketRouteCloseServerEvent = tObject({
code: tOptional(tNumber),
reason: tOptional(tString),
wasClean: tBoolean,
});
scheme.WebSocketRouteConnectParams = tOptional(tObject({}));
scheme.WebSocketRouteConnectResult = tOptional(tObject({}));
scheme.WebSocketRouteEnsureOpenedParams = tOptional(tObject({}));
@ -2154,11 +2163,18 @@ scheme.WebSocketRouteSendToServerParams = tObject({
isBase64: tBoolean,
});
scheme.WebSocketRouteSendToServerResult = tOptional(tObject({}));
scheme.WebSocketRouteCloseParams = tObject({
scheme.WebSocketRouteClosePageParams = tObject({
code: tOptional(tNumber),
reason: tOptional(tString),
wasClean: tBoolean,
});
scheme.WebSocketRouteCloseResult = tOptional(tObject({}));
scheme.WebSocketRouteClosePageResult = tOptional(tObject({}));
scheme.WebSocketRouteCloseServerParams = tObject({
code: tOptional(tNumber),
reason: tOptional(tString),
wasClean: tBoolean,
});
scheme.WebSocketRouteCloseServerResult = tOptional(tObject({}));
scheme.ResourceTiming = tObject({
startTime: tNumber,
domainLookupStart: tNumber,

View File

@ -44,14 +44,14 @@ export class WebSocketRouteDispatcher extends Dispatcher<{ guid: string }, chann
// from the mock websocket, so pretend like it was closed.
eventsHelper.addEventListener(frame._page, Page.Events.InternalFrameNavigatedToNewDocument, (frame: Frame) => {
if (frame === this._frame)
this._onClose();
this._executionContextGone();
}),
eventsHelper.addEventListener(frame._page, Page.Events.FrameDetached, (frame: Frame) => {
if (frame === this._frame)
this._onClose();
this._executionContextGone();
}),
eventsHelper.addEventListener(frame._page, Page.Events.Close, () => this._onClose()),
eventsHelper.addEventListener(frame._page, Page.Events.Crash, () => this._onClose()),
eventsHelper.addEventListener(frame._page, Page.Events.Close, () => this._executionContextGone()),
eventsHelper.addEventListener(frame._page, Page.Events.Crash, () => this._executionContextGone()),
);
WebSocketRouteDispatcher._idToDispatcher.set(this._id, this);
(scope as any)._dispatchEvent('webSocketRoute', { webSocketRoute: this });
@ -84,8 +84,10 @@ export class WebSocketRouteDispatcher extends Dispatcher<{ guid: string }, chann
dispatcher?._dispatchEvent('messageFromPage', { message: payload.data.data, isBase64: payload.data.isBase64 });
if (payload.type === 'onMessageFromServer')
dispatcher?._dispatchEvent('messageFromServer', { message: payload.data.data, isBase64: payload.data.isBase64 });
if (payload.type === 'onClose')
dispatcher?._onClose();
if (payload.type === 'onClosePage')
dispatcher?._dispatchEvent('closePage', { code: payload.code, reason: payload.reason, wasClean: payload.wasClean });
if (payload.type === 'onCloseServer')
dispatcher?._dispatchEvent('closeServer', { code: payload.code, reason: payload.reason, wasClean: payload.wasClean });
});
}
@ -117,8 +119,12 @@ export class WebSocketRouteDispatcher extends Dispatcher<{ guid: string }, chann
await this._evaluateAPIRequest({ id: this._id, type: 'sendToServer', data: { data: params.message, isBase64: params.isBase64 } });
}
async close(params: channels.WebSocketRouteCloseParams) {
await this._evaluateAPIRequest({ id: this._id, type: 'close', code: params.code, reason: params.reason, wasClean: true });
async closePage(params: channels.WebSocketRouteClosePageParams) {
await this._evaluateAPIRequest({ id: this._id, type: 'closePage', code: params.code, reason: params.reason, wasClean: params.wasClean });
}
async closeServer(params: channels.WebSocketRouteCloseServerParams) {
await this._evaluateAPIRequest({ id: this._id, type: 'closeServer', code: params.code, reason: params.reason, wasClean: params.wasClean });
}
private async _evaluateAPIRequest(request: ws.APIRequest) {
@ -129,14 +135,14 @@ export class WebSocketRouteDispatcher extends Dispatcher<{ guid: string }, chann
WebSocketRouteDispatcher._idToDispatcher.delete(this._id);
}
_onClose() {
// We could enter here twice upon page closure:
private _executionContextGone() {
// We could enter here after being disposed upon page closure:
// - first from the recursive dispose inintiated by PageDispatcher;
// - then from our own page.on('close') listener.
if (this._disposed)
return;
this._dispatchEvent('close');
this._dispose();
if (!this._disposed) {
this._dispatchEvent('closePage', { wasClean: true });
this._dispatchEvent('closeServer', { wasClean: true });
}
}
}

View File

@ -19,17 +19,19 @@ export type WSData = { data: string, isBase64: boolean };
export type OnCreatePayload = { type: 'onCreate', id: string, url: string };
export type OnMessageFromPagePayload = { type: 'onMessageFromPage', id: string, data: WSData };
export type OnClosePayload = { type: 'onClose', id: string, code: number | undefined, reason: string | undefined, wasClean: boolean };
export type OnClosePagePayload = { type: 'onClosePage', id: string, code: number | undefined, reason: string | undefined, wasClean: boolean };
export type OnMessageFromServerPayload = { type: 'onMessageFromServer', id: string, data: WSData };
export type BindingPayload = OnCreatePayload | OnMessageFromPagePayload | OnMessageFromServerPayload | OnClosePayload;
export type OnCloseServerPayload = { type: 'onCloseServer', id: string, code: number | undefined, reason: string | undefined, wasClean: boolean };
export type BindingPayload = OnCreatePayload | OnMessageFromPagePayload | OnMessageFromServerPayload | OnClosePagePayload | OnCloseServerPayload;
export type ConnectRequest = { type: 'connect', id: string };
export type PassthroughRequest = { type: 'passthrough', id: string };
export type EnsureOpenedRequest = { type: 'ensureOpened', id: string };
export type SendToPageRequest = { type: 'sendToPage', id: string, data: WSData };
export type SendToServerRequest = { type: 'sendToServer', id: string, data: WSData };
export type CloseRequest = { type: 'close', id: string, code: number | undefined, reason: string | undefined, wasClean: boolean };
export type APIRequest = ConnectRequest | PassthroughRequest | EnsureOpenedRequest | SendToPageRequest | SendToServerRequest | CloseRequest;
export type ClosePageRequest = { type: 'closePage', id: string, code: number | undefined, reason: string | undefined, wasClean: boolean };
export type CloseServerRequest = { type: 'closeServer', id: string, code: number | undefined, reason: string | undefined, wasClean: boolean };
export type APIRequest = ConnectRequest | PassthroughRequest | EnsureOpenedRequest | SendToPageRequest | SendToServerRequest | ClosePageRequest | CloseServerRequest;
// eslint-disable-next-line no-restricted-globals
type GlobalThis = typeof globalThis;
@ -98,10 +100,12 @@ export function inject(globalThis: GlobalThis) {
ws._apiEnsureOpened();
if (request.type === 'sendToPage')
ws._apiSendToPage(dataToMessage(request.data, ws.binaryType));
if (request.type === 'close')
ws._apiClose(request.code, request.reason, request.wasClean);
if (request.type === 'closePage')
ws._apiClosePage(request.code, request.reason, request.wasClean);
if (request.type === 'sendToServer')
ws._apiSendToServer(dataToMessage(request.data, ws.binaryType));
if (request.type === 'closeServer')
ws._apiCloseServer(request.code, request.reason, request.wasClean);
};
class WebSocketMock extends EventTarget {
@ -214,21 +218,23 @@ export function inject(globalThis: GlobalThis) {
throw new DOMException(`Failed to execute 'send' on 'WebSocket': Still in CONNECTING state.`);
if (this.readyState !== WebSocketMock.OPEN)
throw new DOMException(`WebSocket is already in CLOSING or CLOSED state.`);
if (this._passthrough)
if (this._passthrough) {
if (this._ws)
this._apiSendToServer(message);
else
} else {
messageToData(message, data => binding({ type: 'onMessageFromPage', id: this._id, data }));
}
}
close(code?: number, reason?: string): void {
if (code !== undefined && code !== 1000 && (code < 3000 || code > 4999))
throw new DOMException(`Failed to execute 'close' on 'WebSocket': The close code must be either 1000, or between 3000 and 4999. ${code} is neither.`);
if (this.readyState === WebSocketMock.OPEN || this.readyState === WebSocketMock.CONNECTING)
this.readyState = WebSocketMock.CLOSING;
if (this._ws)
this._ws.close(code, reason);
if (this._passthrough)
this._apiCloseServer(code, reason, true);
else
this._onWSClose(code, reason, true);
binding({ type: 'onClosePage', id: this._id, code, reason, wasClean: true });
}
// --- methods called from the routing API ---
@ -297,16 +303,26 @@ export function inject(globalThis: GlobalThis) {
this._apiConnect();
}
_apiClose(code: number | undefined, reason: string | undefined, wasClean: boolean) {
if (this.readyState !== WebSocketMock.CLOSED) {
_apiCloseServer(code: number | undefined, reason: string | undefined, wasClean: boolean) {
if (!this._ws) {
// Short-curcuit when there is no server.
this._onWSClose(code, reason, wasClean);
return;
}
if (this._ws.readyState === WebSocketMock.CONNECTING || this._ws.readyState === WebSocketMock.OPEN)
this._ws.close(code, reason);
}
_apiClosePage(code: number | undefined, reason: string | undefined, wasClean: boolean) {
if (this.readyState === WebSocketMock.CLOSED)
return;
this.readyState = WebSocketMock.CLOSED;
this.dispatchEvent(new CloseEvent('close', { code, reason, wasClean, cancelable: true }));
}
// Immediately close the real WS and imitate that it has closed.
this._ws?.close(code, reason);
this._cleanupWS();
binding({ type: 'onClose', id: this._id, code, reason, wasClean });
idToWebSocket.delete(this._id);
this._maybeCleanup();
if (this._passthrough)
this._apiCloseServer(code, reason, wasClean);
else
binding({ type: 'onClosePage', id: this._id, code, reason, wasClean });
}
// --- internals ---
@ -319,18 +335,11 @@ export function inject(globalThis: GlobalThis) {
}
private _onWSClose(code: number | undefined, reason: string | undefined, wasClean: boolean) {
this._cleanupWS();
if (this.readyState !== WebSocketMock.CLOSED) {
this.readyState = WebSocketMock.CLOSED;
this.dispatchEvent(new CloseEvent('close', { code, reason, wasClean, cancelable: true }));
}
binding({ type: 'onClose', id: this._id, code, reason, wasClean });
idToWebSocket.delete(this._id);
}
private _cleanupWS() {
if (!this._ws)
return;
if (this._passthrough)
this._apiClosePage(code, reason, wasClean);
else
binding({ type: 'onCloseServer', id: this._id, code, reason, wasClean });
if (this._ws) {
this._ws.onopen = null;
this._ws.onclose = null;
this._ws.onmessage = null;
@ -338,6 +347,13 @@ export function inject(globalThis: GlobalThis) {
this._ws = undefined;
this._wsBufferedMessages = [];
}
this._maybeCleanup();
}
private _maybeCleanup() {
if (this.readyState === WebSocketMock.CLOSED && !this._ws)
idToWebSocket.delete(this._id);
}
}
globalThis.WebSocket = class WebSocket extends WebSocketMock {};
}

View File

@ -4044,17 +4044,15 @@ export interface Page {
*
* **Usage**
*
* Below is an example of a simple handler that blocks some websocket messages. See
* Below is an example of a simple mock that responds to a single message. See
* [WebSocketRoute](https://playwright.dev/docs/api/class-websocketroute) for more details and examples.
*
* ```js
* await page.routeWebSocket('/ws', async ws => {
* ws.routeSend(message => {
* if (message === 'to-be-blocked')
* return;
* ws.send(message);
* await page.routeWebSocket('/ws', ws => {
* ws.onMessage(message => {
* if (message === 'request')
* ws.send('response');
* });
* await ws.connect();
* });
* ```
*
@ -15349,16 +15347,77 @@ export interface CDPSession {
* Whenever a [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) route is set up with
* [page.routeWebSocket(url, handler)](https://playwright.dev/docs/api/class-page#page-route-web-socket) or
* [browserContext.routeWebSocket(url, handler)](https://playwright.dev/docs/api/class-browsercontext#browser-context-route-web-socket),
* the `WebSocketRoute` object allows to handle the WebSocket.
* the `WebSocketRoute` object allows to handle the WebSocket, like an actual server would do.
*
* By default, the routed WebSocket will not actually connect to the server. This way, you can mock entire
* communcation over the WebSocket. Here is an example that responds to a `"query"` with a `"result"`.
* **Mocking**
*
* By default, the routed WebSocket will not connect to the server. This way, you can mock entire communcation over
* the WebSocket. Here is an example that responds to a `"request"` with a `"response"`.
*
* ```js
* await page.routeWebSocket('/ws', async ws => {
* ws.routeSend(message => {
* if (message === 'query')
* ws.receive('result');
* await page.routeWebSocket('/ws', ws => {
* ws.onMessage(message => {
* if (message === 'request')
* ws.send('response');
* });
* });
* ```
*
* Since we do not call
* [webSocketRoute.connectToServer()](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-connect-to-server)
* inside the WebSocket route handler, Playwright assumes that WebSocket will be mocked, and opens the WebSocket
* inside the page automatically.
*
* **Intercepting**
*
* Alternatively, you may want to connect to the actual server, but intercept messages in-between and modify or block
* them. Calling
* [webSocketRoute.connectToServer()](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-connect-to-server)
* returns a server-side `WebSocketRoute` instance that you can send messages to, or handle incoming messages.
*
* Below is an example that modifies some messages sent by the page to the server. Messages sent from the server to
* the page are left intact, relying on the default forwarding.
*
* ```js
* await page.routeWebSocket('/ws', ws => {
* const server = ws.connectToServer();
* ws.onMessage(message => {
* if (message === 'request')
* server.send('request2');
* else
* server.send(message);
* });
* });
* ```
*
* After connecting to the server, all **messages are forwarded** between the page and the server by default.
*
* However, if you call
* [webSocketRoute.onMessage(handler)](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-on-message)
* on the original route, messages from the page to the server **will not be forwarded** anymore, but should instead
* be handled by the
* [`handler`](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-on-message-option-handler).
*
* Similarly, calling
* [webSocketRoute.onMessage(handler)](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-on-message)
* on the server-side WebSocket will **stop forwarding messages** from the server to the page, and
* [`handler`](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-on-message-option-handler) should
* take care of them.
*
* The following example blocks some messages in both directions. Since it calls
* [webSocketRoute.onMessage(handler)](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-on-message)
* in both directions, there is no automatic forwarding at all.
*
* ```js
* await page.routeWebSocket('/ws', ws => {
* const server = ws.connectToServer();
* ws.onMessage(message => {
* if (message !== 'blocked-from-the-page')
* server.send(message);
* });
* server.onMessage(message => {
* if (message !== 'blocked-from-the-server')
* ws.send(message);
* });
* });
* ```
@ -15366,63 +15425,38 @@ export interface CDPSession {
*/
export interface WebSocketRoute {
/**
* This method allows to route messages that are sent by `WebSocket.send()` call in the page, instead of actually
* sending them to the server. Once this method is called, sent messages **are not** automatically forwarded to the
* server - you should do that manually by calling
* [webSocketRoute.send(message)](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-send).
* This method allows to handle messages that are sent by the WebSocket, either from the page or from the server.
*
* Calling this method again times will override the handler with a new one.
* @param handler Handler function to route sent messages.
*/
routeSend(handler: (message: string | Buffer) => any): void;
/**
* This method allows to route messages that are received by the
* [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) object in the page from the server. This
* method only makes sense if you are also calling
* [webSocketRoute.connect()](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-connect).
* When called on the original WebSocket route, this method handles messages sent from the page. You can handle this
* messages by responding to them with
* [webSocketRoute.send(message)](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-send),
* forwarding them to the server-side connection returned by
* [webSocketRoute.connectToServer()](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-connect-to-server)
* or do something else.
*
* Once this method is called, received messages are not automatically dispatched to the
* [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) object in the page - you should do that
* manually by calling
* [webSocketRoute.receive(message)](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-receive).
* Once this method is called, messages are not automatically forwarded to the server or to the page - you should do
* that manually by calling
* [webSocketRoute.send(message)](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-send). See
* examples at the top for more details.
*
* Calling this method again times will override the handler with a new one.
* @param handler Handler function to route received messages.
* Calling this method again will override the handler with a new one.
* @param handler Function that will handle messages.
*/
routeReceive(handler: (message: string | Buffer) => any): void;
onMessage(handler: (message: string | Buffer) => any): void;
/**
* Emitted when the [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) closes.
* Allows to handle [`WebSocket.close`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close).
*
* By default, closing one side of the connection, either in the page or on the server, will close the other side.
* However, when
* [webSocketRoute.onClose(handler)](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-on-close)
* handler is set up, the default forwarding of closure is disabled, and handler should take care of it.
* @param handler Function that will handle WebSocket closure. Received an optional
* [close code](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close#code) and an optional
* [close reason](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close#reason).
*/
on(event: 'close', listener: () => any): this;
onClose(handler: (code: number | undefined, reason: string | undefined) => any): void;
/**
* Adds an event listener that will be automatically removed after it is triggered once. See `addListener` for more information about this event.
*/
once(event: 'close', listener: () => any): this;
/**
* Emitted when the [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) closes.
*/
addListener(event: 'close', listener: () => any): this;
/**
* Removes an event listener added by `on` or `addListener`.
*/
removeListener(event: 'close', listener: () => any): this;
/**
* Removes an event listener added by `on` or `addListener`.
*/
off(event: 'close', listener: () => any): this;
/**
* Emitted when the [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) closes.
*/
prependListener(event: 'close', listener: () => any): this;
/**
* Closes the server connection and the [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket)
* object in the page.
* Closes one side of the WebSocket connection.
* @param options
*/
close(options?: {
@ -15439,28 +15473,28 @@ export interface WebSocketRoute {
/**
* By default, routed WebSocket does not connect to the server, so you can mock entire WebSocket communication. This
* method connects to the actual WebSocket server, giving the ability to send and receive messages from the server.
* method connects to the actual WebSocket server, and returns the server-side
* [WebSocketRoute](https://playwright.dev/docs/api/class-websocketroute) instance, giving the ability to send and
* receive messages from the server.
*
* Once connected:
* - Messages received from the server will be automatically dispatched to the
* [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) object in the page, unless
* [webSocketRoute.routeReceive(handler)](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-route-receive)
* is called.
* - Messages sent by the `WebSocket.send()` call in the page will be automatically sent to the server, unless
* [webSocketRoute.routeSend(handler)](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-route-send)
* is called.
* Once connected to the server:
* - Messages received from the server will be **automatically forwarded** to the WebSocket in the page, unless
* [webSocketRoute.onMessage(handler)](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-on-message)
* is called on the server-side `WebSocketRoute`.
* - Messages sent by the [`WebSocket.send()`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/send) call
* in the page will be **automatically forwarded** to the server, unless
* [webSocketRoute.onMessage(handler)](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-on-message)
* is called on the original `WebSocketRoute`.
*
* See examples at the top for more details.
*/
connect(): Promise<void>;
connectToServer(): WebSocketRoute;
/**
* Dispatches a message to the [`WebSocket`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) object in the
* page, like it was received from the server.
* @param message Message to receive.
*/
receive(message: string|Buffer): void;
/**
* Sends a message to the server, like it was sent in the page with `WebSocket.send()`.
* Sends a message to the WebSocket. When called on the original WebSocket, sends the message to the page. When called
* on the result of
* [webSocketRoute.connectToServer()](https://playwright.dev/docs/api/class-websocketroute#web-socket-route-connect-to-server),
* sends the message to the server. See examples at the top for more details.
* @param message Message to send.
*/
send(message: string|Buffer): void;

View File

@ -3810,7 +3810,8 @@ export type WebSocketRouteInitializer = {
export interface WebSocketRouteEventTarget {
on(event: 'messageFromPage', callback: (params: WebSocketRouteMessageFromPageEvent) => void): this;
on(event: 'messageFromServer', callback: (params: WebSocketRouteMessageFromServerEvent) => void): this;
on(event: 'close', callback: (params: WebSocketRouteCloseEvent) => void): this;
on(event: 'closePage', callback: (params: WebSocketRouteClosePageEvent) => void): this;
on(event: 'closeServer', callback: (params: WebSocketRouteCloseServerEvent) => void): this;
}
export interface WebSocketRouteChannel extends WebSocketRouteEventTarget, Channel {
_type_WebSocketRoute: boolean;
@ -3818,7 +3819,8 @@ export interface WebSocketRouteChannel extends WebSocketRouteEventTarget, Channe
ensureOpened(params?: WebSocketRouteEnsureOpenedParams, metadata?: CallMetadata): Promise<WebSocketRouteEnsureOpenedResult>;
sendToPage(params: WebSocketRouteSendToPageParams, metadata?: CallMetadata): Promise<WebSocketRouteSendToPageResult>;
sendToServer(params: WebSocketRouteSendToServerParams, metadata?: CallMetadata): Promise<WebSocketRouteSendToServerResult>;
close(params: WebSocketRouteCloseParams, metadata?: CallMetadata): Promise<WebSocketRouteCloseResult>;
closePage(params: WebSocketRouteClosePageParams, metadata?: CallMetadata): Promise<WebSocketRouteClosePageResult>;
closeServer(params: WebSocketRouteCloseServerParams, metadata?: CallMetadata): Promise<WebSocketRouteCloseServerResult>;
}
export type WebSocketRouteMessageFromPageEvent = {
message: string,
@ -3828,7 +3830,16 @@ export type WebSocketRouteMessageFromServerEvent = {
message: string,
isBase64: boolean,
};
export type WebSocketRouteCloseEvent = {};
export type WebSocketRouteClosePageEvent = {
code?: number,
reason?: string,
wasClean: boolean,
};
export type WebSocketRouteCloseServerEvent = {
code?: number,
reason?: string,
wasClean: boolean,
};
export type WebSocketRouteConnectParams = {};
export type WebSocketRouteConnectOptions = {};
export type WebSocketRouteConnectResult = void;
@ -3851,20 +3862,32 @@ export type WebSocketRouteSendToServerOptions = {
};
export type WebSocketRouteSendToServerResult = void;
export type WebSocketRouteCloseParams = {
export type WebSocketRouteClosePageParams = {
code?: number,
reason?: string,
wasClean: boolean,
};
export type WebSocketRouteClosePageOptions = {
code?: number,
reason?: string,
};
export type WebSocketRouteCloseOptions = {
export type WebSocketRouteClosePageResult = void;
export type WebSocketRouteCloseServerParams = {
code?: number,
reason?: string,
wasClean: boolean,
};
export type WebSocketRouteCloseServerOptions = {
code?: number,
reason?: string,
};
export type WebSocketRouteCloseResult = void;
export type WebSocketRouteCloseServerResult = void;
export interface WebSocketRouteEvents {
'messageFromPage': WebSocketRouteMessageFromPageEvent;
'messageFromServer': WebSocketRouteMessageFromServerEvent;
'close': WebSocketRouteCloseEvent;
'closePage': WebSocketRouteClosePageEvent;
'closeServer': WebSocketRouteCloseServerEvent;
}
export type ResourceTiming = {

View File

@ -2987,10 +2987,17 @@ WebSocketRoute:
message: string
isBase64: boolean
close:
closePage:
parameters:
code: number?
reason: string?
wasClean: boolean
closeServer:
parameters:
code: number?
reason: string?
wasClean: boolean
events:
@ -3004,7 +3011,17 @@ WebSocketRoute:
message: string
isBase64: boolean
close:
closePage:
parameters:
code: number?
reason: string?
wasClean: boolean
closeServer:
parameters:
code: number?
reason: string?
wasClean: boolean
ResourceTiming:

View File

@ -62,10 +62,10 @@ for (const mock of ['no-mock', 'no-match', 'pass-through']) {
if (mock === 'no-match') {
await page.routeWebSocket(/zzz/, () => {});
} else if (mock === 'pass-through') {
await page.routeWebSocket(/.*/, async ws => {
ws.routeSend(message => ws.send(message));
ws.routeReceive(message => ws.receive(message));
await ws.connect();
await page.routeWebSocket(/.*/, ws => {
const server = ws.connectToServer();
ws.onMessage(message => server.send(message));
server.onMessage(message => ws.send(message));
});
}
});
@ -196,9 +196,9 @@ for (const mock of ['no-mock', 'no-match', 'pass-through']) {
test('should work with ws.close', async ({ page, server }) => {
const { promise, resolve } = withResolvers<WebSocketRoute>();
await page.routeWebSocket(/.*/, async route => {
await route.connect();
resolve(route);
await page.routeWebSocket(/.*/, async ws => {
ws.connectToServer();
resolve(ws);
});
const wsPromise = server.waitForWebSocket();
@ -206,7 +206,7 @@ test('should work with ws.close', async ({ page, server }) => {
const ws = await wsPromise;
const route = await promise;
route.receive('hello');
route.send('hello');
await expect.poll(() => page.evaluate(() => window.log)).toEqual([
'open',
`message: data=hello origin=ws://localhost:${server.PORT} lastEventId=`,
@ -224,12 +224,12 @@ test('should work with ws.close', async ({ page, server }) => {
test('should pattern match', async ({ page, server }) => {
await page.routeWebSocket(/.*\/ws$/, async ws => {
await ws.connect();
ws.connectToServer();
});
await page.routeWebSocket('**/mock-ws', ws => {
ws.routeSend(message => {
ws.receive('mock-response');
ws.onMessage(message => {
ws.send('mock-response');
});
});
@ -260,33 +260,33 @@ test('should pattern match', async ({ page, server }) => {
test('should work with server', async ({ page, server }) => {
const { promise, resolve } = withResolvers<WebSocketRoute>();
await page.routeWebSocket(/.*/, async route => {
route.routeSend(message => {
await page.routeWebSocket(/.*/, async ws => {
const server = ws.connectToServer();
ws.onMessage(message => {
switch (message) {
case 'to-respond':
route.receive('response');
ws.send('response');
return;
case 'to-block':
return;
case 'to-modify':
route.send('modified');
server.send('modified');
return;
}
route.send(message);
server.send(message);
});
route.routeReceive(message => {
server.onMessage(message => {
switch (message) {
case 'to-block':
return;
case 'to-modify':
route.receive('modified');
ws.send('modified');
return;
}
route.receive(message);
ws.send(message);
});
await route.connect();
route.send('fake');
resolve(route);
server.send('fake');
resolve(ws);
});
const wsPromise = server.waitForWebSocket();
@ -324,7 +324,7 @@ test('should work with server', async ({ page, server }) => {
]);
const route = await promise;
route.receive('another');
route.send('another');
await expect.poll(() => page.evaluate(() => window.log)).toEqual([
'open',
`message: data=modified origin=ws://localhost:${server.PORT} lastEventId=`,
@ -346,15 +346,15 @@ test('should work with server', async ({ page, server }) => {
test('should work without server', async ({ page, server }) => {
const { promise, resolve } = withResolvers<WebSocketRoute>();
await page.routeWebSocket(/.*/, route => {
route.routeSend(message => {
await page.routeWebSocket(/.*/, ws => {
ws.onMessage(message => {
switch (message) {
case 'to-respond':
route.receive('response');
ws.send('response');
return;
}
});
resolve(route);
resolve(ws);
});
await setupWS(page, server.PORT, 'blob');
@ -373,7 +373,7 @@ test('should work without server', async ({ page, server }) => {
]);
const route = await promise;
route.receive('another');
route.send('another');
await route.close({ code: 3008, reason: 'oops' });
await expect.poll(() => page.evaluate(() => window.log)).toEqual([
@ -387,68 +387,68 @@ test('should work without server', async ({ page, server }) => {
test('should emit close upon frame navigation', async ({ page, server }) => {
const { promise, resolve } = withResolvers<WebSocketRoute>();
await page.routeWebSocket(/.*/, async route => {
await route.connect();
resolve(route);
await page.routeWebSocket(/.*/, async ws => {
ws.connectToServer();
resolve(ws);
});
await setupWS(page, server.PORT, 'blob');
const route = await promise;
route.receive('hello');
route.send('hello');
await expect.poll(() => page.evaluate(() => window.log)).toEqual([
'open',
`message: data=hello origin=ws://localhost:${server.PORT} lastEventId=`,
]);
const closedPromise = new Promise<void>(f => route.addListener('close', f));
const closedPromise = new Promise<void>(f => route.onClose(() => f()));
await page.goto(server.EMPTY_PAGE);
await closedPromise;
});
test('should emit close upon frame detach', async ({ page, server }) => {
const { promise, resolve } = withResolvers<WebSocketRoute>();
await page.routeWebSocket(/.*/, async route => {
await route.connect();
resolve(route);
await page.routeWebSocket(/.*/, async ws => {
ws.connectToServer();
resolve(ws);
});
const frame = await attachFrame(page, 'frame1', server.EMPTY_PAGE);
await setupWS(frame, server.PORT, 'blob');
const route = await promise;
route.receive('hello');
route.send('hello');
await expect.poll(() => frame.evaluate(() => window.log)).toEqual([
'open',
`message: data=hello origin=ws://localhost:${server.PORT} lastEventId=`,
]);
const closedPromise = new Promise<void>(f => route.addListener('close', f));
const closedPromise = new Promise<void>(f => route.onClose(() => f()));
await detachFrame(page, 'frame1');
await closedPromise;
});
test('should route on context', async ({ page, server }) => {
await page.routeWebSocket(/ws1/, ws => {
ws.routeSend(message => {
ws.receive('page-mock-1');
ws.onMessage(message => {
ws.send('page-mock-1');
});
});
await page.routeWebSocket(/ws1/, ws => {
ws.routeSend(message => {
ws.receive('page-mock-2');
ws.onMessage(message => {
ws.send('page-mock-2');
});
});
await page.context().routeWebSocket(/.*/, ws => {
ws.routeSend(message => {
ws.receive('context-mock-1');
ws.onMessage(message => {
ws.send('context-mock-1');
});
ws.routeSend(message => {
ws.receive('context-mock-2');
ws.onMessage(message => {
ws.send('context-mock-2');
});
});
@ -470,9 +470,9 @@ test('should route on context', async ({ page, server }) => {
test('should not throw after page closure', async ({ page, server }) => {
const { promise, resolve } = withResolvers<WebSocketRoute>();
await page.routeWebSocket(/.*/, async route => {
await route.connect();
resolve(route);
await page.routeWebSocket(/.*/, async ws => {
ws.connectToServer();
resolve(ws);
});
await setupWS(page, server.PORT, 'blob');
@ -480,6 +480,31 @@ test('should not throw after page closure', async ({ page, server }) => {
const route = await promise;
await Promise.all([
page.close(),
route.receive('hello'),
route.send('hello'),
]);
});
test('should not throw with empty handler', async ({ page, server }) => {
await page.routeWebSocket(/.*/, () => {});
await setupWS(page, server.PORT, 'blob');
await expect.poll(() => page.evaluate(() => window.log)).toEqual(['open']);
await page.evaluate(() => window.ws.send('hi'));
await page.evaluate(() => window.ws.send('hi2'));
await expect.poll(() => page.evaluate(() => window.log)).toEqual(['open']);
});
test('should throw when connecting twice', async ({ page, server }) => {
const { promise, resolve } = withResolvers<Error>();
await page.routeWebSocket(/.*/, ws => {
ws.connectToServer();
try {
ws.connectToServer();
} catch (e) {
resolve(e);
}
});
await setupWS(page, server.PORT, 'blob');
const error = await promise;
expect(error.message).toContain('Already connected to the server');
});

View File

@ -226,8 +226,8 @@ export interface CDPSession {
}
export interface WebSocketRoute {
routeSend(handler: (message: string | Buffer) => any): void;
routeReceive(handler: (message: string | Buffer) => any): void;
onMessage(handler: (message: string | Buffer) => any): void;
onClose(handler: (code: number | undefined, reason: string | undefined) => any): void;
}
type DeviceDescriptor = {