lib/capybara/sessionkeeper.rb in capybara-sessionkeeper-0.1.5 vs lib/capybara/sessionkeeper.rb in capybara-sessionkeeper-0.1.6

@@ -1,7 +1,8 @@
 require 'capybara'
 require "capybara/sessionkeeper/version"
+require 'yaml'
 
 module Capybara
   module Sessionkeeper
     class CookieError < StandardError; end
 
@@ -11,21 +12,30 @@
       File.open(path, 'wb') {|f| f.puts(data) }
       path
     end
 
     def restore_cookies(path = nil)
-      raise CookieError, "visit must be performed to restore cookies" if ['data:,', 'about:blank'].include?(current_url)
       path ||= find_latest_cookie_file
       return nil if path.nil?
       data = File.open(path, 'rb') {|f| f.read }
-      Marshal.load(data).each do |d|
+      restore_cookies_from_data(data)
+    end
+
+    def restore_cookies_from_data(data, options = {})
+      raise CookieError, "visit must be performed to restore cookies" if ['data:,', 'about:blank'].include?(current_url)
+      cookies = %w[yml yaml].include?(options[:format]) ? YAML.load(data) : Marshal.load(data)
+      cookies.each do |d|
         begin
           driver.browser.manage.add_cookie d
         rescue => e
           skip_invalid_cookie_domain_error(e)
         end
       end
       driver.browser.manage.all_cookies
+    end
+
+    def cookies_to_yaml
+      YAML.dump driver.browser.manage.all_cookies
     end
 
     def cookie_file_extension
       'cookies.txt'
     end