Parse tests using prism

railties/lib/rails/test_unit/test_parser.rb

@@ -1,5 +1,59 @@
 # frozen_string_literal: true
 
+begin
+  require "prism"
+rescue LoadError
+  # If Prism isn't available (because of using an older Ruby version) then we'll
+  # define a fallback parser using ripper.
+end
+
+if defined?(Prism)
+  module Rails
+    module TestUnit
+      # Parse a test file to extract the line ranges of all tests in both
+      # method-style (def test_foo) and declarative-style (test "foo" do)
+      class TestParser < Prism::Visitor
+        # Helper to translate a method object into the path and line range where
+        # the method was defined.
+        def self.definition_for(method)
+          filepath, start_line = method.source_location
+          Prism.parse_file(filepath).value.accept(new(ranges = {}))
+          (end_line = ranges[start_line]) && [filepath, (start_line..end_line)]
+        end
+
+        def self.definitions_for(source, filepath)
+          Prism.parse(source, filepath: filepath).value.accept(new(ranges = {}))
+          ranges
+        end
+
+        attr_reader :ranges
+
+        def initialize(ranges)
+          @ranges = ranges
+        end
+
+        def visit_def_node(node)
+          if node.name.start_with?("test")
+            ranges[node.location.start_line] = node.location.end_line
+          end
+          super
+        end
+
+        def visit_call_node(node)
+          if node.name == :test
+            ranges[node.location.start_line] = node.location.end_line
+          end
+          super
+        end
+      end
+    end
+  end
+
+  # If we have Prism, then we don't need to define the fallback parser using
+  # ripper.
+  return
+end
+
 require "ripper"
 
 module Rails

railties/test/test_unit/test_parser_test.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require "active_support/deprecator"
 require "active_support/test_case"
 require "active_support/testing/autorun"
 require "rails/test_unit/test_parser"
@@ -42,7 +43,7 @@ class TestParserTest < ActiveSupport::TestCase
       end
     RUBY
 
-    parser = Rails::TestUnit::TestParser.new(example_test, "example_test.rb")
+    actual_map = Rails::TestUnit::TestParser.definitions_for(example_test, "example_test.rb")
     expected_map = {
       4 => 8,   # test_method
       10 => 10, # test_oneline
@@ -53,6 +54,6 @@ class TestParserTest < ActiveSupport::TestCase
       27 => 27, # declarative oneilne do
       29 => 32  # declarative multiline w/braces
     }
-    assert_equal expected_map, parser.parse
+    assert_equal expected_map, actual_map
   end
 end