Support Bean Overrides with AOT and native image

This set of commits introduces AOT and native image support for the new
Bean Override feature in the Spring TestContext Framework -- for
example, for using @⁠TestBean and @⁠MockitoBean in AOT mode on the JVM
as well as in a GraalVM native image.

Note, however, that @⁠MockitoBean has currently only been tested in a
native image when mocking interfaces and using Mockito's
ProxyMockMaker, along with a custom runtime hints.

Closes gh-32933

Author: sbrannen

Diff for: spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverride.java

@@ -22,6 +22,8 @@
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 
+import org.springframework.aot.hint.annotation.Reflective;
+
 /**
  * Mark a composed annotation as eligible for Bean Override processing.
  *
@@ -37,11 +39,13 @@
  * {@link org.springframework.test.context.bean.override.mockito.MockitoSpyBean @MockitoSpyBean}.
  *
  * @author Simon Baslé
+ * @author Sam Brannen
  * @since 6.2
  */
 @Retention(RetentionPolicy.RUNTIME)
 @Target(ElementType.ANNOTATION_TYPE)
 @Documented
+@Reflective(BeanOverrideReflectiveProcessor.class)
 public @interface BeanOverride {
 
 	/**

Diff for: spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideBeanFactoryPostProcessor.java

@@ -35,6 +35,7 @@
 import org.springframework.beans.factory.support.DefaultBeanNameGenerator;
 import org.springframework.beans.factory.support.DefaultListableBeanFactory;
 import org.springframework.beans.factory.support.RootBeanDefinition;
+import org.springframework.context.aot.AbstractAotProcessor;
 import org.springframework.core.Ordered;
 import org.springframework.core.ResolvableType;
 import org.springframework.lang.Nullable;
@@ -105,6 +106,12 @@ private void registerBeanOverride(ConfigurableListableBeanFactory beanFactory, O
 	private void replaceDefinition(ConfigurableListableBeanFactory beanFactory, OverrideMetadata overrideMetadata,
 			boolean enforceExistingDefinition) {
 
+		// NOTE: This method supports 3 distinct scenarios which must be accounted for.
+		//
+		// 1) JVM runtime
+		// 2) AOT processing
+		// 3) AOT runtime
+
 		if (!(beanFactory instanceof BeanDefinitionRegistry registry)) {
 			throw new IllegalStateException("Cannot process bean override with a BeanFactory " +
 					"that doesn't implement BeanDefinitionRegistry: " + beanFactory.getClass().getName());
@@ -147,12 +154,24 @@ else if (enforceExistingDefinition) {
 
 		if (existingBeanDefinition != null) {
 			// Validate the existing bean definition.
+			//
+			// Applies during "JVM runtime", "AOT processing", and "AOT runtime".
 			validateBeanDefinition(beanFactory, beanName);
 		}
+		else if (Boolean.getBoolean(AbstractAotProcessor.AOT_PROCESSING)) {
+			// There was no existing bean definition, but during "AOT processing" we
+			// do not register the "pseudo" bean definition since our AOT support
+			// cannot automatically convert that to a functional bean definition for
+			// use at "AOT runtime". Furthermore, by not registering a bean definition
+			// for a nonexistent bean, we allow the "JVM runtime" and "AOT runtime"
+			// to operate the same in the following else-block.
+		}
 		else {
 			// There was no existing bean definition, so we register the "pseudo" bean
 			// definition to ensure that a suitable bean definition exists for the given
 			// bean name for proper autowiring candidate resolution.
+			//
+			// Applies during "JVM runtime" and "AOT runtime".
 			registry.registerBeanDefinition(beanName, pseudoBeanDefinition);
 		}
 
@@ -163,6 +182,11 @@ else if (enforceExistingDefinition) {
 		// Now we have an instance (the override) that we can register. At this stage, we don't
 		// expect a singleton instance to be present. If for some reason a singleton instance
 		// already exists, the following will throw an exception.
+		//
+		// As a bonus, by manually registering a singleton during "AOT processing", we allow
+		// GenericApplicationContext's preDetermineBeanType() method to transparently register
+		// runtime hints for a proxy generated by the above createOverride() invocation --
+		// for example, when @MockitoBean creates a mock based on a JDK dynamic proxy.
 		beanFactory.registerSingleton(beanName, override);
 	}
 

Diff for: spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideContextCustomizer.java

@@ -17,13 +17,8 @@
 package org.springframework.test.context.bean.override;
 
 import java.util.Set;
-import java.util.function.Consumer;
 
-import org.springframework.beans.factory.config.BeanDefinition;
-import org.springframework.beans.factory.config.ConstructorArgumentValues;
-import org.springframework.beans.factory.config.RuntimeBeanReference;
-import org.springframework.beans.factory.support.BeanDefinitionRegistry;
-import org.springframework.beans.factory.support.RootBeanDefinition;
+import org.springframework.beans.factory.config.ConfigurableBeanFactory;
 import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.test.context.ContextCustomizer;
 import org.springframework.test.context.MergedContextConfiguration;
@@ -34,6 +29,7 @@
  *
  * @author Simon Baslé
  * @author Stephane Nicoll
+ * @author Sam Brannen
  * @since 6.2
  */
 class BeanOverrideContextCustomizer implements ContextCustomizer {
@@ -56,43 +52,25 @@ class BeanOverrideContextCustomizer implements ContextCustomizer {
 
 	@Override
 	public void customizeContext(ConfigurableApplicationContext context, MergedContextConfiguration mergedConfig) {
-		if (!(context instanceof BeanDefinitionRegistry registry)) {
-			throw new IllegalStateException("Cannot process bean overrides with an ApplicationContext " +
-					"that doesn't implement BeanDefinitionRegistry: " + context.getClass().getName());
-		}
-		registerInfrastructure(registry);
+		ConfigurableBeanFactory beanFactory = context.getBeanFactory();
+		// Since all three Bean Override infrastructure beans are never injected as
+		// dependencies into other beans within the ApplicationContext, it is sufficient
+		// to register them as manual singleton instances. In addition, registration of
+		// the BeanOverrideBeanFactoryPostProcessor as a singleton is a requirement for
+		// AOT processing, since a bean definition cannot be generated for the
+		// Set<OverrideMetadata> argument that it accepts in its constructor.
+		BeanOverrideRegistrar beanOverrideRegistrar = new BeanOverrideRegistrar(beanFactory);
+		beanFactory.registerSingleton(REGISTRAR_BEAN_NAME, beanOverrideRegistrar);
+		beanFactory.registerSingleton(INFRASTRUCTURE_BEAN_NAME,
+				new BeanOverrideBeanFactoryPostProcessor(this.metadata, beanOverrideRegistrar));
+		beanFactory.registerSingleton(EARLY_INFRASTRUCTURE_BEAN_NAME,
+				new WrapEarlyBeanPostProcessor(beanOverrideRegistrar));
 	}
 
 	Set<OverrideMetadata> getMetadata() {
 		return this.metadata;
 	}
 
-	private void registerInfrastructure(BeanDefinitionRegistry registry) {
-		addInfrastructureBeanDefinition(registry, BeanOverrideRegistrar.class, REGISTRAR_BEAN_NAME,
-				constructorArgs -> {});
-
-		RuntimeBeanReference registrarReference = new RuntimeBeanReference(REGISTRAR_BEAN_NAME);
-		addInfrastructureBeanDefinition(registry, WrapEarlyBeanPostProcessor.class, EARLY_INFRASTRUCTURE_BEAN_NAME,
-				constructorArgs -> constructorArgs.addIndexedArgumentValue(0, registrarReference));
-		addInfrastructureBeanDefinition(registry, BeanOverrideBeanFactoryPostProcessor.class, INFRASTRUCTURE_BEAN_NAME,
-				constructorArgs -> {
-					constructorArgs.addIndexedArgumentValue(0, this.metadata);
-					constructorArgs.addIndexedArgumentValue(1, registrarReference);
-				});
-	}
-
-	private void addInfrastructureBeanDefinition(BeanDefinitionRegistry registry,
-			Class<?> clazz, String beanName, Consumer<ConstructorArgumentValues> constructorArgumentsConsumer) {
-
-		if (!registry.containsBeanDefinition(beanName)) {
-			RootBeanDefinition definition = new RootBeanDefinition(clazz);
-			definition.setRole(BeanDefinition.ROLE_INFRASTRUCTURE);
-			ConstructorArgumentValues constructorArguments = definition.getConstructorArgumentValues();
-			constructorArgumentsConsumer.accept(constructorArguments);
-			registry.registerBeanDefinition(beanName, definition);
-		}
-	}
-
 	@Override
 	public boolean equals(Object other) {
 		if (other == this) {

Diff for: spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideReflectiveProcessor.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2002-2024 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.context.bean.override;
+
+import java.lang.reflect.AnnotatedElement;
+
+import org.springframework.aot.hint.ReflectionHints;
+import org.springframework.aot.hint.annotation.ReflectiveProcessor;
+import org.springframework.core.annotation.MergedAnnotation;
+import org.springframework.core.annotation.MergedAnnotations;
+
+import static org.springframework.aot.hint.MemberCategory.INVOKE_DECLARED_CONSTRUCTORS;
+
+/**
+ * {@link ReflectiveProcessor} that processes {@link BeanOverride @BeanOverride}
+ * annotations.
+ *
+ * @author Sam Brannen
+ * @since 6.2
+ */
+class BeanOverrideReflectiveProcessor implements ReflectiveProcessor {
+
+	@Override
+	public void registerReflectionHints(ReflectionHints hints, AnnotatedElement element) {
+		MergedAnnotations.from(element)
+				.get(BeanOverride.class)
+				.synthesize(MergedAnnotation::isPresent)
+				.map(BeanOverride::value)
+				.ifPresent(clazz -> hints.registerType(clazz, INVOKE_DECLARED_CONSTRUCTORS));
+	}
+
+}

Diff for: spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideRegistrar.java

@@ -22,10 +22,7 @@
 
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanCreationException;
-import org.springframework.beans.factory.BeanFactory;
-import org.springframework.beans.factory.BeanFactoryAware;
 import org.springframework.beans.factory.config.ConfigurableBeanFactory;
-import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.ReflectionUtils;
 import org.springframework.util.StringUtils;
@@ -36,36 +33,31 @@
  * for test execution listeners.
  *
  * @author Simon Baslé
+ * @author Sam Brannen
  * @since 6.2
  */
-class BeanOverrideRegistrar implements BeanFactoryAware {
+class BeanOverrideRegistrar {
 
 	private final Map<OverrideMetadata, String> beanNameRegistry = new HashMap<>();
 
 	private final Map<String, OverrideMetadata> earlyOverrideMetadata = new HashMap<>();
 
-	@Nullable
-	private ConfigurableBeanFactory beanFactory;
+	private final ConfigurableBeanFactory beanFactory;
 
 
-	@Override
-	public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
-		if (!(beanFactory instanceof ConfigurableBeanFactory cbf)) {
-			throw new IllegalStateException("Cannot process bean override with a BeanFactory " +
-					"that doesn't implement ConfigurableBeanFactory: " + beanFactory.getClass().getName());
-		}
-		this.beanFactory = cbf;
+	BeanOverrideRegistrar(ConfigurableBeanFactory beanFactory) {
+		Assert.notNull(beanFactory, "ConfigurableBeanFactory must not be null");
+		this.beanFactory = beanFactory;
 	}
 
 	/**
-	 * Check {@link #markWrapEarly(OverrideMetadata, String) early override}
+	 * Check {@linkplain #markWrapEarly(OverrideMetadata, String) early override}
 	 * records and use the {@link OverrideMetadata} to create an override
-	 * instance from the provided bean, if relevant.
+	 * instance based on the provided bean, if relevant.
 	 */
 	Object wrapIfNecessary(Object bean, String beanName) throws BeansException {
 		OverrideMetadata metadata = this.earlyOverrideMetadata.get(beanName);
 		if (metadata != null && metadata.getStrategy() == BeanOverrideStrategy.WRAP_BEAN) {
-			Assert.state(this.beanFactory != null, "ConfigurableBeanFactory must not be null");
 			bean = metadata.createOverride(beanName, null, bean);
 			metadata.track(bean, this.beanFactory);
 		}
@@ -99,7 +91,6 @@ private void inject(Field field, Object target, String beanName) {
 		try {
 			ReflectionUtils.makeAccessible(field);
 			Object existingValue = ReflectionUtils.getField(field, target);
-			Assert.state(this.beanFactory != null, "ConfigurableBeanFactory must not be null");
 			Object bean = this.beanFactory.getBean(beanName, field.getType());
 			if (existingValue == bean) {
 				return;

Diff for: spring-test/src/main/java/org/springframework/test/context/bean/override/mockito/MockitoAnnotationDetector.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2002-2024 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.context.bean.override.mockito;
+
+import java.lang.annotation.Annotation;
+import java.lang.reflect.AnnotatedElement;
+import java.util.Arrays;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Predicate;
+
+import org.springframework.util.ReflectionUtils;
+
+/**
+ * Utility class that detects {@code org.mockito} annotations as well as the
+ * annotations in this package (like {@link MockitoBeanSettings @MockitoBeanSettings}).
+ *
+ * @author Simon Baslé
+ * @author Sam Brannen
+ */
+abstract class MockitoAnnotationDetector {
+
+	private static final String MOCKITO_BEAN_PACKAGE = MockitoBeanSettings.class.getPackageName();
+
+	private static final String ORG_MOCKITO_PACKAGE = "org.mockito";
+
+	private static final Predicate<Annotation> isMockitoAnnotation = annotation -> {
+			String packageName = annotation.annotationType().getPackageName();
+			return (packageName.startsWith(MOCKITO_BEAN_PACKAGE) ||
+					packageName.startsWith(ORG_MOCKITO_PACKAGE));
+		};
+
+	static boolean hasMockitoAnnotations(Class<?> testClass) {
+		if (isAnnotated(testClass)) {
+			return true;
+		}
+		// TODO Ideally we should short-circuit the search once we've found a Mockito annotation,
+		// since there's no need to continue searching additional fields or further up the class
+		// hierarchy; however, that is not possible with ReflectionUtils#doWithFields. Plus, the
+		// previous invocation of isAnnotated(testClass) only finds annotations declared directly
+		// on the test class. So, we'll likely need a completely different approach that combines
+		// the "test class/interface is annotated?" and "field is annotated?" checks in a single
+		// search algorithm.
+		AtomicBoolean found = new AtomicBoolean();
+		ReflectionUtils.doWithFields(testClass, field -> found.set(true), MockitoAnnotationDetector::isAnnotated);
+		return found.get();
+	}
+
+	private static boolean isAnnotated(AnnotatedElement annotatedElement) {
+		return Arrays.stream(annotatedElement.getAnnotations()).anyMatch(isMockitoAnnotation);
+	}
+
+}

Diff for: spring-test/src/main/java/org/springframework/test/context/bean/override/mockito/MockitoResetTestExecutionListener.java

@@ -29,7 +29,6 @@
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ConfigurableApplicationContext;
-import org.springframework.core.NativeDetector;
 import org.springframework.core.Ordered;
 import org.springframework.lang.Nullable;
 import org.springframework.test.context.TestContext;
@@ -58,14 +57,16 @@ public int getOrder() {
 
 	@Override
 	public void beforeTestMethod(TestContext testContext) throws Exception {
-		if (MockitoTestExecutionListener.mockitoPresent && !NativeDetector.inNativeImage()) {
+		Class<?> testClass = testContext.getTestClass();
+		if (MockitoTestExecutionListener.mockitoPresent && MockitoAnnotationDetector.hasMockitoAnnotations(testClass)) {
 			resetMocks(testContext.getApplicationContext(), MockReset.BEFORE);
 		}
 	}
 
 	@Override
 	public void afterTestMethod(TestContext testContext) throws Exception {
-		if (MockitoTestExecutionListener.mockitoPresent && !NativeDetector.inNativeImage()) {
+		Class<?> testClass = testContext.getTestClass();
+		if (MockitoTestExecutionListener.mockitoPresent && MockitoAnnotationDetector.hasMockitoAnnotations(testClass)) {
 			resetMocks(testContext.getApplicationContext(), MockReset.AFTER);
 		}
 	}