Add dependencies in JVMTestSuite dependency block

rewrite-gradle/src/main/java/org/openrewrite/gradle/AddDependency.java

@@ -27,14 +27,17 @@
 import org.openrewrite.java.marker.JavaProject;
 import org.openrewrite.java.marker.JavaSourceSet;
 import org.openrewrite.java.search.UsesType;
+import org.openrewrite.java.tree.Expression;
 import org.openrewrite.java.tree.J;
 import org.openrewrite.java.tree.JavaSourceFile;
+import org.openrewrite.java.tree.Statement;
 import org.openrewrite.maven.table.MavenMetadataFailures;
 import org.openrewrite.semver.Semver;
 
 import java.util.*;
+import java.util.function.Predicate;
 
-import static java.util.Collections.*;
+import static java.util.Collections.singletonList;
 import static java.util.Objects.requireNonNull;
 
 @Value
@@ -56,25 +59,25 @@ public class AddDependency extends ScanningRecipe<AddDependency.Scanned> {
 
     @Option(displayName = "Version",
             description = "An exact version number or node-style semver selector used to select the version number. " +
-                          "You can also use `latest.release` for the latest available version and `latest.patch` if " +
-                          "the current version is a valid semantic version. For more details, you can look at the documentation " +
-                          "page of [version selectors](https://docs.openrewrite.org/reference/dependency-version-selectors).",
+                    "You can also use `latest.release` for the latest available version and `latest.patch` if " +
+                    "the current version is a valid semantic version. For more details, you can look at the documentation " +
+                    "page of [version selectors](https://docs.openrewrite.org/reference/dependency-version-selectors).",
             example = "29.X",
             required = false)
     @Nullable
     String version;
 
     @Option(displayName = "Version pattern",
             description = "Allows version selection to be extended beyond the original Node Semver semantics. So for example, " +
-                          "Setting 'version' to \"25-29\" can be paired with a metadata pattern of \"-jre\" to select Guava 29.0-jre",
+                    "Setting 'version' to \"25-29\" can be paired with a metadata pattern of \"-jre\" to select Guava 29.0-jre",
             example = "-jre",
             required = false)
     @Nullable
     String versionPattern;
 
     @Option(displayName = "Configuration",
             description = "A configuration to use when it is not what can be inferred from usage. Most of the time this will be left empty, but " +
-                          "is used when adding a new as of yet unused dependency.",
+                    "is used when adding a new as of yet unused dependency.",
             example = "implementation",
             required = false)
     @Nullable
@@ -103,7 +106,7 @@ public class AddDependency extends ScanningRecipe<AddDependency.Scanned> {
 
     @Option(displayName = "Family pattern",
             description = "A pattern, applied to groupIds, used to determine which other dependencies should have aligned version numbers. " +
-                          "Accepts '*' as a wildcard character.",
+                    "Accepts '*' as a wildcard character.",
             example = "com.fasterxml.jackson*",
             required = false)
     @Nullable
@@ -143,6 +146,7 @@ public Validated<Object> validate() {
     public static class Scanned {
         boolean usingType;
         Map<JavaProject, Set<String>> configurationsByProject = new HashMap<>();
+        Map<JavaProject, Set<String>> customJvmTestSuitesWithDependencies = new HashMap<>();
     }
 
     @Override
@@ -156,16 +160,75 @@ public TreeVisitor<?, ExecutionContext> getScanner(Scanned acc) {
 
             @Nullable
             UsesType<ExecutionContext> usesType = null;
+
             private boolean usesType(SourceFile sourceFile, ExecutionContext ctx) {
-                if(onlyIfUsing == null) {
+                if (onlyIfUsing == null) {
                     return true;
                 }
-                if(usesType == null) {
+                if (usesType == null) {
                     usesType = new UsesType<>(onlyIfUsing, true);
                 }
                 return usesType.isAcceptable(sourceFile, ctx) && usesType.visit(sourceFile, ctx) != sourceFile;
             }
 
+            private Set<String> jvmTestSuites(SourceFile sourceFile, ExecutionContext ctx) {
+                HashSet<String> customJVMSuites = new HashSet<>();
+                new GroovyIsoVisitor<ExecutionContext>() {
+
+                    private boolean isJVMTestSuitesBlock(J.MethodInvocation methodInvocation) {
+                        if ("suites".equals(methodInvocation.getSimpleName())) {
+                            if (getCursor().getParent() != null) {
+                                J.MethodInvocation parentBlock = getCursor().getParent().firstEnclosing(J.MethodInvocation.class);
+                                return parentBlock != null && "testing".equals(parentBlock.getSimpleName());
+                            }
+                        }
+                        return false;
+                    }
+
+                    boolean definesDependencies(J.MethodInvocation suite) {
+                        for (Expression suiteDefinition : suite.getArguments()) {
+                            if (suiteDefinition instanceof J.Lambda) {
+                                for (Statement statement : ((J.Block) ((J.Lambda) suiteDefinition).getBody()).getStatements()) {
+                                    if (statement instanceof J.Return) {
+                                        Expression expression = ((J.Return) statement).getExpression();
+                                        if (expression instanceof J.MethodInvocation) {
+                                            if ("dependencies".equals(((J.MethodInvocation) expression).getSimpleName())) {
+                                                return true;
+                                            }
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                        return false;
+                    }
+
+                    @Override
+                    public J.MethodInvocation visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx) {
+                        if (isJVMTestSuitesBlock(method)) {
+                            for (Expression suiteDefinition : method.getArguments()) {
+                                if (suiteDefinition instanceof J.Lambda) {
+                                    for (Statement statement : ((J.Block) ((J.Lambda) suiteDefinition).getBody()).getStatements()) {
+                                        if (statement instanceof J.Return) {
+                                            Expression expression = ((J.Return) statement).getExpression();
+                                            if (expression instanceof J.MethodInvocation) {
+                                                J.MethodInvocation suite = (J.MethodInvocation) expression;
+                                                if (definesDependencies(suite)) {
+                                                    customJVMSuites.add(suite.getSimpleName());
+                                                }
+                                            }
+                                        }
+                                    }
+                                }
+                            }
+                            return method;
+                        }
+                        return super.visitMethodInvocation(method, ctx);
+                    }
+                }.visit(sourceFile, ctx);
+                return customJVMSuites;
+            }
+
             @Override
             public @Nullable Tree visit(@Nullable Tree tree, ExecutionContext ctx) {
                 if (!(tree instanceof SourceFile)) {
@@ -176,6 +239,9 @@ private boolean usesType(SourceFile sourceFile, ExecutionContext ctx) {
                     if (usesType(sourceFile, ctx)) {
                         acc.usingType = true;
                     }
+
+                    acc.customJvmTestSuitesWithDependencies.computeIfAbsent(javaProject, k -> jvmTestSuites(sourceFile, ctx));
+
                     Set<String> configurations = acc.configurationsByProject.computeIfAbsent(javaProject, ignored -> new HashSet<>());
                     sourceFile.getMarkers().findFirst(JavaSourceSet.class).ifPresent(sourceSet ->
                             configurations.add("main".equals(sourceSet.getName()) ? "implementation" : sourceSet.getName() + "Implementation"));
@@ -248,8 +314,13 @@ public TreeVisitor<?, ExecutionContext> getVisitor(Scanned acc) {
 
                         G.CompilationUnit g = (G.CompilationUnit) s;
                         for (String resolvedConfiguration : resolvedConfigurations) {
-                            g = (G.CompilationUnit) new AddDependencyVisitor(groupId, artifactId, version, versionPattern, resolvedConfiguration,
-                                    classifier, extension, metadataFailures, this::isTopLevel).visitNonNull(g, ctx);
+                            if (!targetsCustomJVMTestSuite(resolvedConfiguration, acc.customJvmTestSuitesWithDependencies.get(jp))) {
+                                g = (G.CompilationUnit) new AddDependencyVisitor(groupId, artifactId, version, versionPattern, resolvedConfiguration,
+                                        classifier, extension, metadataFailures, this::isTopLevel).visitNonNull(g, ctx);
+                            } else {
+                                g = (G.CompilationUnit) new AddDependencyVisitor(groupId, artifactId, version, versionPattern, "implementation",
+                                        classifier, extension, metadataFailures, isMatchingJVMTestSuite(resolvedConfiguration)).visitNonNull(g, ctx);
+                            }
                         }
 
                         return g;
@@ -258,7 +329,36 @@ public TreeVisitor<?, ExecutionContext> getVisitor(Scanned acc) {
                     private boolean isTopLevel(Cursor cursor) {
                         return cursor.getParentOrThrow().firstEnclosing(J.MethodInvocation.class) == null;
                     }
+
+                    private Predicate<Cursor> isMatchingJVMTestSuite(String resolvedConfiguration) {
+                        return cursor -> {
+                            String sourceSet = resolvedConfiguration.substring(0, resolvedConfiguration.length() - 14);
+                            J.MethodInvocation methodInvocation = cursor.getParentOrThrow().firstEnclosing(J.MethodInvocation.class);
+                            return methodInvocation != null && sourceSet.equals(methodInvocation.getSimpleName());
+                        };
+                    }
+
+                    private final Set<String> gradleStandardConfigurations = new HashSet<>(Arrays.asList(
+                            "api",
+                            "implementation",
+                            "compileOnly",
+                            "compileOnlyApi",
+                            "runtimeOnly",
+                            "testImplementation",
+                            "testCompileOnly",
+                            "testRuntimeOnly"));
+
+                    boolean targetsCustomJVMTestSuite(String configuration, Set<String> customJvmTestSuites) {
+                        if (gradleStandardConfigurations.contains(configuration) || "default".equals(configuration)) {
+                            return false;
+                        }
+
+                        String sourceSet = configuration.substring(0, configuration.length() - 14);
+                        return customJvmTestSuites.contains(sourceSet);
+                    }
                 })
         );
     }
+
+
 }