From 32f6e6bee8064287e5299e7328252a002bd19298 Mon Sep 17 00:00:00 2001
From: Fabien Huitelec <fabien@huitelec.fr>
Date: Sat, 17 Sep 2016 10:49:43 +0200
Subject: [PATCH] Add closeEntityManagerOnRollback property on
 DoctrineEntityManager

---
 src/Doctrine/DoctrineEntityManager.php        | 12 ++++++++++-
 .../Doctrine/DoctrineEntityManagerSpec.php    | 20 +++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/src/Doctrine/DoctrineEntityManager.php b/src/Doctrine/DoctrineEntityManager.php
index 9d82d21..f55b4ce 100644
--- a/src/Doctrine/DoctrineEntityManager.php
+++ b/src/Doctrine/DoctrineEntityManager.php
@@ -14,15 +14,21 @@ final class DoctrineEntityManager implements Transactional
      * @var EntityManagerInterface
      */
     private $entityManager;
+    /**
+     * @var bool
+     */
+    private $closeEntityManagerOnRollback;
 
     /**
      * Constructor.
      *
      * @param EntityManagerInterface $entityManager
+     * @param bool $closeEntityManagerOnRollback
      */
-    public function __construct(EntityManagerInterface $entityManager)
+    public function __construct(EntityManagerInterface $entityManager, $closeEntityManagerOnRollback = false)
     {
         $this->entityManager = $entityManager;
+        $this->closeEntityManagerOnRollback = $closeEntityManagerOnRollback;
     }
 
     /**
@@ -64,5 +70,9 @@ public function rollback()
         } catch (\Exception $e) {
             throw new RollbackException('Cannot rollback Doctrine ORM transaction', $e->getCode(), $e);
         }
+
+        if ($this->closeEntityManagerOnRollback) {
+            $this->entityManager->close();
+        }
     }
 }
diff --git a/tests/spec/Doctrine/DoctrineEntityManagerSpec.php b/tests/spec/Doctrine/DoctrineEntityManagerSpec.php
index 1065472..f1975b3 100644
--- a/tests/spec/Doctrine/DoctrineEntityManagerSpec.php
+++ b/tests/spec/Doctrine/DoctrineEntityManagerSpec.php
@@ -75,4 +75,24 @@ function it_should_throw_an_exception_if_em_transaction_rollback_failed(EntityMa
         $this->shouldThrow('\RemiSan\TransactionManager\Exception\RollbackException')
              ->duringRollback();
     }
+
+    function it_should_close_em_if_em_transaction_rollback_is_on_closeEntityManagerOnRollback_mode(EntityManagerInterface $entityManager)
+    {
+        $this->beConstructedWith($entityManager, true);
+
+        $entityManager->rollback()->shouldBeCalledTimes(1);
+        $entityManager->close()->shouldBeCalledTimes(1);
+
+        $this->rollback();
+    }
+
+    function it_should_not_close_em_if_em_transaction_rollback_is_not_on_closeEntityManagerOnRollback_mode(EntityManagerInterface $entityManager)
+    {
+        $this->beConstructedWith($entityManager, false);
+
+        $entityManager->rollback()->shouldBeCalledTimes(1);
+        $entityManager->close()->shouldBeCalledTimes(0);
+
+        $this->rollback();
+    }
 }