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}