Add headers(Consumer<HttpHeaders>) to RequestEntity and ResponseEntity

Closes gh-23404

spring-web/src/main/java/org/springframework/http/RequestEntity.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2019 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.
@@ -22,6 +22,7 @@ import java.nio.charset.Charset;
 import java.time.Instant;
 import java.time.ZonedDateTime;
 import java.util.Arrays;
+import java.util.function.Consumer;
 
 import org.springframework.lang.Nullable;
 import org.springframework.util.MultiValueMap;
@@ -315,6 +316,27 @@ public class RequestEntity<T> extends HttpEntity<T> {
 		 */
 		B 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
+		 * @since 5.2
+		 * @see HttpHeaders#add(String, String)
+		 */
+		B headers(@Nullable HttpHeaders headers);
+
+		/**
+		 * Manipulate this entity's headers with the given consumer. The
+		 * headers provided to the consumer are "live", so that the consumer can be used to
+		 * {@linkplain HttpHeaders#set(String, String) overwrite} existing header values,
+		 * {@linkplain HttpHeaders#remove(Object) remove} values, or use any of the other
+		 * {@link HttpHeaders} methods.
+		 * @param headersConsumer a function that consumes the {@code HttpHeaders}
+		 * @return this builder
+		 * @since 5.2
+		 */
+		B headers(Consumer<HttpHeaders> headersConsumer);
+
 		/**
 		 * Set the list of acceptable {@linkplain MediaType media types}, as
 		 * specified by the {@code Accept} header.
@@ -430,6 +452,20 @@ public class RequestEntity<T> extends HttpEntity<T> {
 			return this;
 		}
 
+		@Override
+		public BodyBuilder headers(@Nullable HttpHeaders headers) {
+			if (headers != null) {
+				this.headers.putAll(headers);
+			}
+			return this;
+		}
+
+		@Override
+		public BodyBuilder headers(Consumer<HttpHeaders> headersConsumer) {
+			headersConsumer.accept(this.headers);
+			return this;
+		}
+
 		@Override
 		public BodyBuilder accept(MediaType... acceptableMediaTypes) {
 			this.headers.setAccept(Arrays.asList(acceptableMediaTypes));

spring-web/src/main/java/org/springframework/http/ResponseEntity.java

@@ -23,6 +23,7 @@ import java.util.Arrays;
 import java.util.LinkedHashSet;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.Consumer;
 
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
@@ -333,6 +334,18 @@ public class ResponseEntity<T> extends HttpEntity<T> {
 		 */
 		B headers(@Nullable HttpHeaders headers);
 
+		/**
+		 * Manipulate this entity's headers with the given consumer. The
+		 * headers provided to the consumer are "live", so that the consumer can be used to
+		 * {@linkplain HttpHeaders#set(String, String) overwrite} existing header values,
+		 * {@linkplain HttpHeaders#remove(Object) remove} values, or use any of the other
+		 * {@link HttpHeaders} methods.
+		 * @param headersConsumer a function that consumes the {@code HttpHeaders}
+		 * @return this builder
+		 * @since 5.2
+		 */
+		B headers(Consumer<HttpHeaders> headersConsumer);
+
 		/**
 		 * Set the set of allowed {@link HttpMethod HTTP methods}, as specified
 		 * by the {@code Allow} header.
@@ -481,6 +494,12 @@ public class ResponseEntity<T> extends HttpEntity<T> {
 			return this;
 		}
 
+		@Override
+		public BodyBuilder headers(Consumer<HttpHeaders> headersConsumer) {
+			headersConsumer.accept(this.headers);
+			return this;
+		}
+
 		@Override
 		public BodyBuilder allow(HttpMethod... allowedMethods) {
 			this.headers.setAllow(new LinkedHashSet<>(Arrays.asList(allowedMethods)));

spring-web/src/test/java/org/springframework/http/RequestEntityTests.java

@@ -106,6 +106,7 @@ public class RequestEntityTests {
 				ifNoneMatch(ifNoneMatch).
 				contentLength(contentLength).
 				contentType(contentType).
+				headers(headers -> assertThat(headers).hasSize(6)).
 				build();
 
 		assertThat(responseEntity).isNotNull();

spring-web/src/test/java/org/springframework/http/ResponseEntityTests.java

@@ -171,6 +171,7 @@ public class ResponseEntityTests {
 				location(location).
 				contentLength(contentLength).
 				contentType(contentType).
+				headers(headers -> assertThat(headers).hasSize(5)).
 				build();
 
 		assertThat(responseEntity).isNotNull();
@@ -219,7 +220,7 @@ public class ResponseEntityTests {
 		ResponseEntity<Void> responseEntityWithEmptyHeaders =
 				ResponseEntity.ok().headers(new HttpHeaders()).build();
 		ResponseEntity<Void> responseEntityWithNullHeaders =
-				ResponseEntity.ok().headers(null).build();
+				ResponseEntity.ok().headers((HttpHeaders) null).build();
 
 		assertThat(responseEntityWithEmptyHeaders.getStatusCode()).isEqualTo(HttpStatus.OK);
 		assertThat(responseEntityWithEmptyHeaders.getHeaders().isEmpty()).isTrue();