From d3d3ccd7e3d69208cb023cf78f69177f911deccd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?fn=20=E2=8C=83=20=E2=8C=A5?=
 <70830482+FnControlOption@users.noreply.github.com>
Date: Sun, 24 Oct 2021 11:56:46 -0700
Subject: [PATCH] Add `expect_correction` and `expect_no_corrections`

---
 src/ameba/spec/annotated_source.cr |  3 ++-
 src/ameba/spec/expect_issue.cr     | 38 +++++++++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/src/ameba/spec/annotated_source.cr b/src/ameba/spec/annotated_source.cr
index 26d5271e..76081fe7 100644
--- a/src/ameba/spec/annotated_source.cr
+++ b/src/ameba/spec/annotated_source.cr
@@ -15,7 +15,8 @@ class Ameba::Spec::AnnotatedSource
   def self.parse(annotated_code)
     lines = [] of String
     annotations = [] of {Int32, String, String}
-    annotated_code.each_line do |code_line|
+    code_lines = annotated_code.split('\n') # must preserve trailing newline
+    code_lines.each do |code_line|
       if (annotation_match = ANNOTATION_PATTERN_1.match(code_line))
         message_index = annotation_match.end
         prefix = code_line[0...message_index]
diff --git a/src/ameba/spec/expect_issue.cr b/src/ameba/spec/expect_issue.cr
index fd8d1871..3571b164 100644
--- a/src/ameba/spec/expect_issue.cr
+++ b/src/ameba/spec/expect_issue.cr
@@ -47,6 +47,8 @@ require "./util"
 module Ameba::Spec::ExpectIssue
   include Spec::Util
 
+  class_property source : Source?
+
   def expect_issue(rules : Rule::Base | Enumerable(Rule::Base),
                    annotated_code : String,
                    path = "",
@@ -79,6 +81,39 @@ module Ameba::Spec::ExpectIssue
     end
   end
 
+  def expect_correction(correction, *, file = __FILE__, line = __LINE__)
+    source = ExpectIssue.source
+    raise "`expect_correction` must follow `expect_issue`" unless source
+
+    corrected_code = Source::Corrector.correct(source) # TODO: recursive
+    raise "Use `expect_no_corrections` if the code will not change" unless corrected_code
+    return if correction == corrected_code
+
+    fail <<-MSG, file, line
+      Expected correction:
+
+      #{correction}
+
+      Got:
+
+      #{corrected_code}
+      MSG
+  end
+
+  def expect_no_corrections(*, file = __FILE__, line = __LINE__)
+    source = ExpectIssue.source
+    raise "`expect_no_corrections` must follow `expect_offense`" unless source
+
+    corrected_code = Source::Corrector.correct(source)
+    return unless corrected_code
+
+    fail <<-MSG, file, line
+      Expected no corrections, but got:
+
+      #{corrected_code}
+      MSG
+  end
+
   def expect_no_issues(rules : Rule::Base | Enumerable(Rule::Base),
                        code : String,
                        path = "",
@@ -87,7 +122,7 @@ module Ameba::Spec::ExpectIssue
                        file = __FILE__,
                        line = __LINE__)
     code = normalize_code(code) if normalize
-    lines = code.lines
+    lines = code.split('\n') # must preserve trailing newline
     actual_annotations = actual_annotations(rules, code, path, lines)
     unless actual_annotations.to_s == code
       fail <<-MSG, file, line
@@ -100,6 +135,7 @@ module Ameba::Spec::ExpectIssue
 
   private def actual_annotations(rules, code, path, lines)
     source = Source.new(code, path, normalize: false) # already normalized
+    ExpectIssue.source = source
     if rules.is_a?(Enumerable)
       rules.each(&.catch(source))
     else