Do not incorrectly strip method calls from assertion arguments (#918)

Author: timtebeek

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

@@ -34,7 +34,6 @@
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Objects;
 import java.util.Set;
 
 import static java.util.Collections.singleton;
@@ -170,7 +169,7 @@ private String getStringTemplateAndAppendArguments(J.MethodInvocation assertThat
                 }
 
                 // 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));
+                arguments.add(assertThatArgumentIsEmpty ? methodToReplaceArgument : assertThatArgument);
 
                 // Special case for Path.of() assertions
                 if ("java.nio.file.Path".equals(requiredType) && dedicatedAssertion.contains("Raw") &&
@@ -181,17 +180,6 @@ private String getStringTemplateAndAppendArguments(J.MethodInvocation assertThat
 
                 return "assertThat(#{any()}).%s(#{any()})";
             }
-
-            private 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 (chainedAssertMatcher.matches(assertThatArgument)) {
-                    return Objects.requireNonNull(((J.MethodInvocation) assertThatArgument).getSelect());
-                }
-                return assertThatArgument;
-            }
         };
     }
 }

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

@@ -269,6 +269,55 @@ void biscuits() {
         }
     }
 
+    @Issue("https://github.com/openrewrite/rewrite-testing-frameworks/issues/496")
+    @Test
+    void assertTrueContainsWithMethodCallArgPreservesMethodCall() {
+        rewriteRun(
+          spec -> spec.parser(JavaParser.fromJavaVersion().classpathFromResources(new InMemoryExecutionContext(), "junit-jupiter-api-5")),
+          //language=java
+          java(
+            """
+              import java.util.Locale;
+              import java.util.Set;
+
+              import static org.junit.jupiter.api.Assertions.assertTrue;
+
+              class MyTest {
+                  static class UnderTest {
+                      Locale asLocale() {
+                          return Locale.US;
+                      }
+                  }
+
+                  void test(Set<Locale> localeSet) {
+                      UnderTest underTest = new UnderTest();
+                      assertTrue(localeSet.contains(underTest.asLocale()));
+                  }
+              }
+              """,
+            """
+              import java.util.Locale;
+              import java.util.Set;
+
+              import static org.assertj.core.api.Assertions.assertThat;
+
+              class MyTest {
+                  static class UnderTest {
+                      Locale asLocale() {
+                          return Locale.US;
+                      }
+                  }
+
+                  void test(Set<Locale> localeSet) {
+                      UnderTest underTest = new UnderTest();
+                      assertThat(localeSet).contains(underTest.asLocale());
+                  }
+              }
+              """
+          )
+        );
+    }
+
     /**
      * Chained AssertJ assertions should be simplified to the corresponding dedicated assertion, as
      * per <a

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

@@ -711,6 +711,34 @@ void testMethod(Map<String, ?> map) {
         );
     }
 
+    @Test
+    void argumentMatchingChainedAssertionIsNotStripped() {
+        rewriteRun(
+          spec -> spec.recipe(new SimplifyChainedAssertJAssertion("equals", "isTrue", "isEqualTo", "java.lang.Object")),
+          //language=java
+          java(
+            """
+              import static org.assertj.core.api.Assertions.assertThat;
+
+              class MyTest {
+                  void testMethod(Object obj1, Object obj2, Object obj3) {
+                      assertThat(obj1.equals(obj2.equals(obj3))).isTrue();
+                  }
+              }
+              """,
+            """
+              import static org.assertj.core.api.Assertions.assertThat;
+
+              class MyTest {
+                  void testMethod(Object obj1, Object obj2, Object obj3) {
+                      assertThat(obj1).isEqualTo(obj2.equals(obj3));
+                  }
+              }
+              """
+          )
+        );
+    }
+
     @Test
     void sizeIsEqualToLongLiteralIsNotConverted() {
         rewriteRun(