Extracted save_and_open_page related functionality to a module and included it in SeleniumSession as well as the standard webrat session. Also added save_and_open_screengrab method to SeleniumSession.

lib/webrat/core.rb

@@ -11,3 +11,4 @@ require "webrat/core/elements/select_option"
 require "webrat/core/session"
 require "webrat/core/methods"
 require "webrat/core/matchers"
+require "webrat/core/save_and_open_page"

lib/webrat/core/save_and_open_page.rb

@@ -0,0 +1,50 @@
+module Webrat
+  module SaveAndOpenPage
+    # Saves the page out to RAILS_ROOT/tmp/ and opens it in the default
+    # web browser if on OS X. Useful for debugging.
+    # 
+    # Example:
+    #   save_and_open_page
+    def save_and_open_page
+      return unless File.exist?(saved_page_dir)
+
+      filename = "#{saved_page_dir}/webrat-#{Time.now.to_i}.html"
+      
+      File.open(filename, "w") do |f|
+        f.write rewrite_css_and_image_references(response_body)
+      end
+
+      open_in_browser(filename)
+    end
+    
+    def open_in_browser(path) # :nodoc
+      platform = ruby_platform
+      if platform =~ /cygwin/ || platform =~ /win32/
+        `rundll32 url.dll,FileProtocolHandler #{path.gsub("/", "\\\\")}`
+      elsif platform =~ /darwin/
+        `open #{path}`
+      end
+    end
+    
+    def rewrite_css_and_image_references(response_html) # :nodoc:
+      return response_html unless doc_root
+      response_html.gsub(/"\/(stylesheets|images)/, doc_root + '/\1')
+    end
+    
+    def saved_page_dir #:nodoc:
+      File.expand_path(".")
+    end
+    
+    def doc_root #:nodoc:
+      nil
+    end
+  
+  private
+
+    # accessor for testing
+    def ruby_platform
+      RUBY_PLATFORM
+    end
+    
+  end
+end

lib/webrat/core/session.rb

@@ -2,6 +2,7 @@ require "forwardable"
 require "ostruct"
 
 require "webrat/core/mime"
+require "webrat/core/save_and_open_page"
 
 module Webrat
   # A page load or form submission returned an unsuccessful response code (500-599)
@@ -30,6 +31,7 @@ module Webrat
   class Session
     extend Forwardable
     include Logging
+    include SaveAndOpenPage
     
     attr_reader :current_url
     attr_reader :elements
@@ -43,23 +45,6 @@ module Webrat
       
       reset
     end
-
-    # Saves the page out to RAILS_ROOT/tmp/ and opens it in the default
-    # web browser if on OS X. Useful for debugging.
-    # 
-    # Example:
-    #   save_and_open_page
-    def save_and_open_page
-      return unless File.exist?(saved_page_dir)
-
-      filename = "#{saved_page_dir}/webrat-#{Time.now.to_i}.html"
-      
-      File.open(filename, "w") do |f|
-        f.write rewrite_css_and_image_references(response_body)
-      end
-
-      open_in_browser(filename)
-    end
     
     def current_dom #:nodoc:
       current_scope.dom
@@ -78,10 +63,6 @@ module Webrat
       nil
     end
     
-    def saved_page_dir #:nodoc:
-      File.expand_path(".")
-    end
-
     def header(key, value)
       @custom_headers[key] = value
     end
@@ -172,20 +153,6 @@ module Webrat
     end
     
     webrat_deprecate :visits, :visit
-    
-    def open_in_browser(path) # :nodoc
-      platform = ruby_platform
-      if platform =~ /cygwin/ || platform =~ /win32/
-        `rundll32 url.dll,FileProtocolHandler #{path.gsub("/", "\\\\")}`
-      elsif platform =~ /darwin/
-        `open #{path}`
-      end
-    end
-    
-    def rewrite_css_and_image_references(response_html) # :nodoc:
-      return response_html unless doc_root
-      response_html.gsub(/"\/(stylesheets|images)/, doc_root + '/\1')
-    end
 
     # Subclasses can override this to show error messages without html
     def formatted_error #:nodoc:
@@ -247,9 +214,5 @@ module Webrat
       @_page_scope  = nil
     end
     
-    # accessor for testing
-    def ruby_platform
-      RUBY_PLATFORM
-    end
   end
 end

lib/webrat/selenium.rb

@@ -90,7 +90,10 @@ module Webrat
         def wait_for(*args, &block)
           webrat_session.wait_for(*args, &block)
         end
-        
+
+        def save_and_open_screengrab
+          webrat_session.save_and_open_screengrab
+        end
       end
     end
   end

lib/webrat/selenium/selenium_session.rb

@@ -1,3 +1,5 @@
+require "webrat/core/save_and_open_page"
+
 module Webrat
   class TimeoutError < WebratError
   end
@@ -17,6 +19,7 @@ module Webrat
   end
   
   class SeleniumSession
+    include Webrat::SaveAndOpenPage
     
     def initialize(*args) # :nodoc:
     end
@@ -160,6 +163,20 @@ module Webrat
     
     webrat_deprecate :browser, :selenium
     
+    
+    def save_and_open_screengrab
+      return unless File.exist?(saved_page_dir)
+      
+      filename = "#{saved_page_dir}/webrat-#{Time.now.to_i}.png"
+         
+      if $browser.chrome_backend?
+        $browser.capture_entire_page_screenshot(filename, '')
+      else
+        $browser.capture_screenshot(filename)
+      end
+        open_in_browser(filename)
+    end
+    
   protected
     
     def setup #:nodoc: