Migrate AOT tests to use GeneratedClasses and refine/polish AOT APIs

Migrate all AOT tests to make use of `GeneratedClasses` rather than
directly generating Java files. This commit also refines and polishes
AOT APIs to being greater consistency.

Specifically:

	- The `MethodGenerator` interface has been removed in favor of
	  working directly with `GeneratedMethods`.
	- The visibility of several constructors and methods has been
	  reduced to package-private.
	- The `using(...)` and `builder` methods have been removed in
	  favor of setting the `Consumer` callbacks directly as
	  constructor arguments.
	- Variable names for builders are now named `type` or `method`
	  depending on what they're building.

Closes gh-28831

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

@@ -21,13 +21,11 @@ import java.util.function.BiFunction;
 
 import javax.lang.model.element.Modifier;
 
-import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.aot.generate.DefaultGenerationContext;
 import org.springframework.aot.generate.InMemoryGeneratedFiles;
 import org.springframework.aot.generate.MethodReference;
-import org.springframework.aot.hint.RuntimeHints;
 import org.springframework.aot.hint.predicate.RuntimeHintsPredicates;
 import org.springframework.aot.test.generator.compile.CompileWithTargetClassAccess;
 import org.springframework.aot.test.generator.compile.Compiled;
@@ -41,10 +39,8 @@ import org.springframework.core.env.Environment;
 import org.springframework.core.env.StandardEnvironment;
 import org.springframework.core.testfixture.aot.generate.TestGenerationContext;
 import org.springframework.javapoet.CodeBlock;
-import org.springframework.javapoet.JavaFile;
 import org.springframework.javapoet.MethodSpec;
 import org.springframework.javapoet.ParameterizedTypeName;
