[rb] Support using custom element classes

rb/lib/selenium/webdriver/remote/bridge.rb

@@ -33,7 +33,7 @@ class Bridge
 
         class << self
           attr_reader :extra_commands
-          attr_writer :locator_converter
+          attr_writer :element_class, :locator_converter
 
           def add_command(name, verb, url, &block)
             @extra_commands ||= {}
@@ -44,6 +44,10 @@ def add_command(name, verb, url, &block)
           def locator_converter
             @locator_converter ||= LocatorConverter.new
           end
+
+          def element_class
+            @element_class ||= Element
+          end
         end
 
         #
@@ -432,7 +436,7 @@ def submit_element(element)
                    "e.initEvent('submit', true, true);\n" \
                    "if (form.dispatchEvent(e)) { HTMLFormElement.prototype.submit.call(form) }\n"
 
-          execute_script(script, Element::ELEMENT_KEY => element)
+          execute_script(script, Bridge.element_class::ELEMENT_KEY => element)
         rescue Error::JavascriptError
           raise Error::UnsupportedOperationError, 'To submit an element, it must be nested inside a form element'
         end
@@ -519,7 +523,7 @@ def element_value_of_css_property(element, prop)
         #
 
         def active_element
-          Element.new self, element_id_from(execute(:get_active_element))
+          Bridge.element_class.new self, element_id_from(execute(:get_active_element))
         end
 
         alias switch_to_active_element active_element
@@ -539,7 +543,7 @@ def find_element_by(how, what, parent_ref = [])
                  execute :find_element, {}, {using: how, value: what.to_s}
                end
 
-          Element.new self, element_id_from(id)
+          Bridge.element_class.new self, element_id_from(id)
         end
 
         def find_elements_by(how, what, parent_ref = [])
@@ -557,7 +561,7 @@ def find_elements_by(how, what, parent_ref = [])
                   execute :find_elements, {}, {using: how, value: what.to_s}
                 end
 
-          ids.map { |id| Element.new self, element_id_from(id) }
+          ids.map { |id| Bridge.element_class.new self, element_id_from(id) }
         end
 
         def shadow_root(element)
@@ -631,7 +635,7 @@ def escaper
         end
 
         def commands(command)
-          command_list[command]|| Bridge.extra_commands[command]
+          command_list[command] || Bridge.extra_commands[command]
         end
 
         def unwrap_script_result(arg)
@@ -640,7 +644,7 @@ def unwrap_script_result(arg)
             arg.map { |e| unwrap_script_result(e) }
           when Hash
             element_id = element_id_from(arg)
-            return Element.new(self, element_id) if element_id
+            return Bridge.element_class.new(self, element_id) if element_id
 
             shadow_root_id = shadow_root_id_from(arg)
             return ShadowRoot.new self, shadow_root_id if shadow_root_id
@@ -652,7 +656,7 @@ def unwrap_script_result(arg)
         end
 
         def element_id_from(id)
-          id['ELEMENT'] || id[Element::ELEMENT_KEY]
+          id['ELEMENT'] || id[Bridge.element_class::ELEMENT_KEY]
         end
 
         def shadow_root_id_from(id)
@@ -663,7 +667,6 @@ def prepare_capabilities_payload(capabilities)
           capabilities = {alwaysMatch: capabilities} if !capabilities['alwaysMatch'] && !capabilities['firstMatch']
           {capabilities: capabilities}
         end
-
       end # Bridge
     end # Remote
   end # WebDriver

rb/spec/unit/selenium/webdriver/remote/bridge_spec.rb

@@ -141,6 +141,72 @@ module Remote
             expect { bridge.quit }.not_to raise_error
           end
         end
+
+        describe 'finding elements' do
+          let(:http) { WebDriver::Remote::Http::Default.new }
+          let(:bridge) { described_class.new(http_client: http, url: 'http://localhost') }
+
+          before do
+            allow(http).to receive(:request)
+              .with(:post, URI('http://localhost/session'), any_args)
+              .and_return('status' => 200, 'value' => {'sessionId' => 'foo', 'capabilities' => {}})
+            bridge.create_session({})
+          end
+
+          describe '#find_element_by' do
+            before do
+              allow(http).to receive(:request)
+                .with(:post, URI('http://localhost/session/foo/element'), any_args)
+                .and_return('status' => 200, 'value' => {Element::ELEMENT_KEY => 'bar'})
+            end
+
+            it 'returns an element' do
+              expect(bridge.find_element_by(:id, 'test', nil)).to be_an_instance_of(Element)
+            end
+
+            context 'when custom element class is used' do
+              before do
+                stub_const('MyCustomElement', Class.new(Selenium::WebDriver::Element))
+                described_class.element_class = MyCustomElement
+              end
+
+              after do
+                described_class.element_class = nil
+              end
+
+              it 'returns a custom element' do
+                expect(bridge.find_element_by(:id, 'test', nil)).to be_an_instance_of(MyCustomElement)
+              end
+            end
+          end
+
+          describe '#find_elements_by' do
+            before do
+              allow(http).to receive(:request)
+                .with(:post, URI('http://localhost/session/foo/elements'), any_args)
+                .and_return('status' => 200, 'value' => [{Element::ELEMENT_KEY => 'bar'}])
+            end
+
+            it 'returns an element' do
+              expect(bridge.find_elements_by(:id, 'test', nil)).to all(be_an_instance_of(Element))
+            end
+
+            context 'when custom element class is used' do
+              before do
+                stub_const('MyCustomElement', Class.new(Selenium::WebDriver::Element))
+                described_class.element_class = MyCustomElement
+              end
+
+              after do
+                described_class.element_class = nil
+              end
+
+              it 'returns a custom element' do
+                expect(bridge.find_elements_by(:id, 'test', nil)).to all(be_an_instance_of(MyCustomElement))
+              end
+            end
+          end
+        end
       end
     end # Remote
   end # WebDriver