Consider optionality when finding uncommon fields in subsection extraction

Fixes gh-573

Author: wilkinsona

spring-restdocs-core/src/main/java/org/springframework/restdocs/payload/AbstractFieldsSnippet.java

@@ -151,7 +151,8 @@ protected Map<String, Object> createModel(Operation operation) {
 		}
 		MediaType contentType = getContentType(operation);
 		if (this.subsectionExtractor != null) {
-			content = verifyContent(this.subsectionExtractor.extractSubsection(content, contentType));
+			content = verifyContent(
+					this.subsectionExtractor.extractSubsection(content, contentType, this.fieldDescriptors));
 		}
 		ContentHandler contentHandler = ContentHandler.forContentWithDescriptors(content, contentType,
 				this.fieldDescriptors);

spring-restdocs-core/src/main/java/org/springframework/restdocs/payload/FieldPathPayloadSubsectionExtractor.java

@@ -17,9 +17,11 @@
 package org.springframework.restdocs.payload;
 
 import java.io.IOException;
-import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.SerializationFeature;
@@ -71,28 +73,39 @@ protected FieldPathPayloadSubsectionExtractor(String fieldPath, String subsectio
 
 	@Override
 	public byte[] extractSubsection(byte[] payload, MediaType contentType) {
+		return extractSubsection(payload, contentType, Collections.emptyList());
+	}
+
+	@Override
+	public byte[] extractSubsection(byte[] payload, MediaType contentType, List<FieldDescriptor> descriptors) {
 		try {
 			ExtractedField extractedField = new JsonFieldProcessor().extract(this.fieldPath,
 					objectMapper.readValue(payload, Object.class));
 			Object value = extractedField.getValue();
 			if (value == ExtractedField.ABSENT) {
 				throw new PayloadHandlingException(this.fieldPath + " does not identify a section of the payload");
 			}
+			Map<JsonFieldPath, FieldDescriptor> descriptorsByPath = descriptors.stream()
+					.collect(Collectors.toMap(
+							(descriptor) -> JsonFieldPath.compile(this.fieldPath + "." + descriptor.getPath()),
+							this::prependFieldPath));
 			if (value instanceof List) {
 				List<?> extractedList = (List<?>) value;
-				Set<String> uncommonPaths = JsonFieldPaths.from(extractedList).getUncommon();
+				JsonContentHandler contentHandler = new JsonContentHandler(payload, descriptorsByPath.values());
+				Set<JsonFieldPath> uncommonPaths = JsonFieldPaths.from(extractedList).getUncommon().stream()
+						.map((path) -> JsonFieldPath.compile(this.fieldPath + "." + path)).filter((path) -> {
+							FieldDescriptor descriptorForPath = descriptorsByPath.getOrDefault(path,
+									new FieldDescriptor(path.toString()));
+							return contentHandler.isMissing(descriptorForPath);
+						}).collect(Collectors.toSet());
 				if (uncommonPaths.isEmpty()) {
 					value = extractedList.get(0);
 				}
 				else {
 					String message = this.fieldPath + " identifies multiple sections of "
 							+ "the payload and they do not have a common structure. The "
-							+ "following uncommon paths were found: ";
-					List<String> prefixedPaths = new ArrayList<>();
-					for (String uncommonPath : uncommonPaths) {
-						prefixedPaths.add(this.fieldPath + "." + uncommonPath);
-					}
-					message += prefixedPaths;
+							+ "following non-optional uncommon paths were found: ";
+					message += uncommonPaths;
 					throw new PayloadHandlingException(message);
 				}
 			}
@@ -103,6 +116,14 @@ public byte[] extractSubsection(byte[] payload, MediaType contentType) {
 		}
 	}
 
+	private FieldDescriptor prependFieldPath(FieldDescriptor original) {
+		FieldDescriptor prefixed = new FieldDescriptor(this.fieldPath + "." + original.getPath());
+		if (original.isOptional()) {
+			prefixed.optional();
+		}
+		return prefixed;
+	}
+
 	@Override
 	public String getSubsectionId() {
 		return this.subsectionId;

spring-restdocs-core/src/main/java/org/springframework/restdocs/payload/JsonContentHandler.java

@@ -44,9 +44,9 @@ class JsonContentHandler implements ContentHandler {
 
 	private final byte[] rawContent;
 
-	private final List<FieldDescriptor> fieldDescriptors;
+	private final Collection<FieldDescriptor> fieldDescriptors;
 
-	JsonContentHandler(byte[] content, List<FieldDescriptor> fieldDescriptors) {
+	JsonContentHandler(byte[] content, Collection<FieldDescriptor> fieldDescriptors) {
 		this.rawContent = content;
 		this.fieldDescriptors = fieldDescriptors;
 		readContent();
@@ -55,22 +55,26 @@ class JsonContentHandler implements ContentHandler {
 	@Override
 	public List<FieldDescriptor> findMissingFields() {
 		List<FieldDescriptor> missingFields = new ArrayList<>();
-		Object payload = readContent();
 		for (FieldDescriptor fieldDescriptor : this.fieldDescriptors) {
-			if (!fieldDescriptor.isOptional() && !this.fieldProcessor.hasField(fieldDescriptor.getPath(), payload)
-					&& !isNestedBeneathMissingOptionalField(fieldDescriptor, payload)) {
+			if (isMissing(fieldDescriptor)) {
 				missingFields.add(fieldDescriptor);
 			}
 		}
 
 		return missingFields;
 	}
 
-	private boolean isNestedBeneathMissingOptionalField(FieldDescriptor missing, Object payload) {
+	boolean isMissing(FieldDescriptor descriptor) {
+		Object payload = readContent();
+		return !descriptor.isOptional() && !this.fieldProcessor.hasField(descriptor.getPath(), payload)
+				&& !isNestedBeneathMissingOptionalField(descriptor, payload);
+	}
+
+	private boolean isNestedBeneathMissingOptionalField(FieldDescriptor descriptor, Object payload) {
 		List<FieldDescriptor> candidates = new ArrayList<>(this.fieldDescriptors);
-		candidates.remove(missing);
+		candidates.remove(descriptor);
 		for (FieldDescriptor candidate : candidates) {
-			if (candidate.isOptional() && missing.getPath().startsWith(candidate.getPath())
+			if (candidate.isOptional() && descriptor.getPath().startsWith(candidate.getPath())
 					&& isMissing(candidate, payload)) {
 				return true;
 			}

spring-restdocs-core/src/main/java/org/springframework/restdocs/payload/JsonFieldPath.java

@@ -55,6 +55,26 @@ List<String> getSegments() {
 		return this.segments;
 	}
 
+	@Override
+	public boolean equals(Object obj) {
+		if (this == obj) {
+			return true;
+		}
+		if (obj == null) {
+			return false;
+		}
+		if (getClass() != obj.getClass()) {
+			return false;
+		}
+		JsonFieldPath other = (JsonFieldPath) obj;
+		return this.segments.equals(other.segments);
+	}
+
+	@Override
+	public int hashCode() {
+		return this.segments.hashCode();
+	}
+
 	@Override
 	public String toString() {
 		return this.rawPath;

spring-restdocs-core/src/main/java/org/springframework/restdocs/payload/PayloadSubsectionExtractor.java

@@ -16,6 +16,8 @@
 
 package org.springframework.restdocs.payload;
 
+import java.util.List;
+
 import org.springframework.http.MediaType;
 
 /**
@@ -36,6 +38,19 @@ public interface PayloadSubsectionExtractor<T extends PayloadSubsectionExtractor
 	 */
 	byte[] extractSubsection(byte[] payload, MediaType contentType);
 
+	/**
+	 * Extracts a subsection of the given {@code payload} that has the given
+	 * {@code contentType} and that is described by the given {@code descriptors}.
+	 * @param payload the payload
+	 * @param contentType the content type of the payload
+	 * @param descriptors descriptors that describe the payload
+	 * @return the subsection of the payload
+	 * @since 2.0.4
+	 */
+	default byte[] extractSubsection(byte[] payload, MediaType contentType, List<FieldDescriptor> descriptors) {
+		return extractSubsection(payload, contentType);
+	}
+
 	/**
 	 * Returns an identifier for the subsection that this extractor will extract.
 	 * @return the identifier

spring-restdocs-core/src/test/java/org/springframework/restdocs/payload/FieldPathPayloadSubsectionExtractorTests.java

@@ -17,6 +17,7 @@
 package org.springframework.restdocs.payload;
 
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.Map;
 
 import com.fasterxml.jackson.core.JsonParseException;
@@ -101,11 +102,35 @@ public void extractMapSubsectionWithCommonStructureFromMultiElementArrayInAJsonM
 	public void extractMapSubsectionWithVaryingStructureFromMultiElementArrayInAJsonMap()
 			throws JsonParseException, JsonMappingException, IOException {
 		this.thrown.expect(PayloadHandlingException.class);
-		this.thrown.expectMessage("The following uncommon paths were found: [a.[].b.d]");
+		this.thrown.expectMessage("The following non-optional uncommon paths were found: [a.[].b.d]");
 		new FieldPathPayloadSubsectionExtractor("a.[].b").extractSubsection(
 				"{\"a\":[{\"b\":{\"c\":5}},{\"b\":{\"c\":6, \"d\": 7}}]}".getBytes(), MediaType.APPLICATION_JSON);
 	}
 
+	@Test
+	@SuppressWarnings("unchecked")
+	public void extractMapSubsectionWithVaryingStructureDueToOptionalFieldsFromMultiElementArrayInAJsonMap()
+			throws JsonParseException, JsonMappingException, IOException {
+		byte[] extractedPayload = new FieldPathPayloadSubsectionExtractor("a.[].b").extractSubsection(
+				"{\"a\":[{\"b\":{\"c\":5}},{\"b\":{\"c\":6, \"d\": 7}}]}".getBytes(), MediaType.APPLICATION_JSON,
+				Arrays.asList(new FieldDescriptor("d").optional()));
+		Map<String, Object> extracted = new ObjectMapper().readValue(extractedPayload, Map.class);
+		assertThat(extracted.size()).isEqualTo(1);
+		assertThat(extracted).containsOnlyKeys("c");
+	}
+
+	@Test
+	@SuppressWarnings("unchecked")
+	public void extractMapSubsectionWithVaryingStructureDueToOptionalParentFieldsFromMultiElementArrayInAJsonMap()
+			throws JsonParseException, JsonMappingException, IOException {
+		byte[] extractedPayload = new FieldPathPayloadSubsectionExtractor("a.[].b").extractSubsection(
+				"{\"a\":[{\"b\":{\"c\":5}},{\"b\":{\"c\":6, \"d\": { \"e\": 7}}}]}".getBytes(),
+				MediaType.APPLICATION_JSON, Arrays.asList(new FieldDescriptor("d").optional()));
+		Map<String, Object> extracted = new ObjectMapper().readValue(extractedPayload, Map.class);
+		assertThat(extracted.size()).isEqualTo(1);
+		assertThat(extracted).containsOnlyKeys("c");
+	}
+
 	@Test
 	public void extractedSubsectionIsPrettyPrintedWhenInputIsPrettyPrinted()
 			throws JsonParseException, JsonMappingException, JsonProcessingException, IOException {