Improve writing in mock reactive request and response

Before this change the write Publisher was saved and Mono.empty()
returned from the write metohd which did not properly implement
the write contract since no writing ("consuming") was done.

This can be a problem in some cases. For example the request may appear
to succeed even if the publisher produces an error later when
subscribed to later after request handling completes.

This commit introduces a writeHandler function in the mock request and
response. By default it "writes" by consuming the content immediately,
which allows it to return a Mono<Void> that properly reflects when
writing is done, and it also caches the data so it may be replayed
later for test assertions.

For streaming scenario a custom writeHandler may be registered which
allows the custom handling to determine how long to stream before
cancelling so request handling may complete.

Issue: SPR-14590

spring-test/src/main/java/org/springframework/mock/http/client/reactive/MockClientHttpRequest.java

@@ -17,6 +17,7 @@
 package org.springframework.mock.http.client.reactive;
 
 import java.net.URI;
+import java.util.function.Function;
 
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
@@ -28,6 +29,7 @@ import org.springframework.core.io.buffer.DefaultDataBufferFactory;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.client.reactive.AbstractClientHttpRequest;
 import org.springframework.http.client.reactive.ClientHttpRequest;
+import org.springframework.util.Assert;
 import org.springframework.web.util.UriComponentsBuilder;
 
 /**
@@ -44,7 +46,11 @@ public class MockClientHttpRequest extends AbstractClientHttpRequest {
 
 	private final DataBufferFactory bufferFactory = new DefaultDataBufferFactory();
 
-	private Flux<DataBuffer> body;
+	private Flux<DataBuffer> body = Flux.error(
+			new IllegalStateException("The body is not set. " +
+					"Did handling complete with success? Is a custom \"writeHandler\" configured?"));
+
+	private Function<Flux<DataBuffer>, Mono<Void>> writeHandler = initDefaultWriteHandler();
 
 
 	public MockClientHttpRequest(HttpMethod httpMethod, String urlTemplate, Object... vars) {
@@ -56,6 +62,13 @@ public class MockClientHttpRequest extends AbstractClientHttpRequest {
 		this.url = url;
 	}
 
+	private Function<Flux<DataBuffer>, Mono<Void>> initDefaultWriteHandler() {
+		return body -> {
+			this.body = body.cache();
+			return this.body.then();
+		};
+	}
+
 
 	@Override
 	public HttpMethod getMethod() {
@@ -72,22 +85,27 @@ public class MockClientHttpRequest extends AbstractClientHttpRequest {
 		return this.bufferFactory;
 	}
 
+	/**
+	 * Return the request body, or an error stream if the body was never set
+	 * or when {@link #setWriteHandler} is configured.
+	 */
 	public Flux<DataBuffer> getBody() {
 		return this.body;
 	}
 
-	@Override
+	/**
-	public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
+	 * Configure a custom handler for writing the request body.
-		this.body = Flux.from(body);
+	 *
-		return doCommit(() -> {
+	 * <p>The default write handler consumes and caches the request body so it
-			this.body = Flux.from(body);
+	 * may be accessed subsequently, e.g. in test assertions. Use this property
-			return Mono.empty();
+	 * when the request body is an infinite stream.
-		});
+	 *
-	}
+	 * @param writeHandler the write handler to use returning {@code Mono<Void>}
-
+	 * when the body has been "written" (i.e. consumed).
-	@Override
+	 */
-	public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
+	public void setWriteHandler(Function<Flux<DataBuffer>, Mono<Void>> writeHandler) {
-		return writeWith(Flux.from(body).flatMap(p -> p));
+		Assert.notNull(writeHandler, "'writeHandler' is required");
+		this.writeHandler = writeHandler;
 	}
 
 	@Override
@@ -98,9 +116,19 @@ public class MockClientHttpRequest extends AbstractClientHttpRequest {
 	protected void applyCookies() {
 	}
 
+	@Override
+	public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
+		return doCommit(() -> Mono.defer(() -> this.writeHandler.apply(Flux.from(body))));
+	}
+
+	@Override
+	public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
+		return writeWith(Flux.from(body).flatMap(p -> p));
+	}
+
 	@Override
 	public Mono<Void> setComplete() {
-		return doCommit(Mono::empty);
+		return writeWith(Flux.empty());
 	}
 
 }

