SimplifyChainedAssertJAssertion drops method calls in replacement (#629)

* Add test for maintaining method call after replacement

* Only get select when method name matches chained assertion

* Reuse CHAINED_ASSERT_MATCHER to check argument unpack

---------

Co-authored-by: Tim te Beek <tim@moderne.io>

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

@@ -30,10 +30,7 @@
 import org.openrewrite.java.tree.J;
 import org.openrewrite.java.tree.TypeUtils;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-import java.util.Set;
+import java.util.*;
 
 @AllArgsConstructor
 @NoArgsConstructor
@@ -163,19 +160,16 @@ private String getStringTemplateAndAppendArguments(J.MethodInvocation assertThat
 
             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;
+        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 (CHAINED_ASSERT_MATCHER.matches(assertThatArgument)) {
+                return Objects.requireNonNull(((J.MethodInvocation) assertThatArgument).getSelect());
             }
+            return assertThatArgument;
         }
-        return assertThatArgument;
     }
 }