[#2768] Update existing test

Now we test with and without transactions.
And the test will throw an exception if context and connections
aren't handled correctly.

Author: DavideD

hibernate-reactive-core/src/test/java/org/hibernate/reactive/MultithreadedInsertionWithLazyConnectionTest.java

@@ -8,18 +8,20 @@
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.function.BiFunction;
 
 import org.hibernate.SessionFactory;
 import org.hibernate.boot.registry.StandardServiceRegistry;
 import org.hibernate.boot.registry.StandardServiceRegistryBuilder;
 import org.hibernate.cfg.Configuration;
+import org.hibernate.reactive.annotations.DisabledFor;
 import org.hibernate.reactive.provider.ReactiveServiceRegistryBuilder;
 import org.hibernate.reactive.stage.Stage;
 import org.hibernate.reactive.util.impl.CompletionStages;
 import org.hibernate.reactive.vertx.VertxInstance;
 
-import org.junit.jupiter.api.AfterAll;
-import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -41,6 +43,7 @@
 import static org.assertj.core.api.Assertions.fail;
 import static org.hibernate.cfg.AvailableSettings.SHOW_SQL;
 import static org.hibernate.reactive.BaseReactiveTest.setDefaultProperties;
+import static org.hibernate.reactive.containers.DatabaseConfiguration.DBType.DB2;
 import static org.hibernate.reactive.provider.Settings.POOL_CONNECT_TIMEOUT;
 import static org.hibernate.reactive.util.impl.CompletionStages.failedFuture;
 import static org.hibernate.reactive.util.impl.CompletionStages.loop;
@@ -101,8 +104,8 @@ public class MultithreadedInsertionWithLazyConnectionTest {
 	private static Vertx vertx;
 	private static SessionFactory sessionFactory;
 
-	@BeforeAll
-	public static void setupSessionFactory() {
+	@BeforeEach
+	public void setupSessionFactory() {
 		vertx = Vertx.vertx( getVertxOptions() );
 		Configuration configuration = new Configuration();
 		setDefaultProperties( configuration );
@@ -130,8 +133,8 @@ private static VertxOptions getVertxOptions() {
 		return vertxOptions;
 	}
 
-	@AfterAll
-	public static void closeSessionFactory() {
+	@AfterEach
+	public void closeSessionFactory() {
 		stageSessionFactory.close();
 	}
 
@@ -140,8 +143,33 @@ public void testIdentityGenerator(VertxTestContext context) {
 		final DeploymentOptions deploymentOptions = new DeploymentOptions();
 		deploymentOptions.setInstances( N_THREADS );
 
+		// We are not using transactions on purpose here, because this approach will cause a context switch
+		// and an assertion error if things aren't handled correctly.
+		// See Hibernate Reactive issue #2768: https://github.com/hibernate/hibernate-reactive/issues/2768
 		vertx
-				.deployVerticle( InsertEntitiesVerticle::new, deploymentOptions )
+				.deployVerticle( () -> new InsertEntitiesVerticle( (s, entity) -> s
+						.persist( entity )
+						.thenCompose( v -> s.flush() )
+						.thenAccept( v -> s.clear() ) ), deploymentOptions
+				)
+				.onSuccess( res -> {
+					endLatch.waitForEveryone();
+					context.completeNow();
+				} )
+				.onFailure( context::failNow )
+				.eventually( () -> vertx.close() );
+	}
+
+	@Test
+	@DisabledFor(value = DB2, reason = "Exception: IllegalStateException: Needed to have 6 in buffer but only had 0")
+	public void testIdentityGeneratorWithTransaction(VertxTestContext context) {
+		final DeploymentOptions deploymentOptions = new DeploymentOptions();
+		deploymentOptions.setInstances( N_THREADS );
+		vertx
+				.deployVerticle(
+						() -> new InsertEntitiesVerticle( (s, entity) -> s
+								.withTransaction( t -> s.persist( entity ) ) ), deploymentOptions
+				)
 				.onSuccess( res -> {
 					endLatch.waitForEveryone();
 					context.completeNow();
@@ -152,9 +180,12 @@ public void testIdentityGenerator(VertxTestContext context) {
 
 	private static class InsertEntitiesVerticle extends AbstractVerticle {
 
+		final BiFunction<Stage.Session, EntityWithGeneratedId, CompletionStage<Void>> insertFun;
+
 		int sequentialOperation = 0;
 
-		public InsertEntitiesVerticle() {
+		public InsertEntitiesVerticle(BiFunction<Stage.Session, EntityWithGeneratedId, CompletionStage<Void>> insertFun) {
+			this.insertFun = insertFun;
 		}
 
 		@Override
@@ -196,9 +227,8 @@ private CompletionStage<Void> storeEntity(Stage.Session s) {
 			final int localVerticleOperationSequence = sequentialOperation++;
 			final EntityWithGeneratedId entity = new EntityWithGeneratedId();
 			entity.name = beforeOperationThread + "__" + localVerticleOperationSequence;
-
-			return s
-					.withTransaction( t -> s.persist( entity ) )
+			return insertFun
+					.apply( s, entity )
 					.thenCompose( v -> beforeOperationThread != Thread.currentThread()
 							? failedFuture( new IllegalStateException( "Detected an unexpected switch of carrier threads!" ) )
 							: voidFuture() );