Support intermediate methods in SimplifyChainedAssertJAssertion (#812) (#962)

Walk past .as(), .describedAs(), .withFailMessage(), and
.overridingErrorMessage() when matching chained assertions, then splice
them back into the transformed result.

Closes #812

Author: timtebeek

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

@@ -33,6 +33,7 @@
 import org.openrewrite.java.tree.TypeUtils;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Set;
 
@@ -85,6 +86,12 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
         MethodMatcher assertThatMatcher = new MethodMatcher("org.assertj.core.api.Assertions assertThat(..)");
         MethodMatcher chainedAssertMatcher = new MethodMatcher("java..* " + chainedAssertion + "(..)");
         MethodMatcher assertToReplace = new MethodMatcher("org.assertj.core.api.* " + this.assertToReplace + "(..)");
+        List<MethodMatcher> intermediateMatchers = Arrays.asList(
+                new MethodMatcher("org.assertj.core.api.* as(..)"),
+                new MethodMatcher("org.assertj.core.api.* describedAs(..)"),
+                new MethodMatcher("org.assertj.core.api.* withFailMessage(..)"),
+                new MethodMatcher("org.assertj.core.api.* overridingErrorMessage(..)")
+        );
 
         return new JavaIsoVisitor<ExecutionContext>() {
             @Override
@@ -96,9 +103,19 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation methodInvocat
                     return mi;
                 }
 
-                // assertThat has method call
-                J.MethodInvocation assertThat = (J.MethodInvocation) mi.getSelect();
-                if (!assertThatMatcher.matches(assertThat) || !(assertThat.getArguments().get(0) instanceof J.MethodInvocation)) {
+                // Walk past intermediate methods (as, describedAs, etc.) to find assertThat
+                List<J.MethodInvocation> intermediates = new ArrayList<>();
+                J.MethodInvocation current = (J.MethodInvocation) mi.getSelect();
+                while (!assertThatMatcher.matches(current)) {
+                    if (isIntermediate(current) && current.getSelect() instanceof J.MethodInvocation) {
+                        intermediates.add(current);
+                        current = (J.MethodInvocation) current.getSelect();
+                    } else {
+                        return mi;
+                    }
+                }
+                J.MethodInvocation assertThat = current;
+                if (!(assertThat.getArguments().get(0) instanceof J.MethodInvocation)) {
                     return mi;
                 }
 
@@ -142,11 +159,30 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation methodInvocat
                 arguments.add(actual);
 
                 String template = getStringTemplateAndAppendArguments(assertThatArg, mi, arguments);
-                return JavaTemplate.builder(String.format(template, dedicatedAssertion))
+                J.MethodInvocation result = JavaTemplate.builder(String.format(template, dedicatedAssertion))
                         .contextSensitive()
                         .javaParser(JavaParser.fromJavaVersion().classpathFromResources(ctx, "junit-jupiter-api-5", "assertj-core-3"))
                         .build()
                         .apply(getCursor(), mi.getCoordinates().replace(), arguments.toArray());
+
+                // Splice intermediate methods (as, describedAs, etc.) back into the chain
+                if (!intermediates.isEmpty()) {
+                    Expression chain = result.getSelect();
+                    for (int i = intermediates.size() - 1; i >= 0; i--) {
+                        chain = intermediates.get(i).withSelect(chain);
+                    }
+                    result = result.withSelect(chain);
+                }
+                return result;
+            }
+
+            private boolean isIntermediate(J.MethodInvocation method) {
+                for (MethodMatcher matcher : intermediateMatchers) {
+                    if (matcher.matches(method)) {
+                        return true;
+                    }
+                }
+                return false;
             }
 
             private String getStringTemplateAndAppendArguments(J.MethodInvocation assertThatArg, J.MethodInvocation methodToReplace, List<Expression> arguments) {

src/test/java/org/openrewrite/java/testing/assertj/SimplifyChainedAssertJAssertionTest.java

@@ -15,7 +15,6 @@
  */
 package org.openrewrite.java.testing.assertj;
 
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
 import org.openrewrite.DocumentExample;
@@ -62,7 +61,6 @@ void testMethod() {
         );
     }
 
-    @Disabled(".as(reason) is not yet supported")
     @Test
     void stringIsEmptyDescribedAs() {
         rewriteRun(
@@ -91,6 +89,94 @@ void testMethod(String actual) {
         );
     }
 
+    @Test
+    void collectionSizeWithAs() {
+        rewriteRun(
+          spec -> spec.recipe(new SimplifyChainedAssertJAssertion("size", "isEqualTo", "hasSize", "java.util.Collection")),
+          //language=java
+          java(
+            """
+              import java.util.List;
+
+              import static org.assertj.core.api.Assertions.assertThat;
+
+              class MyTest {
+                  void testMethod(List<String> list) {
+                      assertThat(list.size()).as("Expected size to match").isEqualTo(5);
+                  }
+              }
+              """,
+            """
+              import java.util.List;
+
+              import static org.assertj.core.api.Assertions.assertThat;
+
+              class MyTest {
+                  void testMethod(List<String> list) {
+                      assertThat(list).as("Expected size to match").hasSize(5);
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void stringIsEmptyWithDescribedAs() {
+        rewriteRun(
+          spec -> spec.recipe(new SimplifyChainedAssertJAssertion("isEmpty", "isTrue", "isEmpty", "java.lang.String")),
+          //language=java
+          java(
+            """
+              import static org.assertj.core.api.Assertions.assertThat;
+
+              class MyTest {
+                  void testMethod(String actual) {
+                      assertThat(actual.isEmpty()).describedAs("Reason").isTrue();
+                  }
+              }
+              """,
+            """
+              import static org.assertj.core.api.Assertions.assertThat;
+
+              class MyTest {
+                  void testMethod(String actual) {
+                      assertThat(actual).describedAs("Reason").isEmpty();
+                  }
+              }
+              """
+          )
+        );
+    }
+
+    @Test
+    void multipleIntermediates() {
+        rewriteRun(
+          spec -> spec.recipe(new SimplifyChainedAssertJAssertion("isEmpty", "isTrue", "isEmpty", "java.lang.String")),
+          //language=java
+          java(
+            """
+              import static org.assertj.core.api.Assertions.assertThat;
+
+              class MyTest {
+                  void testMethod(String actual) {
+                      assertThat(actual.isEmpty()).as("check").withFailMessage("failed").isTrue();
+                  }
+              }
+              """,
+            """
+              import static org.assertj.core.api.Assertions.assertThat;
+
+              class MyTest {
+                  void testMethod(String actual) {
+                      assertThat(actual).as("check").withFailMessage("failed").isEmpty();
+                  }
+              }
+              """
+          )
+        );
+    }
+
     @Test
     void chainedRecipes() {
         rewriteRun(