Force the flush of the EntityManager before the assertion if a transaction is running in the test

Author: Lemick

src/main/java/com/lemick/integration/spring/HibernateStatementCountTestListener.java

@@ -6,9 +6,11 @@
 import com.lemick.assertions.HibernateStatementAssertionsProvider;
 import com.lemick.integration.hibernate.HibernateStatementCountInspector;
 import com.lemick.integration.hibernate.HibernateStatistics;
+import jakarta.persistence.EntityManager;
 import org.springframework.core.Ordered;
 import org.springframework.test.context.TestContext;
 import org.springframework.test.context.TestExecutionListener;
+import org.springframework.test.context.transaction.TestTransaction;
 
 import java.util.List;
 import java.util.function.Supplier;
@@ -17,6 +19,8 @@ public class HibernateStatementCountTestListener implements TestExecutionListene
 
     private HibernateStatementAssertionsProvider hibernateStatementAssertionsProvider = new HibernateStatementAssertionsProvider();
     private Supplier<HibernateStatistics> statisticsSupplier = HibernateStatementCountInspector::getStatistics;
+    private Supplier<Boolean> transactionAvailabilitySupplier = TestTransaction::isActive;
+
 
     @Override
     public void beforeTestMethod(TestContext testContext) {
@@ -27,18 +31,38 @@ public void beforeTestMethod(TestContext testContext) {
     public void afterTestMethod(TestContext testContext) {
         AssertHibernateSQLStatementCount annotation = testContext.getTestMethod().getAnnotation(AssertHibernateSQLStatementCount.class);
         if (annotation != null) {
-            HibernateStatementAssertionResults assertionResults = new HibernateStatementAssertionResults(List.of(
-                    hibernateStatementAssertionsProvider.generateSelectStatementAssertion(annotation.selects(), statisticsSupplier),
-                    hibernateStatementAssertionsProvider.generateUpdateStatementAssertion(annotation.updates(), statisticsSupplier),
-                    hibernateStatementAssertionsProvider.generateInsertStatementAssertion(annotation.inserts(), statisticsSupplier),
-                    hibernateStatementAssertionsProvider.generateDeleteStatementAssertion(annotation.deletes(), statisticsSupplier)
-            ));
-            assertionResults.validate();
+            flushExistingPersistenceContext(testContext, transactionAvailabilitySupplier);
+            doStatementCountEvaluation(annotation);
         }
     }
 
+
+    /**
+     * Low precedence for executing before {@link org.springframework.test.context.transaction.TransactionalTestExecutionListener}
+     * closes the transaction and have the ability to flush the EntityManager
+     */
     @Override
     public int getOrder() {
-        return Integer.MIN_VALUE;
+        return Ordered.LOWEST_PRECEDENCE;
+    }
+
+    private void flushExistingPersistenceContext(TestContext testContext, Supplier<Boolean> transactionAvailabilitySupplier) {
+        if (transactionAvailabilitySupplier.get()) {
+            testContext.getApplicationContext()
+                    .getAutowireCapableBeanFactory()
+                    .getBean(EntityManager.class)
+                    .flush();
+        }
+    }
+
+    private void doStatementCountEvaluation(AssertHibernateSQLStatementCount annotation) {
+        HibernateStatementAssertionResults assertionResults = new HibernateStatementAssertionResults(List.of(
+                hibernateStatementAssertionsProvider.generateSelectStatementAssertion(annotation.selects(), statisticsSupplier),
+                hibernateStatementAssertionsProvider.generateUpdateStatementAssertion(annotation.updates(), statisticsSupplier),
+                hibernateStatementAssertionsProvider.generateInsertStatementAssertion(annotation.inserts(), statisticsSupplier),
+                hibernateStatementAssertionsProvider.generateDeleteStatementAssertion(annotation.deletes(), statisticsSupplier)
+        ));
+        assertionResults.validate();
     }
+
 }

src/test/java/com/lemick/integration/spring/HibernateStatementCountTestListenerTest.java

@@ -4,8 +4,11 @@
 import com.lemick.assertions.HibernateStatementAssertionResult;
 import com.lemick.assertions.HibernateStatementAssertionsProvider;
 import com.lemick.integration.hibernate.HibernateStatistics;
+import jakarta.persistence.EntityManager;
+import jakarta.transaction.Transactional;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Answers;
 import org.mockito.InjectMocks;
 import org.mockito.Mock;
 import org.mockito.Spy;
@@ -19,7 +22,7 @@
 @ExtendWith(MockitoExtension.class)
 class HibernateStatementCountTestListenerTest {
 
-    public class FakeClass {
+    public static class FakeClass {
         @AssertHibernateSQLStatementCount(inserts = 1, deletes = 2, selects = 3, updates = 4)
         public void annotatedMethod() {
 
@@ -40,6 +43,9 @@ public void notAnnotatedMethod() {
     Supplier<HibernateStatistics> statisticsSupplier;
 
     @Mock
+    Supplier<Boolean> transactionAvailabilitySupplier;
+
+    @Mock(answer = Answers.RETURNS_DEEP_STUBS)
     TestContext testContext;
 
     @Test
@@ -61,7 +67,7 @@ public void _beforeTestMethod_with_annotation() {
     }
 
     @Test
-    public void _afterTestMethod_no_annotation() throws NoSuchMethodException {
+    public void _afterTestMethod_without_annotation() throws NoSuchMethodException {
         when(testContext.getTestMethod()).thenReturn(FakeClass.class.getMethod("notAnnotatedMethod"));
 
         model.afterTestMethod(testContext);
@@ -73,8 +79,31 @@ public void _afterTestMethod_no_annotation() throws NoSuchMethodException {
     }
 
     @Test
-    public void _afterTestMethod_with_annotation() throws NoSuchMethodException {
+    public void _afterTestMethod_with_annotation_without_transaction() throws NoSuchMethodException {
+        when(testContext.getTestMethod()).thenReturn(FakeClass.class.getMethod("annotatedMethod"));
+        when(transactionAvailabilitySupplier.get()).thenReturn(false);
+
+        HibernateStatementAssertionResult statementAssertionResult = mock(HibernateStatementAssertionResult.class);
+        when(hibernateStatementAssertUtils.generateInsertStatementAssertion(anyInt(), any())).thenReturn(statementAssertionResult);
+        when(hibernateStatementAssertUtils.generateSelectStatementAssertion(anyInt(), any())).thenReturn(statementAssertionResult);
+        when(hibernateStatementAssertUtils.generateUpdateStatementAssertion(anyInt(), any())).thenReturn(statementAssertionResult);
+        when(hibernateStatementAssertUtils.generateDeleteStatementAssertion(anyInt(), any())).thenReturn(statementAssertionResult);
+
+        model.afterTestMethod(testContext);
+
+        verify(testContext, times(0).description("EntityManager is NOT flushed")).getApplicationContext();
+        verify(hibernateStatementAssertUtils, description("assert method is called")).generateSelectStatementAssertion(3, statisticsSupplier);
+        verify(hibernateStatementAssertUtils, description("assert method is called")).generateUpdateStatementAssertion(4, statisticsSupplier);
+        verify(hibernateStatementAssertUtils, description("assert method is called")).generateInsertStatementAssertion(1, statisticsSupplier);
+        verify(hibernateStatementAssertUtils, description("assert method is called")).generateDeleteStatementAssertion(2, statisticsSupplier);
+    }
+
+    @Test
+    public void _afterTestMethod_with_annotation_with_transaction() throws NoSuchMethodException {
+        EntityManager entityManager = mock(EntityManager.class);
         when(testContext.getTestMethod()).thenReturn(FakeClass.class.getMethod("annotatedMethod"));
+        when(testContext.getApplicationContext().getAutowireCapableBeanFactory().getBean(EntityManager.class)).thenReturn(entityManager);
+        when(transactionAvailabilitySupplier.get()).thenReturn(true);
 
         HibernateStatementAssertionResult statementAssertionResult = mock(HibernateStatementAssertionResult.class);
         when(hibernateStatementAssertUtils.generateInsertStatementAssertion(anyInt(), any())).thenReturn(statementAssertionResult);
@@ -84,6 +113,7 @@ public void _afterTestMethod_with_annotation() throws NoSuchMethodException {
 
         model.afterTestMethod(testContext);
 
+        verify(entityManager, description("EntityManager is flushed")).flush();
         verify(hibernateStatementAssertUtils, description("assert method is called")).generateSelectStatementAssertion(3, statisticsSupplier);
         verify(hibernateStatementAssertUtils, description("assert method is called")).generateUpdateStatementAssertion(4, statisticsSupplier);
         verify(hibernateStatementAssertUtils, description("assert method is called")).generateInsertStatementAssertion(1, statisticsSupplier);