DATAMONGO-1118 - Polishing.

Created dedicated prepareMapKey(…) method to chain calls to potentiallyConvertMapKey(…) and potentiallyEscapeMapKey(…) and make sure they always get applied in combination.

Fixed initial map creation for DBRefs to apply the fixed behavior, too.

Original pull request: spring-projects#260.

Author: odrotbohm

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/convert/MappingMongoConverter.java

@@ -591,7 +591,7 @@ protected DBObject createMap(Map<Object, Object> map, MongoPersistentProperty pr
 
 			if (conversions.isSimpleType(key.getClass())) {
 
-				String simpleKey = potentiallyEscapeMapKey(key.toString());
+				String simpleKey = prepareMapKey(key.toString());
 				dbObject.put(simpleKey, value != null ? createDBRef(value, property) : null);
 
 			} else {
@@ -643,11 +643,13 @@ private BasicDBList writeCollectionInternal(Collection<?> source, TypeInformatio
 	protected DBObject writeMapInternal(Map<Object, Object> obj, DBObject dbo, TypeInformation<?> propertyType) {
 
 		for (Map.Entry<Object, Object> entry : obj.entrySet()) {
+
 			Object key = entry.getKey();
 			Object val = entry.getValue();
+
 			if (conversions.isSimpleType(key.getClass())) {
 
-				String simpleKey = potentiallyEscapeMapKey(potentiallyConvertMapKey(key));
+				String simpleKey = prepareMapKey(key);
 				if (val == null || conversions.isSimpleType(val.getClass())) {
 					writeSimpleInternal(val, dbo, simpleKey);
 				} else if (val instanceof Collection || val.getClass().isArray()) {
@@ -668,17 +670,19 @@ protected DBObject writeMapInternal(Map<Object, Object> obj, DBObject dbo, TypeI
 		return dbo;
 	}
 
-	private String potentiallyConvertMapKey(Object key) {
-
-		if (key instanceof String) {
-			return (String) key;
-		}
+	/**
+	 * Prepares the given {@link Map} key to be converted into a {@link String}. Will invoke potentially registered custom
+	 * conversions and escape dots from the result as they're not supported as {@link Map} key in MongoDB.
+	 * 
+	 * @param key must not be {@literal null}.
+	 * @return
+	 */
+	private String prepareMapKey(Object key) {
 
-		if (conversions.hasCustomWriteTarget(key.getClass(), String.class)) {
-			return (String) getPotentiallyConvertedSimpleWrite(key);
-		}
+		Assert.notNull(key, "Map key must not be null!");
 
-		return key.toString();
+		String convertedKey = potentiallyConvertMapKey(key);
+		return potentiallyEscapeMapKey(convertedKey);
 	}
 
 	/**
@@ -703,6 +707,22 @@ protected String potentiallyEscapeMapKey(String source) {
 		return source.replaceAll("\\.", mapKeyDotReplacement);
 	}
 
+	/**
+	 * Returns a {@link String} representation of the given {@link Map} key
+	 * 
+	 * @param key
+	 * @return
+	 */
+	private String potentiallyConvertMapKey(Object key) {
+
+		if (key instanceof String) {
+			return (String) key;
+		}
+
+		return conversions.hasCustomWriteTarget(key.getClass(), String.class) ? (String) getPotentiallyConvertedSimpleWrite(key)
+				: key.toString();
+	}
+
 	/**
 	 * Translates the map key replacements in the given key just read with a dot in case a map key replacement has been
 	 * configured.

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/convert/MappingMongoConverterUnitTests.java

@@ -77,6 +77,7 @@
 import org.springframework.data.mongodb.core.DBObjectTestUtils;
 import org.springframework.data.mongodb.core.convert.DBObjectAccessorUnitTests.NestedType;
 import org.springframework.data.mongodb.core.convert.DBObjectAccessorUnitTests.ProjectingType;
+import org.springframework.data.mongodb.core.convert.MappingMongoConverterUnitTests.ClassWithMapUsingEnumAsKey.FooBarEnum;
 import org.springframework.data.mongodb.core.geo.Sphere;
 import org.springframework.data.mongodb.core.mapping.Document;
 import org.springframework.data.mongodb.core.mapping.Field;
@@ -1967,7 +1968,6 @@ public void convertsJava8DateTimeTypesToDateAndBack() {
 		assertThat(converter.read(TypeWithLocalDateTime.class, result).date, is(reference));
 	}
 
-
 	/**
 	 * @see DATAMONGO-1128
 	 */
@@ -2009,6 +2009,7 @@ public void readsOptionalsCorrectly() {
 	 * @see DATAMONGO-1118
 	 */
 	@Test
+	@SuppressWarnings("unchecked")
 	public void convertsMapKeyUsingCustomConverterForAndBackwards() {
 
 		MappingMongoConverter converter = new MappingMongoConverter(resolver, mappingContext);
@@ -2017,13 +2018,13 @@ public void convertsMapKeyUsingCustomConverterForAndBackwards() {
 		converter.afterPropertiesSet();
 
 		ClassWithMapUsingEnumAsKey source = new ClassWithMapUsingEnumAsKey();
-		source.map = new HashMap<MappingMongoConverterUnitTests.FooBarEnum, String>();
+		source.map = new HashMap<FooBarEnum, String>();
 		source.map.put(FooBarEnum.FOO, "wohoo");
 
 		DBObject target = new BasicDBObject();
 		converter.write(source, target);
 
-		assertThat(converter.read(ClassWithMapUsingEnumAsKey.class, target).map, equalTo(source.map));
+		assertThat(converter.read(ClassWithMapUsingEnumAsKey.class, target).map, is(source.map));
 	}
 
 	/**
@@ -2037,7 +2038,7 @@ public void writesMapKeyUsingCustomConverter() {
 		converter.afterPropertiesSet();
 
 		ClassWithMapUsingEnumAsKey source = new ClassWithMapUsingEnumAsKey();
-		source.map = new HashMap<MappingMongoConverterUnitTests.FooBarEnum, String>();
+		source.map = new HashMap<FooBarEnum, String>();
 		source.map.put(FooBarEnum.FOO, "spring");
 		source.map.put(FooBarEnum.BAR, "data");
 
@@ -2370,11 +2371,12 @@ static class TypeWithOptional {
 		Optional<LocalDateTime> localDateTime = Optional.empty();
 	}
 
-	static enum FooBarEnum {
-		FOO, BAR;
-	}
-
 	static class ClassWithMapUsingEnumAsKey {
+
+		static enum FooBarEnum {
+			FOO, BAR;
+		}
+
 		Map<FooBarEnum, String> map;
 	}
 
@@ -2383,13 +2385,13 @@ static class FooBarEnumToStringConverter implements Converter<FooBarEnum, String
 
 		@Override
 		public String convert(FooBarEnum source) {
+
 			if (source == null) {
 				return null;
 			}
 
 			return FooBarEnum.FOO.equals(source) ? "foo-enum-value" : "bar-enum-value";
 		}
-
 	}
 
 	@ReadingConverter
@@ -2411,8 +2413,5 @@ public FooBarEnum convert(String source) {
 
 			throw new ConversionNotSupportedException(source, String.class, null);
 		}
-
 	}
-
-
 }