2 files changed, 39 insertions, 35 deletions
diff --git a/test/prism/ruby/ripper_test.rb b/test/prism/ruby/ripper_test.rb
index c8d259135f..a89a9503b9 100644
--- a/test/prism/ruby/ripper_test.rb
+++ b/test/prism/ruby/ripper_test.rb
@@ -136,7 +136,7 @@ module Prism
       assert_equal(expected, lexer.parse[0].to_a)
       assert_equal(lexer.parse[0].to_a, lexer.scan[0].to_a)
 
-      assert_equal(%i[on_int on_sp on_op], Translation::Ripper::Lexer.new("1 +").lex.map(&:event))
+      assert_equal(%i[on_int on_sp on_op], Translation::Ripper::Lexer.new("1 +").lex.map { |token| token[1] })
       assert_raise(SyntaxError) { Translation::Ripper::Lexer.new("1 +").lex(raise_errors: true) }
     end
 
@@ -169,13 +169,13 @@ module Prism
       # Prism emits tokens by their order in the code, not in parse order
       ripper.sort_by! { |elem| elem[0] }
 
-      [prism.size, ripper.size].max.times do |i|
-        expected = ripper[i]
-        actual = prism[i]
+      [prism.size, ripper.size].max.times do |index|
+        expected = ripper[index]
+        actual = prism[index]
 
-        # Since tokens related to heredocs are not emitted in the same order,
-        # the state also doesn't line up.
-        if expected && actual && expected[1] == :on_heredoc_end && actual[1] == :on_heredoc_end
+        # There are some tokens that have slightly different state that do not
+        # effect the parse tree, so they may not match.
+        if expected && actual && expected[1] == actual[1] && %i[on_comment on_heredoc_end on_embexpr_end on_sp].include?(expected[1])
           expected[3] = actual[3] = nil
         end
 
diff --git a/test/prism/ruby/source_test.rb b/test/prism/ruby/source_test.rb
index afd2825765..f7cf4fe83a 100644
--- a/test/prism/ruby/source_test.rb
+++ b/test/prism/ruby/source_test.rb
@@ -4,44 +4,48 @@ require_relative "../test_helper"
 
 module Prism
   class SourceTest < TestCase
-    def test_line_to_byte_offset
-      parse_result = Prism.parse(<<~SRC)
+    def test_byte_offset
+      source = Prism.parse(<<~SRC).source
       abcd
       efgh
       ijkl
       SRC
-      source = parse_result.source
-
-      assert_equal 0, source.line_to_byte_offset(1)
-      assert_equal 5, source.line_to_byte_offset(2)
-      assert_equal 10, source.line_to_byte_offset(3)
-      assert_equal 15, source.line_to_byte_offset(4)
-      e = assert_raise(ArgumentError) { source.line_to_byte_offset(5) }
-      assert_equal "line 5 is out of range", e.message
-      e = assert_raise(ArgumentError) { source.line_to_byte_offset(0) }
-      assert_equal "line 0 is out of range", e.message
-      e = assert_raise(ArgumentError) { source.line_to_byte_offset(-1) }
-      assert_equal "line -1 is out of range", e.message
+
+      assert_equal 0, source.byte_offset(1, 0)
+      assert_equal 5, source.byte_offset(2, 0)
+      assert_equal 10, source.byte_offset(3, 0)
+      assert_equal 15, source.byte_offset(4, 0)
+
+      error = assert_raise(ArgumentError) { source.byte_offset(5, 0) }
+      assert_equal "line 5 is out of range", error.message
+
+      error = assert_raise(ArgumentError) { source.byte_offset(0, 0) }
+      assert_equal "line 0 is out of range", error.message
+
+      error = assert_raise(ArgumentError) { source.byte_offset(-1, 0) }
+      assert_equal "line -1 is out of range", error.message
     end
 
-    def test_line_to_byte_offset_with_start_line
-      parse_result = Prism.parse(<<~SRC, line: 11)
+    def test_byte_offset_with_start_line
+      source = Prism.parse(<<~SRC, line: 11).source
       abcd
       efgh
       ijkl
       SRC
-      source = parse_result.source
-
-      assert_equal 0, source.line_to_byte_offset(11)
-      assert_equal 5, source.line_to_byte_offset(12)
-      assert_equal 10, source.line_to_byte_offset(13)
-      assert_equal 15, source.line_to_byte_offset(14)
-      e = assert_raise(ArgumentError) { source.line_to_byte_offset(15) }
-      assert_equal "line 15 is out of range", e.message
-      e = assert_raise(ArgumentError) { source.line_to_byte_offset(10) }
-      assert_equal "line 10 is out of range", e.message
-      e = assert_raise(ArgumentError) { source.line_to_byte_offset(9) }
-      assert_equal "line 9 is out of range", e.message
+
+      assert_equal 0, source.byte_offset(11, 0)
+      assert_equal 5, source.byte_offset(12, 0)
+      assert_equal 10, source.byte_offset(13, 0)
+      assert_equal 15, source.byte_offset(14, 0)
+
+      error = assert_raise(ArgumentError) { source.byte_offset(15, 0) }
+      assert_equal "line 15 is out of range", error.message
+
+      error = assert_raise(ArgumentError) { source.byte_offset(10, 0) }
+      assert_equal "line 10 is out of range", error.message
+
+      error = assert_raise(ArgumentError) { source.byte_offset(9, 0) }
+      assert_equal "line 9 is out of range", error.message
     end
   end
 end