Refactor WebTestClient exchange mutator support

This commit factors ServerWebExchange mutator support out of
WebTestClient in favor of an independent public class called
MockServerExchangeMutator which implements WebFilter and can be
applied to the WebTestClient as any other WebFilter.

The MockServerExchangeMutator also exposes a method to apply
a client-side filter for "per request" mutators. See the Javadoc
of the MockServerExchangeMutator.

Issue: SPR-15570

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

@@ -20,9 +20,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
-import java.util.function.UnaryOperator;
 
-import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebFilter;
 import org.springframework.web.server.adapter.WebHttpHandlerBuilder;
 
@@ -35,22 +33,9 @@ import org.springframework.web.server.adapter.WebHttpHandlerBuilder;
 abstract class AbstractMockServerSpec<B extends WebTestClient.MockServerSpec<B>>
 		implements WebTestClient.MockServerSpec<B> {
 
-	private final ExchangeMutatingWebFilter exchangeMutatingWebFilter = new ExchangeMutatingWebFilter();
-
 	private final List<WebFilter> filters = new ArrayList<>(4);
 
 
-	AbstractMockServerSpec() {
-		this.filters.add(this.exchangeMutatingWebFilter);
-	}
-
-
-	@Override
-	public <T extends B> T exchangeMutator(UnaryOperator<ServerWebExchange> mutator) {
-		this.exchangeMutatingWebFilter.registerGlobalMutator(mutator);
-		return self();
-	}
-
 	@Override
 	public <T extends B> T webFilter(WebFilter... filter) {
 		this.filters.addAll(Arrays.asList(filter));
@@ -67,7 +52,7 @@ abstract class AbstractMockServerSpec<B extends WebTestClient.MockServerSpec<B>>
 	public WebTestClient.Builder configureClient() {
 		WebHttpHandlerBuilder builder = initHttpHandlerBuilder();
 		filtersInReverse().forEach(builder::prependFilter);
-		return new DefaultWebTestClientBuilder(builder.build(), this.exchangeMutatingWebFilter);
+		return new DefaultWebTestClientBuilder(builder.build());
 	}
 
 	/**

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

@@ -27,7 +27,6 @@ import java.util.Optional;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Consumer;
 import java.util.function.Function;
-import java.util.function.UnaryOperator;
 
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
@@ -48,7 +47,6 @@ import org.springframework.web.reactive.function.BodyInserter;
 import org.springframework.web.reactive.function.client.ClientResponse;
 import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
 import org.springframework.web.reactive.function.client.WebClient;
-import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.util.UriBuilder;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -69,28 +67,21 @@ class DefaultWebTestClient implements WebTestClient {
 
 	private final WiretapConnector wiretapConnector;
 
-	private final ExchangeMutatingWebFilter exchangeMutatingWebFilter;
-
 	private final Duration timeout;
 
 	private final AtomicLong requestIndex = new AtomicLong();
 
 
-	DefaultWebTestClient(WebClient.Builder webClientBuilder, ClientHttpConnector connector,
-			ExchangeMutatingWebFilter filter, Duration timeout) {
-
-		Assert.notNull(webClientBuilder, "WebClient.Builder is required");
-
+	DefaultWebTestClient(WebClient.Builder clientBuilder, ClientHttpConnector connector, Duration timeout) {
+		Assert.notNull(clientBuilder, "WebClient.Builder is required");
 		this.wiretapConnector = new WiretapConnector(connector);
-		this.webClient = webClientBuilder.clientConnector(this.wiretapConnector).build();
-		this.exchangeMutatingWebFilter = (filter != null ? filter : new ExchangeMutatingWebFilter());
+		this.webClient = clientBuilder.clientConnector(this.wiretapConnector).build();
 		this.timeout = (timeout != null ? timeout : Duration.ofSeconds(5));
 	}
 
 	private DefaultWebTestClient(DefaultWebTestClient webTestClient, ExchangeFilterFunction filter) {
 		this.webClient = webTestClient.webClient.filter(filter);
 		this.wiretapConnector = webTestClient.wiretapConnector;
-		this.exchangeMutatingWebFilter = webTestClient.exchangeMutatingWebFilter;
 		this.timeout = webTestClient.timeout;
 	}
 
@@ -147,20 +138,6 @@ class DefaultWebTestClient implements WebTestClient {
 		return new DefaultWebTestClient(this, filter);
 	}
 
-	@Override
-	public WebTestClient exchangeMutator(UnaryOperator<ServerWebExchange> mutator) {
-
-		Assert.notNull(this.exchangeMutatingWebFilter,
-				"This option is applicable only for tests without an actual running server");
-
-		return filter((request, next) -> {
-			String requestId = request.headers().getFirst(WiretapConnector.REQUEST_ID_HEADER_NAME);
-			Assert.notNull(requestId, "No request-id header");
-			this.exchangeMutatingWebFilter.registerPerRequestMutator(requestId, mutator);
-			return next.exchange(request);
-		});
-	}
-
 
 	@SuppressWarnings("unchecked")
 	private class DefaultUriSpec<S extends RequestHeadersSpec<?>> implements UriSpec<S> {

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

@@ -37,8 +37,6 @@ class DefaultWebTestClientBuilder implements WebTestClient.Builder {
 
 	private final ClientHttpConnector connector;
 
-	private final ExchangeMutatingWebFilter exchangeMutatingWebFilter;
-
 	private Duration responseTimeout;
 
 
@@ -48,12 +46,10 @@ class DefaultWebTestClientBuilder implements WebTestClient.Builder {
 
 	DefaultWebTestClientBuilder(ClientHttpConnector connector) {
 		this.connector = connector;
-		this.exchangeMutatingWebFilter = null;
 	}
 
-	DefaultWebTestClientBuilder(HttpHandler httpHandler, ExchangeMutatingWebFilter exchangeMutatingWebFilter) {
+	DefaultWebTestClientBuilder(HttpHandler httpHandler) {
 		this.connector = new HttpHandlerConnector(httpHandler);
-		this.exchangeMutatingWebFilter = exchangeMutatingWebFilter;
 	}
 
 
@@ -95,8 +91,7 @@ class DefaultWebTestClientBuilder implements WebTestClient.Builder {
 
 	@Override
 	public WebTestClient build() {
-		return new DefaultWebTestClient(this.webClientBuilder, this.connector,
-				this.exchangeMutatingWebFilter, this.responseTimeout);
+		return new DefaultWebTestClient(this.webClientBuilder, this.connector, this.responseTimeout);
 	}
 
 }

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

@@ -1,82 +0,0 @@
-/*
- * 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.Map;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.function.Function;
-import java.util.function.UnaryOperator;
-
-import reactor.core.publisher.Mono;
-
-import org.springframework.util.Assert;
-import org.springframework.web.server.ServerWebExchange;
-import org.springframework.web.server.WebFilter;
-import org.springframework.web.server.WebFilterChain;
-
-/**
- * WebFilter for applying global and per-request transformations to a
- * {@link ServerWebExchange}.
- *
- * @author Rossen Stoyanchev
- * @since 5.0
- */
-class ExchangeMutatingWebFilter implements WebFilter {
-
-	private static final Function<ServerWebExchange, ServerWebExchange> NO_OP_MUTATOR = e -> e;
-
-
-	private volatile Function<ServerWebExchange, ServerWebExchange> globalMutator = NO_OP_MUTATOR;
-
-	private final Map<String, Function<ServerWebExchange, ServerWebExchange>> perRequestMutators =
-			new ConcurrentHashMap<>(4);
-
-
-	/**
-	 * Register a global transformation function to apply to all requests.
-	 * @param mutator the transformation function
-	 */
-	public void registerGlobalMutator(UnaryOperator<ServerWebExchange> mutator) {
-		Assert.notNull(mutator, "'mutator' is required");
-		this.globalMutator = this.globalMutator.andThen(mutator);
-	}
-
-	/**
-	 * Register a per-request transformation function.
-	 * @param requestId the "request-id" header value identifying the request
-	 * @param mutator the transformation function
-	 */
-	public void registerPerRequestMutator(String requestId, UnaryOperator<ServerWebExchange> mutator) {
-		this.perRequestMutators.compute(requestId,
-				(s, value) -> value != null ? value.andThen(mutator) : mutator);
-	}
-
-
-	@Override
-	public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
-		exchange = this.globalMutator.apply(exchange);
-		exchange = getMutatorFor(exchange).apply(exchange);
-		return chain.filter(exchange);
-	}
-
-	private Function<ServerWebExchange, ServerWebExchange> getMutatorFor(ServerWebExchange exchange) {
-		String id = WiretapConnector.getRequestIdHeader(exchange.getRequest().getHeaders());
-		Function<ServerWebExchange, ServerWebExchange> mutator = this.perRequestMutators.remove(id);
-		return mutator != null ? mutator : NO_OP_MUTATOR;
-	}
-
-}

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

@@ -0,0 +1,119 @@
+/*
+ * 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.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+import java.util.function.UnaryOperator;
+
+import reactor.core.publisher.Mono;
+
+import org.springframework.util.Assert;
+import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.server.WebFilter;
+import org.springframework.web.server.WebFilterChain;
+
+/**
+ * Built-in {@link WebFilter} for applying {@code ServerWebExchange}
+ * transformations during requests from the {@code WebTestClient} to a mock
+ * server -- i.e. when one of the following is in use:
+ * <ul>
+ * <li>{@link WebTestClient#bindToController},
+ * <li>{@link WebTestClient#bindToRouterFunction}
+ * <li>{@link WebTestClient#bindToApplicationContext}.
+ * </ul>
+ *
+ * <p>Example of registering a "global" transformation:
+ * <pre class="code">
+ *
+ * MockServerExchangeMutator mutator = new MockServerExchangeMutator(exchange -> ...);
+ * WebTestClient client = WebTestClient.bindToController(new MyController()).webFilter(mutator).build()
+ * </pre>
+ *
+ * <p>Example of registering "per client" transformations:
+ * <pre class="code">
+ *
+ * MockServerExchangeMutator mutator = new MockServerExchangeMutator(exchange -> ...);
+ * WebTestClient client = WebTestClient.bindToController(new MyController()).webFilter(mutator).build()
+ *
+ * WebTestClient clientA = mutator.filterClient(client, exchange -> ...);
+ * // Use client A...
+ *
+ * WebTestClient clientB = mutator.filterClient(client, exchange -> ...);
+ * // Use client B...
+ * </pre>
+ *
+ * @author Rossen Stoyanchev
+ * @since 5.0
+ */
+public class MockServerExchangeMutator implements WebFilter {
+
+	private final Function<ServerWebExchange, ServerWebExchange> mutator;
+
+	private final Map<String, Function<ServerWebExchange, ServerWebExchange>> perRequestMutators =
+			new ConcurrentHashMap<>(4);
+
+
+	public MockServerExchangeMutator(Function<ServerWebExchange, ServerWebExchange> mutator) {
+		Assert.notNull(mutator, "'mutator' is required");
+		this.mutator = mutator;
+	}
+
+
+	@Override
+	public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
+		return chain.filter(getMutatorsFor(exchange).apply(exchange));
+	}
+
+	private Function<ServerWebExchange, ServerWebExchange> getMutatorsFor(ServerWebExchange exchange) {
+		String id = WiretapConnector.getRequestIdHeader(exchange.getRequest().getHeaders());
+		Function<ServerWebExchange, ServerWebExchange> m = this.perRequestMutators.remove(id);
+		return (m != null ? this.mutator.andThen(m) : this.mutator);
+	}
+
+
+	/**
+	 * Apply a filter to the given client in order to apply
+	 * {@code ServerWebExchange} transformations only to requests executed
+	 * through the returned client instance. See examples in the
+	 * {@link MockServerExchangeMutator class-level Javadoc}.
+	 *
+	 * @param mutator the per-request mutator to use
+	 * @param mutators additional per-request mutators to use
+	 * @return a new client instance filtered with {@link WebTestClient#filter}
+	 */
+	@SafeVarargs
+	public final WebTestClient filterClient(WebTestClient client,
+			UnaryOperator<ServerWebExchange> mutator, UnaryOperator<ServerWebExchange>... mutators) {
+
+		return client.filter((request, next) -> {
+			String id = request.headers().getFirst(WiretapConnector.REQUEST_ID_HEADER_NAME);
+			Assert.notNull(id, "No request-id header");
+			registerPerRequestMutator(id, mutator);
+			for (UnaryOperator<ServerWebExchange> current : mutators) {
+				registerPerRequestMutator(id, current);
+			}
+			return next.exchange(request);
+		});
+	}
+
+	private void registerPerRequestMutator(String id, UnaryOperator<ServerWebExchange> m) {
+		this.perRequestMutators.compute(id, (s, value) -> value != null ? value.andThen(m) : m);
+	}
+
+}

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

