WebTestClient assert response body with Consumer<B>

Issue: SPR-15421

Author: rstoyanchev

spring-core/src/main/java/org/springframework/core/codec/ResourceDecoder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2016 the original author or authors.
+ * Copyright 2002-2017 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.
@@ -57,6 +57,13 @@ public boolean canDecode(ResolvableType elementType, MimeType mimeType) {
 	public Flux<Resource> decode(Publisher<DataBuffer> inputStream, ResolvableType elementType,
 			MimeType mimeType, Map<String, Object> hints) {
 
+		return Flux.from(decodeToMono(inputStream, elementType, mimeType, hints));
+	}
+
+	@Override
+	public Mono<Resource> decodeToMono(Publisher<DataBuffer> inputStream, ResolvableType elementType,
+			MimeType mimeType, Map<String, Object> hints) {
+
 		Class<?> clazz = elementType.getRawClass();
 
 		Mono<byte[]> byteArray = Flux.from(inputStream).
@@ -70,13 +77,13 @@ public Flux<Resource> decode(Publisher<DataBuffer> inputStream, ResolvableType e
 
 
 		if (InputStreamResource.class.equals(clazz)) {
-			return Flux.from(byteArray.map(ByteArrayInputStream::new).map(InputStreamResource::new));
+			return Mono.from(byteArray.map(ByteArrayInputStream::new).map(InputStreamResource::new));
 		}
 		else if (clazz.isAssignableFrom(ByteArrayResource.class)) {
-			return Flux.from(byteArray.map(ByteArrayResource::new));
+			return Mono.from(byteArray.map(ByteArrayResource::new));
 		}
 		else {
-			return Flux.error(new IllegalStateException("Unsupported resource class: " + clazz));
+			return Mono.error(new IllegalStateException("Unsupported resource class: " + clazz));
 		}
 	}
 

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

@@ -23,7 +23,9 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.UnaryOperator;
 
@@ -32,13 +34,14 @@
 import reactor.core.publisher.Mono;
 
 import org.springframework.core.ResolvableType;
-import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.ByteArrayResource;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.MediaType;
 import org.springframework.http.client.reactive.ClientHttpConnector;
 import org.springframework.http.client.reactive.ClientHttpRequest;
 import org.springframework.util.Assert;
+import org.springframework.util.MimeType;
 import org.springframework.util.MultiValueMap;
 import org.springframework.web.reactive.function.BodyInserter;
 import org.springframework.web.reactive.function.client.ClientResponse;
@@ -47,9 +50,9 @@
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.util.UriBuilder;
 
+import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.springframework.test.util.AssertionErrors.assertEquals;
 import static org.springframework.test.util.AssertionErrors.assertTrue;
-import static org.springframework.web.reactive.function.BodyExtractors.toDataBuffers;
 import static org.springframework.web.reactive.function.BodyExtractors.toFlux;
 import static org.springframework.web.reactive.function.BodyExtractors.toMono;
 
@@ -292,6 +295,7 @@ private DefaultResponseSpec toResponseSpec(Mono<ClientResponse> mono) {
 			ExchangeResult exchangeResult = wiretapConnector.claimRequest(this.requestId);
 			return new DefaultResponseSpec(exchangeResult, clientResponse, getTimeout());
 		}
+
 	}
 
 
@@ -326,11 +330,12 @@ public <T> FluxExchangeResult<T> decodeToFlux(ResolvableType elementType) {
 			return new FluxExchangeResult<>(this, body, this.timeout);
 		}
 
-		public EntityExchangeResult<Void> decodeToEmpty() {
-			DataBuffer buffer = this.response.body(toDataBuffers()).blockFirst(this.timeout);
-			assertWithDiagnostics(() -> assertTrue("Expected empty body", buffer == null));
-			return new EntityExchangeResult<>(this, null);
+		public EntityExchangeResult<byte[]> decodeToByteArray() {
+			ByteArrayResource resource = this.response.body(toMono(ByteArrayResource.class)).block(this.timeout);
+			byte[] body = (resource != null ? resource.getByteArray() : null);
+			return new EntityExchangeResult<>(this, body);
 		}
+
 	}
 
 
@@ -375,7 +380,7 @@ public <E> ListBodySpec<E> expectBodyList(ResolvableType elementType) {
 
 		@Override
 		public BodyContentSpec expectBody() {
-			return new DefaultBodyContentSpec(this.result);
+			return new DefaultBodyContentSpec(this.result.decodeToByteArray());
 		}
 
 		@Override
@@ -387,6 +392,7 @@ public <T> FluxExchangeResult<T> returnResult(Class<T> elementType) {
 		public <T> FluxExchangeResult<T> returnResult(ResolvableType elementType) {
 			return this.result.decodeToFlux(elementType);
 		}
+
 	}
 
 
@@ -406,11 +412,18 @@ protected EntityExchangeResult<B> getResult() {
 
 		@Override
 		public <T extends S> T isEqualTo(B expected) {
-			Object actual = this.result.getResponseBody();
+			B actual = this.result.getResponseBody();
 			this.result.assertWithDiagnostics(() -> assertEquals("Response body", expected, actual));
 			return self();
 		}
 
+		@Override
+		public <T extends S> T consumeWith(Consumer<B> consumer) {
+			B actual = this.result.getResponseBody();
+			this.result.assertWithDiagnostics(() -> consumer.accept(actual));
+			return self();
+		}
+
 		@SuppressWarnings("unchecked")
 		private <T extends S> T self() {
 			return (T) this;
@@ -420,6 +433,7 @@ private <T extends S> T self() {
 		public EntityExchangeResult<B> returnResult() {
 			return this.result;
 		}
+
 	}
 
 
@@ -465,23 +479,55 @@ public ListBodySpec<E> doesNotContain(E... elements) {
 		public EntityExchangeResult<List<E>> returnResult() {
 			return getResult();
 		}
+
 	}
 
 
 	private static class DefaultBodyContentSpec implements BodyContentSpec {
 
-		private final UndecodedExchangeResult result;
+		private final EntityExchangeResult<byte[]> result;
 
+		private final boolean isEmpty;
 
-		DefaultBodyContentSpec(UndecodedExchangeResult result) {
+
+		DefaultBodyContentSpec(EntityExchangeResult<byte[]> result) {
 			this.result = result;
+			this.isEmpty = (result.getResponseBody() == null);
 		}
 
 
 		@Override
 		public EntityExchangeResult<Void> isEmpty() {
-			return this.result.decodeToEmpty();
+			this.result.assertWithDiagnostics(() -> assertTrue("Expected empty body", this.isEmpty));
+			return new EntityExchangeResult<>(this.result, null);
 		}
+
+		@Override
+		public BodyContentSpec consumeAsStringWith(Consumer<String> consumer) {
+			this.result.assertWithDiagnostics(() -> consumer.accept(getBodyAsString()));
+			return this;
+		}
+
+		private String getBodyAsString() {
+			if (this.isEmpty) {
+				return null;
+			}
+			MediaType mediaType = this.result.getResponseHeaders().getContentType();
+			Charset charset = Optional.ofNullable(mediaType).map(MimeType::getCharset).orElse(UTF_8);
+			return new String(this.result.getResponseBody(), charset);
+		}
+
+		@Override
+		public BodyContentSpec consumeWith(Consumer<byte[]> consumer) {
+			this.result.assertWithDiagnostics(() -> consumer.accept(this.result.getResponseBody()));
+			return this;
+		}
+
+		@Override
+		public EntityExchangeResult<byte[]> returnResult() {
+			return this.result;
+		}
+
 	}
 
 }

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

@@ -486,7 +486,7 @@ interface RequestBodySpec extends RequestHeadersSpec<RequestBodySpec> {
 	}
 
 	/**
-	 * Specification for processing the response and applying expectations.
+	 * Spec for declaring expectations on the response.
 	 */
 	interface ResponseSpec {
 
@@ -544,7 +544,7 @@ interface ResponseSpec {
 	}
 
 	/**
-	 * Specification for asserting a response body decoded to a single Object.
+	 * Spec for expectations on the response body decoded to a single Object.
 	 */
 	interface BodySpec<B, S extends BodySpec<B, S>> {
 
@@ -553,6 +553,11 @@ interface BodySpec<B, S extends BodySpec<B, S>> {
 		 */
 		<T extends S> T isEqualTo(B expected);
 
+		/**
+		 * Assert the extracted body with the given {@link Consumer}.
+		 */
+		<T extends S> T consumeWith(Consumer<B> consumer);
+
 		/**
 		 * Return the exchange result with the decoded body.
 		 */
@@ -561,7 +566,7 @@ interface BodySpec<B, S extends BodySpec<B, S>> {
 	}
 
 	/**
-	 * Specification for asserting a response body decoded to a List.
+	 * Spec for expectations on the response body decoded to a List.
 	 */
 	interface ListBodySpec<E> extends BodySpec<List<E>, ListBodySpec<E>> {
 
@@ -587,6 +592,9 @@ interface ListBodySpec<E> extends BodySpec<List<E>, ListBodySpec<E>> {
 
 	}
 
+	/**
+	 * Spec for expectations on the response body content.
+	 */
 	interface BodyContentSpec {
 
 		/**
@@ -595,6 +603,27 @@ interface BodyContentSpec {
 		 */
 		EntityExchangeResult<Void> isEmpty();
 
+		/**
+		 * Assert the response body content converted to a String with the given
+		 * {@link Consumer}. The String is created with the {@link Charset} from
+		 * the "content-type" response header or {@code UTF-8} otherwise.
+		 * @param consumer the consumer for the response body; the input String
+		 * may be {@code null} if there was no response body.
+		 */
+		BodyContentSpec consumeAsStringWith(Consumer<String> consumer);
+
+		/**
+		 * Assert the response body content with the given {@link Consumer}.
+		 * @param consumer the consumer for the response body; the input
+		 * {@code byte[]} may be {@code null} if there was no response body.
+		 */
+		BodyContentSpec consumeWith(Consumer<byte[]> consumer);
+
+		/**
+		 * Return the exchange result with body content as {@code byte[]}.
+		 */
+		EntityExchangeResult<byte[]> returnResult();
+
 	}
 
 }

spring-test/src/test/java/org/springframework/test/web/reactive/server/samples/ResponseEntityTests.java

@@ -42,10 +42,12 @@
 
 import static java.time.Duration.ofMillis;
 import static org.hamcrest.CoreMatchers.endsWith;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThat;
 import static org.springframework.core.ResolvableType.forClassWithGenerics;
 import static org.springframework.http.MediaType.TEXT_EVENT_STREAM;
 
+
 /**
  * Annotated controllers accepting and returning typed Objects.
  *
@@ -66,6 +68,15 @@ public void entity() throws Exception {
 				.expectBody(Person.class).isEqualTo(new Person("John"));
 	}
 
+	@Test
+	public void entityWithConsumer() throws Exception {
+		this.client.get().uri("/persons/John")
+				.exchange()
+				.expectStatus().isOk()
+				.expectHeader().contentType(MediaType.APPLICATION_JSON_UTF8)
+				.expectBody(Person.class).consumeWith(p -> assertEquals(new Person("John"), p));
+	}
+
 	@Test
 	public void entityList() throws Exception {
 

spring-test/src/test/java/org/springframework/test/web/reactive/server/samples/bind/ApplicationContextTests.java

@@ -34,6 +34,8 @@
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebFilter;
 
+import static org.junit.Assert.assertEquals;
+
 /**
  * Binding to server infrastructure declared in a Spring ApplicationContext.
  *
@@ -58,14 +60,23 @@ public void setUp() throws Exception {
 				.build();
 	}
 
+
 	@Test
-	public void basic() throws Exception {
+	public void bodyContent() throws Exception {
 		this.client.get().uri("/principal")
 				.exchange()
 				.expectStatus().isOk()
 				.expectBody(String.class).isEqualTo("Hello Mr. Pablo!");
 	}
 
+	@Test
+	public void bodyContentWithConsumer() throws Exception {
+		this.client.get().uri("/principal")
+				.exchange()
+				.expectStatus().isOk()
+				.expectBody().consumeAsStringWith(body -> assertEquals("Hello Mr. Pablo!", body));
+	}
+
 	@Test
 	public void perRequestExchangeMutator() throws Exception {
 		this.client.exchangeMutator(principal("Giovanni"))

spring-test/src/test/java/org/springframework/test/web/reactive/server/samples/bind/ControllerTests.java

@@ -29,6 +29,8 @@
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebFilter;
 
+import static org.junit.Assert.assertEquals;
+
 /**
  * Bind to annotated controllers.
  *
@@ -44,13 +46,21 @@ public class ControllerTests {
 
 
 	@Test
-	public void basic() throws Exception {
+	public void bodyContent() throws Exception {
 		this.client.get().uri("/principal")
 				.exchange()
 				.expectStatus().isOk()
 				.expectBody(String.class).isEqualTo("Hello Mr. Pablo!");
 	}
 
+	@Test
+	public void bodyContentWithConsumer() throws Exception {
+		this.client.get().uri("/principal")
+				.exchange()
+				.expectStatus().isOk()
+				.expectBody().consumeAsStringWith(body -> assertEquals("Hello Mr. Pablo!", body));
+	}
+
 	@Test
 	public void perRequestExchangeMutator() throws Exception {
 		this.client.exchangeMutator(principal("Giovanni"))