Use context sensitive templates where argument can be custom

Fixes #491

Author: timtebeek

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

@@ -67,6 +67,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
 
             if (args.size() == 2) {
                 return JavaTemplate.builder("assertThat(#{any()}).isEqualTo(#{any()});")
+                        .contextSensitive()
                         .staticImports("org.assertj.core.api.Assertions.assertThat")
                         .javaParser(JavaParser.fromJavaVersion().classpathFromResources(ctx, "assertj-core-3.24"))
                         .build()
@@ -77,6 +78,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
                         JavaTemplate.builder("assertThat(#{any()}).as(#{any(String)}).isEqualTo(#{any()});") :
                         JavaTemplate.builder("assertThat(#{any()}).as(#{any(java.util.function.Supplier)}).isEqualTo(#{any()});");
                 return template
+                        .contextSensitive()
                         .staticImports("org.assertj.core.api.Assertions.assertThat")
                         .imports("java.util.function.Supplier")
                         .javaParser(JavaParser.fromJavaVersion().classpathFromResources(ctx, "assertj-core-3.24"))
@@ -91,6 +93,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
             } else if (args.size() == 3) {
                 maybeAddImport("org.assertj.core.api.Assertions", "within");
                 return JavaTemplate.builder("assertThat(#{any()}).isCloseTo(#{any()}, within(#{any()}));")
+                        .contextSensitive()
                         .staticImports("org.assertj.core.api.Assertions.assertThat", "org.assertj.core.api.Assertions.within")
                         .javaParser(JavaParser.fromJavaVersion().classpathFromResources(ctx, "assertj-core-3.24"))
                         .build()
@@ -106,6 +109,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
                     JavaTemplate.builder("assertThat(#{any()}).as(#{any(String)}).isCloseTo(#{any()}, within(#{any()}));") :
                     JavaTemplate.builder("assertThat(#{any()}).as(#{any(java.util.function.Supplier)}).isCloseTo(#{any()}, within(#{any()}));");
             return template
+                    .contextSensitive()
                     .staticImports("org.assertj.core.api.Assertions.assertThat", "org.assertj.core.api.Assertions.within")
                     .imports("java.util.function.Supplier")
                     .javaParser(JavaParser.fromJavaVersion().classpathFromResources(ctx, "assertj-core-3.24"))

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

@@ -74,6 +74,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
 
             if (args.size() == 2) {
                 method = JavaTemplate.builder("assertThat(#{any()}).isNotEqualTo(#{any()});")
+                        .contextSensitive()
                         .staticImports("org.assertj.core.api.Assertions.assertThat")
                         .javaParser(assertionsParser(ctx))
                         .build()
@@ -92,6 +93,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
 
 
                 method = template
+                        .contextSensitive()
                         .staticImports("org.assertj.core.api.Assertions.assertThat")
                         .javaParser(assertionsParser(ctx))
                         .build()
@@ -104,6 +106,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
                         );
             } else if (args.size() == 3) {
                 method = JavaTemplate.builder("assertThat(#{any()}).isNotCloseTo(#{any()}, within(#{any()}));")
+                        .contextSensitive()
                         .staticImports("org.assertj.core.api.Assertions.assertThat", "org.assertj.core.api.Assertions.within")
                         .javaParser(assertionsParser(ctx))
                         .build()
@@ -123,6 +126,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
                         JavaTemplate.builder("assertThat(#{any()}).as(#{any(java.util.function.Supplier)}).isNotCloseTo(#{any()}, within(#{any()}));");
 
                 method = template
+                        .contextSensitive()
                         .staticImports("org.assertj.core.api.Assertions.assertThat", "org.assertj.core.api.Assertions.within")
                         .javaParser(assertionsParser(ctx))
                         .build()

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

@@ -71,6 +71,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
 
             if (args.size() == 1) {
                 method = JavaTemplate.builder("assertThat(#{any()}).isNotNull();")
+                        .contextSensitive()
                         .staticImports("org.assertj.core.api.Assertions.assertThat")
                         .javaParser(assertionsParser(ctx))
                         .build()
@@ -88,6 +89,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
                         JavaTemplate.builder("assertThat(#{any()}).as(#{any(java.util.function.Supplier)}).isNotNull();");
 
                 method = template
+                        .contextSensitive()
                         .staticImports("org.assertj.core.api.Assertions.assertThat")
                         .javaParser(assertionsParser(ctx))
                         .build()

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

@@ -71,6 +71,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
 
             if (args.size() == 1) {
                 method = JavaTemplate.builder("assertThat(#{any()}).isNull();")
+                        .contextSensitive()
                         .staticImports("org.assertj.core.api.Assertions.assertThat")
                         .javaParser(assertionsParser(ctx))
                         .build()
@@ -87,6 +88,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
                         JavaTemplate.builder("assertThat(#{any()}).as(#{any(java.util.function.Supplier)}).isNull();");
 
                 method = template
+                        .contextSensitive()
                         .staticImports("org.assertj.core.api.Assertions.assertThat")
                         .javaParser(assertionsParser(ctx))
                         .build()

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

@@ -72,6 +72,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
 
             if (args.size() == 2) {
                 method = JavaTemplate.builder("assertThat(#{any()}).isSameAs(#{any()});")
+                        .contextSensitive()
                         .staticImports("org.assertj.core.api.Assertions.assertThat")
                         .javaParser(assertionsParser(ctx))
                         .build()
@@ -89,6 +90,7 @@ public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, Execu
                         JavaTemplate.builder("assertThat(#{any()}).as(#{any(java.util.function.Supplier)}).isSameAs(#{any()});");
 
                 method = template
+                        .contextSensitive()
                         .staticImports("org.assertj.core.api.Assertions.assertThat")
                         .javaParser(assertionsParser(ctx))
                         .build()

src/test/java/org/openrewrite/java/testing/assertj/JUnitAssertNotNullToAssertThatTest.java

@@ -18,6 +18,7 @@
 import org.junit.jupiter.api.Test;
 import org.openrewrite.DocumentExample;
 import org.openrewrite.InMemoryExecutionContext;
+import org.openrewrite.Issue;
 import org.openrewrite.java.JavaParser;
 import org.openrewrite.test.RecipeSpec;
 import org.openrewrite.test.RewriteTest;
@@ -241,4 +242,44 @@ private String notification() {
           )
         );
     }
+
+    @Test
+    @Issue("https://github.com/openrewrite/rewrite-testing-frameworks/issues/491")
+    void importAddedForCustomArguments() {
+        rewriteRun(
+          //language=java
+          java(
+            """
+              import org.junit.jupiter.api.Test;
+              
+              import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+              class TTest {
+
+                  class A {}
+
+                  @Test
+                  public void testClass() {
+                      assertNotNull(new A());
+                  }
+              }
+              """,
+            """
+              import org.junit.jupiter.api.Test;
+              
+              import static org.assertj.core.api.Assertions.assertThat;
+
+              class TTest {
+
+                  class A {}
+
+                  @Test
+                  public void testClass() {
+                      assertThat(new A()).isNotNull();
+                  }
+              }
+              """
+          )
+        );
+    }
 }