@@ -24,7 +24,6 @@ import java.util.List;
 import java.util.Map;
 import java.util.function.Consumer;
 import java.util.function.Function;
-import java.util.function.UnaryOperator;
 
 import org.reactivestreams.Publisher;
 
@@ -51,7 +50,6 @@ import org.springframework.web.reactive.function.client.WebClient;
 import org.springframework.web.reactive.function.server.HandlerStrategies;
 import org.springframework.web.reactive.function.server.RouterFunction;
 import org.springframework.web.reactive.result.method.annotation.ArgumentResolverConfigurer;
-import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebFilter;
 import org.springframework.web.util.UriBuilder;
 import org.springframework.web.util.UriBuilderFactory;
@@ -128,16 +126,6 @@ public interface WebTestClient {
 	 */
 	WebTestClient filter(ExchangeFilterFunction filterFunction);
 
-	/**
-	 * Filter the client applying the given transformation function on the
-	 * {@code ServerWebExchange} to every request.
-	 * <p><strong>Note:</strong> this option is applicable only when testing
-	 * without an actual running server.
-	 * @param mutator the transformation function
-	 * @return the filtered client
-	 */
-	WebTestClient exchangeMutator(UnaryOperator<ServerWebExchange> mutator);
-
 
 	// Static, factory methods
 
@@ -180,7 +168,7 @@ public interface WebTestClient {
 	 * @return the {@link WebTestClient} builder
 	 */
 	static Builder bindToHttpHandler(HttpHandler httpHandler) {
-		return new DefaultWebTestClientBuilder(httpHandler, null);
+		return new DefaultWebTestClientBuilder(httpHandler);
 	}
 
 	/**
@@ -198,16 +186,15 @@ public interface WebTestClient {
 	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 mutator the transforming function.
-		 * @see ServerWebExchange#mutate()
-		 */
-		<T extends B> T exchangeMutator(UnaryOperator<ServerWebExchange> mutator);
-
-		/**
-		 * Configure {@link WebFilter}'s for server request processing.
+		 * Register one or more {@link WebFilter} instances to apply to the
+		 * mock server.
+		 *
+		 * <p>This could be used for example to apply {@code ServerWebExchange}
+		 * transformations such as setting the Principal (for all requests or a
+		 * subset) via {@link MockServerExchangeMutator}.
+		 *
 		 * @param filter one or more filters
+		 * @see MockServerExchangeMutator
 		 */
 		<T extends B> T webFilter(WebFilter... filter);
 

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

@@ -26,6 +26,7 @@ import reactor.core.publisher.Mono;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.test.web.reactive.server.MockServerExchangeMutator;
 import org.springframework.test.web.reactive.server.WebTestClient;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.RequestAttribute;
@@ -46,6 +47,8 @@ public class ApplicationContextTests {
 
 	private WebTestClient client;
 
+	private MockServerExchangeMutator exchangeMutator;
+
 
 	@Before
 	public void setUp() throws Exception {
@@ -54,9 +57,15 @@ public class ApplicationContextTests {
 		context.register(WebConfig.class);
 		context.refresh();
 
+		this.exchangeMutator = new MockServerExchangeMutator(principal("Pablo"));
+
+		WebFilter userPrefixFilter = (exchange, chain) -> {
+			Mono<Principal> user = exchange.getPrincipal().map(p -> new TestUser("Mr. " + p.getName()));
+			return chain.filter(exchange.mutate().principal(user).build());
+		};
+
 		this.client = WebTestClient.bindToApplicationContext(context)
-				.exchangeMutator(principal("Pablo"))
-				.webFilter(prefixFilter("Mr."))
+				.webFilter(this.exchangeMutator, userPrefixFilter)
 				.build();
 	}
 
@@ -79,7 +88,7 @@ public class ApplicationContextTests {
 
 	@Test
 	public void perRequestExchangeMutator() throws Exception {
-		this.client.exchangeMutator(principal("Giovanni"))
+		this.exchangeMutator.filterClient(this.client, principal("Giovanni"))
 				.get().uri("/principal")
 				.exchange()
 				.expectStatus().isOk()
@@ -88,9 +97,8 @@ public class ApplicationContextTests {
 
 	@Test
 	public void perRequestMultipleExchangeMutators() throws Exception {
-		this.client
-				.exchangeMutator(attribute("attr1", "foo"))
-				.exchangeMutator(attribute("attr2", "bar"))
+		this.exchangeMutator
+				.filterClient(this.client, attribute("attr1", "foo"), attribute("attr2", "bar"))
 				.get().uri("/attributes")
 				.exchange()
 				.expectStatus().isOk()
@@ -102,13 +110,6 @@ public class ApplicationContextTests {
 		return exchange -> exchange.mutate().principal(Mono.just(new TestUser(userName))).build();
 	}
 
-	private WebFilter prefixFilter(String prefix) {
-		return (exchange, chain) -> {
-			Mono<Principal> user = exchange.getPrincipal().map(p -> new TestUser(prefix + " " + p.getName()));
-			return chain.filter(exchange.mutate().principal(user).build());
-		};
-	}
-
 	private UnaryOperator<ServerWebExchange> attribute(String attrName, String attrValue) {
 		return exchange -> {
 			exchange.getAttributes().put(attrName, attrValue);

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

@@ -19,9 +19,11 @@ package org.springframework.test.web.reactive.server.samples.bind;
 import java.security.Principal;
 import java.util.function.UnaryOperator;
 
+import org.junit.Before;
 import org.junit.Test;
 import reactor.core.publisher.Mono;
 
+import org.springframework.test.web.reactive.server.MockServerExchangeMutator;
 import org.springframework.test.web.reactive.server.WebTestClient;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.RequestAttribute;
@@ -39,11 +41,25 @@ import static org.junit.Assert.assertEquals;
  */
 public class ControllerTests {
 
-	private final WebTestClient client = WebTestClient.bindToController(new TestController())
-			.exchangeMutator(principal("Pablo"))
-			.webFilter(prefixFilter("Mr."))
-			.build();
+	private WebTestClient client;
 
+	private MockServerExchangeMutator exchangeMutator;
+
+
+	@Before
+	public void setUp() throws Exception {
+
+		this.exchangeMutator = new MockServerExchangeMutator(principal("Pablo"));
+
+		WebFilter userPrefixFilter = (exchange, chain) -> {
+			Mono<Principal> user = exchange.getPrincipal().map(p -> new TestUser("Mr. " + p.getName()));
+			return chain.filter(exchange.mutate().principal(user).build());
+		};
+
+		this.client = WebTestClient.bindToController(new TestController())
+				.webFilter(this.exchangeMutator, userPrefixFilter)
+				.build();
+	}
 
 	@Test
 	public void bodyContent() throws Exception {
@@ -63,7 +79,7 @@ public class ControllerTests {
 
 	@Test
 	public void perRequestExchangeMutator() throws Exception {
-		this.client.exchangeMutator(principal("Giovanni"))
+		this.exchangeMutator.filterClient(this.client, principal("Giovanni"))
 				.get().uri("/principal")
 				.exchange()
 				.expectStatus().isOk()
@@ -72,9 +88,8 @@ public class ControllerTests {
 
 	@Test
 	public void perRequestMultipleExchangeMutators() throws Exception {
-		this.client
-				.exchangeMutator(attribute("attr1", "foo"))
-				.exchangeMutator(attribute("attr2", "bar"))
+		this.exchangeMutator
+				.filterClient(this.client, attribute("attr1", "foo"), attribute("attr2", "bar"))
 				.get().uri("/attributes")
 				.exchange()
 				.expectStatus().isOk()
@@ -86,13 +101,6 @@ public class ControllerTests {
 		return exchange -> exchange.mutate().principal(Mono.just(new TestUser(userName))).build();
 	}
 
-	private WebFilter prefixFilter(String prefix) {
-		return (exchange, chain) -> {
-			Mono<Principal> user = exchange.getPrincipal().map(p -> new TestUser(prefix + " " + p.getName()));
-			return chain.filter(exchange.mutate().principal(user).build());
-		};
-	}
-
 	private UnaryOperator<ServerWebExchange> attribute(String attrName, String attrValue) {
 		return exchange -> {
 			exchange.getAttributes().put(attrName, attrValue);