Reduce duplication and migrate closeTo with three arguments (#389)

* Reduce duplication and allow closeTo with three arguments

* Extract a method to get arguments template

Author: timtebeek

src/main/java/org/openrewrite/java/testing/hamcrest/HamcrestMatcherToAssertJ.java

@@ -17,6 +17,7 @@
 
 import lombok.AllArgsConstructor;
 import lombok.NoArgsConstructor;
+import org.jetbrains.annotations.NotNull;
 import org.openrewrite.*;
 import org.openrewrite.internal.lang.Nullable;
 import org.openrewrite.java.JavaIsoVisitor;
@@ -75,76 +76,64 @@ private class MigrateToAssertJVisitor extends JavaIsoVisitor<ExecutionContext> {
         public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx) {
             J.MethodInvocation mi = super.visitMethodInvocation(method, ctx);
             if (assertThatMatcher.matches(mi)) {
-                if (mi.getArguments().size() == 2) {
-                    return handleTwoArgumentCase(mi, ctx);
-                }
-                if (mi.getArguments().size() == 3) {
-                    return handleThreeArgumentCase(mi, ctx);
-                }
+                return replace(mi, ctx);
             }
             return mi;
         }
 
-        private J.MethodInvocation handleTwoArgumentCase(J.MethodInvocation mi, ExecutionContext ctx) {
-            Expression actualArgument = mi.getArguments().get(0);
-            Expression matcherArgument = mi.getArguments().get(1);
+        private J.MethodInvocation replace(J.MethodInvocation mi, ExecutionContext ctx) {
+            List<Expression> mia = mi.getArguments();
+            Expression reasonArgument = mia.size() == 3 ? mia.get(0) : null;
+            Expression actualArgument = mia.get(mia.size() - 2);
+            Expression matcherArgument = mia.get(mia.size() - 1);
             if (!matchersMatcher.matches(matcherArgument) || subMatcher.matches(matcherArgument)) {
                 return mi;
             }
             String actual = typeToIndicator(actualArgument.getType());
-            List<Expression> originalArguments = ((J.MethodInvocation) matcherArgument).getArguments().stream()
-                    .filter(a -> !(a instanceof J.Empty))
-                    .collect(Collectors.toList());
-            String argumentsTemplate = originalArguments.stream()
-                    .map(a -> typeToIndicator(a.getType()))
-                    .collect(Collectors.joining(", "));
-            argumentsTemplate = applySpecialCases((J.MethodInvocation) matcherArgument, argumentsTemplate);
-
-            JavaTemplate template = JavaTemplate.builder(String.format("assertThat(%s).%s(%s)",
-                            actual, assertion, argumentsTemplate))
+            J.MethodInvocation matcherArgumentMethod = (J.MethodInvocation) matcherArgument;
+            JavaTemplate template = JavaTemplate.builder(String.format(
+                            "assertThat(%s)" +
+                            (reasonArgument != null ? ".as(#{any(String)})" : "") +
+                            ".%s(%s)",
+                            actual, assertion, getArgumentsTemplate(matcherArgumentMethod)))
                     .javaParser(JavaParser.fromJavaVersion().classpathFromResources(ctx, "assertj-core-3.24"))
-                    .staticImports("org.assertj.core.api.Assertions.assertThat", "org.assertj.core.api.Assertions.within")
+                    .staticImports(
+                            "org.assertj.core.api.Assertions.assertThat",
+                            "org.assertj.core.api.Assertions.within")
                     .build();
             maybeAddImport("org.assertj.core.api.Assertions", "assertThat");
             maybeAddImport("org.assertj.core.api.Assertions", "within");
             maybeRemoveImport("org.hamcrest.Matchers." + matcher);
+            maybeRemoveImport("org.hamcrest.MatcherAssert");
             maybeRemoveImport("org.hamcrest.MatcherAssert.assertThat");
 
             List<Expression> templateArguments = new ArrayList<>();
             templateArguments.add(actualArgument);
-            templateArguments.addAll(originalArguments);
+            if (reasonArgument != null) {
+                templateArguments.add(reasonArgument);
+            }
+            for (Expression originalArgument : matcherArgumentMethod.getArguments()) {
+                if (!(originalArgument instanceof J.Empty)) {
+                    templateArguments.add(originalArgument);
+                }
+            }
             return template.apply(getCursor(), mi.getCoordinates().replace(), templateArguments.toArray());
         }
 
-        private J.MethodInvocation handleThreeArgumentCase(J.MethodInvocation mi, ExecutionContext ctx) {
-            Expression reasonArgument = mi.getArguments().get(0);
-            Expression actualArgument = mi.getArguments().get(1);
-            Expression matcherArgument = mi.getArguments().get(2);
-            if (!matchersMatcher.matches(matcherArgument) || subMatcher.matches(matcherArgument)) {
-                return mi;
+        private final MethodMatcher CLOSE_TO_MATCHER = new MethodMatcher("org.hamcrest.Matchers closeTo(..)");
+
+        @NotNull
+        private String getArgumentsTemplate(J.MethodInvocation matcherArgument) {
+            List<Expression> methodArguments = matcherArgument.getArguments();
+            if (CLOSE_TO_MATCHER.matches(matcherArgument)) {
+                return String.format("%s, within(%s)",
+                        typeToIndicator(methodArguments.get(0).getType()),
+                        typeToIndicator(methodArguments.get(1).getType()));
             }
-            String actual = typeToIndicator(actualArgument.getType());
-            List<Expression> originalArguments = ((J.MethodInvocation) matcherArgument).getArguments().stream()
+            return methodArguments.stream()
                     .filter(a -> !(a instanceof J.Empty))
-                    .collect(Collectors.toList());
-            String argumentsTemplate = originalArguments.stream()
                     .map(a -> typeToIndicator(a.getType()))
                     .collect(Collectors.joining(", "));
-            JavaTemplate template = JavaTemplate.builder(String.format("assertThat(%s).as(#{any(String)}).%s(%s)",
-                            actual, assertion, argumentsTemplate))
-                    .javaParser(JavaParser.fromJavaVersion().classpathFromResources(ctx, "assertj-core-3.24"))
-                    .staticImports("org.assertj.core.api.Assertions.assertThat")
-                    .build();
-            maybeAddImport("org.assertj.core.api.Assertions", "assertThat");
-            maybeRemoveImport("org.hamcrest.Matchers." + matcher);
-            maybeRemoveImport("org.hamcrest.MatcherAssert");
-            maybeRemoveImport("org.hamcrest.MatcherAssert.assertThat");
-
-            List<Expression> templateArguments = new ArrayList<>();
-            templateArguments.add(actualArgument);
-            templateArguments.add(reasonArgument);
-            templateArguments.addAll(originalArguments);
-            return template.apply(getCursor(), mi.getCoordinates().replace(), templateArguments.toArray());
         }
 
         private String typeToIndicator(JavaType type) {
@@ -159,24 +148,5 @@ private String typeToIndicator(JavaType type) {
                 return String.format("#{any(%s)}", str);
             }
         }
-
-        private String applySpecialCases(J.MethodInvocation mi, String template) {
-            final MethodMatcher CLOSE_TO_MATCHER = new MethodMatcher("org.hamcrest.Matchers closeTo(..)");
-            String[] splitTemplate = template.split(", ");
-
-            if (CLOSE_TO_MATCHER.matches(mi)) {
-                List<String> newTemplateArr = new ArrayList<>();
-                for (int i = 0; i < splitTemplate.length; i++) {
-                    // within needs to placed on the second argument of isCloseTo
-                    if (i == 1) {
-                        newTemplateArr.add(String.format("within(%s)", splitTemplate[i]));
-                        continue;
-                    }
-                    newTemplateArr.add(splitTemplate[i]);
-                }
-                return String.join(", ", newTemplateArr);
-            }
-            return template;
-        }
     }
 }