From 9c12833e160aa6d964d1c323c0af8d8847f01e12 Mon Sep 17 00:00:00 2001
From: Phillip Webb <pwebb@vmware.com>
Date: Tue, 26 Apr 2022 21:42:35 -0700
Subject: [PATCH] Add withFiles support to TestCompiler

Update `TestCompiler` so that `withFiles` can be used to quickly
configure the compiler with `InMemoryGeneratedFiles`.

See gh-28414
---
 .../test/generator/compile/TestCompiler.java  | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/spring-core-test/src/main/java/org/springframework/aot/test/generator/compile/TestCompiler.java b/spring-core-test/src/main/java/org/springframework/aot/test/generator/compile/TestCompiler.java
index 063096a96cb..d862a8aa1f2 100644
--- a/spring-core-test/src/main/java/org/springframework/aot/test/generator/compile/TestCompiler.java
+++ b/spring-core-test/src/main/java/org/springframework/aot/test/generator/compile/TestCompiler.java
@@ -17,6 +17,7 @@
 package org.springframework.aot.test.generator.compile;
 
 import java.io.PrintStream;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Locale;
 import java.util.function.Consumer;
@@ -29,6 +30,8 @@ import javax.tools.JavaFileObject;
 import javax.tools.StandardJavaFileManager;
 import javax.tools.ToolProvider;
 
+import org.springframework.aot.generate.GeneratedFiles.Kind;
+import org.springframework.aot.generate.InMemoryGeneratedFiles;
 import org.springframework.aot.test.generator.file.ResourceFile;
 import org.springframework.aot.test.generator.file.ResourceFiles;
 import org.springframework.aot.test.generator.file.SourceFile;
@@ -83,6 +86,23 @@ public final class TestCompiler {
 				ResourceFiles.none());
 	}
 
+	/**
+	 * Return a new {@link TestCompiler} instance with additional generated
+	 * source and resource files.
+	 * @param generatedFiles the generated files to add
+	 * @return a new {@link TestCompiler} instance
+	 */
+	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)));
+		return withSources(sourceFiles).withResources(resourceFiles);
+	}
+
 	/**
 	 * Return a new {@link TestCompiler} instance with addition source files.
 	 * @param sourceFiles the additional source files