1 files changed, 9 insertions, 21 deletions
diff --git a/actionpack/test/dispatch/system_testing/screenshot_helper_test.rb b/actionpack/test/dispatch/system_testing/screenshot_helper_test.rb
index d6b501b3ac..a83818fd80 100644
--- a/actionpack/test/dispatch/system_testing/screenshot_helper_test.rb
+++ b/actionpack/test/dispatch/system_testing/screenshot_helper_test.rb
@@ -4,13 +4,13 @@ require "capybara/dsl"
 
 class ScreenshotHelperTest < ActiveSupport::TestCase
   test "image path is saved in tmp directory" do
-    new_test = ActionDispatch::SystemTestCase.new("x")
+    new_test = DrivenBySeleniumWithChrome.new("x")
 
     assert_equal "tmp/screenshots/x.png", new_test.send(:image_path)
   end
 
   test "image path includes failures text if test did not pass" do
-    new_test = ActionDispatch::SystemTestCase.new("x")
+    new_test = DrivenBySeleniumWithChrome.new("x")
 
     new_test.stub :passed?, false do
       assert_equal "tmp/screenshots/failures_x.png", new_test.send(:image_path)
@@ -18,7 +18,7 @@ class ScreenshotHelperTest < ActiveSupport::TestCase
   end
 
   test "image path does not include failures text if test skipped" do
-    new_test = ActionDispatch::SystemTestCase.new("x")
+    new_test = DrivenBySeleniumWithChrome.new("x")
 
     new_test.stub :passed?, false do
       new_test.stub :skipped?, true do
@@ -26,28 +26,16 @@ class ScreenshotHelperTest < ActiveSupport::TestCase
       end
     end
   end
+end
 
+class RackTestScreenshotsTest < DrivenByRackTest
   test "rack_test driver does not support screenshot" do
-    begin
-      original_driver = Capybara.current_driver
-      Capybara.current_driver = :rack_test
-
-      new_test = ActionDispatch::SystemTestCase.new("x")
-      assert_not new_test.send(:supports_screenshot?)
-    ensure
-      Capybara.current_driver = original_driver
-    end
+    assert_not self.send(:supports_screenshot?)
   end
+end
 
+class SeleniumScreenshotsTest < DrivenBySeleniumWithChrome
   test "selenium driver supports screenshot" do
-    begin
-      original_driver = Capybara.current_driver
-      Capybara.current_driver = :selenium
-
-      new_test = ActionDispatch::SystemTestCase.new("x")
-      assert new_test.send(:supports_screenshot?)
-    ensure
-      Capybara.current_driver = original_driver
-    end
+    assert self.send(:supports_screenshot?)
   end
 end