1 files changed, 32 insertions, 28 deletions

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