Setup MongoMappingContext setInitialEntitySet

Update MongoDataAutoConfiguration to set the MongoMappingContext
initialEntitySet by scanning for @document or @Persistent classes
from AutoConfigurationPackages.

Fixes gh-2107

Author: odrotbohm

spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/mongo/MongoDataAutoConfiguration.java

@@ -17,19 +17,30 @@
 package org.springframework.boot.autoconfigure.mongo;
 
 import java.net.UnknownHostException;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
 
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.NoSuchBeanDefinitionException;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.config.BeanDefinition;
+import org.springframework.boot.autoconfigure.AutoConfigurationPackages;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.ClassPathScanningCandidateComponentProvider;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.core.env.Environment;
+import org.springframework.core.io.ResourceLoader;
+import org.springframework.core.type.filter.AnnotationTypeFilter;
 import org.springframework.dao.DataAccessException;
 import org.springframework.dao.support.PersistenceExceptionTranslator;
+import org.springframework.data.annotation.Persistent;
 import org.springframework.data.authentication.UserCredentials;
 import org.springframework.data.mongodb.MongoDbFactory;
 import org.springframework.data.mongodb.core.MongoTemplate;
@@ -39,9 +50,11 @@
 import org.springframework.data.mongodb.core.convert.DefaultDbRefResolver;
 import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
 import org.springframework.data.mongodb.core.convert.MongoConverter;
+import org.springframework.data.mongodb.core.mapping.Document;
 import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
 import org.springframework.data.mongodb.gridfs.GridFsTemplate;
 import org.springframework.util.Assert;
+import org.springframework.util.ClassUtils;
 import org.springframework.util.StringUtils;
 
 import com.mongodb.DB;
@@ -71,6 +84,12 @@ public class MongoDataAutoConfiguration {
 	@Autowired
 	private MongoProperties properties;
 
+	@Autowired
+	private Environment environment;
+
+	@Autowired
+	private ResourceLoader resourceLoader;
+
 	@Bean
 	@ConditionalOnMissingBean
 	public MongoDbFactory mongoDbFactory(Mongo mongo) throws Exception {
@@ -111,8 +130,42 @@ public MappingMongoConverter mappingMongoConverter(MongoDbFactory factory,
 
 	@Bean
 	@ConditionalOnMissingBean
-	public MongoMappingContext mongoMappingContext() {
-		return new MongoMappingContext();
+	public MongoMappingContext mongoMappingContext(BeanFactory beanFactory)
+			throws ClassNotFoundException {
+		MongoMappingContext context = new MongoMappingContext();
+		context.setInitialEntitySet(getInitialEntitySet(beanFactory));
+		return context;
+	}
+
+	private Set<Class<?>> getInitialEntitySet(BeanFactory beanFactory)
+			throws ClassNotFoundException {
+		Set<Class<?>> entitySet = new HashSet<Class<?>>();
+		ClassPathScanningCandidateComponentProvider scanner = new ClassPathScanningCandidateComponentProvider(
+				false);
+		scanner.setEnvironment(this.environment);
+		scanner.setResourceLoader(this.resourceLoader);
+		scanner.addIncludeFilter(new AnnotationTypeFilter(Document.class));
+		scanner.addIncludeFilter(new AnnotationTypeFilter(Persistent.class));
+		for (String basePackage : getMappingBasePackages(beanFactory)) {
+			if (StringUtils.hasText(basePackage)) {
+				for (BeanDefinition candidate : scanner
+						.findCandidateComponents(basePackage)) {
+					entitySet.add(ClassUtils.forName(candidate.getBeanClassName(),
+							MongoDataAutoConfiguration.class.getClassLoader()));
+				}
+			}
+		}
+		return entitySet;
+	}
+
+	private static Collection<String> getMappingBasePackages(BeanFactory beanFactory) {
+		try {
+			return AutoConfigurationPackages.get(beanFactory);
+		}
+		catch (IllegalStateException ex) {
+			// no auto-configuration package registered yet
+			return Collections.emptyList();
+		}
 	}
 
 	@Bean

spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/mongo/MongoDataAutoConfigurationTests.java

@@ -17,28 +17,37 @@
 package org.springframework.boot.autoconfigure.mongo;
 
 import java.util.Arrays;
+import java.util.Set;
 
+import org.hamcrest.Matchers;
 import org.junit.After;
 import org.junit.Test;
+import org.springframework.boot.autoconfigure.AutoConfigurationPackages;
 import org.springframework.boot.autoconfigure.PropertyPlaceholderAutoConfiguration;
+import org.springframework.boot.autoconfigure.data.mongo.city.City;
 import org.springframework.boot.test.EnvironmentTestUtils;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.core.convert.converter.Converter;
 import org.springframework.data.mongodb.core.MongoTemplate;
 import org.springframework.data.mongodb.core.convert.CustomConversions;
+import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
 import org.springframework.data.mongodb.gridfs.GridFsTemplate;
+import org.springframework.test.util.ReflectionTestUtils;
 
 import com.mongodb.Mongo;
 
+import static org.hamcrest.Matchers.hasSize;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 
 /**
  * Tests for {@link MongoDataAutoConfiguration}.
  *
  * @author Josh Long
+ * @author Oliver Gierke
  */
 public class MongoDataAutoConfigurationTests {
 
@@ -82,14 +91,34 @@ public void customConversions() throws Exception {
 				.canConvert(Mongo.class, Boolean.class));
 	}
 
+	@Test
+	public void usesAutoConfigurationPackageToPickUpDocumentTypes() {
+		this.context = new AnnotationConfigApplicationContext();
+		String cityPackage = City.class.getPackage().getName();
+		AutoConfigurationPackages.register(this.context, cityPackage);
+		this.context.register(MongoAutoConfiguration.class,
+				MongoDataAutoConfiguration.class);
+		this.context.refresh();
+		assertDomainTypesDiscovered(this.context.getBean(MongoMappingContext.class),
+				City.class);
+	}
+
+	@SuppressWarnings({ "unchecked", "rawtypes" })
+	private static void assertDomainTypesDiscovered(MongoMappingContext mappingContext,
+			Class<?>... types) {
+		Set<Class> initialEntitySet = (Set<Class>) ReflectionTestUtils.getField(
+				mappingContext, "initialEntitySet");
+		assertThat(initialEntitySet, hasSize(types.length));
+		assertThat(initialEntitySet, Matchers.<Class> hasItems(types));
+	}
+
 	@Configuration
 	static class CustomConversionsConfig {
 
 		@Bean
 		public CustomConversions customConversions() {
 			return new CustomConversions(Arrays.asList(new MyConverter()));
 		}
-
 	}
 
 	private static class MyConverter implements Converter<Mongo, Boolean> {
@@ -100,4 +129,5 @@ public Boolean convert(Mongo source) {
 		}
 
 	}
+
 }