1 files changed, 67 insertions, 6 deletions
diff --git a/test/prism/result/warnings_test.rb b/test/prism/result/warnings_test.rb
index c147616a6a..8ccaec74ed 100644
--- a/test/prism/result/warnings_test.rb
+++ b/test/prism/result/warnings_test.rb
@@ -71,6 +71,64 @@ module Prism
       assert_warning("_ = 1.0e100000", "out of range")
     end
 
+    def test_indentation_mismatch
+      assert_warning("if true\n  end", "mismatched indentations at 'end' with 'if'")
+      assert_warning("if true\n  elsif true\nend", "mismatched indentations at 'elsif' with 'if'")
+      assert_warning("if true\n  else\nend", "mismatched indentations at 'else' with 'if'", "mismatched indentations at 'end' with 'else'")
+
+      assert_warning("unless true\n  end", "mismatched indentations at 'end' with 'unless'")
+      assert_warning("unless true\n  else\nend", "mismatched indentations at 'else' with 'unless'", "mismatched indentations at 'end' with 'else'")
+
+      assert_warning("while true\n  end", "mismatched indentations at 'end' with 'while'")
+      assert_warning("until true\n  end", "mismatched indentations at 'end' with 'until'")
+
+      assert_warning("begin\n  end", "mismatched indentations at 'end' with 'begin'")
+      assert_warning("begin\n  rescue\nend", "mismatched indentations at 'rescue' with 'begin'")
+      assert_warning("begin\n  ensure\nend", "mismatched indentations at 'ensure' with 'begin'")
+      assert_warning("begin\nrescue\n  else\nend", "mismatched indentations at 'else' with 'begin'", "mismatched indentations at 'end' with 'else'")
+      assert_warning("begin\n  rescue\n    ensure\n      end", "mismatched indentations at 'rescue' with 'begin'", "mismatched indentations at 'ensure' with 'begin'", "mismatched indentations at 'end' with 'begin'");
+
+      assert_warning("def foo\n  end", "mismatched indentations at 'end' with 'def'")
+      assert_warning("def foo\n  rescue\nend", "mismatched indentations at 'rescue' with 'def'")
+      assert_warning("def foo\n  ensure\nend", "mismatched indentations at 'ensure' with 'def'")
+      assert_warning("def foo\nrescue\n  else\nend", "mismatched indentations at 'else' with 'def'", "mismatched indentations at 'end' with 'else'")
+      assert_warning("def foo\n  rescue\n    ensure\n      end", "mismatched indentations at 'rescue' with 'def'", "mismatched indentations at 'ensure' with 'def'", "mismatched indentations at 'end' with 'def'");
+
+      assert_warning("class Foo\n  end", "mismatched indentations at 'end' with 'class'")
+      assert_warning("class Foo\n  rescue\nend", "mismatched indentations at 'rescue' with 'class'")
+      assert_warning("class Foo\n  ensure\nend", "mismatched indentations at 'ensure' with 'class'")
+      assert_warning("class Foo\nrescue\n  else\nend", "mismatched indentations at 'else' with 'class'", "mismatched indentations at 'end' with 'else'")
+      assert_warning("class Foo\n  rescue\n    ensure\n      end", "mismatched indentations at 'rescue' with 'class'", "mismatched indentations at 'ensure' with 'class'", "mismatched indentations at 'end' with 'class'");
+
+      assert_warning("module Foo\n  end", "mismatched indentations at 'end' with 'module'")
+      assert_warning("module Foo\n  rescue\nend", "mismatched indentations at 'rescue' with 'module'")
+      assert_warning("module Foo\n  ensure\nend", "mismatched indentations at 'ensure' with 'module'")
+      assert_warning("module Foo\nrescue\n  else\nend", "mismatched indentations at 'else' with 'module'", "mismatched indentations at 'end' with 'else'")
+      assert_warning("module Foo\n  rescue\n    ensure\n      end", "mismatched indentations at 'rescue' with 'module'", "mismatched indentations at 'ensure' with 'module'", "mismatched indentations at 'end' with 'module'");
+
+      assert_warning("class << foo\n  end", "mismatched indentations at 'end' with 'class'")
+      assert_warning("class << foo\n  rescue\nend", "mismatched indentations at 'rescue' with 'class'")
+      assert_warning("class << foo\n  ensure\nend", "mismatched indentations at 'ensure' with 'class'")
+      assert_warning("class << foo\nrescue\n  else\nend", "mismatched indentations at 'else' with 'class'", "mismatched indentations at 'end' with 'else'")
+      assert_warning("class << foo\n  rescue\n    ensure\n      end", "mismatched indentations at 'rescue' with 'class'", "mismatched indentations at 'ensure' with 'class'", "mismatched indentations at 'end' with 'class'");
+
+      assert_warning("case 1; when 2\n  end", "mismatched indentations at 'end' with 'case'")
+      assert_warning("case 1; in 2\n  end", "mismatched indentations at 'end' with 'case'")
+
+      assert_warning("-> {\n  }", "mismatched indentations at '}' with '->'")
+      assert_warning("-> do\n  end", "mismatched indentations at 'end' with '->'")
+      assert_warning("-> do\n  rescue\nend", "mismatched indentations at 'rescue' with '->'")
+      assert_warning("-> do\n  ensure\nend", "mismatched indentations at 'ensure' with '->'")
+      assert_warning("-> do\nrescue\n  else\nend", "mismatched indentations at 'else' with '->'", "mismatched indentations at 'end' with 'else'")
+      assert_warning("-> do\n  rescue\n    ensure\n      end", "mismatched indentations at 'rescue' with '->'", "mismatched indentations at 'ensure' with '->'", "mismatched indentations at 'end' with '->'");
+      assert_warning("foo do\nrescue\n  else\nend", "mismatched indentations at 'end' with 'else'")
+
+      refute_warning("class Foo; end") # same line
+      refute_warning("; class Foo\nend") # non whitespace on opening line
+      refute_warning("\tclass Foo\n        end") # tab stop matches space
+      refute_warning("    \tclass Foo\n        end") # tab stop matches space
+    end
+
     def test_integer_in_flip_flop
       assert_warning("1 if 2..foo", "integer")
     end
