From e73eb458b4cd081ddfb92031abdde65a87fc06fe Mon Sep 17 00:00:00 2001
From: Cedomir Igaly <cedomir.igaly@e-dama.net>
Date: Thu, 16 May 2024 10:12:42 +0200
Subject: [PATCH] HHH-18106 - Handling source element in @WithClasses
 annotation

---
 .../processor/test/util/CompilationRunner.java |  6 +++++-
 .../test/util/CompilationStatement.java        | 18 +++++++++++++++---
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/tooling/metamodel-generator/src/test/java/org/hibernate/processor/test/util/CompilationRunner.java b/tooling/metamodel-generator/src/test/java/org/hibernate/processor/test/util/CompilationRunner.java
index 4ba8dfee96..69b271d3b6 100644
--- a/tooling/metamodel-generator/src/test/java/org/hibernate/processor/test/util/CompilationRunner.java
+++ b/tooling/metamodel-generator/src/test/java/org/hibernate/processor/test/util/CompilationRunner.java
@@ -26,6 +26,7 @@ import org.junit.runners.model.Statement;
 public class CompilationRunner extends BlockJUnit4ClassRunner {
 	private final List<Class<?>> testEntities;
 	private final List<Class<?>> preCompileEntities;
+	private final List<String> sources;
 	private final List<String> mappingFiles;
 	private final Map<String, String> processorOptions;
 	private final String packageName;
@@ -36,6 +37,7 @@ public class CompilationRunner extends BlockJUnit4ClassRunner {
 		super( clazz );
 		this.testEntities = new ArrayList<>();
 		this.preCompileEntities = new ArrayList<>();
+		this.sources = new ArrayList<>();
 		this.mappingFiles = new ArrayList<>();
 		this.processorOptions = new HashMap<>();
 		Package pkg = clazz.getPackage();
@@ -64,6 +66,7 @@ public class CompilationRunner extends BlockJUnit4ClassRunner {
 				getTestClass().getJavaClass(),
 				testEntities,
 				preCompileEntities,
+				sources,
 				mappingFiles,
 				processorOptions,
 				ignoreCompilationErrors
@@ -74,6 +77,7 @@ public class CompilationRunner extends BlockJUnit4ClassRunner {
 		if ( withClasses != null ) {
 			Collections.addAll( testEntities, withClasses.value() );
 			Collections.addAll( preCompileEntities, withClasses.preCompile() );
+			Collections.addAll( sources, withClasses.sources() );
 		}
 	}
 
@@ -116,7 +120,7 @@ public class CompilationRunner extends BlockJUnit4ClassRunner {
 	}
 
 	private boolean annotationProcessorNeedsToRun() {
-		return !testEntities.isEmpty() || !mappingFiles.isEmpty();
+		return !testEntities.isEmpty() || !sources.isEmpty() || !mappingFiles.isEmpty();
 	}
 }
 
diff --git a/tooling/metamodel-generator/src/test/java/org/hibernate/processor/test/util/CompilationStatement.java b/tooling/metamodel-generator/src/test/java/org/hibernate/processor/test/util/CompilationStatement.java
index 64d20245ed..087e56cf3b 100644
--- a/tooling/metamodel-generator/src/test/java/org/hibernate/processor/test/util/CompilationStatement.java
+++ b/tooling/metamodel-generator/src/test/java/org/hibernate/processor/test/util/CompilationStatement.java
@@ -42,6 +42,7 @@ public class CompilationStatement extends Statement {
 	private final Class<?> testClass;
 	private final List<Class<?>> testEntities;
 	private final List<Class<?>> preCompileEntities;
+	private final List<String> sources;
 	private final List<String> xmlMappingFiles;
 	private final Map<String, String> processorOptions;
 	private final boolean ignoreCompilationErrors;
@@ -51,6 +52,7 @@ public class CompilationStatement extends Statement {
 			Class<?> testClass,
 			List<Class<?>> testEntities,
 			List<Class<?>> proCompileEntities,
+			List<String> sources,
 			List<String> xmlMappingFiles,
 			Map<String, String> processorOptions,
 			boolean ignoreCompilationErrors) {
@@ -58,6 +60,7 @@ public class CompilationStatement extends Statement {
 		this.testClass = testClass;
 		this.testEntities = testEntities;
 		this.preCompileEntities = proCompileEntities;
+		this.sources = sources;
 		this.xmlMappingFiles = xmlMappingFiles;
 		this.processorOptions = processorOptions;
 		this.ignoreCompilationErrors = ignoreCompilationErrors;
@@ -68,11 +71,11 @@ public class CompilationStatement extends Statement {
 	public void evaluate() throws Throwable {
 		// some test needs to compile some classes prior to the actual classes under test
 		if ( !preCompileEntities.isEmpty() ) {
-			compile( getCompilationUnits( preCompileEntities ) );
+			compile( getCompilationUnits( preCompileEntities, null ) );
 		}
 
 		// now we compile the actual test classes
-		compile( getCompilationUnits( testEntities ) );
+		compile( getCompilationUnits( testEntities, sources ) );
 
 		if ( !ignoreCompilationErrors ) {
 			TestUtil.assertNoCompilationError( compilationDiagnostics );
@@ -81,12 +84,21 @@ public class CompilationStatement extends Statement {
 		originalStatement.evaluate();
 	}
 
-	private List<File> getCompilationUnits(List<Class<?>> classesToCompile) {
+	private List<File> getCompilationUnits(List<Class<?>> classesToCompile, List<String> sources) {
 		List<File> javaFiles = new ArrayList<File>();
 		for ( Class<?> testClass : classesToCompile ) {
 			String pathToSource = getPathToSource( testClass );
 			javaFiles.add( new File( pathToSource ) );
 		}
+		if ( sources != null ) {
+			final var resourcesBaseDir = TestUtil.getResourcesBaseDir( testClass );
+			for ( String source : sources ) {
+				javaFiles.add(
+						new File( resourcesBaseDir,
+										 source.replace( PACKAGE_SEPARATOR, File.separator ) + ".java" ) );
+			}
+
+		}
 		return javaFiles;
 	}