Align withBean methods with ApplicationContext

Rework `AbstractApplicationContextRunner.withBean` methods to
align signatures as much as possible with those provided by
the `ApplicationContext`.

Also update the implementation to use a dedicate member
variable rather than adding initializers.

Closes gh-16011

Author: philwebb

spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/flyway/FlywayEndpointTests.java

@@ -41,7 +41,7 @@ public class FlywayEndpointTests {
 	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
 			.withConfiguration(AutoConfigurations.of(FlywayAutoConfiguration.class))
 			.withUserConfiguration(EmbeddedDataSourceConfiguration.class)
-			.withBean("endpoint", FlywayEndpoint.class, FlywayEndpoint::new);
+			.withBean("endpoint", FlywayEndpoint.class);
 
 	@Test
 	public void flywayReportIsProduced() {

spring-boot-project/spring-boot-test/src/main/java/org/springframework/boot/test/context/runner/AbstractApplicationContextRunner.java

@@ -19,9 +19,12 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
+import org.springframework.beans.factory.config.BeanDefinition;
+import org.springframework.beans.factory.config.BeanDefinitionCustomizer;
 import org.springframework.beans.factory.support.BeanNameGenerator;
 import org.springframework.boot.context.annotation.Configurations;
 import org.springframework.boot.context.annotation.UserConfigurations;
@@ -112,6 +115,8 @@ public abstract class AbstractApplicationContextRunner<SELF extends AbstractAppl
 
 	private final ApplicationContext parent;
 
+	private final List<BeanRegistration<?>> beanRegistrations;
+
 	private final List<Configurations> configurations;
 
 	/**
@@ -120,7 +125,8 @@ public abstract class AbstractApplicationContextRunner<SELF extends AbstractAppl
 	 */
 	protected AbstractApplicationContextRunner(Supplier<C> contextFactory) {
 		this(contextFactory, Collections.emptyList(), TestPropertyValues.empty(),
-				TestPropertyValues.empty(), null, null, Collections.emptyList());
+				TestPropertyValues.empty(), null, null, Collections.emptyList(),
+				Collections.emptyList());
 	}
 
 	/**
@@ -131,12 +137,14 @@ protected AbstractApplicationContextRunner(Supplier<C> contextFactory) {
 	 * @param systemProperties the system properties
 	 * @param classLoader the class loader
 	 * @param parent the parent
+	 * @param beanRegistrations the bean registrations
 	 * @param configurations the configuration
 	 */
 	protected AbstractApplicationContextRunner(Supplier<C> contextFactory,
 			List<ApplicationContextInitializer<? super C>> initializers,
 			TestPropertyValues environmentProperties, TestPropertyValues systemProperties,
 			ClassLoader classLoader, ApplicationContext parent,
+			List<BeanRegistration<?>> beanRegistrations,
 			List<Configurations> configurations) {
 		Assert.notNull(contextFactory, "ContextFactory must not be null");
 		Assert.notNull(environmentProperties, "EnvironmentProperties must not be null");
@@ -149,6 +157,7 @@ protected AbstractApplicationContextRunner(Supplier<C> contextFactory,
 		this.systemProperties = systemProperties;
 		this.classLoader = classLoader;
 		this.parent = parent;
+		this.beanRegistrations = Collections.unmodifiableList(beanRegistrations);
 		this.configurations = Collections.unmodifiableList(configurations);
 	}
 
@@ -162,7 +171,7 @@ public SELF withInitializer(ApplicationContextInitializer<? super C> initializer
 		Assert.notNull(initializer, "Initializer must not be null");
 		return newInstance(this.contextFactory, add(this.initializers, initializer),
 				this.environmentProperties, this.systemProperties, this.classLoader,
-				this.parent, this.configurations);
+				this.parent, this.beanRegistrations, this.configurations);
 	}
 
 	/**
@@ -178,7 +187,8 @@ public SELF withInitializer(ApplicationContextInitializer<? super C> initializer
 	public SELF withPropertyValues(String... pairs) {
 		return newInstance(this.contextFactory, this.initializers,
 				this.environmentProperties.and(pairs), this.systemProperties,
-				this.classLoader, this.parent, this.configurations);
+				this.classLoader, this.parent, this.beanRegistrations,
+				this.configurations);
 	}
 
 	/**
@@ -194,7 +204,8 @@ public SELF withPropertyValues(String... pairs) {
 	public SELF withSystemProperties(String... pairs) {
 		return newInstance(this.contextFactory, this.initializers,
 				this.environmentProperties, this.systemProperties.and(pairs),
-				this.classLoader, this.parent, this.configurations);
+				this.classLoader, this.parent, this.beanRegistrations,
+				this.configurations);
 	}
 
 	/**
@@ -207,7 +218,7 @@ public SELF withSystemProperties(String... pairs) {
 	public SELF withClassLoader(ClassLoader classLoader) {
 		return newInstance(this.contextFactory, this.initializers,
 				this.environmentProperties, this.systemProperties, classLoader,
-				this.parent, this.configurations);
+				this.parent, this.beanRegistrations, this.configurations);
 	}
 
 	/**
@@ -219,7 +230,7 @@ public SELF withClassLoader(ClassLoader classLoader) {
 	public SELF withParent(ApplicationContext parent) {
 		return newInstance(this.contextFactory, this.initializers,
 				this.environmentProperties, this.systemProperties, this.classLoader,
-				parent, this.configurations);
+				parent, this.beanRegistrations, this.configurations);
 	}
 
 	/**
@@ -229,75 +240,81 @@ public SELF withParent(ApplicationContext parent) {
 	 * <p>
 	 * Such beans are registered after regular {@linkplain #withUserConfiguration(Class[])
 	 * user configurations} in the order of registration.
-	 * @param beanType the type of the bean
-	 * @param beanDefinition a supplier for the bean
+	 * @param type the type of the bean
+	 * @param constructorArgs custom argument values to be fed into Spring's constructor
+	 * resolution algorithm, resolving either all arguments or just specific ones, with
+	 * the rest to be resolved through regular autowiring (may be {@code null} or empty)
 	 * @param <T> the type of the bean
 	 * @return a new instance with the updated bean
 	 */
-	public <T> SELF withBean(Class<T> beanType, Supplier<T> beanDefinition) {
-		return withBean(null, beanType, beanDefinition);
+	public <T> SELF withBean(Class<T> type, Object... constructorArgs) {
+		return withBean(null, type, constructorArgs);
 	}
 
 	/**
-	 * Register the specified user bean with the {@link ApplicationContext}. The bean name
-	 * is generated from the configured {@link BeanNameGenerator} on the underlying
-	 * context.
+	 * Register the specified user bean with the {@link ApplicationContext}.
 	 * <p>
 	 * Such beans are registered after regular {@linkplain #withUserConfiguration(Class[])
 	 * user configurations} in the order of registration.
-	 * @param beanType the type of the bean
-	 * @param beanDefinition a function that accepts the context and return the bean
+	 * @param name the bean name or {@code null} to use a generated name
+	 * @param type the type of the bean
+	 * @param constructorArgs custom argument values to be fed into Spring's constructor
+	 * resolution algorithm, resolving either all arguments or just specific ones, with
+	 * the rest to be resolved through regular autowiring (may be {@code null} or empty)
 	 * @param <T> the type of the bean
 	 * @return a new instance with the updated bean
 	 */
-	public <T> SELF withBean(Class<T> beanType, Function<? super C, T> beanDefinition) {
-		return withBean(null, beanType, beanDefinition);
+	public <T> SELF withBean(String name, Class<T> type, Object... constructorArgs) {
+		return newInstance(this.contextFactory, this.initializers,
+				this.environmentProperties, this.systemProperties, this.classLoader,
+				this.parent,
+				add(this.beanRegistrations,
+						new BeanRegistration<>(name, type, constructorArgs)),
+				this.configurations);
 	}
 
 	/**
-	 * Register the specified user bean with the {@link ApplicationContext}. If no bean
-	 * name is provided, a default one is generated from the configured
-	 * {@link BeanNameGenerator} on the underlying context.
+	 * Register the specified user bean with the {@link ApplicationContext}. The bean name
+	 * is generated from the configured {@link BeanNameGenerator} on the underlying
+	 * context.
 	 * <p>
 	 * Such beans are registered after regular {@linkplain #withUserConfiguration(Class[])
 	 * user configurations} in the order of registration.
-	 * @param beanName the name of the bean (may be {@code null})
-	 * @param beanType the type of the bean
-	 * @param beanDefinition a supplier for the bean
+	 * @param type the type of the bean
+	 * @param supplier a supplier for the bean
+	 * @param customizers one or more callbacks for customizing the factory's
+	 * {@link BeanDefinition}, e.g. setting a lazy-init or primary flag
 	 * @param <T> the type of the bean
 	 * @return a new instance with the updated bean
 	 */
-	public <T> SELF withBean(String beanName, Class<T> beanType,
-			Supplier<T> beanDefinition) {
-		return withBean(beanName, beanType, (context) -> beanDefinition.get());
+	public <T> SELF withBean(Class<T> type, Supplier<T> supplier,
+			BeanDefinitionCustomizer... customizers) {
+		return withBean(null, type, supplier, customizers);
 	}
 
 	/**
-	 * Register the specified user bean with the {@link ApplicationContext}. If no bean
-	 * name is provided, a default one is generated from the configured
-	 * {@link BeanNameGenerator} on the underlying context.
+	 * Register the specified user bean with the {@link ApplicationContext}. The bean name
+	 * is generated from the configured {@link BeanNameGenerator} on the underlying
+	 * context.
 	 * <p>
 	 * Such beans are registered after regular {@linkplain #withUserConfiguration(Class[])
 	 * user configurations} in the order of registration.
-	 * @param beanName the name of the bean (may be {@code null})
-	 * @param beanType the type of the bean
-	 * @param beanDefinition a function that accepts the context and return the bean
+	 * @param name the bean name or {@code null} to use a generated name
+	 * @param type the type of the bean
+	 * @param supplier a supplier for the bean
+	 * @param customizers one or more callbacks for customizing the factory's
+	 * {@link BeanDefinition}, e.g. setting a lazy-init or primary flag
 	 * @param <T> the type of the bean
 	 * @return a new instance with the updated bean
 	 */
-	public <T> SELF withBean(String beanName, Class<T> beanType,
-			Function<? super C, T> beanDefinition) {
-		return withInitializer(
-				beanDefinitionRegistrar(beanName, beanType, beanDefinition));
-	}
-
-	private <T> ApplicationContextInitializer<? super C> beanDefinitionRegistrar(
-			String beanName, Class<T> beanType, Function<? super C, T> beanDefinition) {
-		return (context) -> {
-			Assert.isInstanceOf(GenericApplicationContext.class, context);
-			((GenericApplicationContext) context).registerBean(beanName, beanType,
-					() -> beanDefinition.apply(context));
-		};
+	public <T> SELF withBean(String name, Class<T> type, Supplier<T> supplier,
+			BeanDefinitionCustomizer... customizers) {
+		return newInstance(this.contextFactory, this.initializers,
+				this.environmentProperties, this.systemProperties, this.classLoader,
+				this.parent,
+				add(this.beanRegistrations,
+						new BeanRegistration<>(name, type, supplier, customizers)),
+				this.configurations);
 	}
 
 	/**
@@ -319,7 +336,8 @@ public SELF withConfiguration(Configurations configurations) {
 		Assert.notNull(configurations, "Configurations must not be null");
 		return newInstance(this.contextFactory, this.initializers,
 				this.environmentProperties, this.systemProperties, this.classLoader,
-				this.parent, add(this.configurations, configurations));
+				this.parent, this.beanRegistrations,
+				add(this.configurations, configurations));
 	}
 
 	/**
@@ -342,6 +360,7 @@ protected abstract SELF newInstance(Supplier<C> contextFactory,
 			List<ApplicationContextInitializer<? super C>> initializers,
 			TestPropertyValues environmentProperties, TestPropertyValues systemProperties,
 			ClassLoader classLoader, ApplicationContext parent,
+			List<BeanRegistration<?>> beanRegistrations,
 			List<Configurations> configurations);
 
 	/**
@@ -416,6 +435,7 @@ private void configureContext(C context) {
 		if (classes.length > 0) {
 			((AnnotationConfigRegistry) context).register(classes);
 		}
+		this.beanRegistrations.forEach((registration) -> registration.apply(context));
 		this.initializers.forEach((initializer) -> initializer.initialize(context));
 		context.refresh();
 	}
@@ -434,4 +454,31 @@ private <E extends Throwable> void rethrow(Throwable e) throws E {
 		throw (E) e;
 	}
 
+	/**
+	 * A Bean registration to be applied when the context loaded.
+	 *
+	 * @param <T> the bean type
+	 */
+	protected final class BeanRegistration<T> {
+
+		Consumer<GenericApplicationContext> registrar;
+
+		public BeanRegistration(String name, Class<T> type, Object... constructorArgs) {
+			this.registrar = (context) -> context.registerBean(name, type,
+					constructorArgs);
+		}
+
+		public BeanRegistration(String name, Class<T> type, Supplier<T> supplier,
+				BeanDefinitionCustomizer... customizers) {
+			this.registrar = (context) -> context.registerBean(name, type, supplier,
+					customizers);
+		}
+
+		public void apply(ConfigurableApplicationContext context) {
+			Assert.isInstanceOf(GenericApplicationContext.class, context);
+			this.registrar.accept(((GenericApplicationContext) context));
+		}
+
+	}
+
 }

spring-boot-project/spring-boot-test/src/main/java/org/springframework/boot/test/context/runner/ApplicationContextRunner.java

@@ -64,9 +64,10 @@ private ApplicationContextRunner(
 			List<ApplicationContextInitializer<? super ConfigurableApplicationContext>> initializers,
 			TestPropertyValues environmentProperties, TestPropertyValues systemProperties,
 			ClassLoader classLoader, ApplicationContext parent,
+			List<BeanRegistration<?>> beanRegistrations,
 			List<Configurations> configurations) {
 		super(contextFactory, initializers, environmentProperties, systemProperties,
-				classLoader, parent, configurations);
+				classLoader, parent, beanRegistrations, configurations);
 	}
 
 	@Override
@@ -75,10 +76,11 @@ protected ApplicationContextRunner newInstance(
 			List<ApplicationContextInitializer<? super ConfigurableApplicationContext>> initializers,
 			TestPropertyValues environmentProperties, TestPropertyValues systemProperties,
 			ClassLoader classLoader, ApplicationContext parent,
+			List<BeanRegistration<?>> beanRegistrations,
 			List<Configurations> configurations) {
 		return new ApplicationContextRunner(contextFactory, initializers,
 				environmentProperties, systemProperties, classLoader, parent,
-				configurations);
+				beanRegistrations, configurations);
 	}
 
 }

spring-boot-project/spring-boot-test/src/main/java/org/springframework/boot/test/context/runner/ReactiveWebApplicationContextRunner.java

@@ -64,9 +64,10 @@ private ReactiveWebApplicationContextRunner(
 			List<ApplicationContextInitializer<? super ConfigurableReactiveWebApplicationContext>> initializers,
 			TestPropertyValues environmentProperties, TestPropertyValues systemProperties,
 			ClassLoader classLoader, ApplicationContext parent,
+			List<BeanRegistration<?>> beanRegistrations,
 			List<Configurations> configurations) {
 		super(contextFactory, initializers, environmentProperties, systemProperties,
-				classLoader, parent, configurations);
+				classLoader, parent, beanRegistrations, configurations);
 	}
 
 	@Override
