From 1897ec54f202f35d0aab2b0b1a39b734374dd4bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Roland=20Hu=C3=9F?= Date: Thu, 4 Feb 2016 11:55:03 +0100 Subject: [PATCH] Add support for removing headers that match a Pattern A new preprocessor method `removeMatchingHeaders` has been introduced which takes one or more regular expression patterns (as Strings). Any header that matches a pattern is removed. Closes gh-195 --- .../customizing-requests-and-responses.adoc | 2 + .../HeaderRemovingOperationPreprocessor.java | 40 +++++++++++++++---- .../operation/preprocess/Preprocessors.java | 17 +++++++- ...derRemovingOperationPreprocessorTests.java | 32 +++++++++++++-- 4 files changed, 80 insertions(+), 11 deletions(-) diff --git a/docs/src/docs/asciidoc/customizing-requests-and-responses.adoc b/docs/src/docs/asciidoc/customizing-requests-and-responses.adoc index 90010638..6011df8d 100644 --- a/docs/src/docs/asciidoc/customizing-requests-and-responses.adoc +++ b/docs/src/docs/asciidoc/customizing-requests-and-responses.adoc @@ -105,6 +105,8 @@ different replacement can also be specified if you wish. `removeHeaders` on `Preprocessors` removes any occurrences of the named headers from the request or response. +`removeMatchingHeaders` on `Preprocessors` applies the given patterns on every header and +removes them when matching. [[customizing-requests-and-responses-preprocessors-replace-patterns]] 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 7766af57..689acc8f 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 @@ -19,6 +19,7 @@ package org.springframework.restdocs.operation.preprocess; import java.util.Arrays; import java.util.HashSet; import java.util.Set; +import java.util.regex.Pattern; import org.springframework.http.HttpHeaders; import org.springframework.restdocs.operation.OperationRequest; @@ -27,20 +28,28 @@ import org.springframework.restdocs.operation.OperationResponse; import org.springframework.restdocs.operation.OperationResponseFactory; /** - * An {@link OperationPreprocessor} that removes headers. + * An {@link OperationPreprocessor} that removes headers. The headers to remove are + * provided as constructor arguments and can be either plain string or patterns to match + * against the headers found * * @author Andy Wilkinson */ class HeaderRemovingOperationPreprocessor implements OperationPreprocessor { private final OperationRequestFactory requestFactory = new OperationRequestFactory(); - private final OperationResponseFactory responseFactory = new OperationResponseFactory(); - private final Set headersToRemove; + private final Set plainHeadersToRemove; + private final Set patternHeadersToRemove; - HeaderRemovingOperationPreprocessor(String... headersToRemove) { - this.headersToRemove = new HashSet<>(Arrays.asList(headersToRemove)); + HeaderRemovingOperationPreprocessor(String ... headersToRemove) { + this.plainHeadersToRemove = new HashSet<>(Arrays.asList(headersToRemove)); + this.patternHeadersToRemove = null; + } + + HeaderRemovingOperationPreprocessor(Pattern ... patternHeadersToRemove) { + this.plainHeadersToRemove = null; + this.patternHeadersToRemove = new HashSet<>(Arrays.asList(patternHeadersToRemove)); } @Override @@ -58,8 +67,25 @@ class HeaderRemovingOperationPreprocessor implements OperationPreprocessor { private HttpHeaders removeHeaders(HttpHeaders originalHeaders) { HttpHeaders processedHeaders = new HttpHeaders(); processedHeaders.putAll(originalHeaders); - for (String headerToRemove : this.headersToRemove) { - processedHeaders.remove(headerToRemove); + if (this.plainHeadersToRemove != null) { + for (String headerToRemove : this.plainHeadersToRemove) { + processedHeaders.remove(headerToRemove); + } + } + else { + Set toRemove = new HashSet<>(); + for (String headerToCheck : originalHeaders.keySet()) { + for (Pattern pattern : this.patternHeadersToRemove) { + if (pattern.matcher(headerToCheck).matches()) { + toRemove.add(headerToCheck); + } + } + } + // Remove afterwards to avoid side effects when removing while iterating over + // the set keys : + for (String headerToRemove : toRemove) { + processedHeaders.remove(headerToRemove); + } } return processedHeaders; } 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 b4aa590d..bacb4fd7 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 @@ -76,13 +76,28 @@ public final class Preprocessors { * Returns an {@code OperationPreprocessor} that will remove headers from the request * or response. * - * @param headersToRemove the names of the headers to remove + * @param headersToRemove the names of the headers to remove. * @return the preprocessor */ public static OperationPreprocessor removeHeaders(String... headersToRemove) { return new HeaderRemovingOperationPreprocessor(headersToRemove); } + /** + * Returns an {@code OperationPreprocessor} that will remove headers from the request + * or response based on a pattern match. + * + * @param headerPatternsToRemove pattern for the header names to remove. Every matchig header will be removed. + * @return the preprocessor + */ + public static OperationPreprocessor removeMatchingHeaders(String... headerPatternsToRemove) { + Pattern[] patterns = new Pattern[headerPatternsToRemove.length]; + for (int i = 0; i < headerPatternsToRemove.length; i++) { + patterns[i] = Pattern.compile(headerPatternsToRemove[i]); + } + return new HeaderRemovingOperationPreprocessor(patterns); + } + /** * Returns an {@code OperationPreprocessor} that will mask the href of hypermedia * links in the request or response. diff --git a/spring-restdocs-core/src/test/java/org/springframework/restdocs/operation/preprocess/HeaderRemovingOperationPreprocessorTests.java b/spring-restdocs-core/src/test/java/org/springframework/restdocs/operation/preprocess/HeaderRemovingOperationPreprocessorTests.java index 634ef616..99a66135 100644 --- a/spring-restdocs-core/src/test/java/org/springframework/restdocs/operation/preprocess/HeaderRemovingOperationPreprocessorTests.java +++ b/spring-restdocs-core/src/test/java/org/springframework/restdocs/operation/preprocess/HeaderRemovingOperationPreprocessorTests.java @@ -19,6 +19,7 @@ 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; @@ -66,18 +67,43 @@ public class HeaderRemovingOperationPreprocessorTests { @Test public void modifyResponseHeaders() { - OperationResponse response = this.responseFactory.create(HttpStatus.OK, - getHttpHeaders(), new byte[0]); + OperationResponse response = createResponse(); OperationResponse preprocessed = this.preprocessor.preprocess(response); assertThat(preprocessed.getHeaders().size(), is(equalTo(1))); assertThat(preprocessed.getHeaders(), hasEntry("a", Arrays.asList("alpha"))); } - private HttpHeaders getHttpHeaders() { + @Test + public void modifyWithPattern() { + OperationResponse response = createResponse("content-length", "1234"); + HeaderRemovingOperationPreprocessor processor = + new HeaderRemovingOperationPreprocessor(Pattern.compile("co.*le(.)gth]")); + OperationResponse preprocessed = processor.preprocess(response); + assertThat(preprocessed.getHeaders().size(), is(equalTo(2))); + assertThat(preprocessed.getHeaders(), hasEntry("a", Arrays.asList("alpha"))); + assertThat(preprocessed.getHeaders(), hasEntry("b", Arrays.asList("bravo", "banana"))); + } + + @Test + public void removeAllHeaders() { + HeaderRemovingOperationPreprocessor processor = + new HeaderRemovingOperationPreprocessor(Pattern.compile(".*")); + OperationResponse preprocessed = processor.preprocess(createResponse()); + assertThat(preprocessed.getHeaders().size(), is(equalTo(0))); + } + + private OperationResponse createResponse(String ... extraHeaders) { + return this.responseFactory.create(HttpStatus.OK, getHttpHeaders(extraHeaders), new byte[0]); + } + + private HttpHeaders getHttpHeaders(String ... extraHeaders) { HttpHeaders httpHeaders = new HttpHeaders(); httpHeaders.add("a", "alpha"); httpHeaders.add("b", "bravo"); httpHeaders.add("b", "banana"); + for (int i = 0; i < extraHeaders.length; i += 2) { + httpHeaders.add(extraHeaders[i], extraHeaders[i + 1]); + } return httpHeaders; }