DATAJPA-576 - Improvements in JpaMappingContext setup in configuration.

Refactored JpaMappingContext to be able to work with multiple instances of Metamodel. This removes the need to wire a dedicated one in JpaMetamodelContextFactoryBean and thus makes it independent of a particular EntityManager(Factory). This significantly eases the implementation of the configuration integration.

src/main/java/org/springframework/data/jpa/mapping/JpaMetamodelMappingContext.java

@@ -17,7 +17,10 @@ package org.springframework.data.jpa.mapping;
 
 import java.beans.PropertyDescriptor;
 import java.lang.reflect.Field;
+import java.util.Collections;
+import java.util.Set;
 
+import javax.persistence.metamodel.ManagedType;
 import javax.persistence.metamodel.Metamodel;
 
 import org.springframework.data.mapping.context.AbstractMappingContext;
@@ -35,17 +38,22 @@ import org.springframework.util.Assert;
 public class JpaMetamodelMappingContext extends
 		AbstractMappingContext<JpaPersistentEntityImpl<?>, JpaPersistentProperty> {
 
-	private final Metamodel model;
+	private final Set<Metamodel> models;
 
 	/**
 	 * Creates a new JPA {@link Metamodel} based {@link MappingContext}.
 	 * 
-	 * @param model must not be {@literal null}.
+	 * @param models must not be {@literal null} or empty.
 	 */
-	public JpaMetamodelMappingContext(Metamodel model) {
+	public JpaMetamodelMappingContext(Set<Metamodel> models) {
 
-		Assert.notNull(model, "JPA Metamodel must not be null!");
-		this.model = model;
+		Assert.notNull(models, "JPA metamodel must not be null!");
+		Assert.notEmpty(models, "At least one JPA metamodel must be present!");
+		this.models = models;
+	}
+
+	public JpaMetamodelMappingContext(Metamodel model) {
+		this(Collections.singleton(model));
 	}
 
 	/* 
@@ -64,7 +72,9 @@ public class JpaMetamodelMappingContext extends
 	@Override
 	protected JpaPersistentProperty createPersistentProperty(Field field, PropertyDescriptor descriptor,
 			JpaPersistentEntityImpl<?> owner, SimpleTypeHolder simpleTypeHolder) {
-		return new JpaPersistentPropertyImpl(model, field, descriptor, owner, simpleTypeHolder);
+
+		Metamodel metamodel = getMetamodelFor(owner.getType());
+		return new JpaPersistentPropertyImpl(metamodel, field, descriptor, owner, simpleTypeHolder);
 	}
 
 	/* 
@@ -73,12 +83,35 @@ public class JpaMetamodelMappingContext extends
 	 */
 	@Override
 	protected boolean shouldCreatePersistentEntityFor(TypeInformation<?> type) {
+		return getMetamodelFor(type.getType()) != null;
+	}
 
-		try {
-			model.managedType(type.getType());
-			return true;
-		} catch (IllegalArgumentException o_O) {
-			return false;
+	/**
+	 * Returns the {@link Metamodel} aware of the given type.
+	 * 
+	 * @param type
+	 * @return
+	 */
+	private Metamodel getMetamodelFor(Class<?> type) {
+
+		for (Metamodel model : models) {
+
+			try {
+				model.managedType(type);
+				return model;
+			} catch (IllegalArgumentException o_O) {
+
+				// Fall back to inspect *all* managed types manually as Metamodel.managedType(…) only
+				// returns for entities, embeddables and managed supperclasses.
+
+				for (ManagedType<?> managedType : model.getManagedTypes()) {
+					if (type.equals(managedType.getJavaType())) {
+						return model;
+					}
+				}
+			}
 		}
+
+		return null;
 	}
 }

src/main/java/org/springframework/data/jpa/mapping/JpaPersistentPropertyImpl.java

@@ -38,6 +38,7 @@ import org.springframework.data.mapping.Association;
 import org.springframework.data.mapping.PersistentEntity;
 import org.springframework.data.mapping.model.AnnotationBasedPersistentProperty;
 import org.springframework.data.mapping.model.SimpleTypeHolder;
+import org.springframework.util.Assert;
 
 /**
  * {@link JpaPersistentProperty} implementation usind a JPA {@link Metamodel}.
@@ -84,6 +85,9 @@ class JpaPersistentPropertyImpl extends AnnotationBasedPersistentProperty<JpaPer
 			PersistentEntity<?, JpaPersistentProperty> owner, SimpleTypeHolder simpleTypeHolder) {
 
 		super(field, propertyDescriptor, owner, simpleTypeHolder);
+
+		Assert.notNull(metamodel, "Metamodel must not be null!");
+
 		this.metamodel = metamodel;
 	}
 

src/main/java/org/springframework/data/jpa/repository/config/AuditingBeanDefinitionParser.java

@@ -41,7 +41,7 @@ public class AuditingBeanDefinitionParser implements BeanDefinitionParser {
 	private static final String AUDITING_BFPP_CLASS_NAME = "org.springframework.data.jpa.domain.support.AuditingBeanFactoryPostProcessor";
 
 	private final AuditingHandlerBeanDefinitionParser auditingHandlerParser = new AuditingHandlerBeanDefinitionParser(
-			JpaRepositoryConfigExtension.JPA_MAPPING_CONTEXT_BEAN_NAME);
+			BeanDefinitionNames.JPA_MAPPING_CONTEXT_BEAN_NAME);
 	private final SpringConfiguredBeanDefinitionParser springConfiguredParser = new SpringConfiguredBeanDefinitionParser();
 
 	/*

src/main/java/org/springframework/data/jpa/repository/config/BeanDefinitionNames.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2014 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.jpa.repository.config;
+
+/**
+ * Helper class to manage bean definition names in a single place.
+ * 
+ * @author Oliver Gierke
+ */
+class BeanDefinitionNames {
+
+	public static final String JPA_MAPPING_CONTEXT_BEAN_NAME = "jpaMapppingContext";
+}

src/main/java/org/springframework/data/jpa/repository/config/JpaAuditingRegistrar.java

@@ -16,6 +16,7 @@
 package org.springframework.data.jpa.repository.config;
 
 import static org.springframework.data.jpa.domain.support.AuditingBeanFactoryPostProcessor.*;
+import static org.springframework.data.jpa.repository.config.BeanDefinitionNames.*;
 
 import java.lang.annotation.Annotation;
 
@@ -29,6 +30,7 @@ import org.springframework.beans.factory.support.RootBeanDefinition;
 import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
 import org.springframework.core.type.AnnotationMetadata;
 import org.springframework.data.auditing.config.AuditingBeanDefinitionRegistrarSupport;
+import org.springframework.data.auditing.config.AuditingConfiguration;
 import org.springframework.data.config.ParsingUtils;
 import org.springframework.data.jpa.domain.support.AuditingBeanFactoryPostProcessor;
 import org.springframework.data.jpa.domain.support.AuditingEntityListener;
@@ -62,6 +64,17 @@ class JpaAuditingRegistrar extends AuditingBeanDefinitionRegistrarSupport {
 		return "jpaAuditingHandler";
 	}
 
+	/* 
+	 * (non-Javadoc)
+	 * @see org.springframework.data.auditing.config.AuditingBeanDefinitionRegistrarSupport#getAuditHandlerBeanDefinitionBuilder(org.springframework.data.auditing.config.AuditingConfiguration)
+	 */
+	@Override
+	protected BeanDefinitionBuilder getAuditHandlerBeanDefinitionBuilder(AuditingConfiguration configuration) {
+
+		BeanDefinitionBuilder builder = super.getAuditHandlerBeanDefinitionBuilder(configuration);
+		return builder.addConstructorArgReference(JPA_MAPPING_CONTEXT_BEAN_NAME);
+	}
+
 	/* 
 	 * (non-Javadoc)
 	 * @see org.springframework.data.auditing.config.AuditingBeanDefinitionRegistrarSupport#registerBeanDefinitions(org.springframework.core.type.AnnotationMetadata, org.springframework.beans.factory.support.BeanDefinitionRegistry)
@@ -86,6 +99,11 @@ class JpaAuditingRegistrar extends AuditingBeanDefinitionRegistrarSupport {
 	protected void registerAuditListenerBeanDefinition(BeanDefinition auditingHandlerDefinition,
 			BeanDefinitionRegistry registry) {
 
+		if (!registry.containsBeanDefinition(JPA_MAPPING_CONTEXT_BEAN_NAME)) {
+			registry.registerBeanDefinition(JPA_MAPPING_CONTEXT_BEAN_NAME, //
+					new RootBeanDefinition(JpaMetamodelMappingContextFactoryBean.class));
+		}
+
 		BeanDefinitionBuilder builder = BeanDefinitionBuilder.rootBeanDefinition(AuditingEntityListener.class);
 		builder.addPropertyValue("auditingHandler",
 				ParsingUtils.getObjectFactoryBeanDefinition(getAuditingHandlerBeanName(), null));

src/main/java/org/springframework/data/jpa/repository/config/JpaMetamodelMappingContextFactoryBean.java

@@ -0,0 +1,95 @@
+package org.springframework.data.jpa.repository.config;
+
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Set;
+
+import javax.persistence.EntityManagerFactory;
+import javax.persistence.metamodel.ManagedType;
+import javax.persistence.metamodel.Metamodel;
+
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.BeanFactoryUtils;
+import org.springframework.beans.factory.FactoryBean;
+import org.springframework.beans.factory.ListableBeanFactory;
+import org.springframework.beans.factory.config.AbstractFactoryBean;
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationContextAware;
+import org.springframework.data.jpa.mapping.JpaMetamodelMappingContext;
+
+/**
+ * {@link FactoryBean} to setup {@link JpaMetamodelMappingContext} instances from Spring configuration.
+ * 
+ * @author Oliver Gierke
+ * @since 1.6
+ */
+class JpaMetamodelMappingContextFactoryBean extends AbstractFactoryBean<JpaMetamodelMappingContext> implements
+		ApplicationContextAware {
+
+	private ListableBeanFactory beanFactory;
+
+	/* 
+	 * (non-Javadoc)
+	 * @see org.springframework.context.ApplicationContextAware#setApplicationContext(org.springframework.context.ApplicationContext)
+	 */
+	@Override
+	public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
+		this.beanFactory = applicationContext;
+	}
+
+	/* 
+	 * (non-Javadoc)
+	 * @see org.springframework.beans.factory.config.AbstractFactoryBean#getObjectType()
+	 */
+	@Override
+	public Class<?> getObjectType() {
+		return JpaMetamodelMappingContext.class;
+	}
+
+	/* 
+	 * (non-Javadoc)
+	 * @see org.springframework.beans.factory.config.AbstractFactoryBean#createInstance()
+	 */
+	@Override
+	protected JpaMetamodelMappingContext createInstance() throws Exception {
+
+		Set<Metamodel> models = getMetamodels();
+		Set<Class<?>> entitySources = new HashSet<Class<?>>();
+
+		for (Metamodel metamodel : models) {
+
+			for (ManagedType<?> type : metamodel.getManagedTypes()) {
+
+				Class<?> javaType = type.getJavaType();
+
+				if (javaType != null) {
+					entitySources.add(javaType);
+				}
+			}
+		}
+
+		JpaMetamodelMappingContext context = new JpaMetamodelMappingContext(models);
+		context.setInitialEntitySet(entitySources);
+		context.initialize();
+
+		return context;
+	}
+
+	/**
+	 * Obtains all {@link Metamodel} instances of the current {@link ApplicationContext}.
+	 * 
+	 * @return
+	 */
+	private Set<Metamodel> getMetamodels() {
+
+		Collection<EntityManagerFactory> factories = BeanFactoryUtils.beansOfTypeIncludingAncestors(beanFactory,
+				EntityManagerFactory.class).values();
+		Set<Metamodel> metamodels = new HashSet<Metamodel>(factories.size());
+
+		for (EntityManagerFactory emf : factories) {
+			metamodels.add(emf.getMetamodel());
+		}
+
+		return metamodels;
+	}
+}

src/main/java/org/springframework/data/jpa/repository/config/JpaRepositoryConfigExtension.java

@@ -15,18 +15,11 @@
  */
 package org.springframework.data.jpa.repository.config;
 
-import java.util.HashSet;
-import java.util.Set;
+import static org.springframework.data.jpa.repository.config.BeanDefinitionNames.*;
 
-import javax.persistence.EntityManager;
 import javax.persistence.PersistenceContext;
 import javax.persistence.PersistenceUnit;
-import javax.persistence.metamodel.ManagedType;
-import javax.persistence.metamodel.Metamodel;
 
-import org.springframework.beans.factory.FactoryBean;
-import org.springframework.beans.factory.config.AbstractFactoryBean;
-import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.support.AbstractBeanDefinition;
 import org.springframework.beans.factory.support.BeanDefinitionBuilder;
 import org.springframework.beans.factory.support.BeanDefinitionRegistry;
@@ -34,13 +27,11 @@ import org.springframework.beans.factory.support.RootBeanDefinition;
 import org.springframework.context.annotation.AnnotationConfigUtils;
 import org.springframework.dao.DataAccessException;
 import org.springframework.dao.annotation.PersistenceExceptionTranslationPostProcessor;
-import org.springframework.data.jpa.mapping.JpaMetamodelMappingContext;
 import org.springframework.data.jpa.repository.support.EntityManagerBeanDefinitionRegistrarPostProcessor;
 import org.springframework.data.jpa.repository.support.JpaRepositoryFactoryBean;
 import org.springframework.data.repository.config.RepositoryConfigurationExtensionSupport;
 import org.springframework.data.repository.config.RepositoryConfigurationSource;
 import org.springframework.orm.jpa.support.PersistenceAnnotationBeanPostProcessor;
-import org.springframework.util.Assert;
 
 /**
  * JPA specific configuration extension parsing custom attributes from the XML namespace and
@@ -57,8 +48,6 @@ import org.springframework.util.Assert;
  */
 public class JpaRepositoryConfigExtension extends RepositoryConfigurationExtensionSupport {
 
-	public static final String JPA_MAPPING_CONTEXT_BEAN_NAME = "jpaMapppingContext";
-
 	private static final Class<?> PAB_POST_PROCESSOR = PersistenceAnnotationBeanPostProcessor.class;
 	private static final String DEFAULT_TRANSACTION_MANAGER_BEAN_NAME = "transactionManager";
 
@@ -89,24 +78,33 @@ public class JpaRepositoryConfigExtension extends RepositoryConfigurationExtensi
 		String transactionManagerRef = source.getAttribute("transactionManagerRef");
 		builder.addPropertyValue("transactionManager",
 				transactionManagerRef == null ? DEFAULT_TRANSACTION_MANAGER_BEAN_NAME : transactionManagerRef);
-
-		String entityManagerFactoryRef = getEntityManagerFactoryRef(source);
-
-		if (entityManagerFactoryRef != null) {
-			builder.addPropertyValue("entityManager", getEntityManagerBeanDefinitionFor(entityManagerFactoryRef, source));
-		}
-
+		builder.addPropertyValue("entityManager", getEntityManagerBeanDefinitionFor(source, source.getSource()));
 		builder.addPropertyReference("mappingContext", JPA_MAPPING_CONTEXT_BEAN_NAME);
 	}
 
-	/**
-	 * @param source
-	 * @return
+	/* 
+	 * (non-Javadoc)
+	 * @see org.springframework.data.repository.config.RepositoryConfigurationExtensionSupport#registerBeansForRoot(org.springframework.beans.factory.support.BeanDefinitionRegistry, org.springframework.data.repository.config.RepositoryConfigurationSource)
 	 */
-	private String getEntityManagerFactoryRef(RepositoryConfigurationSource source) {
+	@Override
+	public void registerBeansForRoot(BeanDefinitionRegistry registry, RepositoryConfigurationSource config) {
 
-		String entityManagerFactoryRef = source.getAttribute("entityManagerFactoryRef");
-		return entityManagerFactoryRef == null ? "entityManagerFactory" : entityManagerFactoryRef;
+		super.registerBeansForRoot(registry, config);
+
+		Object source = config.getSource();
+		registerWithSourceAndGeneratedBeanName(registry, new RootBeanDefinition(
+				EntityManagerBeanDefinitionRegistrarPostProcessor.class), source);
+
+		if (!registry.containsBeanDefinition(JPA_MAPPING_CONTEXT_BEAN_NAME)) {
+			registry.registerBeanDefinition(JPA_MAPPING_CONTEXT_BEAN_NAME, //
+					new RootBeanDefinition(JpaMetamodelMappingContextFactoryBean.class));
+		}
+
+		if (!hasBean(PAB_POST_PROCESSOR, registry)
+				&& !registry.containsBeanDefinition(AnnotationConfigUtils.PERSISTENCE_ANNOTATION_PROCESSOR_BEAN_NAME)) {
+
+			registerWithSourceAndGeneratedBeanName(registry, new RootBeanDefinition(PAB_POST_PROCESSOR), source);
+		}
 	}
 
 	/**
@@ -117,12 +115,13 @@ public class JpaRepositoryConfigExtension extends RepositoryConfigurationExtensi
 	 * @param source
 	 * @return
 	 */
-	private BeanDefinition getEntityManagerBeanDefinitionFor(String entityManagerFactoryBeanName, Object source) {
+	private static AbstractBeanDefinition getEntityManagerBeanDefinitionFor(RepositoryConfigurationSource config,
+			Object source) {
 
 		BeanDefinitionBuilder builder = BeanDefinitionBuilder
 				.rootBeanDefinition("org.springframework.orm.jpa.SharedEntityManagerCreator");
 		builder.setFactoryMethod("createSharedEntityManager");
-		builder.addConstructorArgReference(entityManagerFactoryBeanName);
+		builder.addConstructorArgReference(getEntityManagerBeanRef(config));
 
 		AbstractBeanDefinition bean = builder.getRawBeanDefinition();
 		bean.setSource(source);
@@ -130,105 +129,9 @@ public class JpaRepositoryConfigExtension extends RepositoryConfigurationExtensi
 		return bean;
 	}
 
-	/* 
-	 * (non-Javadoc)
-	 * @see org.springframework.data.repository.config.RepositoryConfigurationExtensionSupport#registerBeansForRoot(org.springframework.beans.factory.support.BeanDefinitionRegistry, org.springframework.data.repository.config.RepositoryConfigurationSource)
-	 */
-	@Override
-	public void registerBeansForRoot(BeanDefinitionRegistry registry, RepositoryConfigurationSource configurationSource) {
+	private static String getEntityManagerBeanRef(RepositoryConfigurationSource config) {
 
-		super.registerBeansForRoot(registry, configurationSource);
-
-		Object source = configurationSource.getSource();
-		registerWithSourceAndGeneratedBeanName(registry, new RootBeanDefinition(
-				EntityManagerBeanDefinitionRegistrarPostProcessor.class), source);
-
-		BeanDefinition entityManagerBeanDefinitionFor = getEntityManagerBeanDefinitionFor(
-				getEntityManagerFactoryRef(configurationSource), source);
-
-		BeanDefinitionBuilder builder = BeanDefinitionBuilder
-				.rootBeanDefinition(JpaMetamodelMappingContextFactoryBean.class);
-		builder.addPropertyValue("entityManager", entityManagerBeanDefinitionFor);
-
-		AbstractBeanDefinition definition = builder.getBeanDefinition();
-		definition.setSource(source);
-		registry.registerBeanDefinition(JPA_MAPPING_CONTEXT_BEAN_NAME, definition);
-
-		if (!hasBean(PAB_POST_PROCESSOR, registry)
-				&& !registry.containsBeanDefinition(AnnotationConfigUtils.PERSISTENCE_ANNOTATION_PROCESSOR_BEAN_NAME)) {
-
-			registerWithSourceAndGeneratedBeanName(registry, new RootBeanDefinition(PAB_POST_PROCESSOR), source);
-		}
-	}
-
-	/**
-	 * {@link FactoryBean} to setup {@link JpaMetamodelMappingContext} instances from Spring configuration.
-	 * 
-	 * @author Oliver Gierke
-	 * @since 1.6
-	 */
-	static class JpaMetamodelMappingContextFactoryBean extends AbstractFactoryBean<JpaMetamodelMappingContext> {
-
-		private EntityManager entityManager;
-
-		/**
-		 * Configures the {@link EntityManager} to use to create the {@link JpaMetamodelMappingContext}.
-		 * 
-		 * @param entityManager must not be {@literal null}.
-		 */
-		public void setEntityManager(EntityManager entityManager) {
-
-			Assert.notNull(entityManager, "EntityManager must not be null!");
-			this.entityManager = entityManager;
-		}
-
-		/* 
-		 * (non-Javadoc)
-		 * @see org.springframework.beans.factory.config.AbstractFactoryBean#getObjectType()
-		 */
-		@Override
-		public Class<?> getObjectType() {
-			return JpaMetamodelMappingContext.class;
-		}
-
-		/* 
-		 * (non-Javadoc)
-		 * @see org.springframework.beans.factory.config.AbstractFactoryBean#createInstance()
-		 */
-		@Override
-		protected JpaMetamodelMappingContext createInstance() throws Exception {
-
-			Metamodel metamodel = entityManager.getMetamodel();
-
-			Set<ManagedType<?>> managedTypes = metamodel.getManagedTypes();
-			Set<Class<?>> entitySources = new HashSet<Class<?>>(managedTypes.size());
-
-			for (ManagedType<?> type : managedTypes) {
-
-				Class<?> javaType = type.getJavaType();
-
-				if (javaType != null) {
-					entitySources.add(javaType);
-				}
-			}
-
-			JpaMetamodelMappingContext context = new JpaMetamodelMappingContext(metamodel);
-			context.setInitialEntitySet(entitySources);
-			context.initialize();
-
-			return context;
-		}
-
-		/* 
-		 * (non-Javadoc)
-		 * @see org.springframework.beans.factory.config.AbstractFactoryBean#afterPropertiesSet()
-		 */
-		@Override
-		public void afterPropertiesSet() throws Exception {
-
-			Assert.notNull(entityManager, "EntityManager must not be null!");
-
-			super.afterPropertiesSet();
-		}
+		String entityManagerFactoryRef = config == null ? null : config.getAttribute("entityManagerFactoryRef");
+		return entityManagerFactoryRef == null ? "entityManagerFactory" : entityManagerFactoryRef;
 	}
 }

src/test/java/org/springframework/data/jpa/repository/config/JpaRepositoryConfigExtensionUnitTests.java

@@ -23,7 +23,7 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.Set;
 
-import javax.persistence.EntityManager;
+import javax.persistence.EntityManagerFactory;
 import javax.persistence.metamodel.ManagedType;
 import javax.persistence.metamodel.Metamodel;
 
@@ -38,9 +38,8 @@ import org.springframework.beans.factory.NoSuchBeanDefinitionException;
 import org.springframework.beans.factory.support.BeanDefinitionReaderUtils;
 import org.springframework.beans.factory.support.DefaultListableBeanFactory;
 import org.springframework.beans.factory.support.RootBeanDefinition;
+import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.AnnotationConfigUtils;
-import org.springframework.data.jpa.mapping.JpaMetamodelMappingContext;
-import org.springframework.data.jpa.repository.config.JpaRepositoryConfigExtension.JpaMetamodelMappingContextFactoryBean;
 import org.springframework.data.repository.config.RepositoryConfigurationExtension;
 import org.springframework.data.repository.config.RepositoryConfigurationSource;
 import org.springframework.orm.jpa.support.PersistenceAnnotationBeanPostProcessor;
@@ -102,20 +101,21 @@ public class JpaRepositoryConfigExtensionUnitTests {
 	@Test
 	public void guardsAgainstNullJavaTypesReturnedFromJpaMetamodel() throws Exception {
 
-		EntityManager em = mock(EntityManager.class);
+		ApplicationContext context = mock(ApplicationContext.class);
+		EntityManagerFactory emf = mock(EntityManagerFactory.class);
 		Metamodel metamodel = mock(Metamodel.class);
 		ManagedType<?> managedType = mock(ManagedType.class);
 
 		Set<ManagedType<?>> managedTypes = Collections.<ManagedType<?>> singleton(managedType);
 
-		when(em.getMetamodel()).thenReturn(metamodel);
+		when(context.getBeansOfType(EntityManagerFactory.class)).thenReturn(Collections.singletonMap("emf", emf));
+		when(emf.getMetamodel()).thenReturn(metamodel);
 		when(metamodel.getManagedTypes()).thenReturn(managedTypes);
 
-		JpaMetamodelMappingContextFactoryBean factoryBean = new JpaRepositoryConfigExtension.JpaMetamodelMappingContextFactoryBean();
-		factoryBean.setEntityManager(em);
+		JpaMetamodelMappingContextFactoryBean factoryBean = new JpaMetamodelMappingContextFactoryBean();
+		factoryBean.setApplicationContext(context);
 
-		JpaMetamodelMappingContext context = factoryBean.createInstance();
-		context.afterPropertiesSet();
+		factoryBean.createInstance().afterPropertiesSet();
 	}
 
 	private void assertOnlyOnePersistenceAnnotationBeanPostProcessorRegistered(DefaultListableBeanFactory factory,