fix(connect): make selectors.register work in connected browser

src/browserServerImpl.ts

@@ -20,13 +20,15 @@ import * as ws from 'ws';
 import { Browser } from './server/browser';
 import { ChildProcess } from 'child_process';
 import { EventEmitter } from 'ws';
-import { DispatcherScope, DispatcherConnection } from './dispatchers/dispatcher';
+import { Dispatcher, DispatcherScope, DispatcherConnection } from './dispatchers/dispatcher';
 import { BrowserDispatcher } from './dispatchers/browserDispatcher';
 import { BrowserContextDispatcher } from './dispatchers/browserContextDispatcher';
-import { BrowserNewContextParams, BrowserContextChannel } from './protocol/channels';
+import * as channels from './protocol/channels';
 import { BrowserServerLauncher, BrowserServer } from './client/browserType';
 import { envObjectToArray } from './client/clientHelper';
 import { createGuid } from './utils/utils';
+import { SelectorsDispatcher } from './dispatchers/selectorsDispatcher';
+import { Selectors } from './server/selectors';
 
 export class BrowserServerLauncherImpl implements BrowserServerLauncher {
   private _browserType: BrowserTypeBase;
@@ -105,7 +107,10 @@ export class BrowserServerImpl extends EventEmitter implements BrowserServer {
       connection.dispatch(JSON.parse(Buffer.from(message).toString()));
     });
     socket.on('error', () => {});
