diff options
Diffstat (limited to 'remote/test/puppeteer/packages/puppeteer-core/src/injected/CustomQuerySelector.ts')
-rw-r--r-- | remote/test/puppeteer/packages/puppeteer-core/src/injected/CustomQuerySelector.ts | 59 |
1 files changed, 59 insertions, 0 deletions
diff --git a/remote/test/puppeteer/packages/puppeteer-core/src/injected/CustomQuerySelector.ts b/remote/test/puppeteer/packages/puppeteer-core/src/injected/CustomQuerySelector.ts new file mode 100644 index 0000000000..ccd041deea --- /dev/null +++ b/remote/test/puppeteer/packages/puppeteer-core/src/injected/CustomQuerySelector.ts @@ -0,0 +1,59 @@ +/** + * @license + * Copyright 2023 Google Inc. + * SPDX-License-Identifier: Apache-2.0 + */ + +import type {CustomQueryHandler} from '../common/CustomQueryHandler.js'; +import type {Awaitable, AwaitableIterable} from '../common/types.js'; + +export interface CustomQuerySelector { + querySelector(root: Node, selector: string): Awaitable<Node | null>; + querySelectorAll(root: Node, selector: string): AwaitableIterable<Node>; +} + +/** + * This class mimics the injected {@link CustomQuerySelectorRegistry}. + */ +class CustomQuerySelectorRegistry { + #selectors = new Map<string, CustomQuerySelector>(); + + register(name: string, handler: CustomQueryHandler): void { + if (!handler.queryOne && handler.queryAll) { + const querySelectorAll = handler.queryAll; + handler.queryOne = (node, selector) => { + for (const result of querySelectorAll(node, selector)) { + return result; + } + return null; + }; + } else if (handler.queryOne && !handler.queryAll) { + const querySelector = handler.queryOne; + handler.queryAll = (node, selector) => { + const result = querySelector(node, selector); + return result ? [result] : []; + }; + } else if (!handler.queryOne || !handler.queryAll) { + throw new Error('At least one query method must be defined.'); + } + + this.#selectors.set(name, { + querySelector: handler.queryOne, + querySelectorAll: handler.queryAll!, + }); + } + + unregister(name: string): void { + this.#selectors.delete(name); + } + + get(name: string): CustomQuerySelector | undefined { + return this.#selectors.get(name); + } + + clear() { + this.#selectors.clear(); + } +} + +export const customQuerySelectors = new CustomQuerySelectorRegistry(); |