Support for @ControllerAdvice in WebFlux

Issue: SPR-15132

spring-test/src/main/java/org/springframework/test/web/reactive/server/DefaultControllerSpec.java

@@ -15,6 +15,7 @@
  */
 package org.springframework.test.web.reactive.server;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Optional;
@@ -44,6 +45,8 @@ class DefaultControllerSpec implements WebTestClient.ControllerSpec {
 
 	private final List<Object> controllers;
 
+	private final List<Object> controllerAdvice = new ArrayList<>(8);
+
 	private final TestWebFluxConfigurer configurer = new TestWebFluxConfigurer();
 
 
@@ -53,6 +56,12 @@ class DefaultControllerSpec implements WebTestClient.ControllerSpec {
 	}
 
 
+	@Override
+	public DefaultControllerSpec controllerAdvice(Object... controllerAdvice) {
+		this.controllerAdvice.addAll(Arrays.asList(controllerAdvice));
+		return this;
+	}
+
 	@Override
 	public DefaultControllerSpec contentTypeResolver(Consumer<RequestedContentTypeResolverBuilder> consumer) {
 		this.configurer.contentTypeResolverConsumer = consumer;
@@ -103,12 +112,17 @@ class DefaultControllerSpec implements WebTestClient.ControllerSpec {
 
 	@Override
 	public WebTestClient.Builder configureClient() {
+		return WebTestClient.bindToApplicationContext(createApplicationContext());
+	}
+
+	protected AnnotationConfigApplicationContext createApplicationContext() {
 		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
 		this.controllers.forEach(controller -> registerBean(context, controller));
+		this.controllerAdvice.forEach(advice -> registerBean(context, advice));
 		context.register(DelegatingWebFluxConfiguration.class);
 		context.registerBean(WebFluxConfigurer.class, () -> this.configurer);
 		context.refresh();
-		return WebTestClient.bindToApplicationContext(context);
+		return context;
 	}
 
 	@SuppressWarnings("unchecked")

spring-test/src/main/java/org/springframework/test/web/reactive/server/WebTestClient.java

@@ -179,6 +179,13 @@ public interface WebTestClient {
 	 */
 	interface ControllerSpec {
 
+		/**
+		 * Register one or more
+		 * {@link org.springframework.web.bind.annotation.ControllerAdvice
+		 * ControllerAdvice} instances to be used in tests.
+		 */
+		ControllerSpec controllerAdvice(Object... controllerAdvice);
+
 		/**
 		 * Customize content type resolution.
 		 * @see WebFluxConfigurer#configureContentTypeResolver

spring-test/src/test/java/org/springframework/test/web/reactive/server/DefaultControllerSpecTests.java

@@ -0,0 +1,73 @@
+/*
+ * Copyright 2002-2017 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
+ *
+ *      http://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.test.web.reactive.server;
+
+import org.junit.Test;
+
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+
+import static org.junit.Assert.assertSame;
+
+/**
+ * Unit tests for {@link DefaultControllerSpec}.
+ * @author Rossen Stoyanchev
+ */
+public class DefaultControllerSpecTests {
+
+	@Test
+	public void controllers() throws Exception {
+		OneController controller1 = new OneController();
+		SecondController controller2 = new SecondController();
+
+		TestControllerSpec spec = new TestControllerSpec(controller1, controller2);
+		ApplicationContext context = spec.createApplicationContext();
+
+		assertSame(controller1, context.getBean(OneController.class));
+		assertSame(controller2, context.getBean(SecondController.class));
+	}
+
+	@Test
+	public void controllerAdvice() throws Exception {
+		OneControllerAdvice advice = new OneControllerAdvice();
+
+		TestControllerSpec spec = new TestControllerSpec(new OneController());
+		spec.controllerAdvice(advice);
+		ApplicationContext context = spec.createApplicationContext();
+
+		assertSame(advice, context.getBean(OneControllerAdvice.class));
+	}
+
+	private static class OneController {}
+
+	private static class SecondController {}
+
+	private static class OneControllerAdvice {}
+
+
+	private static class TestControllerSpec extends DefaultControllerSpec {
+
+		TestControllerSpec(Object... controllers) {
+			super(controllers);
+		}
+
+		@Override
+		public AnnotationConfigApplicationContext createApplicationContext() {
+			return super.createApplicationContext();
+		}
+	}
+
+}