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

Author: inpink

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,9 +46,11 @@
  * @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();
@@ -83,6 +88,8 @@ public void postProcessTestInstance(Object testInstance, ExtensionContext contex
 
 	@Override
 	public void beforeEach(ExtensionContext context) {
+		handleDropData(context, DropDataTiming.BEFORE_EACH);
+
 		final Optional<SessionFactory> sfAnnRef = AnnotationSupport.findAnnotation(
 				context.getRequiredTestMethod(),
 				SessionFactory.class
@@ -231,6 +238,42 @@ 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 SessionFactoryScope scope = findSessionFactoryScope(testInstance, context);
+
+			final Optional<SessionFactory> sfAnnRef = AnnotationSupport.findAnnotation(
+					context.getRequiredTestClass(),
+					SessionFactory.class
+			);
+
+			if (sfAnnRef.isPresent()) {
+				DropDataTiming configuredTiming = sfAnnRef.get().dropTestData();
+				if (configuredTiming == timing) {
+					scope.dropData();
+				}
+			}
+		} 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;