If no driver matching capybara.current_driver can be found, fall back to looking up by class

lib/capybara-screenshot.rb

@@ -148,9 +148,11 @@ def self.capybara_tmp_path=(path)
     driver.render(path)
   end
 
-  register_driver(:rack_test) do |driver, path|
+  block = proc do |driver, path|
     :not_supported
   end
+  register_driver(:rack_test, &block)
+  register_driver(:'Rack::Test', &block)
 
   register_driver(:mechanize) do |driver, path|
     :not_supported
@@ -163,6 +165,7 @@ def self.capybara_tmp_path=(path)
   register_driver :selenium, &selenium_block
   register_driver :selenium_chrome, &selenium_block
   register_driver :selenium_chrome_headless, &selenium_block
+  register_driver :'Capybara::Selenium::Driver', &selenium_block
 
   register_driver(:poltergeist) do |driver, path|
     driver.render(path, :full => true)
@@ -195,9 +198,11 @@ def self.capybara_tmp_path=(path)
     driver.save_screenshot(path)
   end
 
-  register_driver(:cuprite) do |driver, path|
-    driver.render(path, :full => true)
+  block = proc do |driver, path|
+    driver.render(path, full: true)
   end
+  register_driver(:cuprite, &block)
+  register_driver(:'Capybara::Cuprite::Driver', &block)
 end
 
 # Register filename prefix formatters

lib/capybara-screenshot/saver.rb

@@ -63,10 +63,12 @@ def save_html
       def save_screenshot
         path = screenshot_path
         clear_save_path do
-          result = Capybara::Screenshot.registered_drivers.fetch(capybara.current_driver) { |driver_name|
-            warn "capybara-screenshot could not detect a screenshot driver for '#{capybara.current_driver}'. Saving with default with unknown results."
-            Capybara::Screenshot.registered_drivers[:default]
-          }.call(page.driver, path)
+          driver = Capybara::Screenshot.registered_drivers[capybara.current_driver] ||
+            Capybara::Screenshot.registered_drivers.fetch(page.driver&.class&.to_sym) { |_|
+              warn "capybara-screenshot could not detect a screenshot driver for '#{capybara.current_driver}' or #{page.driver.class}. Saving with default with unknown results."
+              Capybara::Screenshot.registered_drivers[:default]
+            }
+          result = driver.call(page.driver, path)
           @screenshot_saved = result != :not_supported
         end
         run_callbacks :after_save_screenshot, screenshot_path if screenshot_saved?

spec/support/general_helpers.rb

@@ -0,0 +1,31 @@
+# Copied from activesupport/lib/active_support/inflector/methods.rb
+def constantize(camel_cased_word)
+  names = camel_cased_word.split("::".freeze)
+
+  # Trigger a built-in NameError exception including the ill-formed constant in the message.
+  Object.const_get(camel_cased_word) if names.empty?
+
+  # Remove the first blank element in case of '::ClassName' notation.
+  names.shift if names.size > 1 && names.first.empty?
+
+  names.inject(Object) do |constant, name|
+    if constant == Object
+      constant.const_get(name)
+    else
+      candidate = constant.const_get(name)
+      next candidate if constant.const_defined?(name, false)
+      next candidate unless Object.const_defined?(name)
+
+      # Go down the ancestors to check if it is owned directly. The check
+      # stops when we reach Object or the end of ancestors tree.
+      constant = constant.ancestors.inject(constant) do |const, ancestor|
+        break const    if ancestor == Object
+        break ancestor if ancestor.const_defined?(name, false)
+        const
+      end
+
+      # owner is in Object, so raise
+      constant.const_get(name, false)
+    end
+  end
+end

spec/unit/saver_spec.rb

@@ -15,12 +15,25 @@
     Timecop.freeze(Time.local(2012, 6, 7, 8, 9, 10, 0))
   end
 
