From 8d8e874b072b17fc763f33d08e51c046b7435244 Mon Sep 17 00:00:00 2001
From: Alex Rudenko <OrKoN@users.noreply.github.com>
Date: Tue, 21 Dec 2021 09:53:20 +0100
Subject: [PATCH] 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 |  4 ++--
 src/common/DOMWorld.ts         | 31 ++++++++++---------------------
 src/common/JSHandle.ts         | 15 ++++++++++++---
 test/ariaqueryhandler.spec.ts  | 10 ++++++++++
 4 files changed, 34 insertions(+), 26 deletions(-)

diff --git a/src/common/AriaQueryHandler.ts b/src/common/AriaQueryHandler.ts
index cb3fd341dcde6..fd06cfe2b2114 100644
--- a/src/common/AriaQueryHandler.ts
+++ b/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;
     },
   };
diff --git a/src/common/DOMWorld.ts b/src/common/DOMWorld.ts
index c9adf8ff68269..cfacc28734666 100644
--- a/src/common/DOMWorld.ts
+++ b/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_;
     }
diff --git a/src/common/JSHandle.ts b/src/common/JSHandle.ts
index 43b33235377ab..e6358fba96ef6 100644
--- a/src/common/JSHandle.ts
+++ b/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 {
diff --git a/test/ariaqueryhandler.spec.ts b/test/ariaqueryhandler.spec.ts
index 2773520853716..504706af88a2d 100644
--- a/test/ariaqueryhandler.spec.ts
+++ b/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);