Allow additional context interfaces to be defined for testing

Update `AssertableApplicationContext` and `ApplicationContextRunner`
implementations to support additional `ApplicationContext` interfaces.

Closes gh-42369

Author: philwebb

spring-boot-project/spring-boot-test/src/main/java/org/springframework/boot/test/context/assertj/ApplicationContextAssertProvider.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2022 the original author or authors.
+ * Copyright 2012-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.
@@ -18,12 +18,14 @@
 
 import java.io.Closeable;
 import java.lang.reflect.Proxy;
+import java.util.Arrays;
 import java.util.function.Supplier;
 
 import org.assertj.core.api.AssertProvider;
 
 import org.springframework.context.ApplicationContext;
 import org.springframework.util.Assert;
+import org.springframework.util.ObjectUtils;
 
 /**
  * An {@link ApplicationContext} that additionally supports AssertJ style assertions. Can
@@ -101,16 +103,46 @@ public interface ApplicationContextAssertProvider<C extends ApplicationContext>
 	 * {@link ApplicationContext} or throw an exception if the context fails to start.
 	 * @return a {@link ApplicationContextAssertProvider} instance
 	 */
-	@SuppressWarnings("unchecked")
 	static <T extends ApplicationContextAssertProvider<C>, C extends ApplicationContext> T get(Class<T> type,
 			Class<? extends C> contextType, Supplier<? extends C> contextSupplier) {
+		return get(type, contextType, contextSupplier, new Class<?>[0]);
+	}
+
+	/**
+	 * Factory method to create a new {@link ApplicationContextAssertProvider} instance.
+	 * @param <T> the assert provider type
+	 * @param <C> the context type
+	 * @param type the type of {@link ApplicationContextAssertProvider} required (must be
+	 * an interface)
+	 * @param contextType the type of {@link ApplicationContext} being managed (must be an
+	 * interface)
+	 * @param contextSupplier a supplier that will either return a fully configured
+	 * {@link ApplicationContext} or throw an exception if the context fails to start.
+	 * @param additionalContextInterfaces and additional context interfaces to add to the
+	 * proxy
+	 * @return a {@link ApplicationContextAssertProvider} instance
+	 * @since 3.4.0
+	 */
+	@SuppressWarnings("unchecked")
+	static <T extends ApplicationContextAssertProvider<C>, C extends ApplicationContext> T get(Class<T> type,
+			Class<? extends C> contextType, Supplier<? extends C> contextSupplier,
+			Class<?>... additionalContextInterfaces) {
 		Assert.notNull(type, "Type must not be null");
 		Assert.isTrue(type.isInterface(), "Type must be an interface");
 		Assert.notNull(contextType, "ContextType must not be null");
 		Assert.isTrue(contextType.isInterface(), "ContextType must be an interface");
-		Class<?>[] interfaces = { type, contextType };
+		Class<?>[] interfaces = merge(new Class<?>[] { type, contextType }, additionalContextInterfaces);
 		return (T) Proxy.newProxyInstance(Thread.currentThread().getContextClassLoader(), interfaces,
 				new AssertProviderApplicationContextInvocationHandler(contextType, contextSupplier));
 	}
 
+	private static Class<?>[] merge(Class<?>[] classes, Class<?>[] additional) {
+		if (ObjectUtils.isEmpty(additional)) {
+			return classes;
+		}
+		Class<?>[] result = Arrays.copyOf(classes, classes.length + additional.length);
+		System.arraycopy(additional, 0, result, classes.length, additional.length);
+		return result;
+	}
+
 }

spring-boot-project/spring-boot-test/src/main/java/org/springframework/boot/test/context/assertj/AssertableApplicationContext.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2019 the original author or authors.
+ * Copyright 2012-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.
@@ -49,4 +49,20 @@ static AssertableApplicationContext get(Supplier<? extends ConfigurableApplicati
 				ConfigurableApplicationContext.class, contextSupplier);
 	}
 
