Break dependency between TestCompiler and AOT

This commit improves `TestCompiler` with a `with` function that allows
to customize a test compiler instance. Rather than `TestCompiler`
knowing about `TestGenerationContext`, the latter implements the
function so that it can be passed as is.

See gh-29175

spring-aop/src/test/java/org/springframework/aop/scope/ScopedProxyBeanRegistrationAotProcessorTests.java

@@ -151,7 +151,7 @@ private void compile(BiConsumer<DefaultListableBeanFactory, Compiled> result) {
 					.build());
 		});
 		this.generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(this.generationContext.getGeneratedFiles()).compile(compiled -> {
+		TestCompiler.forSystem().with(this.generationContext).compile(compiled -> {
 			DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
 			freshBeanFactory.setBeanClassLoader(compiled.getClassLoader());
 			compiled.getInstance(Consumer.class).accept(freshBeanFactory);

spring-beans/src/test/java/org/springframework/beans/factory/annotation/AutowiredAnnotationBeanRegistrationAotContributionTests.java

@@ -176,7 +176,7 @@ private void compile(RegisteredBean registeredBean,
 
 		});
 		this.generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(this.generationContext.getGeneratedFiles()).compile(compiled ->
+		TestCompiler.forSystem().with(this.generationContext).compile(compiled ->
 				result.accept(compiled.getInstance(BiFunction.class), compiled));
 	}
 

spring-beans/src/test/java/org/springframework/beans/factory/aot/BeanDefinitionMethodGeneratorTests.java

@@ -425,7 +425,7 @@ private void compile(MethodReference method,
 					.addCode("return $L;", methodInvocation).build());
 		});
 		this.generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(this.generationContext.getGeneratedFiles()).compile(compiled ->
+		TestCompiler.forSystem().with(this.generationContext).compile(compiled ->
 				result.accept((RootBeanDefinition) compiled.getInstance(Supplier.class).get(), compiled));
 	}
 

spring-beans/src/test/java/org/springframework/beans/factory/aot/BeanDefinitionPropertiesCodeGeneratorTests.java

@@ -429,7 +429,7 @@ private void compile(
 					.addStatement("return beanDefinition").build());
 		});
 		this.generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(this.generationContext.getGeneratedFiles()).compile(compiled -> {
+		TestCompiler.forSystem().with(this.generationContext).compile(compiled -> {
 			RootBeanDefinition suppliedBeanDefinition = (RootBeanDefinition) compiled
 					.getInstance(Supplier.class).get();
 			result.accept(suppliedBeanDefinition, compiled);

spring-beans/src/test/java/org/springframework/beans/factory/aot/BeanDefinitionPropertyValueCodeGeneratorTests.java

@@ -76,7 +76,7 @@ private void compile(Object value, BiConsumer<Object, Compiled> result) {
 					.returns(Object.class).addStatement("return $L", generatedCode).build());
 		});
 		generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(generationContext.getGeneratedFiles()).compile(compiled ->
+		TestCompiler.forSystem().with(generationContext).compile(compiled ->
 				result.accept(compiled.getInstance(Supplier.class).get(), compiled));
 	}
 

spring-beans/src/test/java/org/springframework/beans/factory/aot/BeanRegistrationsAotContributionTests.java

@@ -167,7 +167,7 @@ private void compile(
 					.build());
 		});
 		this.generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(this.generationContext.getGeneratedFiles()).compile(compiled ->
+		TestCompiler.forSystem().with(this.generationContext).compile(compiled ->
 				result.accept(compiled.getInstance(Consumer.class), compiled));
 	}
 

spring-beans/src/test/java/org/springframework/beans/factory/aot/InstanceSupplierCodeGeneratorTests.java

@@ -317,7 +317,7 @@ private void compile(DefaultListableBeanFactory beanFactory,
 					.addStatement("return $L", generatedCode).build());
 		});
 		this.generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(this.generationContext.getGeneratedFiles()).compile(compiled ->
+		TestCompiler.forSystem().with(this.generationContext).compile(compiled ->
 				result.accept((InstanceSupplier<?>) compiled.getInstance(Supplier.class).get(), compiled));
 	}
 

spring-context/src/test/java/org/springframework/context/annotation/ConfigurationClassPostProcessorAotContributionTests.java

@@ -165,7 +165,7 @@ private void compile(BiConsumer<Consumer<DefaultListableBeanFactory>, Compiled>
 						.build());
 			});
 			generationContext.writeGeneratedContent();
-			TestCompiler.forSystem().withFiles(generationContext.getGeneratedFiles()).compile(compiled ->
+			TestCompiler.forSystem().with(generationContext).compile(compiled ->
 					result.accept(compiled.getInstance(Consumer.class), compiled));
 		}
 
