From c42f5efe0185ce30d2ff79c1de2d01002bb00f28 Mon Sep 17 00:00:00 2001
From: Oliver Drotbohm <odrotbohm@pivotal.io>
Date: Fri, 30 Aug 2019 12:11:55 +0200
Subject: [PATCH] DATACMNS-1547 - Properly remove partially populated
 PersistentEntity instances from cache.

We now remove partially populated PersistentEntity instances from the cache held in AbstractMappingContext as the creation could fail for other RuntimeExceptions other than a MappingException and we wouldn't want to keep the instances around in any case.

Slight refactorings in the unit tests so that we can easily create MappingContext instances that reject certain types with certain exceptions.
---
 .../context/AbstractMappingContext.java       |   2 +-
 .../AbstractMappingContextUnitTests.java      | 100 ++++++++++++++----
 2 files changed, 78 insertions(+), 24 deletions(-)

diff --git a/src/main/java/org/springframework/data/mapping/context/AbstractMappingContext.java b/src/main/java/org/springframework/data/mapping/context/AbstractMappingContext.java
index fb2c92600..806179ada 100644
--- a/src/main/java/org/springframework/data/mapping/context/AbstractMappingContext.java
+++ b/src/main/java/org/springframework/data/mapping/context/AbstractMappingContext.java
@@ -380,7 +380,7 @@ public abstract class AbstractMappingContext<E extends MutablePersistentEntity<?
 					entity.setPersistentPropertyAccessorFactory(persistentPropertyAccessorFactory);
 				}
 
-			} catch (MappingException e) {
+			} catch (RuntimeException e) {
 				persistentEntities.remove(typeInformation);
 				throw e;
 			}
diff --git a/src/test/java/org/springframework/data/mapping/context/AbstractMappingContextUnitTests.java b/src/test/java/org/springframework/data/mapping/context/AbstractMappingContextUnitTests.java
index 64f7d7426..b49482a78 100755
--- a/src/test/java/org/springframework/data/mapping/context/AbstractMappingContextUnitTests.java
+++ b/src/test/java/org/springframework/data/mapping/context/AbstractMappingContextUnitTests.java
@@ -16,19 +16,26 @@
 package org.springframework.data.mapping.context;
 
 import static org.assertj.core.api.Assertions.*;
+import static org.mockito.ArgumentMatchers.*;
 import static org.mockito.Mockito.*;
 
 import groovy.lang.MetaClass;
+import lombok.AccessLevel;
+import lombok.EqualsAndHashCode;
+import lombok.RequiredArgsConstructor;
+import lombok.Value;
 
 import java.time.LocalDateTime;
+import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.TreeMap;
+import java.util.function.Supplier;
 
 import org.junit.Before;
 import org.junit.Test;
-
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationEvent;
 import org.springframework.data.annotation.Id;
@@ -60,29 +67,10 @@ public class AbstractMappingContextUnitTests {
 	@Test // DATACMNS-92
 	public void doesNotAddInvalidEntity() {
 
-		context = new SampleMappingContext() {
-			@Override
-			@SuppressWarnings("unchecked")
-			protected <S> BasicPersistentEntity<Object, SamplePersistentProperty> createPersistentEntity(
-					TypeInformation<S> typeInformation) {
-				return new BasicPersistentEntity<Object, SamplePersistentProperty>((TypeInformation<Object>) typeInformation) {
-					@Override
-					public void verify() {
-						if (Unsupported.class.isAssignableFrom(getType())) {
-							throw new MappingException("Unsupported type!");
-						}
-					}
-				};
-			}
-		};
+		context = TypeRejectingMappingContext.rejecting(() -> new MappingException("Not supported!"), Unsupported.class);
 
-		try {
-			context.getPersistentEntity(Unsupported.class);
-		} catch (MappingException e) {
-			// expected
-		}
-
-		assertThatExceptionOfType(MappingException.class).isThrownBy(() -> context.getPersistentEntity(Unsupported.class));
+		assertThatExceptionOfType(MappingException.class) //
+				.isThrownBy(() -> context.getPersistentEntity(Unsupported.class));
 	}
 
 	@Test
@@ -213,6 +201,21 @@ public class AbstractMappingContextUnitTests {
 		assertThat(context.getPersistentEntity(property)).isNull();
 	}
 
+	@Test // DATACMNS-1574
+	public void cleansUpCacheForRuntimeException() {
+
+		TypeRejectingMappingContext context = TypeRejectingMappingContext.rejecting(() -> new RuntimeException(),
+				Unsupported.class);
+
+		assertThatExceptionOfType(RuntimeException.class) //
+				.isThrownBy(() -> context.getPersistentEntity(Unsupported.class));
+
+		// Second lookup still throws the exception as the temporarily created entity was not cached
+
+		assertThatExceptionOfType(RuntimeException.class) //
+				.isThrownBy(() -> context.getPersistentEntity(Unsupported.class));
+	}
+
 	private static void assertHasEntityFor(Class<?> type, SampleMappingContext context, boolean expected) {
 
 		boolean found = false;
@@ -252,4 +255,55 @@ public class AbstractMappingContextUnitTests {
 	static class Extension extends Base {
 		@Id String foo;
 	}
+
+	/**
+	 * Extension of {@link SampleMappingContext} to reject the creation of certain types with a configurable exception.
+	 *
+	 * @author Oliver Drotbohm
+	 */
+	@Value
+	@EqualsAndHashCode(callSuper = false)
+	@RequiredArgsConstructor(access = AccessLevel.PRIVATE)
+	private static class TypeRejectingMappingContext extends SampleMappingContext {
+
+		Supplier<? extends RuntimeException> exception;
+		Collection<Class<?>> rejectedTypes;
+
+		/**
+		 * Creates a new {@link TypeRejectingMappingContext} producing the given exceptions if any of the given types is
+		 * encountered.
+		 *
+		 * @param <T>
+		 * @param exception must not be {@literal null}.
+		 * @param types must not be {@literal null}.
+		 * @return
+		 */
+		public static <T extends RuntimeException> TypeRejectingMappingContext rejecting(Supplier<T> exception,
+				Class<?>... types) {
+			return new TypeRejectingMappingContext(exception, Arrays.asList(types));
+		}
+
+		/*
+		 * (non-Javadoc)
+		 * @see org.springframework.data.mapping.context.SampleMappingContext#createPersistentEntity(org.springframework.data.util.TypeInformation)
+		 */
+		@Override
+		protected <S> BasicPersistentEntity<Object, SamplePersistentProperty> createPersistentEntity(
+				TypeInformation<S> typeInformation) {
+
+			return new BasicPersistentEntity<Object, SamplePersistentProperty>((TypeInformation<Object>) typeInformation) {
+
+				/*
+				 * (non-Javadoc)
+				 * @see org.springframework.data.mapping.model.BasicPersistentEntity#verify()
+				 */
+				@Override
+				public void verify() {
+					if (rejectedTypes.stream().anyMatch(it -> it.isAssignableFrom(getType()))) {
+						throw exception.get();
+					}
+				}
+			};
+		}
+	}
 }