Only simplify chained AssertJ assertions with a single argument

openrewrite · Oct 11, 2024 · e8fca5a · e8fca5a
1 parent 8a95cb8
commit e8fca5a
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/src/main/java/org/openrewrite/java/testing/assertj/SimplifyChainedAssertJAssertion.java b/src/main/java/org/openrewrite/java/testing/assertj/SimplifyChainedAssertJAssertion.java
@@ -97,7 +97,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation methodInvocat
             J.MethodInvocation mi = super.visitMethodInvocation(methodInvocation, ctx);
 
             // assert has correct assertion
-            if (!ASSERT_TO_REPLACE.matches(mi)) {
+            if (!ASSERT_TO_REPLACE.matches(mi) || mi.getArguments().size() != 1) {
                 return mi;
             }
 

diff --git a/src/test/java/org/openrewrite/java/testing/assertj/SimplifyChainedAssertJAssertionTest.java b/src/test/java/org/openrewrite/java/testing/assertj/SimplifyChainedAssertJAssertionTest.java
@@ -319,7 +319,6 @@ void testMethod() {
         );
     }
 
-
     @Test
     void mapMethodDealsWithTwoArguments() {
         rewriteRun(
@@ -366,6 +365,28 @@ Map<String, String> getMap() {
         );
     }
 
+    @Test
+    void keySetContainsWithMultipleArguments() {
+        rewriteRun(
+          spec -> spec.recipe(new SimplifyChainedAssertJAssertion("keySet", "contains", "containsKey", "java.util.Map")),
+          //language=java
+          java(
+            """
+              import java.util.Map;
+
+              import static org.assertj.core.api.Assertions.assertThat;
+
+              class MyTest {
+                  void testMethod(Map<String, String> map) {
+                      // we don't yet support `containsKeys`
+                      assertThat(map.keySet()).contains("a", "b", "c");
+                  }
+              }
+              """
+          )
+        );
+    }
+
     @Test
     void isNotEmptyTest() {
         rewriteRun(