+	/**
+	 * Factory method to create a new {@link AssertableApplicationContext} instance.
+	 * @param contextSupplier a supplier that will either return a fully configured
+	 * {@link ConfigurableApplicationContext} or throw an exception if the context fails
+	 * to start.
+	 * @param additionalContextInterfaces and additional context interfaces to add to the
+	 * proxy
+	 * @return an {@link AssertableApplicationContext} instance
+	 * @since 3.4.0
+	 */
+	static AssertableApplicationContext get(Supplier<? extends ConfigurableApplicationContext> contextSupplier,
+			Class<?>... additionalContextInterfaces) {
+		return ApplicationContextAssertProvider.get(AssertableApplicationContext.class,
+				ConfigurableApplicationContext.class, contextSupplier, additionalContextInterfaces);
+	}
+
 }

spring-boot-project/spring-boot-test/src/main/java/org/springframework/boot/test/context/assertj/AssertableReactiveWebApplicationContext.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2019 the original author or authors.
+ * Copyright 2012-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.
@@ -51,4 +51,22 @@ static AssertableReactiveWebApplicationContext get(
 				ConfigurableReactiveWebApplicationContext.class, contextSupplier);
 	}
 
+	/**
+	 * Factory method to create a new {@link AssertableReactiveWebApplicationContext}
+	 * instance.
+	 * @param contextSupplier a supplier that will either return a fully configured
+	 * {@link ConfigurableReactiveWebApplicationContext} or throw an exception if the
+	 * context fails to start.
+	 * @param additionalContextInterfaces and additional context interfaces to add to the
+	 * proxy
+	 * @return a {@link AssertableReactiveWebApplicationContext} instance
+	 * @since 3.4.0
+	 */
+	static AssertableReactiveWebApplicationContext get(
+			Supplier<? extends ConfigurableReactiveWebApplicationContext> contextSupplier,
+			Class<?>... additionalContextInterfaces) {
+		return ApplicationContextAssertProvider.get(AssertableReactiveWebApplicationContext.class,
+				ConfigurableReactiveWebApplicationContext.class, contextSupplier, additionalContextInterfaces);
+	}
+
 }

spring-boot-project/spring-boot-test/src/main/java/org/springframework/boot/test/context/assertj/AssertableWebApplicationContext.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2019 the original author or authors.
+ * Copyright 2012-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.
@@ -49,4 +49,20 @@ static AssertableWebApplicationContext get(Supplier<? extends ConfigurableWebApp
 				ConfigurableWebApplicationContext.class, contextSupplier);
 	}
 
+	/**
+	 * Factory method to create a new {@link AssertableWebApplicationContext} instance.
+	 * @param contextSupplier a supplier that will either return a fully configured
+	 * {@link ConfigurableWebApplicationContext} or throw an exception if the context
+	 * fails to start.
+	 * @param additionalContextInterfaces and additional context interfaces to add to the
+	 * proxy
+	 * @return a {@link AssertableWebApplicationContext} instance
+	 * @since 3.4.0
+	 */
+	static AssertableWebApplicationContext get(Supplier<? extends ConfigurableWebApplicationContext> contextSupplier,
+			Class<?>... additionalContextInterfaces) {
+		return ApplicationContextAssertProvider.get(AssertableWebApplicationContext.class,
+				ConfigurableWebApplicationContext.class, contextSupplier, additionalContextInterfaces);
+	}
+
 }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2023 the original author or authors.
