From c9b35131a125bc1a2510509665697bcb316b954e Mon Sep 17 00:00:00 2001
From: Alex Rodionov <p0deje@gmail.com>
Date: Sat, 1 Apr 2023 18:53:52 -0500
Subject: [PATCH 1/6] [rb] Support registering extra headers in HTTP client

---
 .../selenium/webdriver/remote/http/common.rb  |  6 ++++-
 .../webdriver/remote/http/common_spec.rb      | 23 +++++++++++++++----
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/rb/lib/selenium/webdriver/remote/http/common.rb b/rb/lib/selenium/webdriver/remote/http/common.rb
index 10ba4083caf73..3b46c14c9cf98 100644
--- a/rb/lib/selenium/webdriver/remote/http/common.rb
+++ b/rb/lib/selenium/webdriver/remote/http/common.rb
@@ -30,6 +30,10 @@ class Common
             'User-Agent' => "selenium/#{WebDriver::VERSION} (ruby #{Platform.os})"
           }.freeze
 
+          class << self
+            attr_accessor :extra_headers
+          end
+
           attr_writer :server_url
 
           def quit_errors
@@ -42,7 +46,7 @@ def close
 
           def call(verb, url, command_hash)
             url      = server_url.merge(url) unless url.is_a?(URI)
-            headers  = DEFAULT_HEADERS.dup
+            headers  = DEFAULT_HEADERS.merge(Common.extra_headers || {}).dup
             headers['Cache-Control'] = 'no-cache' if verb == :get
 
             if command_hash
diff --git a/rb/spec/unit/selenium/webdriver/remote/http/common_spec.rb b/rb/spec/unit/selenium/webdriver/remote/http/common_spec.rb
index bda3254e340d5..598b73e5d9634 100644
--- a/rb/spec/unit/selenium/webdriver/remote/http/common_spec.rb
+++ b/rb/spec/unit/selenium/webdriver/remote/http/common_spec.rb
@@ -24,11 +24,19 @@ module WebDriver
     module Remote
       module Http
         describe Common do
-          it 'sends non-empty body header for POST requests without command data' do
+          subject(:common) do
             common = described_class.new
             common.server_url = URI.parse('http://server')
             allow(common).to receive(:request)
 
+            common
+          end
+
+          after do
+            described_class.extra_headers = {}
+          end
+
+          it 'sends non-empty body header for POST requests without command data' do
             common.call(:post, 'clear', nil)
 
             expect(common).to have_received(:request)
@@ -37,10 +45,7 @@ module Http
           end
 
           it 'sends a standard User-Agent by default' do
