Support global Consumer<ExchangeResult> in WebTestClient

Closes spring-projectsgh-26662

spring-test/src/main/java/org/springframework/test/web/reactive/server/DefaultWebTestClient.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -79,6 +79,8 @@ class DefaultWebTestClient implements WebTestClient {
 	@Nullable
 	private final MultiValueMap<String, String> defaultCookies;
 
+	private final Consumer<EntityExchangeResult<?>> entityResultConsumer;
+
 	private final Duration responseTimeout;
 
 	private final DefaultWebTestClientBuilder builder;
@@ -89,13 +91,15 @@ class DefaultWebTestClient implements WebTestClient {
 	DefaultWebTestClient(ClientHttpConnector connector,
 			Function<ClientHttpConnector, ExchangeFunction> exchangeFactory, UriBuilderFactory uriBuilderFactory,
 			@Nullable HttpHeaders headers, @Nullable MultiValueMap<String, String> cookies,
+			Consumer<EntityExchangeResult<?>> entityResultConsumer,
 			@Nullable Duration responseTimeout, DefaultWebTestClientBuilder clientBuilder) {
 
 		this.wiretapConnector = new WiretapConnector(connector);
 		this.exchangeFunction = exchangeFactory.apply(this.wiretapConnector);
 		this.uriBuilderFactory = uriBuilderFactory;
 		this.defaultHeaders = headers;
 		this.defaultCookies = cookies;
+		this.entityResultConsumer = entityResultConsumer;
 		this.responseTimeout = (responseTimeout != null ? responseTimeout : Duration.ofSeconds(5));
 		this.builder = clientBuilder;
 	}
@@ -357,7 +361,8 @@ public ResponseSpec exchange() {
 			ExchangeResult result = wiretapConnector.getExchangeResult(
 					this.requestId, this.uriTemplate, getResponseTimeout());
 
-			return new DefaultResponseSpec(result, response, getResponseTimeout());
+			return new DefaultResponseSpec(result, response,
+					DefaultWebTestClient.this.entityResultConsumer, getResponseTimeout());
 		}
 
 		private ClientRequest.Builder initRequestBuilder() {
@@ -408,12 +413,19 @@ private static class DefaultResponseSpec implements ResponseSpec {
 
 		private final ClientResponse response;
 
+		private final Consumer<EntityExchangeResult<?>> entityResultConsumer;
+
 		private final Duration timeout;
 
 
-		DefaultResponseSpec(ExchangeResult exchangeResult, ClientResponse response, Duration timeout) {
+		DefaultResponseSpec(
+				ExchangeResult exchangeResult, ClientResponse response,
+				Consumer<EntityExchangeResult<?>> entityResultConsumer,
+				Duration timeout) {
+
 			this.exchangeResult = exchangeResult;
 			this.response = response;
+			this.entityResultConsumer = entityResultConsumer;
 			this.timeout = timeout;
 		}
 
@@ -435,14 +447,14 @@ public CookieAssertions expectCookie() {
 		@Override
 		public <B> BodySpec<B, ?> expectBody(Class<B> bodyType) {
 			B body = this.response.bodyToMono(bodyType).block(this.timeout);
-			EntityExchangeResult<B> entityResult = new EntityExchangeResult<>(this.exchangeResult, body);
+			EntityExchangeResult<B> entityResult = initEntityExchangeResult(body);
 			return new DefaultBodySpec<>(entityResult);
 		}
 
 		@Override
 		public <B> BodySpec<B, ?> expectBody(ParameterizedTypeReference<B> bodyType) {
 			B body = this.response.bodyToMono(bodyType).block(this.timeout);
-			EntityExchangeResult<B> entityResult = new EntityExchangeResult<>(this.exchangeResult, body);
+			EntityExchangeResult<B> entityResult = initEntityExchangeResult(body);
 			return new DefaultBodySpec<>(entityResult);
 		}
 
@@ -459,18 +471,24 @@ public <E> ListBodySpec<E> expectBodyList(ParameterizedTypeReference<E> elementT
 
 		private <E> ListBodySpec<E> getListBodySpec(Flux<E> flux) {
 			List<E> body = flux.collectList().block(this.timeout);
-			EntityExchangeResult<List<E>> entityResult = new EntityExchangeResult<>(this.exchangeResult, body);
+			EntityExchangeResult<List<E>> entityResult = initEntityExchangeResult(body);
 			return new DefaultListBodySpec<>(entityResult);
 		}
 
 		@Override
 		public BodyContentSpec expectBody() {
 			ByteArrayResource resource = this.response.bodyToMono(ByteArrayResource.class).block(this.timeout);
 			byte[] body = (resource != null ? resource.getByteArray() : null);
-			EntityExchangeResult<byte[]> entityResult = new EntityExchangeResult<>(this.exchangeResult, body);
+			EntityExchangeResult<byte[]> entityResult = initEntityExchangeResult(body);
 			return new DefaultBodyContentSpec(entityResult);
 		}
 
+		private  <B> EntityExchangeResult<B> initEntityExchangeResult(@Nullable B body) {
+			EntityExchangeResult<B> result = new EntityExchangeResult<>(this.exchangeResult, body);
+			result.assertWithDiagnostics(() -> this.entityResultConsumer.accept(result));
+			return result;
+		}
+
 		@Override
 		public <T> FluxExchangeResult<T> returnResult(Class<T> elementClass) {
 			Flux<T> body;

spring-test/src/main/java/org/springframework/test/web/reactive/server/DefaultWebTestClientBuilder.java

@@ -92,6 +92,8 @@ class DefaultWebTestClientBuilder implements WebTestClient.Builder {
 	@Nullable
 	private List<ExchangeFilterFunction> filters;
 
+	private Consumer<EntityExchangeResult<?>> entityResultConsumer = result -> {};
+
 	@Nullable
 	private ExchangeStrategies strategies;
 
@@ -149,6 +151,7 @@ class DefaultWebTestClientBuilder implements WebTestClient.Builder {
 		this.defaultCookies = (other.defaultCookies != null ?
 				new LinkedMultiValueMap<>(other.defaultCookies) : null);
 		this.filters = (other.filters != null ? new ArrayList<>(other.filters) : null);
+		this.entityResultConsumer = other.entityResultConsumer;
 		this.strategies = other.strategies;
 		this.strategiesConfigurers = (other.strategiesConfigurers != null ?
 				new ArrayList<>(other.strategiesConfigurers) : null);
@@ -207,7 +210,7 @@ private MultiValueMap<String, String> initCookies() {
 
 	@Override
 	public WebTestClient.Builder filter(ExchangeFilterFunction filter) {
-		Assert.notNull(filter, "ExchangeFilterFunction must not be null");
+		Assert.notNull(filter, "ExchangeFilterFunction is required");
 		initFilters().add(filter);
 		return this;
 	}
@@ -225,6 +228,13 @@ private List<ExchangeFilterFunction> initFilters() {
 		return this.filters;
 	}
 
+	@Override
+	public WebTestClient.Builder entityExchangeResultConsumer(Consumer<EntityExchangeResult<?>> entityResultConsumer) {
+		Assert.notNull(entityResultConsumer, "`entityResultConsumer` is required");
+		this.entityResultConsumer = this.entityResultConsumer.andThen(entityResultConsumer);
+		return this;
+	}
+
 	@Override
 	public WebTestClient.Builder codecs(Consumer<ClientCodecConfigurer> configurer) {
 		if (this.strategiesConfigurers == null) {
@@ -287,7 +297,7 @@ public WebTestClient build() {
 		return new DefaultWebTestClient(connectorToUse, exchangeFactory, initUriBuilderFactory(),
 				this.defaultHeaders != null ? HttpHeaders.readOnlyHttpHeaders(this.defaultHeaders) : null,
 				this.defaultCookies != null ? CollectionUtils.unmodifiableMultiValueMap(this.defaultCookies) : null,
-				this.responseTimeout, new DefaultWebTestClientBuilder(this));
+				this.entityResultConsumer, this.responseTimeout, new DefaultWebTestClientBuilder(this));
 	}
 
 	private static ClientHttpConnector initConnector() {

spring-test/src/main/java/org/springframework/test/web/reactive/server/WebTestClient.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -438,6 +438,33 @@ interface Builder {
 		 */
 		Builder filters(Consumer<List<ExchangeFilterFunction>> filtersConsumer);
 
+		/**
+		 * Configure an {@code EntityExchangeResult} callback that is invoked
+		 * every time after a response is fully decoded to a single entity, to a
+		 * List of entities, or to a byte[]. In effect, equivalent to each and
+		 * all of the below but registered once, globally:
+		 * <pre>
+		 * client.get().uri("/accounts/1")
+		 *         .exchange()
+		 *         .expectBody(Person.class).consumeWith(exchangeResult -> ... ));
+		 *
+		 * client.get().uri("/accounts")
+		 *         .exchange()
+		 *         .expectBodyList(Person.class).consumeWith(exchangeResult -> ... ));
+		 *
+		 * client.get().uri("/accounts/1")
+		 *         .exchange()
+		 *         .expectBody().consumeWith(exchangeResult -> ... ));
+		 * </pre>
+		 * <p>Note that the configured consumer does not apply to responses
+		 * decoded to {@code Flux<T>} which can be consumed outside the workflow
+		 * of the test client, for example via {@code reactor.test.StepVerifier}.
+		 * @param consumer the consumer to apply to entity responses
+		 * @return the builder
+		 * @since 5.3.5
+		 */
+		Builder entityExchangeResultConsumer(Consumer<EntityExchangeResult<?>> consumer);
+
 		/**
 		 * Configure the codecs for the {@code WebClient} in the
 		 * {@link #exchangeStrategies(ExchangeStrategies) underlying}