From 5bbb458dc5e6402d90fc6df9018e434524ab4bcb Mon Sep 17 00:00:00 2001
From: Mark Paluch <mpaluch@pivotal.io>
Date: Mon, 19 Nov 2018 10:15:00 +0100
Subject: [PATCH] DATACMNS-1421 - Lookup most specific wither method for a
 property.

We now inspect all methods that match the wither method pattern (name, accepting a single argument) to find the most specific method returning the actual entity type.

Previously, we attempted to find a method only considering name and argument properties and not the return type. This lookup strategy could find a method returning the entity super type that isn't assignable to the entity type.

Original pull request: #323.
---
 .../data/mapping/model/Property.java          | 19 +++---
 .../data/mapping/model/PropertyUnitTests.java | 58 +++++++++++++++++++
 2 files changed, 70 insertions(+), 7 deletions(-)

diff --git a/src/main/java/org/springframework/data/mapping/model/Property.java b/src/main/java/org/springframework/data/mapping/model/Property.java
index a15e0e6dc..62470a868 100644
--- a/src/main/java/org/springframework/data/mapping/model/Property.java
+++ b/src/main/java/org/springframework/data/mapping/model/Property.java
@@ -22,6 +22,7 @@ import java.beans.PropertyDescriptor;
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
 import org.springframework.data.util.Lazy;
@@ -79,19 +80,23 @@ public class Property {
 				.filter(it -> type.getParameterTypes(it).get(0).getType().isAssignableFrom(getType()));
 
 		this.wither = Lazy.of(() -> findWither(type, getName(), getType()));
-
 	}
 
 	private static Optional<Method> findWither(TypeInformation<?> owner, String propertyName, Class<?> rawType) {
 
-		Method method = ReflectionUtils.findMethod(owner.getType(),
-				String.format("with%s", StringUtils.capitalize(propertyName)), rawType);
+		AtomicReference<Method> resultHolder = new AtomicReference<>();
+		String methodName = String.format("with%s", StringUtils.capitalize(propertyName));
 
-		if (method != null && owner.isAssignableFrom(owner.getReturnType(method))) {
-			return Optional.of(method);
-		}
+		ReflectionUtils.doWithMethods(owner.getType(), it -> {
 
-		return Optional.empty();
+			if (owner.isAssignableFrom(owner.getReturnType(it))) {
+				resultHolder.set(it);
+			}
+		}, it -> it.getName().equals(methodName) && it.getParameterCount() == 1
+				&& it.getParameterTypes()[0].equals(rawType));
+
+		Method method = resultHolder.get();
+		return method != null ? Optional.of(method) : Optional.empty();
 	}
 
 	/**
diff --git a/src/test/java/org/springframework/data/mapping/model/PropertyUnitTests.java b/src/test/java/org/springframework/data/mapping/model/PropertyUnitTests.java
index 00c3723cb..d00afb668 100644
--- a/src/test/java/org/springframework/data/mapping/model/PropertyUnitTests.java
+++ b/src/test/java/org/springframework/data/mapping/model/PropertyUnitTests.java
@@ -54,8 +54,31 @@ public class PropertyUnitTests {
 		});
 	}
 
+	@Test // DATACMNS-1421
+	public void shouldDiscoverDerivedWitherMethod() {
+
+		Property property = Property.of(ClassTypeInformation.from(DerivedWitherClass.class),
+				ReflectionUtils.findField(DerivedWitherClass.class, "id"));
+
+		assertThat(property.getWither()).isPresent().hasValueSatisfying(actual -> {
+			assertThat(actual.getName()).isEqualTo("withId");
+			assertThat(actual.getReturnType()).isEqualTo(DerivedWitherClass.class);
+			assertThat(actual.getDeclaringClass()).isEqualTo(DerivedWitherClass.class);
+		});
+	}
+
+	@Test // DATACMNS-1421
+	public void shouldNotDiscoverWitherMethodWithIncompatibleReturnType() {
+
+		Property property = Property.of(ClassTypeInformation.from(AnotherLevel.class),
+				ReflectionUtils.findField(AnotherLevel.class, "id"));
+
+		assertThat(property.getWither()).isEmpty();
+	}
+
 	@Value
 	static class ImmutableType {
+
 		String id;
 		String name;
 
@@ -75,7 +98,42 @@ public class PropertyUnitTests {
 	@Value
 	@Wither
 	static class WitherType {
+
 		String id;
 		String name;
 	}
+
+	static abstract class WitherBaseClass {
+
+		abstract WitherBaseClass withId(String id);
+	}
+
+	static abstract class WitherIntermediateClass extends WitherBaseClass {
+
+		abstract WitherIntermediateClass withId(String id);
+	}
+
+	static class DerivedWitherClass extends WitherIntermediateClass {
+
+		private final String id;
+
+		protected DerivedWitherClass(String id) {
+			this.id = id;
+		}
+
+		DerivedWitherClass withId(String id) {
+			return new DerivedWitherClass(id);
+		}
+	}
+
+	static class AnotherLevel extends DerivedWitherClass {
+
+		private AnotherLevel(String id) {
+			super(id);
+		}
+
+		DerivedWitherClass withId(String id) {
+			return new AnotherLevel(id);
+		}
+	}
 }