diff --git a/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/MigrateDataProvider.java b/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/MigrateDataProvider.java
index 5a8e9654..e4e1408c 100644
--- a/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/MigrateDataProvider.java
+++ b/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/MigrateDataProvider.java
@@ -80,12 +80,12 @@ private static class WrapDataProviderMethod extends JavaIsoVisitor<ExecutionCont
         public J.ClassDeclaration visitClassDeclaration(J.ClassDeclaration classDecl, org.openrewrite.ExecutionContext ctx) {
             classDecl = super.visitClassDeclaration(classDecl, ctx);
 
-            Set<J.MethodDeclaration> dataProviders = FindAnnotatedMethods.find(classDecl, DATA_PROVIDER_MATCHER);
+            final Set<J.MethodDeclaration> dataProviders = FindAnnotatedMethods.find(classDecl, DATA_PROVIDER_MATCHER);
 
             // for each add a Wrapper that translates to Jupiter method source
             for (J.MethodDeclaration provider : dataProviders) {
-                String providerMethodName = provider.getSimpleName();
-                String providerName = FindAnnotation.find(provider, DATA_PROVIDER_MATCHER).stream().findAny()
+                final String providerMethodName = provider.getSimpleName();
+                final String providerName = FindAnnotation.find(provider, DATA_PROVIDER_MATCHER).stream().findAny()
                         .flatMap(j -> AnnotationArguments.extractLiteral(j, "name", String.class))
                         .orElse(providerMethodName);
 
@@ -109,35 +109,33 @@ public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, Ex
             method = super.visitMethodDeclaration(method, ctx);
 
             // if @ParameterizedTest is used, skip
-            Optional<J.Annotation> parameterizedTestAnnotation = FindAnnotation.findFirst(method, new AnnotationMatcher("@org.junit.jupiter.params.ParameterizedTest"));
+            final Optional<J.Annotation> parameterizedTestAnnotation = 
+                FindAnnotation.findFirst(method, new AnnotationMatcher("@org.junit.jupiter.params.ParameterizedTest"));
             if (parameterizedTestAnnotation.isPresent()) {
                 return method;
             }
 
             // if no TestNG @Test present, skip
-            Optional<J.Annotation> testNgAnnotation = FindAnnotation.findFirst(method, new AnnotationMatcher("@org.testng.annotations.Test"));
+            final Optional<J.Annotation> testNgAnnotation = 
+                FindAnnotation.findFirst(method, new AnnotationMatcher("@org.testng.annotations.Test"));
             if (!testNgAnnotation.isPresent()) {
                 return method;
             }
 
             // determine if a parameterized test is applicable
-            Optional<String> dataProviderMethodName = AnnotationArguments.extractLiteral(testNgAnnotation.get(), "dataProvider", String.class);
+            final Optional<String> dataProviderMethodName = 
+                AnnotationArguments.extractLiteral(testNgAnnotation.get(), "dataProvider", String.class);
             if (!dataProviderMethodName.isPresent()) {
                 return method;
             }
 
-            JavaCoordinates addAnnotationCoordinate = method.getCoordinates().addAnnotation((a, b) -> 1);
-
-            method = JavaTemplate
+            maybeAddImport("org.junit.jupiter.params.ParameterizedTest");
+            return JavaTemplate
                     .builder("@ParameterizedTest")
                     .javaParser(JavaParser.fromJavaVersion().classpath("junit-jupiter-params"))
                     .imports("org.junit.jupiter.params.ParameterizedTest")
                     .build()
-                    .apply(getCursor(), addAnnotationCoordinate);
-
-            maybeAddImport("org.junit.jupiter.params.ParameterizedTest");
-
-            return method;
+                    .apply(getCursor(), method.getCoordinates().addAnnotation((a, b) -> 1));
         }
     }
 
@@ -147,25 +145,25 @@ public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, Ex
             method = super.visitMethodDeclaration(method, ctx);
 
             // if @MethodSource is used, skip
-            Optional<J.Annotation> methodSourceAnnotation = FindAnnotation.findFirst(method, new AnnotationMatcher("@org.junit.jupiter.params.provider.MethodSource"));
+            final Optional<J.Annotation> methodSourceAnnotation = FindAnnotation.findFirst(method, new AnnotationMatcher("@org.junit.jupiter.params.provider.MethodSource"));
             if (methodSourceAnnotation.isPresent()) {
                 return method;
             }
 
             // if no testng annotation is present, skip
-            Optional<J.Annotation> testNgAnnotation = FindAnnotation.findFirst(method, new AnnotationMatcher("@org.testng.annotations.Test"));
+            final Optional<J.Annotation> testNgAnnotation = FindAnnotation.findFirst(method, new AnnotationMatcher("@org.testng.annotations.Test"));
             if (!testNgAnnotation.isPresent()) {
                 return method;
             }
 
             // determine Provider name, if not present skip!
