Merge pull request #195 from Roland Huß

* gh-195:
  Polish pattern-based header removal contribution
  Add support for removing headers that match a Pattern

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

@@ -102,9 +102,11 @@ different replacement can also be specified if you wish.
 [[customizing-requests-and-responses-preprocessors-remove-headers]]
 ==== Removing headers
 
-`removeHeaders` on `Preprocessors` removes any occurrences of the named headers
-from the request or response.
+`removeHeaders` on `Preprocessors` removes any headers from the request or response where
+the name is equal to any of the given header names.
 
+`removeMatchingHeaders` on `Preprocessors` removes any headers from the request or
+response where the name matches any of the given regular expression patterns.
 
 
 [[customizing-requests-and-responses-preprocessors-replace-patterns]]

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

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2014-2016 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.restdocs.operation.preprocess;
+
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * A {@link HeaderFilter} that excludes a header if its name is an exact match.
+ *
+ * @author Andy Wilkinson
+ */
+class ExactMatchHeaderFilter implements HeaderFilter {
+
+	private final Set<String> headersToExclude;
+
+	ExactMatchHeaderFilter(String... headersToExclude) {
+		this.headersToExclude = new HashSet<>(Arrays.asList(headersToExclude));
+	}
+
+	@Override
+	public boolean excludeHeader(String name) {
+		return this.headersToExclude.contains(name);
+	}
+
+}

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

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2014-2016 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.restdocs.operation.preprocess;
+
+/**
+ * A strategy for determining whether or not a header should be excluded.
+ *
+ * @author Andy Wilkinson
+ */
+interface HeaderFilter {
+
+	/**
+	 * Called to determine whether a header should be excluded. Return {@code true} to
+	 * exclude a header, otherwise {@code false}.
+	 *
+	 * @param name the name of the header
+	 * @return {@code true} to exclude the header, otherwise {@code false}
+	 */
+	boolean excludeHeader(String name);
+
+}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2014-2015 the original author or authors.
+ * Copyright 2014-2016 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.
@@ -16,9 +16,7 @@
 
 package org.springframework.restdocs.operation.preprocess;
 
-import java.util.Arrays;
-import java.util.HashSet;
-import java.util.Set;
+import java.util.Iterator;
 
 import org.springframework.http.HttpHeaders;
 import org.springframework.restdocs.operation.OperationRequest;
@@ -27,7 +25,9 @@ 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
  */
@@ -37,10 +37,10 @@ class HeaderRemovingOperationPreprocessor implements OperationPreprocessor {
 
 	private final OperationResponseFactory responseFactory = new OperationResponseFactory();
 
-	private final Set<String> headersToRemove;
+	private final HeaderFilter headerFilter;
 
-	HeaderRemovingOperationPreprocessor(String... headersToRemove) {
-		this.headersToRemove = new HashSet<>(Arrays.asList(headersToRemove));
+	HeaderRemovingOperationPreprocessor(HeaderFilter headerFilter) {
+		this.headerFilter = headerFilter;
 	}
 
 	@Override
@@ -58,8 +58,11 @@ class HeaderRemovingOperationPreprocessor implements OperationPreprocessor {
 	private HttpHeaders removeHeaders(HttpHeaders originalHeaders) {
 		HttpHeaders processedHeaders = new HttpHeaders();
 		processedHeaders.putAll(originalHeaders);
-		for (String headerToRemove : this.headersToRemove) {
-			processedHeaders.remove(headerToRemove);
+		Iterator<String> headers = processedHeaders.keySet().iterator();
+		while (headers.hasNext()) {
+			if (this.headerFilter.excludeHeader(headers.next())) {
+				headers.remove();
+			}
 		}
 		return processedHeaders;
 	}

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

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2014-2016 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.restdocs.operation.preprocess;
+
+import java.util.HashSet;
+import java.util.Set;
+import java.util.regex.Pattern;
+
+/**
+ * A {@link HeaderFilter} that excludes a header if its name matches a {@link Pattern}.
+ *
+ * @author Andy Wilkinson
+ * @author Roland Huss
+ */
+class PatternMatchHeaderFilter implements HeaderFilter {
+
+	private Set<Pattern> exclusionPatterns;
+
+	PatternMatchHeaderFilter(String... exclusionPatterns) {
+		this.exclusionPatterns = new HashSet<>();
+		for (String exclusionPattern : exclusionPatterns) {
+			this.exclusionPatterns.add(Pattern.compile(exclusionPattern));
+		}
+	}
+
+	@Override
+	public boolean excludeHeader(String name) {
+		for (Pattern pattern : this.exclusionPatterns) {
+			if (pattern.matcher(name).matches()) {
+				return true;
+			}
+		}
+		return false;
+	}
+
+}

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

@@ -30,6 +30,7 @@ import org.springframework.restdocs.operation.OperationResponse;
  * documented.
  *
  * @author Andy Wilkinson
+ * @author Roland Huss
  */
 public final class Preprocessors {
 
@@ -73,14 +74,32 @@ public final class Preprocessors {
 	}
 
 	/**
-	 * Returns an {@code OperationPreprocessor} that will remove headers from the request
-	 * or response.
+	 * Returns an {@code OperationPreprocessor} that will remove any header from the
+	 * request or response with a name that is equal to one of the given
+	 * {@code headersToRemove}.
 	 *
-	 * @param headersToRemove the names of the headers to remove
+	 * @param headerNames the header names
 	 * @return the preprocessor
+	 * @see String#equals(Object)
 	 */
-	public static OperationPreprocessor removeHeaders(String... headersToRemove) {
-		return new HeaderRemovingOperationPreprocessor(headersToRemove);
+	public static OperationPreprocessor removeHeaders(String... headerNames) {
+		return new HeaderRemovingOperationPreprocessor(new ExactMatchHeaderFilter(
+				headerNames));
+	}
+
+	/**
+	 * Returns an {@code OperationPreprocessor} that will remove any headers from the
+	 * request or response with a name that matches one of the given
+	 * {@code headerNamePatterns} regular expressions.
+	 *
+	 * @param headerNamePatterns the header name patterns
+	 * @return the preprocessor
+	 * @see java.util.regex.Matcher#matches()
+	 */
+	public static OperationPreprocessor removeMatchingHeaders(
+			String... headerNamePatterns) {
+		return new HeaderRemovingOperationPreprocessor(new PatternMatchHeaderFilter(
+				headerNamePatterns));
 	}
 
 	/**

spring-restdocs-core/src/test/java/org/springframework/restdocs/operation/preprocess/HeaderRemovingOperationPreprocessorTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2014-2015 the original author or authors.
+ * Copyright 2014-2016 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.
@@ -40,7 +40,7 @@ import static org.junit.Assert.assertThat;
  * Tests for {@link HeaderRemovingOperationPreprocessorTests}.
  *
  * @author Andy Wilkinson
- *
+ * @author Roland Huss
  */
 public class HeaderRemovingOperationPreprocessorTests {
 
@@ -49,7 +49,7 @@ public class HeaderRemovingOperationPreprocessorTests {
 	private final OperationResponseFactory responseFactory = new OperationResponseFactory();
 
 	private final HeaderRemovingOperationPreprocessor preprocessor = new HeaderRemovingOperationPreprocessor(
-			"b");
+			new ExactMatchHeaderFilter("b"));
 
 	@Test
 	public void modifyRequestHeaders() {
@@ -66,18 +66,45 @@ 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(
+				new PatternMatchHeaderFilter("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(
+				new PatternMatchHeaderFilter(".*"));
+		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;
 	}