HHH-17824 - Extend the use of @Jpa to test methods

Signed-off-by: Jan Schatteman <jschatte@redhat.com>

hibernate-core/src/main/java/org/hibernate/jpa/boot/internal/EntityManagerFactoryBuilderImpl.java

@@ -310,8 +310,6 @@ public class EntityManagerFactoryBuilderImpl implements EntityManagerFactoryBuil
 			}
 
 			// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-			// push back class transformation to the environment; for the time being this only has any effect in EE
-			// container situations, calling back into PersistenceUnitInfo#addClassTransformer
 
 			final boolean dirtyTrackingEnabled;
 			Object propertyValue = configurationValues.remove( ENHANCER_ENABLE_DIRTY_TRACKING );
@@ -344,6 +342,8 @@ public class EntityManagerFactoryBuilderImpl implements EntityManagerFactoryBuil
 						associationManagementEnabled
 				);
 
+				// push back class transformation to the environment; for the time being this only has any effect in EE
+				// container situations, calling back into PersistenceUnitInfo#addClassTransformer
 				persistenceUnit.pushClassTransformer( enhancementContext );
 				final ClassTransformer classTransformer = persistenceUnit.getClassTransformer();
 				if ( classTransformer != null ) {

hibernate-testing/src/main/java/org/hibernate/testing/orm/junit/EntityManagerFactoryExtension.java

@@ -36,7 +36,7 @@ import org.hibernate.testing.orm.domain.DomainModelDescriptor;
 import org.hibernate.testing.orm.domain.StandardDomainModel;
 import org.hibernate.testing.orm.jpa.PersistenceUnitInfoImpl;
 import org.hibernate.testing.util.ServiceRegistryUtil;
-import org.junit.jupiter.api.extension.AfterAllCallback;
+import org.junit.jupiter.api.extension.BeforeEachCallback;
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.jupiter.api.extension.TestExecutionExceptionHandler;
 import org.junit.jupiter.api.extension.TestInstancePostProcessor;
@@ -54,19 +54,26 @@ import org.jboss.logging.Logger;
  * @see SessionFactoryExtension
  */
 public class EntityManagerFactoryExtension
-		implements TestInstancePostProcessor, AfterAllCallback, TestExecutionExceptionHandler {
+		implements TestInstancePostProcessor, BeforeEachCallback, TestExecutionExceptionHandler {
 
 	private static final Logger log = Logger.getLogger( EntityManagerFactoryExtension.class );
 	private static final String EMF_KEY = EntityManagerFactoryScope.class.getName();
 
-	private static ExtensionContext.Store locateExtensionStore(Object testInstance, ExtensionContext context) {
+	private static ExtensionContext.Store locateExtensionStore(Object testScope, ExtensionContext context) {
-		return JUnitHelper.locateExtensionStore( EntityManagerFactoryExtension.class, context, testInstance );
+		return JUnitHelper.locateExtensionStore( EntityManagerFactoryExtension.class, context, testScope );
 	}
 
 	public static EntityManagerFactoryScope findEntityManagerFactoryScope(
-			Object testInstance,
+			Object testScope,
+			Optional<Jpa> emfAnnWrapper,
 			ExtensionContext context) {
-		final ExtensionContext.Store store = locateExtensionStore( testInstance, context );
+
+		if ( emfAnnWrapper.isEmpty() ) {
+			// No annotation on the test class, should be on the test methods
+			return null;
+		}
+
+		final ExtensionContext.Store store = locateExtensionStore( testScope, context );
 		final EntityManagerFactoryScope existing = (EntityManagerFactoryScope) store.get( EMF_KEY );
 		if ( existing != null ) {
 			return existing;
@@ -75,12 +82,7 @@ public class EntityManagerFactoryExtension
 		if ( !context.getElement().isPresent() ) {
 			throw new RuntimeException( "Unable to determine how to handle given ExtensionContext : " + context.getDisplayName() );
 		}
-
+		final Jpa emfAnn = emfAnnWrapper.get();
-		final Optional<Jpa> emfAnnWrapper = AnnotationSupport.findAnnotation(
-				context.getElement().get(),
-				Jpa.class
-		);
-		final Jpa emfAnn = emfAnnWrapper.orElseThrow( () -> new RuntimeException( "Could not locate @EntityManagerFactory" ) );
 
 		final PersistenceUnitInfoImpl pui = new PersistenceUnitInfoImpl( emfAnn.persistenceUnitName() );
 		( (Map<Object, Object>) Environment.getProperties() ).forEach(
@@ -196,7 +198,7 @@ public class EntityManagerFactoryExtension
 		ServiceRegistryUtil.applySettings( integrationSettings );
 		final EntityManagerFactoryScopeImpl scope = new EntityManagerFactoryScopeImpl( pui, integrationSettings );
 
-		locateExtensionStore( testInstance, context ).put( EMF_KEY, scope );
+		store.put( EMF_KEY, scope );
 
 		return scope;
 	}
@@ -267,29 +269,31 @@ public class EntityManagerFactoryExtension
 	}
 
 	@Override
-	public void postProcessTestInstance(Object testInstance, ExtensionContext context) {
+	public void beforeEach(ExtensionContext context) {
-		log.tracef( "#postProcessTestInstance(%s, %s)", testInstance, context.getDisplayName() );
+		log.tracef( "#beforeEach(%s)", context.getDisplayName() );
+		final Optional<Jpa> emfAnnWrapper = AnnotationSupport.findAnnotation(
+				context.getRequiredTestMethod(),
+				Jpa.class
+		);
 
-		findEntityManagerFactoryScope( testInstance, context );
+		if ( emfAnnWrapper.isEmpty() ) {
+			// assume the annotation is defined on the class-level...
+			return;
+		}
+
+		findEntityManagerFactoryScope( context.getRequiredTestMethod(), emfAnnWrapper, context );
 	}
 
 	@Override
-	public void afterAll(ExtensionContext context) {
+	public void postProcessTestInstance(Object testInstance, ExtensionContext context) {
-		log.tracef( "#afterAll(%s)", context.getDisplayName() );
+		log.tracef( "#postProcessTestInstance(%s, %s)", testInstance, context.getDisplayName() );
 
-		final Object testInstance = context.getRequiredTestInstance();
+		final Optional<Jpa> emfAnnWrapper = AnnotationSupport.findAnnotation(
+				context.getRequiredTestClass(),
+				Jpa.class
+		);
 
-		if ( testInstance instanceof SessionFactoryScopeAware ) {
+		findEntityManagerFactoryScope( testInstance, emfAnnWrapper, context );
-			( (SessionFactoryScopeAware) testInstance ).injectSessionFactoryScope( null );
-		}
-
-		final EntityManagerFactoryScopeImpl removed = (EntityManagerFactoryScopeImpl) locateExtensionStore(
-				testInstance,
-				context
-		).remove( EMF_KEY );
-		if ( removed != null ) {
-			removed.close();
-		}
 	}
 
 	@Override

hibernate-testing/src/main/java/org/hibernate/testing/orm/junit/EntityManagerFactoryParameterResolver.java

@@ -6,12 +6,15 @@
  */
 package org.hibernate.testing.orm.junit;
 
+import java.util.Optional;
+
 import jakarta.persistence.EntityManagerFactory;
 
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.jupiter.api.extension.ParameterContext;
 import org.junit.jupiter.api.extension.ParameterResolutionException;
 import org.junit.jupiter.api.extension.ParameterResolver;
+import org.junit.platform.commons.support.AnnotationSupport;
 
 import static org.hibernate.testing.orm.junit.EntityManagerFactoryExtension.findEntityManagerFactoryScope;
 
@@ -34,8 +37,31 @@ public class EntityManagerFactoryParameterResolver implements ParameterResolver
 	public Object resolveParameter(
 			ParameterContext parameterContext,
 			ExtensionContext extensionContext) throws ParameterResolutionException {
+
+		// Fall back on the test class annotation in case the method isn't annotated or we're in a @Before/@After method
+		Optional<Jpa> emfAnnWrapper = AnnotationSupport.findAnnotation(
+				extensionContext.getRequiredTestClass(),
+				Jpa.class
+		);
+		Object testScope = extensionContext.getRequiredTestInstance();
+
+		// coming from a @Test
+		if (parameterContext.getDeclaringExecutable() instanceof java.lang.reflect.Method && !extensionContext.getTestMethod().isEmpty()) {
+
+			Optional<Jpa> testEmfAnnWrapper = AnnotationSupport.findAnnotation(
+					extensionContext.getRequiredTestMethod(),
+					Jpa.class
+			);
+			// @Jpa on the test, so override the class annotation
+			if ( !testEmfAnnWrapper.isEmpty() ) {
+				testScope = extensionContext.getRequiredTestMethod();
+				emfAnnWrapper = testEmfAnnWrapper;
+			}
+		}
+
 		final EntityManagerFactoryScope scope = findEntityManagerFactoryScope(
-				extensionContext.getRequiredTestInstance(),
+				testScope,
+				emfAnnWrapper,
 				extensionContext
 		);
 

hibernate-testing/src/main/java/org/hibernate/testing/orm/junit/Jpa.java

@@ -30,7 +30,7 @@ import org.junit.jupiter.api.extension.ExtendWith;
  * @author Steve Ebersole
  */
 @Inherited
-@Target( ElementType.TYPE )
+@Target( {ElementType.TYPE, ElementType.METHOD} )
 @Retention( RetentionPolicy.RUNTIME )
 
 @TestInstance( TestInstance.Lifecycle.PER_CLASS )

hibernate-testing/src/test/java/org/hibernate/testing/annotations/BasicEntityManagerFactoryScopeTests.java

@@ -33,4 +33,4 @@ public class BasicEntityManagerFactoryScopeTests {
 		);
 	}
 
-}
+}

hibernate-testing/src/test/java/org/hibernate/testing/annotations/methods/EntityManagerFactoryScopeTesting.java

@@ -0,0 +1,119 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or http://www.gnu.org/licenses/lgpl-2.1.html
+ */
+package org.hibernate.testing.annotations.methods;
+
+import java.util.Set;
+
+import org.hibernate.dialect.H2Dialect;
+
+import org.hibernate.testing.annotations.AnEntity;
+import org.hibernate.testing.annotations.AnotherEntity;
+import org.hibernate.testing.orm.junit.EntityManagerFactoryScope;
+import org.hibernate.testing.orm.junit.Jpa;
+import org.hibernate.testing.orm.junit.RequiresDialect;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+import jakarta.persistence.metamodel.EntityType;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.notNullValue;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+@RequiresDialect(H2Dialect.class)
+@Jpa(
+		annotatedClasses = {
+				AnEntity.class
+		}
+)
+public class EntityManagerFactoryScopeTesting {
+
+	@BeforeAll
+	public void setup(EntityManagerFactoryScope scope) {
+		scope.inTransaction(
+				entityManager -> {
+					AnEntity ae = new AnEntity(1, "AnEntity_1");
+					entityManager.persist( ae );
+				}
+		);
+	}
+
+	@AfterAll
+	public void tearDown(EntityManagerFactoryScope scope) {
+		scope.inTransaction(
+				entityManager -> {
+					entityManager.createQuery( "delete from AnEntity" ).executeUpdate();
+				}
+		);
+	}
+
+	@Test
+	public void testBasicUsage(EntityManagerFactoryScope scope) {
+		assertThat( scope, notNullValue() );
+		assertThat( scope.getEntityManagerFactory(), notNullValue() );
+		// check we can use the EMF to create EMs
+		scope.inTransaction(
+				(session) -> session.createQuery( "select a from AnEntity a" ).getResultList()
+		);
+	}
+
+	@Test
+	public void nonAnnotatedMethodTest(EntityManagerFactoryScope scope) {
+		Set<EntityType<?>> entities = scope.getEntityManagerFactory().getMetamodel().getEntities();
+		assertEquals( 1, entities.size() );
+		assertEquals( "AnEntity", entities.iterator().next().getName() );
+		assertEquals( Boolean.FALSE, scope.getEntityManagerFactory().getProperties().get( "hibernate.jpa.compliance.query" ) );
+		scope.inEntityManager(
+				entityManager -> {
+					AnEntity ae = entityManager.find( AnEntity.class, 1 );
+					assertNotNull( ae );
+					assertEquals( 1, ae.getId() );
+					assertEquals( "AnEntity_1", ae.getName() );
+				}
+		);
+	}
+
+	@Jpa(
+			annotatedClasses = AnotherEntity.class,
+			queryComplianceEnabled = true
+	)
+	@Test
+	public void annotatedMethodTest(EntityManagerFactoryScope scope) {
+		assertThat( scope, notNullValue() );
+		assertThat( scope.getEntityManagerFactory(), notNullValue() );
+		Set<EntityType<?>> entities = scope.getEntityManagerFactory().getMetamodel().getEntities();
+		assertEquals( 1, entities.size() );
+		assertEquals( "AnotherEntity", entities.iterator().next().getName() );
+		assertEquals( Boolean.TRUE, scope.getEntityManagerFactory().getProperties().get( "hibernate.jpa.compliance.query" ) );
+		scope.inTransaction(
+				entityManager -> {
+					AnotherEntity aoe = new AnotherEntity( 2, "AnotherEntity_1" );
+					entityManager.persist( aoe );
+				}
+		);
+		scope.inTransaction(
+				entityManager -> {
+					AnotherEntity aoe = entityManager.find( AnotherEntity.class, 2 );
+					assertNotNull( aoe );
+					assertEquals( 2, aoe.getId() );
+					assertEquals( "AnotherEntity_1", aoe.getName() );
+				}
+		);
+		Assertions.assertThrows(
+				IllegalArgumentException.class,
+				() -> scope.inTransaction(
+						entityManager -> {
+							AnEntity ae = entityManager.find( AnEntity.class, 1 );
+						}
+				)
+		);
+	}
+
+}

hibernate-testing/src/test/java/org/hibernate/testing/annotations/methods/MoreEntityManagerFactoryScopeTesting.java

@@ -0,0 +1,82 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or http://www.gnu.org/licenses/lgpl-2.1.html
+ */
+package org.hibernate.testing.annotations.methods;
+
+import java.util.Set;
+
+import org.hibernate.dialect.H2Dialect;
+
+import org.hibernate.testing.annotations.AnEntity;
+import org.hibernate.testing.annotations.AnotherEntity;
+import org.hibernate.testing.orm.junit.EntityManagerFactoryScope;
+import org.hibernate.testing.orm.junit.Jpa;
+import org.hibernate.testing.orm.junit.RequiresDialect;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import jakarta.persistence.metamodel.EntityType;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.notNullValue;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+@RequiresDialect(H2Dialect.class)
+public class MoreEntityManagerFactoryScopeTesting {
+
+	@Jpa(
+			annotatedClasses = {
+					AnEntity.class
+			}
+	)
+	@Test
+	public void testBasicUsage(EntityManagerFactoryScope scope) {
+		assertThat( scope, notNullValue() );
+		assertThat( scope.getEntityManagerFactory(), notNullValue() );
+		// check we can use the EMF to create EMs
+		scope.inTransaction(
+				(session) -> session.createQuery( "select a from AnEntity a" ).getResultList()
+		);
+	}
+
+	@Jpa(
+			annotatedClasses = AnotherEntity.class,
+			queryComplianceEnabled = true
+	)
+	@Test
+	public void annotatedMethodTest(EntityManagerFactoryScope scope) {
+		assertThat( scope, notNullValue() );
+		assertThat( scope.getEntityManagerFactory(), notNullValue() );
+		Set<EntityType<?>> entities = scope.getEntityManagerFactory().getMetamodel().getEntities();
+		assertEquals( 1, entities.size() );
+		assertEquals( "AnotherEntity", entities.iterator().next().getName() );
+		assertEquals( Boolean.TRUE, scope.getEntityManagerFactory().getProperties().get( "hibernate.jpa.compliance.query" ) );
+		scope.inTransaction(
+				entityManager -> {
+					AnotherEntity aoe = new AnotherEntity( 2, "AnotherEntity_1" );
+					entityManager.persist( aoe );
+				}
+		);
+		scope.inTransaction(
+				entityManager -> {
+					AnotherEntity aoe = entityManager.find( AnotherEntity.class, 2 );
+					assertNotNull( aoe );
+					assertEquals( 2, aoe.getId() );
+					assertEquals( "AnotherEntity_1", aoe.getName() );
+				}
+		);
+		Assertions.assertThrows(
+				IllegalArgumentException.class,
+				() -> scope.inTransaction(
+						entityManager -> {
+							AnEntity ae = entityManager.find( AnEntity.class, 1 );
+						}
+				)
+		);
+	}
+
+}