Merge branch '5.1.x'

Author: poutsma

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

@@ -26,6 +26,7 @@
 import java.util.Map;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.function.IntPredicate;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 
@@ -416,6 +417,8 @@ public HttpHeaders getHeaders() {
 
 	private static class DefaultResponseSpec implements ResponseSpec {
 
+		private static final IntPredicate STATUS_CODE_ERROR = value -> value >= 400;
+
 		private final Mono<ClientResponse> responseMono;
 
 		private final Supplier<HttpRequest> requestSupplier;
@@ -425,17 +428,30 @@ private static class DefaultResponseSpec implements ResponseSpec {
 		DefaultResponseSpec(Mono<ClientResponse> responseMono, Supplier<HttpRequest> requestSupplier) {
 			this.responseMono = responseMono;
 			this.requestSupplier = requestSupplier;
-			this.statusHandlers.add(new StatusHandler(HttpStatus::isError, ClientResponse::createException));
+			this.statusHandlers.add(new StatusHandler(STATUS_CODE_ERROR, ClientResponse::createException));
 		}
 
 		@Override
 		public ResponseSpec onStatus(Predicate<HttpStatus> statusPredicate,
 				Function<ClientResponse, Mono<? extends Throwable>> exceptionFunction) {
+			return onRawStatus(toIntPredicate(statusPredicate), exceptionFunction);
+		}
 
-			Assert.notNull(statusPredicate, "StatusPredicate must not be null");
+		private static IntPredicate toIntPredicate(Predicate<HttpStatus> predicate) {
+			return value -> {
+				HttpStatus status = HttpStatus.resolve(value);
+				return (status != null) && predicate.test(status);
+			};
+		}
+
+		@Override
+		public ResponseSpec onRawStatus(IntPredicate statusCodePredicate,
+				Function<ClientResponse, Mono<? extends Throwable>> exceptionFunction) {
+
+			Assert.notNull(statusCodePredicate, "StatusCodePredicate must not be null");
 			Assert.notNull(exceptionFunction, "Function must not be null");
 
-			this.statusHandlers.add(0, new StatusHandler(statusPredicate, exceptionFunction));
+			this.statusHandlers.add(0, new StatusHandler(statusCodePredicate, exceptionFunction));
 			return this;
 		}
 
@@ -452,17 +468,12 @@ public <T> Mono<T> bodyToMono(ParameterizedTypeReference<T> elementTypeRef) {
 		}
 
 		private <T> Mono<T> handleBodyMono(ClientResponse response, Mono<T> bodyPublisher) {
-			if (HttpStatus.resolve(response.rawStatusCode()) != null) {
-				Mono<T> result = statusHandlers(response);
-				if (result != null) {
-					return result.switchIfEmpty(bodyPublisher);
-				}
-				else {
-					return bodyPublisher;
-				}
+			Mono<T> result = statusHandlers(response);
+			if (result != null) {
+				return result.switchIfEmpty(bodyPublisher);
 			}
 			else {
-				return response.createException().flatMap(Mono::error);
+				return bodyPublisher;
 			}
 		}
 
@@ -479,24 +490,20 @@ public <T> Flux<T> bodyToFlux(ParameterizedTypeReference<T> elementTypeRef) {
 		}
 
 		private <T> Publisher<T> handleBodyFlux(ClientResponse response, Flux<T> bodyPublisher) {
-			if (HttpStatus.resolve(response.rawStatusCode()) != null) {
-				Mono<T> result = statusHandlers(response);
-				if (result != null) {
-					return result.flux().switchIfEmpty(bodyPublisher);
-				}
-				else {
-					return bodyPublisher;
-				}
+			Mono<T> result = statusHandlers(response);
+			if (result != null) {
+				return result.flux().switchIfEmpty(bodyPublisher);
 			}
 			else {
-				return response.createException().flatMap(Mono::error);
+				return bodyPublisher;
 			}
 		}
 
 		@Nullable
 		private <T> Mono<T> statusHandlers(ClientResponse response) {
+			int statusCode = response.rawStatusCode();
 			for (StatusHandler handler : this.statusHandlers) {
-				if (handler.test(response.statusCode())) {
+				if (handler.test(statusCode)) {
 					Mono<? extends Throwable> exMono;
 					try {
 						exMono = handler.apply(response);
@@ -508,7 +515,7 @@ private <T> Mono<T> statusHandlers(ClientResponse response) {
 					}
 					Mono<T> result = exMono.flatMap(Mono::error);
 					HttpRequest request = this.requestSupplier.get();
-					return insertCheckpoint(result, response.statusCode(), request);
+					return insertCheckpoint(result, statusCode, request);
 				}
 			}
 			return null;
@@ -522,10 +529,10 @@ private <T> Mono<T> drainBody(ClientResponse response, Throwable ex) {
 					.onErrorResume(ex2 -> Mono.empty()).thenReturn(ex);
 		}
 
-		private <T> Mono<T> insertCheckpoint(Mono<T> result, HttpStatus status, HttpRequest request) {
+		private <T> Mono<T> insertCheckpoint(Mono<T> result, int statusCode, HttpRequest request) {
 			String httpMethod = request.getMethodValue();
 			URI uri = request.getURI();
-			String description = status + " from " + httpMethod + " " + uri + " [DefaultWebClient]";
+			String description = statusCode + " from " + httpMethod + " " + uri + " [DefaultWebClient]";
 			return result.checkpoint(description);
 		}
 
@@ -558,24 +565,25 @@ public <T> Mono<ResponseEntity<List<T>>> toEntityList(ParameterizedTypeReference
 
 		private static class StatusHandler {
 
-			private final Predicate<HttpStatus> predicate;
+			private final IntPredicate predicate;
 
 			private final Function<ClientResponse, Mono<? extends Throwable>> exceptionFunction;
 
-			public StatusHandler(Predicate<HttpStatus> predicate,
+			public StatusHandler(IntPredicate predicate,
 					Function<ClientResponse, Mono<? extends Throwable>> exceptionFunction) {
 
 				this.predicate = predicate;
 				this.exceptionFunction = exceptionFunction;
 			}
 
-			public boolean test(HttpStatus status) {
+			public boolean test(int status) {
 				return this.predicate.test(status);
 			}
 
 			public Mono<? extends Throwable> apply(ClientResponse response) {
 				return this.exceptionFunction.apply(response);
 			}
+
 		}
 	}
 

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

@@ -23,6 +23,7 @@
 import java.util.Map;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.function.IntPredicate;
 import java.util.function.Predicate;
 
 import org.reactivestreams.Publisher;
@@ -691,6 +692,24 @@ interface ResponseSpec {
 		ResponseSpec onStatus(Predicate<HttpStatus> statusPredicate,
 				Function<ClientResponse, Mono<? extends Throwable>> exceptionFunction);
 
+		/**
+		 * Register a custom error function that gets invoked when the given raw status code
+		 * predicate applies. The exception returned from the function will be returned from
+		 * {@link #bodyToMono(Class)} and {@link #bodyToFlux(Class)}.
+		 * <p>By default, an error handler is registered that throws a
+		 * {@link WebClientResponseException} when the response status code is 4xx or 5xx.
+		 * @param statusCodePredicate a predicate of the raw status code that indicates
+		 * whether {@code exceptionFunction} applies.
+		 * <p><strong>NOTE:</strong> if the response is expected to have content,
+		 * the exceptionFunction should consume it. If not, the content will be
+		 * automatically drained to ensure resources are released.
+		 * @param exceptionFunction the function that returns the exception
+		 * @return this builder
+		 * @since 5.1.9
+		 */
+		ResponseSpec onRawStatus(IntPredicate statusCodePredicate,
+				Function<ClientResponse, Mono<? extends Throwable>> exceptionFunction);
+
 		/**
 		 * Extract the body to a {@code Mono}. By default, if the response has status code 4xx or
 		 * 5xx, the {@code Mono} will contain a {@link WebClientException}. This can be overridden

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

@@ -730,6 +730,28 @@ public void shouldApplyCustomStatusHandler() {
 		});
 	}
 
+	@Test
+	public void shouldApplyCustomRawStatusHandler() {
+		prepareResponse(response -> response.setResponseCode(500)
+				.setHeader("Content-Type", "text/plain").setBody("Internal Server error"));
+
+		Mono<String> result = this.webClient.get()
+				.uri("/greeting?name=Spring")
+				.retrieve()
+				.onRawStatus(value -> value >= 500 && value < 600, response -> Mono.just(new MyException("500 error!")))
+				.bodyToMono(String.class);
+
+		StepVerifier.create(result)
+				.expectError(MyException.class)
+				.verify(Duration.ofSeconds(3));
+
+		expectRequestCount(1);
+		expectRequest(request -> {
+			assertThat(request.getHeader(HttpHeaders.ACCEPT)).isEqualTo("*/*");
+			assertThat(request.getPath()).isEqualTo("/greeting?name=Spring");
+		});
+	}
+
 	@Test
 	public void shouldApplyCustomStatusHandlerParameterizedTypeReference() {
 		prepareResponse(response -> response.setResponseCode(500)