fix: make sure ElementHandle.waitForSelector is evaluated in the right context (#7843)

So it appears that all bindings are added to the secondary world and all
evaluations are also running there. ElementHandle.evaluate is returning
handles from the main world though. Therefore, we need to be careful
and adopt handles to the right context before doing waitForSelector
So it appears that all bindings are added to the secondary world and all
evaluations are also running there. ElementHandle.evaluate is returning
handles from the main world though. Therefore, we need to be careful
and adopt handles to the right context before doing waitForSelector.

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,7 +381,7 @@ 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;
@@ -389,10 +389,19 @@ export class ElementHandle<
       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);