From 251570fd8dfa3bf67086150c392c3ec42b6816ac Mon Sep 17 00:00:00 2001
From: "yuuji.yaginuma" <yuuji.yaginuma@gmail.com>
Date: Mon, 5 Apr 2021 12:09:53 +0900
Subject: [PATCH] Add support for billy's selenium drivers

`#109` supported billy, but only added poltergeist driver.
This patch supports more drivers(related with selenium).

Ref: https://github.com/oesmith/puffing-billy/blob/fc12d19cfa27182463c5c300acacd517a557d339/lib/billy/browsers/capybara.rb#L48-L96
---
 lib/capybara-screenshot.rb |  5 ++++
 spec/unit/saver_spec.rb    | 60 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/lib/capybara-screenshot.rb b/lib/capybara-screenshot.rb
index 7d32560..c8ce431 100644
--- a/lib/capybara-screenshot.rb
+++ b/lib/capybara-screenshot.rb
@@ -172,6 +172,11 @@ def self.capybara_tmp_path=(path)
     driver.render(path, :full => true)
   end
 
+  register_driver :selenium_billy, &selenium_block
+  register_driver :selenium_headless_billy, &selenium_block
+  register_driver :selenium_chrome_billy, &selenium_block
+  register_driver :selenium_chrome_headless_billy, &selenium_block
+
   webkit_block = proc do |driver, path|
     if driver.respond_to?(:save_screenshot)
       driver.save_screenshot(path, webkit_options)
diff --git a/spec/unit/saver_spec.rb b/spec/unit/saver_spec.rb
index 7408b41..6fcdb74 100644
--- a/spec/unit/saver_spec.rb
+++ b/spec/unit/saver_spec.rb
@@ -298,6 +298,66 @@
     end
   end
 
+  describe "with selenium_billy driver" do
+    before do
+      allow(capybara_mock).to receive(:current_driver).and_return(:selenium_billy)
+    end
+
+    it 'saves via browser' do
+      browser_mock = double('browser')
+      expect(driver_mock).to receive(:browser).and_return(browser_mock)
+      expect(browser_mock).to receive(:save_screenshot).with(screenshot_path)
+
+      saver.save
+      expect(saver).to be_screenshot_saved
+    end
+  end
+
+  describe "with selenium_headless_billy driver" do
+    before do
+      allow(capybara_mock).to receive(:current_driver).and_return(:selenium_headless_billy)
+    end
+
+    it 'saves via browser' do
+      browser_mock = double('browser')
+      expect(driver_mock).to receive(:browser).and_return(browser_mock)
+      expect(browser_mock).to receive(:save_screenshot).with(screenshot_path)
+
+      saver.save
+      expect(saver).to be_screenshot_saved
+    end
+  end
+
+  describe "with selenium_chrome_billy driver" do
+    before do
+      allow(capybara_mock).to receive(:current_driver).and_return(:selenium_chrome_billy)
+    end
+
+    it 'saves via browser' do
+      browser_mock = double('browser')
+      expect(driver_mock).to receive(:browser).and_return(browser_mock)
+      expect(browser_mock).to receive(:save_screenshot).with(screenshot_path)
+
+      saver.save
+      expect(saver).to be_screenshot_saved
+    end
+  end
+
+  describe "with selenium_chrome_headless_billy driver" do
+    before do
+      allow(capybara_mock).to receive(:current_driver).and_return(:selenium_chrome_headless_billy)
+    end
+
+    it 'saves via browser' do
+      browser_mock = double('browser')
+      expect(driver_mock).to receive(:browser).and_return(browser_mock)
+      expect(browser_mock).to receive(:save_screenshot).with(screenshot_path)
+
+      saver.save
+      expect(saver).to be_screenshot_saved
+    end
+  end
+
   describe "with cuprite driver" do
     before do
       allow(capybara_mock).to receive(:current_driver).and_return(:cuprite)