Copy ClientResponseEntity::toEntity* methods to ResponseSpec

This commit copies the toEntity and toEntityList methods from
ClientResponse to ResponseSpec, so that it is possible to retrieve
a ResponseEntity when using retrieve().

Closes gh-22368

Author: poutsma

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

@@ -157,31 +157,22 @@ public <T> Flux<T> bodyToFlux(ParameterizedTypeReference<T> elementTypeRef) {
 
 	@Override
 	public <T> Mono<ResponseEntity<T>> toEntity(Class<T> bodyType) {
-		return toEntityInternal(bodyToMono(bodyType));
+		return WebClientUtils.toEntity(this, bodyToMono(bodyType));
 	}
 
 	@Override
 	public <T> Mono<ResponseEntity<T>> toEntity(ParameterizedTypeReference<T> bodyTypeReference) {
-		return toEntityInternal(bodyToMono(bodyTypeReference));
-	}
-
-	private <T> Mono<ResponseEntity<T>> toEntityInternal(Mono<T> bodyMono) {
-		HttpHeaders headers = headers().asHttpHeaders();
-		int status = rawStatusCode();
-		return bodyMono
-				.map(body -> createEntity(body, headers, status))
-				.switchIfEmpty(Mono.defer(
-						() -> Mono.just(createEntity(headers, status))));
+		return WebClientUtils.toEntity(this, bodyToMono(bodyTypeReference));
 	}
 
 	@Override
 	public <T> Mono<ResponseEntity<List<T>>> toEntityList(Class<T> elementClass) {
-		return toEntityListInternal(bodyToFlux(elementClass));
+		return WebClientUtils.toEntityList(this, bodyToFlux(elementClass));
 	}
 
 	@Override
 	public <T> Mono<ResponseEntity<List<T>>> toEntityList(ParameterizedTypeReference<T> elementTypeRef) {
-		return toEntityListInternal(bodyToFlux(elementTypeRef));
+		return WebClientUtils.toEntityList(this, bodyToFlux(elementTypeRef));
 	}
 
 	@Override
@@ -224,29 +215,6 @@ HttpRequest request() {
 		return this.requestSupplier.get();
 	}
 
-	private <T> Mono<ResponseEntity<List<T>>> toEntityListInternal(Flux<T> bodyFlux) {
-		HttpHeaders headers = headers().asHttpHeaders();
-		int status = rawStatusCode();
-		return bodyFlux
-				.collectList()
-				.map(body -> createEntity(body, headers, status));
-	}
-
-	private <T> ResponseEntity<T> createEntity(HttpHeaders headers, int status) {
-		HttpStatus resolvedStatus = HttpStatus.resolve(status);
-		return resolvedStatus != null
-				? new ResponseEntity<>(headers, resolvedStatus)
-				: ResponseEntity.status(status).headers(headers).build();
-	}
-
-	private <T> ResponseEntity<T> createEntity(T body, HttpHeaders headers, int status) {
-		HttpStatus resolvedStatus = HttpStatus.resolve(status);
-		return resolvedStatus != null
-				? new ResponseEntity<>(body, headers, resolvedStatus)
-				: ResponseEntity.status(status).headers(headers).body(body);
-	}
-
-
 	private class DefaultHeaders implements Headers {
 
 		private HttpHeaders delegate() {

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

@@ -39,6 +39,7 @@
 import org.springframework.http.HttpRequest;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
+import org.springframework.http.ResponseEntity;
 import org.springframework.http.client.reactive.ClientHttpRequest;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
@@ -528,6 +529,32 @@ private <T> Mono<T> insertCheckpoint(Mono<T> result, HttpStatus status, HttpRequ
 			return result.checkpoint(description);
 		}
 
+		@Override
+		public <T> Mono<ResponseEntity<T>> toEntity(Class<T> bodyClass) {
+			return this.responseMono.flatMap(response ->
+					WebClientUtils.toEntity(response, handleBodyMono(response, response.bodyToMono(bodyClass))));
+		}
+
+		@Override
+		public <T> Mono<ResponseEntity<T>> toEntity(ParameterizedTypeReference<T> bodyTypeReference) {
+			return this.responseMono.flatMap(response ->
+					WebClientUtils.toEntity(response,
+							handleBodyMono(response, response.bodyToMono(bodyTypeReference))));
+		}
+
+		@Override
+		public <T> Mono<ResponseEntity<List<T>>> toEntityList(Class<T> elementClass) {
+			return this.responseMono.flatMap(response ->
+					WebClientUtils.toEntityList(response,
+							handleBodyFlux(response, response.bodyToFlux(elementClass))));
+		}
+
+		@Override
+		public <T> Mono<ResponseEntity<List<T>>> toEntityList(ParameterizedTypeReference<T> elementTypeRef) {
+			return this.responseMono.flatMap(response ->
+					WebClientUtils.toEntityList(response,
+							handleBodyFlux(response, response.bodyToFlux(elementTypeRef))));
+		}
 
 		private static class StatusHandler {
 

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

@@ -35,6 +35,7 @@
 import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
+import org.springframework.http.ResponseEntity;
 import org.springframework.http.client.reactive.ClientHttpConnector;
 import org.springframework.http.client.reactive.ClientHttpRequest;
 import org.springframework.util.MultiValueMap;
@@ -734,6 +735,48 @@ ResponseSpec onStatus(Predicate<HttpStatus> statusPredicate,
 		 */
 		<T> Flux<T> bodyToFlux(ParameterizedTypeReference<T> elementTypeRef);
 
+		/**
+		 * Return the response as a delayed {@code ResponseEntity}. By default, if the response has
+		 * status code 4xx or 5xx, the {@code Mono} will contain a {@link WebClientException}. This
+		 * can be overridden with {@link #onStatus(Predicate, Function)}.
+		 * @param bodyClass the expected response body type
+		 * @param <T> response body type
+		 * @return {@code Mono} with the {@code ResponseEntity}
+		 */
+		<T> Mono<ResponseEntity<T>> toEntity(Class<T> bodyClass);
+
+		/**
+		 * Return the response as a delayed {@code ResponseEntity}. By default, if the response has
+		 * status code 4xx or 5xx, the {@code Mono} will contain a {@link WebClientException}. This
+		 * can be overridden with {@link #onStatus(Predicate, Function)}.
+		 * @param bodyTypeReference a type reference describing the expected response body type
+		 * @param <T> response body type
+		 * @return {@code Mono} with the {@code ResponseEntity}
+		 */
+		<T> Mono<ResponseEntity<T>> toEntity(ParameterizedTypeReference<T> bodyTypeReference);
+
+		/**
+		 * 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
+		 * {@link WebClientException}. This can be overridden with
+		 * {@link #onStatus(Predicate, Function)}.
+		 * @param elementClass the expected response body list element class
+		 * @param <T> the type of elements in the list
+		 * @return {@code Mono} with the list of {@code ResponseEntity}s
+		 */
+		<T> Mono<ResponseEntity<List<T>>> toEntityList(Class<T> elementClass);
+
+		/**
+		 * 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
+		 * {@link WebClientException}. This can be overridden with
+		 * {@link #onStatus(Predicate, Function)}.
+		 * @param elementTypeRef the expected response body list element reference type
+		 * @param <T> the type of elements in the list
+		 * @return {@code Mono} with the list of {@code ResponseEntity}s
+		 */
+		<T> Mono<ResponseEntity<List<T>>> toEntityList(ParameterizedTypeReference<T> elementTypeRef);
+
 	}
 
 

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

@@ -0,0 +1,72 @@
+/*
+ * Copyright 2002-2019 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.web.reactive.function.client;
+
+import java.util.List;
+
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.ResponseEntity;
+import org.springframework.lang.Nullable;
+
+/**
+ * Internal methods shared between {@link DefaultWebClient} and {@link DefaultClientResponse}.
+ *
+ * @author Arjen Poutsma
+ * @since 5.2
+ */
+abstract class WebClientUtils {
+
+	/**
+	 * Create a delayed {@link ResponseEntity} from the given response and body.
+	 */
+	public static  <T> Mono<ResponseEntity<T>> toEntity(ClientResponse response, Mono<T> bodyMono) {
+		return Mono.defer(() -> {
+			HttpHeaders headers = response.headers().asHttpHeaders();
+			int status = response.rawStatusCode();
+			return bodyMono
+					.map(body -> createEntity(body, headers, status))
+					.switchIfEmpty(Mono.defer(
+							() -> Mono.just(createEntity(null, headers, status))));
+		});
+	}
+
+	/**
+	 * Create a delayed {@link ResponseEntity} list from the given response and body.
+	 */
+	public static <T> Mono<ResponseEntity<List<T>>> toEntityList(ClientResponse response, Publisher<T> body) {
+		return Mono.defer(() -> {
+			HttpHeaders headers = response.headers().asHttpHeaders();
+			int status = response.rawStatusCode();
+			return Flux.from(body)
+					.collectList()
+					.map(list -> createEntity(list, headers, status));
+		});
+	}
+
+	private static  <T> ResponseEntity<T> createEntity(@Nullable T body, HttpHeaders headers, int status) {
+		HttpStatus resolvedStatus = HttpStatus.resolve(status);
+		return resolvedStatus != null
+				? new ResponseEntity<>(body, headers, resolvedStatus)
+				: ResponseEntity.status(status).headers(headers).body(body);
+	}
+
+}

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

@@ -219,7 +219,7 @@ public void shouldReceiveJsonAsTypeReferenceString() {
 	}
 
 	@Test
-	public void shouldReceiveJsonAsResponseEntityString() {
+	public void exchangeShouldReceiveJsonAsResponseEntityString() {
 		String content = "{\"bar\":\"barbar\",\"foo\":\"foofoo\"}";
 		prepareResponse(response -> response
 				.setHeader("Content-Type", "application/json").setBody(content));
@@ -246,7 +246,82 @@ public void shouldReceiveJsonAsResponseEntityString() {
 	}
 
 	@Test
-	public void shouldReceiveJsonAsResponseEntityList() {
+	public void retrieveShouldReceiveJsonAsResponseEntityString() {
+		String content = "{\"bar\":\"barbar\",\"foo\":\"foofoo\"}";
+		prepareResponse(response -> response
+				.setHeader("Content-Type", "application/json").setBody(content));
+
+		Mono<ResponseEntity<String>> result = this.webClient.get()
+				.uri("/json").accept(MediaType.APPLICATION_JSON)
+				.retrieve()
+				.toEntity(String.class);
+
+		StepVerifier.create(result)
+				.consumeNextWith(entity -> {
+					assertThat(entity.getStatusCode()).isEqualTo(HttpStatus.OK);
+					assertThat(entity.getHeaders().getContentType()).isEqualTo(MediaType.APPLICATION_JSON);
+					assertThat(entity.getHeaders().getContentLength()).isEqualTo(31);
+					assertThat(entity.getBody()).isEqualTo(content);
+				})
+				.expectComplete().verify(Duration.ofSeconds(3));
+
+		expectRequestCount(1);
+		expectRequest(request -> {
+			assertThat(request.getPath()).isEqualTo("/json");
+			assertThat(request.getHeader(HttpHeaders.ACCEPT)).isEqualTo("application/json");
+		});
+	}
+
+	@Test
+	public void retrieveEntityWithServerError() {
+		prepareResponse(response -> response.setResponseCode(500)
+				.setHeader("Content-Type", "text/plain").setBody("Internal Server error"));
+
+		Mono<ResponseEntity<String>> result = this.webClient.get()
+				.uri("/").accept(MediaType.APPLICATION_JSON)
+				.retrieve()
+				.toEntity(String.class);
+
+		StepVerifier.create(result)
+				.expectError(WebClientResponseException.class)
+				.verify(Duration.ofSeconds(3));
+
+		expectRequestCount(1);
+		expectRequest(request -> {
+			assertThat(request.getPath()).isEqualTo("/");
+			assertThat(request.getHeader(HttpHeaders.ACCEPT)).isEqualTo("application/json");
+		});
+	}
+
+	@Test
+	public void retrieveEntityWithServerErrorStatusHandler() {
+		String content = "Internal Server error";
+		prepareResponse(response -> response.setResponseCode(500)
+				.setHeader("Content-Type", "text/plain").setBody(content));
+
+		Mono<ResponseEntity<String>> result = this.webClient.get()
+				.uri("/").accept(MediaType.APPLICATION_JSON)
+				.retrieve()
+				.onStatus(HttpStatus::is5xxServerError, response -> Mono.empty())// use normal response
+				.toEntity(String.class);
+
+		StepVerifier.create(result)
+				.consumeNextWith(entity -> {
+					assertThat(entity.getStatusCode()).isEqualTo(HttpStatus.INTERNAL_SERVER_ERROR);
+					assertThat(entity.getBody()).isEqualTo(content);
+				})
+				.expectComplete()
+				.verify(Duration.ofSeconds(3));
+
+		expectRequestCount(1);
+		expectRequest(request -> {
+			assertThat(request.getPath()).isEqualTo("/");
+			assertThat(request.getHeader(HttpHeaders.ACCEPT)).isEqualTo("application/json");
+		});
+	}
+
+	@Test
+	public void exchangeShouldReceiveJsonAsResponseEntityList() {
 		String content = "[{\"bar\":\"bar1\",\"foo\":\"foo1\"}, {\"bar\":\"bar2\",\"foo\":\"foo2\"}]";
 		prepareResponse(response -> response
 				.setHeader("Content-Type", "application/json").setBody(content));
@@ -274,6 +349,57 @@ public void shouldReceiveJsonAsResponseEntityList() {
 		});
 	}
 
+	@Test
+	public void retrieveShouldReceiveJsonAsResponseEntityList() {
+		String content = "[{\"bar\":\"bar1\",\"foo\":\"foo1\"}, {\"bar\":\"bar2\",\"foo\":\"foo2\"}]";
+		prepareResponse(response -> response
+				.setHeader("Content-Type", "application/json").setBody(content));
+
+		Mono<ResponseEntity<List<Pojo>>> result = this.webClient.get()
+				.uri("/json").accept(MediaType.APPLICATION_JSON)
+				.retrieve()
+				.toEntityList(Pojo.class);
+
+		StepVerifier.create(result)
+				.consumeNextWith(entity -> {
+					assertThat(entity.getStatusCode()).isEqualTo(HttpStatus.OK);
+					assertThat(entity.getHeaders().getContentType()).isEqualTo(MediaType.APPLICATION_JSON);
+					assertThat(entity.getHeaders().getContentLength()).isEqualTo(58);
+					Pojo pojo1 = new Pojo("foo1", "bar1");
+					Pojo pojo2 = new Pojo("foo2", "bar2");
+					assertThat(entity.getBody()).isEqualTo(Arrays.asList(pojo1, pojo2));
+				})
+				.expectComplete().verify(Duration.ofSeconds(3));
+
+		expectRequestCount(1);
+		expectRequest(request -> {
+			assertThat(request.getPath()).isEqualTo("/json");
+			assertThat(request.getHeader(HttpHeaders.ACCEPT)).isEqualTo("application/json");
+		});
+	}
+
+	@Test
+	public void retrieveEntityListWithServerError() {
+		prepareResponse(response -> response.setResponseCode(500)
+				.setHeader("Content-Type", "text/plain").setBody("Internal Server error"));
+
+		Mono<ResponseEntity<List<String>>> result = this.webClient.get()
+				.uri("/").accept(MediaType.APPLICATION_JSON)
+				.retrieve()
+				.toEntityList(String.class);
+
+		StepVerifier.create(result)
+				.expectError(WebClientResponseException.class)
+				.verify(Duration.ofSeconds(3));
+
+		expectRequestCount(1);
+		expectRequest(request -> {
+			assertThat(request.getPath()).isEqualTo("/");
+			assertThat(request.getHeader(HttpHeaders.ACCEPT)).isEqualTo("application/json");
+		});
+	}
+
+
 	@Test
 	public void shouldReceiveJsonAsFluxString() {
 		String content = "{\"bar\":\"barbar\",\"foo\":\"foofoo\"}";