-            common = described_class.new
-            common.server_url = URI.parse('http://server')
             user_agent_regexp = %r{\Aselenium/#{WebDriver::VERSION} \(ruby #{Platform.os}\)\z}
-            allow(common).to receive(:request)
 
             common.call(:post, 'session', nil)
 
@@ -48,6 +53,16 @@ module Http
               .with(:post, URI.parse('http://server/session'),
                     hash_including('User-Agent' => a_string_matching(user_agent_regexp)), '{}')
           end
+
+          it 'allows registering extra headers' do
+            described_class.extra_headers = {'Foo' => 'bar'}
+
+            common.call(:post, 'session', nil)
+
+            expect(common).to have_received(:request)
+              .with(:post, URI.parse('http://server/session'),
+                    hash_including('Foo' => 'bar'), '{}')
+          end
         end
       end # Http
     end # Remote

From 82e179de18e8a114d8b6686bb3e7f7c7df76239e Mon Sep 17 00:00:00 2001
From: Alex Rodionov <p0deje@gmail.com>
Date: Sat, 1 Apr 2023 19:02:49 -0500
Subject: [PATCH 2/6] [rb] Support overriding User-Agent in HTTP client

---
 .../selenium/webdriver/remote/http/common.rb  | 20 ++++++++++++++++---
 .../webdriver/remote/http/common_spec.rb      | 13 +++++++++++-
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/rb/lib/selenium/webdriver/remote/http/common.rb b/rb/lib/selenium/webdriver/remote/http/common.rb
index 3b46c14c9cf98..423dc5d7ece14 100644
--- a/rb/lib/selenium/webdriver/remote/http/common.rb
+++ b/rb/lib/selenium/webdriver/remote/http/common.rb
@@ -26,12 +26,16 @@ class Common
           CONTENT_TYPE    = 'application/json'
           DEFAULT_HEADERS = {
             'Accept' => CONTENT_TYPE,
-            'Content-Type' => "#{CONTENT_TYPE}; charset=UTF-8",
-            'User-Agent' => "selenium/#{WebDriver::VERSION} (ruby #{Platform.os})"
+            'Content-Type' => "#{CONTENT_TYPE}; charset=UTF-8"
           }.freeze
 
           class << self
             attr_accessor :extra_headers
+            attr_writer :user_agent
+
+            def user_agent
+              @user_agent ||= "selenium/#{WebDriver::VERSION} (ruby #{Platform.os})"
+            end
           end
 
           attr_writer :server_url
@@ -46,7 +50,7 @@ def close
 
           def call(verb, url, command_hash)
             url      = server_url.merge(url) unless url.is_a?(URI)
-            headers  = DEFAULT_HEADERS.merge(Common.extra_headers || {}).dup
+            headers  = common_headers.dup
             headers['Cache-Control'] = 'no-cache' if verb == :get
 
             if command_hash
@@ -65,6 +69,16 @@ def call(verb, url, command_hash)
 
           private
 
+          def common_headers
+            @common_headers ||= begin
+              headers = DEFAULT_HEADERS.dup
+              headers['User-Agent'] = Common.user_agent
+              headers = headers.merge(Common.extra_headers || {})
+
+              headers
+            end
+          end
+
           def server_url
             return @server_url if @server_url
 
diff --git a/rb/spec/unit/selenium/webdriver/remote/http/common_spec.rb b/rb/spec/unit/selenium/webdriver/remote/http/common_spec.rb
index 598b73e5d9634..6864a94e4d06e 100644
--- a/rb/spec/unit/selenium/webdriver/remote/http/common_spec.rb
+++ b/rb/spec/unit/selenium/webdriver/remote/http/common_spec.rb
@@ -33,7 +33,8 @@ module Http
           end
 
           after do
-            described_class.extra_headers = {}
+            described_class.extra_headers = nil
+            described_class.user_agent = nil
           end
 
           it 'sends non-empty body header for POST requests without command data' do
@@ -63,6 +64,16 @@ module Http
               .with(:post, URI.parse('http://server/session'),
                     hash_including('Foo' => 'bar'), '{}')
           end
+
+          it 'allows overriding default User-Agent' do
+            described_class.user_agent = 'rspec/1.0 (ruby 3.2)'
+
+            common.call(:post, 'session', nil)
+
+            expect(common).to have_received(:request)
+              .with(:post, URI.parse('http://server/session'),
+                    hash_including('User-Agent' => 'rspec/1.0 (ruby 3.2)'), '{}')
+          end
         end
       end # Http
     end # Remote

From e078018407cb00b7e851140f087703d90669e866 Mon Sep 17 00:00:00 2001
From: Alex Rodionov <p0deje@gmail.com>
Date: Sat, 1 Apr 2023 21:51:34 -0500
Subject: [PATCH 3/6] [rb] Support registering extra bridge commands

---
 rb/lib/selenium/webdriver/remote/bridge.rb    | 12 ++++++++-
 .../selenium/webdriver/remote/bridge_spec.rb  | 27 +++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/rb/lib/selenium/webdriver/remote/bridge.rb b/rb/lib/selenium/webdriver/remote/bridge.rb
index 58cc61385c0df..655d331e1fc57 100644
--- a/rb/lib/selenium/webdriver/remote/bridge.rb
+++ b/rb/lib/selenium/webdriver/remote/bridge.rb
@@ -29,6 +29,16 @@ class Bridge
         attr_accessor :http, :file_detector
         attr_reader :capabilities
 
+        class << self
+          attr_reader :extra_commands
+
+          def add_command(name, verb, url, &block)
+            @extra_commands ||= {}
+            @extra_commands[name] = [verb, url]
+            define_method(name, &block)
+          end
+        end
+
         #
         # Initializes the bridge with the given server URL
         # @param [String, URI] url url for the remote server
@@ -612,7 +622,7 @@ def escaper
         end
 
         def commands(command)
-          command_list[command]
+          command_list[command]|| Bridge.extra_commands[command]
         end
 
         def unwrap_script_result(arg)
diff --git a/rb/spec/unit/selenium/webdriver/remote/bridge_spec.rb b/rb/spec/unit/selenium/webdriver/remote/bridge_spec.rb
index 188e51b3f5dbd..cd8d7091561d0 100644
--- a/rb/spec/unit/selenium/webdriver/remote/bridge_spec.rb
+++ b/rb/spec/unit/selenium/webdriver/remote/bridge_spec.rb
@@ -23,6 +23,33 @@ module Selenium
   module WebDriver
     module Remote
       describe Bridge do
+        describe '.add_command' 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(any_args)
+              .and_return('status' => 200, 'value' => {'sessionId' => 'foo', 'capabilities' => {}})
+
+            bridge.create_session({})
+          end
+
+          after do
+            described_class.extra_commands.clear
+          end
+
+          it 'adds new command' do
+            described_class.add_command(:highlight, :get, 'session/:session_id/highlight/:id') do |element|
+              execute :highlight, id: element
+            end
+
+            bridge.highlight('bar')
+            expect(http).to have_received(:request)
+              .with(:get, URI('http://localhost/session/foo/highlight/bar'), any_args)
+          end
+        end
+
         describe '#initialize' do
           it 'raises ArgumentError if passed invalid options' do
             expect { described_class.new(foo: 'bar') }.to raise_error(ArgumentError)

From 724d7dc1558d37d5ea95cf0f3522a6420965d7a7 Mon Sep 17 00:00:00 2001
From: Alex Rodionov <p0deje@gmail.com>
Date: Sun, 4 Jun 2023 16:04:09 -0700
Subject: [PATCH 4/6] [rb] Support overriding default locator conversion

---
 rb/lib/selenium/webdriver/remote/bridge.rb    | 49 +++---------
 .../remote/bridge/locator_converter.rb        | 76 +++++++++++++++++++
 2 files changed, 87 insertions(+), 38 deletions(-)
 create mode 100644 rb/lib/selenium/webdriver/remote/bridge/locator_converter.rb

diff --git a/rb/lib/selenium/webdriver/remote/bridge.rb b/rb/lib/selenium/webdriver/remote/bridge.rb
index 655d331e1fc57..b977a99aab335 100644
--- a/rb/lib/selenium/webdriver/remote/bridge.rb
+++ b/rb/lib/selenium/webdriver/remote/bridge.rb
@@ -22,6 +22,8 @@ module WebDriver
     module Remote
       class Bridge
         autoload :COMMANDS, 'selenium/webdriver/remote/bridge/commands'
+        autoload :LocatorConverter, 'selenium/webdriver/remote/bridge/locator_converter'
+
         include Atoms
 
         PORT = 4444
@@ -31,12 +33,17 @@ class Bridge
 
         class << self
           attr_reader :extra_commands
+          attr_writer :locator_converter
 
           def add_command(name, verb, url, &block)
             @extra_commands ||= {}
             @extra_commands[name] = [verb, url]
             define_method(name, &block)
           end
+
+          def locator_converter
+            @locator_converter ||= LocatorConverter.new
+          end
         end
 
         #
@@ -53,6 +60,8 @@ def initialize(url:, http_client: nil)
           @http = http_client || Http::Default.new
           @http.server_url = uri
           @file_detector = nil
+
+          @locator_converter = self.class.locator_converter
         end
 
         #
@@ -516,7 +525,7 @@ def active_element
         alias switch_to_active_element active_element
 
         def find_element_by(how, what, parent_ref = [])
-          how, what = convert_locator(how, what)
+          how, what = @locator_converter.convert(how, what)
 
           return execute_atom(:findElements, Support::RelativeLocator.new(what).as_json).first if how == 'relative'
 
@@ -534,7 +543,7 @@ def find_element_by(how, what, parent_ref = [])
         end
 
         def find_elements_by(how, what, parent_ref = [])
-          how, what = convert_locator(how, what)
+          how, what = @locator_converter.convert(how, what)
 
           return execute_atom :findElements, Support::RelativeLocator.new(what).as_json if how == 'relative'
 
@@ -655,42 +664,6 @@ def prepare_capabilities_payload(capabilities)
           {capabilities: capabilities}
         end
 
-        def convert_locator(how, what)
-          how = SearchContext::FINDERS[how.to_sym] || how
-
-          case how
-          when 'class name'
-            how = 'css selector'
-            what = ".#{escape_css(what.to_s)}"
-          when 'id'
-            how = 'css selector'
-            what = "##{escape_css(what.to_s)}"
-          when 'name'
-            how = 'css selector'
-            what = "*[name='#{escape_css(what.to_s)}']"
-          end
-
-          if what.is_a?(Hash)
-            what = what.each_with_object({}) do |(h, w), hash|
-              h, w = convert_locator(h.to_s, w)
-              hash[h] = w
-            end
-          end
-
-          [how, what]
-        end
-
-        ESCAPE_CSS_REGEXP = /(['"\\#.:;,!?+<>=~*^$|%&@`{}\-\[\]()])/
-        UNICODE_CODE_POINT = 30
-
-        # Escapes invalid characters in CSS selector.
-        # @see https://mathiasbynens.be/notes/css-escapes
-        def escape_css(string)
-          string = string.gsub(ESCAPE_CSS_REGEXP) { |match| "\\#{match}" }
-          string = "\\#{UNICODE_CODE_POINT + Integer(string[0])} #{string[1..]}" if string[0]&.match?(/[[:digit:]]/)
-
-          string
-        end
       end # Bridge
     end # Remote
   end # WebDriver
diff --git a/rb/lib/selenium/webdriver/remote/bridge/locator_converter.rb b/rb/lib/selenium/webdriver/remote/bridge/locator_converter.rb
new file mode 100644
index 0000000000000..6c1065b9a5a25
--- /dev/null
+++ b/rb/lib/selenium/webdriver/remote/bridge/locator_converter.rb
@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+# Licensed to the Software Freedom Conservancy (SFC) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The SFC licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+module Selenium
+  module WebDriver
+    module Remote
+      class Bridge
+        class LocatorConverter
+          ESCAPE_CSS_REGEXP = /(['"\\#.:;,!?+<>=~*^$|%&@`{}\-\[\]()])/
+          UNICODE_CODE_POINT = 30
+
+          #
+          # Converts a locator to a specification compatible one.
+          # @param [String, Symbol] how
+          # @param [String] what
+          #
+
+          def convert(how, what)
+            how = SearchContext.finders[how.to_sym] || how
+
+            case how
+            when 'class name'
+              how = 'css selector'
+              what = ".#{escape_css(what.to_s)}"
+            when 'id'
+              how = 'css selector'
+              what = "##{escape_css(what.to_s)}"
+            when 'name'
+              how = 'css selector'
+              what = "*[name='#{escape_css(what.to_s)}']"
+            end
+
+            if what.is_a?(Hash)
+              what = what.each_with_object({}) do |(h, w), hash|
+                h, w = convert(h.to_s, w)
+                hash[h] = w
+              end
+            end
+
+            [how, what]
+          end
+
+          private
+
+          #
+          # Escapes invalid characters in CSS selector.
+          # @see https://mathiasbynens.be/notes/css-escapes
+          #
+
+          def escape_css(string)
+            string = string.gsub(ESCAPE_CSS_REGEXP) { |match| "\\#{match}" }
+            string = "\\#{UNICODE_CODE_POINT + Integer(string[0])} #{string[1..]}" if string[0]&.match?(/[[:digit:]]/)
+
+            string
+          end
+        end # LocatorConverter
+      end # Bridge
+    end # Remote
+  end # WebDriver
+end # Selenium

From 5d648ad72fac7b6eb1725468135dea9f6a0d9fdb Mon Sep 17 00:00:00 2001
From: Alex Rodionov <p0deje@gmail.com>
Date: Sun, 4 Jun 2023 16:19:54 -0700
Subject: [PATCH 5/6] [rb] Support registering custom finders for SearchContext

---
 .../selenium/webdriver/common/search_context.rb  | 12 ++++++++++--
 .../selenium/webdriver/search_context_spec.rb    | 16 ++++++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/rb/lib/selenium/webdriver/common/search_context.rb b/rb/lib/selenium/webdriver/common/search_context.rb
index 6bbc405333888..dad7896cbc907 100644
--- a/rb/lib/selenium/webdriver/common/search_context.rb
+++ b/rb/lib/selenium/webdriver/common/search_context.rb
@@ -35,6 +35,14 @@ module SearchContext
         xpath: 'xpath'
       }.freeze
 
+      class << self
+        attr_accessor :extra_finders
+
+        def finders
+          FINDERS.merge(extra_finders || {})
+        end
+      end
+
       #
       # Find the first element matching the given arguments
       #
@@ -57,7 +65,7 @@ module SearchContext
       def find_element(*args)
         how, what = extract_args(args)
 
-        by = FINDERS[how.to_sym]
+        by = SearchContext.finders[how.to_sym]
         raise ArgumentError, "cannot find element by #{how.inspect}" unless by
 
         bridge.find_element_by by, what, ref
@@ -72,7 +80,7 @@ def find_element(*args)
       def find_elements(*args)
         how, what = extract_args(args)
 
-        by = FINDERS[how.to_sym]
+        by = SearchContext.finders[how.to_sym]
         raise ArgumentError, "cannot find elements by #{how.inspect}" unless by
 
         bridge.find_elements_by by, what, ref
diff --git a/rb/spec/unit/selenium/webdriver/search_context_spec.rb b/rb/spec/unit/selenium/webdriver/search_context_spec.rb
index 88cd544b90d57..167fa51a410d5 100644
--- a/rb/spec/unit/selenium/webdriver/search_context_spec.rb
+++ b/rb/spec/unit/selenium/webdriver/search_context_spec.rb
@@ -89,6 +89,22 @@ def initialize(bridge)
           }.to raise_error(ArgumentError, 'cannot find elements by :foo')
         end
       end
+
+      context 'when extra finders are registered' do
+        around do |example|
+          described_class.extra_finders = {accessibility_id: 'accessibility id'}
+          example.call
+        ensure
+          described_class.extra_finders = nil
+        end
+
+        it 'finds element' do
+          allow(bridge).to receive(:find_element_by).with('accessibility id', 'foo', nil).and_return(element)
+
+          expect(search_context.find_element(accessibility_id: 'foo')).to eq(element)
+          expect(bridge).to have_received(:find_element_by).with('accessibility id', 'foo', nil)
+        end
+      end
     end
   end # WebDriver
 end # Selenium

From cae37fa4345778923373c458f20026f66c47cfa7 Mon Sep 17 00:00:00 2001
From: Alex Rodionov <p0deje@gmail.com>
Date: Sun, 4 Feb 2024 08:03:08 -0800
Subject: [PATCH 6/6] [rb] Support using custom element classes

---
 rb/lib/selenium/webdriver/remote/bridge.rb    | 21 +++---
 .../selenium/webdriver/remote/bridge_spec.rb  | 66 +++++++++++++++++++
 2 files changed, 78 insertions(+), 9 deletions(-)

diff --git a/rb/lib/selenium/webdriver/remote/bridge.rb b/rb/lib/selenium/webdriver/remote/bridge.rb
index b977a99aab335..773d3ef37d778 100644
--- a/rb/lib/selenium/webdriver/remote/bridge.rb
+++ b/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
diff --git a/rb/spec/unit/selenium/webdriver/remote/bridge_spec.rb b/rb/spec/unit/selenium/webdriver/remote/bridge_spec.rb
index cd8d7091561d0..cc8bf707045c3 100644
--- a/rb/spec/unit/selenium/webdriver/remote/bridge_spec.rb
+++ b/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