Provide a preprocessor that can modify existing headers

Author: ghoonch4

docs/src/docs/asciidoc/customizing-requests-and-responses.adoc

@@ -125,6 +125,13 @@ Any occurrences that match a regular expression are replaced.
 
 
 
+[[customizing-requests-and-responses-preprocessors-modify-headers]]
+==== Modifying Headers
+
+You can use `modifyHeaders` on `Preprocessors` to add, set, and remove request or response headers.
+
+
+
 [[customizing-requests-and-responses-preprocessors-modify-request-parameters]]
 ==== Modifying Request Parameters
 

spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/preprocess/HeaderRemovingOperationPreprocessor.java

@@ -31,7 +31,9 @@
  * against the headers found
  *
  * @author Andy Wilkinson
+ * @deprecated Use {@link HeadersModifyingOperationPreprocessor} instead
  */
+@Deprecated
 class HeaderRemovingOperationPreprocessor implements OperationPreprocessor {
 
 	private final OperationRequestFactory requestFactory = new OperationRequestFactory();

spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/preprocess/HeadersModifyingOperationPreprocessor.java

@@ -0,0 +1,227 @@
+/*
+ * Copyright 2014-2022 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
+ *
+ *      https://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.restdocs.operation.preprocess;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+import org.springframework.http.HttpHeaders;
+import org.springframework.restdocs.operation.OperationRequest;
+import org.springframework.restdocs.operation.OperationRequestFactory;
+import org.springframework.restdocs.operation.OperationResponse;
+import org.springframework.restdocs.operation.OperationResponseFactory;
+import org.springframework.util.Assert;
+
+/**
+ * An {@link OperationPreprocessor} that can be used to modify a request's
+ * {@link OperationRequest#getHeaders()} by adding, setting, and removing headers.
+ *
+ * @author Jihoon Cha
+ */
+public class HeadersModifyingOperationPreprocessor implements OperationPreprocessor {
+
+	private final OperationRequestFactory requestFactory = new OperationRequestFactory();
+
+	private final OperationResponseFactory responseFactory = new OperationResponseFactory();
+
+	private final List<Modification> modifications = new ArrayList<>();
+
+	@Override
+	public OperationRequest preprocess(OperationRequest request) {
+		HttpHeaders headers = copyHttpHeaders(request.getHeaders());
+		for (Modification modification : this.modifications) {
+			modification.applyTo(headers);
+		}
+		return this.requestFactory.createFrom(request, headers);
+	}
+
+	@Override
+	public OperationResponse preprocess(OperationResponse response) {
+		HttpHeaders headers = copyHttpHeaders(response.getHeaders());
+		for (Modification modification : this.modifications) {
+			modification.applyTo(headers);
+		}
+		return this.responseFactory.createFrom(response, headers);
+	}
+
+	private HttpHeaders copyHttpHeaders(HttpHeaders headers) {
+		HttpHeaders copy = new HttpHeaders();
+		for (String name : headers.keySet()) {
+			List<String> values = headers.get(name);
+			if (values == null) {
+				continue;
+			}
+			copy.put(name, new ArrayList<>(values));
+		}
+		return copy;
+	}
+
+	/**
+	 * Adds a header with the given {@code name} and {@code value}.
+	 * @param name the name
+	 * @param value the value
+	 * @return {@code this}
+	 */
+	public HeadersModifyingOperationPreprocessor add(String name, String value) {
+		this.modifications.add(new AddHeaderModification(name, value));
+		return this;
+	}
+
+	/**
+	 * Sets the header with the given {@code name} to have the given {@code values}.
+	 * @param name the name
+	 * @param values the values
+	 * @return {@code this}
+	 */
+	public HeadersModifyingOperationPreprocessor set(String name, String... values) {
+		Assert.notEmpty(values, "At least one value must be provided");
+		this.modifications.add(new SetHeaderModification(name, Arrays.asList(values)));
+		return this;
+	}
+
+	/**
+	 * Removes the header with the given {@code name}.
+	 * @param name the name of the parameter
+	 * @return {@code this}
+	 */
+	public HeadersModifyingOperationPreprocessor remove(String name) {
+		this.modifications.add(new RemoveHeaderModification(name));
+		return this;
+	}
+
+	/**
+	 * Removes the given {@code value} from the header with the given {@code name}.
+	 * @param name the name
+	 * @param value the value
+	 * @return {@code this}
+	 */
+	public HeadersModifyingOperationPreprocessor remove(String name, String value) {
+		this.modifications.add(new RemoveValueHeaderModification(name, value));
+		return this;
+	}
+
+	/**
+	 * Remove headers that match the given {@code namePattern} regular expression.
+	 * @param namePattern the name pattern
+	 * @return {@code this}
+	 * @see Matcher#matches()
+	 */
+	public HeadersModifyingOperationPreprocessor remove(Pattern namePattern) {
+		this.modifications.add(new RemoveHeadersByNamePatternModification(namePattern));
+		return this;
+	}
+
+	private interface Modification {
+
+		void applyTo(HttpHeaders headers);
+
+	}
+
+	private static final class AddHeaderModification implements Modification {
+
+		private final String name;
+
+		private final String value;
+
+		private AddHeaderModification(String name, String value) {
+			this.name = name;
+			this.value = value;
+		}
+
+		@Override
+		public void applyTo(HttpHeaders headers) {
+			headers.add(this.name, this.value);
+		}
+
+	}
+
+	private static final class SetHeaderModification implements Modification {
+
+		private final String name;
+
+		private final List<String> values;
+
+		private SetHeaderModification(String name, List<String> values) {
+			this.name = name;
+			this.values = values;
+		}
+
+		@Override
+		public void applyTo(HttpHeaders headers) {
+			headers.put(this.name, this.values);
+		}
+
+	}
+
+	private static final class RemoveHeaderModification implements Modification {
+
+		private final String name;
+
+		private RemoveHeaderModification(String name) {
+			this.name = name;
+		}
+
+		@Override
+		public void applyTo(HttpHeaders headers) {
+			headers.remove(this.name);
+		}
+
+	}
+
+	private static final class RemoveValueHeaderModification implements Modification {
+
+		private final String name;
+
+		private final String value;
+
+		private RemoveValueHeaderModification(String name, String value) {
+			this.name = name;
+			this.value = value;
+		}
+
+		@Override
+		public void applyTo(HttpHeaders headers) {
+			List<String> values = headers.get(this.name);
+			if (values != null) {
+				values.remove(this.value);
+				if (values.isEmpty()) {
+					headers.remove(this.name);
+				}
+			}
+		}
+
+	}
+
+	private static final class RemoveHeadersByNamePatternModification implements Modification {
+
+		private final Pattern namePattern;
+
+		private RemoveHeadersByNamePatternModification(Pattern namePattern) {
+			this.namePattern = namePattern;
+		}
+
+		@Override
+		public void applyTo(HttpHeaders headers) {
+			headers.keySet().removeIf((name) -> this.namePattern.matcher(name).matches());
+		}
+
+	}
+
+}

spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/preprocess/Preprocessors.java

@@ -31,6 +31,7 @@
  *
  * @author Andy Wilkinson
  * @author Roland Huss
+ * @author Jihoon Cha
  */
 public final class Preprocessors {
 
@@ -73,8 +74,10 @@ public static OperationPreprocessor prettyPrint() {
 	 * {@code headersToRemove}.
 	 * @param headerNames the header names
 	 * @return the preprocessor
+	 * @deprecated Use {@link #modifyHeaders()} instead
 	 * @see String#equals(Object)
 	 */
+	@Deprecated
 	public static OperationPreprocessor removeHeaders(String... headerNames) {
 		return new HeaderRemovingOperationPreprocessor(new ExactMatchHeaderFilter(headerNames));
 	}
@@ -85,8 +88,10 @@ public static OperationPreprocessor removeHeaders(String... headerNames) {
 	 * {@code headerNamePatterns} regular expressions.
 	 * @param headerNamePatterns the header name patterns
 	 * @return the preprocessor
+	 * @deprecated Use {@link #modifyHeaders()} instead
 	 * @see java.util.regex.Matcher#matches()
 	 */
+	@Deprecated
 	public static OperationPreprocessor removeMatchingHeaders(String... headerNamePatterns) {
 		return new HeaderRemovingOperationPreprocessor(new PatternMatchHeaderFilter(headerNamePatterns));
 	}
@@ -132,6 +137,15 @@ public static ParametersModifyingOperationPreprocessor modifyParameters() {
 		return new ParametersModifyingOperationPreprocessor();
 	}
 
+	/**
+	 * Returns a {@code HeadersModifyingOperationPreprocessor} that can then be configured
+	 * to modify the headers of the request.
+	 * @return the preprocessor
+	 */
+	public static HeadersModifyingOperationPreprocessor modifyHeaders() {
+		return new HeadersModifyingOperationPreprocessor();
+	}
+
 	/**
 	 * Returns a {@code UriModifyingOperationPreprocessor} that will modify URIs in the
 	 * request or response by changing one or more of their host, scheme, and port.