Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
feat: add support for string-based custom queries (#5753)
  • Loading branch information
paullewis committed Apr 30, 2020
1 parent c212126 commit 4a47867
Show file tree
Hide file tree
Showing 5 changed files with 207 additions and 17 deletions.
29 changes: 24 additions & 5 deletions src/DOMWorld.ts
Expand Up @@ -23,6 +23,11 @@ import {ExecutionContext} from './ExecutionContext';
import {TimeoutSettings} from './TimeoutSettings';
import {MouseButtonInput} from './Input';
import {FrameManager, Frame} from './FrameManager';
import {getQueryHandlerAndSelector, QueryHandler} from './QueryHandler';

// This predicateQueryHandler is declared here so that TypeScript knows about it
// when it is used in the predicate function below.
declare const predicateQueryHandler: QueryHandler;

const readFileAsync = helper.promisify(fs.readFile);

Expand Down Expand Up @@ -364,7 +369,7 @@ export class DOMWorld {
polling = 'raf',
timeout = this._timeoutSettings.timeout(),
} = options;
return new WaitTask(this, pageFunction, 'function', polling, timeout, ...args).promise;
return new WaitTask(this, pageFunction, undefined, 'function', polling, timeout, ...args).promise;
}

async title(): Promise<string> {
Expand All @@ -379,7 +384,8 @@ export class DOMWorld {
} = options;
const polling = waitForVisible || waitForHidden ? 'raf' : 'mutation';
const title = `${isXPath ? 'XPath' : 'selector'} "${selectorOrXPath}"${waitForHidden ? ' to be hidden' : ''}`;
const waitTask = new WaitTask(this, predicate, title, polling, timeout, selectorOrXPath, isXPath, waitForVisible, waitForHidden);
const {updatedSelector, queryHandler} = getQueryHandlerAndSelector(selectorOrXPath, (element, selector) => document.querySelector(selector));
const waitTask = new WaitTask(this, predicate, queryHandler, title, polling, timeout, updatedSelector, isXPath, waitForVisible, waitForHidden);
const handle = await waitTask.promise;
if (!handle.asElement()) {
await handle.dispose();
Expand All @@ -397,7 +403,7 @@ export class DOMWorld {
function predicate(selectorOrXPath: string, isXPath: boolean, waitForVisible: boolean, waitForHidden: boolean): Node | null | boolean {
const node = isXPath
? document.evaluate(selectorOrXPath, document, null, XPathResult.FIRST_ORDERED_NODE_TYPE, null).singleNodeValue
: document.querySelector(selectorOrXPath);
: predicateQueryHandler ? predicateQueryHandler(document, selectorOrXPath) as Element : document.querySelector(selectorOrXPath);
if (!node)
return waitForHidden;
if (!waitForVisible && !waitForHidden)
Expand Down Expand Up @@ -430,18 +436,31 @@ class WaitTask {
_timeoutTimer?: NodeJS.Timeout;
_terminated = false;

constructor(domWorld: DOMWorld, predicateBody: Function | string, title: string, polling: string | number, timeout: number, ...args: unknown[]) {
constructor(domWorld: DOMWorld, predicateBody: Function | string, predicateQueryHandlerBody: Function | string | undefined, title: string, polling: string | number, timeout: number, ...args: unknown[]) {
if (helper.isString(polling))
assert(polling === 'raf' || polling === 'mutation', 'Unknown polling option: ' + polling);
else if (helper.isNumber(polling))
assert(polling > 0, 'Cannot poll with non-positive interval: ' + polling);
else
throw new Error('Unknown polling options: ' + polling);

function getPredicateBody(predicateBody: Function | string, predicateQueryHandlerBody: Function | string) {
if (helper.isString(predicateBody))
return `return (${predicateBody});`;
if (predicateQueryHandlerBody) {
return `
return (function wrapper(args) {
const predicateQueryHandler = ${predicateQueryHandlerBody};
return (${predicateBody})(...args);
})(args);`;
}
return `return (${predicateBody})(...args);`;
}

this._domWorld = domWorld;
this._polling = polling;
this._timeout = timeout;
this._predicateBody = helper.isString(predicateBody) ? 'return (' + predicateBody + ')' : 'return (' + predicateBody + ')(...args)';
this._predicateBody = getPredicateBody(predicateBody, predicateQueryHandlerBody);
this._args = args;
this._runCount = 0;
domWorld._waitTasks.add(this);
Expand Down
24 changes: 12 additions & 12 deletions src/JSHandle.ts
Expand Up @@ -19,6 +19,7 @@ import {ExecutionContext} from './ExecutionContext';
import {CDPSession} from './Connection';
import {KeyInput} from './USKeyboardLayout';
import {FrameManager, Frame} from './FrameManager';
import {getQueryHandlerAndSelector} from './QueryHandler';

interface BoxModel {
content: Array<{x: number; y: number}>;
Expand Down Expand Up @@ -427,10 +428,10 @@ export class ElementHandle extends JSHandle {
}

async $(selector: string): Promise<ElementHandle | null> {
const handle = await this.evaluateHandle(
(element, selector) => element.querySelector(selector),
selector
);
const defaultHandler = (element: Element, selector: string) => element.querySelector(selector);
const {updatedSelector, queryHandler} = getQueryHandlerAndSelector(selector, defaultHandler);

const handle = await this.evaluateHandle(queryHandler, updatedSelector);
const element = handle.asElement();
if (element)
return element;
Expand All @@ -443,10 +444,10 @@ export class ElementHandle extends JSHandle {
* @return {!Promise<!Array<!ElementHandle>>}
*/
async $$(selector: string): Promise<ElementHandle[]> {
const arrayHandle = await this.evaluateHandle(
(element, selector) => element.querySelectorAll(selector),
selector
);
const defaultHandler = (element: Element, selector: string) => element.querySelectorAll(selector);
const {updatedSelector, queryHandler} = getQueryHandlerAndSelector(selector, defaultHandler);

const arrayHandle = await this.evaluateHandle(queryHandler, updatedSelector);
const properties = await arrayHandle.getProperties();
await arrayHandle.dispose();
const result = [];
Expand All @@ -468,11 +469,10 @@ export class ElementHandle extends JSHandle {
}

async $$eval<ReturnType extends any>(selector: string, pageFunction: Function | string, ...args: unknown[]): Promise<ReturnType> {
const arrayHandle = await this.evaluateHandle(
(element, selector) => Array.from(element.querySelectorAll(selector)),
selector
);
const defaultHandler = (element: Element, selector: string) => Array.from(element.querySelectorAll(selector));
const {updatedSelector, queryHandler} = getQueryHandlerAndSelector(selector, defaultHandler);

const arrayHandle = await this.evaluateHandle(queryHandler, updatedSelector);
const result = await arrayHandle.evaluate<ReturnType>(pageFunction, ...args);
await arrayHandle.dispose();
return result;
Expand Down
24 changes: 24 additions & 0 deletions src/Puppeteer.js
Expand Up @@ -20,6 +20,7 @@ const DeviceDescriptors = require('./DeviceDescriptors');
// Import used as typedef
// eslint-disable-next-line no-unused-vars
const {Browser} = require('./Browser');
const QueryHandler = require('./QueryHandler');

module.exports = class {
/**
Expand Down Expand Up @@ -147,4 +148,27 @@ module.exports = class {
createBrowserFetcher(options) {
return new BrowserFetcher(this._projectRoot, options);
}

/**
* @param {string} name
* @param {!Function} queryHandler
*/
__experimental_registerCustomQueryHandler(name, queryHandler) {
QueryHandler.registerCustomQueryHandler(name, queryHandler);
}

/**
* @param {string} name
*/
__experimental_unregisterCustomQueryHandler(name) {
QueryHandler.unregisterCustomQueryHandler(name);
}

__experimental_customQueryHandlers() {
return QueryHandler.customQueryHandlers();
}

__experimental_clearQueryHandlers() {
QueryHandler.clearQueryHandlers();
}
};
74 changes: 74 additions & 0 deletions src/QueryHandler.ts
@@ -0,0 +1,74 @@
/**
* Copyright 2020 Google Inc. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

export interface QueryHandler {
(element: Element | Document, selector: string): Element | Element[] | NodeListOf<Element>;
}

const _customQueryHandlers = new Map<string, QueryHandler>();

export function registerCustomQueryHandler(name: string, handler: Function): void {
if (_customQueryHandlers.get(name))
throw new Error(`A custom query handler named "${name}" already exists`);

const isValidName = /^[a-zA-Z]+$/.test(name);
if (!isValidName)
throw new Error(`Custom query handler names may only contain [a-zA-Z]`);

_customQueryHandlers.set(name, handler as QueryHandler);
}

/**
* @param {string} name
*/
export function unregisterCustomQueryHandler(name: string): void {
_customQueryHandlers.delete(name);
}

export function customQueryHandlers(): Map<string, QueryHandler> {
return _customQueryHandlers;
}

export function clearQueryHandlers(): void {
_customQueryHandlers.clear();
}

export function getQueryHandlerAndSelector(selector: string, defaultQueryHandler: QueryHandler):
{ updatedSelector: string; queryHandler: QueryHandler} {
const hasCustomQueryHandler = /^[a-zA-Z]+\//.test(selector);
if (!hasCustomQueryHandler)
return {updatedSelector: selector, queryHandler: defaultQueryHandler};

const index = selector.indexOf('/');
const name = selector.slice(0, index);
const updatedSelector = selector.slice(index + 1);
const queryHandler = customQueryHandlers().get(name);
if (!queryHandler)
throw new Error(`Query set to use "${name}", but no query handler of that name was found`);

return {
updatedSelector,
queryHandler
};
}

module.exports = {
registerCustomQueryHandler,
unregisterCustomQueryHandler,
customQueryHandlers,
getQueryHandlerAndSelector,
clearQueryHandlers
};
73 changes: 73 additions & 0 deletions test/elementhandle.spec.js
Expand Up @@ -248,4 +248,77 @@ describe('ElementHandle specs', function() {
}
});
});

describe('Custom queries', function() {
this.afterEach(() => {
const {puppeteer} = getTestState();
puppeteer.__experimental_clearQueryHandlers();
});
it('should register and unregister', async() => {
const {page, puppeteer} = getTestState();
await page.setContent('<div id="not-foo"></div><div id="foo"></div>');

// Register.
puppeteer.__experimental_registerCustomQueryHandler('getById', (element, selector) => document.querySelector(`[id="${selector}"]`));
const element = await page.$('getById/foo');
expect(await page.evaluate(element => element.id, element)).toBe('foo');

// Unregister.
puppeteer.__experimental_unregisterCustomQueryHandler('getById');
try {
await page.$('getById/foo');
expect.fail('Custom query handler not set - throw expected');
} catch (error) {
expect(error).toStrictEqual(new Error('Query set to use "getById", but no query handler of that name was found'));
}
});
it('should throw with invalid query names', () => {
try {
const {puppeteer} = getTestState();
puppeteer.__experimental_registerCustomQueryHandler('1/2/3', (element, selector) => {});
expect.fail('Custom query handler name was invalid - throw expected');
} catch (error) {
expect(error).toStrictEqual(new Error('Custom query handler names may only contain [a-zA-Z]'));
}
});
it('should work for multiple elements', async() => {
const {page, puppeteer} = getTestState();
await page.setContent('<div id="not-foo"></div><div class="foo">Foo1</div><div class="foo baz">Foo2</div>');
puppeteer.__experimental_registerCustomQueryHandler('getByClass', (element, selector) => document.querySelectorAll(`.${selector}`));
const elements = await page.$$('getByClass/foo');
const classNames = await Promise.all(elements.map(async element => await page.evaluate(element => element.className, element)));

expect(classNames).toStrictEqual(['foo', 'foo baz']);
});
it('should eval correctly', async() => {
const {page, puppeteer} = getTestState();
await page.setContent('<div id="not-foo"></div><div class="foo">Foo1</div><div class="foo baz">Foo2</div>');
puppeteer.__experimental_registerCustomQueryHandler('getByClass', (element, selector) => document.querySelectorAll(`.${selector}`));
const elements = await page.$$eval('getByClass/foo', divs => divs.length);

expect(elements).toBe(2);
});
it('should wait correctly with waitForSelector', async() => {
const {page, puppeteer} = getTestState();
puppeteer.__experimental_registerCustomQueryHandler('getByClass', (element, selector) => element.querySelector(`.${selector}`));
const waitFor = page.waitForSelector('getByClass/foo');

// Set the page content after the waitFor has been started.
await page.setContent('<div id="not-foo"></div><div class="foo">Foo1</div>');
const element = await waitFor;

expect(element).toBeDefined();
});
it('should wait correctly with waitFor', async() => {
const {page, puppeteer} = getTestState();
puppeteer.__experimental_registerCustomQueryHandler('getByClass', (element, selector) => element.querySelector(`.${selector}`));
const waitFor = page.waitFor('getByClass/foo');

// Set the page content after the waitFor has been started.
await page.setContent('<div id="not-foo"></div><div class="foo">Foo1</div>');
const element = await waitFor;

expect(element).toBeDefined();
});
});
});

0 comments on commit 4a47867

Please sign in to comment.