fix(routeWebSocket): should work after context reuse (#34165)

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<void> {
     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);
   }

packages/playwright-core/src/server/dispatchers/browserContextDispatcher.ts

@@ -288,7 +288,7 @@ export class BrowserContextDispatcher extends Dispatcher<BrowserContext, channel
   async setWebSocketInterceptionPatterns(params: channels.PageSetWebSocketInterceptionPatternsParams, metadata: CallMetadata): Promise<void> {
     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<channels.BrowserContextStorageStateResult> {

packages/playwright-core/src/server/dispatchers/pageDispatcher.ts

@@ -191,7 +191,7 @@ export class PageDispatcher extends Dispatcher<Page, channels.PageChannel, Brows
   async setWebSocketInterceptionPatterns(params: channels.PageSetWebSocketInterceptionPatternsParams, metadata: CallMetadata): Promise<void> {
     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<channels.PageExpectScreenshotResult> {

packages/playwright-core/src/server/dispatchers/webSocketRouteDispatcher.ts

@@ -18,17 +18,14 @@ 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';
 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<BrowserContextDispatcher>(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);
     }
   }
 

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;
   }
 }
 

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<BrowserContext> }>({
   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`]);
+});