Also remove EnableBatchProcessing import

Fixes #406
openrewrite · Aug 7, 2023 · 5514dc5 · 5514dc5
1 parent 3c42c95
commit 5514dc5
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/src/main/java/org/openrewrite/java/spring/boot3/RemoveEnableBatchProcessing.java b/src/main/java/org/openrewrite/java/spring/boot3/RemoveEnableBatchProcessing.java
@@ -16,11 +16,13 @@
 package org.openrewrite.java.spring.boot3;
 
 import org.openrewrite.ExecutionContext;
+import org.openrewrite.Preconditions;
 import org.openrewrite.Recipe;
 import org.openrewrite.TreeVisitor;
 import org.openrewrite.internal.ListUtils;
 import org.openrewrite.java.JavaIsoVisitor;
 import org.openrewrite.java.search.FindAnnotations;
+import org.openrewrite.java.search.UsesType;
 import org.openrewrite.java.tree.J;
 import org.openrewrite.java.tree.TypeUtils;
 
@@ -36,23 +38,25 @@ public String getDescription() {
         return "Add or remove the `@EnableBatchProcessing` annotation from a Spring Boot application.";
     }
 
+    private static final String ENABLE_BATCH_PROCESSING = "org.springframework.batch.core.configuration.annotation.EnableBatchProcessing";
+
     @Override
     public TreeVisitor<?, ExecutionContext> getVisitor() {
-        return new JavaIsoVisitor<ExecutionContext>() {
-
+        return Preconditions.check(new UsesType<>(ENABLE_BATCH_PROCESSING, true), new JavaIsoVisitor<ExecutionContext>() {
             @Override
             public J.ClassDeclaration visitClassDeclaration(J.ClassDeclaration classDecl, ExecutionContext ctx) {
                 if (!FindAnnotations.find(classDecl, "@org.springframework.boot.autoconfigure.SpringBootApplication").isEmpty() &&
-                    !FindAnnotations.find(classDecl, "@org.springframework.batch.core.configuration.annotation.EnableBatchProcessing").isEmpty()) {
+                    !FindAnnotations.find(classDecl, "@" + ENABLE_BATCH_PROCESSING).isEmpty()) {
                     return classDecl.withLeadingAnnotations(ListUtils.map(classDecl.getLeadingAnnotations(), a -> {
-                        if (TypeUtils.isOfClassType(a.getType(), "org.springframework.batch.core.configuration.annotation.EnableBatchProcessing")) {
+                        if (TypeUtils.isOfClassType(a.getType(), ENABLE_BATCH_PROCESSING)) {
+                            maybeRemoveImport(ENABLE_BATCH_PROCESSING);
                             return null;
                         }
                         return a;
                     }));
                 }
                 return super.visitClassDeclaration(classDecl, ctx);
             }
-        };
+        });
     }
 }
diff --git a/...pringBoot_3_0/java/org/openrewrite/java/spring/boot3/RemoveEnableBatchProcessingTest.java b/...pringBoot_3_0/java/org/openrewrite/java/spring/boot3/RemoveEnableBatchProcessingTest.java
@@ -23,7 +23,7 @@
 
 import static org.openrewrite.java.Assertions.java;
 
-public class RemoveEnableBatchProcessingTest implements RewriteTest {
+class RemoveEnableBatchProcessingTest implements RewriteTest {
 
     @Override
     public void defaults(RecipeSpec spec) {
@@ -47,7 +47,6 @@ public class Application {
               }
               """,
             """
-              import org.springframework.batch.core.configuration.annotation.EnableBatchProcessing;
               import org.springframework.boot.autoconfigure.SpringBootApplication;
                             
               @SpringBootApplication