-    const browser = new ConnectedBrowser(connection.rootDispatcher(), this._browser);
+    const selectors = new Selectors();
+    const scope = connection.rootDispatcher();
+    const browser = new ConnectedBrowser(scope, this._browser, selectors);
+    new RemoteBrowserDispatcher(scope, browser, selectors);
     socket.on('close', () => {
       // Avoid sending any more messages over closed socket.
       connection.onmessage = () => {};
@@ -115,17 +120,30 @@ export class BrowserServerImpl extends EventEmitter implements BrowserServer {
   }
 }
 
+class RemoteBrowserDispatcher extends Dispatcher<{}, channels.RemoteBrowserInitializer> implements channels.PlaywrightChannel {
+  constructor(scope: DispatcherScope, browser: ConnectedBrowser, selectors: Selectors) {
+    super(scope, {}, 'RemoteBrowser', {
+      selectors: new SelectorsDispatcher(scope, selectors),
+      browser,
+    }, false, 'remoteBrowser');
+  }
+}
+
 class ConnectedBrowser extends BrowserDispatcher {
   private _contexts: BrowserContextDispatcher[] = [];
+  private _selectors: Selectors;
   _closed = false;
 
-  constructor(scope: DispatcherScope, browser: Browser) {
-    super(scope, browser, 'connectedBrowser');
+  constructor(scope: DispatcherScope, browser: Browser, selectors: Selectors) {
+    super(scope, browser);
+    this._selectors = selectors;
   }
 
-  async newContext(params: BrowserNewContextParams): Promise<{ context: BrowserContextChannel }> {
+  async newContext(params: channels.BrowserNewContextParams): Promise<{ context: channels.BrowserContextChannel }> {
     const result = await super.newContext(params);
-    this._contexts.push(result.context as BrowserContextDispatcher);
+    const dispatcher = result.context as BrowserContextDispatcher;
+    dispatcher._object._setSelectors(this._selectors);
+    this._contexts.push(dispatcher);
     return result;
   }
 

src/client/browserType.ts

@@ -28,6 +28,7 @@ import { ChildProcess } from 'child_process';
 import { envObjectToArray } from './clientHelper';
 import { validateHeaders } from './network';
 import { assert, makeWaitForNextTask, headersObjectToArray } from '../utils/utils';
+import { SelectorsOwner, sharedSelectors } from './selectors';
 
 export interface BrowserServerLauncher {
   launchServer(options?: LaunchServerOptions): Promise<BrowserServer>;
@@ -144,7 +145,13 @@ export class BrowserType extends ChannelOwner<channels.BrowserTypeChannel, chann
           }
         }
         ws.addEventListener('open', async () => {
-          const browser = (await connection.waitForObjectWithKnownName('connectedBrowser')) as Browser;
+          const remoteBrowser = await connection.waitForObjectWithKnownName('remoteBrowser') as RemoteBrowser;
+
+          // Inherit shared selectors for connected browser.
+          const selectorsOwner = SelectorsOwner.from(remoteBrowser._initializer.selectors);
+          sharedSelectors._addChannel(selectorsOwner);
+
+          const browser = Browser.from(remoteBrowser._initializer.browser);
           browser._logger = logger;
           browser._isRemote = true;
           const closeListener = () => {
@@ -158,6 +165,7 @@ export class BrowserType extends ChannelOwner<channels.BrowserTypeChannel, chann
           };
           ws.addEventListener('close', closeListener);
           browser.on(Events.Browser.Disconnected, () => {
+            sharedSelectors._removeChannel(selectorsOwner);
             ws.removeEventListener('close', closeListener);
             ws.close();
           });
@@ -171,3 +179,6 @@ export class BrowserType extends ChannelOwner<channels.BrowserTypeChannel, chann
     }, logger);
   }
 }
+
+export class RemoteBrowser extends ChannelOwner<channels.RemoteBrowserChannel, channels.RemoteBrowserInitializer> {
+}

src/client/connection.ts

@@ -16,7 +16,7 @@
 
 import { Browser } from './browser';
 import { BrowserContext } from './browserContext';
-import { BrowserType } from './browserType';
+import { BrowserType, RemoteBrowser } from './browserType';
 import { ChannelOwner } from './channelOwner';
 import { ElementHandle } from './elementHandle';
 import { Frame } from './frame';
@@ -34,12 +34,12 @@ import { Electron, ElectronApplication } from './electron';
 import * as channels from '../protocol/channels';
 import { ChromiumBrowser } from './chromiumBrowser';
 import { ChromiumBrowserContext } from './chromiumBrowserContext';
-import { Selectors } from './selectors';
 import { Stream } from './stream';
 import { createScheme, Validator, ValidationError } from '../protocol/validator';
 import { WebKitBrowser } from './webkitBrowser';
 import { FirefoxBrowser } from './firefoxBrowser';
 import { debugLogger } from '../utils/debugLogger';
+import { SelectorsOwner } from './selectors';
 
 class Root extends ChannelOwner<channels.Channel, {}> {
   constructor(connection: Connection) {
@@ -195,20 +195,23 @@ export class Connection {
       case 'Playwright':
         result = new Playwright(parent, type, guid, initializer);
         break;
+      case 'RemoteBrowser':
+        result = new RemoteBrowser(parent, type, guid, initializer);
+        break;
       case 'Request':
         result = new Request(parent, type, guid, initializer);
         break;
-      case 'Stream':
-        result = new Stream(parent, type, guid, initializer);
-        break;
       case 'Response':
         result = new Response(parent, type, guid, initializer);
         break;
       case 'Route':
         result = new Route(parent, type, guid, initializer);
         break;
+      case 'Stream':
+        result = new Stream(parent, type, guid, initializer);
+        break;
       case 'Selectors':
-        result = new Selectors(parent, type, guid, initializer);
+        result = new SelectorsOwner(parent, type, guid, initializer);
         break;
       case 'Worker':
         result = new Worker(parent, type, guid, initializer);

src/client/elementHandle.ts

@@ -233,6 +233,10 @@ export class ElementHandle<T extends Node = Node> extends JSHandle<T> {
       return ElementHandle.fromNullable(result.element) as ElementHandle<Element> | null;
     });
   }
+
+  async _createSelectorForTest(name: string): Promise<string | undefined> {
+    return (await this._elementChannel.createSelectorForTest({ name })).value;
+  }
 }
 
 export function convertSelectOptionValues(values: string | ElementHandle | SelectOption | string[] | ElementHandle[] | SelectOption[] | null): { elements?: channels.ElementHandleChannel[], options?: SelectOption[] } {

src/client/playwright.ts

@@ -17,7 +17,7 @@
 import * as channels from '../protocol/channels';
 import { BrowserType } from './browserType';
 import { ChannelOwner } from './channelOwner';
-import { Selectors } from './selectors';
+import { Selectors, SelectorsOwner, sharedSelectors } from './selectors';
 import { Electron } from './electron';
 import { TimeoutError } from '../utils/errors';
 import { Size } from './types';
@@ -49,7 +49,8 @@ export class Playwright extends ChannelOwner<channels.PlaywrightChannel, channel
     this.devices = {};
     for (const { name, descriptor } of initializer.deviceDescriptors)
       this.devices[name] = descriptor;
-    this.selectors = Selectors.from(initializer.selectors);
+    this.selectors = sharedSelectors;
     this.errors = { TimeoutError };
+    this.selectors._addChannel(SelectorsOwner.from(initializer.selectors));
   }
 }

src/client/selectors.ts

@@ -14,26 +14,39 @@
  * limitations under the License.
  */
 
