From 4cf9aeb151a9f01d51b4450a1b8998d21e0f0ebc Mon Sep 17 00:00:00 2001
From: Alex Rodionov <p0deje@gmail.com>
Date: Sat, 1 Apr 2023 21:51:34 -0500
Subject: [PATCH] [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)