AddTimeUnitArgument recipe

src/main/java/org/openrewrite/java/apache/httpclient5/AddTimeUnitArgument.java

@@ -15,21 +15,22 @@
  */
 package org.openrewrite.java.apache.httpclient5;
 
-import lombok.AllArgsConstructor;
 import lombok.EqualsAndHashCode;
-import lombok.RequiredArgsConstructor;
 import lombok.Value;
 import org.openrewrite.ExecutionContext;
 import org.openrewrite.Option;
 import org.openrewrite.Recipe;
 import org.openrewrite.TreeVisitor;
+import org.openrewrite.internal.StringUtils;
 import org.openrewrite.internal.lang.Nullable;
 import org.openrewrite.java.JavaIsoVisitor;
 import org.openrewrite.java.JavaParser;
 import org.openrewrite.java.JavaTemplate;
 import org.openrewrite.java.MethodMatcher;
 import org.openrewrite.java.tree.J;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
 
 
@@ -68,21 +69,25 @@ public String getDescription() {
     public TreeVisitor<?, ExecutionContext> getVisitor() {
         return new JavaIsoVisitor<ExecutionContext>() {
             final MethodMatcher matcher = new MethodMatcher(methodPattern);
-            final JavaTemplate template = JavaTemplate.builder("#{any(long)}, TimeUnit.#{}")
-                    .contextSensitive()
-                    .javaParser(JavaParser.fromJavaVersion().classpath("httpclient5", "httpcore5"))
-                    .imports("java.util.concurrent.TimeUnit")
-                    .build();
 
             @Override
             public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, ExecutionContext executionContext) {
                 J.MethodInvocation m = super.visitMethodInvocation(method, executionContext);
-                if (matcher.matches(method) && method.getArguments().size() == 1) {
+                if (matcher.matches(m)) {
+                    JavaTemplate template = JavaTemplate
+                            .builder(StringUtils.repeat("#{any()}, ", m.getArguments().size()) + "TimeUnit.#{}")
+                            .contextSensitive()
+                            .javaParser(JavaParser.fromJavaVersion().classpath("httpclient5", "httpcore5"))
+                            .imports("java.util.concurrent.TimeUnit")
+                            .build();
+
+                    List<Object> arguments = new ArrayList<>(m.getArguments());
+                    arguments.add(timeUnit != null ? timeUnit : TimeUnit.MILLISECONDS);
+
                     m = template.apply(
                             updateCursor(m),
                             m.getCoordinates().replaceArguments(),
-                            m.getArguments().get(0),
-                            timeUnit != null ? timeUnit : TimeUnit.MILLISECONDS
+                            arguments.toArray(new Object[0])
                     );
                     maybeAddImport("java.util.concurrent.TimeUnit");
                 }

src/test/java/org/openrewrite/java/apache/httpclient5/AddTimeUnitArgumentTest.java

@@ -117,7 +117,7 @@ void test() {
     }
 
     @Test
-    void doesNotModifyMethodsWithMoreThanOneArgument() {
+    void doesModifyMethodsWithMoreThanOneArgument() {
         rewriteRun(
           spec -> spec.recipe(new AddTimeUnitArgument("A method(int, float)", null)),
           //language=java
@@ -126,16 +126,19 @@ void doesNotModifyMethodsWithMoreThanOneArgument() {
 
             class A {
                 private long value;
+                private float foo;
                 private TimeUnit timeunit;
 
                 A method(int value, float foo) {
                     this.value = value;
+                    this.foo = foo;
                     this.timeunit = TimeUnit.MILLISECONDS;
                     return this;
                 }
 
                 A method(long value, float foo, TimeUnit timeunit) {
                     this.value = value;
+                    this.foo = foo;
                     this.timeunit = timeunit;
                     return this;
                 }
@@ -149,6 +152,15 @@ void test() {
                     a.method(100, 1.0f);
                 }
             }
+            """, """
+            import java.util.concurrent.TimeUnit;
+
+            class B {
+                void test() {
+                    A a = new A();
+                    a.method(100, 1.0f, TimeUnit.MILLISECONDS);
+                }
+            }
             """)
         );
     }