From 51e02c2911f211ecb21ca64c0a4671d5f6bf176a Mon Sep 17 00:00:00 2001
From: Arjen Poutsma <apoutsma@pivotal.io>
Date: Tue, 4 Jul 2017 16:17:06 +0200
Subject: [PATCH] Add status-based error filter

This commit introduces an ExchangeFilterFunction that throws an
exception given a HTTP status predicate.

Issue: SPR-15724
---
 .../client/ExchangeFilterFunctions.java       | 33 +++++++++++++-
 .../client/ExchangeFilterFunctionsTests.java  | 44 +++++++++++++++++++
 2 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/ExchangeFilterFunctions.java b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/ExchangeFilterFunctions.java
index 0f8481f4c7c..9053fdb9c95 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/ExchangeFilterFunctions.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/function/client/ExchangeFilterFunctions.java
@@ -21,10 +21,12 @@ import java.nio.charset.StandardCharsets;
 import java.util.Base64;
 import java.util.Optional;
 import java.util.function.Function;
+import java.util.function.Predicate;
 
 import reactor.core.publisher.Mono;
 
 import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpStatus;
 import org.springframework.util.Assert;
 
 /**
@@ -110,7 +112,35 @@ public abstract class ExchangeFilterFunctions {
 		return "Basic " + encodedCredentials;
 	}
 
-	private static class Credentials {
+	/**
+	 * Return a filter that returns a given {@link Throwable} as response if the given
+	 * {@link HttpStatus} predicate matches.
+	 * @param statusPredicate the predicate that should match the
+	 * {@linkplain ClientResponse#statusCode() response status}
+	 * @param exceptionFunction the function that returns the exception
+	 * @return the {@link ExchangeFilterFunction} that returns the given exception if the predicate
+	 * matches
+	 */
+	public static ExchangeFilterFunction statusError(Predicate<HttpStatus> statusPredicate,
+			Function<ClientResponse, ? extends Throwable> exceptionFunction) {
+
+		Assert.notNull(statusPredicate, "'statusPredicate' must not be null");
+		Assert.notNull(exceptionFunction, "'exceptionFunction' must not be null");
+
+		return ExchangeFilterFunction.ofResponseProcessor(
+				clientResponse -> {
+					if (statusPredicate.test(clientResponse.statusCode())) {
+						return Mono.error(exceptionFunction.apply(clientResponse));
+					}
+					else {
+						return Mono.just(clientResponse);
+					}
+				}
+		);
+	}
+
+
+	private static final class Credentials {
 
 		private String username;
 
@@ -128,5 +158,4 @@ public abstract class ExchangeFilterFunctions {
 
 	}
 
-
 }
diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/ExchangeFilterFunctionsTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/ExchangeFilterFunctionsTests.java
index 3aaecbb0a72..5e4266207f4 100644
--- a/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/ExchangeFilterFunctionsTests.java
+++ b/spring-webflux/src/test/java/org/springframework/web/reactive/function/client/ExchangeFilterFunctionsTests.java
@@ -20,8 +20,10 @@ import java.net.URI;
 
 import org.junit.Test;
 import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
 
 import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpStatus;
 
 import static org.junit.Assert.*;
 import static org.mockito.Mockito.*;
@@ -133,4 +135,46 @@ public class ExchangeFilterFunctionsTests {
 		assertEquals(response, result);
 	}
 
+	@Test
+	public void statusHandlerMatch() throws Exception {
+		ClientRequest request = ClientRequest.method(GET, URI.create("http://example.com")).build();
+		ClientResponse response = mock(ClientResponse.class);
+		when(response.statusCode()).thenReturn(HttpStatus.NOT_FOUND);
+
+		ExchangeFunction exchange = r -> Mono.just(response);
+
+		ExchangeFilterFunction errorHandler = ExchangeFilterFunctions.statusError(
+				HttpStatus::is4xxClientError, r -> new MyException());
+
+		Mono<ClientResponse> result = errorHandler.filter(request, exchange);
+
+		StepVerifier.create(result)
+				.expectError(MyException.class)
+				.verify();
+	}
+
+	@Test
+	public void statusHandlerNoMatch() throws Exception {
+		ClientRequest request = ClientRequest.method(GET, URI.create("http://example.com")).build();
+		ClientResponse response = mock(ClientResponse.class);
+		when(response.statusCode()).thenReturn(HttpStatus.NOT_FOUND);
+
+		ExchangeFunction exchange = r -> Mono.just(response);
+
+		ExchangeFilterFunction errorHandler = ExchangeFilterFunctions.statusError(
+				HttpStatus::is5xxServerError, r -> new MyException());
+
+		Mono<ClientResponse> result = errorHandler.filter(request, exchange);
+
+		StepVerifier.create(result)
+				.expectNext(response)
+				.expectComplete()
+				.verify();
+	}
+
+	@SuppressWarnings("serial")
+	private static class MyException extends Exception {
+
+	}
+
 }