Migrate some more tests to the new ContextLoader

spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/reactive/HttpHandlerAutoConfigurationTests.java

@@ -16,12 +16,9 @@
 
 package org.springframework.boot.autoconfigure.web.reactive;
 
-import org.junit.Rule;
 import org.junit.Test;
-import org.junit.rules.ExpectedException;
 
-import org.springframework.boot.test.util.TestPropertyValues;
-import org.springframework.boot.web.reactive.context.GenericReactiveWebApplicationContext;
+import org.springframework.boot.test.context.ContextLoader;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.http.server.reactive.HttpHandler;
@@ -37,36 +34,27 @@ import static org.assertj.core.api.Assertions.assertThat;
  *
  * @author Brian Clozel
  * @author Stephane Nicoll
+ * @author Andy Wilkinson
  */
 public class HttpHandlerAutoConfigurationTests {
 
-	@Rule
-	public ExpectedException thrown = ExpectedException.none();
-
-	private GenericReactiveWebApplicationContext context;
+	private final ContextLoader contextLoader = new ContextLoader().webReactive()
+			.autoConfig(HttpHandlerAutoConfiguration.class);
 
 	@Test
 	public void shouldNotProcessIfExistingHttpHandler() {
-		load(CustomHttpHandler.class);
-		assertThat(this.context.getBeansOfType(HttpHandler.class)).hasSize(1);
-		assertThat(this.context.getBean(HttpHandler.class))
-				.isSameAs(this.context.getBean("customHttpHandler"));
+		this.contextLoader.config(CustomHttpHandler.class).load(context -> {
+			assertThat(context.getBeansOfType(HttpHandler.class)).hasSize(1);
+			assertThat(context.getBean(HttpHandler.class))
+					.isSameAs(context.getBean("customHttpHandler"));
+		});
 	}
 
 	@Test
 	public void shouldConfigureHttpHandlerAnnotation() {
-		load(WebFluxAutoConfiguration.class);
-		assertThat(this.context.getBeansOfType(HttpHandler.class).size()).isEqualTo(1);
-	}
-
-	private void load(Class<?> config, String... environment) {
-		this.context = new GenericReactiveWebApplicationContext();
-		TestPropertyValues.of(environment).applyTo(this.context);
-		if (this.context != null) {
-			this.context.register(config);
-		}
-		this.context.register(HttpHandlerAutoConfiguration.class);
-		this.context.refresh();
+		this.contextLoader.autoConfig(WebFluxAutoConfiguration.class).load(context -> {
+			assertThat(context.getBeansOfType(HttpHandler.class).size()).isEqualTo(1);
+		});
 	}
 
 	@Configuration

spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/webservices/WebServicesAutoConfigurationTests.java

@@ -16,17 +16,14 @@
 
 package org.springframework.boot.autoconfigure.webservices;
 
-import org.junit.After;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 
 import org.springframework.beans.factory.BeanCreationException;
-import org.springframework.boot.test.util.TestPropertyValues;
+import org.springframework.boot.test.context.ContextLoader;
 import org.springframework.boot.web.servlet.ServletRegistrationBean;
-import org.springframework.mock.web.MockServletContext;
 import org.springframework.test.util.ReflectionTestUtils;
-import org.springframework.web.context.support.AnnotationConfigWebApplicationContext;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -35,79 +32,73 @@ import static org.assertj.core.api.Assertions.assertThat;
  *
  * @author Vedran Pavic
  * @author Stephane Nicoll
+ * @author Andy Wilkinson
  */
 public class WebServicesAutoConfigurationTests {
 
+	private final ContextLoader contextLoader = new ContextLoader().webServlet()
+			.autoConfig(WebServicesAutoConfiguration.class);
+
 	@Rule
 	public ExpectedException thrown = ExpectedException.none();
 
-	private AnnotationConfigWebApplicationContext context;
-
-	@After
-	public void close() {
-		if (this.context != null) {
-			this.context.close();
-		}
-	}
-
 	@Test
 	public void defaultConfiguration() {
-		load(WebServicesAutoConfiguration.class);
-		assertThat(this.context.getBeansOfType(ServletRegistrationBean.class)).hasSize(1);
+		this.contextLoader.load(context -> {
+			assertThat(context.getBeansOfType(ServletRegistrationBean.class)).hasSize(1);
+		});
 	}
 
 	@Test
 	public void customPathMustBeginWithASlash() {
-		this.thrown.expect(BeanCreationException.class);
-		this.thrown.expectMessage("Failed to bind properties under 'spring.webservices'");
-		load(WebServicesAutoConfiguration.class, "spring.webservices.path=invalid");
-	}
-
-	@Test
-	public void customPathWithTrailingSlash() {
-		load(WebServicesAutoConfiguration.class, "spring.webservices.path=/valid/");
-		ServletRegistrationBean<?> servletRegistrationBean = this.context
-				.getBean(ServletRegistrationBean.class);
-		assertThat(servletRegistrationBean.getUrlMappings()).contains("/valid/*");
+		this.contextLoader.env("spring.webservices.path=invalid")
+				.loadAndFail(BeanCreationException.class, (ex) -> {
+					System.out.println(ex.getMessage());
+					assertThat(ex.getMessage()).contains(
+							"Failed to bind properties under 'spring.webservices'");
+				});
 	}
 
 	@Test
 	public void customPath() {
-		load(WebServicesAutoConfiguration.class, "spring.webservices.path=/valid");
-		assertThat(this.context.getBeansOfType(ServletRegistrationBean.class)).hasSize(1);
-		ServletRegistrationBean<?> servletRegistrationBean = this.context
-				.getBean(ServletRegistrationBean.class);
-		assertThat(servletRegistrationBean.getUrlMappings()).contains("/valid/*");
+		this.contextLoader.env("spring.webservices.path=/valid").load(context -> {
+			ServletRegistrationBean<?> servletRegistrationBean = context
+					.getBean(ServletRegistrationBean.class);
+			assertThat(servletRegistrationBean.getUrlMappings()).contains("/valid/*");
+		});
+	}
+
+	@Test
+	public void customPathWithTrailingSlash() {
+		this.contextLoader.env("spring.webservices.path=/valid/").load(context -> {
+			ServletRegistrationBean<?> servletRegistrationBean = context
+					.getBean(ServletRegistrationBean.class);
+			assertThat(servletRegistrationBean.getUrlMappings()).contains("/valid/*");
+		});
 	}
 
 	@Test
 	public void customLoadOnStartup() {
-		load(WebServicesAutoConfiguration.class,
-				"spring.webservices.servlet.load-on-startup=1");
-		ServletRegistrationBean<?> registrationBean = this.context
-				.getBean(ServletRegistrationBean.class);
-		assertThat(ReflectionTestUtils.getField(registrationBean, "loadOnStartup"))
-				.isEqualTo(1);
+		this.contextLoader.env("spring.webservices.servlet.load-on-startup=1")
+				.load(context -> {
+					ServletRegistrationBean<?> registrationBean = context
+							.getBean(ServletRegistrationBean.class);
+					assertThat(ReflectionTestUtils.getField(registrationBean,
+							"loadOnStartup")).isEqualTo(1);
+				});
 	}
 
 	@Test
 	public void customInitParameters() {
-		load(WebServicesAutoConfiguration.class,
-				"spring.webservices.servlet.init.key1=value1",
-				"spring.webservices.servlet.init.key2=value2");
-		ServletRegistrationBean<?> registrationBean = this.context
-				.getBean(ServletRegistrationBean.class);
-		assertThat(registrationBean.getInitParameters()).containsEntry("key1", "value1");
-		assertThat(registrationBean.getInitParameters()).containsEntry("key2", "value2");
-	}
-
-	private void load(Class<?> config, String... environment) {
-		AnnotationConfigWebApplicationContext context = new AnnotationConfigWebApplicationContext();
-		context.setServletContext(new MockServletContext());
-		TestPropertyValues.of(environment).applyTo(context);
-		context.register(config);
-		context.refresh();
-		this.context = context;
+		this.contextLoader.env("spring.webservices.servlet.init.key1=value1",
+				"spring.webservices.servlet.init.key2=value2").load(context -> {
+					ServletRegistrationBean<?> registrationBean = context
+							.getBean(ServletRegistrationBean.class);
+					assertThat(registrationBean.getInitParameters()).containsEntry("key1",
+							"value1");
+					assertThat(registrationBean.getInitParameters()).containsEntry("key2",
+							"value2");
+				});
 	}
 
 }

spring-boot-test-autoconfigure/src/test/java/org/springframework/boot/test/autoconfigure/jdbc/TestDatabaseAutoConfigurationTests.java

@@ -18,12 +18,9 @@ package org.springframework.boot.test.autoconfigure.jdbc;
 
 import javax.sql.DataSource;
 
-import org.junit.After;
 import org.junit.Test;
 
-import org.springframework.boot.test.util.TestPropertyValues;
-import org.springframework.context.ConfigurableApplicationContext;
-import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.boot.test.context.ContextLoader;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.jdbc.core.JdbcTemplate;
@@ -36,52 +33,33 @@ import static org.assertj.core.api.Assertions.assertThat;
  * Tests for {@link TestDatabaseAutoConfiguration}.
  *
  * @author Stephane Nicoll
+ * @author Andy Wilkinson
  */
 public class TestDatabaseAutoConfigurationTests {
 
-	private ConfigurableApplicationContext context;
-
-	@After
-	public void closeContext() {
-		if (this.context != null) {
-			this.context.close();
-		}
-	}
+	private final ContextLoader contextLoader = new ContextLoader()
+			.autoConfig(TestDatabaseAutoConfiguration.class);
 
 	@Test
 	public void replaceWithNoDataSourceAvailable() {
-		load(null);
-		assertThat(this.context.getBeansOfType(DataSource.class)).isEmpty();
+		this.contextLoader.load(context -> {
+			assertThat(context.getBeansOfType(DataSource.class)).isEmpty();
+		});
 	}
 
 	@Test
 	public void replaceWithUniqueDatabase() {
-		load(ExistingDataSourceConfiguration.class);
-		DataSource datasource = this.context.getBean(DataSource.class);
-		JdbcTemplate jdbcTemplate = new JdbcTemplate(datasource);
-		jdbcTemplate.execute("create table example (id int, name varchar);");
-		try (ConfigurableApplicationContext anotherContext = doLoad(
-				ExistingDataSourceConfiguration.class)) {
-			DataSource anotherDatasource = anotherContext.getBean(DataSource.class);
-			JdbcTemplate anotherJdbcTemplate = new JdbcTemplate(anotherDatasource);
-			anotherJdbcTemplate.execute("create table example (id int, name varchar);");
-		}
-	}
-
-	private void load(Class<?> config, String... environment) {
-		this.context = doLoad(config, environment);
-	}
-
-	private ConfigurableApplicationContext doLoad(Class<?> config,
-			String... environment) {
-		AnnotationConfigApplicationContext ctx = new AnnotationConfigApplicationContext();
-		if (config != null) {
-			ctx.register(config);
-		}
-		ctx.register(TestDatabaseAutoConfiguration.class);
-		TestPropertyValues.of(environment).applyTo(ctx);
-		ctx.refresh();
-		return ctx;
+		this.contextLoader.config(ExistingDataSourceConfiguration.class).load(context -> {
+			DataSource datasource = context.getBean(DataSource.class);
+			JdbcTemplate jdbcTemplate = new JdbcTemplate(datasource);
+			jdbcTemplate.execute("create table example (id int, name varchar);");
+			this.contextLoader.load(anotherContext -> {
+				DataSource anotherDatasource = anotherContext.getBean(DataSource.class);
+				JdbcTemplate anotherJdbcTemplate = new JdbcTemplate(anotherDatasource);
+				anotherJdbcTemplate
+						.execute("create table example (id int, name varchar);");
+			});
+		});
 	}
 
 	@Configuration

spring-boot-test-autoconfigure/src/test/java/org/springframework/boot/test/autoconfigure/orm/jpa/TestDatabaseAutoConfigurationNoEmbeddedTests.java

@@ -18,17 +18,14 @@ package org.springframework.boot.test.autoconfigure.orm.jpa;
 
 import javax.sql.DataSource;
 
-import org.junit.After;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
 import org.springframework.beans.factory.BeanCreationException;
 import org.springframework.boot.test.autoconfigure.jdbc.TestDatabaseAutoConfiguration;
-import org.springframework.boot.test.util.TestPropertyValues;
+import org.springframework.boot.test.context.ContextLoader;
 import org.springframework.boot.testsupport.runner.classpath.ClassPathExclusions;
 import org.springframework.boot.testsupport.runner.classpath.ModifiedClassPathRunner;
-import org.springframework.context.ConfigurableApplicationContext;
-import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 
@@ -40,26 +37,19 @@ import static org.mockito.Mockito.mock;
  * available.
  *
  * @author Stephane Nicoll
+ * @author Andy Wilkinson
  */
 @RunWith(ModifiedClassPathRunner.class)
 @ClassPathExclusions({ "h2-*.jar", "hsqldb-*.jar", "derby-*.jar" })
 public class TestDatabaseAutoConfigurationNoEmbeddedTests {
 
-	private ConfigurableApplicationContext context;
-
-	@After
-	public void closeContext() {
-		if (this.context != null) {
-			this.context.close();
-		}
-	}
+	private final ContextLoader contextLoader = new ContextLoader()
+			.config(ExistingDataSourceConfiguration.class)
+			.autoConfig(TestDatabaseAutoConfiguration.class);
 
 	@Test
 	public void applyAnyReplace() {
-		try {
-			load(ExistingDataSourceConfiguration.class);
-		}
-		catch (BeanCreationException ex) {
+		this.contextLoader.loadAndFail(BeanCreationException.class, ex -> {
 			String message = ex.getMessage();
 			assertThat(message).contains(
 					"Failed to replace DataSource with an embedded database for tests.");
@@ -68,30 +58,16 @@ public class TestDatabaseAutoConfigurationNoEmbeddedTests {
 							+ "classpath");
 			assertThat(message).contains(
 					"or tune the replace attribute of @AutoconfigureTestDatabase.");
-		}
+		});
 	}
 
 	@Test
 	public void applyNoReplace() {
-		load(ExistingDataSourceConfiguration.class, "spring.test.database.replace=NONE");
-		assertThat(this.context.getBeansOfType(DataSource.class)).hasSize(1);
-		assertThat(this.context.getBean(DataSource.class))
-				.isSameAs(this.context.getBean("myCustomDataSource"));
-	}
-
-	public void load(Class<?> config, String... environment) {
-		this.context = doLoad(config, environment);
-	}
-
-	public ConfigurableApplicationContext doLoad(Class<?> config, String... environment) {
-		AnnotationConfigApplicationContext ctx = new AnnotationConfigApplicationContext();
-		if (config != null) {
-			ctx.register(config);
-		}
-		ctx.register(TestDatabaseAutoConfiguration.class);
-		TestPropertyValues.of(environment).applyTo(ctx);
-		ctx.refresh();
-		return ctx;
+		this.contextLoader.env("spring.test.database.replace=NONE").load(context -> {
+			assertThat(context.getBeansOfType(DataSource.class)).hasSize(1);
+			assertThat(context.getBean(DataSource.class))
+					.isSameAs(context.getBean("myCustomDataSource"));
+		});
 	}
 
 	@Configuration

spring-boot-test/src/main/java/org/springframework/boot/test/context/ContextLoader.java

@@ -26,13 +26,19 @@ import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 import org.springframework.boot.test.util.TestPropertyValues;
+import org.springframework.boot.web.reactive.context.GenericReactiveWebApplicationContext;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.context.annotation.AnnotationConfigRegistry;
+import org.springframework.core.io.DefaultResourceLoader;
+import org.springframework.mock.web.MockServletContext;
 import org.springframework.util.Assert;
 import org.springframework.util.ObjectUtils;
+import org.springframework.web.context.support.AnnotationConfigWebApplicationContext;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -96,6 +102,8 @@ public class ContextLoader {
 
 	private final LinkedList<Class<?>> autoConfigurations = new LinkedList<>();
 
+	private Supplier<ConfigurableApplicationContext> contextSupplier = () -> new AnnotationConfigApplicationContext();
+
 	private ClassLoader classLoader;
 
 	/**
@@ -182,6 +190,32 @@ public class ContextLoader {
 		return this;
 	}
 
+	/**
+	 * Configures the loader to create an {@link ApplicationContext} suitable for use in a
+	 * reactive web application.
+	 * @return this instance
+	 */
+	public ContextLoader webReactive() {
+		this.contextSupplier = () -> {
+			return new GenericReactiveWebApplicationContext();
+		};
+		return this;
+	}
+
+	/**
+	 * Configures the loader to create an {@link ApplicationContext} suitable for use in a
+	 * servlet web application.
+	 * @return this instance
+	 */
+	public ContextLoader webServlet() {
+		this.contextSupplier = () -> {
+			AnnotationConfigWebApplicationContext context = new AnnotationConfigWebApplicationContext();
+			context.setServletContext(new MockServletContext());
+			return context;
+		};
+		return this;
+	}
+
 	/**
 	 * Create and refresh a new {@link ApplicationContext} based on the current state of
 	 * this loader. The context is consumed by the specified {@link ContextConsumer} and
@@ -239,25 +273,26 @@ public class ContextLoader {
 	}
 
 	private ConfigurableApplicationContext createApplicationContext() {
-		AnnotationConfigApplicationContext ctx = new AnnotationConfigApplicationContext();
+		ConfigurableApplicationContext context = ContextLoader.this.contextSupplier.get();
 		if (this.classLoader != null) {
-			ctx.setClassLoader(this.classLoader);
+			((DefaultResourceLoader) context).setClassLoader(this.classLoader);
 		}
 		if (!ObjectUtils.isEmpty(this.env)) {
 			TestPropertyValues.of(this.env.toArray(new String[this.env.size()]))
-					.applyTo(ctx);
+					.applyTo(context);
 		}
+		AnnotationConfigRegistry registry = ((AnnotationConfigRegistry) context);
 		if (!ObjectUtils.isEmpty(this.userConfigurations)) {
-			ctx.register(this.userConfigurations
+			registry.register(this.userConfigurations
 					.toArray(new Class<?>[this.userConfigurations.size()]));
 		}
 		if (!ObjectUtils.isEmpty(this.autoConfigurations)) {
 			LinkedHashSet<Class<?>> linkedHashSet = new LinkedHashSet<>(
 					this.autoConfigurations);
-			ctx.register(
+			registry.register(
 					linkedHashSet.toArray(new Class<?>[this.autoConfigurations.size()]));
 		}
-		return ctx;
+		return context;
 	}
 
 	/**