Polish pattern-based header removal contribution

Closes gh-191

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

@@ -102,11 +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` applies the given patterns on every header and
-removes them when matching.
+`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,10 +16,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 java.util.Iterator;
 
 import org.springframework.http.HttpHeaders;
 import org.springframework.restdocs.operation.OperationRequest;
@@ -37,19 +34,13 @@ import org.springframework.restdocs.operation.OperationResponseFactory;
 class HeaderRemovingOperationPreprocessor implements OperationPreprocessor {
 
 	private final OperationRequestFactory requestFactory = new OperationRequestFactory();
+
 	private final OperationResponseFactory responseFactory = new OperationResponseFactory();
 
-	private final Set<String> plainHeadersToRemove;
-	private final Set<Pattern> patternHeadersToRemove;
+	private final HeaderFilter headerFilter;
 
-	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));
+	HeaderRemovingOperationPreprocessor(HeaderFilter headerFilter) {
+		this.headerFilter = headerFilter;
 	}
 
 	@Override
@@ -67,24 +58,10 @@ class HeaderRemovingOperationPreprocessor implements OperationPreprocessor {
 	private HttpHeaders removeHeaders(HttpHeaders originalHeaders) {
 		HttpHeaders processedHeaders = new HttpHeaders();
 		processedHeaders.putAll(originalHeaders);
-		if (this.plainHeadersToRemove != null) {
-			for (String headerToRemove : this.plainHeadersToRemove) {
-				processedHeaders.remove(headerToRemove);
-			}
-		}
-		else {
-			Set<String> 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);
+		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,29 +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 headers from the request
-	 * or response based on a pattern match.
+	 * 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 headerPatternsToRemove pattern for the header names to remove. Every matchig header will be removed.
+	 * @param headerNamePatterns the header name patterns
 	 * @return the preprocessor
+	 * @see java.util.regex.Matcher#matches()
 	 */
-	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);
+	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.
@@ -19,7 +19,6 @@ 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;
@@ -41,7 +40,7 @@ import static org.junit.Assert.assertThat;
  * Tests for {@link HeaderRemovingOperationPreprocessorTests}.
  *
  * @author Andy Wilkinson
- *
+ * @author Roland Huss
  */
 public class HeaderRemovingOperationPreprocessorTests {
 
@@ -50,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() {
@@ -76,27 +75,29 @@ public class HeaderRemovingOperationPreprocessorTests {
 	@Test
 	public void modifyWithPattern() {
 		OperationResponse response = createResponse("content-length", "1234");
-		HeaderRemovingOperationPreprocessor processor =
-			new HeaderRemovingOperationPreprocessor(Pattern.compile("co.*le(.)gth]"));
+		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")));
+		assertThat(preprocessed.getHeaders(),
+				hasEntry("b", Arrays.asList("bravo", "banana")));
 	}
 
 	@Test
 	public void removeAllHeaders() {
-		HeaderRemovingOperationPreprocessor processor =
-			new HeaderRemovingOperationPreprocessor(Pattern.compile(".*"));
+		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 OperationResponse createResponse(String... extraHeaders) {
+		return this.responseFactory.create(HttpStatus.OK, getHttpHeaders(extraHeaders),
+				new byte[0]);
 	}
 
-	private HttpHeaders getHttpHeaders(String ... extraHeaders) {
+	private HttpHeaders getHttpHeaders(String... extraHeaders) {
 		HttpHeaders httpHeaders = new HttpHeaders();
 		httpHeaders.add("a", "alpha");
 		httpHeaders.add("b", "bravo");