Added test for #877 BatchingConnection use-case (#925)

Co-authored-by: Gavin King <[email protected]>
hibernate · Aug 3, 2021 · 3e4ecc0 · 3e4ecc0
1 parent a11ea01
commit 3e4ecc0
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 0 deletions.
diff --git a/...ate-reactive-core/src/main/java/org/hibernate/reactive/mutiny/impl/MutinySessionImpl.java b/...ate-reactive-core/src/main/java/org/hibernate/reactive/mutiny/impl/MutinySessionImpl.java
@@ -18,6 +18,7 @@
 import org.hibernate.reactive.logging.impl.Log;
 import org.hibernate.reactive.logging.impl.LoggerFactory;
 import org.hibernate.reactive.mutiny.Mutiny;
+import org.hibernate.reactive.pool.ReactiveConnection;
 import org.hibernate.reactive.session.Criteria;
 import org.hibernate.reactive.session.ReactiveSession;
 
@@ -84,6 +85,10 @@ public <T> T getReference(Class<T> entityClass, Object id) {
 		return delegate.getReference( entityClass, id );
 	}
 
+	public ReactiveConnection getReactiveConnection() {
+		return delegate.getReactiveConnection();
+	}
+
 	@Override
 	public <T> T getReference(T entity) {
 		return delegate.getReference( delegate.getEntityClass(entity), delegate.getEntityId(entity) );

diff --git a/...rnate-reactive-core/src/main/java/org/hibernate/reactive/stage/impl/StageSessionImpl.java b/...rnate-reactive-core/src/main/java/org/hibernate/reactive/stage/impl/StageSessionImpl.java
@@ -17,6 +17,7 @@
 import org.hibernate.reactive.engine.ReactiveActionQueue;
 import org.hibernate.reactive.logging.impl.Log;
 import org.hibernate.reactive.logging.impl.LoggerFactory;
+import org.hibernate.reactive.pool.ReactiveConnection;
 import org.hibernate.reactive.session.Criteria;
 import org.hibernate.reactive.session.ReactiveSession;
 import org.hibernate.reactive.stage.Stage;
@@ -76,6 +77,10 @@ public <T> CompletionStage<T> unproxy(T association) {
 		return stage( v -> delegate.reactiveFetch(association, true) );
 	}
 
+	public ReactiveConnection getReactiveConnection() {
+		return delegate.getReactiveConnection();
+	}
+
 	@Override
 	public <T> T getReference(Class<T> entityClass, Object id) {
 		//it's important that this method does not hit the database!

diff --git a/hibernate-reactive-core/src/test/java/org/hibernate/reactive/ReactiveSessionTest.java b/hibernate-reactive-core/src/test/java/org/hibernate/reactive/ReactiveSessionTest.java
@@ -6,11 +6,16 @@
 package org.hibernate.reactive;
 
 import io.vertx.ext.unit.TestContext;
+
 import org.hibernate.LockMode;
 import org.hibernate.cfg.AvailableSettings;
 import org.hibernate.cfg.Configuration;
 import org.hibernate.reactive.common.AffectedEntities;
+import org.hibernate.reactive.mutiny.Mutiny;
+import org.hibernate.reactive.mutiny.impl.MutinySessionImpl;
+import org.hibernate.reactive.pool.BatchingConnection;
 import org.hibernate.reactive.stage.Stage;
+import org.hibernate.reactive.stage.impl.StageSessionImpl;
 
 import org.junit.After;
 import org.junit.Test;
@@ -26,6 +31,7 @@
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.hibernate.reactive.util.impl.CompletionStages.voidFuture;
 
 public class ReactiveSessionTest extends BaseReactiveTest {
@@ -700,6 +706,28 @@ public void testBatching(TestContext context) {
 		);
 	}
 
+	@Test
+	public void testBatchingConnection() {
+		Stage.Session session = getSessionFactory().openSession();
+		try {
+			assertThat(((StageSessionImpl)session).getReactiveConnection()).isInstanceOf( BatchingConnection.class );
+		}
+		finally {
+			closeSession( session );
+		}
+	}
+
+	@Test
+	public void testBatchingConnectionMutiny() {
+		Mutiny.Session session = getMutinySessionFactory().openSession();
+		try {
+			assertThat(((MutinySessionImpl)session).getReactiveConnection() ).isInstanceOf( BatchingConnection.class );
+		}
+		finally {
+			closeSession( session );
+		}
+	}
+
 	@Test
 	public void testSessionWithNativeAffectedEntities(TestContext context) {
 		GuineaPig pig = new GuineaPig(3,"Rorshach");