From 794693525f00f9e0ed076727535210315f1a1397 Mon Sep 17 00:00:00 2001
From: Arjen Poutsma <apoutsma@pivotal.io>
Date: Thu, 28 Jun 2018 14:09:34 +0200
Subject: [PATCH] Polishing

---
 .../server/RouterFunctionBuilder.java         | 28 ++++++++---
 .../function/server/RouterFunctions.java      | 50 ++++++++++++++++---
 .../server/RouterFunctionBuilderTests.java    |  6 +--
 3 files changed, 67 insertions(+), 17 deletions(-)

diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RouterFunctionBuilder.java b/spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RouterFunctionBuilder.java
index d88ac793d0..8191b0a372 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RouterFunctionBuilder.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RouterFunctionBuilder.java
@@ -21,6 +21,7 @@ import java.util.List;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.function.Supplier;
 
 import reactor.core.publisher.Mono;
@@ -164,7 +165,7 @@ class RouterFunctionBuilder implements RouterFunctions.Builder {
 	}
 
 	@Override
-	public RouterFunctions.Builder before(
+	public RouterFunctions.Builder filterBefore(
 			Function<ServerRequest, Mono<ServerRequest>> requestProcessor) {
 
 		Assert.notNull(requestProcessor, "Function must not be null");
@@ -172,21 +173,32 @@ class RouterFunctionBuilder implements RouterFunctions.Builder {
 	}
 
 	@Override
-	public RouterFunctions.Builder after(
+	public RouterFunctions.Builder filterAfter(
 			BiFunction<ServerRequest, ServerResponse, Mono<ServerResponse>> responseProcessor) {
 		return filter((request, next) -> next.handle(request)
 				.flatMap(serverResponse -> responseProcessor.apply(request, serverResponse)));
 	}
 
 	@Override
-	public <T extends Throwable> RouterFunctions.Builder exception(
-			Class<T> exceptionType,
-			BiFunction<T, ServerRequest, Mono<ServerResponse>> fallback) {
-		Assert.notNull(exceptionType, "'exceptionType' must not be null");
-		Assert.notNull(fallback, "'fallback' must not be null");
+	public RouterFunctions.Builder filterException(Predicate<? super Throwable> predicate,
+			BiFunction<? super Throwable, ServerRequest, Mono<ServerResponse>> responseProvider) {
+
+		Assert.notNull(predicate, "'exceptionType' must not be null");
+		Assert.notNull(responseProvider, "'fallback' must not be null");
 
 		return filter((request, next) -> next.handle(request)
-				.onErrorResume(exceptionType, t -> fallback.apply(t, request)));
+				.onErrorResume(predicate, t -> responseProvider.apply(t, request)));
+	}
+
+	@Override
+	public <T extends Throwable> RouterFunctions.Builder filterException(
+			Class<T> exceptionType,
+			BiFunction<? super T, ServerRequest, Mono<ServerResponse>> responseProvider) {
+		Assert.notNull(exceptionType, "'exceptionType' must not be null");
+		Assert.notNull(responseProvider, "'fallback' must not be null");
+
+		return filter((request, next) -> next.handle(request)
+				.onErrorResume(exceptionType, t -> responseProvider.apply(t, request)));
 	}
 
 	@Override
diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RouterFunctions.java b/spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RouterFunctions.java
index 2b24b923c1..2c5d94f7f8 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RouterFunctions.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RouterFunctions.java
@@ -23,6 +23,7 @@ import java.util.Map;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.function.Supplier;
 
 import org.apache.commons.logging.Log;
@@ -555,7 +556,7 @@ public abstract class RouterFunctions {
 		 * RouterFunction&lt;ServerResponse&gt; filteredRoute =
 		 *   RouterFunctions.builder()
 		 *     .routeGet("/user", this::listUsers)
-		 *     .before(request -> {
+		 *     .filterBefore(request -> {
 		 *       log(request);
 		 *       return Mono.just(request);
 		 *     })
@@ -564,7 +565,7 @@ public abstract class RouterFunctions {
 		 * @param requestProcessor a function that transforms the request
 		 * @return this builder
 		 */
-		Builder before(Function<ServerRequest, Mono<ServerRequest>> requestProcessor);
+		Builder filterBefore(Function<ServerRequest, Mono<ServerRequest>> requestProcessor);
 
 		/**
 		 * Filters the response for all routes created by this builder with the given response
@@ -576,7 +577,7 @@ public abstract class RouterFunctions {
 		 * RouterFunction&lt;ServerResponse&gt; filteredRoute =
 		 *   RouterFunctions.builder()
 		 *     .routeGet("/user", this::listUsers)
-		 *     .after((request, response) -> {
+		 *     .filterAfter((request, response) -> {
 		 *       log(response);
 		 *       return Mono.just(response);
 		 *     })
@@ -585,10 +586,47 @@ public abstract class RouterFunctions {
 		 * @param responseProcessor a function that transforms the response
 		 * @return this builder
 		 */
-		Builder after(BiFunction<ServerRequest, ServerResponse, Mono<ServerResponse>> responseProcessor);
+		Builder filterAfter(BiFunction<ServerRequest, ServerResponse, Mono<ServerResponse>> responseProcessor);
 
-		<T extends Throwable> Builder exception(Class<T> exceptionType,
-				BiFunction<T, ServerRequest, Mono<ServerResponse>> fallback);
+		/**
+		 * Filters all exceptions that match the predicate by applying the given response provider
+		 * function.
+		 * <p>For instance, the following example creates a filter that returns a 500 response
+		 * status when an {@code IllegalStateException} occurs.
+		 * <pre class="code">
+		 * RouterFunction&lt;ServerResponse&gt; filteredRoute =
+		 *   RouterFunctions.builder()
+		 *     .routeGet("/user", this::listUsers)
+		 *     .filterException(e -> e instanceof IllegalStateException,
+		 *       (e, request) -> ServerResponse.status(HttpStatus.INTERNAL_SERVER_ERROR).build())
+		 *     .build();
+		 * </pre>
+		 * @param predicate the type of exception to filter
+		 * @param responseProvider a function that creates a response
+		 * @return this builder
+		 */
+		Builder filterException(Predicate<? super Throwable> predicate,
+				BiFunction<? super  Throwable, ServerRequest, Mono<ServerResponse>> responseProvider);
+
+		/**
+		 * Filters all exceptions of the given type by applying the given response provider
+		 * function.
+		 * <p>For instance, the following example creates a filter that returns a 500 response
+		 * status when an {@code IllegalStateException} occurs.
+		 * <pre class="code">
+		 * RouterFunction&lt;ServerResponse&gt; filteredRoute =
+		 *   RouterFunctions.builder()
+		 *     .routeGet("/user", this::listUsers)
+		 *     .filterException(IllegalStateException.class,
+		 *       (e, request) -> ServerResponse.status(HttpStatus.INTERNAL_SERVER_ERROR).build())
+		 *     .build();
+		 * </pre>
+		 * @param exceptionType the type of exception to filter
+		 * @param responseProvider a function that creates a response
+		 * @return this builder
+		 */
+		<T extends Throwable> Builder filterException(Class<T> exceptionType,
+				BiFunction<? super T, ServerRequest, Mono<ServerResponse>> responseProvider);
 
 		/**
 		 * Builds the {@code RouterFunction}. All created routes are
diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/function/server/RouterFunctionBuilderTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/function/server/RouterFunctionBuilderTests.java
index 4e014ffe90..e1208c1c72 100644
--- a/spring-webflux/src/test/java/org/springframework/web/reactive/function/server/RouterFunctionBuilderTests.java
+++ b/spring-webflux/src/test/java/org/springframework/web/reactive/function/server/RouterFunctionBuilderTests.java
@@ -102,12 +102,12 @@ public class RouterFunctionBuilderTests {
 		RouterFunction<?> route = RouterFunctions.builder()
 				.routeGet("/foo", request -> ServerResponse.ok().build())
 				.routeGet("/bar", request -> Mono.error(new IllegalStateException()))
-				.before(request -> {
+				.filterBefore(request -> {
 					int count = filterCount.getAndIncrement();
 					assertEquals(0, count);
 					return Mono.just(request);
 				})
-				.after((request, response) -> {
+				.filterAfter((request, response) -> {
 					int count = filterCount.getAndIncrement();
 					assertEquals(3, count);
 					return Mono.just(response);
@@ -120,7 +120,7 @@ public class RouterFunctionBuilderTests {
 					assertEquals(2, count);
 					return responseMono;
 				})
-				.exception(IllegalStateException.class, (e, request) -> ServerResponse.status(HttpStatus.INTERNAL_SERVER_ERROR).build())
+				.filterException(IllegalStateException.class, (e, request) -> ServerResponse.status(HttpStatus.INTERNAL_SERVER_ERROR).build())
 				.build();
 
 		MockServerRequest fooRequest = MockServerRequest.builder().