From a95cf0731774661e7639a39abac195f991a2f2f6 Mon Sep 17 00:00:00 2001
From: Arjen Poutsma <apoutsma@pivotal.io>
Date: Tue, 20 Jun 2017 17:06:41 +0200
Subject: [PATCH] Add WebTestClient.mutate()

This commit introduces a WebTestClient.mutate() method,
returning a WebTestClient.Builder.

Issue: SPR-15657
---
 .../reactive/server/DefaultWebTestClient.java  |  9 ++++++++-
 .../server/DefaultWebTestClientBuilder.java    | 18 ++++++++++++++----
 .../web/reactive/server/WebTestClient.java     |  5 +++++
 .../web/reactive/server/WiretapConnector.java  |  3 +++
 4 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/spring-test/src/main/java/org/springframework/test/web/reactive/server/DefaultWebTestClient.java b/spring-test/src/main/java/org/springframework/test/web/reactive/server/DefaultWebTestClient.java
index 34f09097212f..1fe573acb417 100644
--- a/spring-test/src/main/java/org/springframework/test/web/reactive/server/DefaultWebTestClient.java
+++ b/spring-test/src/main/java/org/springframework/test/web/reactive/server/DefaultWebTestClient.java
@@ -74,7 +74,8 @@ class DefaultWebTestClient implements WebTestClient {
 	private final AtomicLong requestIndex = new AtomicLong();
 
 
-	DefaultWebTestClient(WebClient.Builder clientBuilder, ClientHttpConnector connector, @Nullable Duration timeout) {
+	DefaultWebTestClient(WebClient.Builder clientBuilder, ClientHttpConnector connector,
+			@Nullable Duration timeout) {
 		Assert.notNull(clientBuilder, "WebClient.Builder is required");
 		this.wiretapConnector = new WiretapConnector(connector);
 		this.webClient = clientBuilder.clientConnector(this.wiretapConnector).build();
@@ -122,6 +123,12 @@ public UriSpec<RequestHeadersSpec<?>> options() {
 		return toUriSpec(wc -> wc.method(HttpMethod.OPTIONS));
 	}
 
+	@Override
+	public Builder mutate() {
+		return new DefaultWebTestClientBuilder(this.wiretapConnector.getDelegate(),
+				this.webClient.mutate(), this.timeout);
+	}
+
 	private <S extends RequestHeadersSpec<?>> UriSpec<S> toUriSpec(
 			Function<WebClient, WebClient.UriSpec<WebClient.RequestBodySpec>> function) {
 
diff --git a/spring-test/src/main/java/org/springframework/test/web/reactive/server/DefaultWebTestClientBuilder.java b/spring-test/src/main/java/org/springframework/test/web/reactive/server/DefaultWebTestClientBuilder.java
index edddbe8ba19a..b02194c86242 100644
--- a/spring-test/src/main/java/org/springframework/test/web/reactive/server/DefaultWebTestClientBuilder.java
+++ b/spring-test/src/main/java/org/springframework/test/web/reactive/server/DefaultWebTestClientBuilder.java
@@ -24,6 +24,7 @@
 import org.springframework.http.client.reactive.ClientHttpConnector;
 import org.springframework.http.client.reactive.ReactorClientHttpConnector;
 import org.springframework.http.server.reactive.HttpHandler;
+import org.springframework.lang.Nullable;
 import org.springframework.util.MultiValueMap;
 import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
 import org.springframework.web.reactive.function.client.ExchangeStrategies;
@@ -38,7 +39,7 @@
  */
 class DefaultWebTestClientBuilder implements WebTestClient.Builder {
 
-	private final WebClient.Builder webClientBuilder = WebClient.builder();
+	private final WebClient.Builder webClientBuilder;
 
 	private final ClientHttpConnector connector;
 
@@ -49,12 +50,21 @@ class DefaultWebTestClientBuilder implements WebTestClient.Builder {
 		this(new ReactorClientHttpConnector());
 	}
 
+	DefaultWebTestClientBuilder(HttpHandler httpHandler) {
+		this(new HttpHandlerConnector(httpHandler));
+	}
+
 	DefaultWebTestClientBuilder(ClientHttpConnector connector) {
-		this.connector = connector;
+		this(connector, null, null);
 	}
 
-	DefaultWebTestClientBuilder(HttpHandler httpHandler) {
-		this.connector = new HttpHandlerConnector(httpHandler);
+	DefaultWebTestClientBuilder(ClientHttpConnector connector,
+			@Nullable WebClient.Builder webClientBuilder,
+			@Nullable Duration responseTimeout) {
+
+		this.connector = connector;
+		this.webClientBuilder = (webClientBuilder != null ? webClientBuilder : WebClient.builder());
+		this.responseTimeout = responseTimeout;
 	}
 
 
diff --git a/spring-test/src/main/java/org/springframework/test/web/reactive/server/WebTestClient.java b/spring-test/src/main/java/org/springframework/test/web/reactive/server/WebTestClient.java
index 043c8eae32e0..b72fa0dfc3f4 100644
--- a/spring-test/src/main/java/org/springframework/test/web/reactive/server/WebTestClient.java
+++ b/spring-test/src/main/java/org/springframework/test/web/reactive/server/WebTestClient.java
@@ -127,6 +127,11 @@ public interface WebTestClient {
 	UriSpec<RequestHeadersSpec<?>> options();
 
 
+	/**
+	 * Return a builder to mutate properties of this web test client.
+	 */
+	Builder mutate();
+
 	// Static, factory methods
 
 	/**
diff --git a/spring-test/src/main/java/org/springframework/test/web/reactive/server/WiretapConnector.java b/spring-test/src/main/java/org/springframework/test/web/reactive/server/WiretapConnector.java
index a99154639a92..95dbb8f4b43d 100644
--- a/spring-test/src/main/java/org/springframework/test/web/reactive/server/WiretapConnector.java
+++ b/spring-test/src/main/java/org/springframework/test/web/reactive/server/WiretapConnector.java
@@ -50,6 +50,9 @@ class WiretapConnector implements ClientHttpConnector {
 		this.delegate = delegate;
 	}
 
+	public ClientHttpConnector getDelegate() {
+		return this.delegate;
+	}
 
 	@Override
 	public Mono<ClientHttpResponse> connect(HttpMethod method, URI uri,