Support for mutating ServerWebExchange in WebTestClient

This commit adds a common base class for server-less setup with the
option to configure a transformation function on the
ServerWebExchange for every request.

The transformation is applied through a WebFilter. As a result the
RouterFunction setup is now invoked behind a DispatcherHandler with
a HandlerMapping + HandlerAdapter.

Issue: SPR-15250

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

@@ -0,0 +1,69 @@
+/*
+ * 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 java.util.function.Function;
+
+import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.server.adapter.WebHttpHandlerBuilder;
+
+/**
+ * Base class for implementations of {@link WebTestClient.MockServerSpec}.
+ *
+ * @author Rossen Stoyanchev
+ * @since 5.0
+ */
+abstract class AbstractMockServerSpec<B extends WebTestClient.MockServerSpec<B>>
+		implements WebTestClient.MockServerSpec<B> {
+
+	private Function<ServerWebExchange, ServerWebExchange> exchangeMutator;
+
+
+	@Override
+	public <T extends B> T exchangeMutator(Function<ServerWebExchange, ServerWebExchange> mutator) {
+		this.exchangeMutator = this.exchangeMutator != null ? this.exchangeMutator.andThen(mutator) : mutator;
+		return self();
+	}
+
+	@SuppressWarnings("unchecked")
+	protected <T extends B> T self() {
+		return (T) this;
+	}
+
+
+	@Override
+	public WebTestClient.Builder configureClient() {
+
+		WebHttpHandlerBuilder handlerBuilder = createHttpHandlerBuilder();
+
+		if (this.exchangeMutator != null) {
+			handlerBuilder.prependFilter((exchange, chain) -> {
+				exchange = this.exchangeMutator.apply(exchange);
+				return chain.filter(exchange);
+			});
+		}
+
+		return new DefaultWebTestClientBuilder(handlerBuilder.build());
+	}
+
+	protected abstract WebHttpHandlerBuilder createHttpHandlerBuilder();
+
+	@Override
+	public WebTestClient build() {
+		return configureClient().build();
+	}
+
+}

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

@@ -0,0 +1,45 @@
+/*
+ * 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.springframework.context.ApplicationContext;
+import org.springframework.util.Assert;
+import org.springframework.web.server.adapter.WebHttpHandlerBuilder;
+
+/**
+ * Spec for setting up server-less testing by detecting components in an
+ * {@link ApplicationContext}.
+ *
+ * @author Rossen Stoyanchev
+ * @since 5.0
+ */
+class ApplicationContextSpec extends AbstractMockServerSpec<ApplicationContextSpec> {
+
+	private final ApplicationContext applicationContext;
+
+
+	ApplicationContextSpec(ApplicationContext applicationContext) {
+		Assert.notNull(applicationContext, "ApplicationContext is required");
+		this.applicationContext = applicationContext;
+	}
+
+
+	@Override
+	protected WebHttpHandlerBuilder createHttpHandlerBuilder() {
+		return WebHttpHandlerBuilder.applicationContext(this.applicationContext);
+	}
+
+}

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

@@ -21,6 +21,7 @@ import java.util.List;
 import java.util.Optional;
 import java.util.function.Consumer;
 
+import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.format.FormatterRegistry;
 import org.springframework.http.codec.HttpMessageReader;
@@ -34,6 +35,7 @@ import org.springframework.web.reactive.config.DelegatingWebFluxConfiguration;
 import org.springframework.web.reactive.config.PathMatchConfigurer;
 import org.springframework.web.reactive.config.ViewResolverRegistry;
 import org.springframework.web.reactive.config.WebFluxConfigurer;
+import org.springframework.web.server.adapter.WebHttpHandlerBuilder;
 
 /**
  * Default implementation of {@link WebTestClient.ControllerSpec}.
@@ -41,7 +43,8 @@ import org.springframework.web.reactive.config.WebFluxConfigurer;
  * @author Rossen Stoyanchev
  * @since 5.0
  */
