Introduce TestGenerationContext

This commit polishes DefaultGenerationContext to make the method
that flushes generated classes more explicit. It now throws an
IOException and TestGenerationContext has been updated to handle
that to ease its use in code that can't throw such an exception.

As this use case is likely to happen outside the Spring Framework,
this commit adds such a convenience to spring-test as well.

Closes gh-28877

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

@@ -25,8 +25,6 @@ import javax.lang.model.element.Modifier;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.aop.framework.AopInfrastructureBean;
-import org.springframework.aot.generate.DefaultGenerationContext;
-import org.springframework.aot.generate.InMemoryGeneratedFiles;
 import org.springframework.aot.generate.MethodReference;
 import org.springframework.aot.test.generator.compile.Compiled;
 import org.springframework.aot.test.generator.compile.TestCompiler;
@@ -64,9 +62,7 @@ class ScopedProxyBeanRegistrationAotProcessorTests {
 
 	private final TestBeanRegistrationsAotProcessor processor;
 
-	private final InMemoryGeneratedFiles generatedFiles;
+	private final TestGenerationContext generationContext;
-
-	private final DefaultGenerationContext generationContext;
 
 	private final MockBeanFactoryInitializationCode beanFactoryInitializationCode;
 
@@ -74,8 +70,7 @@ class ScopedProxyBeanRegistrationAotProcessorTests {
 	ScopedProxyBeanRegistrationAotProcessorTests() {
 		this.beanFactory = new DefaultListableBeanFactory();
 		this.processor = new TestBeanRegistrationsAotProcessor();
-		this.generatedFiles = new InMemoryGeneratedFiles();
+		this.generationContext = new TestGenerationContext();
-		this.generationContext = new TestGenerationContext(this.generatedFiles);
 		this.beanFactoryInitializationCode = new MockBeanFactoryInitializationCode(this.generationContext);
 	}
 
@@ -152,7 +147,7 @@ class ScopedProxyBeanRegistrationAotProcessorTests {
 					.build());
 		});
 		this.generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(this.generatedFiles).compile(compiled -> {
+		TestCompiler.forSystem().withFiles(this.generationContext.getGeneratedFiles()).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

@@ -23,8 +23,6 @@ import javax.lang.model.element.Modifier;
 
 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.predicate.RuntimeHintsPredicates;
 import org.springframework.aot.test.generator.compile.CompileWithTargetClassAccess;
@@ -53,9 +51,7 @@ import static org.assertj.core.api.Assertions.assertThat;
  */
 class AutowiredAnnotationBeanRegistrationAotContributionTests {
 
-	private final InMemoryGeneratedFiles generatedFiles;
+	private final TestGenerationContext generationContext;
-
-	private final DefaultGenerationContext generationContext;
 
 	private final MockBeanRegistrationCode beanRegistrationCode;
 
@@ -63,8 +59,7 @@ class AutowiredAnnotationBeanRegistrationAotContributionTests {
 
 
 	AutowiredAnnotationBeanRegistrationAotContributionTests() {
-		this.generatedFiles = new InMemoryGeneratedFiles();
+		this.generationContext = new TestGenerationContext();
-		this.generationContext = new TestGenerationContext(this.generatedFiles);
 		this.beanRegistrationCode = new MockBeanRegistrationCode(this.generationContext);
 		this.beanFactory = new DefaultListableBeanFactory();
 	}
@@ -177,7 +172,7 @@ class AutowiredAnnotationBeanRegistrationAotContributionTests {
 
 		});
 		this.generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(this.generatedFiles).compile(compiled ->
+		TestCompiler.forSystem().withFiles(this.generationContext.getGeneratedFiles()).compile(compiled ->
 				result.accept(compiled.getInstance(BiFunction.class), compiled));
 	}
 

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

@@ -27,10 +27,8 @@ import javax.lang.model.element.Modifier;
 
 import org.junit.jupiter.api.Test;
 
-import org.springframework.aot.generate.DefaultGenerationContext;
 import org.springframework.aot.generate.GeneratedMethod;
 import org.springframework.aot.generate.GenerationContext;
-import org.springframework.aot.generate.InMemoryGeneratedFiles;
 import org.springframework.aot.generate.MethodReference;
 import org.springframework.aot.test.generator.compile.CompileWithTargetClassAccess;
 import org.springframework.aot.test.generator.compile.Compiled;
@@ -64,9 +62,7 @@ import static org.assertj.core.api.Assertions.assertThat;
  */
 class BeanDefinitionMethodGeneratorTests {
 
-	private final InMemoryGeneratedFiles generatedFiles;
+	private final TestGenerationContext generationContext;
-
-	private final DefaultGenerationContext generationContext;
 
 	private final DefaultListableBeanFactory beanFactory;
 
@@ -76,8 +72,7 @@ class BeanDefinitionMethodGeneratorTests {
 
 
 	BeanDefinitionMethodGeneratorTests() {
-		this.generatedFiles = new InMemoryGeneratedFiles();
+		this.generationContext = new TestGenerationContext();
-		this.generationContext = new TestGenerationContext(this.generatedFiles);
 		this.beanFactory = new DefaultListableBeanFactory();
 		this.methodGeneratorFactory = new BeanDefinitionMethodGeneratorFactory(
 				new AotFactoriesLoader(this.beanFactory, new MockSpringFactoriesLoader()));
@@ -412,7 +407,7 @@ class BeanDefinitionMethodGeneratorTests {
 					.addCode("return $L;", method.toInvokeCodeBlock()).build());
 		});
 		this.generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(this.generatedFiles).printFiles(System.out).compile(compiled ->
+		TestCompiler.forSystem().withFiles(this.generationContext.getGeneratedFiles()).compile(compiled ->
 				result.accept((RootBeanDefinition) compiled.getInstance(Supplier.class).get(), compiled));
 	}
 

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

@@ -27,9 +27,7 @@ import javax.lang.model.element.Modifier;
 
 import org.junit.jupiter.api.Test;
 
-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;
@@ -63,9 +61,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 
 	private final RootBeanDefinition beanDefinition = new RootBeanDefinition();
 
-	private final InMemoryGeneratedFiles generatedFiles = new InMemoryGeneratedFiles();
+	private final TestGenerationContext generationContext = new TestGenerationContext();
-
-	private final DefaultGenerationContext generationContext = new TestGenerationContext(this.generatedFiles);
 
 	@Test
 	void setPrimaryWhenFalse() {
@@ -434,7 +430,7 @@ class BeanDefinitionPropertiesCodeGeneratorTests {
 					.addStatement("return beanDefinition").build());
 		});
 		this.generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(this.generatedFiles).compile(compiled -> {
+		TestCompiler.forSystem().withFiles(this.generationContext.getGeneratedFiles()).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

@@ -33,9 +33,7 @@ import javax.lang.model.element.Modifier;
 import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
 
-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.beans.factory.config.BeanReference;
@@ -64,8 +62,7 @@ import static org.assertj.core.api.Assertions.assertThat;
 class BeanDefinitionPropertyValueCodeGeneratorTests {
 
 	private void compile(Object value, BiConsumer<Object, Compiled> result) {
-		InMemoryGeneratedFiles generatedFiles = new InMemoryGeneratedFiles();
+		TestGenerationContext generationContext = new TestGenerationContext();
-		DefaultGenerationContext generationContext = new TestGenerationContext(generatedFiles);
 		DeferredTypeBuilder typeBuilder = new DeferredTypeBuilder();
 		GeneratedClass generatedClass = generationContext.getGeneratedClasses().addForFeature("TestCode", typeBuilder);
 		CodeBlock generatedCode = new BeanDefinitionPropertyValueCodeGenerator(
@@ -78,7 +75,7 @@ class BeanDefinitionPropertyValueCodeGeneratorTests {
 					.returns(Object.class).addStatement("return $L", generatedCode).build());
 		});
 		generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(generatedFiles).compile(compiled ->
+		TestCompiler.forSystem().withFiles(generationContext.getGeneratedFiles()).compile(compiled ->
 				result.accept(compiled.getInstance(Supplier.class).get(), compiled));
 	}
 
@@ -468,7 +465,7 @@ class BeanDefinitionPropertyValueCodeGeneratorTests {
 		@Test
 		void generatedWhenBeanNameReference() {
 			RuntimeBeanNameReference beanReference = new RuntimeBeanNameReference("test");
-			compile(beanReference, (instance, compiler) ->  {
+			compile(beanReference, (instance, compiler) -> {
 				RuntimeBeanReference actual = (RuntimeBeanReference) instance;
 				assertThat(actual.getBeanName()).isEqualTo(beanReference.getBeanName());
 			});
@@ -477,7 +474,7 @@ class BeanDefinitionPropertyValueCodeGeneratorTests {
 		@Test
 		void generatedWhenBeanReferenceByName() {
 			RuntimeBeanReference beanReference = new RuntimeBeanReference("test");
-			compile(beanReference, (instance, compiler) ->  {
+			compile(beanReference, (instance, compiler) -> {
 				RuntimeBeanReference actual = (RuntimeBeanReference) instance;
 				assertThat(actual.getBeanName()).isEqualTo(beanReference.getBeanName());
 				assertThat(actual.getBeanType()).isEqualTo(beanReference.getBeanType());

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

@@ -29,9 +29,7 @@ import javax.lang.model.element.Modifier;
 import org.junit.jupiter.api.Test;
 
 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.aot.generate.MethodReference;
 import org.springframework.aot.test.generator.compile.Compiled;
 import org.springframework.aot.test.generator.compile.TestCompiler;
@@ -61,9 +59,7 @@ class BeanRegistrationsAotContributionTests {
 
 	private DefaultListableBeanFactory beanFactory;
 
-	private final InMemoryGeneratedFiles generatedFiles;
+	private TestGenerationContext generationContext;
-
-	private DefaultGenerationContext generationContext;
 
 	private final BeanDefinitionMethodGeneratorFactory methodGeneratorFactory;
 
@@ -73,8 +69,7 @@ class BeanRegistrationsAotContributionTests {
 	BeanRegistrationsAotContributionTests() {
 		this.springFactoriesLoader = new MockSpringFactoriesLoader();
 		this.beanFactory = new DefaultListableBeanFactory();
-		this.generatedFiles = new InMemoryGeneratedFiles();
+		this.generationContext = new TestGenerationContext();
-		this.generationContext = new TestGenerationContext(this.generatedFiles);
 		this.methodGeneratorFactory = new BeanDefinitionMethodGeneratorFactory(
 				new AotFactoriesLoader(this.beanFactory, this.springFactoriesLoader));
 		this.beanFactoryInitializationCode = new MockBeanFactoryInitializationCode(this.generationContext);
@@ -102,8 +97,8 @@ class BeanRegistrationsAotContributionTests {
 
 	@Test
 	void applyToWhenHasNameGeneratesPrefixedFeatureName() {
-		this.generationContext = new DefaultGenerationContext(
+		this.generationContext = new TestGenerationContext(
-				new ClassNameGenerator(TestTarget.class, "Management"), this.generatedFiles);
+				new ClassNameGenerator(TestTarget.class, "Management"));
 		this.beanFactoryInitializationCode = new MockBeanFactoryInitializationCode(this.generationContext);
 		Map<String, BeanDefinitionMethodGenerator> registrations = new LinkedHashMap<>();
 		RegisteredBean registeredBean = registerBean(
@@ -170,7 +165,7 @@ class BeanRegistrationsAotContributionTests {
 					.build());
 		});
 		this.generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(this.generatedFiles).printFiles(System.out).compile(compiled ->
+		TestCompiler.forSystem().withFiles(this.generationContext.getGeneratedFiles()).compile(compiled ->
 				result.accept(compiled.getInstance(Consumer.class), compiled));
 	}
 

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

@@ -25,9 +25,7 @@ import javax.lang.model.element.Modifier;
 import org.assertj.core.api.ThrowingConsumer;
 import org.junit.jupiter.api.Test;
 
-import org.springframework.aot.generate.DefaultGenerationContext;
 import org.springframework.aot.generate.GeneratedClass;
-import org.springframework.aot.generate.InMemoryGeneratedFiles;
 import org.springframework.aot.hint.ExecutableHint;
 import org.springframework.aot.hint.ExecutableMode;
 import org.springframework.aot.hint.ReflectionHints;
@@ -68,14 +66,11 @@ import static org.assertj.core.api.Assertions.assertThat;
  */
 class InstanceSupplierCodeGeneratorTests {
 
-	private final InMemoryGeneratedFiles generatedFiles;
+	private final TestGenerationContext generationContext;
-
-	private final DefaultGenerationContext generationContext;
 
 
 	InstanceSupplierCodeGeneratorTests() {
-		this.generatedFiles = new InMemoryGeneratedFiles();
+		this.generationContext = new TestGenerationContext();
-		this.generationContext = new TestGenerationContext(this.generatedFiles);
 	}
 
 
@@ -323,8 +318,8 @@ class InstanceSupplierCodeGeneratorTests {
 					.addStatement("return $L", generatedCode).build());
 		});
 		this.generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(this.generatedFiles).printFiles(System.out).compile(compiled ->
+		TestCompiler.forSystem().withFiles(this.generationContext.getGeneratedFiles()).compile(compiled ->
-			result.accept((InstanceSupplier<?>) compiled.getInstance(Supplier.class).get(), compiled));
+				result.accept((InstanceSupplier<?>) compiled.getInstance(Supplier.class).get(), compiled));
 	}
 
 }

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

@@ -24,8 +24,6 @@ import javax.lang.model.element.Modifier;
 import org.assertj.core.api.InstanceOfAssertFactories;
 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.ResourcePatternHint;
 import org.springframework.aot.test.generator.compile.Compiled;
@@ -54,16 +52,13 @@ import static org.assertj.core.api.Assertions.entry;
  */
 class ConfigurationClassPostProcessorAotContributionTests {
 
-	private final InMemoryGeneratedFiles generatedFiles;
+	private final TestGenerationContext generationContext;
-
-	private final DefaultGenerationContext generationContext;
 
 	private final MockBeanFactoryInitializationCode beanFactoryInitializationCode;
 
 
 	ConfigurationClassPostProcessorAotContributionTests() {
-		this.generatedFiles = new InMemoryGeneratedFiles();
+		this.generationContext = new TestGenerationContext();
-		this.generationContext = new TestGenerationContext(this.generatedFiles);
 		this.beanFactoryInitializationCode = new MockBeanFactoryInitializationCode(this.generationContext);
 	}
 
@@ -123,7 +118,7 @@ class ConfigurationClassPostProcessorAotContributionTests {
 					.build());
 		});
 		this.generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(this.generatedFiles).compile(compiled ->
+		TestCompiler.forSystem().withFiles(this.generationContext.getGeneratedFiles()).compile(compiled ->
 				result.accept(compiled.getInstance(Consumer.class), compiled));
 	}
 

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

@@ -20,8 +20,6 @@ import java.util.function.BiConsumer;
 
 import org.junit.jupiter.api.Test;
 
-import org.springframework.aot.generate.DefaultGenerationContext;
-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.beans.BeansException;
@@ -187,11 +185,10 @@ class ApplicationContextAotGeneratorTests {
 	private void testCompiledResult(GenericApplicationContext applicationContext,
 			BiConsumer<ApplicationContextInitializer<GenericApplicationContext>, Compiled> result) {
 		ApplicationContextAotGenerator generator = new ApplicationContextAotGenerator();
-		InMemoryGeneratedFiles generatedFiles = new InMemoryGeneratedFiles();
+		TestGenerationContext generationContext = new TestGenerationContext();
-		DefaultGenerationContext generationContext = new TestGenerationContext(generatedFiles);
 		generator.generateApplicationContext(applicationContext, generationContext);
 		generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(generatedFiles).compile(compiled ->
+		TestCompiler.forSystem().withFiles(generationContext.getGeneratedFiles()).compile(compiled ->
 				result.accept(compiled.getInstance(ApplicationContextInitializer.class), compiled));
 	}
 

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

@@ -25,7 +25,6 @@ import java.lang.annotation.Target;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.aot.generate.GenerationContext;
-import org.springframework.aot.generate.InMemoryGeneratedFiles;
 import org.springframework.aot.hint.MemberCategory;
 import org.springframework.aot.hint.RuntimeHints;
 import org.springframework.aot.hint.TypeReference;
@@ -54,7 +53,7 @@ class ReflectiveProcessorBeanRegistrationAotProcessorTests {
 
 	private final ReflectiveProcessorBeanRegistrationAotProcessor processor = new ReflectiveProcessorBeanRegistrationAotProcessor();
 
-	private final GenerationContext generationContext = new TestGenerationContext(new InMemoryGeneratedFiles());
+	private final GenerationContext generationContext = new TestGenerationContext();
 
 	@Test
 	void shouldIgnoreNonAnnotatedType() {

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

@@ -26,7 +26,6 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.aot.generate.GenerationContext;
-import org.springframework.aot.generate.InMemoryGeneratedFiles;
 import org.springframework.aot.hint.ResourceBundleHint;
 import org.springframework.aot.hint.RuntimeHints;
 import org.springframework.aot.hint.RuntimeHintsRegistrar;
@@ -56,7 +55,7 @@ class RuntimeHintsBeanFactoryInitializationAotProcessorTests {
 
 	@BeforeEach
 	void setup() {
-		this.generationContext = new TestGenerationContext(new InMemoryGeneratedFiles());
+		this.generationContext = new TestGenerationContext();
 		this.generator = new ApplicationContextAotGenerator();
 	}
 

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

@@ -20,8 +20,6 @@ import java.util.function.BiConsumer;
 
 import org.junit.jupiter.api.Test;
 
-import org.springframework.aot.generate.DefaultGenerationContext;
-import org.springframework.aot.generate.InMemoryGeneratedFiles;
 import org.springframework.aot.hint.MemberCategory;
 import org.springframework.aot.hint.RuntimeHints;
 import org.springframework.aot.hint.RuntimeHintsRegistrar;
@@ -110,11 +108,10 @@ class ApplicationContextAotGeneratorRuntimeHintsTests {
 	@SuppressWarnings({"rawtypes", "unchecked"})
 	private void compile(GenericApplicationContext applicationContext, BiConsumer<RuntimeHints, RuntimeHintsInvocations> initializationResult) {
 		ApplicationContextAotGenerator generator = new ApplicationContextAotGenerator();
-		InMemoryGeneratedFiles generatedFiles = new InMemoryGeneratedFiles();
+		TestGenerationContext generationContext = new TestGenerationContext();
-		DefaultGenerationContext generationContext = new TestGenerationContext(generatedFiles);
 		generator.generateApplicationContext(applicationContext, generationContext);
 		generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(generatedFiles).compile(compiled -> {
+		TestCompiler.forSystem().withFiles(generationContext.getGeneratedFiles()).compile(compiled -> {
 			ApplicationContextInitializer instance = compiled.getInstance(ApplicationContextInitializer.class);
 			GenericApplicationContext freshContext = new GenericApplicationContext();
 			RuntimeHintsInvocations recordedInvocations = RuntimeHintsRecorder.record(() -> {

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

@@ -27,6 +27,10 @@ import org.springframework.util.Assert;
 /**
  * Default {@link GenerationContext} implementation.
  *
+ * <p>Generated classes are flushed out using {@link #writeGeneratedContent()}
+ * and should be called once the generation process using this instance has
+ * completed.
+ *
  * @author Phillip Webb
  * @author Stephane Nicoll
  * @since 6.0
@@ -104,13 +108,8 @@ public class DefaultGenerationContext implements GenerationContext {
 	/**
 	 * Write any generated content out to the generated files.
 	 */
-	public void writeGeneratedContent() {
+	public void writeGeneratedContent() throws IOException {
-		try {
+		this.generatedClasses.writeTo(this.generatedFiles);
-			this.generatedClasses.writeTo(this.generatedFiles);
-		}
-		catch (IOException ex) {
-			throw new IllegalStateException(ex);
-		}
 	}
 
 }

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

@@ -16,6 +16,7 @@
 
 package org.springframework.aot.generate;
 
+import java.io.IOException;
 import java.util.function.Consumer;
 
 import org.junit.jupiter.api.Test;
@@ -118,7 +119,7 @@ class DefaultGenerationContextTests {
 	}
 
 	@Test
-	void withNameKeepsTrackOfAllGeneratedFiles() {
+	void withNameKeepsTrackOfAllGeneratedFiles() throws IOException {
 		DefaultGenerationContext context = new DefaultGenerationContext(
 				new ClassNameGenerator(TestTarget.class), this.generatedFiles);
 		context.getGeneratedClasses().addForFeature("Test", typeSpecCustomizer);
@@ -132,7 +133,7 @@ class DefaultGenerationContextTests {
 	}
 
 	@Test
-	void withNameGeneratesUniqueName() {
+	void withNameGeneratesUniqueName() throws IOException {
 		DefaultGenerationContext context = new DefaultGenerationContext(
 				new ClassNameGenerator(Object.class), this.generatedFiles);
 		context.withName("Test").getGeneratedClasses()

spring-core/src/testFixtures/java/org/springframework/core/testfixture/aot/generate/TestGenerationContext.java

@@ -16,25 +16,43 @@
 
 package org.springframework.core.testfixture.aot.generate;
 
+import java.io.IOException;
+
 import org.springframework.aot.generate.ClassNameGenerator;
 import org.springframework.aot.generate.DefaultGenerationContext;
-import org.springframework.aot.generate.GeneratedFiles;
 import org.springframework.aot.generate.GenerationContext;
 import org.springframework.aot.generate.InMemoryGeneratedFiles;
 
 /**
  * Test {@link GenerationContext} implementation that uses
- * {@link TestTarget} as the main target.
+ * {@link InMemoryGeneratedFiles} and provides a convenient
+ * {@link TestTarget} by default.
  *
  * @author Stephane Nicoll
  */
 public class TestGenerationContext extends DefaultGenerationContext {
 
-	public TestGenerationContext(GeneratedFiles generatedFiles) {
+	public TestGenerationContext(ClassNameGenerator classNameGenerator) {
-		super(new ClassNameGenerator(TestTarget.class), generatedFiles);
+		super(classNameGenerator, new InMemoryGeneratedFiles());
 	}
 
 	public TestGenerationContext() {
-		this(new InMemoryGeneratedFiles());
+		this(new ClassNameGenerator(TestTarget.class));
 	}
+
+	@Override
+	public InMemoryGeneratedFiles getGeneratedFiles() {
+		return (InMemoryGeneratedFiles) super.getGeneratedFiles();
+	}
+
+	@Override
+	public void writeGeneratedContent() {
+		try {
+			super.writeGeneratedContent();
+		}
+		catch (IOException ex) {
+			throw new IllegalStateException(ex);
+		}
+	}
+
 }

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

@@ -32,8 +32,6 @@ import org.assertj.core.api.InstanceOfAssertFactories;
 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.hint.TypeReference;
 import org.springframework.aot.test.generator.compile.CompileWithTargetClassAccess;
 import org.springframework.aot.test.generator.compile.Compiled;
@@ -60,15 +58,12 @@ class PersistenceAnnotationBeanPostProcessorAotContributionTests {
 
 	private DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
 
-	private InMemoryGeneratedFiles generatedFiles;
+	private TestGenerationContext generationContext;
-
-	private DefaultGenerationContext generationContext;
 
 	@BeforeEach
 	void setup() {
 		this.beanFactory = new DefaultListableBeanFactory();
-		this.generatedFiles = new InMemoryGeneratedFiles();
+		this.generationContext = new TestGenerationContext();
-		this.generationContext = new TestGenerationContext(generatedFiles);
 	}
 
 	@Test
@@ -185,7 +180,7 @@ class PersistenceAnnotationBeanPostProcessorAotContributionTests {
 		BeanRegistrationCode beanRegistrationCode = mock(BeanRegistrationCode.class);
 		contribution.applyTo(generationContext, beanRegistrationCode);
 		generationContext.writeGeneratedContent();
-		TestCompiler.forSystem().withFiles(generatedFiles)
+		TestCompiler.forSystem().withFiles(generationContext.getGeneratedFiles())
 				.compile(compiled -> result.accept(new Invoker(compiled), compiled));
 	}
 

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

@@ -0,0 +1,67 @@
+/*
+ * 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.test.aot.generate;
+
+import java.io.IOException;
+
+import org.springframework.aot.generate.ClassNameGenerator;
+import org.springframework.aot.generate.DefaultGenerationContext;
+import org.springframework.aot.generate.GenerationContext;
+import org.springframework.aot.generate.InMemoryGeneratedFiles;
+
+/**
+ * {@link GenerationContext} test implementation that uses
+ * {@link InMemoryGeneratedFiles} by default, with a convenient override of
+ * {@link #writeGeneratedContent()} that does not throw {@link IOException}.
+ *
+ * @author Stephane Nicoll
+ * @since 6.0
+ */
+public class TestGenerationContext extends DefaultGenerationContext {
+
+	/**
+	 * Create an instance using the specified {@link ClassNameGenerator}.
+	 * @param classNameGenerator the class name generator to use.
+	 */
+	public TestGenerationContext(ClassNameGenerator classNameGenerator) {
+		super(classNameGenerator, new InMemoryGeneratedFiles());
+	}
+
+	/**
+	 * Create an instance using the specified {@code target}.
+	 * @param target the default target class to use
+	 */
+	public TestGenerationContext(Class<?> target) {
+		this(new ClassNameGenerator(target));
+	}
+
+	@Override
+	public InMemoryGeneratedFiles getGeneratedFiles() {
+		return (InMemoryGeneratedFiles) super.getGeneratedFiles();
+	}
+
+	@Override
+	public void writeGeneratedContent() {
+		try {
+			super.writeGeneratedContent();
+		}
+		catch (IOException ex) {
+			throw new IllegalStateException(ex);
+		}
+	}
+
+}

spring-test/src/main/java/org/springframework/test/aot/generate/package-info.java

@@ -0,0 +1,9 @@
+/**
+ * Test support for core AOT classes.
+ */
+@NonNullApi
+@NonNullFields
+package org.springframework.test.aot.generate;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;