Isolate JavaTemplate issue with generic type parameters

- For openrewrite/rewrite-spring#284
openrewrite · Aug 2, 2024 · 45ca436 · 45ca436
1 parent 65fb0cf
commit 45ca436
Showing 1 changed file with 62 additions and 12 deletions.
diff --git a/rewrite-java-test/src/test/java/org/openrewrite/java/JavaTemplateTest.java b/rewrite-java-test/src/test/java/org/openrewrite/java/JavaTemplateTest.java
@@ -1200,19 +1200,69 @@ void finalMethodParameter() {
         rewriteRun(
           spec -> spec.recipe(new ReplaceAnnotation("@org.jetbrains.annotations.NotNull", "@lombok.NonNull", null)),
           java(
-                """
-                    import org.jetbrains.annotations.NotNull;
-                    
-                    class A {
-                        String testMethod(@NotNull final String test) {}
-                    }
-                    """, """
-                    import lombok.NonNull;
-                    
-                    class A {
-                        String testMethod(@NonNull final String test) {}
+            """
+              import org.jetbrains.annotations.NotNull;
+              
+              class A {
+                  String testMethod(@NotNull final String test) {}
+              }
+              """, """
+              import lombok.NonNull;
+              
+              class A {
+                  String testMethod(@NonNull final String test) {}
+              }
+              """)
+        );
+    }
+
+    @Test
+    @Issue("https://github.com/openrewrite/rewrite-spring/pull/284")
+    void replaceMethodInChainFollowedByGenericTypeParameters() {
+        rewriteRun(
+          spec -> spec.recipe(toRecipe(() -> new JavaVisitor<>() {
+                @Override
+                public J visitMethodInvocation(J.MethodInvocation method, ExecutionContext executionContext) {
+                    if (new MethodMatcher("batch.StepBuilder create()").matches(method)) {
+                        return JavaTemplate.builder("new StepBuilder()")
+                          .contextSensitive()
+                          .build()
+                          .apply(getCursor(), method.getCoordinates().replace());
                     }
-                    """)
+                    return super.visitMethodInvocation(method, executionContext);
+                }
+            }))
+            .parser(JavaParser.fromJavaVersion().dependsOn(
+                """
+                  package batch;
+                  public class StepBuilder {
+                      public static StepBuilder create() { return new StepBuilder(); }
+                      public StepBuilder() {}
+                      public <T> T method() { return null; }
+                  }
+                  """
+              )
+            ),
+          java(
+            """
+              import batch.StepBuilder;
+              class Foo {
+                  void test() {
+                      StepBuilder.create()
+                          .<String>method();
+                  }
+              }
+              """,
+            """
+              import batch.StepBuilder;
+              class Foo {
+                  void test() {
+                      new StepBuilder()
+                          .<String>method();
+                  }
+              }
+              """
+          )
         );
     }
 }