@@ -275,14 +333,17 @@ module Prism
 
     private
 
-    def assert_warning(source, message)
+    def assert_warning(source, *messages)
       warnings = Prism.parse(source).warnings
+      assert_equal messages.length, warnings.length, "Expected #{messages.length} warning(s) in #{source.inspect}, got #{warnings.map(&:message).inspect}"
 
-      assert_equal 1, warnings.length, "Expected only one warning in #{source.inspect}, got #{warnings.map(&:message).inspect}"
-      assert_include warnings.first.message, message
+      warnings.zip(messages).each do |warning, message|
+        assert_include warning.message, message
+      end
 
       if defined?(RubyVM::AbstractSyntaxTree)
-        assert_include capture_warning { RubyVM::AbstractSyntaxTree.parse(source) }, message
+        stderr = capture_stderr { RubyVM::AbstractSyntaxTree.parse(source) }
+        messages.each { |message| assert_include stderr, message }
       end
     end
 
@@ -290,11 +351,11 @@ module Prism
       assert_empty Prism.parse(source, **options).warnings
 
       if compare && defined?(RubyVM::AbstractSyntaxTree)
-        assert_empty capture_warning { RubyVM::AbstractSyntaxTree.parse(source) }
+        assert_empty capture_stderr { RubyVM::AbstractSyntaxTree.parse(source) }
       end
     end
 
-    def capture_warning
+    def capture_stderr
       stderr, $stderr, verbose, $VERBOSE = $stderr, StringIO.new, $VERBOSE, true
 
       begin