+  shared_context 'driver based on class' do |class_name|
+    before do
+      allow(capybara_mock).to receive(:current_driver).and_return(:some_unrecognized_alias)
+      expect(driver_mock).to receive(:class).and_return(class_name)
+      stub_const(class_name, driver_mock)
+      expect(page_mock).to receive(:driver).and_return(constantize(class_name).new)
+    end
+  end
+
   let(:capybara_root) { '/tmp' }
   let(:timestamp) { '2012-06-07-08-09-10.000' }
   let(:file_basename) { "screenshot_#{timestamp}" }
   let(:screenshot_path) { "#{capybara_root}/#{file_basename}.png" }
 
-  let(:driver_mock) { double('Capybara driver').as_null_object }
+  let(:driver_mock) {
+    double('Capybara driver').as_null_object.tap do |m|
+      allow(m).to receive(:class).and_return('MockCapybaraDriver')
+    end
+  }
   let(:page_mock) {
     double('Capybara session page', :body => 'body', :driver => driver_mock).as_null_object.tap do |m|
       allow(m).to receive(:current_path).and_return('/')
@@ -34,6 +47,8 @@
 
   let(:saver) { Capybara::Screenshot::Saver.new(capybara_mock, page_mock) }
 
+  #═════════════════════════════════════════════════════════════════════════════════════════════════
+
   context 'html filename with Capybara Version 1' do
     before do
       stub_const("Capybara::VERSION", '1')
@@ -257,11 +272,9 @@
     end
   end
 
-  describe "with selenium driver" do
-    before do
-      allow(capybara_mock).to receive(:current_driver).and_return(:selenium)
-    end
+  #═════════════════════════════════════════════════════════════════════════════════════════════════
 
+  shared_examples_for "with selenium driver" do
     it 'saves via browser' do
       browser_mock = double('browser')
       expect(driver_mock).to receive(:browser).and_return(browser_mock)
@@ -272,6 +285,22 @@
     end
   end
 
+  describe "with selenium driver" do
+    before do
+      allow(capybara_mock).to receive(:current_driver).and_return(:selenium)
+    end
+
+    it_behaves_like 'with selenium driver'
+  end
+
+  describe "with selenium driver based on class" do
+    include_context 'driver based on class', 'Capybara::Selenium::Driver'
+
+    it_behaves_like 'with selenium driver'
+  end
+
+  #═════════════════════════════════════════════════════════════════════════════════════════════════
+
   describe "with poltergeist driver" do
     before do
       allow(capybara_mock).to receive(:current_driver).and_return(:poltergeist)
@@ -298,11 +327,9 @@
     end
   end
 
-  describe "with cuprite driver" do
-    before do
-      allow(capybara_mock).to receive(:current_driver).and_return(:cuprite)
-    end
+  #═════════════════════════════════════════════════════════════════════════════════════════════════
 
+  shared_examples_for "with cuprite driver" do
     it 'saves driver render with :full => true' do
       expect(driver_mock).to receive(:render).with(screenshot_path, {:full => true})
 
@@ -311,6 +338,22 @@
     end
   end
 
+  describe "with cuprite driver" do
+    before do
+      allow(capybara_mock).to receive(:current_driver).and_return(:cuprite)
+    end
+
+    it_behaves_like 'with cuprite driver'
+  end
+
+  describe "with cuprite driver based on class" do
+    include_context 'driver based on class', 'Capybara::Cuprite::Driver'
+
+    it_behaves_like 'with cuprite driver'
+  end
+
+  #═════════════════════════════════════════════════════════════════════════════════════════════════
+
   describe "with webkit driver" do
     before do
       allow(capybara_mock).to receive(:current_driver).and_return(:webkit)
@@ -395,6 +438,8 @@
     end
   end
 
+  #═════════════════════════════════════════════════════════════════════════════════════════════════
+
   describe "with unknown driver" do
     before do
       allow(capybara_mock).to receive(:current_driver).and_return(:unknown)