-import org.springframework.javapoet.TypeSpec;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -57,25 +53,23 @@ import static org.assertj.core.api.Assertions.assertThat;
  */
 class AutowiredAnnotationBeanRegistrationAotContributionTests {
 
-	private InMemoryGeneratedFiles generatedFiles;
+	private final InMemoryGeneratedFiles generatedFiles;
 
-	private DefaultGenerationContext generationContext;
+	private final DefaultGenerationContext generationContext;
 
-	private RuntimeHints runtimeHints;
+	private final MockBeanRegistrationCode beanRegistrationCode;
 
-	private MockBeanRegistrationCode beanRegistrationCode;
+	private final DefaultListableBeanFactory beanFactory;
 
-	private DefaultListableBeanFactory beanFactory;
 
-	@BeforeEach
-	void setup() {
+	AutowiredAnnotationBeanRegistrationAotContributionTests() {
 		this.generatedFiles = new InMemoryGeneratedFiles();
 		this.generationContext = new TestGenerationContext(this.generatedFiles);
-		this.runtimeHints = this.generationContext.getRuntimeHints();
-		this.beanRegistrationCode = new MockBeanRegistrationCode();
+		this.beanRegistrationCode = new MockBeanRegistrationCode(this.generationContext);
 		this.beanFactory = new DefaultListableBeanFactory();
 	}
 
+
 	@Test
 	void contributeWhenPrivateFieldInjectionInjectsUsingReflection() {
 		Environment environment = new StandardEnvironment();
@@ -84,8 +78,8 @@ class AutowiredAnnotationBeanRegistrationAotContributionTests {
 				PrivateFieldInjectionSample.class);
 		assertThat(RuntimeHintsPredicates.reflection()
 				.onField(PrivateFieldInjectionSample.class, "environment").allowWrite())
-				.accepts(this.runtimeHints);
-		testCompiledResult(registeredBean, (postProcessor, compiled) -> {
+				.accepts(this.generationContext.getRuntimeHints());
+		compile(registeredBean, (postProcessor, compiled) -> {
 			PrivateFieldInjectionSample instance = new PrivateFieldInjectionSample();
 			postProcessor.apply(registeredBean, instance);
 			assertThat(instance).extracting("environment").isSameAs(environment);
@@ -103,8 +97,8 @@ class AutowiredAnnotationBeanRegistrationAotContributionTests {
 				PackagePrivateFieldInjectionSample.class);
 		assertThat(RuntimeHintsPredicates.reflection()
 				.onField(PackagePrivateFieldInjectionSample.class, "environment").allowWrite())
-				.accepts(this.runtimeHints);
-		testCompiledResult(registeredBean, (postProcessor, compiled) -> {
+				.accepts(this.generationContext.getRuntimeHints());
+		compile(registeredBean, (postProcessor, compiled) -> {
 			PackagePrivateFieldInjectionSample instance = new PackagePrivateFieldInjectionSample();
 			postProcessor.apply(registeredBean, instance);
 			assertThat(instance).extracting("environment").isSameAs(environment);
@@ -121,8 +115,8 @@ class AutowiredAnnotationBeanRegistrationAotContributionTests {
 				PrivateMethodInjectionSample.class);
 		assertThat(RuntimeHintsPredicates.reflection()
 				.onMethod(PrivateMethodInjectionSample.class, "setTestBean").invoke())
-				.accepts(this.runtimeHints);
-		testCompiledResult(registeredBean, (postProcessor, compiled) -> {
+				.accepts(this.generationContext.getRuntimeHints());
+		compile(registeredBean, (postProcessor, compiled) -> {
 			PrivateMethodInjectionSample instance = new PrivateMethodInjectionSample();
 			postProcessor.apply(registeredBean, instance);
 			assertThat(instance).extracting("environment").isSameAs(environment);
@@ -140,8 +134,8 @@ class AutowiredAnnotationBeanRegistrationAotContributionTests {
 				PackagePrivateMethodInjectionSample.class);
 		assertThat(RuntimeHintsPredicates.reflection()
 				.onMethod(PackagePrivateMethodInjectionSample.class, "setTestBean").introspect())
-				.accepts(this.runtimeHints);
-		testCompiledResult(registeredBean, (postProcessor, compiled) -> {
+				.accepts(this.generationContext.getRuntimeHints());
+		compile(registeredBean, (postProcessor, compiled) -> {
 			PackagePrivateMethodInjectionSample instance = new PackagePrivateMethodInjectionSample();
 			postProcessor.apply(registeredBean, instance);
 			assertThat(instance).extracting("environment").isSameAs(environment);
@@ -167,29 +161,24 @@ class AutowiredAnnotationBeanRegistrationAotContributionTests {
 	}
 
 	@SuppressWarnings("unchecked")
-	private void testCompiledResult(RegisteredBean registeredBean,
+	private void compile(RegisteredBean registeredBean,
 			BiConsumer<BiFunction<RegisteredBean, Object, Object>, Compiled> result) {
-		this.generationContext.writeGeneratedContent();
-		JavaFile javaFile = createJavaFile(registeredBean.getBeanClass());
-		TestCompiler.forSystem().withFiles(this.generatedFiles).compile(javaFile::writeTo,
-				compiled -> result.accept(compiled.getInstance(BiFunction.class),
-						compiled));
-	}
+		Class<?> target = registeredBean.getBeanClass();
+		MethodReference methodReference = this.beanRegistrationCode.getInstancePostProcessors().get(0);
+		this.beanRegistrationCode.getTypeBuilder().set(type -> {
+			type.addModifiers(Modifier.PUBLIC);
+			type.addSuperinterface(ParameterizedTypeName.get(BiFunction.class, RegisteredBean.class, target, target));
+			type.addMethod(MethodSpec.methodBuilder("apply")
+					.addModifiers(Modifier.PUBLIC)
+					.addParameter(RegisteredBean.class, "registeredBean")
+					.addParameter(target, "instance").returns(target)
+					.addStatement("return $L", methodReference.toInvokeCodeBlock(CodeBlock.of("registeredBean"), CodeBlock.of("instance")))
+					.build());
 
-	private JavaFile createJavaFile(Class<?> target) {
-		MethodReference methodReference = this.beanRegistrationCode.getInstancePostProcessors()
-				.get(0);
-		TypeSpec.Builder builder = TypeSpec.classBuilder("TestPostProcessor");
-		builder.addModifiers(Modifier.PUBLIC);
-		builder.addSuperinterface(ParameterizedTypeName.get(BiFunction.class,
-				RegisteredBean.class, target, target));
-		builder.addMethod(MethodSpec.methodBuilder("apply").addModifiers(Modifier.PUBLIC)
-				.addParameter(RegisteredBean.class, "registeredBean")
-				.addParameter(target, "instance").returns(target)
-				.addStatement("return $L", methodReference.toInvokeCodeBlock(
-						CodeBlock.of("registeredBean"), CodeBlock.of("instance")))
-				.build());
-		return JavaFile.builder("__", builder.build()).build();
+		});
+		this.generationContext.writeGeneratedContent();
+		TestCompiler.forSystem().withFiles(this.generatedFiles).compile(compiled ->
+				result.accept(compiled.getInstance(BiFunction.class), compiled));
 	}
 
 }

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

@@ -25,7 +25,6 @@ import java.util.function.Supplier;
 
 import javax.lang.model.element.Modifier;
 
-import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.aot.generate.DefaultGenerationContext;
@@ -51,12 +50,9 @@ import org.springframework.beans.testfixture.beans.factory.aot.MockBeanRegistrat
 import org.springframework.core.ResolvableType;
 import org.springframework.core.mock.MockSpringFactoriesLoader;
 import org.springframework.core.testfixture.aot.generate.TestGenerationContext;
-import org.springframework.javapoet.ClassName;
 import org.springframework.javapoet.CodeBlock;
-import org.springframework.javapoet.JavaFile;
 import org.springframework.javapoet.MethodSpec;
 import org.springframework.javapoet.ParameterizedTypeName;
-import org.springframework.javapoet.TypeSpec;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -68,27 +64,27 @@ import static org.assertj.core.api.Assertions.assertThat;
  */
 class BeanDefinitionMethodGeneratorTests {
 
-	private InMemoryGeneratedFiles generatedFiles;
+	private final InMemoryGeneratedFiles generatedFiles;
 
-	private DefaultGenerationContext generationContext;
+	private final DefaultGenerationContext generationContext;
 
-	private DefaultListableBeanFactory beanFactory;
+	private final DefaultListableBeanFactory beanFactory;
 
-	private MockBeanRegistrationsCode beanRegistrationsCode;
+	private final MockBeanRegistrationsCode beanRegistrationsCode;
 
-	private BeanDefinitionMethodGeneratorFactory methodGeneratorFactory;
+	private final BeanDefinitionMethodGeneratorFactory methodGeneratorFactory;
 
-	@BeforeEach
-	void setup() {
+
+	BeanDefinitionMethodGeneratorTests() {
 		this.generatedFiles = new InMemoryGeneratedFiles();
 		this.generationContext = new TestGenerationContext(this.generatedFiles);
 		this.beanFactory = new DefaultListableBeanFactory();
 		this.methodGeneratorFactory = new BeanDefinitionMethodGeneratorFactory(
 				new AotFactoriesLoader(this.beanFactory, new MockSpringFactoriesLoader()));
-		this.beanRegistrationsCode = new MockBeanRegistrationsCode(
-				ClassName.get("__", "Registration"));
+		this.beanRegistrationsCode = new MockBeanRegistrationsCode(this.generationContext);
 	}
 
+
 	@Test
 	void generateBeanDefinitionMethodGeneratesMethod() {
 		RegisteredBean registeredBean = registerBean(
@@ -98,7 +94,7 @@ class BeanDefinitionMethodGeneratorTests {
 				Collections.emptyList());
 		MethodReference method = generator.generateBeanDefinitionMethod(
 				this.generationContext, this.beanRegistrationsCode);
-		testCompiledResult(method, (actual, compiled) -> {
+		compile(method, (actual, compiled) -> {
 			SourceFile sourceFile = compiled.getSourceFile(".*BeanDefinitions");
 			assertThat(sourceFile).contains("Get the bean definition for 'testBean'");
 			assertThat(sourceFile).contains("beanType = TestBean.class");
@@ -116,7 +112,7 @@ class BeanDefinitionMethodGeneratorTests {
 				Collections.emptyList());
 		MethodReference method = generator.generateBeanDefinitionMethod(
 				this.generationContext, this.beanRegistrationsCode);
-		testCompiledResult(method, (actual, compiled) -> {
+		compile(method, (actual, compiled) -> {
 			assertThat(actual.getResolvableType().resolve()).isEqualTo(GenericBean.class);
 			SourceFile sourceFile = compiled.getSourceFile(".*BeanDefinitions");
 			assertThat(sourceFile).contains("Get the bean definition for 'testBean'");
@@ -133,15 +129,13 @@ class BeanDefinitionMethodGeneratorTests {
 				new RootBeanDefinition(TestBean.class));
 		BeanRegistrationAotContribution aotContribution = (generationContext,
 				beanRegistrationCode) -> {
-			GeneratedMethod method = beanRegistrationCode.getMethodGenerator()
-					.generateMethod("postProcess")
-					.using(builder -> builder.addModifiers(Modifier.STATIC)
+			GeneratedMethod generatedMethod = beanRegistrationCode.getMethods().add("postProcess", method ->
+					method.addModifiers(Modifier.STATIC)
 							.addParameter(RegisteredBean.class, "registeredBean")
 							.addParameter(TestBean.class, "testBean")
-							.returns(TestBean.class).addCode("return new $T($S);",
-									TestBean.class, "postprocessed"));
+							.returns(TestBean.class).addCode("return new $T($S);", TestBean.class, "postprocessed"));
 			beanRegistrationCode.addInstancePostProcessor(MethodReference.ofStatic(
-					beanRegistrationCode.getClassName(), method.getName()));
+					beanRegistrationCode.getClassName(), generatedMethod.getName()));
 		};
 		List<BeanRegistrationAotContribution> aotContributions = Collections
 				.singletonList(aotContribution);
@@ -149,7 +143,7 @@ class BeanDefinitionMethodGeneratorTests {
 				this.methodGeneratorFactory, registeredBean, null, aotContributions);
 		MethodReference method = generator.generateBeanDefinitionMethod(
 				this.generationContext, this.beanRegistrationsCode);
-		testCompiledResult(method, (actual, compiled) -> {
+		compile(method, (actual, compiled) -> {
 			assertThat(actual.getBeanClass()).isEqualTo(TestBean.class);
 			InstanceSupplier<?> supplier = (InstanceSupplier<?>) actual
 					.getInstanceSupplier();
@@ -175,7 +169,7 @@ class BeanDefinitionMethodGeneratorTests {
 				this.methodGeneratorFactory, registeredBean, null, aotContributions);
 		MethodReference method = generator.generateBeanDefinitionMethod(
 				this.generationContext, this.beanRegistrationsCode);
-		testCompiledResult(method, (actual, compiled) -> {
+		compile(method, (actual, compiled) -> {
 			assertThat(actual.getBeanClass()).isEqualTo(TestBean.class);
 			SourceFile sourceFile = compiled.getSourceFile(".*BeanDefinitions");
 			assertThat(sourceFile).contains("I am custom");
@@ -215,7 +209,7 @@ class BeanDefinitionMethodGeneratorTests {
 				aotContributions);
 		MethodReference method = generator.generateBeanDefinitionMethod(
 				this.generationContext, this.beanRegistrationsCode);
-		testCompiledResult(method, (actual, compiled) -> {
+		compile(method, (actual, compiled) -> {
 			assertThat(actual.getAttribute("a")).isEqualTo("A");
 			assertThat(actual.getAttribute("b")).isNull();
 		});
@@ -248,7 +242,7 @@ class BeanDefinitionMethodGeneratorTests {
 				Collections.emptyList());
 		MethodReference method = generator.generateBeanDefinitionMethod(
 				this.generationContext, this.beanRegistrationsCode);
-		testCompiledResult(method, (actual, compiled) -> {
+		compile(method, (actual, compiled) -> {
 			assertThat(compiled.getSourceFile(".*BeanDefinitions"))
 					.contains("Get the inner-bean definition for 'testInnerBean'");
 			assertThat(actual).isInstanceOf(RootBeanDefinition.class);
@@ -269,7 +263,7 @@ class BeanDefinitionMethodGeneratorTests {
 				Collections.emptyList());
 		MethodReference method = generator.generateBeanDefinitionMethod(
 				this.generationContext, this.beanRegistrationsCode);
-		testCompiledResult(method, (actual, compiled) -> {
+		compile(method, (actual, compiled) -> {
 			RootBeanDefinition actualInnerBeanDefinition = (RootBeanDefinition) actual
 					.getPropertyValues().get("name");
 			assertThat(actualInnerBeanDefinition.isPrimary()).isTrue();
@@ -303,7 +297,7 @@ class BeanDefinitionMethodGeneratorTests {
 				Collections.emptyList());
 		MethodReference method = generator.generateBeanDefinitionMethod(
 				this.generationContext, this.beanRegistrationsCode);
-		testCompiledResult(method, (actual, compiled) -> {
+		compile(method, (actual, compiled) -> {
 			RootBeanDefinition actualInnerBeanDefinition = (RootBeanDefinition) actual
 					.getConstructorArgumentValues()
 					.getIndexedArgumentValue(0, RootBeanDefinition.class).getValue();
@@ -328,15 +322,14 @@ class BeanDefinitionMethodGeneratorTests {
 		RegisteredBean registeredBean = registerBean(
 				new RootBeanDefinition(TestBean.class));
 		List<BeanRegistrationAotContribution> aotContributions = new ArrayList<>();
-		aotContributions
-				.add((generationContext, beanRegistrationCode) -> beanRegistrationCode
-						.getMethodGenerator().generateMethod("aotContributedMethod")
-						.using(builder -> builder.addComment("Example Contribution")));
+		aotContributions.add((generationContext, beanRegistrationCode) ->
+				beanRegistrationCode.getMethods().add("aotContributedMethod", method ->
+						method.addComment("Example Contribution")));
 		BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(
 				this.methodGeneratorFactory, registeredBean, null, aotContributions);
 		MethodReference method = generator.generateBeanDefinitionMethod(
 				this.generationContext, this.beanRegistrationsCode);
-		testCompiledResult(method, (actual, compiled) -> {
+		compile(method, (actual, compiled) -> {
 			SourceFile sourceFile = compiled.getSourceFile(".*BeanDefinitions");
 			assertThat(sourceFile).contains("AotContributedMethod()");
 			assertThat(sourceFile).contains("Example Contribution");
@@ -353,7 +346,7 @@ class BeanDefinitionMethodGeneratorTests {
 				Collections.emptyList());
 		MethodReference method = generator.generateBeanDefinitionMethod(
 				this.generationContext, this.beanRegistrationsCode);
-		testCompiledResult(method, (actual, compiled) -> {
+		compile(method, (actual, compiled) -> {
 			DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
 			freshBeanFactory.registerBeanDefinition("test", actual);
 			Object bean = freshBeanFactory.getBean("test");
@@ -369,27 +362,19 @@ class BeanDefinitionMethodGeneratorTests {
 		return RegisteredBean.of(this.beanFactory, beanName);
 	}
 
-	private void testCompiledResult(MethodReference method,
+	private void compile(MethodReference method,
 			BiConsumer<RootBeanDefinition, Compiled> result) {
+		this.beanRegistrationsCode.getTypeBuilder().set(type -> {
+			type.addModifiers(Modifier.PUBLIC);
+			type.addSuperinterface(ParameterizedTypeName.get(Supplier.class, BeanDefinition.class));
+			type.addMethod(MethodSpec.methodBuilder("get")
+					.addModifiers(Modifier.PUBLIC)
+					.returns(BeanDefinition.class)
+					.addCode("return $L;", method.toInvokeCodeBlock()).build());
+		});
 		this.generationContext.writeGeneratedContent();
-		JavaFile javaFile = generateJavaFile(method);
-		TestCompiler.forSystem().withFiles(this.generatedFiles).compile(
-				javaFile::writeTo, compiled -> result.accept(
-						(RootBeanDefinition) compiled.getInstance(Supplier.class).get(),
-						compiled));
-	}
-
-	private JavaFile generateJavaFile(MethodReference method) {
-		TypeSpec.Builder builder = TypeSpec.classBuilder("Registration");
-		builder.addModifiers(Modifier.PUBLIC);
-		builder.addSuperinterface(
-				ParameterizedTypeName.get(Supplier.class, BeanDefinition.class));
-		builder.addMethod(MethodSpec.methodBuilder("get").addModifiers(Modifier.PUBLIC)
-				.returns(BeanDefinition.class)
-				.addCode("return $L;", method.toInvokeCodeBlock()).build());
-		this.beanRegistrationsCode.getMethodGenerator()
-				.doWithMethodSpecs(builder::addMethod);
-		return JavaFile.builder("__", builder.build()).build();
+		TestCompiler.forSystem().withFiles(this.generatedFiles).compile(compiled ->
+				result.accept((RootBeanDefinition) compiled.getInstance(Supplier.class).get(), compiled));
 	}
 
 }

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

@@ -27,8 +27,9 @@ import javax.lang.model.element.Modifier;
 
 import org.junit.jupiter.api.Test;
 
-import org.springframework.aot.generate.GeneratedMethods;
-import org.springframework.aot.hint.RuntimeHints;
+import org.springframework.aot.generate.DefaultGenerationContext;
+import org.springframework.aot.generate.GeneratedClass;
+import org.springframework.aot.generate.InMemoryGeneratedFiles;
 import org.springframework.aot.hint.predicate.RuntimeHintsPredicates;
 import org.springframework.aot.test.generator.compile.Compiled;
 import org.springframework.aot.test.generator.compile.TestCompiler;
@@ -43,11 +44,11 @@ import org.springframework.beans.factory.support.ManagedList;
 import org.springframework.beans.factory.support.ManagedMap;
 import org.springframework.beans.factory.support.ManagedSet;
 import org.springframework.beans.factory.support.RootBeanDefinition;
+import org.springframework.beans.testfixture.beans.factory.aot.DeferredTypeBuilder;
+import org.springframework.core.testfixture.aot.generate.TestGenerationContext;
 import org.springframework.javapoet.CodeBlock;
-import org.springframework.javapoet.JavaFile;
 import org.springframework.javapoet.MethodSpec;
 import org.springframework.javapoet.ParameterizedTypeName;
-import org.springframework.javapoet.TypeSpec;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -61,18 +62,14 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 
 	private final RootBeanDefinition beanDefinition = new RootBeanDefinition();
 
-	private final GeneratedMethods generatedMethods = new GeneratedMethods();
-
-	private final RuntimeHints hints = new RuntimeHints();
-
-	private BeanDefinitionPropertiesCodeGenerator generator = new BeanDefinitionPropertiesCodeGenerator(
-			this.hints, attribute -> true, this.generatedMethods, (name, value) -> null);
+	private final InMemoryGeneratedFiles generatedFiles = new InMemoryGeneratedFiles();
 
+	private final DefaultGenerationContext generationContext = new TestGenerationContext(this.generatedFiles);
 
 	@Test
 	void setPrimaryWhenFalse() {
 		this.beanDefinition.setPrimary(false);
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(compiled.getSourceFile()).doesNotContain("setPrimary");
 			assertThat(actual.isPrimary()).isFalse();
 		});
@@ -81,13 +78,13 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	@Test
 	void setPrimaryWhenTrue() {
 		this.beanDefinition.setPrimary(true);
-		testCompiledResult((actual, compiled) -> assertThat(actual.isPrimary()).isTrue());
+		compile((actual, compiled) -> assertThat(actual.isPrimary()).isTrue());
 	}
 
 	@Test
 	void setScopeWhenEmptyString() {
 		this.beanDefinition.setScope("");
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(compiled.getSourceFile()).doesNotContain("setScope");
 			assertThat(actual.getScope()).isEmpty();
 		});
@@ -96,7 +93,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	@Test
 	void setScopeWhenSingleton() {
 		this.beanDefinition.setScope("singleton");
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(compiled.getSourceFile()).doesNotContain("setScope");
 			assertThat(actual.getScope()).isEmpty();
 		});
@@ -105,14 +102,14 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	@Test
 	void setScopeWhenOther() {
 		this.beanDefinition.setScope("prototype");
-		testCompiledResult((actual, compiled) -> assertThat(actual.getScope())
+		compile((actual, compiled) -> assertThat(actual.getScope())
 				.isEqualTo("prototype"));
 	}
 
 	@Test
 	void setDependsOnWhenEmpty() {
 		this.beanDefinition.setDependsOn();
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(compiled.getSourceFile()).doesNotContain("setDependsOn");
 			assertThat(actual.getDependsOn()).isNull();
 		});
@@ -121,13 +118,13 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	@Test
 	void setDependsOnWhenNotEmpty() {
 		this.beanDefinition.setDependsOn("a", "b", "c");
-		testCompiledResult((actual, compiled) -> assertThat(actual.getDependsOn())
+		compile((actual, compiled) -> assertThat(actual.getDependsOn())
 				.containsExactly("a", "b", "c"));
 	}
 
 	@Test
 	void setLazyInitWhenNoSet() {
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(compiled.getSourceFile()).doesNotContain("setLazyInit");
 			assertThat(actual.isLazyInit()).isFalse();
 			assertThat(actual.getLazyInit()).isNull();
@@ -137,7 +134,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	@Test
 	void setLazyInitWhenFalse() {
 		this.beanDefinition.setLazyInit(false);
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(actual.isLazyInit()).isFalse();
 			assertThat(actual.getLazyInit()).isFalse();
 		});
@@ -146,7 +143,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	@Test
 	void setLazyInitWhenTrue() {
 		this.beanDefinition.setLazyInit(true);
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(actual.isLazyInit()).isTrue();
 			assertThat(actual.getLazyInit()).isTrue();
 		});
@@ -155,14 +152,14 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	@Test
 	void setAutowireCandidateWhenFalse() {
 		this.beanDefinition.setAutowireCandidate(false);
-		testCompiledResult(
+		compile(
 				(actual, compiled) -> assertThat(actual.isAutowireCandidate()).isFalse());
 	}
 
 	@Test
 	void setAutowireCandidateWhenTrue() {
 		this.beanDefinition.setAutowireCandidate(true);
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(compiled.getSourceFile()).doesNotContain("setAutowireCandidate");
 			assertThat(actual.isAutowireCandidate()).isTrue();
 		});
@@ -171,7 +168,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	@Test
 	void setSyntheticWhenFalse() {
 		this.beanDefinition.setSynthetic(false);
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(compiled.getSourceFile()).doesNotContain("setSynthetic");
 			assertThat(actual.isSynthetic()).isFalse();
 		});
@@ -180,14 +177,14 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	@Test
 	void setSyntheticWhenTrue() {
 		this.beanDefinition.setSynthetic(true);
-		testCompiledResult(
+		compile(
 				(actual, compiled) -> assertThat(actual.isSynthetic()).isTrue());
 	}
 
 	@Test
 	void setRoleWhenApplication() {
 		this.beanDefinition.setRole(BeanDefinition.ROLE_APPLICATION);
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(compiled.getSourceFile()).doesNotContain("setRole");
 			assertThat(actual.getRole()).isEqualTo(BeanDefinition.ROLE_APPLICATION);
 		});
@@ -196,7 +193,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	@Test
 	void setRoleWhenInfrastructure() {
 		this.beanDefinition.setRole(BeanDefinition.ROLE_INFRASTRUCTURE);
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(compiled.getSourceFile())
 					.contains("setRole(BeanDefinition.ROLE_INFRASTRUCTURE);");
 			assertThat(actual.getRole()).isEqualTo(BeanDefinition.ROLE_INFRASTRUCTURE);
@@ -206,7 +203,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	@Test
 	void setRoleWhenSupport() {
 		this.beanDefinition.setRole(BeanDefinition.ROLE_SUPPORT);
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(compiled.getSourceFile())
 					.contains("setRole(BeanDefinition.ROLE_SUPPORT);");
 			assertThat(actual.getRole()).isEqualTo(BeanDefinition.ROLE_SUPPORT);
@@ -216,7 +213,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	@Test
 	void setRoleWhenOther() {
 		this.beanDefinition.setRole(999);
-		testCompiledResult(
+		compile(
 				(actual, compiled) -> assertThat(actual.getRole()).isEqualTo(999));
 	}
 
@@ -224,7 +221,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	void setInitMethodWhenSingleInitMethod() {
 		this.beanDefinition.setTargetType(InitDestroyBean.class);
 		this.beanDefinition.setInitMethodName("i1");
-		testCompiledResult((actual, compiled) -> assertThat(actual.getInitMethodNames())
+		compile((actual, compiled) -> assertThat(actual.getInitMethodNames())
 				.containsExactly("i1"));
 		String[] methodNames = { "i1" };
 		assertHasMethodInvokeHints(InitDestroyBean.class, methodNames);
@@ -234,14 +231,14 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	void setInitMethodWhenSingleInferredInitMethod() {
 		this.beanDefinition.setTargetType(InitDestroyBean.class);
 		this.beanDefinition.setInitMethodName(AbstractBeanDefinition.INFER_METHOD);
-		testCompiledResult((actual, compiled) -> assertThat(actual.getInitMethodNames()).isNull());
+		compile((actual, compiled) -> assertThat(actual.getInitMethodNames()).isNull());
 	}
 
 	@Test
 	void setInitMethodWhenMultipleInitMethods() {
 		this.beanDefinition.setTargetType(InitDestroyBean.class);
 		this.beanDefinition.setInitMethodNames("i1", "i2");
-		testCompiledResult((actual, compiled) -> assertThat(actual.getInitMethodNames())
+		compile((actual, compiled) -> assertThat(actual.getInitMethodNames())
 				.containsExactly("i1", "i2"));
 		String[] methodNames = { "i1", "i2" };
 		assertHasMethodInvokeHints(InitDestroyBean.class, methodNames);
@@ -251,7 +248,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	void setDestroyMethodWhenDestroyInitMethod() {
 		this.beanDefinition.setTargetType(InitDestroyBean.class);
 		this.beanDefinition.setDestroyMethodName("d1");
-		testCompiledResult(
+		compile(
 				(actual, compiled) -> assertThat(actual.getDestroyMethodNames())
 						.containsExactly("d1"));
 		String[] methodNames = { "d1" };
@@ -262,14 +259,14 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	void setDestroyMethodWhenSingleInferredInitMethod() {
 		this.beanDefinition.setTargetType(InitDestroyBean.class);
 		this.beanDefinition.setDestroyMethodName(AbstractBeanDefinition.INFER_METHOD);
-		testCompiledResult((actual, compiled) -> assertThat(actual.getDestroyMethodNames()).isNull());
+		compile((actual, compiled) -> assertThat(actual.getDestroyMethodNames()).isNull());
 	}
 
 	@Test
 	void setDestroyMethodWhenMultipleDestroyMethods() {
 		this.beanDefinition.setTargetType(InitDestroyBean.class);
 		this.beanDefinition.setDestroyMethodNames("d1", "d2");
-		testCompiledResult(
+		compile(
 				(actual, compiled) -> assertThat(actual.getDestroyMethodNames())
 						.containsExactly("d1", "d2"));
 		String[] methodNames = { "d1", "d2" };
@@ -277,8 +274,9 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	}
 
 	private void assertHasMethodInvokeHints(Class<?> beanType, String... methodNames) {
-		assertThat(methodNames).allMatch(methodName ->
-				RuntimeHintsPredicates.reflection().onMethod(beanType, methodName).invoke().test(this.hints));
+		assertThat(methodNames).allMatch(methodName -> RuntimeHintsPredicates.reflection()
+				.onMethod(beanType, methodName).invoke()
+				.test(this.generationContext.getRuntimeHints()));
 	}
 
 	@Test
@@ -289,7 +287,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 				"test");
 		this.beanDefinition.getConstructorArgumentValues().addIndexedArgumentValue(2,
 				123);
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			Map<Integer, ValueHolder> values = actual.getConstructorArgumentValues()
 					.getIndexedArgumentValues();
 			assertThat(values.get(0).getValue()).isEqualTo(String.class);
@@ -303,7 +301,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 		this.beanDefinition.setTargetType(PropertyValuesBean.class);
 		this.beanDefinition.getPropertyValues().add("test", String.class);
 		this.beanDefinition.getPropertyValues().add("spring", "framework");
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(actual.getPropertyValues().get("test")).isEqualTo(String.class);
 			assertThat(actual.getPropertyValues().get("spring")).isEqualTo("framework");
 		});
@@ -315,7 +313,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 	void propertyValuesWhenContainsBeanReference() {
 		this.beanDefinition.getPropertyValues().add("myService",
 				new RuntimeBeanNameReference("test"));
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(actual.getPropertyValues().contains("myService")).isTrue();
 			assertThat(actual.getPropertyValues().get("myService"))
 					.isInstanceOfSatisfying(RuntimeBeanReference.class,
@@ -329,7 +327,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 		ManagedList<Object> managedList = new ManagedList<>();
 		managedList.add(new RuntimeBeanNameReference("test"));
 		this.beanDefinition.getPropertyValues().add("value", managedList);
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			Object value = actual.getPropertyValues().get("value");
 			assertThat(value).isInstanceOf(ManagedList.class);
 			assertThat(((List<?>) value).get(0)).isInstanceOf(BeanReference.class);
@@ -341,7 +339,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 		ManagedSet<Object> managedSet = new ManagedSet<>();
 		managedSet.add(new RuntimeBeanNameReference("test"));
 		this.beanDefinition.getPropertyValues().add("value", managedSet);
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			Object value = actual.getPropertyValues().get("value");
 			assertThat(value).isInstanceOf(ManagedSet.class);
 			assertThat(((Set<?>) value).iterator().next())
@@ -354,7 +352,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 		ManagedMap<String, Object> managedMap = new ManagedMap<>();
 		managedMap.put("test", new RuntimeBeanNameReference("test"));
 		this.beanDefinition.getPropertyValues().add("value", managedMap);
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			Object value = actual.getPropertyValues().get("value");
 			assertThat(value).isInstanceOf(ManagedMap.class);
 			assertThat(((Map<?, ?>) value).get("test")).isInstanceOf(BeanReference.class);
@@ -366,9 +364,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 		this.beanDefinition.setAttribute("a", "A");
 		this.beanDefinition.setAttribute("b", "B");
 		Predicate<String> attributeFilter = attribute -> false;
-		this.generator = new BeanDefinitionPropertiesCodeGenerator(this.hints,
-				attributeFilter, this.generatedMethods, (name, value) -> null);
-		testCompiledResult((actual, compiled) -> {
+		compile(attributeFilter, (actual, compiled) -> {
 			assertThat(compiled.getSourceFile()).doesNotContain("setAttribute");
 			assertThat(actual.getAttribute("a")).isNull();
 			assertThat(actual.getAttribute("b")).isNull();
@@ -380,9 +376,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 		this.beanDefinition.setAttribute("a", "A");
 		this.beanDefinition.setAttribute("b", "B");
 		Predicate<String> attributeFilter = "a"::equals;
-		this.generator = new BeanDefinitionPropertiesCodeGenerator(this.hints,
-				attributeFilter, this.generatedMethods, (name, value) -> null);
-		testCompiledResult(this.beanDefinition, (actual, compiled) -> {
+		compile(attributeFilter, (actual, compiled) -> {
 			assertThat(actual.getAttribute("a")).isEqualTo("A");
 			assertThat(actual.getAttribute("b")).isNull();
 		});
@@ -393,47 +387,43 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 		this.beanDefinition.setPrimary(true);
 		this.beanDefinition.setScope("test");
 		this.beanDefinition.setRole(BeanDefinition.ROLE_SUPPORT);
-		testCompiledResult((actual, compiled) -> {
+		compile((actual, compiled) -> {
 			assertThat(actual.isPrimary()).isTrue();
 			assertThat(actual.getScope()).isEqualTo("test");
 			assertThat(actual.getRole()).isEqualTo(BeanDefinition.ROLE_SUPPORT);
 		});
 	}
 
-	private void testCompiledResult(BiConsumer<RootBeanDefinition, Compiled> result) {
-		testCompiledResult(this.beanDefinition, result);
+	private void compile(BiConsumer<RootBeanDefinition, Compiled> result) {
+		compile(attribute -> true, result);
 	}
 
-	private void testCompiledResult(RootBeanDefinition beanDefinition,
+	private void compile(
+			Predicate<String> attributeFilter,
 			BiConsumer<RootBeanDefinition, Compiled> result) {
-		testCompiledResult(() -> this.generator.generateCode(beanDefinition), result);
-	}
-
-	private void testCompiledResult(Supplier<CodeBlock> codeBlock,
-			BiConsumer<RootBeanDefinition, Compiled> result) {
-		JavaFile javaFile = createJavaFile(codeBlock);
-		TestCompiler.forSystem().compile(javaFile::writeTo, compiled -> {
-			RootBeanDefinition beanDefinition = (RootBeanDefinition) compiled
-					.getInstance(Supplier.class).get();
-			result.accept(beanDefinition, compiled);
+		DeferredTypeBuilder typeBuilder = new DeferredTypeBuilder();
+		GeneratedClass generatedClass = this.generationContext.getGeneratedClasses().addForFeature("TestCode", typeBuilder);
+		BeanDefinitionPropertiesCodeGenerator codeGenerator = new BeanDefinitionPropertiesCodeGenerator(
+				this.generationContext.getRuntimeHints(), attributeFilter,
+				generatedClass.getMethods(), (name, value) -> null);
+		CodeBlock generatedCode = codeGenerator.generateCode(this.beanDefinition);
+		typeBuilder.set(type -> {
+			type.addModifiers(Modifier.PUBLIC);
+			type.addSuperinterface(ParameterizedTypeName.get(Supplier.class, RootBeanDefinition.class));
+			type.addMethod(MethodSpec.methodBuilder("get")
+					.addModifiers(Modifier.PUBLIC)
+					.returns(RootBeanDefinition.class)
+					.addStatement("$T beanDefinition = new $T()", RootBeanDefinition.class, RootBeanDefinition.class)
+					.addStatement("$T beanFactory = new $T()", DefaultListableBeanFactory.class, DefaultListableBeanFactory.class)
+					.addCode(generatedCode)
+					.addStatement("return beanDefinition").build());
+		});
+		this.generationContext.writeGeneratedContent();
+		TestCompiler.forSystem().withFiles(this.generatedFiles).compile(compiled -> {
+			RootBeanDefinition suppliedBeanDefinition = (RootBeanDefinition) compiled
+					.getInstance(Supplier.class).get();
+			result.accept(suppliedBeanDefinition, compiled);
 		});
-	}
-
-	private JavaFile createJavaFile(Supplier<CodeBlock> codeBlock) {
-		TypeSpec.Builder builder = TypeSpec.classBuilder("BeanSupplier");
-		builder.addModifiers(Modifier.PUBLIC);
-		builder.addSuperinterface(
-				ParameterizedTypeName.get(Supplier.class, RootBeanDefinition.class));
-		builder.addMethod(MethodSpec.methodBuilder("get").addModifiers(Modifier.PUBLIC)
-				.returns(RootBeanDefinition.class)
-				.addStatement("$T beanDefinition = new $T()", RootBeanDefinition.class,
-						RootBeanDefinition.class)
-				.addStatement("$T beanFactory = new $T()",
-						DefaultListableBeanFactory.class,
-						DefaultListableBeanFactory.class)
-				.addCode(codeBlock.get()).addStatement("return beanDefinition").build());
-		this.generatedMethods.doWithMethodSpecs(builder::addMethod);
-		return JavaFile.builder("com.example", builder.build()).build();
 	}
 
 	static class InitDestroyBean {

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

@@ -33,21 +33,22 @@ import javax.lang.model.element.Modifier;
 import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
 
-import org.springframework.aot.generate.GeneratedMethods;
+import org.springframework.aot.generate.DefaultGenerationContext;
+import org.springframework.aot.generate.GeneratedClass;
+import org.springframework.aot.generate.InMemoryGeneratedFiles;
 import org.springframework.aot.test.generator.compile.Compiled;
 import org.springframework.aot.test.generator.compile.TestCompiler;
-import org.springframework.aot.test.generator.file.SourceFile;
 import org.springframework.beans.factory.config.BeanReference;
 import org.springframework.beans.factory.config.RuntimeBeanNameReference;
 import org.springframework.beans.factory.support.ManagedList;
 import org.springframework.beans.factory.support.ManagedMap;
 import org.springframework.beans.factory.support.ManagedSet;
+import org.springframework.beans.testfixture.beans.factory.aot.DeferredTypeBuilder;
 import org.springframework.core.ResolvableType;
+import org.springframework.core.testfixture.aot.generate.TestGenerationContext;
 import org.springframework.javapoet.CodeBlock;
-import org.springframework.javapoet.JavaFile;
 import org.springframework.javapoet.MethodSpec;
 import org.springframework.javapoet.ParameterizedTypeName;
-import org.springframework.javapoet.TypeSpec;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -61,28 +62,23 @@ import static org.assertj.core.api.Assertions.assertThat;
  */
 class BeanDefinitionPropertyValueCodeGeneratorTests {
 
-	private GeneratedMethods generatedMethods = new GeneratedMethods();
-
-	private BeanDefinitionPropertyValueCodeGenerator instance = new BeanDefinitionPropertyValueCodeGenerator(
-			generatedMethods);
-
 	private void compile(Object value, BiConsumer<Object, Compiled> result) {
-		CodeBlock code = instance.generateCode(value);
-		JavaFile javaFile = createJavaFile(code);
-		TestCompiler.forSystem().compile(SourceFile.of(javaFile::writeTo),
-				compiled -> result.accept(compiled.getInstance(Supplier.class).get(),
-						compiled));
-	}
-
-	private JavaFile createJavaFile(CodeBlock code) {
-		TypeSpec.Builder builder = TypeSpec.classBuilder("InstanceSupplier");
-		builder.addModifiers(Modifier.PUBLIC);
-		builder.addSuperinterface(
-				ParameterizedTypeName.get(Supplier.class, Object.class));
-		builder.addMethod(MethodSpec.methodBuilder("get").addModifiers(Modifier.PUBLIC)
-				.returns(Object.class).addStatement("return $L", code).build());
-		generatedMethods.doWithMethodSpecs(builder::addMethod);
-		return JavaFile.builder("com.example", builder.build()).build();
+		InMemoryGeneratedFiles generatedFiles = new InMemoryGeneratedFiles();
+		DefaultGenerationContext generationContext = new TestGenerationContext(generatedFiles);
+		DeferredTypeBuilder typeBuilder = new DeferredTypeBuilder();
+		GeneratedClass generatedClass = generationContext.getGeneratedClasses().addForFeature("TestCode", typeBuilder);
+		CodeBlock generatedCode = new BeanDefinitionPropertyValueCodeGenerator(
+				generatedClass.getMethods()).generateCode(value);
+		typeBuilder.set(type -> {
+			type.addModifiers(Modifier.PUBLIC);
+			type.addSuperinterface(
+					ParameterizedTypeName.get(Supplier.class, Object.class));
+			type.addMethod(MethodSpec.methodBuilder("get").addModifiers(Modifier.PUBLIC)
+					.returns(Object.class).addStatement("return $L", generatedCode).build());
+		});
+		generationContext.writeGeneratedContent();
+		TestCompiler.forSystem().withFiles(generatedFiles).compile(compiled ->
+				result.accept(compiled.getInstance(Supplier.class).get(), compiled));
 	}
 
 	@Nested

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

@@ -26,7 +26,6 @@ import java.util.function.Consumer;
 
 import javax.lang.model.element.Modifier;
 
-import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.aot.generate.ClassNameGenerator;
@@ -46,10 +45,8 @@ import org.springframework.core.mock.MockSpringFactoriesLoader;
 import org.springframework.core.testfixture.aot.generate.TestGenerationContext;
 import org.springframework.core.testfixture.aot.generate.TestTarget;
 import org.springframework.javapoet.CodeBlock;
-import org.springframework.javapoet.JavaFile;
 import org.springframework.javapoet.MethodSpec;
 import org.springframework.javapoet.ParameterizedTypeName;
-import org.springframework.javapoet.TypeSpec;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -60,28 +57,30 @@ import static org.assertj.core.api.Assertions.assertThat;
  */
 class BeanRegistrationsAotContributionTests {
 
-	private InMemoryGeneratedFiles generatedFiles;
-
-	private DefaultGenerationContext generationContext;
+	private final MockSpringFactoriesLoader springFactoriesLoader;
 
 	private DefaultListableBeanFactory beanFactory;
 
-	private MockSpringFactoriesLoader springFactoriesLoader;
+	private final InMemoryGeneratedFiles generatedFiles;
 
-	private BeanDefinitionMethodGeneratorFactory methodGeneratorFactory;
+	private DefaultGenerationContext generationContext;
 
-	private MockBeanFactoryInitializationCode beanFactoryInitializationCode = new MockBeanFactoryInitializationCode();
+	private final BeanDefinitionMethodGeneratorFactory methodGeneratorFactory;
 
-	@BeforeEach
-	void setup() {
+	private MockBeanFactoryInitializationCode beanFactoryInitializationCode;
+
+
+	BeanRegistrationsAotContributionTests() {
+		this.springFactoriesLoader = new MockSpringFactoriesLoader();
+		this.beanFactory = new DefaultListableBeanFactory();
 		this.generatedFiles = new InMemoryGeneratedFiles();
 		this.generationContext = new TestGenerationContext(this.generatedFiles);
-		this.beanFactory = new DefaultListableBeanFactory();
-		this.springFactoriesLoader = new MockSpringFactoriesLoader();
 		this.methodGeneratorFactory = new BeanDefinitionMethodGeneratorFactory(
 				new AotFactoriesLoader(this.beanFactory, this.springFactoriesLoader));
+		this.beanFactoryInitializationCode = new MockBeanFactoryInitializationCode(this.generationContext);
 	}
 
+
 	@Test
 	void applyToAppliesContribution() {
 		Map<String, BeanDefinitionMethodGenerator> registrations = new LinkedHashMap<>();
@@ -94,7 +93,7 @@ class BeanRegistrationsAotContributionTests {
 		BeanRegistrationsAotContribution contribution = new BeanRegistrationsAotContribution(
 				registrations);
 		contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
-		testCompiledResult((consumer, compiled) -> {
+		compile((consumer, compiled) -> {
 			DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
 			consumer.accept(freshBeanFactory);
 			assertThat(freshBeanFactory.getBean(TestBean.class)).isNotNull();
@@ -105,7 +104,7 @@ class BeanRegistrationsAotContributionTests {
 	void applyToWhenHasNameGeneratesPrefixedFeatureName() {
 		this.generationContext = new DefaultGenerationContext(
 				new ClassNameGenerator(TestTarget.class, "Management"), this.generatedFiles);
-		this.beanFactoryInitializationCode = new MockBeanFactoryInitializationCode();
+		this.beanFactoryInitializationCode = new MockBeanFactoryInitializationCode(this.generationContext);
 		Map<String, BeanDefinitionMethodGenerator> registrations = new LinkedHashMap<>();
 		RegisteredBean registeredBean = registerBean(
 				new RootBeanDefinition(TestBean.class));
@@ -116,7 +115,7 @@ class BeanRegistrationsAotContributionTests {
 		BeanRegistrationsAotContribution contribution = new BeanRegistrationsAotContribution(
 				registrations);
 		contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
-		testCompiledResult((consumer, compiled) -> {
+		compile((consumer, compiled) -> {
 			SourceFile sourceFile = compiled.getSourceFile(".*BeanDefinitions");
 			assertThat(sourceFile.getClassName()).endsWith("__ManagementBeanDefinitions");
 		});
@@ -148,7 +147,7 @@ class BeanRegistrationsAotContributionTests {
 		contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
 		assertThat(beanRegistrationsCodes).hasSize(1);
 		BeanRegistrationsCode actual = beanRegistrationsCodes.get(0);
-		assertThat(actual.getMethodGenerator()).isNotNull();
+		assertThat(actual.getMethods()).isNotNull();
 	}
 
 	private RegisteredBean registerBean(RootBeanDefinition rootBeanDefinition) {
@@ -158,27 +157,21 @@ class BeanRegistrationsAotContributionTests {
 	}
 
 	@SuppressWarnings({ "unchecked", "cast" })
-	private void testCompiledResult(
+	private void compile(
 			BiConsumer<Consumer<DefaultListableBeanFactory>, Compiled> result) {
+		MethodReference methodReference = this.beanFactoryInitializationCode
+				.getInitializers().get(0);
+		this.beanFactoryInitializationCode.getTypeBuilder().set(type -> {
+			type.addModifiers(Modifier.PUBLIC);
+			type.addSuperinterface(ParameterizedTypeName.get(Consumer.class, DefaultListableBeanFactory.class));
+			type.addMethod(MethodSpec.methodBuilder("accept").addModifiers(Modifier.PUBLIC)
+					.addParameter(DefaultListableBeanFactory.class, "beanFactory")
+					.addStatement(methodReference.toInvokeCodeBlock(CodeBlock.of("beanFactory")))
+					.build());
+		});
 		this.generationContext.writeGeneratedContent();
-		JavaFile javaFile = createJavaFile();
-		TestCompiler.forSystem().withFiles(this.generatedFiles).compile(javaFile::writeTo,
-				compiled -> result.accept(compiled.getInstance(Consumer.class),
-						compiled));
-	}
-
-	private JavaFile createJavaFile() {
-		MethodReference initializer = this.beanFactoryInitializationCode.getInitializers()
-				.get(0);
-		TypeSpec.Builder builder = TypeSpec.classBuilder("BeanFactoryConsumer");
-		builder.addModifiers(Modifier.PUBLIC);
-		builder.addSuperinterface(ParameterizedTypeName.get(Consumer.class,
-				DefaultListableBeanFactory.class));
-		builder.addMethod(MethodSpec.methodBuilder("accept").addModifiers(Modifier.PUBLIC)
-				.addParameter(DefaultListableBeanFactory.class, "beanFactory")
-				.addStatement(initializer.toInvokeCodeBlock(CodeBlock.of("beanFactory")))
-				.build());
-		return JavaFile.builder("__", builder.build()).build();
+		TestCompiler.forSystem().withFiles(this.generatedFiles).printFiles(System.out).compile(compiled ->
+				result.accept(compiled.getInstance(Consumer.class), compiled));
 	}
 
 }

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

@@ -23,11 +23,10 @@ import java.util.function.Supplier;
 import javax.lang.model.element.Modifier;
 
 import org.assertj.core.api.ThrowingConsumer;
-import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.aot.generate.DefaultGenerationContext;
-import org.springframework.aot.generate.GeneratedMethods;
+import org.springframework.aot.generate.GeneratedClass;
 import org.springframework.aot.generate.InMemoryGeneratedFiles;
 import org.springframework.aot.hint.ExecutableHint;
 import org.springframework.aot.hint.ExecutableMode;
@@ -43,6 +42,7 @@ import org.springframework.beans.factory.support.RegisteredBean;
 import org.springframework.beans.factory.support.RootBeanDefinition;
 import org.springframework.beans.testfixture.beans.TestBean;
 import org.springframework.beans.testfixture.beans.TestBeanWithPrivateConstructor;
+import org.springframework.beans.testfixture.beans.factory.aot.DeferredTypeBuilder;
 import org.springframework.beans.testfixture.beans.factory.generator.InnerComponentConfiguration;
 import org.springframework.beans.testfixture.beans.factory.generator.InnerComponentConfiguration.EnvironmentAwareComponent;
 import org.springframework.beans.testfixture.beans.factory.generator.InnerComponentConfiguration.NoDependencyComponent;
@@ -53,12 +53,9 @@ import org.springframework.beans.testfixture.beans.factory.generator.factory.Sam
 import org.springframework.beans.testfixture.beans.factory.generator.injection.InjectionComponent;
 import org.springframework.core.env.StandardEnvironment;
 import org.springframework.core.testfixture.aot.generate.TestGenerationContext;
-import org.springframework.javapoet.ClassName;
 import org.springframework.javapoet.CodeBlock;
-import org.springframework.javapoet.JavaFile;
 import org.springframework.javapoet.MethodSpec;
 import org.springframework.javapoet.ParameterizedTypeName;
-import org.springframework.javapoet.TypeSpec;
 import org.springframework.util.ReflectionUtils;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -71,17 +68,14 @@ import static org.assertj.core.api.Assertions.assertThat;
  */
 class InstanceSupplierCodeGeneratorTests {
 
-	private InMemoryGeneratedFiles generatedFiles;
+	private final InMemoryGeneratedFiles generatedFiles;
 
-	private DefaultGenerationContext generationContext;
+	private final DefaultGenerationContext generationContext;
 
 	private boolean allowDirectSupplierShortcut = false;
 
-	private ClassName className = ClassName.get("__", "InstanceSupplierSupplier");
 
-
-	@BeforeEach
-	void setup() {
+	InstanceSupplierCodeGeneratorTests() {
 		this.generatedFiles = new InMemoryGeneratedFiles();
 		this.generationContext = new TestGenerationContext(this.generatedFiles);
 	}
@@ -91,7 +85,7 @@ class InstanceSupplierCodeGeneratorTests {
 	void generateWhenHasDefaultConstructor() {
 		BeanDefinition beanDefinition = new RootBeanDefinition(TestBean.class);
 		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
-		testCompiledResult(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
+		compile(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
 			TestBean bean = getBean(beanFactory, beanDefinition, instanceSupplier);
 			assertThat(bean).isInstanceOf(TestBean.class);
 			assertThat(compiled.getSourceFile())
@@ -106,7 +100,7 @@ class InstanceSupplierCodeGeneratorTests {
 		BeanDefinition beanDefinition = new RootBeanDefinition(InjectionComponent.class);
 		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
 		beanFactory.registerSingleton("injected", "injected");
-		testCompiledResult(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
+		compile(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
 			InjectionComponent bean = getBean(beanFactory, beanDefinition,
 					instanceSupplier);
 			assertThat(bean).isInstanceOf(InjectionComponent.class).extracting("bean")
@@ -122,7 +116,7 @@ class InstanceSupplierCodeGeneratorTests {
 				NoDependencyComponent.class);
 		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
 		beanFactory.registerSingleton("configuration", new InnerComponentConfiguration());
-		testCompiledResult(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
+		compile(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
 			NoDependencyComponent bean = getBean(beanFactory, beanDefinition,
 					instanceSupplier);
 			assertThat(bean).isInstanceOf(NoDependencyComponent.class);
@@ -140,7 +134,7 @@ class InstanceSupplierCodeGeneratorTests {
 		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
 		beanFactory.registerSingleton("configuration", new InnerComponentConfiguration());
 		beanFactory.registerSingleton("environment", new StandardEnvironment());
-		testCompiledResult(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
+		compile(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
 			EnvironmentAwareComponent bean = getBean(beanFactory, beanDefinition,
 					instanceSupplier);
 			assertThat(bean).isInstanceOf(EnvironmentAwareComponent.class);
@@ -157,7 +151,7 @@ class InstanceSupplierCodeGeneratorTests {
 				NumberHolderFactoryBean.class);
 		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
 		beanFactory.registerSingleton("number", 123);
-		testCompiledResult(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
+		compile(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
 			NumberHolder<?> bean = getBean(beanFactory, beanDefinition, instanceSupplier);
 			assertThat(bean).isInstanceOf(NumberHolder.class);
 			assertThat(bean).extracting("number").isNull(); // No property
@@ -173,7 +167,7 @@ class InstanceSupplierCodeGeneratorTests {
 		BeanDefinition beanDefinition = new RootBeanDefinition(
 				TestBeanWithPrivateConstructor.class);
 		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
-		testCompiledResult(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
+		compile(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
 			TestBeanWithPrivateConstructor bean = getBean(beanFactory, beanDefinition,
 					instanceSupplier);
 			assertThat(bean).isInstanceOf(TestBeanWithPrivateConstructor.class);
@@ -192,7 +186,7 @@ class InstanceSupplierCodeGeneratorTests {
 		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
 		beanFactory.registerBeanDefinition("config", BeanDefinitionBuilder
 				.genericBeanDefinition(SimpleConfiguration.class).getBeanDefinition());
-		testCompiledResult(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
+		compile(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
 			String bean = getBean(beanFactory, beanDefinition, instanceSupplier);
 			assertThat(bean).isInstanceOf(String.class);
 			assertThat(bean).isEqualTo("Hello");
@@ -212,7 +206,7 @@ class InstanceSupplierCodeGeneratorTests {
 		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
 		beanFactory.registerBeanDefinition("config", BeanDefinitionBuilder
 				.genericBeanDefinition(SimpleConfiguration.class).getBeanDefinition());
-		testCompiledResult(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
+		compile(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
 			String bean = getBean(beanFactory, beanDefinition, instanceSupplier);
 			assertThat(bean).isInstanceOf(String.class);
 			assertThat(bean).isEqualTo("Hello");
@@ -231,7 +225,7 @@ class InstanceSupplierCodeGeneratorTests {
 		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
 		beanFactory.registerBeanDefinition("config", BeanDefinitionBuilder
 				.genericBeanDefinition(SimpleConfiguration.class).getBeanDefinition());
-		testCompiledResult(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
+		compile(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
 			Integer bean = getBean(beanFactory, beanDefinition, instanceSupplier);
 			assertThat(bean).isInstanceOf(Integer.class);
 			assertThat(bean).isEqualTo(42);
@@ -254,7 +248,7 @@ class InstanceSupplierCodeGeneratorTests {
 				.genericBeanDefinition(SampleFactory.class).getBeanDefinition());
 		beanFactory.registerSingleton("number", 42);
 		beanFactory.registerSingleton("string", "test");
-		testCompiledResult(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
+		compile(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
 			String bean = getBean(beanFactory, beanDefinition, instanceSupplier);
 			assertThat(bean).isInstanceOf(String.class);
 			assertThat(bean).isEqualTo("42test");
@@ -273,7 +267,7 @@ class InstanceSupplierCodeGeneratorTests {
 		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
 		beanFactory.registerBeanDefinition("config", BeanDefinitionBuilder
 				.genericBeanDefinition(SimpleConfiguration.class).getBeanDefinition());
-		testCompiledResult(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
+		compile(beanFactory, beanDefinition, (instanceSupplier, compiled) -> {
 			Integer bean = getBean(beanFactory, beanDefinition, instanceSupplier);
 			assertThat(bean).isInstanceOf(Integer.class);
 			assertThat(bean).isEqualTo(42);
@@ -308,41 +302,30 @@ class InstanceSupplierCodeGeneratorTests {
 	}
 
 	@SuppressWarnings("unchecked")
-	private void testCompiledResult(DefaultListableBeanFactory beanFactory,
+	private void compile(DefaultListableBeanFactory beanFactory,
 			BeanDefinition beanDefinition,
 			BiConsumer<InstanceSupplier<?>, Compiled> result) {
-		this.generationContext.writeGeneratedContent();
-		DefaultListableBeanFactory registrationBeanFactory = new DefaultListableBeanFactory(
-				beanFactory);
-		registrationBeanFactory.registerBeanDefinition("testBean", beanDefinition);
-		RegisteredBean registeredBean = RegisteredBean.of(registrationBeanFactory,
-				"testBean");
-		GeneratedMethods generatedMethods = new GeneratedMethods();
+		DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory(beanFactory);
+		freshBeanFactory.registerBeanDefinition("testBean", beanDefinition);
+		RegisteredBean registeredBean = RegisteredBean.of(freshBeanFactory, "testBean");
+		DeferredTypeBuilder typeBuilder = new DeferredTypeBuilder();
+		GeneratedClass generateClass = this.generationContext.getGeneratedClasses().addForFeature("TestCode", typeBuilder);
 		InstanceSupplierCodeGenerator generator = new InstanceSupplierCodeGenerator(
-				this.generationContext, this.className, generatedMethods,
-				this.allowDirectSupplierShortcut);
-		Executable constructorOrFactoryMethod = ConstructorOrFactoryMethodResolver
-				.resolve(registeredBean);
-		CodeBlock generatedCode = generator.generateCode(registeredBean,
-				constructorOrFactoryMethod);
-		JavaFile javaFile = createJavaFile(generatedCode, generatedMethods);
-		TestCompiler.forSystem().withFiles(this.generatedFiles).compile(javaFile::writeTo,
-				compiled -> result.accept(
-						(InstanceSupplier<?>) compiled.getInstance(Supplier.class).get(),
-						compiled));
-	}
-
-	private JavaFile createJavaFile(CodeBlock generatedCode,
-			GeneratedMethods generatedMethods) {
-		TypeSpec.Builder builder = TypeSpec.classBuilder("InstanceSupplierSupplier");
-		builder.addModifiers(Modifier.PUBLIC);
-		builder.addSuperinterface(
-				ParameterizedTypeName.get(Supplier.class, InstanceSupplier.class));
-		builder.addMethod(MethodSpec.methodBuilder("get").addModifiers(Modifier.PUBLIC)
-				.returns(InstanceSupplier.class).addStatement("return $L", generatedCode)
-				.build());
-		generatedMethods.doWithMethodSpecs(builder::addMethod);
-		return JavaFile.builder("__", builder.build()).build();
+				this.generationContext, generateClass.getName(),
+				generateClass.getMethods(), this.allowDirectSupplierShortcut);
+		Executable constructorOrFactoryMethod = ConstructorOrFactoryMethodResolver.resolve(registeredBean);
+		CodeBlock generatedCode = generator.generateCode(registeredBean, constructorOrFactoryMethod);
+		typeBuilder.set(type -> {
+			type.addModifiers(Modifier.PUBLIC);
+			type.addSuperinterface(ParameterizedTypeName.get(Supplier.class, InstanceSupplier.class));
+			type.addMethod(MethodSpec.methodBuilder("get")
+					.addModifiers(Modifier.PUBLIC)
+					.returns(InstanceSupplier.class)
+					.addStatement("return $L", generatedCode).build());
+		});
+		this.generationContext.writeGeneratedContent();
+		TestCompiler.forSystem().withFiles(this.generatedFiles).compile(compiled ->
+			result.accept((InstanceSupplier<?>) compiled.getInstance(Supplier.class).get(), compiled));
 	}
 
 }