diff --git a/packages/playwright-core/src/server/browserContext.ts b/packages/playwright-core/src/server/browserContext.ts index fc20c52bb5..8a835d3726 100644 --- a/packages/playwright-core/src/server/browserContext.ts +++ b/packages/playwright-core/src/server/browserContext.ts @@ -314,6 +314,10 @@ export abstract class BrowserContext extends SdkObject { return this.doSetHTTPCredentials(httpCredentials); } + hasBinding(name: string) { + return this._pageBindings.has(name); + } + async exposeBinding(name: string, needsHandle: boolean, playwrightBinding: frames.FunctionWithSource): Promise { if (this._pageBindings.has(name)) throw new Error(`Function "${name}" has been already registered`); @@ -414,8 +418,8 @@ export abstract class BrowserContext extends SdkObject { this._options.httpCredentials = { username, password: password || '' }; } - async addInitScript(source: string) { - const initScript = new InitScript(source); + async addInitScript(source: string, name?: string) { + const initScript = new InitScript(source, false /* internal */, name); this.initScripts.push(initScript); await this.doAddInitScript(initScript); } diff --git a/packages/playwright-core/src/server/dispatchers/browserContextDispatcher.ts b/packages/playwright-core/src/server/dispatchers/browserContextDispatcher.ts index c6ffce49f7..3579f4b3bb 100644 --- a/packages/playwright-core/src/server/dispatchers/browserContextDispatcher.ts +++ b/packages/playwright-core/src/server/dispatchers/browserContextDispatcher.ts @@ -288,7 +288,7 @@ export class BrowserContextDispatcher extends Dispatcher { this._webSocketInterceptionPatterns = params.patterns; if (params.patterns.length) - await WebSocketRouteDispatcher.installIfNeeded(this, this._context); + await WebSocketRouteDispatcher.installIfNeeded(this._context); } async storageState(params: channels.BrowserContextStorageStateParams, metadata: CallMetadata): Promise { diff --git a/packages/playwright-core/src/server/dispatchers/pageDispatcher.ts b/packages/playwright-core/src/server/dispatchers/pageDispatcher.ts index 5e3b73a73f..6e8c47929b 100644 --- a/packages/playwright-core/src/server/dispatchers/pageDispatcher.ts +++ b/packages/playwright-core/src/server/dispatchers/pageDispatcher.ts @@ -191,7 +191,7 @@ export class PageDispatcher extends Dispatcher { this._webSocketInterceptionPatterns = params.patterns; if (params.patterns.length) - await WebSocketRouteDispatcher.installIfNeeded(this.parentScope(), this._page); + await WebSocketRouteDispatcher.installIfNeeded(this._page); } async expectScreenshot(params: channels.PageExpectScreenshotParams, metadata: CallMetadata): Promise { diff --git a/packages/playwright-core/src/server/dispatchers/webSocketRouteDispatcher.ts b/packages/playwright-core/src/server/dispatchers/webSocketRouteDispatcher.ts index 87f469b7d0..bcbc89fe03 100644 --- a/packages/playwright-core/src/server/dispatchers/webSocketRouteDispatcher.ts +++ b/packages/playwright-core/src/server/dispatchers/webSocketRouteDispatcher.ts @@ -18,7 +18,7 @@ import type { BrowserContext } from '../browserContext'; import type { Frame } from '../frames'; import { Page } from '../page'; import type * as channels from '@protocol/channels'; -import { Dispatcher } from './dispatcher'; +import { Dispatcher, existingDispatcher } from './dispatcher'; import { createGuid, urlMatches } from '../../utils'; import { PageDispatcher } from './pageDispatcher'; import type { BrowserContextDispatcher } from './browserContextDispatcher'; @@ -26,9 +26,6 @@ import * as webSocketMockSource from '../../generated/webSocketMockSource'; import type * as ws from '../injected/webSocketMock'; import { eventsHelper } from '../../utils/eventsHelper'; -const kBindingInstalledSymbol = Symbol('webSocketRouteBindingInstalled'); -const kInitScriptInstalledSymbol = Symbol('webSocketRouteInitScriptInstalled'); - export class WebSocketRouteDispatcher extends Dispatcher<{ guid: string }, channels.WebSocketRouteChannel, PageDispatcher | BrowserContextDispatcher> implements channels.WebSocketRouteChannel { _type_WebSocketRoute = true; private _id: string; @@ -57,18 +54,18 @@ export class WebSocketRouteDispatcher extends Dispatcher<{ guid: string }, chann (scope as any)._dispatchEvent('webSocketRoute', { webSocketRoute: this }); } - static async installIfNeeded(contextDispatcher: BrowserContextDispatcher, target: Page | BrowserContext) { + static async installIfNeeded(target: Page | BrowserContext) { + const kBindingName = '__pwWebSocketBinding'; const context = target instanceof Page ? target.context() : target; - if (!(context as any)[kBindingInstalledSymbol]) { - (context as any)[kBindingInstalledSymbol] = true; - - await context.exposeBinding('__pwWebSocketBinding', false, (source, payload: ws.BindingPayload) => { + if (!context.hasBinding(kBindingName)) { + await context.exposeBinding(kBindingName, false, (source, payload: ws.BindingPayload) => { if (payload.type === 'onCreate') { - const pageDispatcher = PageDispatcher.fromNullable(contextDispatcher, source.page); + const contextDispatcher = existingDispatcher(context); + const pageDispatcher = contextDispatcher ? PageDispatcher.fromNullable(contextDispatcher, source.page) : undefined; let scope: PageDispatcher | BrowserContextDispatcher | undefined; if (pageDispatcher && matchesPattern(pageDispatcher, context._options.baseURL, payload.url)) scope = pageDispatcher; - else if (matchesPattern(contextDispatcher, context._options.baseURL, payload.url)) + else if (contextDispatcher && matchesPattern(contextDispatcher, context._options.baseURL, payload.url)) scope = contextDispatcher; if (scope) { new WebSocketRouteDispatcher(scope, payload.id, payload.url, source.frame); @@ -91,15 +88,15 @@ export class WebSocketRouteDispatcher extends Dispatcher<{ guid: string }, chann }); } - if (!(target as any)[kInitScriptInstalledSymbol]) { - (target as any)[kInitScriptInstalledSymbol] = true; + const kInitScriptName = 'webSocketMockSource'; + if (!target.initScripts.find(s => s.name === kInitScriptName)) { await target.addInitScript(` (() => { const module = {}; ${webSocketMockSource.source} (module.exports.inject())(globalThis); })(); - `); + `, kInitScriptName); } } diff --git a/packages/playwright-core/src/server/page.ts b/packages/playwright-core/src/server/page.ts index fe483b8347..9b85837b65 100644 --- a/packages/playwright-core/src/server/page.ts +++ b/packages/playwright-core/src/server/page.ts @@ -564,8 +564,8 @@ export class Page extends SdkObject { await this._delegate.bringToFront(); } - async addInitScript(source: string) { - const initScript = new InitScript(source); + async addInitScript(source: string, name?: string) { + const initScript = new InitScript(source, false /* internal */, name); this.initScripts.push(initScript); await this._delegate.addInitScript(initScript); } @@ -953,8 +953,9 @@ function addPageBinding(playwrightBinding: string, bindingName: string, needsHan export class InitScript { readonly source: string; readonly internal: boolean; + readonly name?: string; - constructor(source: string, internal?: boolean) { + constructor(source: string, internal?: boolean, name?: string) { const guid = createGuid(); this.source = `(() => { globalThis.__pwInitScripts = globalThis.__pwInitScripts || {}; @@ -965,6 +966,7 @@ export class InitScript { ${source} })();`; this.internal = !!internal; + this.name = name; } } diff --git a/tests/library/browsercontext-reuse.spec.ts b/tests/library/browsercontext-reuse.spec.ts index 30319f0aad..2d65d472b3 100644 --- a/tests/library/browsercontext-reuse.spec.ts +++ b/tests/library/browsercontext-reuse.spec.ts @@ -15,7 +15,7 @@ */ import { browserTest, expect } from '../config/browserTest'; -import type { BrowserContext } from '@playwright/test'; +import type { BrowserContext, Page } from '@playwright/test'; const test = browserTest.extend<{ reusedContext: () => Promise }>({ reusedContext: async ({ browserType, browser }, use) => { @@ -287,3 +287,42 @@ test('should continue issuing events after closing the reused page', async ({ re ]); } }); + +test('should work with routeWebSocket', async ({ reusedContext, server, browser }, testInfo) => { + async function setup(page: Page, suffix: string) { + await page.routeWebSocket(/ws1/, ws => { + ws.onMessage(message => { + ws.send('page-mock-' + suffix); + }); + }); + await page.context().routeWebSocket(/.*/, ws => { + ws.onMessage(message => { + ws.send('context-mock-' + suffix); + }); + }); + await page.goto('about:blank'); + await page.evaluate(({ port }) => { + window.log = []; + (window as any).ws1 = new WebSocket('ws://localhost:' + port + '/ws1'); + (window as any).ws1.addEventListener('message', event => window.log.push(`ws1:${event.data}`)); + (window as any).ws2 = new WebSocket('ws://localhost:' + port + '/ws2'); + (window as any).ws2.addEventListener('message', event => window.log.push(`ws2:${event.data}`)); + }, { port: server.PORT }); + } + + let context = await reusedContext(); + let page = await context.newPage(); + await setup(page, 'before'); + await page.evaluate(() => (window as any).ws1.send('request')); + await expect.poll(() => page.evaluate(() => window.log)).toEqual([`ws1:page-mock-before`]); + await page.evaluate(() => (window as any).ws2.send('request')); + await expect.poll(() => page.evaluate(() => window.log)).toEqual([`ws1:page-mock-before`, `ws2:context-mock-before`]); + + context = await reusedContext(); + page = context.pages()[0]; + await setup(page, 'after'); + await page.evaluate(() => (window as any).ws1.send('request')); + await expect.poll(() => page.evaluate(() => window.log)).toEqual([`ws1:page-mock-after`]); + await page.evaluate(() => (window as any).ws2.send('request')); + await expect.poll(() => page.evaluate(() => window.log)).toEqual([`ws1:page-mock-after`, `ws2:context-mock-after`]); +});