chore: refactor injected script harness (#2259)

Author: pavelfeldman

src/chromium/crExecutionContext.ts

@@ -49,7 +49,7 @@ export class CRExecutionContext implements js.ExecutionContextDelegate {
       }).catch(rewriteError);
       if (exceptionDetails)
         throw new Error('Evaluation failed: ' + getExceptionMessage(exceptionDetails));
-      return returnByValue ? valueFromRemoteObject(remoteObject) : context._createHandle(remoteObject);
+      return returnByValue ? valueFromRemoteObject(remoteObject) : context.createHandle(remoteObject);
     }
 
     if (typeof pageFunction !== 'function')
@@ -91,7 +91,7 @@ export class CRExecutionContext implements js.ExecutionContextDelegate {
       }).catch(rewriteError);
       if (exceptionDetails)
         throw new Error('Evaluation failed: ' + getExceptionMessage(exceptionDetails));
-      return returnByValue ? valueFromRemoteObject(remoteObject) : context._createHandle(remoteObject);
+      return returnByValue ? valueFromRemoteObject(remoteObject) : context.createHandle(remoteObject);
     } finally {
       dispose();
     }
@@ -122,7 +122,7 @@ export class CRExecutionContext implements js.ExecutionContextDelegate {
     for (const property of response.result) {
       if (!property.enumerable)
         continue;
-      result.set(property.name, handle._context._createHandle(property.value));
+      result.set(property.name, handle._context.createHandle(property.value));
     }
     return result;
   }

src/chromium/crPage.ts

@@ -577,7 +577,7 @@ class FrameSession {
       session.send('Runtime.runIfWaitingForDebugger'),
     ]).catch(logError(this._page));  // This might fail if the target is closed before we initialize.
     session.on('Runtime.consoleAPICalled', event => {
-      const args = event.args.map(o => worker._existingExecutionContext!._createHandle(o));
+      const args = event.args.map(o => worker._existingExecutionContext!.createHandle(o));
       this._page._addConsoleMessage(event.type, args, toConsoleMessageLocation(event.stackTrace));
     });
     session.on('Runtime.exceptionThrown', exception => this._page.emit(Events.Page.PageError, exceptionToError(exception.exceptionDetails)));
@@ -608,7 +608,7 @@ class FrameSession {
       return;
     }
     const context = this._contextIdToContext.get(event.executionContextId)!;
-    const values = event.args.map(arg => context._createHandle(arg));
+    const values = event.args.map(arg => context.createHandle(arg));
     this._page._addConsoleMessage(event.type, values, toConsoleMessageLocation(event.stackTrace));
   }
 
@@ -846,7 +846,7 @@ class FrameSession {
     }).catch(logError(this._page));
     if (!result || result.object.subtype === 'null')
       throw new Error('Unable to adopt element handle from a different document');
-    return to._createHandle(result.object).asElement()!;
+    return to.createHandle(result.object).asElement()!;
   }
 }
 

src/dom.ts

@@ -20,7 +20,8 @@ import * as path from 'path';
 import * as util from 'util';
 import * as frames from './frames';
 import { assert, helper } from './helper';
-import { Injected, InjectedResult } from './injected/injected';
+import InjectedScript from './injected/injectedScript';
+import * as injectedScriptSource from './generated/injectedScriptSource';
 import * as input from './input';
 import * as js from './javascript';
 import { Page } from './page';
@@ -52,29 +53,35 @@ export class FrameExecutionContext extends js.ExecutionContext {
     this.frame = frame;
   }
 
