From e8fca5a0c8c76382e829f056146537008eb5f6d6 Mon Sep 17 00:00:00 2001
From: Tim te Beek <tim@moderne.io>
Date: Fri, 11 Oct 2024 13:46:49 +0200
Subject: [PATCH] Only simplify chained AssertJ assertions with a single
 argument

---
 .../SimplifyChainedAssertJAssertion.java      |  2 +-
 .../SimplifyChainedAssertJAssertionTest.java  | 23 ++++++++++++++++++-
 2 files changed, 23 insertions(+), 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
index e34076fed..ca514103f 100644
--- 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
index 785822116..8a2b7f42e 100644
--- 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(