Adapt to change in Spring Framework

See spring-projects/spring-framework#34513

See gh-44481

Author: snicoll

spring-boot-project/spring-boot-test/src/main/java/org/springframework/boot/test/context/SpringBootContextLoader.java

@@ -19,19 +19,13 @@
 import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.function.Consumer;
 
-import org.springframework.aot.generate.GenerationContext;
 import org.springframework.aot.hint.ExecutableMode;
-import org.springframework.aot.hint.ReflectionHints;
+import org.springframework.aot.hint.RuntimeHints;
 import org.springframework.beans.BeanUtils;
-import org.springframework.beans.factory.aot.BeanFactoryInitializationAotContribution;
-import org.springframework.beans.factory.aot.BeanFactoryInitializationAotProcessor;
-import org.springframework.beans.factory.aot.BeanFactoryInitializationCode;
-import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.boot.ApplicationContextFactory;
 import org.springframework.boot.Banner;
 import org.springframework.boot.ConfigurableBootstrapContext;
@@ -113,28 +107,33 @@ public class SpringBootContextLoader extends AbstractContextLoader implements Ao
 
 	@Override
 	public ApplicationContext loadContext(MergedContextConfiguration mergedConfig) throws Exception {
-		return loadContext(mergedConfig, Mode.STANDARD, null);
+		return loadContext(mergedConfig, Mode.STANDARD, null, null);
 	}
 
 	@Override
