fix: make sure ElementHandle.waitForSelector is evaluated in the right context

src/common/AriaQueryHandler.ts

@@ -92,8 +92,8 @@ const waitFor = async (
   const binding: PageBinding = {
     name: 'ariaQuerySelector',
     pptrFunction: async (selector: string) => {
-      const document = await domWorld._document();
-      const element = await queryOne(document, selector);
+      const root = options.root || (await domWorld._document());
+      const element = await queryOne(root, selector);
       return element;
     },
   };

src/common/DOMWorld.ts

@@ -766,7 +766,7 @@ export class WaitTask {
   _reject: (x: Error) => void;
   _timeoutTimer?: NodeJS.Timeout;
   _terminated = false;
-  _root: ElementHandle;
+  _root: ElementHandle = null;
 
   constructor(options: WaitTaskOptions) {
     if (helper.isString(options.polling))
@@ -838,26 +838,15 @@ export class WaitTask {
     }
     if (this._terminated || runCount !== this._runCount) return;
     try {
-      if (this._root) {
-        success = await this._root.evaluateHandle(
-          waitForPredicatePageFunction,
-          this._predicateBody,
-          this._predicateAcceptsContextElement,
-          this._polling,
-          this._timeout,
-          ...this._args
-        );
-      } else {
-        success = await context.evaluateHandle(
-          waitForPredicatePageFunction,
-          null,
-          this._predicateBody,
-          this._predicateAcceptsContextElement,
-          this._polling,
-          this._timeout,
-          ...this._args
-        );
-      }
+      success = await context.evaluateHandle(
+        waitForPredicatePageFunction,
+        this._root || null,
+        this._predicateBody,
+        this._predicateAcceptsContextElement,
+        this._polling,
+        this._timeout,
+        ...this._args
+      );
     } catch (error_) {
       error = error_;
     }

src/common/JSHandle.ts

@@ -381,18 +381,27 @@ export class ElementHandle<
    * (30 seconds). Pass `0` to disable timeout. The default value can be changed
    * by using the {@link Page.setDefaultTimeout} method.
    */
-  waitForSelector(
+  async waitForSelector(
     selector: string,
     options: {
       visible?: boolean;
       hidden?: boolean;
       timeout?: number;
     } = {}
   ): Promise<ElementHandle | null> {
-    return this._context._world.waitForSelector(selector, {
+    const frame = this._context.frame();
+    const secondaryContext = await frame._secondaryWorld.executionContext();
+    const adoptedRoot = await secondaryContext._adoptElementHandle(this);
+    const handle = await frame._secondaryWorld.waitForSelector(selector, {
       ...options,
-      root: this,
+      root: adoptedRoot,
     });
+    await adoptedRoot.dispose();
+    if (!handle) return null;
+    const mainExecutionContext = await frame._mainWorld.executionContext();
+    const result = await mainExecutionContext._adoptElementHandle(handle);
+    await handle.dispose();
+    return result;
   }
 
   asElement(): ElementHandle<ElementType> | null {

test/ariaqueryhandler.spec.ts

@@ -196,6 +196,16 @@ describeChromeOnly('AriaQueryHandler', () => {
       await page.waitForSelector('aria/[role="button"]');
     });
 
+    it('should work for ElementHandler.waitForSelector', async () => {
+      const { page, server } = getTestState();
+      await page.goto(server.EMPTY_PAGE);
+      await page.evaluate(
+        () => (document.body.innerHTML = `<div><button>test</button></div>`)
+      );
+      const element = await page.$('div');
+      await element.waitForSelector('aria/test');
+    });
+
     it('should persist query handler bindings across reloads', async () => {
       const { page, server } = getTestState();
       await page.goto(server.EMPTY_PAGE);