From f7e32a9c7890f36cd2f4c975b4f851850437ee85 Mon Sep 17 00:00:00 2001
From: Sam Brannen <104798+sbrannen@users.noreply.github.com>
Date: Thu, 10 Oct 2024 12:13:44 +0200
Subject: [PATCH] Revise singleton destruction for Bean Overrides

See gh-33678
---
 .../BeanOverrideBeanFactoryPostProcessor.java | 28 ++++++++-----------
 1 file changed, 12 insertions(+), 16 deletions(-)

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 32a25f9f6e..c024d29360 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
@@ -154,12 +154,6 @@ class BeanOverrideBeanFactoryPostProcessor implements BeanFactoryPostProcessor,
 			}
 		}
 
-		// Ensure we don't have any manually registered singletons registered, since we
-		// register a bean override instance as a manual singleton at the end of this method.
-		if (beanFactory.containsSingleton(beanName)) {
-			destroySingleton(beanFactory, beanName);
-		}
-
 		if (existingBeanDefinition != null) {
 			// Validate the existing bean definition.
 			//
@@ -200,14 +194,19 @@ class BeanOverrideBeanFactoryPostProcessor implements BeanFactoryPostProcessor,
 		overrideMetadata.track(override, beanFactory);
 		this.overrideRegistrar.registerNameForMetadata(overrideMetadata, beanName);
 
-		// 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.
+		// Now we have an instance (the override) that we can manually register as a singleton.
+		//
+		// However, we need to remove any existing singleton instance -- for example, a
+		// manually registered singleton or a singleton that was registered as a side effect
+		// of the isSingleton() check in validateBeanDefinition().
 		//
 		// 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.
+		if (beanFactory.containsSingleton(beanName)) {
+			destroySingleton(beanFactory, beanName);
+		}
 		beanFactory.registerSingleton(beanName, override);
 	}
 
@@ -339,17 +338,14 @@ class BeanOverrideBeanFactoryPostProcessor implements BeanFactoryPostProcessor,
 	private static void validateBeanDefinition(ConfigurableListableBeanFactory beanFactory, String beanName) {
 		Assert.state(beanFactory.isSingleton(beanName),
 				() -> "Unable to override bean '" + beanName + "': only singleton beans can be overridden.");
-
-		// Since the isSingleton() check above may have registered a singleton as a side
-		// effect -- for example, for a FactoryBean -- we need to destroy the singleton,
-		// because we later manually register a bean override instance as a singleton.
-		destroySingleton(beanFactory, beanName);
 	}
 
 	private static void destroySingleton(ConfigurableListableBeanFactory beanFactory, String beanName) {
-		if (beanFactory instanceof DefaultListableBeanFactory dlbf) {
-			dlbf.destroySingleton(beanName);
+		if (!(beanFactory instanceof DefaultListableBeanFactory dlbf)) {
+			throw new IllegalStateException("Cannot process bean override with a BeanFactory " +
+					"that doesn't implement DefaultListableBeanFactory: " + beanFactory.getClass().getName());
 		}
+		dlbf.destroySingleton(beanName);
 	}
 
 }