1 files changed, 34 insertions, 5 deletions
diff --git a/actionpack/lib/action_dispatch/system_testing/driver.rb b/actionpack/lib/action_dispatch/system_testing/driver.rb
index 1a027f2e23..280989a146 100644
--- a/actionpack/lib/action_dispatch/system_testing/driver.rb
+++ b/actionpack/lib/action_dispatch/system_testing/driver.rb
@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module ActionDispatch
   module SystemTesting
     class Driver # :nodoc:
@@ -9,14 +11,14 @@ module ActionDispatch
       end
 
       def use
-        register unless rack_test?
+        register if registerable?
 
         setup
       end
 
       private
-        def rack_test?
-          @name == :rack_test
+        def registerable?
+          [:selenium, :poltergeist, :webkit].include?(@name)
         end
 
         def register
@@ -29,8 +31,35 @@ module ActionDispatch
           end
         end
 
+        def browser_options
+          if @browser == :headless_chrome
+            browser_options = Selenium::WebDriver::Chrome::Options.new
+            browser_options.args << "--headless"
+            browser_options.args << "--disable-gpu"
+
+            @options.merge(options: browser_options)
+          elsif @browser == :headless_firefox
+            browser_options = Selenium::WebDriver::Firefox::Options.new
+            browser_options.args << "-headless"
+
+            @options.merge(options: browser_options)
+          else
+            @options
+          end
+        end
+
+        def browser
+          if @browser == :headless_chrome
+            :chrome
+          elsif @browser == :headless_firefox
+            :firefox
+          else
+            @browser
+          end
+        end
+
         def register_selenium(app)
-          Capybara::Selenium::Driver.new(app, { browser: @browser }.merge(@options)).tap do |driver|
+          Capybara::Selenium::Driver.new(app, { browser: browser }.merge(browser_options)).tap do |driver|
             driver.browser.manage.window.size = Selenium::WebDriver::Dimension.new(*@screen_size)
           end
         end
@@ -41,7 +70,7 @@ module ActionDispatch
 
         def register_webkit(app)
           Capybara::Webkit::Driver.new(app, Capybara::Webkit::Configuration.to_hash.merge(@options)).tap do |driver|
-            driver.resize_window(*@screen_size)
+            driver.resize_window_to(driver.current_window_handle, *@screen_size)
           end
         end