Introduce BeanOverrideHandler.createOverrideInstance() template method

Author: sbrannen

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

@@ -191,8 +191,7 @@ else if (Boolean.getBoolean(AbstractAotProcessor.AOT_PROCESSING)) {
 			registry.registerBeanDefinition(beanName, pseudoBeanDefinition);
 		}
 
-		Object override = handler.createOverrideInstance(beanName, existingBeanDefinition, null);
-		handler.trackOverrideInstance(override, beanFactory);
+		Object override = handler.createOverrideInstance(beanName, existingBeanDefinition, null, beanFactory);
 		this.beanOverrideRegistry.registerBeanOverrideHandler(handler, beanName);
 
 		// Now we have an instance (the override) that we can manually register as a singleton.

spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideHandler.java

@@ -151,6 +151,31 @@ public final BeanOverrideStrategy getStrategy() {
 		return this.strategy;
 	}
 
+	/**
+	 * {@linkplain #createOverrideInstance Create} and
+	 * {@linkplain #trackOverrideInstance track} a bean override instance for an
+	 * existing {@link BeanDefinition} or an existing singleton bean, based on the
+	 * metadata in this {@code BeanOverrideHandler}.
+	 * @param beanName the name of the bean being overridden
+	 * @param existingBeanDefinition an existing bean definition for the supplied
+	 * bean name, or {@code null} if not available or not relevant
+	 * @param existingBeanInstance an existing instance for the supplied bean name
+	 * for wrapping purposes, or {@code null} if not available or not relevant
+	 * @param singletonBeanRegistry a registry in which this handler can store
+	 * tracking state in the form of a singleton bean
+	 * @return the instance with which to override the bean
+	 * @see #trackOverrideInstance(Object, SingletonBeanRegistry)
+	 * @see #createOverrideInstance(String, BeanDefinition, Object)
+	 */
+	final Object createOverrideInstance(
+			String beanName, @Nullable BeanDefinition existingBeanDefinition,
+			@Nullable Object existingBeanInstance, SingletonBeanRegistry singletonBeanRegistry) {
+
+		Object override = createOverrideInstance(beanName, existingBeanDefinition, existingBeanInstance);
+		trackOverrideInstance(override, singletonBeanRegistry);
+		return override;
+	}
+
 	/**
 	 * Create a bean override instance for an existing {@link BeanDefinition} or
 	 * an existing singleton bean, based on the metadata in this
@@ -172,11 +197,11 @@ protected abstract Object createOverrideInstance(String beanName,
 	 * <p>The default implementation does not track the supplied instance, but
 	 * this can be overridden in subclasses as appropriate.
 	 * @param override the bean override instance to track
-	 * @param trackingBeanRegistry a registry in which this handler can store
+	 * @param singletonBeanRegistry a registry in which this handler can store
 	 * tracking state in the form of a singleton bean
 	 * @see #createOverrideInstance(String, BeanDefinition, Object)
 	 */
-	protected void trackOverrideInstance(Object override, SingletonBeanRegistry trackingBeanRegistry) {
+	protected void trackOverrideInstance(Object override, SingletonBeanRegistry singletonBeanRegistry) {
 		// NO-OP
 	}
 

spring-test/src/main/java/org/springframework/test/context/bean/override/BeanOverrideRegistry.java

@@ -78,9 +78,7 @@ Object wrapBeanIfNecessary(Object bean, String beanName) {
 		BeanOverrideHandler handler = this.wrappingBeanOverrideHandlers.get(beanName);
 		Assert.state(handler != null,
 				() -> "Failed to find wrapping BeanOverrideHandler for bean '" + beanName + "'");
-		bean = handler.createOverrideInstance(beanName, null, bean);
-		handler.trackOverrideInstance(bean, this.beanFactory);
-		return bean;
+		return handler.createOverrideInstance(beanName, null, bean, this.beanFactory);
 	}
 
 	void inject(Object target, BeanOverrideHandler handler) {