@@ -297,7 +297,7 @@ private void compile(BiConsumer<Consumer<GenericApplicationContext>, Compiled> r
 						.build());
 			});
 			generationContext.writeGeneratedContent();
-			TestCompiler.forSystem().withFiles(generationContext.getGeneratedFiles()).compile(compiled ->
+			TestCompiler.forSystem().with(generationContext).compile(compiled ->
 					result.accept(compiled.getInstance(Consumer.class), compiled));
 		}
 

spring-context/src/test/java/org/springframework/context/aot/ApplicationContextAotGeneratorTests.java

@@ -315,7 +315,7 @@ private void testCompiledResult(GenericApplicationContext applicationContext,
 	@SuppressWarnings({ "rawtypes", "unchecked" })
 	private void testCompiledResult(TestGenerationContext generationContext,
 			BiConsumer<ApplicationContextInitializer<GenericApplicationContext>, Compiled> result) {
-		TestCompiler.forSystem().withFiles(generationContext.getGeneratedFiles()).compile(compiled ->
+		TestCompiler.forSystem().with(generationContext).compile(compiled ->
 				result.accept(compiled.getInstance(ApplicationContextInitializer.class), compiled));
 	}
 

spring-context/src/test/java/org/springframework/context/generator/ApplicationContextAotGeneratorRuntimeHintsTests.java

@@ -86,7 +86,7 @@ private void compile(GenericApplicationContext applicationContext, BiConsumer<Ru
 		TestGenerationContext generationContext = new TestGenerationContext();
 		generator.processAheadOfTime(applicationContext, generationContext);
 		generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(generationContext.getGeneratedFiles()).compile(compiled -> {
+		TestCompiler.forSystem().with(generationContext).compile(compiled -> {
 			ApplicationContextInitializer instance = compiled.getInstance(ApplicationContextInitializer.class);
 			GenericApplicationContext freshContext = new GenericApplicationContext();
 			RuntimeHintsInvocations recordedInvocations = RuntimeHintsRecorder.record(() -> {

spring-core-test/src/main/java/org/springframework/aot/test/generate/GeneratedFilesTestCompilerUtils.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2002-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.aot.test.generate;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.springframework.aot.generate.GeneratedFiles.Kind;
+import org.springframework.aot.generate.InMemoryGeneratedFiles;
+import org.springframework.core.test.tools.ClassFile;
+import org.springframework.core.test.tools.ResourceFile;
+import org.springframework.core.test.tools.SourceFile;
+import org.springframework.core.test.tools.TestCompiler;
+
+/**
+ * {@link TestCompiler} utilities for generated files.
+ *
+ * @author Stephane Nicoll
+ * @since 6.0
+ */
+public abstract class GeneratedFilesTestCompilerUtils {
+
+	/**
+	 * Apply the specified {@link InMemoryGeneratedFiles} to the specified {@link TestCompiler}.
+	 * @param testCompiler the compiler to configure
+	 * @param generatedFiles the generated files to apply
+	 * @return a new {@link TestCompiler} instance configured with the generated files
+	 */
+	public static TestCompiler configure(TestCompiler testCompiler, InMemoryGeneratedFiles generatedFiles) {
+		List<SourceFile> sourceFiles = new ArrayList<>();
+		generatedFiles.getGeneratedFiles(Kind.SOURCE).forEach(
+				(path, inputStreamSource) -> sourceFiles.add(SourceFile.of(inputStreamSource)));
+		List<ResourceFile> resourceFiles = new ArrayList<>();
+		generatedFiles.getGeneratedFiles(Kind.RESOURCE).forEach(
+				(path, inputStreamSource) -> resourceFiles.add(ResourceFile.of(path, inputStreamSource)));
+		List<ClassFile> classFiles = new ArrayList<>();
+		generatedFiles.getGeneratedFiles(Kind.CLASS).forEach(
+				(path, inputStreamSource) -> classFiles.add(ClassFile.of(
+						ClassFile.toClassName(path), inputStreamSource)));
+		return testCompiler.withSources(sourceFiles).withResources(resourceFiles).withClasses(classFiles);
+	}
+}

spring-core-test/src/main/java/org/springframework/aot/test/generate/TestGenerationContext.java

@@ -16,20 +16,24 @@
 
 package org.springframework.aot.test.generate;
 
+import java.util.function.Function;
+
 import org.springframework.aot.generate.ClassNameGenerator;
 import org.springframework.aot.generate.DefaultGenerationContext;
 import org.springframework.aot.generate.GenerationContext;
 import org.springframework.aot.generate.InMemoryGeneratedFiles;
+import org.springframework.core.test.tools.TestCompiler;
 
 /**
  * {@link GenerationContext} test implementation that uses
- * {@link InMemoryGeneratedFiles}.
+ * {@link InMemoryGeneratedFiles} and can configure a {@link TestCompiler}
+ * instance.
  *
  * @author Stephane Nicoll
  * @author Sam Brannen
  * @since 6.0
  */
-public class TestGenerationContext extends DefaultGenerationContext {
+public class TestGenerationContext extends DefaultGenerationContext implements Function<TestCompiler, TestCompiler> {
 
 	/**
 	 * Create an instance using the specified {@link ClassNameGenerator}.
@@ -60,4 +64,15 @@ public InMemoryGeneratedFiles getGeneratedFiles() {
 		return (InMemoryGeneratedFiles) super.getGeneratedFiles();
 	}
 
+	/**
+	 * Configure the specified {@link TestCompiler} with the state of this context.
+	 * @param testCompiler the compiler to configure
+	 * @return a new {@link TestCompiler} instance configured with the generated files
+	 * @see TestCompiler#with(Function)
+	 */
+	@Override
+	public TestCompiler apply(TestCompiler testCompiler) {
+		return GeneratedFilesTestCompilerUtils.configure(testCompiler, getGeneratedFiles());
+	}
+
 }

spring-core-test/src/main/java/org/springframework/core/test/tools/TestCompiler.java

@@ -23,6 +23,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 import javax.annotation.processing.Processor;
 import javax.tools.Diagnostic;
@@ -33,8 +34,6 @@
 import javax.tools.StandardJavaFileManager;
 import javax.tools.ToolProvider;
 
-import org.springframework.aot.generate.GeneratedFiles.Kind;
-import org.springframework.aot.generate.InMemoryGeneratedFiles;
 import org.springframework.lang.Nullable;
 
 /**
@@ -93,23 +92,12 @@ public static TestCompiler forCompiler(JavaCompiler javaCompiler) {
 	}
 
 	/**
-	 * Create a new {@code TestCompiler} instance with additional generated
-	 * source, resource, and class files.
-	 * @param generatedFiles the generated files to add
-	 * @return a new {@code TestCompiler} instance
+	 * Apply customization to this compiler.
+	 * @param customizer the customizer to call
+	 * @return a new {@code TestCompiler} instance with the customizations applied
 	 */
-	public TestCompiler withFiles(InMemoryGeneratedFiles generatedFiles) {
-		List<SourceFile> sourceFiles = new ArrayList<>();
-		generatedFiles.getGeneratedFiles(Kind.SOURCE).forEach(
-				(path, inputStreamSource) -> sourceFiles.add(SourceFile.of(inputStreamSource)));
-		List<ResourceFile> resourceFiles = new ArrayList<>();
-		generatedFiles.getGeneratedFiles(Kind.RESOURCE).forEach(
-				(path, inputStreamSource) -> resourceFiles.add(ResourceFile.of(path, inputStreamSource)));
-		List<ClassFile> classFiles = new ArrayList<>();
-		generatedFiles.getGeneratedFiles(Kind.CLASS).forEach(
-				(path, inputStreamSource) -> classFiles.add(ClassFile.of(
-						ClassFile.toClassName(path), inputStreamSource)));
-		return withSources(sourceFiles).withResources(resourceFiles).withClasses(classFiles);
+	public TestCompiler with(Function<TestCompiler, TestCompiler> customizer) {
+		return customizer.apply(this);
 	}
 
 	/**

spring-orm/src/test/java/org/springframework/orm/jpa/persistenceunit/PersistenceManagedTypesBeanRegistrationAotProcessorTests.java

@@ -109,7 +109,7 @@ private void compile(GenericApplicationContext applicationContext,
 		TestGenerationContext generationContext = new TestGenerationContext();
 		generator.processAheadOfTime(applicationContext, generationContext);
 		generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(generationContext.getGeneratedFiles()).compile(compiled ->
+		TestCompiler.forSystem().with(generationContext).compile(compiled ->
 				result.accept(compiled.getInstance(ApplicationContextInitializer.class), compiled));
 	}
 

spring-orm/src/test/java/org/springframework/orm/jpa/support/PersistenceAnnotationBeanPostProcessorAotContributionTests.java

@@ -175,7 +175,7 @@ private void testCompile(RegisteredBean registeredBean,
 		BeanRegistrationCode beanRegistrationCode = mock(BeanRegistrationCode.class);
 		contribution.applyTo(generationContext, beanRegistrationCode);
 		generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(generationContext.getGeneratedFiles())
+		TestCompiler.forSystem().with(generationContext)
 				.compile(compiled -> result.accept(new Invoker(compiled), compiled));
 	}
 

spring-test/src/test/java/org/springframework/test/context/aot/AotIntegrationTests.java

@@ -21,6 +21,7 @@
 import java.nio.file.Paths;
 import java.util.List;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Stream;
 
 import org.junit.jupiter.api.Disabled;
@@ -37,6 +38,7 @@
 import org.springframework.aot.AotDetector;
 import org.springframework.aot.generate.GeneratedFiles.Kind;
 import org.springframework.aot.generate.InMemoryGeneratedFiles;
+import org.springframework.aot.test.generate.GeneratedFilesTestCompilerUtils;
 import org.springframework.core.test.tools.CompileWithForkedClassLoader;
 import org.springframework.core.test.tools.TestCompiler;
 import org.springframework.test.context.aot.samples.basic.BasicSpringJupiterSharedConfigTests;
@@ -92,7 +94,7 @@ void endToEndTests() {
 		assertThat(sourceFiles).containsExactlyInAnyOrder(expectedSourceFilesForBasicSpringTests);
 
 		// AOT BUILD-TIME: COMPILATION
-		TestCompiler.forSystem().withFiles(generatedFiles)
+		TestCompiler.forSystem().with(setupGeneratedFiles(generatedFiles))
 			// .printFiles(System.out)
 			.compile(compiled ->
 				// AOT RUN-TIME: EXECUTION
@@ -123,13 +125,17 @@ void endToEndTestsForEntireSpringTestModule() {
 		generator.processAheadOfTime(testClasses.stream());
 
 		// AOT BUILD-TIME: COMPILATION
-		TestCompiler.forSystem().withFiles(generatedFiles)
+		TestCompiler.forSystem().with(setupGeneratedFiles(generatedFiles))
 			// .printFiles(System.out)
 			.compile(compiled ->
 				// AOT RUN-TIME: EXECUTION
 				runTestsInAotMode(testClasses));
 	}
 
+	private static Function<TestCompiler, TestCompiler> setupGeneratedFiles(InMemoryGeneratedFiles generatedFiles) {
+		return testCompiler -> GeneratedFilesTestCompilerUtils.configure(testCompiler, generatedFiles);
+	}
+
 	private static void runTestsInAotMode(List<Class<?>> testClasses) {
 		runTestsInAotMode(-1, testClasses);
 	}

spring-test/src/test/java/org/springframework/test/context/aot/TestContextAotGeneratorTests.java

@@ -20,6 +20,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Stream;
 
 import javax.sql.DataSource;
@@ -33,6 +34,7 @@
 import org.springframework.aot.hint.MemberCategory;
 import org.springframework.aot.hint.RuntimeHints;
 import org.springframework.aot.hint.TypeReference;
+import org.springframework.aot.test.generate.GeneratedFilesTestCompilerUtils;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextInitializer;
 import org.springframework.context.ConfigurableApplicationContext;
@@ -83,6 +85,10 @@
 @CompileWithForkedClassLoader
 class TestContextAotGeneratorTests extends AbstractAotTests {
 
+	private static Function<TestCompiler, TestCompiler> setupGeneratedFiles(InMemoryGeneratedFiles generatedFiles) {
+		return testCompiler -> GeneratedFilesTestCompilerUtils.configure(testCompiler, generatedFiles);
+	}
+
 	/**
 	 * End-to-end tests within the scope of the {@link TestContextAotGenerator}.
 	 *
@@ -110,7 +116,7 @@ void endToEndTests() {
 		List<String> sourceFiles = generatedFiles.getGeneratedFiles(Kind.SOURCE).keySet().stream().toList();
 		assertThat(sourceFiles).containsExactlyInAnyOrder(expectedSourceFiles);
 
-		TestCompiler.forSystem().withFiles(generatedFiles).compile(ThrowingConsumer.of(compiled -> {
+		TestCompiler.forSystem().with(setupGeneratedFiles(generatedFiles)).compile(ThrowingConsumer.of(compiled -> {
 			try {
 				System.setProperty(AotDetector.AOT_ENABLED, "true");
 				AotTestAttributesFactory.reset();
@@ -321,7 +327,7 @@ private void processAheadOfTime(Set<Class<?>> testClasses, ThrowingConsumer<Appl
 		InMemoryGeneratedFiles generatedFiles = new InMemoryGeneratedFiles();
 		TestContextAotGenerator generator = new TestContextAotGenerator(generatedFiles);
 		List<Mapping> mappings = processAheadOfTime(generator, testClasses);
-		TestCompiler.forSystem().withFiles(generatedFiles).compile(ThrowingConsumer.of(compiled -> {
+		TestCompiler.forSystem().with(setupGeneratedFiles(generatedFiles)).compile(ThrowingConsumer.of(compiled -> {
 			for (Mapping mapping : mappings) {
 				MergedContextConfiguration mergedConfig = mapping.mergedConfig();
 				ApplicationContextInitializer<ConfigurableApplicationContext> contextInitializer =