@@ -75,10 +76,11 @@ protected ReactiveWebApplicationContextRunner newInstance(
 			List<ApplicationContextInitializer<? super ConfigurableReactiveWebApplicationContext>> initializers,
 			TestPropertyValues environmentProperties, TestPropertyValues systemProperties,
 			ClassLoader classLoader, ApplicationContext parent,
+			List<BeanRegistration<?>> beanRegistrations,
 			List<Configurations> configurations) {
 		return new ReactiveWebApplicationContextRunner(contextFactory, initializers,
 				environmentProperties, systemProperties, classLoader, parent,
-				configurations);
+				beanRegistrations, configurations);
 	}
 
 }

spring-boot-project/spring-boot-test/src/main/java/org/springframework/boot/test/context/runner/WebApplicationContextRunner.java

@@ -68,9 +68,10 @@ private WebApplicationContextRunner(
 			List<ApplicationContextInitializer<? super ConfigurableWebApplicationContext>> initializers,
 			TestPropertyValues environmentProperties, TestPropertyValues systemProperties,
 			ClassLoader classLoader, ApplicationContext parent,
+			List<BeanRegistration<?>> beanRegistrations,
 			List<Configurations> configurations) {
 		super(contextFactory, initializers, environmentProperties, systemProperties,
-				classLoader, parent, configurations);
+				classLoader, parent, beanRegistrations, configurations);
 	}
 
 	@Override
