diff --git a/spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideBeanFactoryPostProcessor.java b/spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideBeanFactoryPostProcessor.java index a257a07abaf..d85d0a35c53 100644 --- a/spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideBeanFactoryPostProcessor.java +++ b/spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideBeanFactoryPostProcessor.java @@ -19,9 +19,7 @@ package org.springframework.test.context.bean.override; import java.lang.reflect.Field; import java.util.Arrays; import java.util.LinkedHashSet; -import java.util.Map; import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; import org.springframework.aop.scope.ScopedProxyUtils; import org.springframework.beans.BeansException; @@ -32,16 +30,13 @@ import org.springframework.beans.factory.config.BeanDefinition; import org.springframework.beans.factory.config.BeanFactoryPostProcessor; import org.springframework.beans.factory.config.ConfigurableListableBeanFactory; import org.springframework.beans.factory.config.DependencyDescriptor; -import org.springframework.beans.factory.config.SmartInstantiationAwareBeanPostProcessor; import org.springframework.beans.factory.support.BeanDefinitionRegistry; import org.springframework.beans.factory.support.BeanNameGenerator; import org.springframework.beans.factory.support.DefaultBeanNameGenerator; import org.springframework.beans.factory.support.RootBeanDefinition; import org.springframework.core.Ordered; -import org.springframework.core.PriorityOrdered; import org.springframework.core.ResolvableType; import org.springframework.util.Assert; -import org.springframework.util.StringUtils; /** * A {@link BeanFactoryPostProcessor} implementation that processes identified @@ -297,47 +292,4 @@ class BeanOverrideBeanFactoryPostProcessor implements BeanFactoryPostProcessor, () -> "Unable to override bean '" + beanName + "': only singleton beans can be overridden."); } - - static class WrapEarlyBeanPostProcessor implements SmartInstantiationAwareBeanPostProcessor, - PriorityOrdered { - - private final Map earlyReferences = new ConcurrentHashMap<>(16); - - private final BeanOverrideRegistrar overrideRegistrar; - - WrapEarlyBeanPostProcessor(BeanOverrideRegistrar registrar) { - this.overrideRegistrar = registrar; - } - - @Override - public int getOrder() { - return Ordered.HIGHEST_PRECEDENCE; - } - - @Override - public Object getEarlyBeanReference(Object bean, String beanName) throws BeansException { - if (bean instanceof FactoryBean) { - return bean; - } - this.earlyReferences.put(getCacheKey(bean, beanName), bean); - return this.overrideRegistrar.wrapIfNecessary(bean, beanName); - } - - @Override - public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException { - if (bean instanceof FactoryBean) { - return bean; - } - if (this.earlyReferences.remove(getCacheKey(bean, beanName)) != bean) { - return this.overrideRegistrar.wrapIfNecessary(bean, beanName); - } - return bean; - } - - private String getCacheKey(Object bean, String beanName) { - return (StringUtils.hasLength(beanName) ? beanName : bean.getClass().getName()); - } - - } - } 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 d74740aa8b8..acf425e4b59 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 @@ -27,7 +27,6 @@ import org.springframework.beans.factory.support.RootBeanDefinition; import org.springframework.context.ConfigurableApplicationContext; import org.springframework.test.context.ContextCustomizer; import org.springframework.test.context.MergedContextConfiguration; -import org.springframework.test.context.bean.override.BeanOverrideBeanFactoryPostProcessor.WrapEarlyBeanPostProcessor; /** * {@link ContextCustomizer} implementation that registers the necessary diff --git a/spring-test/src/main/java/org/springframework/test/context/bean/override/WrapEarlyBeanPostProcessor.java b/spring-test/src/main/java/org/springframework/test/context/bean/override/WrapEarlyBeanPostProcessor.java new file mode 100644 index 00000000000..4c5ebf5fbfe --- /dev/null +++ b/spring-test/src/main/java/org/springframework/test/context/bean/override/WrapEarlyBeanPostProcessor.java @@ -0,0 +1,77 @@ +/* + * 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.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +import org.springframework.beans.BeansException; +import org.springframework.beans.factory.FactoryBean; +import org.springframework.beans.factory.config.SmartInstantiationAwareBeanPostProcessor; +import org.springframework.core.Ordered; +import org.springframework.core.PriorityOrdered; +import org.springframework.util.StringUtils; + +/** + * {@link SmartInstantiationAwareBeanPostProcessor} implementation that wraps + * beans in order to support the {@link BeanOverrideStrategy#WRAP_BEAN WRAP_BEAN} + * bean override strategy. + * + * @author Simon Baslé + * @author Stephane Nicoll + * @since 6.2 + */ +class WrapEarlyBeanPostProcessor implements SmartInstantiationAwareBeanPostProcessor, PriorityOrdered { + + private final Map earlyReferences = new ConcurrentHashMap<>(16); + + private final BeanOverrideRegistrar overrideRegistrar; + + WrapEarlyBeanPostProcessor(BeanOverrideRegistrar registrar) { + this.overrideRegistrar = registrar; + } + + @Override + public int getOrder() { + return Ordered.HIGHEST_PRECEDENCE; + } + + @Override + public Object getEarlyBeanReference(Object bean, String beanName) throws BeansException { + if (bean instanceof FactoryBean) { + return bean; + } + this.earlyReferences.put(getCacheKey(bean, beanName), bean); + return this.overrideRegistrar.wrapIfNecessary(bean, beanName); + } + + @Override + public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException { + if (bean instanceof FactoryBean) { + return bean; + } + if (this.earlyReferences.remove(getCacheKey(bean, beanName)) != bean) { + return this.overrideRegistrar.wrapIfNecessary(bean, beanName); + } + return bean; + } + + private String getCacheKey(Object bean, String beanName) { + return (StringUtils.hasLength(beanName) ? beanName : bean.getClass().getName()); + } + +}