diff --git a/spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideContextCustomizer.java b/spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideContextCustomizer.java index 07df3d4f08..3753f348ea 100644 --- a/spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideContextCustomizer.java +++ b/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 @@ import org.springframework.test.context.MergedContextConfiguration; * * @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 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 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 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 --git a/spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideRegistrar.java b/spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideRegistrar.java index 3d49207ef0..d9f749a4b4 100644 --- a/spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideRegistrar.java +++ b/spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideRegistrar.java @@ -22,10 +22,7 @@ import java.util.Map; 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 @@ import org.springframework.util.StringUtils; * for test execution listeners. * * @author Simon Baslé + * @author Sam Brannen * @since 6.2 */ -class BeanOverrideRegistrar implements BeanFactoryAware { +class BeanOverrideRegistrar { private final Map beanNameRegistry = new HashMap<>(); private final Map 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 @@ class BeanOverrideRegistrar implements BeanFactoryAware { 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;