Add Gradle plugin support for processing test contexts ahead-of-time

Closes gh-32192

spring-boot-project/spring-boot-tools/spring-boot-gradle-plugin/src/main/java/org/springframework/boot/gradle/plugin/NativeImagePluginAction.java

@@ -52,6 +52,10 @@ class NativeImagePluginAction implements PluginApplicationAction {
 			SourceSet aotSourceSet = sourceSets.getByName(SpringBootAotPlugin.AOT_SOURCE_SET_NAME);
 			project.getTasks().named(NativeImagePlugin.NATIVE_COMPILE_TASK_NAME, BuildNativeImageTask.class,
 					(nativeCompile) -> nativeCompile.getOptions().get().classpath(aotSourceSet.getOutput()));
+			SourceSet aotTestSourceSet = sourceSets.getByName(SpringBootAotPlugin.AOT_TEST_SOURCE_SET_NAME);
+			project.getTasks().named("nativeTestCompile", BuildNativeImageTask.class,
+					(nativeTestCompile) -> nativeTestCompile.getOptions().get()
+							.classpath(aotTestSourceSet.getOutput()));
 		});
 		GraalVMExtension graalVmExtension = project.getExtensions().getByType(GraalVMExtension.class);
 		graalVmExtension.getToolchainDetection().set(false);

spring-boot-project/spring-boot-tools/spring-boot-gradle-plugin/src/main/java/org/springframework/boot/gradle/plugin/SpringBootAotPlugin.java

@@ -35,6 +35,7 @@ import org.gradle.api.tasks.SourceSetContainer;
 import org.gradle.api.tasks.TaskProvider;
 
 import org.springframework.boot.gradle.tasks.aot.ProcessAot;
+import org.springframework.boot.gradle.tasks.aot.ProcessTestAot;
 
 /**
  * Gradle plugin for Spring Boot AOT.
@@ -45,45 +46,57 @@ import org.springframework.boot.gradle.tasks.aot.ProcessAot;
 public class SpringBootAotPlugin implements Plugin<Project> {
 
 	/**
-	 * Name of the {@code aot} {@link SourceSet source set}.
+	 * Name of the main {@code aot} {@link SourceSet source set}.
 	 */
 	public static final String AOT_SOURCE_SET_NAME = "aot";
 
+	/**
+	 * Name of the {@code aotTest} {@link SourceSet source set}.
+	 */
+	public static final String AOT_TEST_SOURCE_SET_NAME = "aotTest";
+
 	/**
 	 * Name of the default {@link ProcessAot} task.
 	 */
 	public static final String PROCESS_AOT_TASK_NAME = "processAot";
 
+	/**
+	 * Name of the default {@link ProcessAot} task.
+	 */
+	public static final String PROCESS_TEST_AOT_TASK_NAME = "processTestAot";
+
 	@Override
 	public void apply(Project project) {
 		PluginContainer plugins = project.getPlugins();
 		plugins.withType(JavaPlugin.class).all((javaPlugin) -> {
 			plugins.withType(SpringBootPlugin.class).all((bootPlugin) -> {
-				SourceSet aotSourceSet = configureAotSourceSet(project);
+				SourceSet aotSourceSet = configureSourceSet(project, "aot", SourceSet.MAIN_SOURCE_SET_NAME);
 				registerProcessAotTask(project, aotSourceSet);
+				SourceSet aotTestSourceSet = configureSourceSet(project, "aotTest", SourceSet.TEST_SOURCE_SET_NAME);
+				registerProcessTestAotTask(project, aotTestSourceSet);
 			});
 		});
 	}
 
