001/* 002 * Copyright 2012-2017 the original author or authors. 003 * 004 * Licensed under the Apache License, Version 2.0 (the "License"); 005 * you may not use this file except in compliance with the License. 006 * You may obtain a copy of the License at 007 * 008 * http://www.apache.org/licenses/LICENSE-2.0 009 * 010 * Unless required by applicable law or agreed to in writing, software 011 * distributed under the License is distributed on an "AS IS" BASIS, 012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 013 * See the License for the specific language governing permissions and 014 * limitations under the License. 015 */ 016 017package org.springframework.boot.test.autoconfigure.orm.jpa; 018 019import javax.persistence.EntityManager; 020import javax.persistence.EntityManagerFactory; 021import javax.persistence.PersistenceUnitUtil; 022 023import org.springframework.orm.jpa.EntityManagerFactoryUtils; 024import org.springframework.util.Assert; 025 026/** 027 * Alternative to {@link EntityManager} for use in JPA tests. Provides a subset of 028 * {@link EntityManager} methods that are useful for tests as well as helper methods for 029 * common testing tasks such as {@link #persistFlushFind(Object) persist/flush/find}. 030 * 031 * @author Phillip Webb 032 * @since 1.4.0 033 */ 034public class TestEntityManager { 035 036 private final EntityManagerFactory entityManagerFactory; 037 038 /** 039 * Create a new {@link TestEntityManager} instance for the given 040 * {@link EntityManagerFactory}. 041 * @param entityManagerFactory the source entity manager factory 042 */ 043 public TestEntityManager(EntityManagerFactory entityManagerFactory) { 044 Assert.notNull(entityManagerFactory, "EntityManagerFactory must not be null"); 045 this.entityManagerFactory = entityManagerFactory; 046 } 047 048 /** 049 * Make an instance managed and persistent then return it's ID. Delegates to 050 * {@link EntityManager#persist(Object)} then {@link #getId(Object)}. 051 * <p> 052 * Helpful when setting up test data in a test: <pre class="code"> 053 * Object entityId = this.testEntityManager.persist(new MyEntity("Spring")); 054 * </pre> 055 * @param entity the source entity 056 * @return the ID of the newly persisted entity 057 */ 058 public Object persistAndGetId(Object entity) { 059 persist(entity); 060 return getId(entity); 061 } 062 063 /** 064 * Make an instance managed and persistent then return it's ID. Delegates to 065 * {@link EntityManager#persist(Object)} then {@link #getId(Object, Class)}. 066 * <p> 067 * Helpful when setting up test data in a test: <pre class="code"> 068 * Long entityId = this.testEntityManager.persist(new MyEntity("Spring"), Long.class); 069 * </pre> 070 * @param <T> the ID type 071 * @param entity the source entity 072 * @param idType the ID type 073 * @return the ID of the newly persisted entity 074 */ 075 public <T> T persistAndGetId(Object entity, Class<T> idType) { 076 persist(entity); 077 return getId(entity, idType); 078 079 } 080 081 /** 082 * Make an instance managed and persistent. Delegates to 083 * {@link EntityManager#persist(Object)} then returns the original source entity. 084 * <p> 085 * Helpful when setting up test data in a test: <pre class="code"> 086 * MyEntity entity = this.testEntityManager.persist(new MyEntity("Spring")); 087 * </pre> 088 * @param <E> the entity type 089 * @param entity the entity to persist 090 * @return the persisted entity 091 */ 092 public <E> E persist(E entity) { 093 getEntityManager().persist(entity); 094 return entity; 095 } 096 097 /** 098 * Make an instance managed and persistent, synchronize the persistence context to the 099 * underlying database and finally find the persisted entity by its ID. Delegates to 100 * {@link #persistAndFlush(Object)} then {@link #find(Class, Object)} with the 101 * {@link #getId(Object) entity ID}. 102 * <p> 103 * Helpful when ensuring that entity data is actually written and read from the 104 * underlying database correctly. 105 * @param <E> the entity type 106 * @param entity the entity to persist 107 * @return the entity found using the ID of the persisted entity 108 */ 109 @SuppressWarnings("unchecked") 110 public <E> E persistFlushFind(E entity) { 111 EntityManager entityManager = getEntityManager(); 112 persistAndFlush(entity); 113 Object id = getId(entity); 114 entityManager.detach(entity); 115 return (E) entityManager.find(entity.getClass(), id); 116 } 117 118 /** 119 * Make an instance managed and persistent then synchronize the persistence context to 120 * the underlying database. Delegates to {@link EntityManager#persist(Object)} then 121 * {@link #flush()} and finally returns the original source entity. 122 * <p> 123 * Helpful when setting up test data in a test: <pre class="code"> 124 * MyEntity entity = this.testEntityManager.persistAndFlush(new MyEntity("Spring")); 125 * </pre> 126 * @param <E> the entity type 127 * @param entity the entity to persist 128 * @return the persisted entity 129 */ 130 public <E> E persistAndFlush(E entity) { 131 persist(entity); 132 flush(); 133 return entity; 134 } 135 136 /** 137 * Merge the state of the given entity into the current persistence context. Delegates 138 * to {@link EntityManager#merge(Object)} 139 * @param <E> the entity type 140 * @param entity the entity to merge 141 * @return the merged entity 142 */ 143 public <E> E merge(E entity) { 144 return getEntityManager().merge(entity); 145 } 146 147 /** 148 * Remove the entity instance. Delegates to {@link EntityManager#remove(Object)} 149 * @param entity the entity to remove 150 */ 151 public void remove(Object entity) { 152 getEntityManager().remove(entity); 153 } 154 155 /** 156 * Find by primary key. Delegates to {@link EntityManager#find(Class, Object)}. 157 * @param <E> the entity type 158 * @param entityClass the entity class 159 * @param primaryKey the entity primary key 160 * @return the found entity or {@code null} if the entity does not exist 161 * @see #getId(Object) 162 */ 163 public <E> E find(Class<E> entityClass, Object primaryKey) { 164 return getEntityManager().find(entityClass, primaryKey); 165 } 166 167 /** 168 * Synchronize the persistence context to the underlying database. Delegates to 169 * {@link EntityManager#flush()}. 170 */ 171 public void flush() { 172 getEntityManager().flush(); 173 } 174 175 /** 176 * Refresh the state of the instance from the database, overwriting changes made to 177 * the entity, if any. Delegates to {@link EntityManager#refresh(Object)}. 178 * @param <E> the entity type 179 * @param entity the entity to refresh 180 * @return the refreshed entity 181 */ 182 public <E> E refresh(E entity) { 183 getEntityManager().refresh(entity); 184 return entity; 185 } 186 187 /** 188 * Clear the persistence context, causing all managed entities to become detached. 189 * Delegates to {@link EntityManager#clear()} 190 */ 191 public void clear() { 192 getEntityManager().clear(); 193 } 194 195 /** 196 * Remove the given entity from the persistence context, causing a managed entity to 197 * become detached. Delegates to {@link EntityManager#detach(Object)}. 198 * @param entity the entity to detach. 199 */ 200 public void detach(Object entity) { 201 getEntityManager().detach(entity); 202 } 203 204 /** 205 * Return the ID of the given entity. Delegates to 206 * {@link PersistenceUnitUtil#getIdentifier(Object)}. 207 * @param entity the source entity 208 * @return the ID of the entity or {@code null} 209 * @see #getId(Object, Class) 210 */ 211 public Object getId(Object entity) { 212 return this.entityManagerFactory.getPersistenceUnitUtil().getIdentifier(entity); 213 } 214 215 /** 216 * Return the ID of the given entity cast to a specific type. Delegates to 217 * {@link PersistenceUnitUtil#getIdentifier(Object)}. 218 * @param <T> the ID type 219 * @param entity the source entity 220 * @param idType the expected ID type 221 * @return the ID of the entity or {@code null} 222 * @see #getId(Object) 223 */ 224 @SuppressWarnings("unchecked") 225 public <T> T getId(Object entity, Class<T> idType) { 226 Object id = getId(entity); 227 Assert.isInstanceOf(idType, id, "ID mismatch:"); 228 return (T) id; 229 } 230 231 /** 232 * Return the underlying {@link EntityManager} that's actually used to perform all 233 * operations. 234 * @return the entity manager 235 */ 236 public final EntityManager getEntityManager() { 237 EntityManager manager = EntityManagerFactoryUtils 238 .getTransactionalEntityManager(this.entityManagerFactory); 239 Assert.state(manager != null, "No transactional EntityManager found"); 240 return manager; 241 } 242 243}