From c01cdac12ca8495cfe47a4ddb0afc172cc58fbcd Mon Sep 17 00:00:00 2001
From: Alex Rudenko <alexrudenko@chromium.org>
Date: Tue, 14 Dec 2021 20:57:53 +0100
Subject: [PATCH] fix: make sure ElementHandle.waitForSelector is evaluated in
 the right context

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         |  3 ++-
 src/common/JSHandle.ts         | 15 ++++++++++++---
 test/ariaqueryhandler.spec.ts  | 10 ++++++++++
 4 files changed, 26 insertions(+), 6 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 54b5dc206418c..14d23d0418f13 100644
--- a/src/common/DOMWorld.ts
+++ b/src/common/DOMWorld.ts
@@ -832,8 +832,9 @@ export class WaitTask {
     if (this._terminated || runCount !== this._runCount) return;
     try {
       if (this._root) {
-        success = await this._root.evaluateHandle(
+        success = await context.evaluateHandle(
           waitForPredicatePageFunction,
+          this._root,
           this._predicateBody,
           this._polling,
           this._timeout,
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);