Introduce EntityResponse

This commit introduces the EntityResponse, an entity-specific subtype of
ServerResponse that exposes the entity itself.

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/DefaultEntityResponseBuilder.java

+/*
+ * 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.web.reactive.function.server;
+
+import java.net.URI;
+import java.time.ZoneId;
+import java.time.ZonedDateTime;
+import java.time.format.DateTimeFormatter;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.LinkedHashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Supplier;
+import java.util.stream.Stream;
+
+import reactor.core.publisher.Mono;
+
+import org.springframework.http.CacheControl;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpMethod;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.MediaType;
+import org.springframework.http.codec.HttpMessageWriter;
+import org.springframework.http.server.reactive.ServerHttpResponse;
+import org.springframework.util.Assert;
+import org.springframework.web.reactive.function.BodyInserter;
+import org.springframework.web.server.ServerWebExchange;
+
+/**
+ * Default {@link EntityResponse.Builder} implementation.
+ *
+ * @author Arjen Poutsma
+ * @since 5.0
+ */
+class DefaultEntityResponseBuilder<T> implements EntityResponse.Builder<T> {
+
+	private final T entity;
+
+	private final BodyInserter<T, ? super ServerHttpResponse> inserter;
+
+	private final HttpHeaders headers = new HttpHeaders();
+
+	private HttpStatus statusCode = HttpStatus.OK;
+
+	private final Map<String, Object> hints = new HashMap<>();
+
+
+
+	public DefaultEntityResponseBuilder(T entity,
+			BodyInserter<T, ? super ServerHttpResponse> inserter) {
+		this.entity = entity;
+		this.inserter = inserter;
+	}
+
+	@Override
+	public EntityResponse.Builder<T> status(HttpStatus status) {
+		Assert.notNull(status, "'status' must not be null");
+		this.statusCode = status;
+		return this;
+	}
+
+	@Override
+	public EntityResponse.Builder<T> header(String headerName, String... headerValues) {
+		for (String headerValue : headerValues) {
+			this.headers.add(headerName, headerValue);
+		}
+		return this;
+	}
+
+	@Override
+	public EntityResponse.Builder<T> headers(HttpHeaders headers) {
+		if (headers != null) {
+			this.headers.putAll(headers);
+		}
+		return this;
+	}
+
+	@Override
+	public EntityResponse.Builder<T> allow(HttpMethod... allowedMethods) {
+		this.headers.setAllow(new LinkedHashSet<>(Arrays.asList(allowedMethods)));
+		return this;
+	}
+
+	@Override
+	public EntityResponse.Builder<T> allow(Set<HttpMethod> allowedMethods) {
+		this.headers.setAllow(allowedMethods);
+		return this;
+	}
+
+	@Override
+	public EntityResponse.Builder<T> contentLength(long contentLength) {
+		this.headers.setContentLength(contentLength);
+		return this;
+	}
+
+	@Override
+	public EntityResponse.Builder<T> contentType(MediaType contentType) {
+		this.headers.setContentType(contentType);
+		return this;
+	}
+
+	@Override
+	public EntityResponse.Builder<T> eTag(String eTag) {
+		if (eTag != null) {
+			if (!eTag.startsWith("\"") && !eTag.startsWith("W/\"")) {
+				eTag = "\"" + eTag;
+			}
+			if (!eTag.endsWith("\"")) {
+				eTag = eTag + "\"";
+			}
+		}
+		this.headers.setETag(eTag);
+		return this;
+	}
+
+	@Override
+	public EntityResponse.Builder<T> hint(String key, Object value) {
+		this.hints.put(key, value);
+		return this;
+	}
+
+	@Override
+	public EntityResponse.Builder<T> lastModified(ZonedDateTime lastModified) {
+		ZonedDateTime gmt = lastModified.withZoneSameInstant(ZoneId.of("GMT"));
+		String headerValue = DateTimeFormatter.RFC_1123_DATE_TIME.format(gmt);
+		this.headers.set(HttpHeaders.LAST_MODIFIED, headerValue);
+		return this;
+	}
+
+	@Override
+	public EntityResponse.Builder<T> location(URI location) {
+		this.headers.setLocation(location);
+		return this;
+	}
+
+	@Override
+	public EntityResponse.Builder<T> cacheControl(CacheControl cacheControl) {
+		String ccValue = cacheControl.getHeaderValue();
+		if (ccValue != null) {
+			this.headers.setCacheControl(cacheControl.getHeaderValue());
+		}
+		return this;
+	}
+
+	@Override
+	public EntityResponse.Builder<T> varyBy(String... requestHeaders) {
+		this.headers.setVary(Arrays.asList(requestHeaders));
+		return this;
+	}
+
+	@Override
+	public Mono<EntityResponse<T>> build() {
+		return Mono.just(new DefaultEntityResponse<T>(this.statusCode, this.headers, this.entity,
+				this.inserter, this.hints));
+	}
+
+	private final static class DefaultEntityResponse<T>
+			extends DefaultServerResponseBuilder.AbstractServerResponse
+			implements EntityResponse<T> {
+
+		private final T entity;
+
+		private final BodyInserter<T, ? super ServerHttpResponse> inserter;
+
+		private final Map<String, Object> hints;
+
+
+		public DefaultEntityResponse(HttpStatus statusCode,
+				HttpHeaders headers, T entity,
+				BodyInserter<T, ? super ServerHttpResponse> inserter,
+				Map<String, Object> hints) {
+			super(statusCode, headers);
+			this.entity = entity;
+			this.inserter = inserter;
+			this.hints = hints;
+		}
+
+
+		@Override
+		public T entity() {
+			return this.entity;
+		}
+
+		@Override
+		public Mono<Void> writeTo(ServerWebExchange exchange, HandlerStrategies strategies) {
+			ServerHttpResponse response = exchange.getResponse();
+			writeStatusAndHeaders(response);
+			return this.inserter.insert(response, new BodyInserter.Context() {
+				@Override
+				public Supplier<Stream<HttpMessageWriter<?>>> messageWriters() {
+					return strategies.messageWriters();
+				}
+
+				@Override
+				public Map<String, Object> hints() {
+					return hints;
+				}
+			});
+		}
+	}
+}

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/DefaultRenderingResponseBuilder.java

