From 2e2d6621586304749372b3cd9888c6955bdf3b53 Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@vmware.com>
Date: Tue, 3 Nov 2020 19:05:46 +0000
Subject: [PATCH] Add toEntityFlux methods

Closes gh-26023
---
 .../function/client/DefaultWebClient.java     | 16 +++++++++
 .../reactive/function/client/WebClient.java   | 24 +++++++++++++
 .../client/WebClientIntegrationTests.java     | 35 ++++++++++++++++++-
 3 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java
index 9431b263b4b..1cf5225376a 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java
@@ -622,6 +622,22 @@ class DefaultWebClient implements WebClient {
 							handleBodyFlux(response, response.bodyToFlux(elementTypeRef))));
 		}
 
+		@Override
+		public <T> Mono<ResponseEntity<Flux<T>>> toEntityFlux(Class<T> elementType) {
+			return this.responseMono.map(response ->
+					ResponseEntity.status(response.rawStatusCode())
+							.headers(response.headers().asHttpHeaders())
+							.body(response.bodyToFlux(elementType)));
+		}
+
+		@Override
+		public <T> Mono<ResponseEntity<Flux<T>>> toEntityFlux(ParameterizedTypeReference<T> elementTypeReference) {
+			return this.responseMono.map(response ->
+					ResponseEntity.status(response.rawStatusCode())
+							.headers(response.headers().asHttpHeaders())
+							.body(response.bodyToFlux(elementTypeReference)));
+		}
+
 		@Override
 		public Mono<ResponseEntity<Void>> toBodilessEntity() {
 			return this.responseMono.flatMap(response ->
diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/WebClient.java b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/WebClient.java
index adc044752ff..d87290ec1e4 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/WebClient.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/WebClient.java
@@ -846,6 +846,30 @@ public interface WebClient {
 		 */
 		<T> Mono<ResponseEntity<T>> toEntity(ParameterizedTypeReference<T> bodyTypeReference);
 
+		/**
+		 * Return a {@code ResponseEntity} with the body decoded to a {@code Flux}
+		 * of elements of the given type. For an error response (status code of
+		 * 4xx or 5xx), the {@code Mono} emits a {@link WebClientException}.
+		 * Use {@link #onStatus(Predicate, Function)} to customize error response
+		 * handling.
+		 * <p><strong>Note:</strong> The {@code Flux} representing the body must
+		 * be subscribed to or else associated resources will not be released.
+		 * @param elementType the type of element to decode the target Flux to
+		 * @param <T> the body element type
+		 * @return the resulting {@code ResponseEntity}
+		 * @since 5.3.1
+		 */
+		<T> Mono<ResponseEntity<Flux<T>>> toEntityFlux(Class<T> elementType);
+
+		/**
+		 * Variant of {@link #toEntity(Class)} with a {@link ParameterizedTypeReference}.
+		 * @param elementTypeReference the type of element to decode the target Flux to
+		 * @param <T> the body element type
+		 * @return the resulting {@code ResponseEntity}
+		 * @since 5.3.1
+		 */
+		<T> Mono<ResponseEntity<Flux<T>>> toEntityFlux(ParameterizedTypeReference<T> elementTypeReference);
+
 		/**
 		 * Return the response as a delayed list of {@code ResponseEntity}s. By default, if the
 		 * response has status code 4xx or 5xx, the {@code Mono} will contain a
diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/WebClientIntegrationTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/WebClientIntegrationTests.java
index 5c61261fb28..4a940669d07 100644
--- a/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/WebClientIntegrationTests.java
+++ b/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/WebClientIntegrationTests.java
@@ -279,7 +279,7 @@ class WebClientIntegrationTests {
 	}
 
 	@ParameterizedWebClientTest
-	void retrieveJsonArrayAsResponseEntity(ClientHttpConnector connector) {
+	void retrieveJsonArrayAsResponseEntityList(ClientHttpConnector connector) {
 		startServer(connector);
 
 		String content = "[{\"bar\":\"bar1\",\"foo\":\"foo1\"}, {\"bar\":\"bar2\",\"foo\":\"foo2\"}]";
@@ -309,6 +309,39 @@ class WebClientIntegrationTests {
 		});
 	}
 
+	@ParameterizedWebClientTest
+	void retrieveJsonArrayAsResponseEntityFlux(ClientHttpConnector connector) {
+		startServer(connector);
+
+		String content = "[{\"bar\":\"bar1\",\"foo\":\"foo1\"}, {\"bar\":\"bar2\",\"foo\":\"foo2\"}]";
+		prepareResponse(response -> response
+				.setHeader("Content-Type", "application/json").setBody(content));
+
+		ResponseEntity<Flux<Pojo>> entity = this.webClient.get()
+				.uri("/json").accept(MediaType.APPLICATION_JSON)
+				.retrieve()
+				.toEntityFlux(Pojo.class)
+				.block(Duration.ofSeconds(3));
+
+		assertThat(entity).isNotNull();
+		assertThat(entity.getStatusCode()).isEqualTo(HttpStatus.OK);
+		assertThat(entity.getHeaders().getContentType()).isEqualTo(MediaType.APPLICATION_JSON);
+		assertThat(entity.getHeaders().getContentLength()).isEqualTo(58);
+
+		assertThat(entity.getBody()).isNotNull();
+		StepVerifier.create(entity.getBody())
+				.expectNext(new Pojo("foo1", "bar1"))
+				.expectNext(new Pojo("foo2", "bar2"))
+				.expectComplete()
+				.verify(Duration.ofSeconds(3));
+
+		expectRequestCount(1);
+		expectRequest(request -> {
+			assertThat(request.getPath()).isEqualTo("/json");
+			assertThat(request.getHeader(HttpHeaders.ACCEPT)).isEqualTo("application/json");
+		});
+	}
+
 	@Test // gh-24788
 	void retrieveJsonArrayAsBodilessEntityShouldReleasesConnection() {