diff --git a/test/support/retry_on_flaky_tests.rb b/test/support/retry_on_flaky_tests.rb
index 4ad11619..27ed7c90 100644
--- a/test/support/retry_on_flaky_tests.rb
+++ b/test/support/retry_on_flaky_tests.rb
@@ -15,16 +15,23 @@ def self.[](*error_classes, max_tries: 3)
         error_classes
       end
 
-      def capture_exceptions
-        tries ||= 1
+      def run_one_method(klass, method_name, reporter)
+        report_result = nil
+        max_tries.times do
+          result = Minitest.run_one_method(klass, method_name)
+          report_result ||= result
+          (report_result = result) and break if result.passed?
 
-        yield
-      rescue => err
-        if (tries +=1) <= max_tries && error_classes.include?(err.class)
-          puts "Got #{err.message} Error, trying again #{tries}. time"
-          retry
-        else
-          super { raise err }
+          break unless retryable_failure?(result)
+        end
+        reporter.record(report_result)
+      end
+
+      def retryable_failure?(result)
+        result.failures.map do |failure|
+          failure.error.to_s
+        end.any? do |failure_msg|
+          error_classes.first {|error_class| failure_msg =~ error_class.name }
         end
       end
     end
diff --git a/test/test_helper.rb b/test/test_helper.rb
index 6be8069c..7f177873 100644
--- a/test/test_helper.rb
+++ b/test/test_helper.rb
@@ -70,7 +70,7 @@
 class ActiveSupport::TestCase
   include CustomAssertions
 
-  prepend RetryOnFlakyTests[
+  extend RetryOnFlakyTests[
     # randomly happening on CI
     Ferrum::PendingConnectionsError,
     # race condition when trying to move mouse to element, can happen e.g. after fade-in/out of modal dialog