spring-test/src/main/java/org/springframework/mock/http/server/reactive/MockServerHttpRequest.java

@@ -219,6 +219,12 @@ public class MockServerHttpRequest extends AbstractServerHttpRequest {
 		 */
 		B cookie(String path, HttpCookie... cookie);
 
+		/**
+		 * Add the given cookies.
+		 * @param cookies the cookies.
+		 */
+		B cookies(MultiValueMap<String, HttpCookie> cookies);
+
 		/**
 		 * Add the given, single header value under the given name.
 		 * @param headerName  the header name
@@ -227,6 +233,12 @@ public class MockServerHttpRequest extends AbstractServerHttpRequest {
 		 */
 		B header(String headerName, String... headerValues);
 
+		/**
+		 * Add the given header values.
+		 * @param headers the header values
+		 */
+		B headers(MultiValueMap<String, String> headers);
+
 		/**
 		 * Set the list of acceptable {@linkplain MediaType media types}, as
 		 * specified by the {@code Accept} header.
@@ -360,6 +372,12 @@ public class MockServerHttpRequest extends AbstractServerHttpRequest {
 			return this;
 		}
 
+		@Override
+		public BodyBuilder cookies(MultiValueMap<String, HttpCookie> cookies) {
+			this.cookies.putAll(cookies);
+			return this;
+		}
+
 		@Override
 		public BodyBuilder header(String headerName, String... headerValues) {
 			for (String headerValue : headerValues) {
@@ -368,6 +386,12 @@ public class MockServerHttpRequest extends AbstractServerHttpRequest {
 			return this;
 		}
 
+		@Override
+		public BodyBuilder headers(MultiValueMap<String, String> headers) {
+			this.headers.putAll(headers);
+			return this;
+		}
+
 		@Override
 		public BodyBuilder accept(MediaType... acceptableMediaTypes) {
 			this.headers.setAccept(Arrays.asList(acceptableMediaTypes));

spring-test/src/main/java/org/springframework/mock/http/server/reactive/MockServerHttpResponse.java

@@ -18,6 +18,7 @@ package org.springframework.mock.http.server.reactive;
 
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
+import java.util.function.Function;
 
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
@@ -38,24 +39,37 @@ import org.springframework.util.Assert;
  */
 public class MockServerHttpResponse extends AbstractServerHttpResponse {
 
-	private Flux<DataBuffer> body;
+	private Flux<DataBuffer> body = Flux.error(
+			new IllegalStateException("The body is not set. " +
+					"Did handling complete with success? Is a custom \"writeHandler\" configured?"));
+
+	private Function<Flux<DataBuffer>, Mono<Void>> writeHandler = initDefaultWriteHandler();
 
 
 	public MockServerHttpResponse() {
 		super(new DefaultDataBufferFactory());
 	}
 
+	private Function<Flux<DataBuffer>, Mono<Void>> initDefaultWriteHandler() {
+		return body -> {
+			this.body = body.cache();
+			return this.body.then();
+		};
+	}
+
 
 	/**
-	 * Return the output Publisher used to write to the response.
+	 * Return the request body, or an error stream if the body was never set
+	 * or when {@link #setWriteHandler} is configured.
 	 */
 	public Flux<DataBuffer> getBody() {
 		return this.body;
 	}
 
 	/**
-	 * Return the response body aggregated and converted to a String using the
+	 * Shortcut method that delegates to {@link #getBody()} and then aggregates
-	 * charset of the Content-Type response or otherwise as "UTF-8".
+	 * the data buffers and converts to a String using the charset of the
+	 * Content-Type header or falling back on "UTF-8" by default.
 	 */
 	public Mono<String> getBodyAsString() {
 		Charset charset = getCharset();
@@ -84,15 +98,19 @@ public class MockServerHttpResponse extends AbstractServerHttpResponse {
 		return (charset != null ? charset : StandardCharsets.UTF_8);
 	}
 
-	@Override
+	/**
-	protected Mono<Void> writeWithInternal(Publisher<? extends DataBuffer> body) {
+	 * Configure a custom handler for writing the request body.
-		this.body = Flux.from(body);
+	 *
-		return Mono.empty();
+	 * <p>The default write handler consumes and caches the request body so it
-	}
+	 * may be accessed subsequently, e.g. in test assertions. Use this property
-
+	 * when the request body is an infinite stream.
-	@Override
+	 *
-	protected Mono<Void> writeAndFlushWithInternal(Publisher<? extends Publisher<? extends DataBuffer>> body) {
+	 * @param writeHandler the write handler to use returning {@code Mono<Void>}
-		return writeWithInternal(Flux.from(body).flatMap(Flux::from));
+	 * when the body has been "written" (i.e. consumed).
+	 */
+	public void setWriteHandler(Function<Flux<DataBuffer>, Mono<Void>> writeHandler) {
+		Assert.notNull(writeHandler, "'writeHandler' is required");
+		this.writeHandler = writeHandler;
 	}
 
 	@Override
@@ -107,4 +125,19 @@ public class MockServerHttpResponse extends AbstractServerHttpResponse {
 	protected void applyCookies() {
 	}
 
+	@Override
+	protected Mono<Void> writeWithInternal(Publisher<? extends DataBuffer> body) {
+		return this.writeHandler.apply(Flux.from(body));
+	}
+
+	@Override
+	protected Mono<Void> writeAndFlushWithInternal(Publisher<? extends Publisher<? extends DataBuffer>> body) {
+		return this.writeHandler.apply(Flux.from(body).concatMap(Flux::from));
+	}
+
+	@Override
+	public Mono<Void> setComplete() {
+		return doCommit(() -> Mono.defer(() -> this.writeHandler.apply(Flux.empty())));
+	}
+
 }

spring-web/src/test/java/org/springframework/mock/http/client/reactive/test/MockClientHttpRequest.java

@@ -17,6 +17,7 @@
 package org.springframework.mock.http.client.reactive.test;
 
 import java.net.URI;
+import java.util.function.Function;
 
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
@@ -28,6 +29,7 @@ import org.springframework.core.io.buffer.DefaultDataBufferFactory;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.client.reactive.AbstractClientHttpRequest;
 import org.springframework.http.client.reactive.ClientHttpRequest;
+import org.springframework.util.Assert;
 import org.springframework.web.util.UriComponentsBuilder;
 
 /**
@@ -44,7 +46,11 @@ public class MockClientHttpRequest extends AbstractClientHttpRequest {
 
 	private final DataBufferFactory bufferFactory = new DefaultDataBufferFactory();
 
-	private Flux<DataBuffer> body;
+	private Flux<DataBuffer> body = Flux.error(
+			new IllegalStateException("The body is not set. " +
+					"Did handling complete with success? Is a custom \"writeHandler\" configured?"));
+
+	private Function<Flux<DataBuffer>, Mono<Void>> writeHandler = initDefaultWriteHandler();
 
 
 	public MockClientHttpRequest(HttpMethod httpMethod, String urlTemplate, Object... vars) {
@@ -56,6 +62,13 @@ public class MockClientHttpRequest extends AbstractClientHttpRequest {
 		this.url = url;
 	}
 
+	private Function<Flux<DataBuffer>, Mono<Void>> initDefaultWriteHandler() {
+		return body -> {
+			this.body = body.cache();
+			return this.body.then();
+		};
+	}
+
 
 	@Override
 	public HttpMethod getMethod() {
@@ -72,22 +85,27 @@ public class MockClientHttpRequest extends AbstractClientHttpRequest {
 		return this.bufferFactory;
 	}
 
+	/**
+	 * Return the request body, or an error stream if the body was never set
+	 * or when {@link #setWriteHandler} is configured.
+	 */
 	public Flux<DataBuffer> getBody() {
 		return this.body;
 	}
 
-	@Override
+	/**
-	public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
+	 * Configure a custom handler for writing the request body.
-		this.body = Flux.from(body);
+	 *
-		return doCommit(() -> {
+	 * <p>The default write handler consumes and caches the request body so it
-			this.body = Flux.from(body);
+	 * may be accessed subsequently, e.g. in test assertions. Use this property
-			return Mono.empty();
+	 * when the request body is an infinite stream.
-		});
+	 *
-	}
+	 * @param writeHandler the write handler to use returning {@code Mono<Void>}
-
+	 * when the body has been "written" (i.e. consumed).
-	@Override
+	 */
-	public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
+	public void setWriteHandler(Function<Flux<DataBuffer>, Mono<Void>> writeHandler) {
-		return writeWith(Flux.from(body).flatMap(p -> p));
+		Assert.notNull(writeHandler, "'writeHandler' is required");
+		this.writeHandler = writeHandler;
 	}
 
 	@Override
@@ -98,9 +116,19 @@ public class MockClientHttpRequest extends AbstractClientHttpRequest {
 	protected void applyCookies() {
 	}
 
+	@Override
+	public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
+		return doCommit(() -> Mono.defer(() -> this.writeHandler.apply(Flux.from(body))));
+	}
+
+	@Override
+	public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
+		return writeWith(Flux.from(body).flatMap(p -> p));
+	}
+
 	@Override
 	public Mono<Void> setComplete() {
-		return doCommit(Mono::empty);
+		return writeWith(Flux.empty());
 	}
 
 }

spring-web/src/test/java/org/springframework/mock/http/server/reactive/test/MockServerHttpResponse.java

@@ -18,6 +18,7 @@ package org.springframework.mock.http.server.reactive.test;
 
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
+import java.util.function.Function;
 
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
@@ -38,24 +39,37 @@ import org.springframework.util.Assert;
  */
 public class MockServerHttpResponse extends AbstractServerHttpResponse {
 
-	private Flux<DataBuffer> body;
+	private Flux<DataBuffer> body = Flux.error(
+			new IllegalStateException("The body is not set. " +
+					"Did handling complete with success? Is a custom \"writeHandler\" configured?"));
+
+	private Function<Flux<DataBuffer>, Mono<Void>> writeHandler = initDefaultWriteHandler();
 
 
 	public MockServerHttpResponse() {
 		super(new DefaultDataBufferFactory());
 	}
 
+	private Function<Flux<DataBuffer>, Mono<Void>> initDefaultWriteHandler() {
+		return body -> {
+			this.body = body.cache();
+			return this.body.then();
+		};
+	}
+
 
 	/**
-	 * Return the output Publisher used to write to the response.
+	 * Return the request body, or an error stream if the body was never set
+	 * or when {@link #setWriteHandler} is configured.
 	 */
 	public Flux<DataBuffer> getBody() {
 		return this.body;
 	}
 
 	/**
-	 * Return the response body aggregated and converted to a String using the
+	 * Shortcut method that delegates to {@link #getBody()} and then aggregates
-	 * charset of the Content-Type response or otherwise as "UTF-8".
+	 * the data buffers and converts to a String using the charset of the
+	 * Content-Type header or falling back on "UTF-8" by default.
 	 */
 	public Mono<String> getBodyAsString() {
 		Charset charset = getCharset();
@@ -84,15 +98,19 @@ public class MockServerHttpResponse extends AbstractServerHttpResponse {
 		return (charset != null ? charset : StandardCharsets.UTF_8);
 	}
 
-	@Override
+	/**
-	protected Mono<Void> writeWithInternal(Publisher<? extends DataBuffer> body) {
+	 * Configure a custom handler for writing the request body.
-		this.body = Flux.from(body);
+	 *
-		return Mono.empty();
+	 * <p>The default write handler consumes and caches the request body so it
-	}
+	 * may be accessed subsequently, e.g. in test assertions. Use this property
-
+	 * when the request body is an infinite stream.
-	@Override
+	 *
-	protected Mono<Void> writeAndFlushWithInternal(Publisher<? extends Publisher<? extends DataBuffer>> body) {
+	 * @param writeHandler the write handler to use returning {@code Mono<Void>}
-		return writeWithInternal(Flux.from(body).flatMap(Flux::from));
+	 * when the body has been "written" (i.e. consumed).
+	 */
+	public void setWriteHandler(Function<Flux<DataBuffer>, Mono<Void>> writeHandler) {
+		Assert.notNull(writeHandler, "'writeHandler' is required");
+		this.writeHandler = writeHandler;
 	}
 
 	@Override
@@ -107,4 +125,19 @@ public class MockServerHttpResponse extends AbstractServerHttpResponse {
 	protected void applyCookies() {
 	}
 
+	@Override
+	protected Mono<Void> writeWithInternal(Publisher<? extends DataBuffer> body) {
+		return this.writeHandler.apply(Flux.from(body));
+	}
+
+	@Override
+	protected Mono<Void> writeAndFlushWithInternal(Publisher<? extends Publisher<? extends DataBuffer>> body) {
+		return this.writeHandler.apply(Flux.from(body).concatMap(Flux::from));
+	}
+
+	@Override
+	public Mono<Void> setComplete() {
+		return doCommit(() -> Mono.defer(() -> this.writeHandler.apply(Flux.empty())));
+	}
+
 }

spring-webflux/src/main/java/org/springframework/web/reactive/result/method/annotation/ResponseEntityResultHandler.java

@@ -146,16 +146,14 @@ public class ResponseEntityResultHandler extends AbstractMessageWriterResultHand
 						.forEach(entry -> responseHeaders.put(entry.getKey(), entry.getValue()));
 			}
 			if(httpEntity.getBody() == null) {
-				exchange.getResponse().setComplete();
+				return exchange.getResponse().setComplete();
-				return Mono.empty();
 			}
 
 			String etag = entityHeaders.getETag();
 			Instant lastModified = Instant.ofEpochMilli(entityHeaders.getLastModified());
 			HttpMethod httpMethod = exchange.getRequest().getMethod();
 			if (SAFE_METHODS.contains(httpMethod) && exchange.checkNotModified(etag, lastModified)) {
-				exchange.getResponse().setComplete();
+				return exchange.getResponse().setComplete();
-				return Mono.empty();
 			}
 
 			return writeBody(httpEntity.getBody(), bodyType, exchange);

spring-webflux/src/test/java/org/springframework/web/reactive/function/client/DefaultClientRequestBuilderTests.java

@@ -23,6 +23,7 @@ import java.util.List;
 
 import org.junit.Test;
 import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
 
 import org.springframework.core.codec.CharSequenceEncoder;
 import org.springframework.core.io.buffer.DataBuffer;
@@ -89,7 +90,7 @@ public class DefaultClientRequestBuilderTests {
 
 		assertEquals("MyValue", request.getHeaders().getFirst("MyKey"));
 		assertEquals("bar", request.getCookies().getFirst("foo").getValue());
-		assertNull(request.getBody());
+		StepVerifier.create(request.getBody()).expectComplete().verify();
 	}
 
 	@Test

spring-webflux/src/test/java/org/springframework/web/reactive/function/server/DefaultServerResponseBuilderTests.java

@@ -36,7 +36,6 @@ import org.springframework.mock.http.server.reactive.test.MockServerHttpResponse
 import org.springframework.web.server.ServerWebExchange;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNull;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -245,8 +244,7 @@ public class DefaultServerResponseBuilderTests {
 
 		assertEquals(HttpStatus.CREATED, response.getStatusCode());
 		assertEquals("MyValue", response.getHeaders().getFirst("MyKey"));
-		assertNull(response.getBody());
+		StepVerifier.create(response.getBody()).expectComplete().verify();
-
 	}
 
 	@Test
@@ -261,7 +259,7 @@ public class DefaultServerResponseBuilderTests {
 
 		result.then(res -> res.writeTo(exchange, strategies)).block();
 
-		assertNull(response.getBody());
+		StepVerifier.create(response.getBody()).expectComplete().verify();
 	}
 
 /*

spring-webflux/src/test/java/org/springframework/web/reactive/function/server/ResourceHandlerFunctionTests.java

@@ -37,7 +37,6 @@ import org.springframework.web.server.session.MockWebSessionManager;
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNull;
 
 /**
  * @author Arjen Poutsma
@@ -150,7 +149,7 @@ TODO: enable when ServerEntityResponse is reintroduced
 		assertEquals(EnumSet.of(HttpMethod.GET, HttpMethod.HEAD, HttpMethod.OPTIONS),
 				mockResponse.getHeaders().getAllow());
 
-		assertNull(mockResponse.getBody());
+		StepVerifier.create(mockResponse.getBody()).expectComplete().verify();
 	}
 
 }

spring-webflux/src/test/java/org/springframework/web/reactive/resource/ResourceWebHandlerTests.java

@@ -125,7 +125,10 @@ public class ResourceWebHandlerTests {
 		assertEquals(headers.getLastModified() / 1000, resourceLastModifiedDate("test/foo.css") / 1000);
 		assertEquals("bytes", headers.getFirst("Accept-Ranges"));
 		assertEquals(1, headers.get("Accept-Ranges").size());
-		assertNull(this.response.getBody());
+
+		StepVerifier.create(this.response.getBody())
+				.expectErrorMatches(ex -> ex.getMessage().startsWith("The body is not set."))
+				.verify();
 	}
 
 	@Test

spring-webflux/src/test/java/org/springframework/web/reactive/result/method/annotation/MessageWriterResultHandlerTests.java

@@ -123,7 +123,8 @@ public class MessageWriterResultHandlerTests {
 		this.resultHandler.writeBody(body, returnType(type), this.exchange).block(Duration.ofSeconds(5));
 
 		assertNull(this.response.getHeaders().get("Content-Type"));
-		assertNull(this.response.getBody());
+		StepVerifier.create(this.response.getBody())
+				.expectErrorMatches(ex -> ex.getMessage().startsWith("The body is not set.")).verify();
 	}
 
 	@Test  // SPR-13135

spring-webflux/src/test/java/org/springframework/web/reactive/result/method/annotation/ResponseEntityResultHandlerTests.java

@@ -62,7 +62,6 @@ import org.springframework.web.server.adapter.DefaultServerWebExchange;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.springframework.core.ResolvableType.forClassWithGenerics;
 import static org.springframework.http.ResponseEntity.notFound;
@@ -156,7 +155,7 @@ public class ResponseEntityResultHandlerTests {
 
 		assertEquals(HttpStatus.NO_CONTENT, this.response.getStatusCode());
 		assertEquals(0, this.response.getHeaders().size());
-		assertNull(this.response.getBody());
+		assertResponseBodyIsEmpty();
 	}
 
 	@Test
@@ -170,7 +169,7 @@ public class ResponseEntityResultHandlerTests {
 		assertEquals(HttpStatus.CREATED, this.response.getStatusCode());
 		assertEquals(1, this.response.getHeaders().size());
 		assertEquals(location, this.response.getHeaders().getLocation());
-		assertNull(this.response.getBody());
+		assertResponseBodyIsEmpty();
 	}
 
 	@Test
@@ -180,7 +179,7 @@ public class ResponseEntityResultHandlerTests {
 		HandlerResult result = handlerResult(returnValue, returnType);
 		this.resultHandler.handleResult(createExchange(), result).block(Duration.ofSeconds(5));
 		assertEquals(HttpStatus.NOT_FOUND, this.response.getStatusCode());
-		assertNull(this.response.getBody());
+		assertResponseBodyIsEmpty();
 	}
 
 	@Test
@@ -311,7 +310,7 @@ public class ResponseEntityResultHandlerTests {
 		this.resultHandler.handleResult(exchange, result).block(Duration.ofSeconds(5));
 
 		assertEquals(HttpStatus.NOT_FOUND, this.response.getStatusCode());
-		assertNull(this.response.getBody());
+		assertResponseBodyIsEmpty();
 	}
 
 
@@ -348,13 +347,19 @@ public class ResponseEntityResultHandlerTests {
 				.verify();
 	}
 
-	private void assertConditionalResponse(HttpStatus status, String body, String etag, Instant lastModified) throws Exception {
+	private void assertResponseBodyIsEmpty() {
+		StepVerifier.create(this.response.getBody()).expectComplete().verify();
+	}
+
+	private void assertConditionalResponse(HttpStatus status, String body, String etag, Instant lastModified)
+			throws Exception {
+
 		assertEquals(status, this.response.getStatusCode());
 		if (body != null) {
 			assertResponseBody(body);
 		}
 		else {
-			assertNull(this.response.getBody());
+			assertResponseBodyIsEmpty();
 		}
 		if (etag != null) {
 			assertEquals(1, this.response.getHeaders().get(HttpHeaders.ETAG).size());