From fa75398921f827e19dfd496112cd80bf79f71ab3 Mon Sep 17 00:00:00 2001
From: Thomas Walpole <twalpole@gmail.com>
Date: Sat, 28 Jan 2017 14:34:44 -0800
Subject: [PATCH] support elements in evaluate_script result structures

---
 lib/capybara/selenium/driver.rb               | 16 ++++++++++++--
 lib/capybara/session.rb                       | 16 ++++++++++++--
 .../spec/session/evaluate_script_spec.rb      |  7 ++++++
 spec/shared_selenium_session.rb               | 22 ++++++++++++++++++-
 4 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/lib/capybara/selenium/driver.rb b/lib/capybara/selenium/driver.rb
index 5e6bca20..54864c0d 100644
--- a/lib/capybara/selenium/driver.rb
+++ b/lib/capybara/selenium/driver.rb
@@ -92,8 +92,7 @@ class Capybara::Selenium::Driver < Capybara::Driver::Base
 
   def evaluate_script(script, *args)
     result = browser.execute_script("return #{script}", *args.map { |arg| arg.is_a?(Capybara::Selenium::Node) ?  arg.native : arg} )
-    result = Capybara::Selenium::Node.new(self, result) if result.is_a? Selenium::WebDriver::Element
-    result
+    unwrap_script_result(result)
   end
 
   def save_screenshot(path, _options={})
@@ -327,4 +326,17 @@ class Capybara::Selenium::Driver < Capybara::Driver::Base
   def silenced_unknown_error_messages
     [ /Error communicating with the remote browser/ ]
   end
+
+  def unwrap_script_result(arg)
+    case arg
+    when Array
+      arg.map { |e| unwrap_script_result(e) }
+    when Hash
+      arg.each { |k, v| arg[k] = unwrap_script_result(v) }
+    when Selenium::WebDriver::Element
+      Capybara::Selenium::Node.new(self, arg)
+    else
+      arg
+    end
+  end
 end
diff --git a/lib/capybara/session.rb b/lib/capybara/session.rb
index fec6a125..785db126 100644
--- a/lib/capybara/session.rb
+++ b/lib/capybara/session.rb
@@ -637,8 +637,7 @@ module Capybara
       else
         driver.evaluate_script(script, *args.map { |arg| arg.is_a?(Capybara::Node::Element) ?  arg.base : arg} )
       end
-      result = Capybara::Node::Element.new(self, result, nil, nil) if result.is_a?(Capybara::Driver::Node)
-      result
+      element_script_result(result)
     end
 
     ##
@@ -844,5 +843,18 @@ module Capybara
     def scopes
       @scopes ||= [nil]
     end
+
+    def element_script_result(arg)
+      case arg
+      when Array
+        arg.map { |e| element_script_result(e) }
+      when Hash
+        arg.each { |k, v| arg[k] = element_script_result(v) }
+      when Capybara::Driver::Node
+        Capybara::Node::Element.new(self, arg, nil, nil)
+      else
+        arg
+      end
+    end
   end
 end
diff --git a/lib/capybara/spec/session/evaluate_script_spec.rb b/lib/capybara/spec/session/evaluate_script_spec.rb
index fdf4be60..e4e96c1f 100644
--- a/lib/capybara/spec/session/evaluate_script_spec.rb
+++ b/lib/capybara/spec/session/evaluate_script_spec.rb
@@ -18,4 +18,11 @@ Capybara::SpecHelper.spec "#evaluate_script", requires: [:js] do
     expect(@session).to have_css('#change', text: 'Doodle Funk')
   end
 
+  it "should support returning elements", requires: [:js, :es_args] do
+    @session.visit('/with_js')
+    el = @session.find(:css, '#change')
+    el = @session.evaluate_script("document.getElementById('change')")
+    expect(el).to be_instance_of(Capybara::Node::Element)
+    expect(el).to eq(@session.find(:css, '#change'))
+  end
 end
diff --git a/spec/shared_selenium_session.rb b/spec/shared_selenium_session.rb
index 5c46add3..11c4cede 100644
--- a/spec/shared_selenium_session.rb
+++ b/spec/shared_selenium_session.rb
@@ -121,11 +121,31 @@ RSpec.shared_examples "Capybara::Session" do |session, mode|
     end
 
     describe "#evaluate_script" do
-      it "can return elements" do
+      it "can return an element" do
         @session.visit('/form')
         element = @session.evaluate_script("document.getElementById('form_title')")
         expect(element).to eq @session.find(:id, 'form_title')
       end
+
+      it "can return arrays of nested elements" do
+        @session.visit('/form')
+        elements = @session.evaluate_script('document.querySelectorAll("#form_city option")')
+        elements.each do |el|
+          expect(el).to be_instance_of Capybara::Node::Element
+        end
+        expect(elements).to eq @session.find(:css, '#form_city').all(:css, 'option').to_a
+      end
+
+      it "can return hashes with elements" do
+        @session.visit('/form')
+        result = @session.evaluate_script("{ a: document.getElementById('form_title'), b: {c: document.querySelectorAll('#form_city option')}}")
+        expect(result).to eq({
+          'a' => @session.find(:id, 'form_title'),
+          'b' => {
+            'c' => @session.find(:css, '#form_city').all(:css, 'option').to_a
+          }
+        })
+      end
     end
   end
 end