-	private SourceSet configureAotSourceSet(Project project) {
+	private SourceSet configureSourceSet(Project project, String newSourceSetName, String existingSourceSetName) {
 		JavaPluginExtension javaPluginExtension = project.getExtensions().getByType(JavaPluginExtension.class);
 		SourceSetContainer sourceSets = javaPluginExtension.getSourceSets();
-		SourceSet main = sourceSets.getByName(SourceSet.MAIN_SOURCE_SET_NAME);
-		SourceSet aotSourceSet = sourceSets.create(AOT_SOURCE_SET_NAME, (aot) -> {
-			aot.getJava().setSrcDirs(List.of("build/generated/aotSources"));
-			aot.getResources().setSrcDirs(List.of("build/generated/aotResources"));
-			aot.setCompileClasspath(aot.getCompileClasspath().plus(main.getOutput()));
-			main.setRuntimeClasspath(main.getRuntimeClasspath().plus(aot.getOutput()));
+		SourceSet existingSourceSet = sourceSets.getByName(existingSourceSetName);
+		return sourceSets.create(newSourceSetName, (sourceSet) -> {
+			sourceSet.getJava().setSrcDirs(List.of("build/generated/" + newSourceSetName + "Sources"));
+			sourceSet.getResources().setSrcDirs(List.of("build/generated/" + newSourceSetName + "Resources"));
+			sourceSet.setCompileClasspath(sourceSet.getCompileClasspath().plus(existingSourceSet.getOutput()));
+			existingSourceSet.setRuntimeClasspath(existingSourceSet.getRuntimeClasspath().plus(sourceSet.getOutput()));
 			ConfigurationContainer configurations = project.getConfigurations();
-			Configuration aotImplementation = configurations.getByName(aot.getImplementationConfigurationName());
-			aotImplementation.extendsFrom(configurations.getByName(main.getImplementationConfigurationName()));
-			aotImplementation.extendsFrom(configurations.getByName(main.getRuntimeOnlyConfigurationName()));
-			configurations.getByName(aot.getCompileClasspathConfigurationName()).attributes((attributes) -> {
+			Configuration implementation = configurations.getByName(sourceSet.getImplementationConfigurationName());
+			implementation
+					.extendsFrom(configurations.getByName(existingSourceSet.getImplementationConfigurationName()));
+			implementation.extendsFrom(configurations.getByName(existingSourceSet.getRuntimeOnlyConfigurationName()));
+			configurations.getByName(sourceSet.getCompileClasspathConfigurationName()).attributes((attributes) -> {
 				configureClassesAndResourcesLibraryElementsAttribute(project, attributes);
 				configureJavaRuntimeUsageAttribute(project, attributes);
 			});
 		});
-		return aotSourceSet;
 	}
 
 	private void configureClassesAndResourcesLibraryElementsAttribute(Project project, AttributeContainer attributes) {
@@ -119,4 +132,27 @@ public class SpringBootAotPlugin implements Plugin<Project> {
 				.configure((processResources) -> processResources.dependsOn(processAot));
 	}
 
+	private void registerProcessTestAotTask(Project project, SourceSet aotTestSourceSet) {
+		JavaPluginExtension javaPluginExtension = project.getExtensions().getByType(JavaPluginExtension.class);
+		SourceSetContainer sourceSets = javaPluginExtension.getSourceSets();
+		SourceSet testSourceSet = sourceSets.getByName(SourceSet.TEST_SOURCE_SET_NAME);
+		TaskProvider<ProcessTestAot> processTestAot = project.getTasks().register(PROCESS_TEST_AOT_TASK_NAME,
+				ProcessTestAot.class, (task) -> {
+					Provider<Directory> generatedClasses = project.getLayout().getBuildDirectory()
+							.dir("generated/aotClasses");
+					aotTestSourceSet.getOutput().dir(generatedClasses);
+					task.setTestClassesDirs(testSourceSet.getOutput().getClassesDirs());
+					task.setClasspath(aotTestSourceSet.getCompileClasspath());
+					task.getSourcesDir().set(aotTestSourceSet.getJava().getSrcDirs().iterator().next());
+					task.getResourcesDir().set(aotTestSourceSet.getResources().getSrcDirs().iterator().next());
+					task.getClassesDir().set(generatedClasses);
+					task.getGroupId().set(project.provider(() -> String.valueOf(project.getGroup())));
+					task.getArtifactId().set(project.provider(() -> project.getName()));
+				});
+		project.getTasks().named(aotTestSourceSet.getCompileJavaTaskName())
+				.configure((compileJava) -> compileJava.dependsOn(processTestAot));
+		project.getTasks().named(aotTestSourceSet.getProcessResourcesTaskName())
+				.configure((processResources) -> processResources.dependsOn(processTestAot));
+	}
+
 }

spring-boot-project/spring-boot-tools/spring-boot-gradle-plugin/src/main/java/org/springframework/boot/gradle/tasks/aot/AbstractAot.java

@@ -0,0 +1,93 @@
+/*
+ * Copyright 2012-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.gradle.tasks.aot;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.gradle.api.file.DirectoryProperty;
+import org.gradle.api.provider.Property;
+import org.gradle.api.tasks.Input;
+import org.gradle.api.tasks.JavaExec;
+import org.gradle.api.tasks.OutputDirectory;
+import org.gradle.work.DisableCachingByDefault;
+
+/**
+ * Specialization of {@link JavaExec} to be used as a base class for tasks that perform
+ * ahead-of-time processing.
+ *
+ * @author Andy Wilkinson
+ * @since 3.0.0
+ */
+@DisableCachingByDefault(because = "Cacheability can only be determined by a concrete implementation")
+public abstract class AbstractAot extends JavaExec {
+
+	private final DirectoryProperty sourcesDir;
+
+	private final DirectoryProperty resourcesDir;
+
+	private final DirectoryProperty classesDir;
+
+	private final Property<String> groupId;
+
+	private final Property<String> artifactId;
+
+	protected AbstractAot() {
+		this.sourcesDir = getProject().getObjects().directoryProperty();
+		this.resourcesDir = getProject().getObjects().directoryProperty();
+		this.classesDir = getProject().getObjects().directoryProperty();
+		this.groupId = getProject().getObjects().property(String.class);
+		this.artifactId = getProject().getObjects().property(String.class);
+	}
+
+	@Input
+	public final Property<String> getGroupId() {
+		return this.groupId;
+	}
+
+	@Input
+	public final Property<String> getArtifactId() {
+		return this.artifactId;
+	}
+
+	@OutputDirectory
+	public final DirectoryProperty getSourcesDir() {
+		return this.sourcesDir;
+	}
+
+	@OutputDirectory
+	public final DirectoryProperty getResourcesDir() {
+		return this.resourcesDir;
+	}
+
+	@OutputDirectory
+	public final DirectoryProperty getClassesDir() {
+		return this.classesDir;
+	}
+
+	List<String> processorArgs() {
+		List<String> args = new ArrayList<>();
+		args.add(getSourcesDir().getAsFile().get().getAbsolutePath());
+		args.add(getResourcesDir().getAsFile().get().getAbsolutePath());
+		args.add(getClassesDir().getAsFile().get().getAbsolutePath());
+		args.add(getGroupId().get());
+		args.add(getArtifactId().get());
+		args.addAll(super.getArgs());
+		return args;
+	}
+
+}

spring-boot-project/spring-boot-tools/spring-boot-gradle-plugin/src/main/java/org/springframework/boot/gradle/tasks/aot/ProcessAot.java

@@ -19,42 +19,25 @@ package org.springframework.boot.gradle.tasks.aot;
 import java.util.ArrayList;
 import java.util.List;
 
-import org.gradle.api.file.DirectoryProperty;
 import org.gradle.api.provider.Property;
 import org.gradle.api.tasks.CacheableTask;
 import org.gradle.api.tasks.Input;
 import org.gradle.api.tasks.JavaExec;
-import org.gradle.api.tasks.OutputDirectory;
 import org.gradle.api.tasks.TaskAction;
 
 /**
- * Custom {@link JavaExec} task for processing code ahead-of-time.
+ * Custom {@link JavaExec} task for processing main code ahead-of-time.
  *
  * @author Andy Wilkinson
  * @since 3.0.0
  */
 @CacheableTask
-public class ProcessAot extends JavaExec {
+public class ProcessAot extends AbstractAot {
 
 	private final Property<String> applicationClass;
 
-	private final DirectoryProperty sourcesDir;
-
-	private final DirectoryProperty resourcesDir;
-
-	private final DirectoryProperty classesDir;
-
-	private final Property<String> groupId;
-
-	private final Property<String> artifactId;
-
 	public ProcessAot() {
 		this.applicationClass = getProject().getObjects().property(String.class);
-		this.sourcesDir = getProject().getObjects().directoryProperty();
-		this.resourcesDir = getProject().getObjects().directoryProperty();
-		this.classesDir = getProject().getObjects().directoryProperty();
-		this.groupId = getProject().getObjects().property(String.class);
-		this.artifactId = getProject().getObjects().property(String.class);
 		getMainClass().set("org.springframework.boot.AotProcessor");
 	}
 
@@ -63,42 +46,12 @@ public class ProcessAot extends JavaExec {
 		return this.applicationClass;
 	}
 
-	@Input
-	public Property<String> getGroupId() {
-		return this.groupId;
-	}
-
-	@Input
-	public Property<String> getArtifactId() {
-		return this.artifactId;
-	}
-
-	@OutputDirectory
-	public DirectoryProperty getSourcesDir() {
-		return this.sourcesDir;
-	}
-
-	@OutputDirectory
-	public DirectoryProperty getResourcesDir() {
-		return this.resourcesDir;
-	}
-
-	@OutputDirectory
-	public DirectoryProperty getClassesDir() {
-		return this.classesDir;
-	}
-
 	@Override
 	@TaskAction
 	public void exec() {
 		List<String> args = new ArrayList<>();
 		args.add(this.applicationClass.get());
-		args.add(this.sourcesDir.getAsFile().get().getAbsolutePath());
-		args.add(this.resourcesDir.getAsFile().get().getAbsolutePath());
-		args.add(this.classesDir.getAsFile().get().getAbsolutePath());
-		args.add(this.groupId.get());
-		args.add(this.artifactId.get());
-		args.addAll(super.getArgs());
+		args.addAll(processorArgs());
 		this.setArgs(args);
 		super.exec();
 	}

spring-boot-project/spring-boot-tools/spring-boot-gradle-plugin/src/main/java/org/springframework/boot/gradle/tasks/aot/ProcessTestAot.java

@@ -0,0 +1,90 @@
+/*
+ * Copyright 2012-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.gradle.tasks.aot;
+
+import java.io.File;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import org.gradle.api.artifacts.Configuration;
+import org.gradle.api.artifacts.Dependency;
+import org.gradle.api.artifacts.DependencySet;
+import org.gradle.api.artifacts.dsl.DependencyHandler;
+import org.gradle.api.file.FileCollection;
+import org.gradle.api.tasks.CacheableTask;
+import org.gradle.api.tasks.Classpath;
+import org.gradle.api.tasks.JavaExec;
+import org.gradle.api.tasks.TaskAction;
+
+import org.springframework.boot.gradle.plugin.SpringBootPlugin;
+
+/**
+ * Custom {@link JavaExec} task for processing test code ahead-of-time.
+ *
+ * @author Andy Wilkinson
+ * @since 3.0.0
+ */
+@CacheableTask
+public class ProcessTestAot extends AbstractAot {
+
+	private final FileCollection junitPlatformLauncher;
+
+	private FileCollection testClassesDirs;
+
+	public ProcessTestAot() {
+		getMainClass().set("org.springframework.test.context.aot.TestAotProcessor");
+		this.junitPlatformLauncher = createJUnitPlatformLauncher();
+	}
+
+	private Configuration createJUnitPlatformLauncher() {
+		Configuration configuration = getProject().getConfigurations().create(getName() + "JUnitPlatformLauncher");
+		DependencyHandler dependencyHandler = getProject().getDependencies();
+		Dependency springBootDependencies = dependencyHandler
+				.create(dependencyHandler.platform(SpringBootPlugin.BOM_COORDINATES));
+		DependencySet dependencies = configuration.getDependencies();
+		dependencies.add(springBootDependencies);
+		dependencies.add(dependencyHandler.create("org.junit.platform:junit-platform-launcher"));
+		return configuration;
+	}
+
+	@Classpath
+	public FileCollection getTestClassesDirs() {
+		return this.testClassesDirs;
+	}
+
+	public void setTestClassesDirs(FileCollection testClassesDirs) {
+		this.testClassesDirs = testClassesDirs;
+	}
+
+	@Classpath
+	FileCollection getJUnitPlatformLauncher() {
+		return this.junitPlatformLauncher;
+	}
+
+	@Override
+	@TaskAction
+	public void exec() {
+		List<String> args = new ArrayList<>();
+		args.add(this.testClassesDirs.getFiles().stream().map(File::getAbsolutePath)
+				.collect(Collectors.joining(File.pathSeparator)));
+		args.addAll(processorArgs());
+		this.setArgs(args);
+		super.exec();
+	}
+
+}

spring-boot-project/spring-boot-tools/spring-boot-gradle-plugin/src/test/java/org/springframework/boot/gradle/plugin/SpringBootAotPluginIntegrationTests.java

@@ -44,12 +44,24 @@ class SpringBootAotPluginIntegrationTests {
 				.contains("processAot exists = false");
 	}
 
+	@TestTemplate
+	void noProcessTestAotTaskWithoutAotPluginApplied() {
+		assertThat(this.gradleBuild.build("taskExists", "-PtaskName=processTestAot").getOutput())
+				.contains("processTestAot exists = false");
+	}
+
 	@TestTemplate
 	void applyingAotPluginCreatesProcessAotTask() {
 		assertThat(this.gradleBuild.build("taskExists", "-PtaskName=processAot").getOutput())
 				.contains("processAot exists = true");
 	}
 
+	@TestTemplate
+	void applyingAotPluginCreatesProcessTestAotTask() {
+		assertThat(this.gradleBuild.build("taskExists", "-PtaskName=processTestAot").getOutput())
+				.contains("processTestAot exists = true");
+	}
+
 	@TestTemplate
 	void processAotHasLibraryResourcesOnItsClasspath() throws IOException {
 		File settings = new File(this.gradleBuild.getProjectDir(), "settings.gradle");
@@ -60,10 +72,26 @@ class SpringBootAotPluginIntegrationTests {
 		assertThat(this.gradleBuild.build("processAotClasspath").getOutput()).contains("library.jar");
 	}
 
+	@TestTemplate
+	void processTestAotHasLibraryResourcesOnItsClasspath() throws IOException {
+		File settings = new File(this.gradleBuild.getProjectDir(), "settings.gradle");
+		Files.write(settings.toPath(), List.of("include 'library'"));
+		File library = new File(this.gradleBuild.getProjectDir(), "library");
+		library.mkdirs();
+		Files.write(library.toPath().resolve("build.gradle"), List.of("plugins {", "    id 'java-library'", "}"));
+		assertThat(this.gradleBuild.build("processTestAotClasspath").getOutput()).contains("library.jar");
+	}
+
 	@TestTemplate
 	void processAotHasTransitiveRuntimeDependenciesOnItsClasspath() {
 		String output = this.gradleBuild.build("processAotClasspath").getOutput();
 		assertThat(output).contains("org.jboss.logging" + File.separatorChar + "jboss-logging");
 	}
 
+	@TestTemplate
+	void processTestAotHasTransitiveRuntimeDependenciesOnItsClasspath() {
+		String output = this.gradleBuild.build("processTestAotClasspath").getOutput();
+		assertThat(output).contains("org.jboss.logging" + File.separatorChar + "jboss-logging");
+	}
+
 }

spring-boot-project/spring-boot-tools/spring-boot-gradle-plugin/src/test/resources/org/springframework/boot/gradle/plugin/SpringBootAotPluginIntegrationTests-applyingAotPluginCreatesProcessTestAotTask.gradle

@@ -0,0 +1,11 @@
+plugins {
+	id 'org.springframework.boot'
+	id 'org.springframework.boot.aot'
+	id 'java'
+}
+
+task('taskExists') {
+	doFirst {
+		println "${taskName} exists = ${tasks.findByName(taskName) != null}"
+	}
+}

spring-boot-project/spring-boot-tools/spring-boot-gradle-plugin/src/test/resources/org/springframework/boot/gradle/plugin/SpringBootAotPluginIntegrationTests-noProcessTestAotTaskWithoutAotPluginApplied.gradle

@@ -0,0 +1,10 @@
+plugins {
+	id 'org.springframework.boot'
+	id 'java'
+}
+
+task('taskExists') {
+	doFirst {
+		println "${taskName} exists = ${tasks.findByName(taskName) != null}"
+	}
+}

spring-boot-project/spring-boot-tools/spring-boot-gradle-plugin/src/test/resources/org/springframework/boot/gradle/plugin/SpringBootAotPluginIntegrationTests-processTestAotHasLibraryResourcesOnItsClasspath.gradle

@@ -0,0 +1,15 @@
+plugins {
+	id 'org.springframework.boot'
+	id 'org.springframework.boot.aot'
+	id 'java'
+}
+
+dependencies {
+	implementation project(":library")
+}
+
+task('processTestAotClasspath') {
+	doFirst {
+		tasks.findByName('processTestAot').classpath.files.each { println it }
+	}
+}

spring-boot-project/spring-boot-tools/spring-boot-gradle-plugin/src/test/resources/org/springframework/boot/gradle/plugin/SpringBootAotPluginIntegrationTests-processTestAotHasTransitiveRuntimeDependenciesOnItsClasspath.gradle

@@ -0,0 +1,19 @@
+plugins {
+	id 'org.springframework.boot'
+	id 'org.springframework.boot.aot'
+	id 'java'
+}
+
+repositories {
+	mavenCentral()
+}
+
+dependencies {
+	implementation "org.hibernate.orm:hibernate-core:6.1.1.Final"
+}
+
+task('processTestAotClasspath') {
+	doFirst {
+		tasks.findByName('processTestAot').classpath.files.each { println it }
+	}
+}