-class DefaultControllerSpec implements WebTestClient.ControllerSpec {
+class DefaultControllerSpec extends AbstractMockServerSpec<WebTestClient.ControllerSpec>
+		implements WebTestClient.ControllerSpec {
 
 	private final List<Object> controllers;
 
@@ -50,7 +53,7 @@ class DefaultControllerSpec implements WebTestClient.ControllerSpec {
 	private final TestWebFluxConfigurer configurer = new TestWebFluxConfigurer();
 
 
-	public DefaultControllerSpec(Object... controllers) {
+	DefaultControllerSpec(Object... controllers) {
 		Assert.isTrue(!ObjectUtils.isEmpty(controllers), "At least one controller is required");
 		this.controllers = Arrays.asList(controllers);
 	}
@@ -110,31 +113,28 @@ class DefaultControllerSpec implements WebTestClient.ControllerSpec {
 		return this;
 	}
 
+
 	@Override
-	public WebTestClient.Builder configureClient() {
-		return WebTestClient.bindToApplicationContext(createApplicationContext());
+	protected WebHttpHandlerBuilder createHttpHandlerBuilder() {
+		return WebHttpHandlerBuilder.applicationContext(initApplicationContext());
 	}
 
-	protected AnnotationConfigApplicationContext createApplicationContext() {
+	private ApplicationContext initApplicationContext() {
 		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
-		this.controllers.forEach(controller -> registerBean(context, controller));
-		this.controllerAdvice.forEach(advice -> registerBean(context, advice));
+		this.controllers.forEach(controller -> {
+			String name = controller.getClass().getName();
+			context.registerBean(name, Object.class, () -> controller);
+		});
+		this.controllerAdvice.forEach(advice -> {
+			String name = advice.getClass().getName();
+			context.registerBean(name, Object.class, () -> advice);
+		});
 		context.register(DelegatingWebFluxConfiguration.class);
 		context.registerBean(WebFluxConfigurer.class, () -> this.configurer);
 		context.refresh();
 		return context;
 	}
 
-	@SuppressWarnings("unchecked")
-	private <T> void registerBean(AnnotationConfigApplicationContext context, T bean) {
-		context.registerBean((Class<T>) bean.getClass(), () -> bean);
-	}
-
-	@Override
-	public WebTestClient build() {
-		return configureClient().build();
-	}
-
 
 	private class TestWebFluxConfigurer implements WebFluxConfigurer {
 

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

@@ -152,7 +152,7 @@ public class ExchangeResult {
 			assertion.run();
 		}
 		catch (AssertionError ex) {
-			throw new AssertionError("Assertion failed on the following exchange:" + this, ex);
+			throw new AssertionError(ex.getMessage() + "\n" + this, ex);
 		}
 	}
 
@@ -168,7 +168,7 @@ public class ExchangeResult {
 				"< " + getStatus() + " " + getStatusReason() + "\n" +
 				"< " + formatHeaders(getResponseHeaders(), "\n< ") + "\n" +
 				"\n" +
-				formatBody(getResponseHeaders().getContentType(), this.response.getRecordedContent()) + "\n\n";
+				formatBody(getResponseHeaders().getContentType(), this.response.getRecordedContent()) +"\n";
 	}
 
 	private String getStatusReason() {

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

@@ -0,0 +1,62 @@
+/*
+ * 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.springframework.context.ApplicationContext;
+import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.web.reactive.DispatcherHandler;
+import org.springframework.web.reactive.HandlerAdapter;
+import org.springframework.web.reactive.HandlerMapping;
+import org.springframework.web.reactive.HandlerResultHandler;
+import org.springframework.web.reactive.function.server.RouterFunction;
+import org.springframework.web.reactive.function.server.RouterFunctions;
+import org.springframework.web.reactive.function.server.support.HandlerFunctionAdapter;
+import org.springframework.web.reactive.function.server.support.ServerResponseResultHandler;
+import org.springframework.web.server.adapter.WebHttpHandlerBuilder;
+
+/**
+ * Spec for setting up server-less testing against a RouterFunction.
+ *
+ * @author Rossen Stoyanchev
+ * @since 5.0
+ */
+public class RouterFunctionSpec extends AbstractMockServerSpec<RouterFunctionSpec> {
+
+	private final RouterFunction<?> routerFunction;
+
+
+	RouterFunctionSpec(RouterFunction<?> routerFunction) {
+		this.routerFunction = routerFunction;
+	}
+
+
+	@Override
+	protected WebHttpHandlerBuilder createHttpHandlerBuilder() {
+		return WebHttpHandlerBuilder.applicationContext(initApplicationContext());
+	}
+
+	@SuppressWarnings("Convert2MethodRef")
+	private ApplicationContext initApplicationContext() {
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
+		context.registerBean("webHandler", DispatcherHandler.class, () -> new DispatcherHandler());
+		context.registerBean(HandlerMapping.class, () -> RouterFunctions.toHandlerMapping(this.routerFunction));
+		context.registerBean(HandlerAdapter.class, () -> new HandlerFunctionAdapter());
+		context.registerBean(HandlerResultHandler.class, () -> new ServerResponseResultHandler());
+		context.refresh();
+		return context;
+	}
+
+}

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

@@ -34,7 +34,6 @@ import org.springframework.http.MediaType;
 import org.springframework.http.client.reactive.ClientHttpRequest;
 import org.springframework.http.codec.HttpMessageReader;
 import org.springframework.http.codec.HttpMessageWriter;
-import org.springframework.http.server.reactive.HttpHandler;
 import org.springframework.util.MultiValueMap;
 import org.springframework.validation.Validator;
 import org.springframework.web.reactive.accept.RequestedContentTypeResolverBuilder;
@@ -48,9 +47,7 @@ import org.springframework.web.reactive.function.client.ExchangeFunction;
 import org.springframework.web.reactive.function.client.ExchangeStrategies;
 import org.springframework.web.reactive.function.client.WebClient;
 import org.springframework.web.reactive.function.server.RouterFunction;
-import org.springframework.web.reactive.function.server.RouterFunctions;
-import org.springframework.web.server.adapter.HttpWebHandlerAdapter;
-import org.springframework.web.server.adapter.WebHttpHandlerBuilder;
+import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.util.UriBuilder;
 import org.springframework.web.util.UriBuilderFactory;
 
@@ -149,9 +146,8 @@ public interface WebTestClient {
 	 * @return the {@link WebTestClient} builder
 	 * @see org.springframework.web.reactive.config.EnableWebFlux
 	 */
-	static Builder bindToApplicationContext(ApplicationContext applicationContext) {
-		HttpHandler httpHandler = WebHttpHandlerBuilder.applicationContext(applicationContext).build();
-		return new DefaultWebTestClientBuilder(httpHandler);
+	static MockServerSpec<?> bindToApplicationContext(ApplicationContext applicationContext) {
+		return new ApplicationContextSpec(applicationContext);
 	}
 
 	/**
@@ -159,9 +155,8 @@ public interface WebTestClient {
 	 * @param routerFunction the RouterFunction to test
 	 * @return the {@link WebTestClient} builder
 	 */
-	static Builder bindToRouterFunction(RouterFunction<?> routerFunction) {
-		HttpWebHandlerAdapter httpHandler = RouterFunctions.toHttpHandler(routerFunction);
-		return new DefaultWebTestClientBuilder(httpHandler);
+	static MockServerSpec<?> bindToRouterFunction(RouterFunction<?> routerFunction) {
+		return new RouterFunctionSpec(routerFunction);
 	}
 
 	/**
@@ -173,11 +168,36 @@ public interface WebTestClient {
 	}
 
 
+	/**
+	 * Base specification for setting up tests without a server.
+	 */
+	interface MockServerSpec<B extends MockServerSpec<B>> {
+
+		/**
+		 * Configure a transformation function on {@code ServerWebExchange} to
+		 * be applied at the start of server-side, request processing.
+		 * @param function the transforming function.
+		 * @see ServerWebExchange#mutate()
+		 */
+		<T extends B> T exchangeMutator(Function<ServerWebExchange, ServerWebExchange> function);
+
+		/**
+		 * Proceed to configure and build the test client.
+		 */
+		Builder configureClient();
+
+		/**
+		 * Shortcut to build the test client.
+		 */
+		WebTestClient build();
+
+	}
+
 	/**
 	 * Specification for customizing controller configuration equivalent to, and
 	 * internally delegating to, a {@link WebFluxConfigurer}.
 	 */
-	interface ControllerSpec {
+	interface ControllerSpec extends MockServerSpec<ControllerSpec> {
 
 		/**
 		 * Register one or more
@@ -234,16 +254,6 @@ public interface WebTestClient {
 		 */
 		ControllerSpec viewResolvers(Consumer<ViewResolverRegistry> consumer);
 
-		/**
-		 * Proceed to configure and build the test client.
-		 */
-		Builder configureClient();
-
-		/**
-		 * Shortcut to build the test client.
-		 */
-		WebTestClient build();
-
 	}
 
 	/**

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

@@ -17,10 +17,11 @@ 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;
+import org.springframework.http.ResponseEntity;
+import org.springframework.web.bind.annotation.ControllerAdvice;
+import org.springframework.web.bind.annotation.ExceptionHandler;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.RestController;
 
 /**
  * Unit tests for {@link DefaultControllerSpec}.
@@ -29,44 +30,47 @@ import static org.junit.Assert.assertSame;
 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));
+	public void controller() throws Exception {
+		new DefaultControllerSpec(new MyController()).build()
+				.get().uri("/")
+				.exchange()
+				.expectStatus().isOk()
+				.expectBody(String.class).value().isEqualTo("Success");
 	}
 
 	@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));
+		new DefaultControllerSpec(new MyController())
+				.controllerAdvice(new MyControllerAdvice())
+				.build()
+				.get().uri("/exception")
+				.exchange()
+				.expectStatus().isBadRequest()
+				.expectBody(String.class).value().isEqualTo("Handled exception");
 	}
 
-	private static class OneController {}
 
-	private static class SecondController {}
+	@RestController
+	private static class MyController {
 
-	private static class OneControllerAdvice {}
-
-
-	private static class TestControllerSpec extends DefaultControllerSpec {
-
-		TestControllerSpec(Object... controllers) {
-			super(controllers);
+		@GetMapping("/")
+		public String handle() {
+			return "Success";
 		}
 
-		@Override
-		public AnnotationConfigApplicationContext createApplicationContext() {
-			return super.createApplicationContext();
+		@GetMapping("/exception")
+		public void handleWithError() {
+			throw new IllegalStateException();
+		}
+
+	}
+
+	@ControllerAdvice
+	private static class MyControllerAdvice {
+
+		@ExceptionHandler
+		public ResponseEntity<String> handle(IllegalStateException ex) {
+			return ResponseEntity.status(400).body("Handled exception");
 		}
 	}
 

spring-test/src/test/java/org/springframework/test/web/reactive/server/samples/bind/ApplicationContextTests.java

@@ -15,8 +15,12 @@
  */
 package org.springframework.test.web.reactive.server.samples.bind;
 
+import java.security.Principal;
+import java.util.function.Function;
+
 import org.junit.Before;
 import org.junit.Test;
+import reactor.core.publisher.Mono;
 
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.annotation.Bean;
@@ -25,6 +29,10 @@ import org.springframework.test.web.reactive.server.WebTestClient;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.reactive.config.EnableWebFlux;
+import org.springframework.web.server.ServerWebExchange;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 /**
  * Binding to server infrastructure declared in a Spring ApplicationContext.
@@ -44,15 +52,26 @@ public class ApplicationContextTests {
 		context.register(WebConfig.class);
 		context.refresh();
 
-		this.client = WebTestClient.bindToApplicationContext(context).build();
+		this.client = WebTestClient.bindToApplicationContext(context)
+				.exchangeMutator(identityMutator("Pablo"))
+				.build();
 	}
 
+	private Function<ServerWebExchange, ServerWebExchange> identityMutator(String userName) {
+		return exchange -> {
+			Principal user = mock(Principal.class);
+			when(user.getName()).thenReturn(userName);
+			return exchange.mutate().principal(Mono.just(user)).build();
+		};
+	}
+
+
 	@Test
 	public void test() throws Exception {
 		this.client.get().uri("/test")
 				.exchange()
 				.expectStatus().isOk()
-				.expectBody(String.class).value().isEqualTo("It works!");
+				.expectBody(String.class).value().isEqualTo("Hello Pablo!");
 	}
 
 
@@ -71,8 +90,8 @@ public class ApplicationContextTests {
 	static class TestController {
 
 		@GetMapping("/test")
-		public String handle() {
-			return "It works!";
+		public String handle(Principal principal) {
+			return "Hello " + principal.getName() + "!";
 		}
 	}
 

spring-test/src/test/java/org/springframework/test/web/reactive/server/samples/bind/ControllerTests.java

@@ -15,12 +15,20 @@
  */
 package org.springframework.test.web.reactive.server.samples.bind;
 
+import java.security.Principal;
+import java.util.function.Function;
+
 import org.junit.Before;
 import org.junit.Test;
+import reactor.core.publisher.Mono;
 
 import org.springframework.test.web.reactive.server.WebTestClient;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.RestController;
+import org.springframework.web.server.ServerWebExchange;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 /**
  * Bind to annotated controllers.
@@ -35,7 +43,18 @@ public class ControllerTests {
 
 	@Before
 	public void setUp() throws Exception {
-		this.client = WebTestClient.bindToController(new TestController()).build();
+
+		this.client = WebTestClient.bindToController(new TestController())
+				.exchangeMutator(identityMutator("Pablo"))
+				.build();
+	}
+
+	private Function<ServerWebExchange, ServerWebExchange> identityMutator(String userName) {
+		return exchange -> {
+			Principal user = mock(Principal.class);
+			when(user.getName()).thenReturn(userName);
+			return exchange.mutate().principal(Mono.just(user)).build();
+		};
 	}
 
 
@@ -44,7 +63,7 @@ public class ControllerTests {
 		this.client.get().uri("/test")
 				.exchange()
 				.expectStatus().isOk()
-				.expectBody(String.class).value().isEqualTo("It works!");
+				.expectBody(String.class).value().isEqualTo("Hello Pablo!");
 	}
 
 
@@ -52,8 +71,9 @@ public class ControllerTests {
 	static class TestController {
 
 		@GetMapping("/test")
-		public String handle() {
-			return "It works!";
+		public String handle(Principal principal) {
+			return "Hello " + principal.getName() + "!";
 		}
 	}
+
 }