Improve access to raw content in WebTestClient

If the content has not been consumed, cause it to be produced, and
wait for a certain amount of time before giving up, so the raw content
can be made available. This can occur when:

1) In a mock server scenario the Flux representing the client request
content is passed directly to the mock server request, but is never
consumed because of an error before the body is read.

2) Test obtains FluxExchangeResult (e.g. for streaming) but instead of
consuming the Flux, it calls getResponseBodyContent() instead.

Issue: SPR-17363

Author: rstoyanchev

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

@@ -300,7 +300,7 @@ private static class DefaultResponseSpec implements ResponseSpec {
 		DefaultResponseSpec(WiretapConnector.Info wiretapInfo, ClientResponse response,
 				@Nullable String uriTemplate, Duration timeout) {
 
-			this.exchangeResult = wiretapInfo.createExchangeResult(uriTemplate);
+			this.exchangeResult = wiretapInfo.createExchangeResult(timeout, uriTemplate);
 			this.response = response;
 			this.timeout = timeout;
 		}
@@ -357,13 +357,13 @@ public BodyContentSpec expectBody() {
 		@Override
 		public <T> FluxExchangeResult<T> returnResult(Class<T> elementType) {
 			Flux<T> body = this.response.bodyToFlux(elementType);
-			return new FluxExchangeResult<>(this.exchangeResult, body, this.timeout);
+			return new FluxExchangeResult<>(this.exchangeResult, body);
 		}
 
 		@Override
 		public <T> FluxExchangeResult<T> returnResult(ParameterizedTypeReference<T> elementType) {
 			Flux<T> body = this.response.bodyToFlux(elementType);
-			return new FluxExchangeResult<>(this.exchangeResult, body, this.timeout);
+			return new FluxExchangeResult<>(this.exchangeResult, body);
 		}
 	}
 

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

@@ -24,7 +24,7 @@
 import java.util.List;
 import java.util.stream.Collectors;
 
-import reactor.core.publisher.MonoProcessor;
+import reactor.core.publisher.Mono;
 
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
@@ -36,7 +36,6 @@
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.MultiValueMap;
-import org.springframework.util.ObjectUtils;
 
 /**
  * Container for request and response details for exchanges performed through
@@ -64,9 +63,11 @@ public class ExchangeResult {
 
 	private final ClientHttpResponse response;
 
-	private final MonoProcessor<byte[]> requestBody;
+	private final Mono<byte[]> requestBody;
 
-	private final MonoProcessor<byte[]> responseBody;
+	private final Mono<byte[]> responseBody;
+
+	private final Duration timeout;
 
 	@Nullable
 	private final String uriTemplate;
@@ -80,11 +81,11 @@ public class ExchangeResult {
 	 * @param response the HTTP response
 	 * @param requestBody capture of serialized request body content
 	 * @param responseBody capture of serialized response body content
+	 * @param timeout how long to wait for content to materialize
 	 * @param uriTemplate the URI template used to set up the request, if any
 	 */
 	ExchangeResult(ClientHttpRequest request, ClientHttpResponse response,
-			MonoProcessor<byte[]> requestBody, MonoProcessor<byte[]> responseBody,
-			@Nullable String uriTemplate) {
+			Mono<byte[]> requestBody, Mono<byte[]> responseBody, Duration timeout, @Nullable String uriTemplate) {
 
 		Assert.notNull(request, "ClientHttpRequest is required");
 		Assert.notNull(response, "ClientHttpResponse is required");
@@ -95,6 +96,7 @@ public class ExchangeResult {
 		this.response = response;
 		this.requestBody = requestBody;
 		this.responseBody = responseBody;
+		this.timeout = timeout;
 		this.uriTemplate = uriTemplate;
 	}
 
@@ -106,6 +108,7 @@ public class ExchangeResult {
 		this.response = other.response;
 		this.requestBody = other.requestBody;
 		this.responseBody = other.responseBody;
+		this.timeout = other.timeout;
 		this.uriTemplate = other.uriTemplate;
 	}
 
@@ -140,14 +143,14 @@ public HttpHeaders getRequestHeaders() {
 	}
 
 	/**
-	 * Return the raw request body content written as a {@code byte[]}.
-	 * @throws IllegalStateException if the request body is not fully written yet.
+	 * Return the raw request body content written through the request.
+	 * <p><strong>Note:</strong> If the request content has not been consumed
+	 * for any reason yet, use of this method will trigger consumption.
+	 * @throws IllegalStateException if the request body is not been fully written.
 	 */
 	@Nullable
 	public byte[] getRequestBodyContent() {
-		MonoProcessor<byte[]> body = this.requestBody;
-		Assert.isTrue(body.isTerminated(), "Request body incomplete.");
-		return body.block(Duration.ZERO);
+		return this.requestBody.block(this.timeout);
 	}
 
 
@@ -173,14 +176,14 @@ public MultiValueMap<String, ResponseCookie> getResponseCookies() {
 	}
 
 	/**
-	 * Return the raw request body content written as a {@code byte[]}.
-	 * @throws IllegalStateException if the response is not fully read yet.
+	 * Return the raw request body content written to the response.
+	 * <p><strong>Note:</strong> If the response content has not been consumed
+	 * yet, use of this method will trigger consumption.
+	 * @throws IllegalStateException if the response is not been fully read.
 	 */
 	@Nullable
 	public byte[] getResponseBodyContent() {
-		MonoProcessor<byte[]> body = this.responseBody;
-		Assert.state(body.isTerminated(), "Response body incomplete");
-		return body.block(Duration.ZERO);
+		return this.responseBody.block(this.timeout);
 	}
 
 
@@ -223,30 +226,25 @@ private String formatHeaders(HttpHeaders headers, String delimiter) {
 				.collect(Collectors.joining(delimiter));
 	}
 
-	private String formatBody(@Nullable MediaType contentType, MonoProcessor<byte[]> body) {
-		if (body.isSuccess()) {
-			byte[] bytes = body.block(Duration.ZERO);
-			if (ObjectUtils.isEmpty(bytes)) {
-				return "No content";
-			}
-			if (contentType == null) {
-				return "Unknown content type (" + bytes.length + " bytes)";
-			}
-			Charset charset = contentType.getCharset();
-			if (charset != null) {
-				return new String(bytes, charset);
-			}
-			if (PRINTABLE_MEDIA_TYPES.stream().anyMatch(contentType::isCompatibleWith)) {
-				return new String(bytes, StandardCharsets.UTF_8);
-			}
-			return "Unknown charset (" + bytes.length + " bytes)";
-		}
-		else if (body.isError()) {
-			return "I/O failure: " + body.getError();
-		}
-		else {
-			return "Content not available yet";
-		}
+	@Nullable
+	private String formatBody(@Nullable MediaType contentType, Mono<byte[]> body) {
+		return body
+				.map(bytes -> {
+					if (contentType == null) {
+						return "Unknown content type (" + bytes.length + " bytes)";
+					}
+					Charset charset = contentType.getCharset();
+					if (charset != null) {
+						return new String(bytes, charset);
+					}
+					if (PRINTABLE_MEDIA_TYPES.stream().anyMatch(contentType::isCompatibleWith)) {
+						return new String(bytes, StandardCharsets.UTF_8);
+					}
+					return "Unknown charset (" + bytes.length + " bytes)";
+				})
+				.defaultIfEmpty("No content")
+				.onErrorResume(ex -> Mono.just("Failed to obtain content: " + ex.getMessage()))
+				.block(this.timeout);
 	}
 
 }

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

@@ -16,13 +16,9 @@
 
 package org.springframework.test.web.reactive.server;
 
-import java.time.Duration;
 import java.util.function.Consumer;
 
 import reactor.core.publisher.Flux;
-import reactor.core.publisher.Mono;
-
-import org.springframework.lang.Nullable;
 
 /**
  * {@code ExchangeResult} variant with the response body decoded as
@@ -35,20 +31,12 @@
  */
 public class FluxExchangeResult<T> extends ExchangeResult {
 
-	private static final IllegalStateException TIMEOUT_ERROR =
-			new IllegalStateException("Response timeout: for infinite streams " +
-					"use getResponseBody() first with explicit cancellation, e.g. via take(n).");
-
-
 	private final Flux<T> body;
 
-	private final Duration timeout;
-
 
-	FluxExchangeResult(ExchangeResult result, Flux<T> body, Duration timeout) {
+	FluxExchangeResult(ExchangeResult result, Flux<T> body) {
 		super(result);
 		this.body = body;
-		this.timeout = timeout;
 	}
 
 
@@ -81,22 +69,6 @@ public Flux<T> getResponseBody() {
 		return this.body;
 	}
 
-	/**
-	 * {@inheritDoc}
-	 * <p><strong>Note:</strong> this method should typically be called after
-	 * the response has been consumed in full via {@link #getResponseBody()}.
-	 * Calling it first will cause the response {@code Flux<T>} to be consumed
-	 * via {@code getResponseBody.ignoreElements()}.
-	 */
-	@Override
-	@Nullable
-	public byte[] getResponseBodyContent() {
-		return this.body.ignoreElements()
-				.timeout(this.timeout, Mono.error(TIMEOUT_ERROR))
-				.then(Mono.defer(() -> Mono.justOrEmpty(super.getResponseBodyContent())))
-				.block();
-	}
-
 	/**
 	 * Invoke the given consumer within {@link #assertWithDiagnostics(Runnable)}
 	 * passing {@code "this"} instance to it. This method allows the following,

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

@@ -17,12 +17,14 @@
 package org.springframework.test.web.reactive.server;
 
 import java.net.URI;
+import java.time.Duration;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
 import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscription;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.MonoProcessor;
@@ -112,9 +114,11 @@ public Info(WiretapClientHttpRequest request, WiretapClientHttpResponse response
 		}
 
 
-		public ExchangeResult createExchangeResult(@Nullable  String uriTemplate) {
+		public ExchangeResult createExchangeResult(Duration timeout, @Nullable String uriTemplate) {
 			return new ExchangeResult(this.request, this.response,
-					this.request.getRecorder().getContent(), this.response.getRecorder().getContent(), uriTemplate);
+					Mono.defer(() -> this.request.getRecorder().getContent()),
+					Mono.defer(() -> this.response.getRecorder().getContent()),
+					timeout, uriTemplate);
 		}
 	}
 
@@ -126,21 +130,21 @@ final static class WiretapRecorder {
 
 		private static final DataBufferFactory bufferFactory = new DefaultDataBufferFactory();
 
-		public static final byte[] EMPTY_CONTENT = new byte[0];
-
 
 		@Nullable
-		private final Publisher<? extends DataBuffer> publisher;
+		private final Flux<? extends DataBuffer> publisher;
 
 		@Nullable
-		private final Publisher<? extends Publisher<? extends DataBuffer>> publisherNested;
+		private final Flux<? extends Publisher<? extends DataBuffer>> publisherNested;
 
 		private final DataBuffer buffer;
 
 		private final MonoProcessor<byte[]> content;
 
+		private volatile boolean subscriberRegistered;
+
 
-		private WiretapRecorder(@Nullable Publisher<? extends DataBuffer> publisher,
+		public WiretapRecorder(@Nullable Publisher<? extends DataBuffer> publisher,
 				@Nullable Publisher<? extends Publisher<? extends DataBuffer>> publisherNested) {
 
 			if (publisher != null && publisherNested != null) {
@@ -149,24 +153,22 @@ private WiretapRecorder(@Nullable Publisher<? extends DataBuffer> publisher,
 
 			this.publisher = publisher != null ?
 					Flux.from(publisher)
+							.doOnSubscribe(this::handleOnSubscribe)
 							.doOnNext(this::handleOnNext)
 							.doOnError(this::handleOnError)
 							.doOnCancel(this::handleOnComplete)
 							.doOnComplete(this::handleOnComplete) : null;
 
 			this.publisherNested = publisherNested != null ?
 					Flux.from(publisherNested)
+							.doOnSubscribe(this::handleOnSubscribe)
 							.map(p -> Flux.from(p).doOnNext(this::handleOnNext).doOnError(this::handleOnError))
 							.doOnError(this::handleOnError)
 							.doOnCancel(this::handleOnComplete)
 							.doOnComplete(this::handleOnComplete) : null;
 
 			this.buffer =  bufferFactory.allocateBuffer();
 			this.content = MonoProcessor.create();
-
-			if (this.publisher == null && this.publisherNested == null) {
-				this.content.onNext(EMPTY_CONTENT);
-			}
 		}
 
 
@@ -180,11 +182,36 @@ public Publisher<? extends Publisher<? extends DataBuffer>> getNestedPublisherTo
 			return this.publisherNested;
 		}
 
-		public MonoProcessor<byte[]> getContent() {
-			return this.content;
+		public Mono<byte[]> getContent() {
+			// No publisher (e.g. request#setComplete)
+			if (this.publisher == null && this.publisherNested == null) {
+				return Mono.empty();
+			}
+			if (this.content.isTerminated()) {
+				return this.content;
+			}
+			if (this.subscriberRegistered) {
+				return Mono.error(new IllegalStateException(
+						"Subscriber registered but content is not yet fully consumed."));
+			}
+			else {
+				// No subscriber, e.g.:
+				// 	- mock server request body never consumed (error before read)
+				//  - FluxExchangeResult#getResponseBodyContent called
+				(this.publisher != null ? this.publisher : this.publisherNested)
+						.onErrorMap(ex -> new IllegalStateException(
+								"Content was not been consumed and " +
+										"an error was raised on attempt to produce it:", ex))
+						.subscribe();
+				return this.content;
+			}
 		}
 
 
+		private void handleOnSubscribe(Subscription subscription) {
+			this.subscriberRegistered = true;
+		}
+
 		private void handleOnNext(DataBuffer nextBuffer) {
 			this.buffer.write(nextBuffer);
 		}

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

@@ -17,6 +17,7 @@
 package org.springframework.test.web.reactive.server;
 
 import java.net.URI;
+import java.time.Duration;
 import java.time.ZoneId;
 import java.time.ZonedDateTime;
 import java.util.concurrent.TimeUnit;
@@ -257,7 +258,7 @@ private HeaderAssertions headerAssertions(HttpHeaders responseHeaders) {
 		MonoProcessor<byte[]> emptyContent = MonoProcessor.create();
 		emptyContent.onComplete();
 
-		ExchangeResult result = new ExchangeResult(request, response, emptyContent, emptyContent, null);
+		ExchangeResult result = new ExchangeResult(request, response, emptyContent, emptyContent, Duration.ZERO, null);
 		return new HeaderAssertions(result, mock(WebTestClient.ResponseSpec.class));
 	}