HHH-19912 Add automatic test data cleanup feature to @SessionFactory annotation

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

@@ -0,0 +1,37 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.testing.orm.junit;
+
+/**
+ * Enumeration of when to drop test data automatically.
+ *
+ * @author inpink
+ */
+public enum DropDataTiming {
+	/**
+	 * Never drop test data automatically
+	 */
+	NEVER,
+
+	/**
+	 * Drop test data before each test method
+	 */
+	BEFORE_EACH,
+
+	/**
+	 * Drop test data after each test method
+	 */
+	AFTER_EACH,
+
+	/**
+	 * Drop test data before all test methods (once per test class)
+	 */
+	BEFORE_ALL,
+
+	/**
+	 * Drop test data after all test methods (once per test class)
+	 */
+	AFTER_ALL
+}

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

@@ -19,6 +19,7 @@
 
 /**
  * @author Steve Ebersole
+ * @author inpink
  */
 @Inherited
 @Target({ElementType.TYPE, ElementType.METHOD})
@@ -56,4 +57,11 @@
 	boolean useCollectingStatementInspector() default false;
 
 	boolean applyCollectionsInDefaultFetchGroup() default true;
+
+	/**
+	 * When to automatically drop test data.
+	 *
+	 * @return the timing for dropping test data
+	 */
+	DropDataTiming dropTestData() default DropDataTiming.NEVER;
 }

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

@@ -22,6 +22,9 @@
 import org.hibernate.tool.schema.spi.SchemaManagementToolCoordinator;
 import org.hibernate.tool.schema.spi.SchemaManagementToolCoordinator.ActionGrouping;
 import org.jboss.logging.Logger;
+import org.junit.jupiter.api.extension.AfterAllCallback;
+import org.junit.jupiter.api.extension.AfterEachCallback;
+import org.junit.jupiter.api.extension.BeforeAllCallback;
 import org.junit.jupiter.api.extension.BeforeEachCallback;
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.jupiter.api.extension.TestExecutionExceptionHandler;
@@ -43,12 +46,15 @@
  * @see DomainModelExtension
  *
  * @author Steve Ebersole
+ * @author inpink
  */
 public class SessionFactoryExtension