@@ -79,10 +80,11 @@ protected WebApplicationContextRunner newInstance(
 			List<ApplicationContextInitializer<? super ConfigurableWebApplicationContext>> initializers,
 			TestPropertyValues environmentProperties, TestPropertyValues systemProperties,
 			ClassLoader classLoader, ApplicationContext parent,
+			List<BeanRegistration<?>> beanRegistrations,
 			List<Configurations> configurations) {
 		return new WebApplicationContextRunner(contextFactory, initializers,
 				environmentProperties, systemProperties, classLoader, parent,
-				configurations);
+				beanRegistrations, configurations);
 	}
 
 	/**

spring-boot-project/spring-boot-test/src/test/java/org/springframework/boot/test/context/runner/AbstractApplicationContextRunnerTests.java

@@ -39,7 +39,6 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 import static org.assertj.core.api.Assertions.assertThatIOException;
-import static org.assertj.core.api.Assertions.entry;
 
 /**
  * Abstract tests for {@link AbstractApplicationContextRunner} implementations.
@@ -167,15 +166,6 @@ public void runWithConfigurationsAndUserBeanShouldRegisterUserBeanLast() {
 				});
 	}
 
-	@Test
-	public void runWithUserBeanShouldHaveAccessToContext() {
-		get().withUserConfiguration(FooConfig.class)
-				.withBean(String.class, (context) -> "Result: " + context.getBean("foo"))
-				.run((context) -> assertThat(context.getBeansOfType(String.class))
-						.containsOnly(entry("foo", "foo"),
-								entry("string", "Result: foo")));
-	}
-
 	@Test
 	public void runWithMultipleConfigurationsShouldRegisterAllConfigurations() {
 		get().withUserConfiguration(FooConfig.class)