Set comments on tree prefix

Using the Comment member on the prefix of tree node gives us more control of the formatting and location of the comments.
aws · May 10, 2024 · c7361a1 · c7361a1
1 parent 55da9ba
commit c7361a1
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 7 deletions.
diff --git a/...ol/src/main/java/software/amazon/awssdk/migration/internal/recipe/AddCommentToMethod.java b/...ol/src/main/java/software/amazon/awssdk/migration/internal/recipe/AddCommentToMethod.java
@@ -17,14 +17,19 @@
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import java.util.ArrayList;
+import java.util.List;
 import org.openrewrite.ExecutionContext;
 import org.openrewrite.Option;
 import org.openrewrite.Recipe;
 import org.openrewrite.TreeVisitor;
 import org.openrewrite.java.JavaIsoVisitor;
 import org.openrewrite.java.MethodMatcher;
+import org.openrewrite.java.tree.Comment;
 import org.openrewrite.java.tree.J;
-import org.openrewrite.marker.SearchResult;
+import org.openrewrite.java.tree.Space;
+import org.openrewrite.java.tree.TextComment;
+import org.openrewrite.marker.Markers;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 
 /**
@@ -69,10 +74,12 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
     private static final class Visitor extends JavaIsoVisitor<ExecutionContext> {
         private final MethodMatcher methodMatcher;
         private final String comment;
+        private final Comment commentToAdd;
 
         Visitor(String methodPattern, String comment) {
             this.methodMatcher = new MethodMatcher(methodPattern, false);
-            this.comment = COMMENT_PREFIX + comment + "\n";
+            this.comment = COMMENT_PREFIX + comment;
+            this.commentToAdd = new TextComment(true, this.comment, "", Markers.EMPTY);
         }
 
         @Override
@@ -83,7 +90,16 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation methodInvocat
                 return method;
             }
 
-            return SearchResult.found(method, comment);
+
+            Space prefix = methodInvocation.getPrefix();
+            List<Comment> comments = new ArrayList<>(prefix.getComments());
+
+            if (comments.contains(commentToAdd)) {
+                return method;
+            }
+
+            comments.add(commentToAdd);
+            return method.withPrefix(prefix.withComments(comments));
         }
     }
 }
diff --git a/...on-tool/src/test/java/software/amazon/awssdk/migration/recipe/AddCommentToMethodTest.java b/...on-tool/src/test/java/software/amazon/awssdk/migration/recipe/AddCommentToMethodTest.java
@@ -49,8 +49,7 @@ void shouldAddCommentToMethod() {
                 + "    \n"
                 + "    void test() {\n"
                 + "        ClientConfiguration clientConfiguration = new ClientConfiguration();\n"
-                + "        /*~~(AWS SDK for Java v2 migration: a comment\n"
-                + ")~~>*/clientConfiguration.setCacheResponseMetadata(false);\n"
+                + "        /*AWS SDK for Java v2 migration: a comment*/clientConfiguration.setCacheResponseMetadata(false);\n"
                 + "    }\n"
                 + "}"
             )
@@ -82,8 +81,7 @@ void hasExistingComments_shouldNotClobber() {
                 + "    void test() {\n"
                 + "        ClientConfiguration clientConfiguration = new ClientConfiguration();\n"
                 + "        // Existing comment \n"
-                + "        /*existing comment*/ /*~~(AWS SDK for Java v2 migration: a comment\n"
-                + ")~~>*/clientConfiguration.setCacheResponseMetadata(false);\n"
+                + "        /*existing comment*/ /*AWS SDK for Java v2 migration: a comment*/clientConfiguration.setCacheResponseMetadata(false);\n"
                 + "    }\n"
                 + "}"
             )