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

rest-notes-spring-data-rest/src/documentation/java/com/example/notes/ApiDocumentation.java

@@ -20,6 +20,7 @@ import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.notNullValue;
 import static org.springframework.restdocs.core.RestDocumentation.document;
 import static org.springframework.restdocs.core.RestDocumentation.linkWithRel;
+import static org.springframework.restdocs.core.RestDocumentation.halLinks;
 import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get;
 import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.patch;
 import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.post;
@@ -95,7 +96,8 @@ public class ApiDocumentation {
 	public void indexExample() throws Exception {
 		document("index-example",
 				this.mockMvc.perform(get("/")).andExpect(status().isOk()))
-				.andDocumentHalLinks(
+				.andDocumentLinks(
+						halLinks(),
 						linkWithRel("notes").description(
 								"The <<resources-notes,Notes resource>>"),
 						linkWithRel("tags").description(
@@ -173,7 +175,8 @@ public class ApiDocumentation {
 				.andExpect(jsonPath("body", is(note.get("body"))))
 				.andExpect(jsonPath("_links.self.href", is(noteLocation)))
 				.andExpect(jsonPath("_links.tags", is(notNullValue())))
-				.andDocumentHalLinks(
+				.andDocumentLinks(
+						halLinks(),
 						linkWithRel("self").description("This <<resources-note,note>>"),
 						linkWithRel("tags").description(
 								"This note's <<resources-note-tags,tags>>"));
@@ -261,7 +264,8 @@ public class ApiDocumentation {
 		document("tag-get-example", this.mockMvc.perform(get(tagLocation)))
 				.andExpect(status().isOk())
 				.andExpect(jsonPath("name", is(tag.get("name"))))
-				.andDocumentHalLinks(
+				.andDocumentLinks(
+						halLinks(),
 						linkWithRel("self").description("This <<resources-tag,tag>>"),
 						linkWithRel("notes")
 								.description(

rest-notes-spring-hateoas/src/documentation/java/com/example/notes/ApiDocumentation.java

@@ -19,6 +19,7 @@ package com.example.notes;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.notNullValue;
 import static org.springframework.restdocs.core.RestDocumentation.document;
+import static org.springframework.restdocs.core.RestDocumentation.halLinks;
 import static org.springframework.restdocs.core.RestDocumentation.linkWithRel;
 import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get;
 import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.patch;
@@ -94,8 +95,9 @@ public class ApiDocumentation {
 	@Test
 	public void indexExample() throws Exception {
 		document("index-example",
-				this.mockMvc.perform(get("/")).andExpect(status().isOk()))
-				.andDocumentHalLinks(
+				this.mockMvc.perform(get("/"))
+				.andExpect(status().isOk()))
+				.andDocumentLinks(halLinks(),
 						linkWithRel("notes").description(
 								"The <<resources-notes,Notes resource>>"),
 						linkWithRel("tags").description(
@@ -171,7 +173,7 @@ public class ApiDocumentation {
 				.andExpect(jsonPath("body", is(note.get("body"))))
 				.andExpect(jsonPath("_links.self.href", is(noteLocation)))
 				.andExpect(jsonPath("_links.note-tags", is(notNullValue())))
-				.andDocumentHalLinks(
+				.andDocumentLinks(halLinks(),
 						linkWithRel("self").description("This <<resources-note,note>>"),
 						linkWithRel("note-tags").description(
 								"This note's <<resources-note-tags,tags>>"));
@@ -259,7 +261,7 @@ public class ApiDocumentation {
 		document("tag-get-example", this.mockMvc.perform(get(tagLocation)))
 				.andExpect(status().isOk())
 				.andExpect(jsonPath("name", is(tag.get("name"))))
-				.andDocumentHalLinks(
+				.andDocumentLinks(halLinks(),
 						linkWithRel("self").description("This <<resources-tag,tag>>"),
 						linkWithRel("tagged-notes")
 								.description(

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();