From 682de39d4c4881799be7f3b2188e38d67d2a0d9b Mon Sep 17 00:00:00 2001
From: Dan Newton <danknewton@hotmail.com>
Date: Fri, 29 Jan 2021 15:29:33 +0000
Subject: [PATCH] ENT-6092 Add flush to `MockServices.withEntityManager`
 (#4180)

---
 .../main/kotlin/net/corda/testing/node/MockServices.kt    | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/testing/node-driver/src/main/kotlin/net/corda/testing/node/MockServices.kt b/testing/node-driver/src/main/kotlin/net/corda/testing/node/MockServices.kt
index 899d36cdc3..efd5813736 100644
--- a/testing/node-driver/src/main/kotlin/net/corda/testing/node/MockServices.kt
+++ b/testing/node-driver/src/main/kotlin/net/corda/testing/node/MockServices.kt
@@ -251,11 +251,15 @@ open class MockServices private constructor(
                 override fun jdbcSession(): Connection = persistence.createSession()
 
                 override fun <T : Any?> withEntityManager(block: EntityManager.() -> T): T {
-                    return block(contextTransaction.entityManager)
+                    return contextTransaction.entityManager.run {
+                        block(this).also { flush () }
+                    }
                 }
 
                 override fun withEntityManager(block: Consumer<EntityManager>) {
-                    return block.accept(contextTransaction.entityManager)
+                    return contextTransaction.entityManager.run {
+                        block.accept(this).also { flush () }
+                    }
                 }
             }
         }