Add consumeWith to FluxExchangeResult

Issue: SPR-15959

Author: rstoyanchev

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

@@ -34,6 +34,7 @@
 
 import org.springframework.core.ParameterizedTypeReference;
 import org.springframework.core.io.ByteArrayResource;
+import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.MediaType;
@@ -46,6 +47,7 @@
 import org.springframework.util.MimeType;
 import org.springframework.util.MultiValueMap;
 import org.springframework.web.reactive.function.BodyExtractor;
+import org.springframework.web.reactive.function.BodyExtractors;
 import org.springframework.web.reactive.function.BodyInserter;
 import org.springframework.web.reactive.function.client.ClientResponse;
 import org.springframework.web.reactive.function.client.WebClient;
@@ -379,6 +381,11 @@ public BodyContentSpec expectBody() {
 			return new DefaultBodyContentSpec(this.result.decodeToByteArray());
 		}
 
+		@Override
+		public FluxExchangeResult<DataBuffer> returnResult() {
+			return this.result.decodeToFlux(BodyExtractors.toDataBuffers());
+		}
+
 		@Override
 		public <T> FluxExchangeResult<T> returnResult(Class<T> elementType) {
 			return this.result.decodeToFlux(toFlux(elementType));

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

@@ -17,6 +17,7 @@
 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;
@@ -97,4 +98,23 @@ public byte[] getResponseBodyContent() {
 				.block();
 	}
 
+	/**
+	 * Invoke the given consumer within {@link #assertWithDiagnostics(Runnable)}
+	 * passing {@code "this"} instance to it. This method allows the following,
+	 * without leaving the {@code WebTestClient} chain of calls:
+	 * <pre class="code">
+	 *	client.get()
+	 * 		.uri("/persons")
+	 * 		.accept(TEXT_EVENT_STREAM)
+	 * 		.exchange()
+	 * 		.expectStatus().isOk()
+	 *	 	.returnResult()
+	 *	 	.consumeWith(result -> assertThat(...);
+	 * </pre>
+	 * @param consumer consumer for {@code "this"} instance
+	 */
+	public void consumeWith(Consumer<FluxExchangeResult<T>> consumer) {
+		assertWithDiagnostics(() -> consumer.accept(this));
+	}
+
 }

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

@@ -29,6 +29,7 @@
 
 import org.springframework.context.ApplicationContext;
 import org.springframework.core.ParameterizedTypeReference;
+import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.format.FormatterRegistry;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
@@ -663,6 +664,15 @@ interface ResponseSpec {
 		 * Variant of {@link #returnResult(Class)} for element types with generics.
 		 */
 		<T> FluxExchangeResult<T> returnResult(ParameterizedTypeReference<T> elementType);
+
+		/**
+		 * Return the exchange result with the body decoded to
+		 * {@code Flux<DataBuffer>}. Use this option for infinite streams and
+		 * consume the stream with the {@code StepVerifier} from the Reactor Add-Ons.
+		 *
+		 * @return
+		 */
+		FluxExchangeResult<DataBuffer> returnResult();
 	}
 
 	/**