-  _adoptIfNeeded(handle: js.JSHandle): Promise<js.JSHandle> | null {
+  adoptIfNeeded(handle: js.JSHandle): Promise<js.JSHandle> | null {
     if (handle instanceof ElementHandle && handle._context !== this)
       return this.frame._page._delegate.adoptElementHandle(handle, this);
     return null;
   }
 
-  async _doEvaluateInternal(returnByValue: boolean, waitForNavigations: boolean, pageFunction: string | Function, ...args: any[]): Promise<any> {
+  async doEvaluateInternal(returnByValue: boolean, waitForNavigations: boolean, pageFunction: string | Function, ...args: any[]): Promise<any> {
     return await this.frame._page._frameManager.waitForSignalsCreatedBy(async () => {
       return this._delegate.evaluate(this, returnByValue, pageFunction, ...args);
     }, Number.MAX_SAFE_INTEGER, waitForNavigations ? undefined : { noWaitAfter: true });
   }
 
-  _createHandle(remoteObject: any): js.JSHandle {
+  createHandle(remoteObject: any): js.JSHandle {
     if (this.frame._page._delegate.isElementHandle(remoteObject))
       return new ElementHandle(this, remoteObject);
-    return super._createHandle(remoteObject);
+    return super.createHandle(remoteObject);
   }
 
-  _injected(): Promise<js.JSHandle<Injected>> {
+  injectedScript(): Promise<js.JSHandle<InjectedScript>> {
     if (!this._injectedPromise) {
-      this._injectedPromise = selectors._prepareEvaluator(this).then(evaluator => {
-        return this.evaluateHandleInternal(evaluator => evaluator.injected, evaluator);
-      });
+      const custom: string[] = [];
+      for (const [name, { source }] of selectors._engines)
+        custom.push(`{ name: '${name}', engine: (${source}) }`);
+      const source = `
+        new (${injectedScriptSource.source})([
+          ${custom.join(',\n')}
+        ])
+      `;
+      this._injectedPromise = this.doEvaluateInternal(false /* returnByValue */, false /* waitForNavigations */, source);
     }
     return this._injectedPromise;
   }
@@ -94,14 +101,14 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
     return this;
   }
 
-  async _evaluateInMain<R, Arg>(pageFunction: types.FuncOn<{ injected: Injected, node: T }, Arg, R>, arg: Arg): Promise<R> {
+  async _evaluateInMain<R, Arg>(pageFunction: types.FuncOn<{ injected: InjectedScript, node: T }, Arg, R>, arg: Arg): Promise<R> {
     const main = await this._context.frame._mainContext();
-    return main._doEvaluateInternal(true /* returnByValue */, true /* waitForNavigations */, pageFunction, { injected: await main._injected(), node: this }, arg);
+    return main.doEvaluateInternal(true /* returnByValue */, true /* waitForNavigations */, pageFunction, { injected: await main.injectedScript(), node: this }, arg);
   }
 
-  async _evaluateInUtility<R, Arg>(pageFunction: types.FuncOn<{ injected: Injected, node: T }, Arg, R>, arg: Arg): Promise<R> {
+  async _evaluateInUtility<R, Arg>(pageFunction: types.FuncOn<{ injected: InjectedScript, node: T }, Arg, R>, arg: Arg): Promise<R> {
     const utility = await this._context.frame._utilityContext();
-    return utility._doEvaluateInternal(true /* returnByValue */, true /* waitForNavigations */, pageFunction, { injected: await utility._injected(), node: this }, arg);
+    return utility.doEvaluateInternal(true /* returnByValue */, true /* waitForNavigations */, pageFunction, { injected: await utility.injectedScript(), node: this }, arg);
   }
 
   async ownerFrame(): Promise<frames.Frame | null> {
@@ -352,7 +359,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
   async setInputFiles(files: string | types.FilePayload | string[] | types.FilePayload[], options?: types.NavigatingActionWaitOptions) {
     this._page._log(inputLog, `elementHandle.setInputFiles(...)`);
     const deadline = this._page._timeoutSettings.computeDeadline(options);
-    const injectedResult = await this._evaluateInUtility(({ node }): InjectedResult<boolean> => {
+    const injectedResult = await this._evaluateInUtility(({ node }): types.InjectedScriptResult<boolean> => {
       if (node.nodeType !== Node.ELEMENT_NODE || (node as Node as Element).tagName !== 'INPUT')
         return { status: 'error', error: 'Node is not an HTMLInputElement' };
       if (!node.isConnected)
@@ -500,7 +507,7 @@ export function toFileTransferPayload(files: types.FilePayload[]): types.FileTra
   }));
 }
 
-function handleInjectedResult<T = undefined>(injectedResult: InjectedResult<T>, timeoutMessage: string): T {
+function handleInjectedResult<T = undefined>(injectedResult: types.InjectedScriptResult<T>, timeoutMessage: string): T {
   if (injectedResult.status === 'notconnected')
     throw new NotConnectedError();
   if (injectedResult.status === 'timeout')

src/firefox/ffExecutionContext.ts

@@ -39,7 +39,7 @@ export class FFExecutionContext implements js.ExecutionContextDelegate {
       checkException(payload.exceptionDetails);
       if (returnByValue)
         return deserializeValue(payload.result!);
-      return context._createHandle(payload.result);
+      return context.createHandle(payload.result);
     }
     if (typeof pageFunction !== 'function')
       throw new Error(`Expected to get |string| or |function| as the first argument, but got "${pageFunction}" instead.`);
@@ -71,7 +71,7 @@ export class FFExecutionContext implements js.ExecutionContextDelegate {
       checkException(payload.exceptionDetails);
       if (returnByValue)
         return deserializeValue(payload.result!);
-      return context._createHandle(payload.result);
+      return context.createHandle(payload.result);
     } finally {
       dispose();
     }
@@ -97,7 +97,7 @@ export class FFExecutionContext implements js.ExecutionContextDelegate {
     });
     const result = new Map();
     for (const property of response.properties)
-      result.set(property.name, handle._context._createHandle(property.value));
+      result.set(property.name, handle._context.createHandle(property.value));
     return result;
   }
 

src/firefox/ffPage.ts

@@ -184,7 +184,7 @@ export class FFPage implements PageDelegate {
   _onConsole(payload: Protocol.Runtime.consolePayload) {
     const {type, args, executionContextId, location} = payload;
     const context = this._contextIdToContext.get(executionContextId)!;
-    this._page._addConsoleMessage(type, args.map(arg => context._createHandle(arg)), location);
+    this._page._addConsoleMessage(type, args.map(arg => context.createHandle(arg)), location);
   }
 
   _onDialogOpened(params: Protocol.Page.dialogOpenedPayload) {
@@ -205,7 +205,7 @@ export class FFPage implements PageDelegate {
   async _onFileChooserOpened(payload: Protocol.Page.fileChooserOpenedPayload) {
     const {executionContextId, element} = payload;
     const context = this._contextIdToContext.get(executionContextId)!;
-    const handle = context._createHandle(element).asElement()!;
+    const handle = context.createHandle(element).asElement()!;
     this._page._onFileChooserOpened(handle);
   }
 
@@ -229,7 +229,7 @@ export class FFPage implements PageDelegate {
     workerSession.on('Runtime.console', event => {
       const {type, args, location} = event;
       const context = worker._existingExecutionContext!;
-      this._page._addConsoleMessage(type, args.map(arg => context._createHandle(arg)), location);
+      this._page._addConsoleMessage(type, args.map(arg => context.createHandle(arg)), location);
     });
     // Note: we receive worker exceptions directly from the page.
   }
@@ -457,7 +457,7 @@ export class FFPage implements PageDelegate {
     });
     if (!result.remoteObject)
       throw new Error('Unable to adopt element handle from a different document');
-    return to._createHandle(result.remoteObject) as dom.ElementHandle<T>;
+    return to.createHandle(result.remoteObject) as dom.ElementHandle<T>;
   }
 
   async getAccessibilityTree(needle?: dom.ElementHandle) {

src/frames.ts

@@ -786,7 +786,7 @@ export class Frame {
     const task = async (context: dom.FrameExecutionContext) => context.evaluateHandleInternal(({ injected, predicateBody, polling, timeout, arg }) => {
       const innerPredicate = new Function('arg', predicateBody);
       return injected.poll(polling, timeout, () => innerPredicate(arg));
-    }, { injected: await context._injected(), predicateBody, polling, timeout: helper.timeUntilDeadline(deadline), arg });
+    }, { injected: await context.injectedScript(), predicateBody, polling, timeout: helper.timeUntilDeadline(deadline), arg });
     return this._scheduleRerunnableTask(task, 'main', deadline) as any as types.SmartHandle<R>;
   }
 

src/injected/injected.ts renamed to src/injected/injectedScript.ts

@@ -15,15 +15,83 @@
  */
 
 import * as types from '../types';
+import { createAttributeEngine } from './attributeSelectorEngine';
+import { createCSSEngine } from './cssSelectorEngine';
+import { SelectorEngine, SelectorRoot } from './selectorEngine';
+import { createTextSelector } from './textSelectorEngine';
+import { XPathEngine } from './xpathSelectorEngine';
 
 type Predicate<T> = () => T;
-export type InjectedResult<T = undefined> =
-  (T extends undefined ? { status: 'success', value?: T} : { status: 'success', value: T }) |
-  { status: 'notconnected' } |
-  { status: 'timeout' } |
-  { status: 'error', error: string };
 
-export class Injected {
+export default class InjectedScript {
+  readonly engines: Map<string, SelectorEngine>;
+
+  constructor(customEngines: { name: string, engine: SelectorEngine}[]) {
+    this.engines = new Map();
+    // Note: keep predefined names in sync with Selectors class.
+    this.engines.set('css', createCSSEngine(true));
+    this.engines.set('css:light', createCSSEngine(false));
+    this.engines.set('xpath', XPathEngine);
+    this.engines.set('xpath:light', XPathEngine);
+    this.engines.set('text', createTextSelector(true));
+    this.engines.set('text:light', createTextSelector(false));
+    this.engines.set('id', createAttributeEngine('id', true));
+    this.engines.set('id:light', createAttributeEngine('id', false));
+    this.engines.set('data-testid', createAttributeEngine('data-testid', true));
+    this.engines.set('data-testid:light', createAttributeEngine('data-testid', false));
+    this.engines.set('data-test-id', createAttributeEngine('data-test-id', true));
+    this.engines.set('data-test-id:light', createAttributeEngine('data-test-id', false));
+    this.engines.set('data-test', createAttributeEngine('data-test', true));
+    this.engines.set('data-test:light', createAttributeEngine('data-test', false));
+    for (const {name, engine} of customEngines)
+      this.engines.set(name, engine);
+  }
+
+  querySelector(selector: types.ParsedSelector, root: Node): Element | undefined {
+    if (!(root as any)['querySelector'])
+      throw new Error('Node is not queryable.');
+    return this._querySelectorRecursively(root as SelectorRoot, selector, 0);
+  }
+
+  private _querySelectorRecursively(root: SelectorRoot, selector: types.ParsedSelector, index: number): Element | undefined {
+    const current = selector.parts[index];
+    if (index === selector.parts.length - 1)
+      return this.engines.get(current.name)!.query(root, current.body);
+    const all = this.engines.get(current.name)!.queryAll(root, current.body);
+    for (const next of all) {
+      const result = this._querySelectorRecursively(next, selector, index + 1);
+      if (result)
+        return selector.capture === index ? next : result;
+    }
+  }
+
+  querySelectorAll(selector: types.ParsedSelector, root: Node): Element[] {
+    if (!(root as any)['querySelectorAll'])
+      throw new Error('Node is not queryable.');
+    const capture = selector.capture === undefined ? selector.parts.length - 1 : selector.capture;
+    // Query all elements up to the capture.
+    const partsToQuerAll = selector.parts.slice(0, capture + 1);
+    // Check they have a descendant matching everything after the capture.
+    const partsToCheckOne = selector.parts.slice(capture + 1);
+    let set = new Set<SelectorRoot>([ root as SelectorRoot ]);
+    for (const { name, body } of partsToQuerAll) {
+      const newSet = new Set<Element>();
+      for (const prev of set) {
+        for (const next of this.engines.get(name)!.queryAll(prev, body)) {
+          if (newSet.has(next))
+            continue;
+          newSet.add(next);
+        }
+      }
+      set = newSet;
+    }
+    const candidates = Array.from(set) as Element[];
+    if (!partsToCheckOne.length)
+      return candidates;
+    const partial = { parts: partsToCheckOne };
+    return candidates.filter(e => !!this._querySelectorRecursively(e, partial, 0));
+  }
+
   isVisible(element: Element): boolean {
     // Note: this logic should be similar to waitForDisplayedAtStablePosition() to avoid surprises.
     if (!element.ownerDocument || !element.ownerDocument.defaultView)
@@ -95,7 +163,7 @@ export class Injected {
     return { left: parseInt(style.borderLeftWidth || '', 10), top: parseInt(style.borderTopWidth || '', 10) };
   }
 
-  selectOptions(node: Node, optionsToSelect: (Node | types.SelectOption)[]): InjectedResult<string[]> {
+  selectOptions(node: Node, optionsToSelect: (Node | types.SelectOption)[]): types.InjectedScriptResult<string[]> {
     if (node.nodeName.toLowerCase() !== 'select')
       return { status: 'error', error: 'Element is not a <select> element.' };
     if (!node.isConnected)
@@ -126,7 +194,7 @@ export class Injected {
     return { status: 'success', value: options.filter(option => option.selected).map(option => option.value) };
   }
 
-  fill(node: Node, value: string): InjectedResult<boolean> {
+  fill(node: Node, value: string): types.InjectedScriptResult<boolean> {
     if (node.nodeType !== Node.ELEMENT_NODE)
       return { status: 'error', error: 'Node is not of type HTMLElement' };
     const element = node as HTMLElement;
@@ -175,7 +243,7 @@ export class Injected {
     return result;
   }
 
-  selectText(node: Node): InjectedResult {
+  selectText(node: Node): types.InjectedScriptResult {
     if (node.nodeType !== Node.ELEMENT_NODE)
       return { status: 'error', error: 'Node is not of type HTMLElement' };
     if (!node.isConnected)
@@ -207,7 +275,7 @@ export class Injected {
     return { status: 'success' };
   }
 
-  focusNode(node: Node): InjectedResult {
+  focusNode(node: Node): types.InjectedScriptResult {
     if (!node.isConnected)
       return { status: 'notconnected' };
     if (!(node as any)['focus'])
@@ -262,7 +330,7 @@ export class Injected {
     input.dispatchEvent(new Event('change', { 'bubbles': true }));
   }
 
-  async waitForDisplayedAtStablePosition(node: Node, rafCount: number, timeout: number): Promise<InjectedResult> {
+  async waitForDisplayedAtStablePosition(node: Node, rafCount: number, timeout: number): Promise<types.InjectedScriptResult> {
     if (!node.isConnected)
       return { status: 'notconnected' };
     const element = node.nodeType === Node.ELEMENT_NODE ? (node as Element) : node.parentElement;
@@ -305,7 +373,7 @@ export class Injected {
     return { status: result === 'notconnected' ? 'notconnected' : (result ? 'success' : 'timeout') };
   }
 
-  checkHitTargetAt(node: Node, point: types.Point): InjectedResult<boolean> {
+  checkHitTargetAt(node: Node, point: types.Point): types.InjectedScriptResult<boolean> {
     let element = node.nodeType === Node.ELEMENT_NODE ? (node as Element) : node.parentElement;
     while (element && window.getComputedStyle(element).pointerEvents === 'none')
       element = element.parentElement;

src/injected/selectorEvaluator.webpack.config.js renamed to src/injected/injectedScript.webpack.config.js

@@ -18,7 +18,7 @@ const path = require('path');
 const InlineSource = require('./webpack-inline-source-plugin.js');
 
 module.exports = {
-  entry: path.join(__dirname, 'selectorEvaluator.ts'),
+  entry: path.join(__dirname, 'injectedScript.ts'),
   devtool: 'source-map',
   module: {
     rules: [
@@ -36,10 +36,10 @@ module.exports = {
     extensions: [ '.tsx', '.ts', '.js' ]
   },
   output: {
-    filename: 'selectorEvaluatorSource.js',
+    filename: 'injectedScriptSource.js',
     path: path.resolve(__dirname, '../../lib/injected/packed')
   },
   plugins: [
-    new InlineSource(path.join(__dirname, '..', 'generated', 'selectorEvaluatorSource.ts')),
+    new InlineSource(path.join(__dirname, '..', 'generated', 'injectedScriptSource.ts')),
   ]
 };