@@ -127,7 +127,8 @@ class DefaultRenderingResponseBuilder implements RenderingResponse.Builder {
 				this.model));
 	}
 
-	static class DefaultRenderingResponse extends DefaultServerResponseBuilder.AbstractServerResponse
+	private final static class DefaultRenderingResponse
+			extends DefaultServerResponseBuilder.AbstractServerResponse
 			implements RenderingResponse {
 
 		private final String name;

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/DefaultServerResponseBuilder.java

@@ -177,14 +177,23 @@ class DefaultServerResponseBuilder implements ServerResponse.BodyBuilder {
 	@Override
 	public <T, P extends Publisher<T>> Mono<ServerResponse> body(P publisher,
 			Class<T> elementClass) {
-		return body(BodyInserters.fromPublisher(publisher, elementClass));
+		Assert.notNull(publisher, "'publisher' must not be null");
+		Assert.notNull(elementClass, "'elementClass' must not be null");
+
+		return new DefaultEntityResponseBuilder<>(publisher,
+				BodyInserters.fromPublisher(publisher, elementClass))
+				.headers(this.headers)
+				.status(this.statusCode)
+				.build()
+				.map(entityResponse -> entityResponse);
 	}
 
 	@Override
 	public <T> Mono<ServerResponse> body(BodyInserter<T, ? super ServerHttpResponse> inserter) {
 		Assert.notNull(inserter, "'inserter' must not be null");
 		return Mono
-				.just(new BodyInserterServerResponse<T>(this.statusCode, this.headers, inserter, this.hints));
+				.just(new BodyInserterServerResponse<>(this.statusCode, this.headers, inserter,
+						this.hints));
 	}
 
 	@Override

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/EntityResponse.java

+/*
+ * 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.web.reactive.function.server;
+
+import java.net.URI;
+import java.time.ZonedDateTime;
+import java.util.Set;
+
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Mono;
+
+import org.springframework.core.ResolvableType;
+import org.springframework.http.CacheControl;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpMethod;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.MediaType;
+import org.springframework.http.codec.json.AbstractJackson2Codec;
+import org.springframework.web.reactive.function.BodyInserters;
+
+/**
+ * Entity-specific subtype of {@link ServerResponse} that exposes entity data.
+ *
+ * @author Arjen Poutsma
+ * @since 5.0
+ */
+public interface EntityResponse<T> extends ServerResponse {
+
+	/**
+	 * Return the entity that makes up this response.
+	 */
+	T entity();
+
+	// Static builder methods
+
+	/**
+	 * Create a builder with the given object.
+	 *
+	 * @param t the object that represents the body of the response
+	 * @param <T> the type of the elements contained in the publisher
+	 * @return the created builder
+	 */
+	static <T> Builder<T> fromObject(T t) {
+		return new DefaultEntityResponseBuilder<>(t, BodyInserters.fromObject(t));
+	}
+
+	/**
+	 * Create a builder with the given publisher.
+	 *
+	 * @param publisher the publisher that represents the body of the response
+	 * @param elementClass the class of elements contained in the publisher
+	 * @param <T> the type of the elements contained in the publisher
+	 * @param <P> the type of the {@code Publisher}
+	 * @return the created builder
+	 */
+	static <T, P extends Publisher<T>> Builder<P> fromPublisher(P publisher, Class<T> elementClass) {
+		return new DefaultEntityResponseBuilder<>(publisher,
+				BodyInserters.fromPublisher(publisher, elementClass));
+	}
+
+	/**
+	 * Create a builder with the given publisher.
+	 *
+	 * @param publisher the publisher that represents the body of the response
+	 * @param elementType the type of elements contained in the publisher
+	 * @param <T> the type of the elements contained in the publisher
+	 * @param <P> the type of the {@code Publisher}
+	 * @return the created builder
+	 */
+	static <T, P extends Publisher<T>> Builder<P> fromPublisher(P publisher, ResolvableType elementType) {
+		return new DefaultEntityResponseBuilder<>(publisher,
+				BodyInserters.fromPublisher(publisher, elementType));
+	}
+
+	/**
+	 * Defines a builder for {@code EntityResponse}.
+	 */
+	interface Builder<T> {
+
+		/**
+		 * Add the given header value(s) under the given name.
+		 *
+		 * @param headerName   the header name
+		 * @param headerValues the header value(s)
+		 * @return this builder
+		 * @see HttpHeaders#add(String, String)
+		 */
+		Builder<T> header(String headerName, String... headerValues);
+
+		/**
+		 * Copy the given headers into the entity's headers map.
+		 *
+		 * @param headers the existing HttpHeaders to copy from
+		 * @return this builder
+		 * @see HttpHeaders#add(String, String)
+		 */
+		Builder<T> headers(HttpHeaders headers);
+
+		/**
+		 * Set the status.
+		 * @param status the response status
+		 * @return this builder
+		 */
+		Builder<T> status(HttpStatus status);
+
+		/**
+		 * Set the set of allowed {@link HttpMethod HTTP methods}, as specified
+		 * by the {@code Allow} header.
+		 *
+		 * @param allowedMethods the allowed methods
+		 * @return this builder
+		 * @see HttpHeaders#setAllow(Set)
+		 */
+		Builder<T> allow(HttpMethod... allowedMethods);
+
+		/**
+		 * Set the set of allowed {@link HttpMethod HTTP methods}, as specified
+		 * by the {@code Allow} header.
+		 *
+		 * @param allowedMethods the allowed methods
+		 * @return this builder
+		 * @see HttpHeaders#setAllow(Set)
+		 */
+		Builder<T> allow(Set<HttpMethod> allowedMethods);
+
+		/**
+		 * Set the entity tag of the body, as specified by the {@code ETag} header.
+		 *
+		 * @param eTag the new entity tag
+		 * @return this builder
+		 * @see HttpHeaders#setETag(String)
+		 */
+		Builder<T> eTag(String eTag);
+
+		/**
+		 * Set the time the resource was last changed, as specified by the
+		 * {@code Last-Modified} header.
+		 * <p>The date should be specified as the number of milliseconds since
+		 * January 1, 1970 GMT.
+		 *
+		 * @param lastModified the last modified date
+		 * @return this builder
+		 * @see HttpHeaders#setLastModified(long)
+		 */
+		Builder<T> lastModified(ZonedDateTime lastModified);
+
+		/**
+		 * Set the location of a resource, as specified by the {@code Location} header.
+		 *
+		 * @param location the location
+		 * @return this builder
+		 * @see HttpHeaders#setLocation(URI)
+		 */
+		Builder<T> location(URI location);
+
+		/**
+		 * Set the caching directives for the resource, as specified by the HTTP 1.1
+		 * {@code Cache-Control} header.
+		 * <p>A {@code CacheControl} instance can be built like
+		 * {@code CacheControl.maxAge(3600).cachePublic().noTransform()}.
+		 *
+		 * @param cacheControl a builder for cache-related HTTP response headers
+		 * @return this builder
+		 * @see <a href="https://tools.ietf.org/html/rfc7234#section-5.2">RFC-7234 Section 5.2</a>
+		 */
+		Builder<T> cacheControl(CacheControl cacheControl);
+
+		/**
+		 * Configure one or more request header names (e.g. "Accept-Language") to
+		 * add to the "Vary" response header to inform clients that the response is
+		 * subject to content negotiation and variances based on the value of the
+		 * given request headers. The configured request header names are added only
+		 * if not already present in the response "Vary" header.
+		 *
+		 * @param requestHeaders request header names
+		 * @return this builder
+		 */
+		Builder<T> varyBy(String... requestHeaders);
+
+		/**
+		 * Set the length of the body in bytes, as specified by the
+		 * {@code Content-Length} header.
+		 *
+		 * @param contentLength the content length
+		 * @return this builder
+		 * @see HttpHeaders#setContentLength(long)
+		 */
+		Builder<T> contentLength(long contentLength);
+
+		/**
+		 * Set the {@linkplain MediaType media type} of the body, as specified by the
+		 * {@code Content-Type} header.
+		 *
+		 * @param contentType the content type
+		 * @return this builder
+		 * @see HttpHeaders#setContentType(MediaType)
+		 */
+		Builder<T> contentType(MediaType contentType);
+
+		/**
+		 * Add a serialization hint like {@link AbstractJackson2Codec#JSON_VIEW_HINT} to
+		 * customize how the body will be serialized.
+		 * @param key the hint key
+		 * @param value the hint value
+		 */
+		Builder<T> hint(String key, Object value);
+
+		/**
+		 * Build the response.
+		 *
+		 * @return the built response
+		 */
+		Mono<EntityResponse<T>> build();
+
+	}
+
+}

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RenderingResponse.java

@@ -57,7 +57,7 @@ public interface RenderingResponse extends ServerResponse {
 	}
 
 	/**
-	 * Defines a builder that adds a body to the response.
+	 * Defines a builder for {@code RenderingResponse}.
 	 */
 	interface Builder {
 

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/ServerResponse.java

 /*
- * 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.
@@ -297,7 +297,6 @@ public interface ServerResponse {
 		/**
 		 * Set the length of the body in bytes, as specified by the
 		 * {@code Content-Length} header.
-		 *
 		 * @param contentLength the content length
 		 * @return this builder
 		 * @see HttpHeaders#setContentLength(long)
@@ -307,7 +306,6 @@ public interface ServerResponse {
 		/**
 		 * Set the {@linkplain MediaType media type} of the body, as specified by the
 		 * {@code Content-Type} header.
-		 *
 		 * @param contentType the content type
 		 * @return this builder
 		 * @see HttpHeaders#setContentType(MediaType)
@@ -317,6 +315,8 @@ public interface ServerResponse {
 		/**
 		 * Add a serialization hint like {@link AbstractJackson2Codec#JSON_VIEW_HINT} to
 		 * customize how the body will be serialized.
+		 * @param key the hint key
+		 * @param value the hint value
 		 */
 		BodyBuilder hint(String key, Object value);
 

spring-webflux/src/test/java/org/springframework/web/reactive/function/server/DefaultEntityResponseBuilderTests.java

+/*
+ * 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.web.reactive.function.server;
+
+import java.nio.ByteBuffer;
+import java.time.ZonedDateTime;
+import java.util.Collections;
+import java.util.EnumSet;
+import java.util.List;
+import java.util.Set;
+import java.util.function.BiFunction;
+
+import org.junit.Test;
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import org.springframework.core.ResolvableType;
+import org.springframework.core.codec.CharSequenceEncoder;
+import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DefaultDataBufferFactory;
+import org.springframework.http.CacheControl;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpMethod;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.MediaType;
+import org.springframework.http.codec.EncoderHttpMessageWriter;
+import org.springframework.http.server.reactive.ServerHttpResponse;
+import org.springframework.mock.http.server.reactive.test.MockServerHttpRequest;
+import org.springframework.mock.http.server.reactive.test.MockServerHttpResponse;
+import org.springframework.web.reactive.function.BodyInserter;
+import org.springframework.web.server.ServerWebExchange;
+import org.springframework.web.server.adapter.DefaultServerWebExchange;
+import org.springframework.web.server.session.MockWebSessionManager;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertSame;
+
+/**
+ * @author Arjen Poutsma
+ */
+public class DefaultEntityResponseBuilderTests {
+
+	@Test
+	public void fromObject() throws Exception {
+		String body = "foo";
+		EntityResponse<String> response = EntityResponse.fromObject(body).build().block();
+		assertSame(body, response.entity());
+	}
+
+	@Test
+	public void fromPublisherClass() throws Exception {
+		Flux<String> body = Flux.just("foo", "bar");
+		EntityResponse<Flux<String>> response = EntityResponse.fromPublisher(body, String.class).build().block();
+		assertSame(body, response.entity());
+	}
+
+	@Test
+	public void fromPublisherResolvableType() throws Exception {
+		Flux<String> body = Flux.just("foo", "bar");
+		ResolvableType type = ResolvableType.forClass(String.class);
+		EntityResponse<Flux<String>> response = EntityResponse.fromPublisher(body, type).build().block();
+		assertSame(body, response.entity());
+	}
+
+	@Test
+	public void status() throws Exception {
+		String body = "foo";
+		Mono<EntityResponse<String>> result = EntityResponse.fromObject(body).status(HttpStatus.CREATED).build();
+		StepVerifier.create(result)
+				.expectNextMatches(response -> HttpStatus.CREATED.equals(response.statusCode()))
+				.expectComplete()
+				.verify();
+	}
+
+	@Test
+	public void allow() throws Exception {
+		String body = "foo";
+		Mono<EntityResponse<String>> result = EntityResponse.fromObject(body).allow(HttpMethod.GET).build();
+		Set<HttpMethod> expected = EnumSet.of(HttpMethod.GET);
+		StepVerifier.create(result)
+				.expectNextMatches(response -> expected.equals(response.headers().getAllow()))
+				.expectComplete()
+				.verify();
+	}
+
+	@Test
+	public void contentLength() throws Exception {
+		String body = "foo";
+		Mono<EntityResponse<String>> result = EntityResponse.fromObject(body).contentLength(42).build();
+		StepVerifier.create(result)
+				.expectNextMatches(response -> Long.valueOf(42).equals(response.headers().getContentLength()))
+				.expectComplete()
+				.verify();
+	}
+
+	@Test
+	public void contentType() throws Exception {
+		String body = "foo";
+		Mono<EntityResponse<String>>
+				result = EntityResponse.fromObject(body).contentType(MediaType.APPLICATION_JSON).build();
+		StepVerifier.create(result)
+				.expectNextMatches(response -> MediaType.APPLICATION_JSON.equals(response.headers().getContentType()))
+				.expectComplete()
+				.verify();
+	}
+
+	@Test
+	public void eTag() throws Exception {
+		String body = "foo";
+		Mono<EntityResponse<String>> result = EntityResponse.fromObject(body).eTag("foo").build();
+		StepVerifier.create(result)
+				.expectNextMatches(response -> "\"foo\"".equals(response.headers().getETag()))
+				.expectComplete()
+				.verify();
+	}
+
+
+	@Test
+	public void lastModified() throws Exception {
+		ZonedDateTime now = ZonedDateTime.now();
+		String body = "foo";
+		Mono<EntityResponse<String>> result = EntityResponse.fromObject(body).lastModified(now).build();
+		Long expected = now.toInstant().toEpochMilli() / 1000;
+		StepVerifier.create(result)
+				.expectNextMatches(response -> expected.equals(response.headers().getLastModified() / 1000))
+				.expectComplete()
+				.verify();
+	}
+
+	@Test
+	public void cacheControlTag() throws Exception {
+		String body = "foo";
+		Mono<EntityResponse<String>>
+				result = EntityResponse.fromObject(body).cacheControl(CacheControl.noCache()).build();
+		StepVerifier.create(result)
+				.expectNextMatches(response -> "no-cache".equals(response.headers().getCacheControl()))
+				.expectComplete()
+				.verify();
+	}
+
+	@Test
+	public void varyBy() throws Exception {
+		String body = "foo";
+		Mono<EntityResponse<String>> result = EntityResponse.fromObject(body).varyBy("foo").build();
+		List<String> expected = Collections.singletonList("foo");
+		StepVerifier.create(result)
+				.expectNextMatches(response -> expected.equals(response.headers().getVary()))
+				.expectComplete()
+				.verify();
+	}
+
+	@Test
+	public void headers() throws Exception {
+		String body = "foo";
+		HttpHeaders headers = new HttpHeaders();
+		Mono<EntityResponse<String>> result = EntityResponse.fromObject(body).headers(headers).build();
+		StepVerifier.create(result)
+				.expectNextMatches(response -> headers.equals(response.headers()))
+				.expectComplete()
+				.verify();
+	}
+
+	@Test
+	public void bodyInserter() throws Exception {
+		String body = "foo";
+		Publisher<String> publisher = Mono.just(body);
+		BiFunction<ServerHttpResponse, BodyInserter.Context, Mono<Void>> writer =
+				(response, strategies) -> {
+					byte[] bodyBytes = body.getBytes(UTF_8);
+					ByteBuffer byteBuffer = ByteBuffer.wrap(bodyBytes);
+					DataBuffer buffer = new DefaultDataBufferFactory().wrap(byteBuffer);
+
+					return response.writeWith(Mono.just(buffer));
+				};
+
+		Mono<EntityResponse<Publisher<String>>> result = EntityResponse.fromPublisher(publisher, String.class).build();
+
+		MockServerHttpRequest request = MockServerHttpRequest.get("http://localhost").build();
+		MockServerHttpResponse mockResponse = new MockServerHttpResponse();
+		ServerWebExchange exchange =
+				new DefaultServerWebExchange(request, mockResponse, new MockWebSessionManager());
+
+		HandlerStrategies strategies = HandlerStrategies.empty().messageWriter(new EncoderHttpMessageWriter<>(new CharSequenceEncoder())).build();
+
+		StepVerifier.create(result)
+				.consumeNextWith(response -> {
+					StepVerifier.create(response.entity())
+							.expectNext(body)
+							.expectComplete()
+							.verify();
+					response.writeTo(exchange, strategies);
+				})
+				.expectComplete()
+				.verify();
+
+		assertNotNull(mockResponse.getBody());
+	}
+
+}
\ No newline at end of file

spring-webflux/src/test/java/org/springframework/web/reactive/function/server/DefaultServerResponseBuilderTests.java

 /*
- * 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.
@@ -262,50 +262,4 @@ public class DefaultServerResponseBuilderTests {
 		StepVerifier.create(response.getBody()).expectComplete().verify();
 	}
 
-/*
-TODO: enable when ServerEntityResponse is reintroduced
-
-	@Test
-	public void bodyInserter() throws Exception {
-		String body = "foo";
-		Publisher<String> publisher = Mono.just(body);
-		BiFunction<ServerHttpResponse, BodyInserter.Context, Mono<Void>> writer =
-				(response, strategies) -> {
-					byte[] bodyBytes = body.getBytes(UTF_8);
-					ByteBuffer byteBuffer = ByteBuffer.wrap(bodyBytes);
-					DataBuffer buffer = new DefaultDataBufferFactory().wrap(byteBuffer);
-
-					return response.writeWith(Mono.just(buffer));
-				};
-
-		Mono<ServerResponse> result = ServerResponse.ok().body(BodyInserter.of(writer, publisher));
-
-		MockServerHttpRequest request =
-				new MockServerHttpRequest(HttpMethod.GET, "http://localhost");
-		MockServerHttpResponse mockResponse = new MockServerHttpResponse();
-		ServerWebExchange exchange =
-				new DefaultServerWebExchange(request, mockResponse, new MockWebSessionManager());
-
-		List<HttpMessageWriter<?>> messageWriters = new ArrayList<>();
-		messageWriters.add(new EncoderHttpMessageWriter<CharSequence>(new CharSequenceEncoder()));
-
-		HandlerStrategies strategies = mock(HandlerStrategies.class);
-		when(strategies.messageWriters()).thenReturn(messageWriters::stream);
-
-		StepVerifier.create(result)
-				.consumeNextWith(response -> {
-					StepVerifier.create(response.body())
-							.expectNext(body)
-							.expectComplete()
-							.verify();
-					response.writeTo(exchange, strategies);
-				})
-				.expectComplete()
-				.verify();
-
-		assertNotNull(mockResponse.getBody());
-	}
-*/
-
-
 }
\ No newline at end of file