lib/page_magic.rb in page_magic-0.8.8 vs lib/page_magic.rb in page_magic-0.8.9

@@ -12,19 +12,37 @@
 require 'page_magic/page_section'
 
 module PageMagic
   class << self
     def session browser=nil, options = {}
-      if browser
-        Capybara.register_driver browser do |app|
-          options[:browser] = browser
-          Capybara::Selenium::Driver.new(app, options)
-        end
-        Session.new(Capybara::Session.new(browser, nil))
+      if browser.is_a?(Hash)
+        Session.new(Capybara::Session.new(:rack_test, browser[:application]))
       else
-        Capybara.reset!
-        Session.new(Capybara.current_session)
+        if browser
+          application = options.delete(:application)
+
+          Capybara.register_driver browser do |app|
+            options[:browser] = browser
+            case browser
+              when :poltergeist
+                require 'capybara/poltergeist'
+                Capybara::Poltergeist::Driver.new(app)
+              when :rack_test
+                Capybara::RackTest::Driver.new(app, options)
+              else
+                require 'watir-webdriver'
+                Capybara::Selenium::Driver.new(app, options)
+            end
+
+          end
+
+          Session.new(Capybara::Session.new(browser, application))
+        else
+          Capybara.reset!
+          Session.new(Capybara.current_session)
+        end
       end
+
     end
 
     def included clazz
       clazz.extend ClassMethods, PageElements
       pages << clazz if clazz.is_a? Class
\ No newline at end of file