diff --git a/docs/src/docs/asciidoc/customizing-requests-and-responses.adoc b/docs/src/docs/asciidoc/customizing-requests-and-responses.adoc index ec8f2932..06776da4 100644 --- a/docs/src/docs/asciidoc/customizing-requests-and-responses.adoc +++ b/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 diff --git a/spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/preprocess/HeaderRemovingOperationPreprocessor.java b/spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/preprocess/HeaderRemovingOperationPreprocessor.java index 01710e99..f8bf8cc9 100644 --- a/spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/preprocess/HeaderRemovingOperationPreprocessor.java +++ b/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(); diff --git a/spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/preprocess/HeadersModifyingOperationPreprocessor.java b/spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/preprocess/HeadersModifyingOperationPreprocessor.java new file mode 100644 index 00000000..fb002dac --- /dev/null +++ b/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 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 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 values; + + private SetHeaderModification(String name, List 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 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()); + } + + } + +} diff --git a/spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/preprocess/Preprocessors.java b/spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/preprocess/Preprocessors.java index 6be9107c..4fb1647e 100644 --- a/spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/preprocess/Preprocessors.java +++ b/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. diff --git a/spring-restdocs-core/src/test/java/org/springframework/restdocs/operation/preprocess/HeadersModifyingOperationPreprocessorTests.java b/spring-restdocs-core/src/test/java/org/springframework/restdocs/operation/preprocess/HeadersModifyingOperationPreprocessorTests.java new file mode 100644 index 00000000..78129895 --- /dev/null +++ b/spring-restdocs-core/src/test/java/org/springframework/restdocs/operation/preprocess/HeadersModifyingOperationPreprocessorTests.java @@ -0,0 +1,155 @@ +/* + * 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.net.URI; +import java.util.Arrays; +import java.util.Collections; +import java.util.regex.Pattern; + +import org.junit.Test; + +import org.springframework.http.HttpHeaders; +import org.springframework.http.HttpMethod; +import org.springframework.http.HttpStatus; +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.restdocs.operation.Parameters; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Tests for {@link HeadersModifyingOperationPreprocessor}. + * + * @author Jihoon Cha + */ +public class HeadersModifyingOperationPreprocessorTests { + + private final HeadersModifyingOperationPreprocessor preprocessor = new HeadersModifyingOperationPreprocessor(); + + @Test + public void addNewHeader() { + HttpHeaders headers = new HttpHeaders(); + OperationPreprocessor preprocessor = this.preprocessor.add("a", "alpha"); + assertThat(preprocessor.preprocess(createRequest(headers)).getHeaders()).containsEntry("a", + Arrays.asList("alpha")); + assertThat(preprocessor.preprocess(createResponse(headers)).getHeaders()).containsEntry("a", + Arrays.asList("alpha")); + } + + @Test + public void addValueToExistingHeader() { + HttpHeaders headers = new HttpHeaders(); + headers.add("a", "apple"); + OperationPreprocessor preprocessor = this.preprocessor.add("a", "alpha"); + assertThat(preprocessor.preprocess(createRequest(headers)).getHeaders()).containsEntry("a", + Arrays.asList("apple", "alpha")); + assertThat(preprocessor.preprocess(createResponse(headers)).getHeaders()).containsEntry("a", + Arrays.asList("apple", "alpha")); + } + + @Test + public void setNewHeader() { + HttpHeaders headers = new HttpHeaders(); + OperationPreprocessor preprocessor = this.preprocessor.set("a", "alpha", "avocado"); + assertThat(preprocessor.preprocess(createRequest(headers)).getHeaders()).containsEntry("a", + Arrays.asList("alpha", "avocado")); + assertThat(preprocessor.preprocess(createResponse(headers)).getHeaders()).containsEntry("a", + Arrays.asList("alpha", "avocado")); + } + + @Test + public void setExistingHeader() { + HttpHeaders headers = new HttpHeaders(); + headers.add("a", "apple"); + OperationPreprocessor preprocessor = this.preprocessor.set("a", "alpha", "avocado"); + assertThat(preprocessor.preprocess(createRequest(headers)).getHeaders()).containsEntry("a", + Arrays.asList("alpha", "avocado")); + assertThat(preprocessor.preprocess(createResponse(headers)).getHeaders()).containsEntry("a", + Arrays.asList("alpha", "avocado")); + } + + @Test + public void removeNonExistentHeader() { + HttpHeaders headers = new HttpHeaders(); + OperationPreprocessor preprocessor = this.preprocessor.remove("a"); + assertThat(preprocessor.preprocess(createRequest(headers)).getHeaders()).doesNotContainKey("a"); + assertThat(preprocessor.preprocess(createResponse(headers)).getHeaders()).doesNotContainKey("a"); + } + + @Test + public void removeHeader() { + HttpHeaders headers = new HttpHeaders(); + headers.add("a", "apple"); + OperationPreprocessor preprocessor = this.preprocessor.remove("a"); + assertThat(preprocessor.preprocess(createRequest(headers)).getHeaders()).doesNotContainKey("a"); + assertThat(preprocessor.preprocess(createResponse(headers)).getHeaders()).doesNotContainKey("a"); + } + + @Test + public void removeHeaderValueForNonExistentHeader() { + HttpHeaders headers = new HttpHeaders(); + OperationPreprocessor preprocessor = this.preprocessor.remove("a", "apple"); + assertThat(preprocessor.preprocess(createRequest(headers)).getHeaders()).doesNotContainKey("a"); + assertThat(preprocessor.preprocess(createResponse(headers)).getHeaders()).doesNotContainKey("a"); + } + + @Test + public void removeHeaderValueWithMultipleValues() { + HttpHeaders headers = new HttpHeaders(); + headers.add("a", "apple"); + headers.add("a", "alpha"); + OperationPreprocessor preprocessor = this.preprocessor.remove("a", "apple"); + assertThat(preprocessor.preprocess(createRequest(headers)).getHeaders()).containsEntry("a", + Arrays.asList("alpha")); + assertThat(preprocessor.preprocess(createResponse(headers)).getHeaders()).containsEntry("a", + Arrays.asList("alpha")); + } + + @Test + public void removeHeaderValueWithSingleValueRemovesEntryEntirely() { + HttpHeaders headers = new HttpHeaders(); + headers.add("a", "apple"); + OperationPreprocessor preprocessor = this.preprocessor.remove("a", "apple"); + assertThat(preprocessor.preprocess(createRequest(headers)).getHeaders()).doesNotContainKey("a"); + assertThat(preprocessor.preprocess(createResponse(headers)).getHeaders()).doesNotContainKey("a"); + } + + @Test + public void removeHeadersByNamePattern() { + HttpHeaders headers = new HttpHeaders(); + headers.add("apple", "apple"); + headers.add("alpha", "alpha"); + headers.add("avocado", "avocado"); + headers.add("bravo", "bravo"); + OperationPreprocessor preprocessor = this.preprocessor.remove(Pattern.compile("^a.*")); + assertThat(preprocessor.preprocess(createRequest(headers)).getHeaders().size()).isEqualTo(2); + assertThat(preprocessor.preprocess(createResponse(headers)).getHeaders().size()).isEqualTo(1); + } + + private OperationRequest createRequest(HttpHeaders headers) { + return new OperationRequestFactory().create(URI.create("http://localhost:8080"), HttpMethod.GET, new byte[0], + headers, new Parameters(), Collections.emptyList()); + } + + private OperationResponse createResponse(HttpHeaders headers) { + return new OperationResponseFactory().create(HttpStatus.OK.value(), headers, new byte[0]); + } + +}