-		implements TestInstancePostProcessor, BeforeEachCallback, TestExecutionExceptionHandler {
+		implements TestInstancePostProcessor, BeforeAllCallback, BeforeEachCallback,
+		AfterEachCallback, AfterAllCallback, TestExecutionExceptionHandler {
 
 	private static final Logger log = Logger.getLogger( SessionFactoryExtension.class );
 	private static final String SESSION_FACTORY_KEY = SessionFactoryScope.class.getName();
+	private static final String DROP_DATA_TIMING_KEY = "DROP_DATA_TIMING";
 
 	/**
 	 * Intended for use from external consumers.  Will never create a scope, just
@@ -77,12 +83,16 @@ public void postProcessTestInstance(Object testInstance, ExtensionContext contex
 				|| SessionFactoryProducer.class.isAssignableFrom( context.getRequiredTestClass() ) ) {
 			final DomainModelScope domainModelScope = DomainModelExtension.getOrCreateDomainModelScope( testInstance, context );
 			final SessionFactoryScope created = createSessionFactoryScope( testInstance, sfAnnRef, domainModelScope, context );
-			locateExtensionStore( testInstance, context ).put( SESSION_FACTORY_KEY, created );
+			final ExtensionContext.Store store = locateExtensionStore( testInstance, context );
+			store.put( SESSION_FACTORY_KEY, created );
+			store.put( DROP_DATA_TIMING_KEY, sfAnnRef.map( SessionFactory::dropTestData ).orElse( DropDataTiming.NEVER ) );
 		}
 	}
 
 	@Override
 	public void beforeEach(ExtensionContext context) {
+		handleDropData(context, DropDataTiming.BEFORE_EACH);
+
 		final Optional<SessionFactory> sfAnnRef = AnnotationSupport.findAnnotation(
 				context.getRequiredTestMethod(),
 				SessionFactory.class
@@ -97,6 +107,7 @@ public void beforeEach(ExtensionContext context) {
 		final DomainModelScope domainModelScope = DomainModelExtension.resolveForMethodLevelSessionFactoryScope( context );
 		final SessionFactoryScope created = createSessionFactoryScope( context.getRequiredTestInstance(), sfAnnRef, domainModelScope, context );
 		final ExtensionContext.Store extensionStore = locateExtensionStore( context.getRequiredTestInstance(), context );
+		extensionStore.put( DROP_DATA_TIMING_KEY, sfAnnRef.map( SessionFactory::dropTestData ).orElse( DropDataTiming.NEVER ) );
 		extensionStore.put( SESSION_FACTORY_KEY, created );
 	}
 
@@ -231,6 +242,39 @@ public void handleTestExecutionException(ExtensionContext context, Throwable thr
 		throw throwable;
 	}
 
+	@Override
+	public void beforeAll(ExtensionContext context) throws Exception {
+		handleDropData(context, DropDataTiming.BEFORE_ALL);
+	}
+
+	@Override
+	public void afterEach(ExtensionContext context) throws Exception {
+		handleDropData(context, DropDataTiming.AFTER_EACH);
+	}
+
+	@Override
+	public void afterAll(ExtensionContext context) throws Exception {
+		handleDropData(context, DropDataTiming.AFTER_ALL);
+	}
+
+	private void handleDropData(ExtensionContext context, DropDataTiming timing) {
+		try {
+			final Object testInstance = context.getRequiredTestInstance();
+			final ExtensionContext.Store store = locateExtensionStore(testInstance, context);
+
+			final DropDataTiming configuredTiming = (DropDataTiming) store.get( DROP_DATA_TIMING_KEY );
+
+			if (configuredTiming != DropDataTiming.NEVER && configuredTiming == timing) {
+				final SessionFactoryScope scope = findSessionFactoryScope(testInstance, context);
+				scope.dropData();
+				log.debugf("Dropped data at timing %s for %s", timing, context.getDisplayName());
+			}
+		}
+		catch (Exception e) {
+			log.debugf("Failed to drop data at timing %s: %s", timing, e.getMessage());
+		}
+	}
+
 	private static class SessionFactoryScopeImpl implements SessionFactoryScope, AutoCloseable {
 		private final DomainModelScope modelScope;
 		private final SessionFactoryProducer producer;

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

@@ -0,0 +1,188 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.testing.annotations.methods;
+
+import org.hibernate.testing.annotations.AnEntity;
+import org.hibernate.testing.orm.junit.DomainModel;
+import org.hibernate.testing.orm.junit.DropDataTiming;
+import org.hibernate.testing.orm.junit.SessionFactory;
+import org.hibernate.testing.orm.junit.SessionFactoryScope;
+import org.junit.jupiter.api.MethodOrderer;
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestMethodOrder;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for dropTestData timing configuration.
+ *
+ * @author inpink
+ */
+public class DropDataTimingTest {
+
+	@Nested
+	@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
+	@DomainModel(annotatedClasses = AnEntity.class)
+	@SessionFactory(dropTestData = DropDataTiming.AFTER_EACH)
+	class AfterEachDropDataTests {
+
+		@Test
+		@Order(1)
+		public void testAfterEachDropData(SessionFactoryScope scope) {
+			scope.inTransaction(session -> {
+				AnEntity entity = new AnEntity(1, "After Each");
+				session.persist(entity);
+			});
+
+			scope.inTransaction(session -> {
+				Long count = session.createQuery("select count(e) from AnEntity e", Long.class)
+						.getSingleResult();
+				assertThat(count).isEqualTo(1L);
+			});
+		}
+
+		@Test
+		@Order(2)
+		public void verifyAfterEachDropDataCleanup(SessionFactoryScope scope) {
+			scope.inTransaction(session -> {
+				Long count = session.createQuery("select count(e) from AnEntity e", Long.class)
+						.getSingleResult();
+				assertThat(count).isEqualTo(0L);
+			});
+		}
+	}
+
+	@Nested
+	@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
+	@DomainModel(annotatedClasses = AnEntity.class)
+	@SessionFactory(dropTestData = DropDataTiming.BEFORE_EACH)
+	class BeforeEachDropDataTests {
+
+		@Test
+		@Order(1)
+		public void prepareBeforeEachDropData(SessionFactoryScope scope) {
+			scope.inTransaction(session -> {
+				AnEntity entity = new AnEntity(2, "Before Each");
+				session.persist(entity);
+			});
+		}
+
+		@Test
+		@Order(2)
+		public void testBeforeEachDropData(SessionFactoryScope scope) {
+			scope.inTransaction(session -> {
+				Long count = session.createQuery("select count(e) from AnEntity e", Long.class)
+						.getSingleResult();
+				assertThat(count).isEqualTo(0L);
+			});
+		}
+	}
+
+	@Nested
+	@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
+	@DomainModel(annotatedClasses = AnEntity.class)
+	@SessionFactory(dropTestData = DropDataTiming.NEVER)
+	class NeverDropDataTests {
+
+		@Test
+		@Order(1)
+		public void prepareNeverDropData(SessionFactoryScope scope) {
+			scope.inTransaction(session -> {
+				AnEntity entity = new AnEntity(3, "Never");
+				session.persist(entity);
+			});
+		}
+
+		@Test
+		@Order(2)
+		public void testNeverDropData(SessionFactoryScope scope) {
+			scope.inTransaction(session -> {
+				Long count = session.createQuery("select count(e) from AnEntity e", Long.class)
+						.getSingleResult();
+				assertThat(count).isGreaterThan(0L);
+			});
+
+		}
+	}
+
+	@Nested
+	@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
+	@DomainModel(annotatedClasses = AnEntity.class)
+	class MethodLevelDropDataTests {
+
+		@Test
+		@Order(1)
+		@SessionFactory(dropTestData = DropDataTiming.AFTER_EACH)
+		public void methodLevelAfterEach(SessionFactoryScope scope) {
+			scope.inTransaction(session -> {
+				AnEntity entity = new AnEntity(10, "Method After Each");
+				session.persist(entity);
+			});
+
+			scope.inTransaction(session -> {
+				Long count = session.createQuery("select count(e) from AnEntity e", Long.class)
+						.getSingleResult();
+				assertThat(count).isEqualTo(1L);
+			});
+		}
+
+		@Test
+		@Order(2)
+		@SessionFactory(dropTestData = DropDataTiming.AFTER_EACH)
+		public void methodLevelAfterEachCleanup(SessionFactoryScope scope) {
+			scope.inTransaction(session -> {
+				Long count = session.createQuery("select count(e) from AnEntity e", Long.class)
+						.getSingleResult();
+				assertThat(count).isEqualTo(0L);
+			});
+		}
+
+		@Test
+		@Order(3)
+		@SessionFactory(dropTestData = DropDataTiming.BEFORE_EACH)
+		public void methodLevelBeforeEachInsert(SessionFactoryScope scope) {
+			scope.inTransaction(session -> {
+				AnEntity entity = new AnEntity(11, "Method Before Each");
+				session.persist(entity);
+			});
+		}
+
+		@Test
+		@Order(4)
+		@SessionFactory(dropTestData = DropDataTiming.BEFORE_EACH)
+		public void methodLevelBeforeEachVerify(SessionFactoryScope scope) {
+			scope.inTransaction(session -> {
+				Long count = session.createQuery("select count(e) from AnEntity e", Long.class)
+						.getSingleResult();
+				assertThat(count).isEqualTo(0L);
+			});
+		}
+
+		@Test
+		@Order(5)
+		@SessionFactory(dropTestData = DropDataTiming.NEVER)
+		public void methodLevelNeverInsert(SessionFactoryScope scope) {
+			scope.inTransaction(session -> {
+				AnEntity entity = new AnEntity(12, "Method Never");
+				session.persist(entity);
+			});
+		}
+
+		@Test
+		@Order(6)
+		@SessionFactory(dropTestData = DropDataTiming.NEVER)
+		public void methodLevelNeverVerify(SessionFactoryScope scope) {
+			scope.inTransaction(session -> {
+				Long count = session.createQuery("select count(e) from AnEntity e", Long.class)
+						.getSingleResult();
+				assertThat(count).isEqualTo(0L);
+			});
+
+			scope.dropData();
+		}
+	}
+}