Applicable tests

Author: jkschneider

src/main/java/org/openrewrite/java/testing/junit5/CleanupJUnitImports.java

@@ -18,7 +18,10 @@
 import org.openrewrite.ExecutionContext;
 import org.openrewrite.Recipe;
 import org.openrewrite.TreeVisitor;
+import org.openrewrite.internal.lang.Nullable;
 import org.openrewrite.java.JavaIsoVisitor;
+import org.openrewrite.java.JavaVisitor;
+import org.openrewrite.java.search.UsesType;
 import org.openrewrite.java.tree.J;
 
 public class CleanupJUnitImports extends Recipe {
@@ -37,6 +40,19 @@ protected TreeVisitor<?, ExecutionContext> getVisitor() {
         return new CleanupJUnitImportsVisitor();
     }
 
+    @Nullable
+    @Override
+    protected JavaVisitor<ExecutionContext> getSingleSourceApplicableTest() {
+        return new JavaVisitor<ExecutionContext>() {
+            @Override
+            public J visitCompilationUnit(J.CompilationUnit cu, ExecutionContext ctx) {
+                doAfterVisit(new UsesType<>("org.junit.*"));
+                doAfterVisit(new UsesType<>("junit.*"));
+                return cu;
+            }
+        };
+    }
+
     public static class CleanupJUnitImportsVisitor extends JavaIsoVisitor<ExecutionContext> {
         @Override
         public J.CompilationUnit visitCompilationUnit(J.CompilationUnit cu, ExecutionContext ctx) {

src/main/java/org/openrewrite/java/testing/junit5/MigrateJUnitTestCase.java

@@ -76,15 +76,14 @@ protected TreeVisitor<?, ExecutionContext> getSingleSourceApplicableTest() {
 
             @Override
             public J.CompilationUnit visitCompilationUnit(J.CompilationUnit cu, ExecutionContext executionContext) {
-                J.CompilationUnit c = cu;
-                for (J.ClassDeclaration clazz : c.getClasses()) {
+                for (J.ClassDeclaration clazz : cu.getClasses()) {
                     if (TypeUtils.isAssignableTo(JavaType.Class.build("junit.framework.TestCase"), clazz.getType())) {
-                        c = c.withMarkers(c.getMarkers().addIfAbsent(FOUND_TYPE));
+                        return cu.withMarkers(cu.getMarkers().addIfAbsent(FOUND_TYPE));
                     }
                 }
 
                 doAfterVisit(new UsesType<>("junit.framework.TestCase"));
-                return c;
+                return cu;
             }
         };
     }

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

@@ -18,8 +18,10 @@
 import org.openrewrite.ExecutionContext;
 import org.openrewrite.Recipe;
 import org.openrewrite.TreeVisitor;
+import org.openrewrite.internal.lang.Nullable;
 import org.openrewrite.java.JavaIsoVisitor;
 import org.openrewrite.java.OrderImports;
+import org.openrewrite.java.search.UsesType;
 import org.openrewrite.java.tree.J;
 
 /**
@@ -41,6 +43,12 @@ protected TreeVisitor<?, ExecutionContext> getVisitor() {
         return new CleanupMockitoImportsVisitor();
     }
 
+    @Nullable
+    @Override
+    protected TreeVisitor<?, ExecutionContext> getSingleSourceApplicableTest() {
+        return new UsesType<>("org.mockito.*");
+    }
+
     public static class CleanupMockitoImportsVisitor extends JavaIsoVisitor<ExecutionContext> {
         @Override
         public J.Import visitImport(J.Import _import, ExecutionContext executionContext) {