1 files changed, 53 insertions, 2 deletions
diff --git a/test/irb/test_ruby_lex.rb b/test/irb/test_ruby_lex.rb
index 25f037ec3b..255a5a1a63 100644
--- a/test/irb/test_ruby_lex.rb
+++ b/test/irb/test_ruby_lex.rb
@@ -7,7 +7,7 @@ module TestIRB
   class TestRubyLex < Test::Unit::TestCase
     Row = Struct.new(:content, :current_line_spaces, :new_line_spaces, :nesting_level)
 
-    class MockIO
+    class MockIO_AutoIndent
       def initialize(params, &assertion)
         @params = params
         @assertion = assertion
@@ -25,7 +25,7 @@ module TestIRB
       byte_pointer = lines.last.length
 
       ruby_lex = RubyLex.new()
-      io = MockIO.new([lines, last_line_index, byte_pointer, add_new_line]) do |auto_indent|
+      io = MockIO_AutoIndent.new([lines, last_line_index, byte_pointer, add_new_line]) do |auto_indent|
         error_message = "Calculated the wrong number of spaces for:\n #{lines.join("\n")}"
         assert_equal(correct_space_count, auto_indent, error_message)
       end
@@ -262,5 +262,56 @@ module TestIRB
         assert_nesting_level(lines, row.nesting_level)
       end
     end
+
+    PromptRow = Struct.new(:prompt, :content)
+
+    class MockIO_DynamicPrompt
+      def initialize(params, &assertion)
+        @params = params
+        @assertion = assertion
+      end
+
+      def dynamic_prompt(&block)
+        result = block.call(@params)
+        @assertion.call(result)
+      end
+    end
+
+    def assert_dynamic_prompt(lines, expected_prompt_list)
+      skip if RUBY_ENGINE == 'truffleruby'
+      ruby_lex = RubyLex.new()
+      io = MockIO_DynamicPrompt.new(lines) do |prompt_list|
+        error_message = "Calculated the wrong number of spaces for:\n #{lines.join("\n")}"
+        assert_equal(expected_prompt_list, prompt_list)
+      end
+      ruby_lex.set_prompt do |ltype, indent, continue, line_no|
+        '%03d:%01d:%1s:%s ' % [line_no, indent, ltype, continue ? '*' : '>']
+      end
+      ruby_lex.set_input(io)
+    end
+
+    def test_dyanmic_prompt
+      input_with_prompt = [
+        PromptRow.new('001:1: :* ', %q(def hoge)),
+        PromptRow.new('002:1: :* ', %q(  3)),
+        PromptRow.new('003:0: :> ', %q(end)),
+      ]
+
+      lines = input_with_prompt.map(&:content)
+      expected_prompt_list = input_with_prompt.map(&:prompt)
+      assert_dynamic_prompt(lines, expected_prompt_list)
+    end
+
+    def test_dyanmic_prompt_with_blank_line
+      input_with_prompt = [
+        PromptRow.new('001:0:]:* ', %q(%w[)),
+        PromptRow.new('002:0:]:* ', %q()),
+        PromptRow.new('003:0: :> ', %q(])),
+      ]
+
+      lines = input_with_prompt.map(&:content)
+      expected_prompt_list = input_with_prompt.map(&:prompt)
+      assert_dynamic_prompt(lines, expected_prompt_list)
+    end
   end
 end