Apply Kotlin style class reference in AddMockitoExtensionIfAnnotationsUsed (#890)

* Apply Kotlin style class reference on Kotlin source files

* Minor improvement

* Use `preVisit` to stop another unnecessary traversal

* Remove unused imports

* Remove excess newline

---------

Co-authored-by: Tim te Beek <tim@moderne.io>

Author: Laurens-W

src/main/java/org/openrewrite/java/testing/mockito/AddMockitoExtensionIfAnnotationsUsed.java

@@ -13,11 +13,12 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 package org.openrewrite.java.testing.mockito;
 
+import org.jspecify.annotations.Nullable;
 import org.openrewrite.ExecutionContext;
 import org.openrewrite.Recipe;
+import org.openrewrite.Tree;
 import org.openrewrite.TreeVisitor;
 import org.openrewrite.java.JavaIsoVisitor;
 import org.openrewrite.java.JavaParser;
@@ -26,6 +27,10 @@
 import org.openrewrite.java.search.FindTypes;
 import org.openrewrite.java.search.IsLikelyTest;
 import org.openrewrite.java.tree.J;
+import org.openrewrite.kotlin.KotlinIsoVisitor;
+import org.openrewrite.kotlin.KotlinParser;
+import org.openrewrite.kotlin.KotlinTemplate;
+import org.openrewrite.kotlin.tree.K;
 
 import static java.util.Comparator.comparing;
 import static org.openrewrite.Preconditions.*;
@@ -60,20 +65,52 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
                         // prevent addition if present
                         not(hasExtendedWithAnnotation),
                         or(hasAnyMockitoAnnotation)),
-                new JavaIsoVisitor<ExecutionContext>() {
+                new TreeVisitor<Tree, ExecutionContext>() {
                     @Override
-                    public J.ClassDeclaration visitClassDeclaration(J.ClassDeclaration classDecl, ExecutionContext ctx) {
-
-                        maybeAddImport("org.mockito.junit.jupiter.MockitoExtension");
-                        maybeAddImport("org.junit.jupiter.api.extension.ExtendWith");
-
-                        return JavaTemplate.builder("@ExtendWith(MockitoExtension.class)")
-                                .imports("org.mockito.junit.jupiter.MockitoExtension")
-                                .imports("org.junit.jupiter.api.extension.ExtendWith")
-                                .javaParser(JavaParser.fromJavaVersion().classpathFromResources(ctx, "junit-jupiter-api", "mockito-junit-jupiter"))
-                                .build()
-                                .apply(getCursor(), classDecl.getCoordinates().addAnnotation(comparing(J.Annotation::getSimpleName)));
+                    public @Nullable Tree preVisit(Tree tree, ExecutionContext ctx) {
+                        stopAfterPreVisit();
+                        if (tree instanceof J.CompilationUnit) {
+                            return getJavaVisitor().visit(tree, ctx);
+                        }
+                        if (tree instanceof K.CompilationUnit) {
+                            return getKotlinVisitor().visit(tree, ctx);
+                        }
+                        return tree;
                     }
                 });
     }
+
+    private JavaIsoVisitor<ExecutionContext> getJavaVisitor() {
+        return new JavaIsoVisitor<ExecutionContext>() {
+            @Override
+            public J.ClassDeclaration visitClassDeclaration(J.ClassDeclaration classDecl, ExecutionContext ctx) {
+                maybeAddImport("org.mockito.junit.jupiter.MockitoExtension");
+                maybeAddImport("org.junit.jupiter.api.extension.ExtendWith");
+
+                return JavaTemplate.builder("@ExtendWith(MockitoExtension.class)")
+                        .imports("org.mockito.junit.jupiter.MockitoExtension")
+                        .imports("org.junit.jupiter.api.extension.ExtendWith")
+                        .javaParser(JavaParser.fromJavaVersion().classpathFromResources(ctx, "junit-jupiter-api", "mockito-junit-jupiter"))
+                        .build()
+                        .apply(getCursor(), classDecl.getCoordinates().addAnnotation(comparing(J.Annotation::getSimpleName)));
+            }
+        };
+    }
+
+    private KotlinIsoVisitor<ExecutionContext> getKotlinVisitor() {
+        return new KotlinIsoVisitor<ExecutionContext>() {
+            @Override
+            public J.ClassDeclaration visitClassDeclaration(J.ClassDeclaration classDecl, ExecutionContext ctx) {
+                maybeAddImport("org.mockito.junit.jupiter.MockitoExtension");
+                maybeAddImport("org.junit.jupiter.api.extension.ExtendWith");
+
+                return KotlinTemplate.builder("@ExtendWith(MockitoExtension::class)")
+                        .imports("org.mockito.junit.jupiter.MockitoExtension")
+                        .imports("org.junit.jupiter.api.extension.ExtendWith")
+                        .parser(KotlinParser.builder().classpathFromResources(ctx, "junit-jupiter-api", "mockito-junit-jupiter"))
+                        .build()
+                        .apply(getCursor(), classDecl.getCoordinates().addAnnotation(comparing(J.Annotation::getSimpleName)));
+            }
+        };
+    }
 }

src/test/java/org/openrewrite/java/testing/mockito/AddMockitoExtensionIfAnnotationsUsedTest.java

@@ -34,7 +34,10 @@ public void defaults(RecipeSpec spec) {
         spec.recipe(new AddMockitoExtensionIfAnnotationsUsed())
           .parser(JavaParser.fromJavaVersion()
             .classpathFromResources(new InMemoryExecutionContext(), "junit-jupiter-api", "mockito-junit-jupiter", "mockito-core")
-            .dependsOn("public class Service {}"));
+            .dependsOn("public class Service {}"))
+          .parser(KotlinParser.builder()
+            .classpathFromResources(new InMemoryExecutionContext(), "junit-jupiter-api", "mockito-junit-jupiter", "mockito-core")
+            .dependsOn("class Service"));
     }
 
     @DocumentExample
@@ -214,4 +217,40 @@ fun test() {}
           )
         );
     }
+
+    @Test
+    void addForMockKotlin() {
+        rewriteRun(
+          //language=java
+          kotlin(
+            """
+              import org.junit.jupiter.api.Test
+              import org.mockito.Mock
+
+              class Test {
+                  @Mock
+                  lateinit var service: Service
+
+                  @Test
+                  fun test() {}
+              }
+              """,
+            """
+              import org.junit.jupiter.api.Test
+              import org.junit.jupiter.api.extension.ExtendWith
+              import org.mockito.Mock
+              import org.mockito.junit.jupiter.MockitoExtension
+
+              @ExtendWith(MockitoExtension::class)
+              class Test {
+                  @Mock
+                  lateinit var service: Service
+
+                  @Test
+                  fun test() {}
+              }
+              """
+          )
+        );
+    }
 }