+import { evaluationScript } from './clientHelper';
 import * as channels from '../protocol/channels';
 import { ChannelOwner } from './channelOwner';
-import { ElementHandle } from './elementHandle';
-import { evaluationScript } from './clientHelper';
 
-export class Selectors extends ChannelOwner<channels.SelectorsChannel, channels.SelectorsInitializer> {
-  static from(selectors: channels.SelectorsChannel): Selectors {
-    return (selectors as any)._object;
+export class Selectors {
+  private _channels = new Set<SelectorsOwner>();
+  private _registrations: channels.SelectorsRegisterParams[] = [];
+
+  async register(name: string, script: string | Function | { path?: string, content?: string }, options: { contentScript?: boolean } = {}): Promise<void> {
+    const source = await evaluationScript(script, undefined, false);
+    const params = { ...options, name, source };
+    for (const channel of this._channels)
+      await channel._channel.register(params);
+    this._registrations.push(params);
   }
 
-  constructor(parent: ChannelOwner, type: string, guid: string, initializer: channels.SelectorsInitializer) {
-    super(parent, type, guid, initializer);
+  _addChannel(channel: SelectorsOwner) {
+    this._channels.add(channel);
+    for (const params of this._registrations) {
+      // This should not fail except for connection closure, but just in case we catch.
+      channel._channel.register(params).catch(e => {});
+    }
   }
 
-  async register(name: string, script: string | Function | { path?: string, content?: string }, options: { contentScript?: boolean } = {}): Promise<void> {
-    const source = await evaluationScript(script, undefined, false);
-    await this._channel.register({ ...options, name, source });
+  _removeChannel(channel: SelectorsOwner) {
+    this._channels.delete(channel);
   }
+}
 
-  async _createSelector(name: string, handle: ElementHandle<Element>): Promise<string | undefined> {
-    return (await this._channel.createSelector({ name, handle: handle._elementChannel })).value;
+export class SelectorsOwner extends ChannelOwner<channels.SelectorsChannel, channels.SelectorsInitializer> {
+  static from(browser: channels.SelectorsChannel): SelectorsOwner {
+    return (browser as any)._object;
   }
 }
+
+export const sharedSelectors = new Selectors();

src/dispatchers/browserDispatcher.ts

@@ -23,8 +23,8 @@ import { CRBrowser } from '../server/chromium/crBrowser';
 import { PageDispatcher } from './pageDispatcher';
 
 export class BrowserDispatcher extends Dispatcher<Browser, channels.BrowserInitializer> implements channels.BrowserChannel {
-  constructor(scope: DispatcherScope, browser: Browser, guid?: string) {
-    super(scope, browser, 'Browser', { version: browser.version(), name: browser._options.name }, true, guid);
+  constructor(scope: DispatcherScope, browser: Browser) {
+    super(scope, browser, 'Browser', { version: browser.version(), name: browser._options.name }, true);
     browser.on(Browser.Events.Disconnected, () => this._didClose());
   }
 

src/dispatchers/elementHandlerDispatcher.ts

@@ -156,4 +156,8 @@ export class ElementHandleDispatcher extends JSHandleDispatcher implements chann
   async waitForSelector(params: channels.ElementHandleWaitForSelectorParams): Promise<channels.ElementHandleWaitForSelectorResult> {
     return { element: ElementHandleDispatcher.createNullable(this._scope, await this._elementHandle.waitForSelector(params.selector, params)) };
   }
+
+  async createSelectorForTest(params: channels.ElementHandleCreateSelectorForTestParams): Promise<channels.ElementHandleCreateSelectorForTestResult> {
+    return { value: await this._elementHandle._page.selectors._createSelector(params.name, this._elementHandle as ElementHandle<Element>) };
+  }
 }

src/dispatchers/playwrightDispatcher.ts

@@ -18,10 +18,10 @@ import { Playwright } from '../server/playwright';
 import * as channels from '../protocol/channels';
 import { BrowserTypeDispatcher } from './browserTypeDispatcher';
 import { Dispatcher, DispatcherScope } from './dispatcher';
-import { SelectorsDispatcher } from './selectorsDispatcher';
 import { Electron } from '../server/electron/electron';
 import { ElectronDispatcher } from './electronDispatcher';
 import { DeviceDescriptors } from '../server/deviceDescriptors';
+import { SelectorsDispatcher } from './selectorsDispatcher';
 
 export class PlaywrightDispatcher extends Dispatcher<Playwright, channels.PlaywrightInitializer> implements channels.PlaywrightChannel {
   constructor(scope: DispatcherScope, playwright: Playwright) {

src/dispatchers/selectorsDispatcher.ts

@@ -17,8 +17,6 @@
 import { Dispatcher, DispatcherScope } from './dispatcher';
 import * as channels from '../protocol/channels';
 import { Selectors } from '../server/selectors';
-import { ElementHandleDispatcher } from './elementHandlerDispatcher';
-import * as dom from '../server/dom';
 
 export class SelectorsDispatcher extends Dispatcher<Selectors, channels.SelectorsInitializer> implements channels.SelectorsChannel {
   constructor(scope: DispatcherScope, selectors: Selectors) {
@@ -28,8 +26,4 @@ export class SelectorsDispatcher extends Dispatcher<Selectors, channels.Selector
   async register(params: channels.SelectorsRegisterParams): Promise<void> {
     await this._object.register(params.name, params.source, params.contentScript);
   }
-
-  async createSelector(params: channels.SelectorsCreateSelectorParams): Promise<channels.SelectorsCreateSelectorResult> {
-    return { value: await this._object._createSelector(params.name, (params.handle as ElementHandleDispatcher)._object as dom.ElementHandle<Element>) };
-  }
 }

src/protocol/channels.ts

@@ -109,11 +109,18 @@ export type PlaywrightInitializer = {
 export interface PlaywrightChannel extends Channel {
 }
 
+// ----------- RemoteBrowser -----------
+export type RemoteBrowserInitializer = {
+  browser: BrowserChannel,
+  selectors: SelectorsChannel,
+};
+export interface RemoteBrowserChannel extends Channel {
+}
+
 // ----------- Selectors -----------
 export type SelectorsInitializer = {};
 export interface SelectorsChannel extends Channel {
   register(params: SelectorsRegisterParams): Promise<SelectorsRegisterResult>;
-  createSelector(params: SelectorsCreateSelectorParams): Promise<SelectorsCreateSelectorResult>;
 }
 export type SelectorsRegisterParams = {
   name: string,
@@ -124,16 +131,6 @@ export type SelectorsRegisterOptions = {
   contentScript?: boolean,
 };
 export type SelectorsRegisterResult = void;
-export type SelectorsCreateSelectorParams = {
-  name: string,
-  handle: ElementHandleChannel,
-};
-export type SelectorsCreateSelectorOptions = {
-
-};
-export type SelectorsCreateSelectorResult = {
-  value?: string,
-};
 
 // ----------- BrowserType -----------
 export type BrowserTypeInitializer = {
@@ -1625,6 +1622,7 @@ export interface ElementHandleChannel extends JSHandleChannel {
   uncheck(params: ElementHandleUncheckParams): Promise<ElementHandleUncheckResult>;
   waitForElementState(params: ElementHandleWaitForElementStateParams): Promise<ElementHandleWaitForElementStateResult>;
   waitForSelector(params: ElementHandleWaitForSelectorParams): Promise<ElementHandleWaitForSelectorResult>;
+  createSelectorForTest(params: ElementHandleCreateSelectorForTestParams): Promise<ElementHandleCreateSelectorForTestResult>;
 }
 export type ElementHandleEvalOnSelectorParams = {
   selector: string,
@@ -1936,6 +1934,15 @@ export type ElementHandleWaitForSelectorOptions = {
 export type ElementHandleWaitForSelectorResult = {
   element?: ElementHandleChannel,
 };
+export type ElementHandleCreateSelectorForTestParams = {
+  name: string,
+};
+export type ElementHandleCreateSelectorForTestOptions = {
+
+};
+export type ElementHandleCreateSelectorForTestResult = {
+  value?: string,
+};
 
 // ----------- Request -----------
 export type RequestInitializer = {

src/protocol/protocol.yml

@@ -146,6 +146,13 @@ Playwright:
     selectors: Selectors
 
 
+RemoteBrowser:
+  type: interface
+
+  initializer:
+    browser: Browser
+    selectors: Selectors
+
 
 Selectors:
   type: interface
@@ -158,14 +165,6 @@ Selectors:
         source: string
         contentScript: boolean?
 
-    createSelector:
-      parameters:
-        name: string
-        handle: ElementHandle
-      returns:
-        value: string?
-
-
 
 BrowserType:
   type: interface
@@ -1606,6 +1605,12 @@ ElementHandle:
       returns:
         element: ElementHandle?
 
+    createSelectorForTest:
+      parameters:
+        name: string
+      returns:
+        value: string?
+
 
 Request:
   type: interface