feat(webdriver): support ARIA selectors (#12315)

Co-authored-by: Nikolay Vitkov <34244704+Lightning00Blade@users.noreply.github.com>

packages/puppeteer-core/src/api/ElementHandle.ts

@@ -10,6 +10,7 @@ import type {Frame} from '../api/Frame.js';
 import {getQueryHandlerAndSelector} from '../common/GetQueryHandler.js';
 import {LazyArg} from '../common/LazyArg.js';
 import type {
+  AwaitableIterable,
   ElementFor,
   EvaluateFuncWith,
   HandleFor,
@@ -873,6 +874,14 @@ export abstract class ElementHandle<
     ...paths: string[]
   ): Promise<void>;
 
+  /**
+   * @internal
+   */
+  abstract queryAXTree(
+    name?: string,
+    role?: string
+  ): AwaitableIterable<ElementHandle<Node>>;
+
   /**
    * This method scrolls element into view if needed, and then uses
    * {@link Touchscreen.tap} to tap in the center of the element.

packages/puppeteer-core/src/bidi/ElementHandle.ts

@@ -7,6 +7,8 @@
 import type * as Bidi from 'chromium-bidi/lib/cjs/protocol/protocol.js';
 
 import {ElementHandle, type AutofillData} from '../api/ElementHandle.js';
+import type {AwaitableIterable} from '../common/types.js';
+import {AsyncIterableUtil} from '../util/AsyncIterableUtil.js';
 import {throwIfDisposed} from '../util/decorators.js';
 
 import type {BidiFrame} from './Frame.js';
@@ -116,4 +118,23 @@ export class BidiElementHandle<
     });
     await this.frame.setFiles(this, files);
   }
+
+  override async *queryAXTree(
+    this: BidiElementHandle<HTMLElement>,
+    name?: string | undefined,
+    role?: string | undefined
+  ): AwaitableIterable<ElementHandle<Node>> {
+    const results = await this.frame.locateNodes(this, {
+      type: 'accessibility',
+      value: {
+        role,
+        name,
+      },
+    });
+
+    return yield* AsyncIterableUtil.map(results, node => {
+      // TODO: maybe change ownership since the default ownership is probably none.
+      return Promise.resolve(BidiElementHandle.from(node, this.realm));
+    });
+  }
 }

packages/puppeteer-core/src/bidi/Frame.ts

@@ -561,6 +561,18 @@ export class BidiFrame extends Frame {
       files
     );
   }
+
+  @throwIfDetached
+  async locateNodes(
+    element: BidiElementHandle,
+    locator: Bidi.BrowsingContext.Locator
+  ): Promise<Bidi.Script.NodeRemoteValue[]> {
+    return await this.browsingContext.locateNodes(
+      locator,
+      // SAFETY: ElementHandles are always remote references.
+      [element.remoteValue() as Bidi.Script.SharedReference]
+    );
+  }
 }
 
 function isConsoleLogEntry(

packages/puppeteer-core/src/bidi/core/BrowsingContext.ts

@@ -601,4 +601,21 @@ export class BrowsingContext extends EventEmitter<{
       })
     );
   }
+
+  @throwIfDisposed<BrowsingContext>(context => {
+    // SAFETY: Disposal implies this exists.
+    return context.#reason!;
+  })
+  async locateNodes(
+    locator: Bidi.BrowsingContext.Locator,
+    startNodes: [Bidi.Script.SharedReference, ...Bidi.Script.SharedReference[]]
+  ): Promise<Bidi.Script.NodeRemoteValue[]> {
+    // TODO: add other locateNodes options if needed.
+    const result = await this.#session.send('browsingContext.locateNodes', {
+      context: this.id,
+      locator,
+      startNodes: startNodes.length ? startNodes : undefined,
+    });
+    return result.result.nodes;
+  }
 }

packages/puppeteer-core/src/bidi/core/Connection.ts

@@ -69,6 +69,10 @@ export interface Commands {
     params: Bidi.BrowsingContext.GetTreeParameters;
     returnType: Bidi.BrowsingContext.GetTreeResult;
   };
+  'browsingContext.locateNodes': {
+    params: Bidi.BrowsingContext.LocateNodesParameters;
+    returnType: Bidi.BrowsingContext.LocateNodesResult;
+  };
   'browsingContext.navigate': {
     params: Bidi.BrowsingContext.NavigateParameters;
     returnType: Bidi.BrowsingContext.NavigateResult;

packages/puppeteer-core/src/cdp/AriaQueryHandler.ts

@@ -4,42 +4,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type {Protocol} from 'devtools-protocol';
-
-import type {CDPSession} from '../api/CDPSession.js';
 import type {ElementHandle} from '../api/ElementHandle.js';
 import {QueryHandler, type QuerySelector} from '../common/QueryHandler.js';
 import type {AwaitableIterable} from '../common/types.js';
 import {assert} from '../util/assert.js';
 import {AsyncIterableUtil} from '../util/AsyncIterableUtil.js';
 
-const NON_ELEMENT_NODE_ROLES = new Set(['StaticText', 'InlineTextBox']);
-
-const queryAXTree = async (
-  client: CDPSession,
-  element: ElementHandle<Node>,
-  accessibleName?: string,
-  role?: string
-): Promise<Protocol.Accessibility.AXNode[]> => {
-  const {nodes} = await client.send('Accessibility.queryAXTree', {
-    objectId: element.id,
-    accessibleName,
-    role,
-  });
-  return nodes.filter((node: Protocol.Accessibility.AXNode) => {
-    if (node.ignored) {
-      return false;
-    }
-    if (!node.role) {
-      return false;
-    }
-    if (NON_ELEMENT_NODE_ROLES.has(node.role.value)) {
-      return false;
-    }
-    return true;
-  });
-};
-
 interface ARIASelector {
   name?: string;
   role?: string;
@@ -105,17 +75,7 @@ export class ARIAQueryHandler extends QueryHandler {
     selector: string
   ): AwaitableIterable<ElementHandle<Node>> {
     const {name, role} = parseARIASelector(selector);
-    const results = await queryAXTree(
-      element.realm.environment.client,
-      element,
-      name,
-      role
-    );
-    yield* AsyncIterableUtil.map(results, node => {
-      return element.realm.adoptBackendNode(node.backendDOMNodeId) as Promise<
-        ElementHandle<Node>
-      >;
-    });
+    yield* element.queryAXTree(name, role);
   }
 
   static override queryOne = async (

packages/puppeteer-core/src/cdp/ElementHandle.ts

@@ -10,8 +10,10 @@ import type {Protocol} from 'devtools-protocol';
 
 import type {CDPSession} from '../api/CDPSession.js';
 import {ElementHandle, type AutofillData} from '../api/ElementHandle.js';
+import type {AwaitableIterable} from '../common/types.js';
 import {debugError} from '../common/util.js';
 import {assert} from '../util/assert.js';
+import {AsyncIterableUtil} from '../util/AsyncIterableUtil.js';
 import {throwIfDisposed} from '../util/decorators.js';
 
 import type {CdpFrame} from './Frame.js';
@@ -19,6 +21,8 @@ import type {FrameManager} from './FrameManager.js';
 import type {IsolatedWorld} from './IsolatedWorld.js';
 import {CdpJSHandle} from './JSHandle.js';
 
+const NON_ELEMENT_NODE_ROLES = new Set(['StaticText', 'InlineTextBox']);
+
 /**
  * The CdpElementHandle extends ElementHandle now to keep compatibility
  * with `instanceof` because of that we need to have methods for
@@ -169,4 +173,34 @@ export class CdpElementHandle<
       card: data.creditCard,
     });
   }
+
+  override async *queryAXTree(
+    name?: string | undefined,
+    role?: string | undefined
+  ): AwaitableIterable<ElementHandle<Node>> {
+    const {nodes} = await this.client.send('Accessibility.queryAXTree', {
+      objectId: this.id,
+      accessibleName: name,
+      role,
+    });
+
+    const results = nodes.filter(node => {
+      if (node.ignored) {
+        return false;
+      }
+      if (!node.role) {
+        return false;
+      }
+      if (NON_ELEMENT_NODE_ROLES.has(node.role.value)) {
+        return false;
+      }
+      return true;
+    });
+
+    return yield* AsyncIterableUtil.map(results, node => {
+      return this.realm.adoptBackendNode(node.backendDOMNodeId) as Promise<
+        ElementHandle<Node>
+      >;
+    });
+  }
 }

test/TestExpectations.json

@@ -4,7 +4,7 @@
     "platforms": ["darwin", "linux", "win32"],
     "parameters": ["firefox"],
     "expectations": ["SKIP"],
-    "comment": "TODO: add a comment explaining why this expectation is required (include links to issues)"
+    "comment": "Firefox crashes when a document is provided as a start node"
   },
   {
     "testIdPattern": "[autofill.spec] *",
@@ -195,6 +195,13 @@
     "expectations": ["SKIP"],
     "comment": "TODO: add a comment explaining why this expectation is required (include links to issues)"
   },
+  {
+    "testIdPattern": "[ariaqueryhandler.spec] AriaQueryHandler queryOne (Chromium web test) should find by role \"heading\"",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["webDriverBiDi"],
+    "expectations": ["FAIL"],
+    "comment": "WebDriver BiDi locateNodes does not support shadow roots so far"
+  },
   {
     "testIdPattern": "[autofill.spec] *",
     "platforms": ["darwin", "linux", "win32"],
@@ -916,13 +923,6 @@
     "expectations": ["FAIL"],
     "comment": "Querying by a11y attributes is not standard behavior"
   },
-  {
-    "testIdPattern": "[ariaqueryhandler.spec] AriaQueryHandler queryOne (Chromium web test) should find by role \"heading\"",
-    "platforms": ["darwin", "linux", "win32"],
-    "parameters": ["firefox", "webDriverBiDi"],
-    "expectations": ["FAIL"],
-    "comment": "Querying by a11y attributes is not standard behavior"
-  },
   {
     "testIdPattern": "[bfcache.spec] BFCache can navigate to a BFCached page containing an OOPIF and a worker",
     "platforms": ["darwin", "linux", "win32"],