+ * Copyright 2012-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.
@@ -106,6 +106,8 @@
  */
 public abstract class AbstractApplicationContextRunner<SELF extends AbstractApplicationContextRunner<SELF, C, A>, C extends ConfigurableApplicationContext, A extends ApplicationContextAssertProvider<C>> {
 
+	private static final Class<?>[] NO_ADDITIONAL_CONTEXT_INTERFACES = {};
+
 	private final RunnerConfiguration<C> runnerConfiguration;
 
 	private final Function<RunnerConfiguration<C>, SELF> instanceFactory;
@@ -115,13 +117,29 @@ public abstract class AbstractApplicationContextRunner<SELF extends AbstractAppl
 	 * @param contextFactory the factory used to create the actual context
 	 * @param instanceFactory the factory used to create new instance of the runner
 	 * @since 2.6.0
+	 * @deprecated since 3.4.0 for removal in 3.6.0 in favor of
+	 * {@link #AbstractApplicationContextRunner(Function, Supplier, Class...)}
 	 */
+	@Deprecated(since = "3.4.0", forRemoval = true)
 	protected AbstractApplicationContextRunner(Supplier<C> contextFactory,
 			Function<RunnerConfiguration<C>, SELF> instanceFactory) {
-		Assert.notNull(contextFactory, "ContextFactory must not be null");
-		Assert.notNull(contextFactory, "RunnerConfiguration must not be null");
-		this.runnerConfiguration = new RunnerConfiguration<>(contextFactory);
+		this(instanceFactory, contextFactory, NO_ADDITIONAL_CONTEXT_INTERFACES);
+	}
+
+	/**
+	 * Create a new {@link AbstractApplicationContextRunner} instance.
+	 * @param instanceFactory the factory used to create new instance of the runner
+	 * @param contextFactory the factory used to create the actual context
+	 * @param additionalContextInterfaces any additional application context interfaces to
+	 * be added to the application context proxy
+	 * @since 3.4.0
+	 */
+	protected AbstractApplicationContextRunner(Function<RunnerConfiguration<C>, SELF> instanceFactory,
+			Supplier<C> contextFactory, Class<?>... additionalContextInterfaces) {
+		Assert.notNull(instanceFactory, "'instanceFactory' must not be null");
+		Assert.notNull(contextFactory, "'contextFactory' must not be null");
 		this.instanceFactory = instanceFactory;
+		this.runnerConfiguration = new RunnerConfiguration<>(contextFactory, additionalContextInterfaces);
 	}
 
 	/**
@@ -386,7 +404,8 @@ private A createAssertableContext(boolean refresh) {
 		ResolvableType resolvableType = ResolvableType.forClass(AbstractApplicationContextRunner.class, getClass());
 		Class<A> assertType = (Class<A>) resolvableType.resolveGeneric(1);
 		Class<C> contextType = (Class<C>) resolvableType.resolveGeneric(2);
-		return ApplicationContextAssertProvider.get(assertType, contextType, () -> createAndLoadContext(refresh));
+		return ApplicationContextAssertProvider.get(assertType, contextType, () -> createAndLoadContext(refresh),
+				this.runnerConfiguration.additionalContextInterfaces);
 	}
 
 	private C createAndLoadContext(boolean refresh) {
@@ -472,6 +491,8 @@ protected static final class RunnerConfiguration<C extends ConfigurableApplicati
 
 		private final Supplier<C> contextFactory;
 
+		private final Class<?>[] additionalContextInterfaces;
+
 		private boolean allowBeanDefinitionOverriding = false;
 
 		private boolean allowCircularReferences = false;
@@ -490,12 +511,14 @@ protected static final class RunnerConfiguration<C extends ConfigurableApplicati
 
 		private List<Configurations> configurations = Collections.emptyList();
 
-		private RunnerConfiguration(Supplier<C> contextFactory) {
+		private RunnerConfiguration(Supplier<C> contextFactory, Class<?>[] additionalContextInterfaces) {
 			this.contextFactory = contextFactory;
+			this.additionalContextInterfaces = additionalContextInterfaces;
 		}
 
 		private RunnerConfiguration(RunnerConfiguration<C> source) {
 			this.contextFactory = source.contextFactory;
+			this.additionalContextInterfaces = source.additionalContextInterfaces;
 			this.allowBeanDefinitionOverriding = source.allowBeanDefinitionOverriding;
 			this.allowCircularReferences = source.allowCircularReferences;
 			this.initializers = source.initializers;

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2022 the original author or authors.
+ * Copyright 2012-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.
@@ -47,10 +47,24 @@ public ApplicationContextRunner() {
 	/**
 	 * Create a new {@link ApplicationContextRunner} instance using the specified
 	 * {@code contextFactory} as the underlying source.
-	 * @param contextFactory a supplier that returns a new instance on each call
+	 * @param contextFactory a supplier that returns a new instance on each call be added
+	 * to the application context proxy
 	 */
 	public ApplicationContextRunner(Supplier<ConfigurableApplicationContext> contextFactory) {
-		super(contextFactory, ApplicationContextRunner::new);
+		super(ApplicationContextRunner::new, contextFactory);
+	}
+
+	/**
+	 * Create a new {@link ApplicationContextRunner} instance using the specified
+	 * {@code contextFactory} as the underlying source.
+	 * @param contextFactory a supplier that returns a new instance on each call
+	 * @param additionalContextInterfaces any additional application context interfaces to
+	 * be added to the application context proxy
+	 * @since 3.4.0
+	 */
+	public ApplicationContextRunner(Supplier<ConfigurableApplicationContext> contextFactory,
+			Class<?>... additionalContextInterfaces) {
+		super(ApplicationContextRunner::new, contextFactory, additionalContextInterfaces);
 	}
 
 	private ApplicationContextRunner(RunnerConfiguration<ConfigurableApplicationContext> runnerConfiguration) {

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2022 the original author or authors.
+ * Copyright 2012-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.
@@ -47,10 +47,25 @@ public ReactiveWebApplicationContextRunner() {
 	/**
 	 * Create a new {@link ApplicationContextRunner} instance using the specified
 	 * {@code contextFactory} as the underlying source.
-	 * @param contextFactory a supplier that returns a new instance on each call
+	 * @param contextFactory a supplier that returns a new instance on each call be added
+	 * to the application context proxy
+	 * @since 3.4.0
 	 */
 	public ReactiveWebApplicationContextRunner(Supplier<ConfigurableReactiveWebApplicationContext> contextFactory) {
-		super(contextFactory, ReactiveWebApplicationContextRunner::new);
+		super(ReactiveWebApplicationContextRunner::new, contextFactory);
+	}
+
+	/**
+	 * Create a new {@link ApplicationContextRunner} instance using the specified
+	 * {@code contextFactory} as the underlying source.
+	 * @param contextFactory a supplier that returns a new instance on each call
+	 * @param additionalContextInterfaces any additional application context interfaces to
+	 * be added to the application context proxy
+	 * @since 3.4.0
+	 */
+	public ReactiveWebApplicationContextRunner(Supplier<ConfigurableReactiveWebApplicationContext> contextFactory,
+			Class<?>... additionalContextInterfaces) {
+		super(ReactiveWebApplicationContextRunner::new, contextFactory, additionalContextInterfaces);
 	}
 
 	private ReactiveWebApplicationContextRunner(

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2022 the original author or authors.
+ * Copyright 2012-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.
@@ -51,10 +51,24 @@ public WebApplicationContextRunner() {
 	/**
 	 * Create a new {@link WebApplicationContextRunner} instance using the specified
 	 * {@code contextFactory} as the underlying source.
-	 * @param contextFactory a supplier that returns a new instance on each call
+	 * @param contextFactory a supplier that returns a new instance on each call be added
+	 * to the application context proxy
 	 */
 	public WebApplicationContextRunner(Supplier<ConfigurableWebApplicationContext> contextFactory) {
-		super(contextFactory, WebApplicationContextRunner::new);
+		super(WebApplicationContextRunner::new, contextFactory);
+	}
+
+	/**
+	 * Create a new {@link WebApplicationContextRunner} instance using the specified
+	 * {@code contextFactory} as the underlying source.
+	 * @param contextFactory a supplier that returns a new instance on each call
+	 * @param additionalContextInterfaces any additional application context interfaces to
+	 * be added to the application context proxy
+	 * @since 3.4.0
+	 */
+	public WebApplicationContextRunner(Supplier<ConfigurableWebApplicationContext> contextFactory,
+			Class<?>... additionalContextInterfaces) {
+		super(WebApplicationContextRunner::new, contextFactory, additionalContextInterfaces);
 	}
 
 	private WebApplicationContextRunner(RunnerConfiguration<ConfigurableWebApplicationContext> configuration) {

spring-boot-project/spring-boot-test/src/test/java/org/springframework/boot/test/context/assertj/AdditionalContextInterface.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2012-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.boot.test.context.assertj;
+
+import org.springframework.context.ApplicationContext;
+
+/**
+ * Tests extra interface that can be applied to an {@link ApplicationContext}
+ *
+ * @author Phillip Webb
+ */
+interface AdditionalContextInterface {
+
+}