Support link formats other than HAL by using a strategy to extract links

spring-restdocs-core/src/main/java/org/springframework/restdocs/core/LinkExtractor.java

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2014 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.core;
+
+import java.util.Map;
+
+public interface LinkExtractor {
+
+	Map<String, Object> extractLinks(Map<String, Object> responseJson);
+
+}

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

@@ -20,6 +20,8 @@ import static org.springframework.restdocs.core.RestDocumentationResultHandlers.
 import static org.springframework.restdocs.core.RestDocumentationResultHandlers.documentCurlRequestAndResponse;
 import static org.springframework.restdocs.core.RestDocumentationResultHandlers.documentCurlResponse;
 
+import java.util.Map;
+
 import org.springframework.test.web.servlet.ResultActions;
 
 public class RestDocumentation {
@@ -35,4 +37,15 @@ public class RestDocumentation {
 	public static LinkDescriptor linkWithRel(String rel) {
 		return new LinkDescriptor(rel);
 	}
+
+	public static LinkExtractor halLinks() {
+		return new LinkExtractor() {
+
+			@SuppressWarnings("unchecked")
+			@Override
+			public Map<String, Object> extractLinks(Map<String, Object> responseJson) {
+				return (Map<String, Object>) responseJson.get("_links");
+			}
+		};
+	}
 }

spring-restdocs-core/src/main/java/org/springframework/restdocs/core/RestDocumentationResultActions.java

@@ -53,9 +53,9 @@ public class RestDocumentationResultActions implements ResultActions {
 		return this.delegate.andReturn();
 	}
 
-	public RestDocumentationResultActions andDocumentHalLinks(LinkDescriptor... descriptors)
+	public RestDocumentationResultActions andDocumentLinks(LinkExtractor linkExtractor, LinkDescriptor... descriptors)
 			throws Exception {
-		this.delegate.andDo(new LinkDocumentingResultHandler(this.outputDir, Arrays
+		this.delegate.andDo(new LinkDocumentingResultHandler(this.outputDir, linkExtractor, Arrays
 				.asList(descriptors)));
 		return this;
 	}

spring-restdocs-core/src/main/java/org/springframework/restdocs/core/RestDocumentationResultHandlers.java

@@ -240,8 +240,11 @@ public abstract class RestDocumentationResultHandlers {
 
 		private final Map<String, LinkDescriptor> descriptorsByRel = new HashMap<String, LinkDescriptor>();
 
-		public LinkDocumentingResultHandler(String outputDir, List<LinkDescriptor> descriptors) {
+		private final LinkExtractor extractor;
+
+		public LinkDocumentingResultHandler(String outputDir, LinkExtractor linkExtractor, List<LinkDescriptor> descriptors) {
 			super(outputDir, "links");
+			this.extractor = linkExtractor;
 			for (LinkDescriptor descriptor: descriptors) {
 				Assert.hasText(descriptor.getRel());
 				Assert.hasText(descriptor.getDescription());
@@ -253,7 +256,7 @@ public abstract class RestDocumentationResultHandlers {
 		@Override
 		void handle(MvcResult result, DocumentationWriter writer) throws Exception {
 			Map<String, Object> json = this.objectMapper.readValue(result.getResponse().getContentAsString(), Map.class);
-			Map<String, Object> links = (Map<String, Object>) json.get("_links");
+			Map<String, Object> links = this.extractor.extractLinks(json);
 
 			Set<String> actualRels = links.keySet();
 			Set<String> expectedRels = this.descriptorsByRel.keySet();