Resolve issues with SimplifyChainedAssertJAssertion (#399)

* fixed issue, all tests pass

* Document failing test

* Do not extract select for method to select argument

* Collapse nested if statement

* Call hasZeroArgument once only

---------

Co-authored-by: Tim te Beek <[email protected]>

Author: AlekSimpson

src/main/java/org/openrewrite/java/testing/assertj/SimplifyChainedAssertJAssertion.java

@@ -79,7 +79,7 @@ public Set<String> getTags() {
     @Override
     public String getDescription() {
         return "Many AssertJ chained assertions have dedicated assertions that function the same. " +
-                "It is best to use the dedicated assertions.";
+               "It is best to use the dedicated assertions.";
     }
 
     @Override
@@ -108,19 +108,28 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation methodInvocat
             }
 
             J.MethodInvocation assertThatArg = (J.MethodInvocation) assertThat.getArguments().get(0);
-            if (!CHAINED_ASSERT_MATCHER.matches(assertThatArg) && !hasZeroArgument(mi)) {
+            if (!CHAINED_ASSERT_MATCHER.matches(assertThatArg)) {
                 return mi;
             }
 
-            // method call has select
-            Expression select = assertThatArg.getSelect() != null ? assertThatArg.getSelect() : assertThatArg;
-            if (!TypeUtils.isAssignableTo(requiredType, select.getType())) {
+            // Extract the actual argument for the new assertThat call
+            Expression actual = assertThatArg.getSelect() != null ? assertThatArg.getSelect() : assertThatArg;
+            if (!TypeUtils.isAssignableTo(requiredType, actual.getType())) {
                 return mi;
             }
+            List<Expression> arguments = new ArrayList<>();
+            arguments.add(actual);
 
-            List<Expression> arguments = new ArrayList<>(Collections.singletonList(select));
-            String template = getTemplate(arguments, assertThatArg, mi);
-            String formattedTemplate = String.format(template, dedicatedAssertion);
+            // Special case for more expressive assertions: assertThat(x.size()).isEqualTo(0) -> isEmpty()
+            if ("size".equals(chainedAssertion) && "isEqualTo".equals(assertToReplace) && hasZeroArgument(mi)) {
+                return applyTemplate("assertThat(#{any()}).isEmpty()", arguments, mi, ctx);
+            }
+
+            String template = getStringTemplateAndAppendArguments(assertThatArg, mi, arguments);
+            return applyTemplate(String.format(template, dedicatedAssertion), arguments, mi, ctx);
+        }
+
+        private J.MethodInvocation applyTemplate(String formattedTemplate, List<Expression> arguments, J.MethodInvocation mi, ExecutionContext ctx) {
             return JavaTemplate.builder(formattedTemplate)
                     .contextSensitive()
                     .javaParser(JavaParser.fromJavaVersion().classpathFromResources(ctx, "junit-jupiter-api-5.9", "assertj-core-3.24"))
@@ -129,33 +138,50 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation methodInvocat
         }
     }
 
-    private String getTemplate(List<Expression> arguments, J.MethodInvocation assertThatArg, J.MethodInvocation methodToReplace) {
-        if (hasZeroArgument(methodToReplace)) {
+    private String getStringTemplateAndAppendArguments(J.MethodInvocation assertThatArg, J.MethodInvocation methodToReplace, List<Expression> arguments) {
+        Expression assertThatArgument = assertThatArg.getArguments().get(0);
+        Expression methodToReplaceArgument = methodToReplace.getArguments().get(0);
+        boolean assertThatArgumentIsEmpty = assertThatArgument instanceof J.Empty;
+        boolean methodToReplaceArgumentIsEmpty = methodToReplaceArgument instanceof J.Empty;
+
+        // If both arguments are empty, then the select is already added to the arguments list, and we use a minimal template
+        if (assertThatArgumentIsEmpty && methodToReplaceArgumentIsEmpty) {
             return "assertThat(#{any()}).%s()";
         }
 
-        if (!(assertThatArg.getArguments().get(0) instanceof J.Empty)
-                && !(methodToReplace.getArguments().get(0) instanceof J.Empty)) {
-            // Note: this should be the only case when more than one argument needs to be handled. When the assertions involve the map functions
-            arguments.add(assertThatArg.getArguments().get(0));
-            arguments.add(methodToReplace.getArguments().get(0));
+        // If both arguments are not empty, then we add both to the arguments to the arguments list, and return a template with two arguments
+        if (!assertThatArgumentIsEmpty && !methodToReplaceArgumentIsEmpty) {
+            // This should only happen for map assertions using a key and value
+            arguments.add(assertThatArgument);
+            arguments.add(methodToReplaceArgument);
             return "assertThat(#{any()}).%s(#{any()}, #{any()})";
         }
 
-        if (!(assertThatArg.getArguments().get(0) instanceof J.Empty)
-                || !(methodToReplace.getArguments().get(0) instanceof J.Empty)) {
-            Expression argumentToAdd = assertThatArg.getArguments().get(0) instanceof J.Empty ? methodToReplace.getArguments().get(0) : assertThatArg.getArguments().get(0);
-            argumentToAdd = argumentToAdd instanceof J.MethodInvocation ? ((J.MethodInvocation) argumentToAdd).getSelect() : argumentToAdd;
-            arguments.add(argumentToAdd);
+        // If either argument is empty, we choose which one to add to the arguments list, and optionally extract the select
+        arguments.add(extractEitherArgument(assertThatArgumentIsEmpty, assertThatArgument, methodToReplaceArgument));
 
-            if (requiredType.equals("java.nio.file.Path") && dedicatedAssertion.contains("Raw")
-                    && TypeUtils.isAssignableTo("java.lang.String", assertThatArg.getArguments().get(0).getType())) {
-                return "assertThat(#{any()}).%s(Path.of(#{any()}))";
-            }
-            return "assertThat(#{any()}).%s(#{any()})";
+        // Special case for Path.of() assertions
+        if ("java.nio.file.Path".equals(requiredType) && dedicatedAssertion.contains("Raw")
+            && TypeUtils.isAssignableTo("java.lang.String", assertThatArgument.getType())) {
+            return "assertThat(#{any()}).%s(Path.of(#{any()}))";
         }
 
-        return "assertThat(#{any()}).%s()";
+        return "assertThat(#{any()}).%s(#{any()})";
+
+    }
+
+    private static Expression extractEitherArgument(boolean assertThatArgumentIsEmpty, Expression assertThatArgument, Expression methodToReplaceArgument) {
+        if (assertThatArgumentIsEmpty) {
+            return methodToReplaceArgument;
+        }
+        // Only on the assertThat argument do we possibly replace the argument with the select; such as list.size() -> list
+        if (assertThatArgument instanceof J.MethodInvocation) {
+            Expression select = ((J.MethodInvocation) assertThatArgument).getSelect();
+            if (select != null) {
+                return select;
+            }
+        }
+        return assertThatArgument;
     }
 
     private boolean hasZeroArgument(J.MethodInvocation method) {

src/main/resources/META-INF/rewrite/assertj.yml

@@ -258,7 +258,7 @@ recipeList:
   - org.openrewrite.java.testing.assertj.SimplifyChainedAssertJAssertion:
       chainedAssertion: size
       assertToReplace: isEqualTo
-      dedicatedAssertion: hasSameSizeAs
+      dedicatedAssertion: hasSize
       requiredType: java.util.Map
   - org.openrewrite.java.testing.assertj.SimplifyChainedAssertJAssertion:
       chainedAssertion: containsKey