1 files changed, 16 insertions, 3 deletions
diff --git a/actionpack/test/dispatch/system_testing/screenshot_helper_test.rb b/actionpack/test/dispatch/system_testing/screenshot_helper_test.rb
index 3b4ea96c4f..a83818fd80 100644
--- a/actionpack/test/dispatch/system_testing/screenshot_helper_test.rb
+++ b/actionpack/test/dispatch/system_testing/screenshot_helper_test.rb
@@ -1,15 +1,16 @@
 require "abstract_unit"
 require "action_dispatch/system_testing/test_helpers/screenshot_helper"
+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)
@@ -17,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,3 +27,15 @@ class ScreenshotHelperTest < ActiveSupport::TestCase
     end
   end
 end
+
+class RackTestScreenshotsTest < DrivenByRackTest
+  test "rack_test driver does not support screenshot" do
+    assert_not self.send(:supports_screenshot?)
+  end
+end
+
+class SeleniumScreenshotsTest < DrivenBySeleniumWithChrome
+  test "selenium driver supports screenshot" do
+    assert self.send(:supports_screenshot?)
+  end
+end