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: #260.

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

@@ -591,7 +591,7 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 
 			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 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 	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 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 		return dbo;
 	}
 
-	private String potentiallyConvertMapKey(Object 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 (key instanceof String) {
-			return (String) key;
-		}
+		Assert.notNull(key, "Map key must not be null!");
 
-		if (conversions.hasCustomWriteTarget(key.getClass(), String.class)) {
-			return (String) getPotentiallyConvertedSimpleWrite(key);
-		}
-
-		return key.toString();
+		String convertedKey = potentiallyConvertMapKey(key);
+		return potentiallyEscapeMapKey(convertedKey);
 	}
 
 	/**
@@ -703,6 +707,22 @@ public class MappingMongoConverter extends AbstractMongoConverter implements App
 		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.mapping.model.MappingInstantiationException;
 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 class MappingMongoConverterUnitTests {
 		assertThat(converter.read(TypeWithLocalDateTime.class, result).date, is(reference));
 	}
 
-
 	/**
 	 * @see DATAMONGO-1128
 	 */
@@ -2009,6 +2009,7 @@ public class MappingMongoConverterUnitTests {
 	 * @see DATAMONGO-1118
 	 */
 	@Test
+	@SuppressWarnings("unchecked")
 	public void convertsMapKeyUsingCustomConverterForAndBackwards() {
 
 		MappingMongoConverter converter = new MappingMongoConverter(resolver, mappingContext);
@@ -2017,13 +2018,13 @@ public class MappingMongoConverterUnitTests {
 		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 class MappingMongoConverterUnitTests {
 		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 @@ public class MappingMongoConverterUnitTests {
 		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 @@ public class MappingMongoConverterUnitTests {
 
 		@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 class MappingMongoConverterUnitTests {
 
 			throw new ConversionNotSupportedException(source, String.class, null);
 		}
-
 	}
-
-
 }