From dc08b17b99744836cff29e7471db54bdf41ca067 Mon Sep 17 00:00:00 2001
From: glaszig <glaszig@gmail.com>
Date: Fri, 22 Apr 2016 22:27:42 +0200
Subject: [PATCH] added after_save_html and after_save_screenshot callbacks

---
 README.md                            | 24 +++++++++++++++
 lib/capybara-screenshot.rb           |  8 +++++
 lib/capybara-screenshot/callbacks.rb | 44 ++++++++++++++++++++++++++++
 lib/capybara-screenshot/saver.rb     |  9 ++++++
 spec/unit/saver_spec.rb              | 34 +++++++++++++++++++++
 5 files changed, 119 insertions(+)
 create mode 100644 lib/capybara-screenshot/callbacks.rb

diff --git a/README.md b/README.md
index ee8bfb6..a86c287 100644
--- a/README.md
+++ b/README.md
@@ -211,6 +211,30 @@ Capybara::Screenshot.prune_strategy = :keep_last_run
 Capybara::Screenshot.prune_strategy = { keep: 20 }
 ```
 
+Callbacks
+---------
+
+You can hook your own logic into callbacks after the html/screenshot has been saved.
+
+```ruby
+# after Saver#save_html
+Capybara::Screenshot.after_save_html do |path|
+  mail = Mail.new do
+    delivery_method :sendmail
+    from     'capybara-screenshot@example.com'
+    to       'dev@example.com'
+    subject  'Capybara Screenshot'
+    add_file File.read path
+  end
+  mail.delivery_method :sendmail
+  mail.deliver
+end
+
+# after Saver#save_screenshot
+Capybara::Screenhot.after_save_screenshot do |path|
+  # ...
+end
+```
 
 Information about screenshots in RSpec output
 ---------------------------------------------
diff --git a/lib/capybara-screenshot.rb b/lib/capybara-screenshot.rb
index 6b84ca2..2d65584 100644
--- a/lib/capybara-screenshot.rb
+++ b/lib/capybara-screenshot.rb
@@ -103,6 +103,14 @@ def self.new_saver(*args)
       return saver
     end
 
+    def self.after_save_html &block
+      Saver.after_save_html &block
+    end
+
+    def self.after_save_screenshot &block
+      Saver.after_save_screenshot &block
+    end
+
     private
 
     # If the path isn't set, default to the current directory
diff --git a/lib/capybara-screenshot/callbacks.rb b/lib/capybara-screenshot/callbacks.rb
new file mode 100644
index 0000000..76a8314
--- /dev/null
+++ b/lib/capybara-screenshot/callbacks.rb
@@ -0,0 +1,44 @@
+module Capybara
+  module Screenshot
+    module Callbacks
+      class CallbackSet < Array
+        def call *args
+          each do |callback|
+            callback.call *args
+          end
+        end
+      end
+
+      module ClassMethods
+        def callbacks
+          @callbacks ||= {}
+        end
+
+        def define_callback name
+          callbacks[name] ||= CallbackSet.new
+
+          define_singleton_method name do |&block|
+            callbacks[name] << block
+          end
+        end
+
+        def run_callbacks name, *args
+          if cb_set = callbacks[name]
+            cb_set.call *args
+          end
+        end
+      end
+
+      module InstanceMethods
+        def run_callbacks name, *args
+          self.class.run_callbacks name, *args
+        end
+      end
+
+      def self.included receiver
+        receiver.extend         ClassMethods
+        receiver.send :include, InstanceMethods
+      end
+    end
+  end
+end
diff --git a/lib/capybara-screenshot/saver.rb b/lib/capybara-screenshot/saver.rb
index 8d89d92..1ca7965 100644
--- a/lib/capybara-screenshot/saver.rb
+++ b/lib/capybara-screenshot/saver.rb
@@ -1,9 +1,16 @@
 require 'capybara-screenshot/helpers'
+require 'capybara-screenshot/callbacks'
 
 module Capybara
   module Screenshot
     class Saver
+      include Capybara::Screenshot::Callbacks
+
+      define_callback :after_save_html
+      define_callback :after_save_screenshot
+
       attr_reader :capybara, :page, :file_base_name
+
       def initialize(capybara, page, html_save=true, filename_prefix='screenshot')
         @capybara, @page, @html_save = capybara, page, html_save
         time_now = Time.now
@@ -36,6 +43,7 @@ def save_html
           end
         end
         @html_saved = true
+        run_callbacks :after_save_html, html_path if html_saved?
       end
 
       def save_screenshot
@@ -47,6 +55,7 @@ def save_screenshot
           }.call(page.driver, path)
           @screenshot_saved = result != :not_supported
         end
+        run_callbacks :after_save_screenshot, screenshot_path if screenshot_saved?
       end
 
       def html_path
diff --git a/spec/unit/saver_spec.rb b/spec/unit/saver_spec.rb
index 8f97c41..f2057e6 100644
--- a/spec/unit/saver_spec.rb
+++ b/spec/unit/saver_spec.rb
@@ -154,6 +154,40 @@
     end
   end
 
+  describe 'callbacks' do
+    let(:saver) { Capybara::Screenshot::Saver.new(capybara_mock, page_mock) }
+
+    before do
+      allow(saver).to receive(:html_path) { 'page.html' }
+      allow(saver).to receive(:screenshot_path) { 'screenshot.png' }
+    end
+
+    before :all do
+      Capybara::Screenshot.after_save_html do |path|
+        puts "after_save_html ran with #{path}"
+      end
+      Capybara::Screenshot.after_save_screenshot do |path|
+        puts "after_save_screenshot ran with #{path}"
+      end
+    end
+
+    after :all do
+      Capybara::Screenshot::Saver.instance_eval { @callbacks = nil }
+    end
+
+    it 'runs after_save_html callbacks' do
+      expect do
+        saver.save
+      end.to output(/after_save_html ran with page\.html/).to_stdout
+    end
+
+    it 'runs after_save_screenshot callbacks' do
+      expect do
+        saver.save
+      end.to output(/after_save_screenshot ran with screenshot\.png/).to_stdout
+    end
+  end
+
   describe "with selenium driver" do
     before do
       allow(capybara_mock).to receive(:current_driver).and_return(:selenium)