Add support for post-processing a response before it is documented

This commit adds support for post-processing a response before it is
documented. Post-processors are provided to pretty-print JSON and XML
responses, remove headers from the response, and mask the hrefs of
Atom- and HAL-formatted links in JSON responses.

Response post-processing is configured prior to configuring the
documentation. For example,

mockMvc.perform(get("/").accept(MediaType.APPLICATION_JSON))
        .andExpect(status().isOk())
		.andDo(modifyResponseTo(prettyPrintContent(), removeHeaders("a"),
				maskLinks()).andDocument("post-processed"));

Closes gh-61
Closes gh-31
Closes gh-54

build.gradle

@@ -59,13 +59,13 @@ project(':spring-restdocs') {
 	cleanEclipseJdt.onlyIf { false }
 
 	dependencies {
+		compile "com.fasterxml.jackson.core:jackson-databind:$jacksonVersion"
 		compile "junit:junit:$junitVersion"
 		compile "org.springframework:spring-test:$springVersion"
 		compile "org.springframework:spring-web:$springVersion"
+		compile "org.springframework:spring-webmvc:$springVersion"
 		compile "javax.servlet:javax.servlet-api:$servletApiVersion"
-		compile "com.fasterxml.jackson.core:jackson-databind:$jacksonVersion"
 		jacoco "org.jacoco:org.jacoco.agent:$jacocoVersion:runtime"
-		testCompile "org.springframework:spring-webmvc:$springVersion"
 		testCompile "org.springframework.hateoas:spring-hateoas:$springHateoasVersion"
 		testCompile "org.mockito:mockito-core:$mockitoVersion"
 		testCompile "org.hamcrest:hamcrest-core:$hamcrestVersion"

spring-restdocs/src/main/java/org/springframework/restdocs/ResponseModifier.java

@@ -0,0 +1,117 @@
+/*
+ * Copyright 2014-2015 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;
+
+import java.lang.reflect.Method;
+import java.util.Arrays;
+import java.util.List;
+
+import org.springframework.cglib.proxy.Enhancer;
+import org.springframework.cglib.proxy.MethodInterceptor;
+import org.springframework.cglib.proxy.MethodProxy;
+import org.springframework.core.BridgeMethodResolver;
+import org.springframework.mock.web.MockHttpServletResponse;
+import org.springframework.restdocs.response.ResponsePostProcessor;
+import org.springframework.test.web.servlet.MvcResult;
+import org.springframework.util.ReflectionUtils;
+
+/**
+ * Modifies the response in an {@link MvcResult} by applying {@link ResponsePostProcessor
+ * ResponsePostProcessors} to it.
+ * 
+ * @see RestDocumentation#modifyResponseTo(ResponsePostProcessor...)
+ * @author Andy Wilkinson
+ */
+public class ResponseModifier {
+
+	private final List<ResponsePostProcessor> postProcessors;
+
+	ResponseModifier(ResponsePostProcessor... postProcessors) {
+		this.postProcessors = Arrays.asList(postProcessors);
+	}
+
+	/**
+	 * Provides a {@link RestDocumentationResultHandler} that can be used to document the
+	 * request and modified result.
+	 * @param outputDir The directory to which the documentation will be written
+	 * @return the result handler that will produce the documentation
+	 */
+	public RestDocumentationResultHandler andDocument(String outputDir) {
+		return new ResponseModifyingRestDocumentationResultHandler(outputDir);
+	}
+
+	class ResponseModifyingRestDocumentationResultHandler extends
+			RestDocumentationResultHandler {
+
+		public ResponseModifyingRestDocumentationResultHandler(String outputDir) {
+			super(outputDir);
+		}
+
+		@Override
+		public void handle(MvcResult result) throws Exception {
+			super.handle(postProcessResponse(result));
+		}
+
+		MvcResult postProcessResponse(MvcResult result) throws Exception {
+			MockHttpServletResponse response = result.getResponse();
+			for (ResponsePostProcessor postProcessor : ResponseModifier.this.postProcessors) {
+				response = postProcessor.postProcess(response);
+			}
+			return decorateResult(result, response);
+		}
+
+		private MvcResult decorateResult(MvcResult result,
+				MockHttpServletResponse response) {
+			Enhancer enhancer = new Enhancer();
+			enhancer.setSuperclass(MvcResult.class);
+			enhancer.setCallback(new GetResponseMethodInterceptor(response, result));
+			return (MvcResult) enhancer.create();
+		}
+
+		private class GetResponseMethodInterceptor implements MethodInterceptor {
+
+			private final MvcResult delegate;
+
+			private final MockHttpServletResponse response;
+
+			private final Method getResponseMethod = findMethod("getResponse");
+
+			private GetResponseMethodInterceptor(MockHttpServletResponse response,
+					MvcResult delegate) {
+				this.delegate = delegate;
+				this.response = response;
+			}
+
+			@Override
+			public Object intercept(Object proxy, Method method, Object[] args,
+					MethodProxy methodProxy) throws Throwable {
+				if (this.getResponseMethod.equals(method)) {
+					return this.response;
+				}
+				return method.invoke(this.delegate, args);
+			}
+
+			private Method findMethod(String methodName) {
+				return BridgeMethodResolver.findBridgedMethod(ReflectionUtils.findMethod(
+						MvcResult.class, methodName));
+			}
+
+		}
+
+	}
+
+}

spring-restdocs/src/main/java/org/springframework/restdocs/RestDocumentation.java

@@ -16,7 +16,10 @@
 
 package org.springframework.restdocs;
 
+import org.springframework.restdocs.response.ResponsePostProcessor;
+import org.springframework.restdocs.response.ResponsePostProcessors;
 import org.springframework.test.web.servlet.MockMvc;
+import org.springframework.test.web.servlet.MvcResult;
 import org.springframework.test.web.servlet.ResultActions;
 
 /**
@@ -42,4 +45,18 @@ public static RestDocumentationResultHandler document(String outputDir) {
 		return new RestDocumentationResultHandler(outputDir);
 	}
 
+	/**
+	 * Enables the modification of the response in a {@link MvcResult} prior to it being
+	 * documented. The modification is performed using the given
+	 * {@code responsePostProcessors}.
+	 * 
+	 * @param responsePostProcessors the post-processors to use to modify the response
+	 * @return the response modifier
+	 * @see ResponsePostProcessors
+	 */
+	public static ResponseModifier modifyResponseTo(
+			ResponsePostProcessor... responsePostProcessors) {
+		return new ResponseModifier(responsePostProcessors);
+	}
+
 }

spring-restdocs/src/main/java/org/springframework/restdocs/RestDocumentationResultHandler.java

@@ -49,7 +49,7 @@ public class RestDocumentationResultHandler implements ResultHandler {
 
 	private final String outputDir;
 
-	private List<ResultHandler> delegates;
+	private List<ResultHandler> delegates = new ArrayList<>();
 
 	RestDocumentationResultHandler(String outputDir) {
 		this.outputDir = outputDir;
@@ -60,13 +60,6 @@ public class RestDocumentationResultHandler implements ResultHandler {
 		this.delegates.add(documentHttpResponse(this.outputDir));
 	}
 
-	@Override
-	public void handle(MvcResult result) throws Exception {
-		for (ResultHandler delegate : this.delegates) {
-			delegate.handle(result);
-		}
-	}
-
 	/**
 	 * Document the links in the response using the given {@code descriptors}. The links
 	 * are extracted from the response based on its content type.
@@ -159,4 +152,12 @@ public RestDocumentationResultHandler withQueryParameters(
 		this.delegates.add(documentQueryParameters(this.outputDir, descriptors));
 		return this;
 	}
+
+	@Override
+	public void handle(MvcResult result) throws Exception {
+		for (ResultHandler delegate : this.delegates) {
+			delegate.handle(result);
+		}
+	}
+
 }

spring-restdocs/src/main/java/org/springframework/restdocs/response/ContentModifyingReponsePostProcessor.java

@@ -0,0 +1,96 @@
+/*
+ * Copyright 2014-2015 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.response;
+
+import java.lang.reflect.Method;
+
+import org.springframework.cglib.proxy.Enhancer;
+import org.springframework.cglib.proxy.MethodInterceptor;
+import org.springframework.cglib.proxy.MethodProxy;
+import org.springframework.core.BridgeMethodResolver;
+import org.springframework.mock.web.MockHttpServletResponse;
+import org.springframework.util.ReflectionUtils;
+
+/**
+ * A base class for {@link ResponsePostProcessor ResponsePostProcessors} that modify the
+ * content of the response.
+ * 
+ * @author Andy Wilkinson
+ */
+public abstract class ContentModifyingReponsePostProcessor implements
+		ResponsePostProcessor {
+
+	@Override
+	public MockHttpServletResponse postProcess(MockHttpServletResponse response)
+			throws Exception {
+		String modifiedContent = modifyContent(response.getContentAsString());
+
+		Enhancer enhancer = new Enhancer();
+		enhancer.setSuperclass(MockHttpServletResponse.class);
+		enhancer.setCallback(new ContentModifyingMethodInterceptor(modifiedContent,
+				response));
+
+		return (MockHttpServletResponse) enhancer.create();
+	}
+
+	/**
+	 * Returns a modified version of the given {@code originalContent}
+	 * 
+	 * @param originalContent the content to modify
+	 * @return the modified content
+	 * @throws Exception if a failure occurs while modifying the content
+	 */
+	protected abstract String modifyContent(String originalContent) throws Exception;
+
+	private static class ContentModifyingMethodInterceptor implements MethodInterceptor {
+
+		private final Method getContentAsStringMethod = findMethod("getContentAsString");
+
+		private final Method getContentAsByteArray = findMethod("getContentAsByteArray");
+
+		private final String modifiedContent;
+
+		private final MockHttpServletResponse delegate;
+
+		public ContentModifyingMethodInterceptor(String modifiedContent,
+				MockHttpServletResponse delegate) {
+			this.modifiedContent = modifiedContent;
+			this.delegate = delegate;
+		}
+
+		@Override
+		public Object intercept(Object proxy, Method method, Object[] args,
+				MethodProxy methodProxy) throws Throwable {
+			if (this.getContentAsStringMethod.equals(method)) {
+				return this.modifiedContent;
+			}
+			if (this.getContentAsByteArray.equals(method)) {
+				throw new UnsupportedOperationException(
+						"Following modification, the response's content should be"
+								+ " accessed as a String");
+			}
+			return method.invoke(this.delegate, args);
+		}
+
+		private static Method findMethod(String methodName) {
+			return BridgeMethodResolver.findBridgedMethod(ReflectionUtils.findMethod(
+					MockHttpServletResponse.class, methodName));
+		}
+
+	}
+
+}