Merge pull request corda#6852 from corda/dan/ENT-6092-port-to-os

ENT-6092 Add flush to `MockServices.withEntityManager` (corda#4180)
steven-tomlinson · Jan 29, 2021 · ade302b · ade302b
2 parents d214f5e + 682de39
commit ade302b
Showing 1 changed file with 6 additions and 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
@@ -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 () }
+                    }
                 }
             }
         }