[rb] Support registering extra headers in HTTP client

SeleniumHQ · May 15, 2024 · 4f72e3f · 4f72e3f
1 parent 6978ea8
commit 4f72e3f
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 5 deletions.
diff --git 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
@@ -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,17 +45,24 @@ 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)
 
             expect(common).to have_received(:request)
               .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