feat(a11y-query): split waitFor logic for selectors and xpath (#6426)

The logic for waitForXPath and waitForSelector is currently very tightly coupled. This commit tries to untangle that relationship. This is the first step towards introducing built-in query handlers that are not executed in the page context (#6307).

src/common/DOMWorld.ts

@@ -37,12 +37,17 @@ import {
 } from './EvalTypes.js';
 import { isNode } from '../environment.js';
 
-// This predicateQueryHandler is declared here so that TypeScript knows about it
+// predicateQueryHandler and checkWaitForOptions are declared here so that
-// when it is used in the predicate function below.
+// TypeScript knows about them when used in the predicate function below.
 declare const predicateQueryHandler: (
   element: Element | Document,
   selector: string
 ) => Element | Element[] | NodeListOf<Element>;
+declare const checkWaitForOptions: (
+  node: Node,
+  waitForVisible: boolean,
+  waitForHidden: boolean
+) => Element | null | boolean;
 
 /**
  * @public
@@ -457,48 +462,9 @@ export class DOMWorld {
     await handle.dispose();
   }
 
-  waitForSelector(
+  async waitForSelector(
     selector: string,
     options: WaitForSelectorOptions
-  ): Promise<ElementHandle | null> {
-    return this._waitForSelectorOrXPath(selector, false, options);
-  }
-
-  waitForXPath(
-    xpath: string,
-    options: WaitForSelectorOptions
-  ): Promise<ElementHandle | null> {
-    return this._waitForSelectorOrXPath(xpath, true, options);
-  }
-
-  waitForFunction(
-    pageFunction: Function | string,
-    options: { polling?: string | number; timeout?: number } = {},
-    ...args: SerializableOrJSHandle[]
-  ): Promise<JSHandle> {
-    const {
-      polling = 'raf',
-      timeout = this._timeoutSettings.timeout(),
-    } = options;
-    return new WaitTask(
-      this,
-      pageFunction,
-      undefined,
-      'function',
-      polling,
-      timeout,
-      ...args
-    ).promise;
-  }
-
-  async title(): Promise<string> {
-    return this.evaluate(() => document.title);
-  }
-
-  private async _waitForSelectorOrXPath(
-    selectorOrXPath: string,
-    isXPath: boolean,
-    options: WaitForSelectorOptions = {}
   ): Promise<ElementHandle | null> {
     const {
       visible: waitForVisible = false,
@@ -506,48 +472,103 @@ export class DOMWorld {
       timeout = this._timeoutSettings.timeout(),
     } = options;
     const polling = waitForVisible || waitForHidden ? 'raf' : 'mutation';
-    const title = `${isXPath ? 'XPath' : 'selector'} "${selectorOrXPath}"${
+    const title = `selector \`${selector}\`${
       waitForHidden ? ' to be hidden' : ''
     }`;
     const { updatedSelector, queryHandler } = getQueryHandlerAndSelector(
-      selectorOrXPath
+      selector
     );
+    function predicate(
+      selector: string,
+      waitForVisible: boolean,
+      waitForHidden: boolean
+    ): Node | null | boolean {
+      const node = predicateQueryHandler
+        ? (predicateQueryHandler(document, selector) as Element)
+        : document.querySelector(selector);
+      return checkWaitForOptions(node, waitForVisible, waitForHidden);
+    }
     const waitTask = new WaitTask(
       this,
-      predicate,
+      this._makePredicateString(predicate, queryHandler.queryOne),
-      queryHandler.queryOne,
       title,
       polling,
       timeout,
       updatedSelector,
-      isXPath,
       waitForVisible,
       waitForHidden
     );
-    const handle = await waitTask.promise;
+    const jsHandle = await waitTask.promise;
-    if (!handle.asElement()) {
+    const elementHandle = jsHandle.asElement();
-      await handle.dispose();
+    if (!elementHandle) {
+      await jsHandle.dispose();
       return null;
     }
-    return handle.asElement();
+    return elementHandle;
+  }
 
+  async waitForXPath(
+    xpath: string,
+    options: WaitForSelectorOptions
+  ): Promise<ElementHandle | null> {
+    const {
+      visible: waitForVisible = false,
+      hidden: waitForHidden = false,
+      timeout = this._timeoutSettings.timeout(),
+    } = options;
+    const polling = waitForVisible || waitForHidden ? 'raf' : 'mutation';
+    const title = `XPath \`${xpath}\`${waitForHidden ? ' to be hidden' : ''}`;
     function predicate(
-      selectorOrXPath: string,
+      xpath: string,
-      isXPath: boolean,
+      waitForVisible: boolean,
+      waitForHidden: boolean
+    ): Node | null | boolean {
+      const node = document.evaluate(
+        xpath,
+        document,
+        null,
+        XPathResult.FIRST_ORDERED_NODE_TYPE,
+        null
+      ).singleNodeValue;
+      return checkWaitForOptions(node, waitForVisible, waitForHidden);
+    }
+    const waitTask = new WaitTask(
+      this,
+      this._makePredicateString(predicate),
+      title,
+      polling,
+      timeout,
+      xpath,
+      waitForVisible,
+      waitForHidden
+    );
+    const jsHandle = await waitTask.promise;
+    const elementHandle = jsHandle.asElement();
+    if (!elementHandle) {
+      await jsHandle.dispose();
+      return null;
+    }
+    return elementHandle;
+  }
+
+  private _makePredicateString(
+    predicate: Function,
+    predicateQueryHandler?: Function
+  ): string {
+    const predicateQueryHandlerDef = predicateQueryHandler
+      ? `const predicateQueryHandler = ${predicateQueryHandler};`
+      : '';
+    return `
+    (() => {
+      ${predicateQueryHandlerDef}
+      const checkWaitForOptions = ${checkWaitForOptions};
+      return (${predicate})(...args)
+    })() `;
+    function checkWaitForOptions(
+      node: Node,
       waitForVisible: boolean,
       waitForHidden: boolean
     ): Node | null | boolean {
-      const node = isXPath
-        ? document.evaluate(
-            selectorOrXPath,
-            document,
-            null,
-            XPathResult.FIRST_ORDERED_NODE_TYPE,
-            null
-          ).singleNodeValue
-        : predicateQueryHandler
-        ? (predicateQueryHandler(document, selectorOrXPath) as Element)
-        : document.querySelector(selectorOrXPath);
       if (!node) return waitForHidden;
       if (!waitForVisible && !waitForHidden) return node;
       const element =
@@ -568,6 +589,29 @@ export class DOMWorld {
       }
     }
   }
+
+  waitForFunction(
+    pageFunction: Function | string,
+    options: { polling?: string | number; timeout?: number } = {},
+    ...args: SerializableOrJSHandle[]
+  ): Promise<JSHandle> {
+    const {
+      polling = 'raf',
+      timeout = this._timeoutSettings.timeout(),
+    } = options;
+    return new WaitTask(
+      this,
+      pageFunction,
+      'function',
+      polling,
+      timeout,
+      ...args
+    ).promise;
+  }
+
+  async title(): Promise<string> {
+    return this.evaluate(() => document.title);
+  }
 }
 
 class WaitTask {
@@ -586,7 +630,6 @@ class WaitTask {
   constructor(
     domWorld: DOMWorld,
     predicateBody: Function | string,
-    predicateQueryHandlerBody: Function | string | undefined,
     title: string,
     polling: string | number,
     timeout: number,
@@ -601,28 +644,15 @@ class WaitTask {
       assert(polling > 0, 'Cannot poll with non-positive interval: ' + polling);
     else throw new Error('Unknown polling options: ' + polling);
 
-    function getPredicateBody(
+    function getPredicateBody(predicateBody: Function | string) {
-      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 = getPredicateBody(
+    this._predicateBody = getPredicateBody(predicateBody);
-      predicateBody,
-      predicateQueryHandlerBody
-    );
     this._args = args;
     this._runCount = 0;
     domWorld._waitTasks.add(this);

test/waittask.spec.ts

@@ -608,7 +608,7 @@ describe('waittask specs', function () {
         .catch((error_) => (error = error_));
       expect(error).toBeTruthy();
       expect(error.message).toContain(
-        'waiting for selector "div" failed: timeout'
+        'waiting for selector `div` failed: timeout'
       );
       expect(error).toBeInstanceOf(puppeteer.errors.TimeoutError);
     });
@@ -622,7 +622,7 @@ describe('waittask specs', function () {
         .catch((error_) => (error = error_));
       expect(error).toBeTruthy();
       expect(error.message).toContain(
-        'waiting for selector "div" to be hidden failed: timeout'
+        'waiting for selector `div` to be hidden failed: timeout'
       );
     });
 
@@ -659,7 +659,7 @@ describe('waittask specs', function () {
       await page
         .waitForSelector('.zombo', { timeout: 10 })
         .catch((error_) => (error = error_));
-      expect(error.stack).toContain('waiting for selector ".zombo" failed');
+      expect(error.stack).toContain('waiting for selector `.zombo` failed');
       // The extension is ts here as Mocha maps back via sourcemaps.
       expect(error.stack).toContain('waittask.spec.ts');
     });
@@ -692,7 +692,7 @@ describe('waittask specs', function () {
         .catch((error_) => (error = error_));
       expect(error).toBeTruthy();
       expect(error.message).toContain(
-        'waiting for XPath "//div" failed: timeout'
+        'waiting for XPath `//div` failed: timeout'
       );
       expect(error).toBeInstanceOf(puppeteer.errors.TimeoutError);
     });