Add ResponseSpec#toEntityFlux overload with BodyExtractor

See gh-26114

Author: HaloFour

spring-webflux/src/main/java/org/springframework/web/reactive/function/client/DefaultWebClient.java

@@ -43,11 +43,13 @@
 import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;
 import org.springframework.http.client.reactive.ClientHttpRequest;
+import org.springframework.http.client.reactive.ClientHttpResponse;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
 import org.springframework.util.LinkedMultiValueMap;
 import org.springframework.util.MultiValueMap;
+import org.springframework.web.reactive.function.BodyExtractor;
 import org.springframework.web.reactive.function.BodyInserter;
 import org.springframework.web.reactive.function.BodyInserters;
 import org.springframework.web.util.UriBuilder;
@@ -599,6 +601,12 @@ public <T> Mono<ResponseEntity<Flux<T>>> toEntityFlux(ParameterizedTypeReference
 					handlerEntityFlux(response, response.bodyToFlux(elementTypeRef)));
 		}
 
+		@Override
+		public <T> Mono<ResponseEntity<Flux<T>>> toEntityFlux(BodyExtractor<Flux<T>, ? super ClientHttpResponse> bodyExtractor) {
+			return this.responseMono.flatMap(response ->
+					handlerEntityFlux(response, response.body(bodyExtractor)));
+		}
+
 		@Override
 		public Mono<ResponseEntity<Void>> toBodilessEntity() {
 			return this.responseMono.flatMap(response ->

spring-webflux/src/main/java/org/springframework/web/reactive/function/client/WebClient.java

@@ -40,8 +40,10 @@
 import org.springframework.http.ResponseEntity;
 import org.springframework.http.client.reactive.ClientHttpConnector;
 import org.springframework.http.client.reactive.ClientHttpRequest;
+import org.springframework.http.client.reactive.ClientHttpResponse;
 import org.springframework.http.codec.ClientCodecConfigurer;
 import org.springframework.util.MultiValueMap;
+import org.springframework.web.reactive.function.BodyExtractor;
 import org.springframework.web.reactive.function.BodyInserter;
 import org.springframework.web.reactive.function.BodyInserters;
 import org.springframework.web.util.DefaultUriBuilderFactory;
@@ -889,14 +891,23 @@ ResponseSpec onRawStatus(IntPredicate statusCodePredicate,
 		<T> Mono<ResponseEntity<Flux<T>>> toEntityFlux(Class<T> elementType);
 
 		/**
-		 * Variant of {@link #toEntity(Class)} with a {@link ParameterizedTypeReference}.
+		 * Variant of {@link #toEntityFlux(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 {@code ResponseEntity}
 		 * @since 5.3.1
 		 */
 		<T> Mono<ResponseEntity<Flux<T>>> toEntityFlux(ParameterizedTypeReference<T> elementTypeReference);
 
+		/**
+		 * Variant of {@link #toEntityFlux(Class)} with a {@link BodyExtractor}.
+		 * @param bodyExtractor the {@code BodyExtractor} that reads from the response
+		 * @param <T> the body element type
+		 * @return the {@code ResponseEntity}
+		 * @since 5.3.2
+		 */
+		<T> Mono<ResponseEntity<Flux<T>>> toEntityFlux(BodyExtractor<Flux<T>, ? super ClientHttpResponse> bodyExtractor);
+
 		/**
 		 * Return a {@code ResponseEntity} without a body. For an error response
 		 * (status code of 4xx or 5xx), the {@code Mono} emits a

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

@@ -31,6 +31,7 @@
 import org.mockito.junit.jupiter.MockitoSettings;
 import org.mockito.quality.Strictness;
 import org.reactivestreams.Publisher;
+import org.springframework.web.reactive.function.BodyExtractors;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
@@ -459,6 +460,7 @@ public void onStatusHandlersApplyForToEntityMethods() {
 		testStatusHandlerForToEntity(spec.toEntityList(new ParameterizedTypeReference<String>() {}));
 		testStatusHandlerForToEntity(spec.toEntityFlux(String.class));
 		testStatusHandlerForToEntity(spec.toEntityFlux(new ParameterizedTypeReference<String>() {}));
+		testStatusHandlerForToEntity(spec.toEntityFlux(BodyExtractors.toFlux(String.class)));
 	}
 
 	private void testStatusHandlerForToEntity(Publisher<?> responsePublisher) {

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

@@ -41,6 +41,7 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.springframework.web.reactive.function.BodyExtractors;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.netty.http.client.HttpClient;
@@ -342,6 +343,39 @@ void retrieveJsonArrayAsResponseEntityFlux(ClientHttpConnector connector) {
 		});
 	}
 
+	@ParameterizedWebClientTest
+	void retrieveJsonArrayAsResponseEntityFluxWithBodyExtractor(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(BodyExtractors.toFlux(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() {