-	public ApplicationContext loadContextForAotProcessing(MergedContextConfiguration mergedConfig) throws Exception {
-		return loadContext(mergedConfig, Mode.AOT_PROCESSING, null);
+	public ApplicationContext loadContextForAotProcessing(MergedContextConfiguration mergedConfig,
+			RuntimeHints runtimeHints) throws Exception {
+		return loadContext(mergedConfig, Mode.AOT_PROCESSING, null, runtimeHints);
 	}
 
 	@Override
 	public ApplicationContext loadContextForAotRuntime(MergedContextConfiguration mergedConfig,
 			ApplicationContextInitializer<ConfigurableApplicationContext> initializer) throws Exception {
-		return loadContext(mergedConfig, Mode.AOT_RUNTIME, initializer);
+		return loadContext(mergedConfig, Mode.AOT_RUNTIME, initializer, null);
 	}
 
 	private ApplicationContext loadContext(MergedContextConfiguration mergedConfig, Mode mode,
-			ApplicationContextInitializer<ConfigurableApplicationContext> initializer) throws Exception {
+			ApplicationContextInitializer<ConfigurableApplicationContext> initializer, RuntimeHints runtimeHints)
+			throws Exception {
 		assertHasClassesOrLocations(mergedConfig);
 		SpringBootTestAnnotation annotation = SpringBootTestAnnotation.get(mergedConfig);
 		String[] args = annotation.getArgs();
 		UseMainMethod useMainMethod = annotation.getUseMainMethod();
 		Method mainMethod = getMainMethod(mergedConfig, useMainMethod);
 		if (mainMethod != null) {
+			if (runtimeHints != null) {
+				runtimeHints.reflection().registerMethod(mainMethod, ExecutableMode.INVOKE);
+			}
 			ContextLoaderHook hook = new ContextLoaderHook(mode, initializer,
 					(application) -> configure(mergedConfig, application));
 			return hook.runMain(() -> ReflectionUtils.invokeMethod(mainMethod, null, new Object[] { args }));
@@ -585,39 +584,4 @@ private ApplicationContext run(ThrowingSupplier<ConfigurableApplicationContext>
 
 	}
 
-	static class MainMethodBeanFactoryInitializationAotProcessor implements BeanFactoryInitializationAotProcessor {
-
-		@Override
-		public BeanFactoryInitializationAotContribution processAheadOfTime(
-				ConfigurableListableBeanFactory beanFactory) {
-			List<Method> mainMethods = new ArrayList<>();
-			for (String beanName : beanFactory.getBeanDefinitionNames()) {
-				Class<?> beanType = beanFactory.getType(beanName);
-				Method mainMethod = findMainMethod(beanType);
-				if (mainMethod != null) {
-					mainMethods.add(mainMethod);
-				}
-			}
-			return !mainMethods.isEmpty() ? new AotContribution(mainMethods) : null;
-		}
-
-		static class AotContribution implements BeanFactoryInitializationAotContribution {
-
-			private final Collection<Method> mainMethods;
-
-			AotContribution(Collection<Method> mainMethods) {
-				this.mainMethods = mainMethods;
-			}
-
-			@Override
-			public void applyTo(GenerationContext generationContext,
-					BeanFactoryInitializationCode beanFactoryInitializationCode) {
-				ReflectionHints reflectionHints = generationContext.getRuntimeHints().reflection();
-				this.mainMethods.forEach((method) -> reflectionHints.registerMethod(method, ExecutableMode.INVOKE));
-			}
-
-		}
-
-	}
-
 }

spring-boot-project/spring-boot-test/src/main/resources/META-INF/spring/aot.factories

@@ -27,14 +27,10 @@
 
 import org.springframework.aot.hint.RuntimeHints;
 import org.springframework.aot.hint.predicate.RuntimeHintsPredicates;
-import org.springframework.aot.test.generate.TestGenerationContext;
 import org.springframework.beans.factory.BeanCreationException;
-import org.springframework.beans.factory.aot.BeanFactoryInitializationAotContribution;
-import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.boot.ApplicationContextFactory;
 import org.springframework.boot.SpringApplication;
 import org.springframework.boot.SpringBootConfiguration;
-import org.springframework.boot.test.context.SpringBootContextLoader.MainMethodBeanFactoryInitializationAotProcessor;
 import org.springframework.boot.test.context.SpringBootTest.UseMainMethod;
 import org.springframework.boot.test.util.TestPropertyValues;
 import org.springframework.boot.web.reactive.context.GenericReactiveWebApplicationContext;
@@ -47,6 +43,7 @@
 import org.springframework.core.env.StandardEnvironment;
 import org.springframework.test.context.ActiveProfiles;
 import org.springframework.test.context.ApplicationContextFailureProcessor;
+import org.springframework.test.context.BootstrapUtils;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.ContextHierarchy;
 import org.springframework.test.context.MergedContextConfiguration;
@@ -59,6 +56,8 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
+import static org.mockito.BDDMockito.then;
+import static org.mockito.Mockito.mock;
 
 /**
  * Tests for {@link SpringBootContextLoader}
@@ -255,30 +254,24 @@ void whenUseMainMethodWithContextHierarchyThrowsException() {
 	}
 
 	@Test
-	void whenMainMethodPresentRegisterReflectionHints() {
-		TestContext testContext = new ExposedTestContextManager(UseMainMethodWhenAvailableAndNoMainMethod.class)
-			.getExposedTestContext();
-		ConfigurableListableBeanFactory beanFactory = (ConfigurableListableBeanFactory) testContext
-			.getApplicationContext()
-			.getAutowireCapableBeanFactory();
-		BeanFactoryInitializationAotContribution aotContribution = new MainMethodBeanFactoryInitializationAotProcessor()
-			.processAheadOfTime(beanFactory);
-		assertThat(aotContribution).isNull();
+	void whenMainMethodNotAvailableReturnsNoAotContribution() throws Exception {
+		SpringBootContextLoader contextLoader = new SpringBootContextLoader();
+		MergedContextConfiguration contextConfiguration = BootstrapUtils
+			.resolveTestContextBootstrapper(UseMainMethodWhenAvailableAndNoMainMethod.class)
+			.buildMergedContextConfiguration();
+		RuntimeHints runtimeHints = mock(RuntimeHints.class);
+		contextLoader.loadContextForAotProcessing(contextConfiguration, runtimeHints);
+		then(runtimeHints).shouldHaveNoInteractions();
 	}
 
 	@Test
-	void whenMainMethodNotAvailableReturnsNoAotContribution() {
-		TestContext testContext = new ExposedTestContextManager(UseMainMethodWhenAvailableAndMainMethod.class)
-			.getExposedTestContext();
-		ConfigurableListableBeanFactory beanFactory = (ConfigurableListableBeanFactory) testContext
-			.getApplicationContext()
-			.getAutowireCapableBeanFactory();
-		BeanFactoryInitializationAotContribution aotContribution = new MainMethodBeanFactoryInitializationAotProcessor()
-			.processAheadOfTime(beanFactory);
-		assertThat(aotContribution).isNotNull();
-		TestGenerationContext generationContext = new TestGenerationContext();
-		aotContribution.applyTo(generationContext, null);
-		RuntimeHints runtimeHints = generationContext.getRuntimeHints();
+	void whenMainMethodPresentRegisterReflectionHints() throws Exception {
+		SpringBootContextLoader contextLoader = new SpringBootContextLoader();
+		MergedContextConfiguration contextConfiguration = BootstrapUtils
+			.resolveTestContextBootstrapper(UseMainMethodWhenAvailableAndMainMethod.class)
+			.buildMergedContextConfiguration();
+		RuntimeHints runtimeHints = new RuntimeHints();
+		contextLoader.loadContextForAotProcessing(contextConfiguration, runtimeHints);
 		assertThat(RuntimeHintsPredicates.reflection().onMethod(ConfigWithMain.class, "main").invoke())
 			.accepts(runtimeHints);
 	}