-            Optional<String> dataProviderMethodName = AnnotationArguments.extractLiteral(testNgAnnotation.get(), "dataProvider", String.class);
+            final Optional<String> dataProviderMethodName = AnnotationArguments.extractLiteral(testNgAnnotation.get(), "dataProvider", String.class);
             if (!dataProviderMethodName.isPresent()) {
                 return method;
             }
 
             // determin provider class or use current class as default
-            String dataProviderClass = AnnotationArguments.extractAssignments(testNgAnnotation.get(), "dataProviderClass").stream()
+            final String dataProviderClass = AnnotationArguments.extractAssignments(testNgAnnotation.get(), "dataProviderClass").stream()
                     .findAny()
                     .map(J.FieldAccess.class::cast)
                     .map(J.FieldAccess::getTarget)
@@ -176,18 +174,19 @@ public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, Ex
                     .orElse(requireNonNull(getCursor().firstEnclosingOrThrow(J.ClassDeclaration.class).getType()).getFullyQualifiedName());
 
             // add MethodSource annotation
-            JavaCoordinates addAnnotationCoordinate = method.getCoordinates().addAnnotation((a, b) -> 1);
-            method = JavaTemplate
+            maybeAddImport("org.junit.jupiter.params.provider.MethodSource");
+            maybeRemoveImport(dataProviderClass);
+            return JavaTemplate
                     .builder("@MethodSource(\"#{}##{}\")")
                     .javaParser(JavaParser.fromJavaVersion().classpath("junit-jupiter-params"))
                     .imports("org.junit.jupiter.params.provider.MethodSource")
                     .build()
-                    .apply(getCursor(), addAnnotationCoordinate, dataProviderClass, dataProviderMethodName.get());
-
-            maybeAddImport("org.junit.jupiter.params.provider.MethodSource");
-            maybeRemoveImport(dataProviderClass);
-
-            return method;
+                    .apply(
+                        getCursor(), 
+                        method.getCoordinates().addAnnotation((a, b) -> 1), 
+                        dataProviderClass, 
+                        dataProviderMethodName.get()
+                    );
         }
     }
 }
diff --git a/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/helper/AnnotationArguments.java b/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/helper/AnnotationArguments.java
index 633e0d8e..3805c4cd 100644
--- a/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/helper/AnnotationArguments.java
+++ b/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/helper/AnnotationArguments.java
@@ -22,8 +22,8 @@ public enum AnnotationArguments {;
      * @param argumentName to extract
      */
     public static List<Expression> extractAssignments(J.Annotation annotation, String argumentName) {
-        List<Expression> arguments = annotation.getArguments();
-
+        
+        final List<Expression> arguments = annotation.getArguments();
         if (arguments == null) {
             return Collections.emptyList();
         }
@@ -46,8 +46,8 @@ public static List<Expression> extractAssignments(J.Annotation annotation, Strin
      * @return the value or Optional#empty
      */
     public static <T> Optional<T> extractLiteral(J.Annotation annotation, String argumentName, Class<T> valueClass) {
-        List<Expression> arguments = annotation.getArguments();
 
+        final List<Expression> arguments = annotation.getArguments();
         if (arguments == null) {
             return Optional.empty();
         }
diff --git a/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/helper/FindAnnotatedMethods.java b/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/helper/FindAnnotatedMethods.java
index 296dbdce..23f1fcd4 100644
--- a/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/helper/FindAnnotatedMethods.java
+++ b/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/helper/FindAnnotatedMethods.java
@@ -35,11 +35,10 @@ public static Set<J.MethodDeclaration> find(J subtree, AnnotationMatcher annotat
 
     @Override
     public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, ExecutionContext ctx) {
-        J.MethodDeclaration m = super.visitMethodDeclaration(method, ctx);
-
-        boolean isAnnotatedWithTargetAnnotation = m.getLeadingAnnotations().stream().anyMatch(annotationMatcher::matches);
-        if (isAnnotatedWithTargetAnnotation) {
-            m = SearchResult.found(m);
+        
+        final J.MethodDeclaration m = super.visitMethodDeclaration(method, ctx);
+        if (m.getLeadingAnnotations().stream().anyMatch(annotationMatcher::matches)) {
+            return SearchResult.found(m);
         }
 
         return m;
diff --git a/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/helper/FindAnnotation.java b/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/helper/FindAnnotation.java
index b9d1aaad..c234734b 100644
--- a/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/helper/FindAnnotation.java
+++ b/src/main/java/io/github/mboegers/openrewrite/testngtojupiter/helper/FindAnnotation.java
@@ -36,10 +36,10 @@ public static Set<J.Annotation> find(J tree, AnnotationMatcher annotationMatcher
 
     @Override
     public J.Annotation visitAnnotation(J.Annotation annotation, ExecutionContext executionContext) {
+        
         annotation = super.visitAnnotation(annotation, executionContext);
-
         if (annotationMatcher.matches(annotation)) {
-            annotation = SearchResult.found(annotation);
+            return SearchResult.found(annotation);
         }
 
         return annotation;