From 866a424963aef890e984a1321465ae87fd2b51cb Mon Sep 17 00:00:00 2001
From: xhuberty <x.huberty@thecodingmachine.com>
Date: Tue, 12 Jan 2016 15:01:34 +0100
Subject: [PATCH] Fix issue when using notify tracking policy with multiple
 flush on entity

---
 lib/Doctrine/ORM/UnitOfWork.php             | 42 ++++++++++++++++++---
 tests/Doctrine/Tests/ORM/UnitOfWorkTest.php | 31 ++++++++++++++-
 2 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/lib/Doctrine/ORM/UnitOfWork.php b/lib/Doctrine/ORM/UnitOfWork.php
index c0b0dd8c3..367d9f4d2 100644
--- a/lib/Doctrine/ORM/UnitOfWork.php
+++ b/lib/Doctrine/ORM/UnitOfWork.php
@@ -324,7 +324,7 @@ class UnitOfWork implements PropertyChangedListener
         }
 
         // Compute changes done since last commit.
-        if ($entity === null) {
+        if (null === $entity) {
             $this->computeChangeSets();
         } elseif (is_object($entity)) {
             $this->computeSingleEntityChangeSet($entity);
@@ -414,17 +414,40 @@ class UnitOfWork implements PropertyChangedListener
 
         $this->dispatchPostFlushEvent();
 
+        $this->postCommitClear($entity);
+    }
+
+    /**
+     * @param null|object|array $entity
+     */
+    private function postCommitClear($entity = null)
+    {
+
         // Clear up
         $this->entityInsertions =
         $this->entityUpdates =
         $this->entityDeletions =
         $this->extraUpdates =
-        $this->entityChangeSets =
         $this->collectionUpdates =
         $this->collectionDeletions =
         $this->visitedCollections =
-        $this->scheduledForSynchronization =
-        $this->orphanRemovals = [];
+        $this->orphanRemovals = array();
+
+        if (null === $entity) {
+            $this->entityChangeSets = $this->scheduledForSynchronization = [];
+            return;
+        }
+
+        if (is_object($entity)) {
+            $entity = [$entity];
+        }
+
+        foreach ($entity as $object) {
+            $oid = spl_object_hash($object);
+            $class = $this->em->getClassMetadata(get_class($object));
+            $this->clearEntityChangeSet($oid);
+            $this->clearScheduledForSynchronization($class, $oid);
+        }
     }
 
     /**
@@ -3099,11 +3122,20 @@ class UnitOfWork implements PropertyChangedListener
      *
      * @return void
      */
-    public function clearEntityChangeSet($oid)
+    private function clearEntityChangeSet($oid)
     {
         $this->entityChangeSets[$oid] = [];
     }
 
+    /**
+     * @param $class
+     * @param string $oid
+     */
+    private function clearScheduledForSynchronization($class, $oid)
+    {
+        unset($this->scheduledForSynchronization[$class->rootEntityName][$oid]);
+    }
+
     /* PropertyChangedListener implementation */
 
     /**
diff --git a/tests/Doctrine/Tests/ORM/UnitOfWorkTest.php b/tests/Doctrine/Tests/ORM/UnitOfWorkTest.php
index ff025df4f..3af8fc5cd 100644
--- a/tests/Doctrine/Tests/ORM/UnitOfWorkTest.php
+++ b/tests/Doctrine/Tests/ORM/UnitOfWorkTest.php
@@ -163,7 +163,8 @@ class UnitOfWorkTest extends OrmTestCase
         $this->_unitOfWork->persist($entity);
 
         $this->_unitOfWork->commit();
-        $this->assertEquals(1, count($persister->getInserts()));
+        $this->assertCount(1, $persister->getInserts());
+
         $persister->reset();
 
         $this->assertTrue($this->_unitOfWork->isInIdentityMap($entity));
@@ -360,6 +361,34 @@ class UnitOfWorkTest extends OrmTestCase
         $this->assertFalse($this->_unitOfWork->isScheduledForInsert($entity2));
     }
 
+    /**
+     * @group 5579
+     */
+    public function testEntityChangeSetNotClearAfterFlushOnEntityOrArrayOfEntity()
+    {
+        // Create and Set first entity
+        $entity1 = new NotifyChangedEntity;
+        $entity1->setData('thedata');
+        $this->_unitOfWork->persist($entity1);
+
+        // Create and Set second entity
+        $entity2 = new NotifyChangedEntity;
+        $entity2->setData('thedata');
+        $this->_unitOfWork->persist($entity2);
+
+        $this->_unitOfWork->commit($entity1);
+        $this->assertCount(1, $this->_unitOfWork->getEntityChangeSet($entity2));
+
+        // Create and Set third entity
+        $entity3 = new NotifyChangedEntity;
+        $entity3->setData('thedata');
+        $this->_unitOfWork->persist($entity3);
+
+        $this->_unitOfWork->commit([$entity1,$entity2]);
+        $this->assertCount(1, $this->_unitOfWork->getEntityChangeSet($entity3));
+
+    }
+
     /**
      * Data Provider
      *