From 75c827fe9b03f2d01a1c81a8cf04a7167ad209f2 Mon Sep 17 00:00:00 2001
From: YusukeIwaki <iwaki@i3-systems.com>
Date: Sat, 1 Jan 2022 22:35:10 +0900
Subject: [PATCH] fix: make sure ElementHandle.waitForSelector is evaluated in
 the right context

---
 lib/puppeteer/aria_query_handler.rb         |  8 ++++---
 lib/puppeteer/element_handle.rb             | 13 ++++++++++-
 lib/puppeteer/wait_task.rb                  | 26 +++++++--------------
 spec/integration/aria_query_handler_spec.rb | 19 ++++++++++++---
 4 files changed, 41 insertions(+), 25 deletions(-)

diff --git a/lib/puppeteer/aria_query_handler.rb b/lib/puppeteer/aria_query_handler.rb
index d67f0100..bde3c1bb 100644
--- a/lib/puppeteer/aria_query_handler.rb
+++ b/lib/puppeteer/aria_query_handler.rb
@@ -36,11 +36,11 @@ def query_one(element, selector)
     end
   end
 
-  def wait_for(dom_world, selector, visible: nil, hidden: nil, timeout: nil)
+  def wait_for(dom_world, selector, visible: nil, hidden: nil, timeout: nil, root: nil)
     # addHandlerToWorld
     binding_function = Puppeteer::DOMWorld::BindingFunction.new(
       name: 'ariaQuerySelector',
-      proc: -> (sel) { query_one(dom_world.send(:document), sel) },
+      proc: -> (sel) { query_one(root || dom_world.send(:document), sel) },
     )
     dom_world.send(:wait_for_selector_in_page,
       '(_, selector) => globalThis.ariaQuerySelector(selector)',
@@ -48,7 +48,9 @@ def wait_for(dom_world, selector, visible: nil, hidden: nil, timeout: nil)
       visible: visible,
       hidden: hidden,
       timeout: timeout,
-      binding_function: binding_function)
+      binding_function: binding_function,
+      root: root,
+    )
   end
 
   def query_all(element, selector)
diff --git a/lib/puppeteer/element_handle.rb b/lib/puppeteer/element_handle.rb
index b15a3cf2..1bab748e 100644
--- a/lib/puppeteer/element_handle.rb
+++ b/lib/puppeteer/element_handle.rb
@@ -58,7 +58,18 @@ def inspect
   # (30 seconds). Pass `0` to disable timeout. The default value can be changed
   # by using the {@link Page.setDefaultTimeout} method.
   def wait_for_selector(selector, visible: nil, hidden: nil, timeout: nil)
-    @context.world.wait_for_selector(selector, visible: visible, hidden: hidden, timeout: timeout, root: self)
+    frame = @context.frame
+
+    secondary_world = frame.secondary_world
+    adopted_root = secondary_world.execution_context.adopt_element_handle(self)
+    handle = secondary_world.wait_for_selector(selector, visible: visible, hidden: hidden, timeout: timeout, root: adopted_root)
+    adopted_root.dispose
+    return nil unless handle
+
+    main_world = frame.main_world
+    result = main_world.execution_context.adopt_element_handle(handle)
+    handle.dispose
+    result
   end
 
   define_async_method :async_wait_for_selector
diff --git a/lib/puppeteer/wait_task.rb b/lib/puppeteer/wait_task.rb
index e1eb8fc3..4a36b027 100644
--- a/lib/puppeteer/wait_task.rb
+++ b/lib/puppeteer/wait_task.rb
@@ -67,24 +67,14 @@ def rerun
     return if @terminated || run_count != @run_count
 
     begin
-      if @root
-        success = @root.evaluate_handle(
-          WAIT_FOR_PREDICATE_PAGE_FUNCTION,
-          @predicate_body,
-          @polling,
-          @timeout,
-          *@args,
-        )
-      else
-        success = context.evaluate_handle(
-          WAIT_FOR_PREDICATE_PAGE_FUNCTION,
-          nil,
-          @predicate_body,
-          @polling,
-          @timeout,
-          *@args,
-        )
-      end
+      success = context.evaluate_handle(
+        WAIT_FOR_PREDICATE_PAGE_FUNCTION,
+        @root,
+        @predicate_body,
+        @polling,
+        @timeout,
+        *@args,
+      )
     rescue => err
       error = err
     end
diff --git a/spec/integration/aria_query_handler_spec.rb b/spec/integration/aria_query_handler_spec.rb
index e10cbf8c..428d14b9 100644
--- a/spec/integration/aria_query_handler_spec.rb
+++ b/spec/integration/aria_query_handler_spec.rb
@@ -93,6 +93,19 @@
       Timeout.timeout(1) { page.wait_for_selector('aria/[role="button"]') }
     end
 
+    it 'should return the element handle' do
+      page.evaluate(<<~JAVASCRIPT)
+        () => (document.body.innerHTML = `<div></div>`)
+      JAVASCRIPT
+      element = page.query_selector('div')
+      inner_element = element.wait_for_selector('aria/test') do
+        element.evaluate(<<~JAVASCRIPT)
+          el => el.innerHTML="<p><button>test</button></p>"
+        JAVASCRIPT
+      end
+      expect(inner_element.evaluate('el => el.outerHTML')).to eq('<button>test</button>')
+    end
+
     it 'should persist query handler bindings across reloads' do
       page.goto(server_empty_page)
       page.evaluate(add_element, 'button')
@@ -372,9 +385,9 @@
     # });
 
     it 'should return the element handle' do
-      promise = page.async_wait_for_selector('aria/zombo')
-      page.content = "<div aria-label='zombo'>anything</div>"
-      result = await promise
+      result = page.wait_for_selector('aria/zombo') do
+        page.content = "<div aria-label='zombo'>anything</div>"
+      end
       expect(result).to be_a(Puppeteer::ElementHandle)
       expect(page.evaluate('(x) => x.textContent', result)).to eq('anything')
     end