From 11d99b84c832c7e63b2856fac6623eb1a116323d Mon Sep 17 00:00:00 2001
From: Thomas Walpole <twalpole@gmail.com>
Date: Fri, 27 Jan 2017 11:11:13 -0800
Subject: [PATCH] Support returning an element from evaluate_script with
 selenium

---
 lib/capybara/selenium/driver.rb | 4 +++-
 lib/capybara/session.rb         | 4 +++-
 spec/shared_selenium_session.rb | 8 ++++++++
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/lib/capybara/selenium/driver.rb b/lib/capybara/selenium/driver.rb
index ca1b62be..5e6bca20 100644
--- a/lib/capybara/selenium/driver.rb
+++ b/lib/capybara/selenium/driver.rb
@@ -91,7 +91,9 @@ class Capybara::Selenium::Driver < Capybara::Driver::Base
   end
 
   def evaluate_script(script, *args)
-    browser.execute_script("return #{script}", *args.map { |arg| arg.is_a?(Capybara::Selenium::Node) ?  arg.native : arg} )
+    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
   end
 
   def save_screenshot(path, _options={})
diff --git a/lib/capybara/session.rb b/lib/capybara/session.rb
index b854ede3..fec6a125 100644
--- a/lib/capybara/session.rb
+++ b/lib/capybara/session.rb
@@ -631,12 +631,14 @@ module Capybara
     #
     def evaluate_script(script, *args)
       @touched = true
-      if driver.method(:evaluate_script).arity == 1
+      result = if driver.method(:evaluate_script).arity == 1
         raise Capybara::NotSupportedByDriverError, "The current driver does not support arguments being passed with execute_script" unless args.empty?
         driver.evaluate_script(script)
       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
     end
 
     ##
diff --git a/spec/shared_selenium_session.rb b/spec/shared_selenium_session.rb
index 453a201b..5c46add3 100644
--- a/spec/shared_selenium_session.rb
+++ b/spec/shared_selenium_session.rb
@@ -119,5 +119,13 @@ RSpec.shared_examples "Capybara::Session" do |session, mode|
         expect(element.path).to eq('/html/body/div[2]/a[1]')
       end
     end
+
+    describe "#evaluate_script" do
+      it "can return elements" do
+        @session.visit('/form')
+        element = @session.evaluate_script("document.getElementById('form_title')")
+        expect(element).to eq @session.find(:id, 'form_title')
+      end
+    end
   end
 end