GH-2262 - Ensure related entities with heterogenous types and a common base class are loaded correctly.

This adds a test for the scenario in the ticket and fixes #2262.

Author: michael-simons

src/test/java/org/springframework/data/neo4j/integration/imperative/InheritanceMappingIT.java

@@ -15,6 +15,16 @@
  */
 package org.springframework.data.neo4j.integration.imperative;
 
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.neo4j.driver.Driver;
@@ -46,13 +56,6 @@
 import org.springframework.transaction.annotation.Transactional;
 import org.springframework.transaction.support.TransactionTemplate;
 
-import java.util.Collection;
-import java.util.Collections;
-import java.util.List;
-import java.util.Optional;
-
-import static org.assertj.core.api.Assertions.assertThat;
-
 /**
  * @author Gerrit Meier
  * @author Michael J. Simons
@@ -105,16 +108,7 @@ void relationshipsShouldHaveCorrectTypes(@Autowired BuildingRepository repositor
 	@Test // GH-2138
 	void collectionsShouldHaveCorrectTypes(@Autowired TerritoryRepository repository) {
 
-		Long territoryId;
-		try (Session session = driver.session(bookmarkCapture.createSessionConfig())) {
-			territoryId = session.run("CREATE (c:Country:BaseTerritory:BaseEntity{nameEn:'country'}) " +
-					"CREATE (c)-[:LINK]->(:Country:BaseTerritory:BaseEntity{nameEn:'anotherCountry', countryProperty:'large'}) " +
-					"CREATE (c)-[:LINK]->(:Continent:BaseTerritory:BaseEntity{nameEn:'continent', continentProperty:'small'}) " +
-					"CREATE (c)-[:LINK]->(:GenericTerritory:BaseTerritory:BaseEntity{nameEn:'generic'}) " +
-					"return id(c) as id").single()
-					.get(0).asLong();
-			bookmarkCapture.seedWith(session.lastBookmark());
-		}
+		Long territoryId = createDivisionAndTerritories().get("territoryId").asLong();
 
 		Inheritance.BaseTerritory territory = repository.findById(territoryId).get();
 
@@ -134,13 +128,7 @@ void collectionsShouldHaveCorrectTypes(@Autowired TerritoryRepository repository
 	@Test // GH-2138
 	void resultCollectionShouldHaveCorrectTypes(@Autowired TerritoryRepository repository) {
 
-		try (Session session = driver.session(bookmarkCapture.createSessionConfig())) {
-			session.run("CREATE (c:Country:BaseTerritory:BaseEntity{nameEn:'country', countryProperty:'baseCountry'}) " +
-					"CREATE (c)-[:LINK]->(:Country:BaseTerritory:BaseEntity{nameEn:'anotherCountry', countryProperty:'large'}) " +
-					"CREATE (c)-[:LINK]->(:Continent:BaseTerritory:BaseEntity{nameEn:'continent', continentProperty:'small'}) " +
-					"CREATE (c)-[:LINK]->(:GenericTerritory:BaseTerritory:BaseEntity{nameEn:'generic'})").consume();
-			bookmarkCapture.seedWith(session.lastBookmark());
-		}
+		createDivisionAndTerritories();
 
 		List<Inheritance.BaseTerritory> territories = repository.findAll();
 
@@ -343,6 +331,100 @@ void mixedInterfaces(@Autowired Neo4jTemplate template) {
 		}
 	}
 
+	@Test // GH-2262
+	void shouldMatchPolymorphicClassesWhenFetchedById(@Autowired DivisionRepository repository) {
+
+		Record divisionAndTerritoryId = createDivisionAndTerritories();
+
+		Optional<Inheritance.Division> optionalDivision = repository.findById(divisionAndTerritoryId.get("divisionId").asLong());
+		assertThat(optionalDivision).isPresent();
+		assertThat(optionalDivision).hasValueSatisfying(twoDifferentClassesHaveBeenLoaded());
+	}
+
+	@Test // GH-2262
+	void shouldMatchPolymorphicClassesWhenFetchingAll(@Autowired DivisionRepository repository) {
+
+		createDivisionAndTerritories();
+
+		List<Inheritance.Division> divisions = repository.findAll();
+		assertThat(divisions).hasSize(1);
+		assertThat(divisions).first().satisfies(twoDifferentClassesHaveBeenLoaded());
+	}
+
+	private Consumer<Inheritance.Division> twoDifferentClassesHaveBeenLoaded() {
+		return d -> {
+			assertThat(d.getIsActiveIn()).hasSize(2);
+			assertThat(d.getIsActiveIn()).extracting(Inheritance.BaseTerritory::getNameEn)
+					.containsExactlyInAnyOrder("anotherCountry", "continent");
+			Map<String, Class> classByName = d.getIsActiveIn().stream()
+					.collect(Collectors.toMap(Inheritance.BaseTerritory::getNameEn, v -> v.getClass()));
+			assertThat(classByName).containsEntry("anotherCountry", Inheritance.Country.class);
+			assertThat(classByName).containsEntry("continent", Inheritance.Continent.class);
+		};
+	}
+
+	@Test // GH-2262
+	void shouldMatchPolymorphicInterfacesWhenFetchedById(@Autowired ParentModelRepository repository) {
+
+		Record record = createRelationsToDifferentImplementations();
+
+		Optional<Inheritance.ParentModel2> optionalDivision = repository.findById(record.get(0).asNode().id());
+		assertThat(optionalDivision).isPresent();
+		assertThat(optionalDivision).hasValueSatisfying(twoDifferentInterfacesHaveBeenLoaded());
+	}
+
+	@Test // GH-2262
+	void shouldMatchPolymorphicInterfacesWhenFetchingAll(@Autowired ParentModelRepository repository) {
+
+		createRelationsToDifferentImplementations();
+
+		List<Inheritance.ParentModel2> divisions = repository.findAll();
+		assertThat(divisions).hasSize(1);
+		assertThat(divisions).first().satisfies(twoDifferentInterfacesHaveBeenLoaded());
+	}
+
+	private Consumer<Inheritance.ParentModel2> twoDifferentInterfacesHaveBeenLoaded() {
+		return d -> {
+			assertThat(d.getIsRelatedTo()).hasSize(2);
+			assertThat(d.getIsRelatedTo()).extracting(Inheritance.SomeInterface3::getName)
+					.containsExactlyInAnyOrder("3a", "3b");
+			Map<String, Class> classByName = d.getIsRelatedTo().stream()
+					.collect(Collectors.toMap(Inheritance.SomeInterface3::getName, v -> v.getClass()));
+			assertThat(classByName).containsEntry("3a", Inheritance.SomeInterfaceImpl3a.class);
+			assertThat(classByName).containsEntry("3b", Inheritance.SomeInterfaceImpl3b.class);
+		};
+	}
+
+	private Record createDivisionAndTerritories() {
+		Record result;
+		try (Session session = driver.session(bookmarkCapture.createSessionConfig())) {
+
+			result = session.run("CREATE (c:Country:BaseTerritory:BaseEntity{nameEn:'country', countryProperty:'baseCountry'}) " +
+								 "CREATE (c)-[:LINK]->(ca:Country:BaseTerritory:BaseEntity{nameEn:'anotherCountry', countryProperty:'large'}) " +
+								 "CREATE (c)-[:LINK]->(cb:Continent:BaseTerritory:BaseEntity{nameEn:'continent', continentProperty:'small'}) " +
+								 "CREATE (c)-[:LINK]->(:GenericTerritory:BaseTerritory:BaseEntity{nameEn:'generic'}) " +
+								 "CREATE (d:Division:BaseEntity{name:'Division'}) " +
+								 "CREATE (d) -[:IS_ACTIVE_IN] -> (ca)" +
+								 "CREATE (d) -[:IS_ACTIVE_IN] -> (cb)" +
+								 "RETURN id(d) as divisionId, id(c) as territoryId").single();
+			bookmarkCapture.seedWith(session.lastBookmark());
+		}
+		return result;
+	}
+
+	private Record createRelationsToDifferentImplementations() {
+		Record result;
+		try (Session session = driver.session(bookmarkCapture.createSessionConfig())) {
+
+			result = session.run("CREATE (p:ParentModel2) " +
+								 "CREATE (p)-[:IS_RELATED_TO]->(:SomeInterface3:SomeInterface3a {name: '3a'}) " +
+								 "CREATE (p)-[:IS_RELATED_TO]->(:SomeInterface3:SomeInterface3b {name: '3b'}) " +
+								 "RETURN p").single();
+			bookmarkCapture.seedWith(session.lastBookmark());
+		}
+		return result;
+	}
+
 	interface PetsRepository extends Neo4jRepository<AbstractPet, Long> {
 
 		@Transactional(readOnly = true)
@@ -355,6 +437,10 @@ interface BuildingRepository extends Neo4jRepository<Inheritance.Building, Long>
 
 	interface TerritoryRepository extends Neo4jRepository<Inheritance.BaseTerritory, Long> {}
 
+	interface DivisionRepository extends Neo4jRepository<Inheritance.Division, Long> {}
+
+	interface ParentModelRepository extends Neo4jRepository<Inheritance.ParentModel2, Long> {}
+
 	@Configuration
 	@EnableNeo4jRepositories(considerNestedRepositories = true)
 	@EnableTransactionManagement

src/test/java/org/springframework/data/neo4j/integration/shared/common/Inheritance.java

@@ -279,6 +279,31 @@ public void setRelated2(SomeInterface3 related2) {
 	}
 	// end::interface3[]
 
+	/**
+	 * A holder for a list of different interface implementations, see GH-2262.
+	 */
+	@Node
+	public static class ParentModel2 {
+
+		@Id
+		@GeneratedValue
+		private Long id;
+
+		private List<SomeInterface3> isRelatedTo;
+
+		public Long getId() {
+			return id;
+		}
+
+		public List<SomeInterface3> getIsRelatedTo() {
+			return isRelatedTo;
+		}
+
+		public void setIsRelatedTo(List<SomeInterface3> isRelatedTo) {
+			this.isRelatedTo = isRelatedTo;
+		}
+	}
+
 	/**
 	 * Interface to get implemented with the one below.
 	 */
@@ -779,6 +804,10 @@ public BaseTerritory(String nameEn) {
 			this.nameEn = nameEn;
 		}
 
+		public String getNameEn() {
+			return nameEn;
+		}
+
 		@Override
 		public boolean equals(Object o) {
 			if (this == o) {
@@ -875,4 +904,25 @@ public int hashCode() {
 			return Objects.hash(super.hashCode(), continentProperty);
 		}
 	}
+
+	/**
+	 * A parent object for some territories, used to test whether those are loaded correct in a polymorphic way.
+	 */
+	@Node
+	public static class Division extends BaseEntity {
+
+		@Relationship
+		List<BaseTerritory> isActiveIn;
+
+		public List<BaseTerritory> getIsActiveIn() {
+			return isActiveIn;
+		}
+
+		public void setIsActiveIn(
+				List<BaseTerritory> isActiveIn) {
+			this.isActiveIn = isActiveIn;
+		}
+	}
+
+
 }