From d10bdcfe0cd2a179fc89bd618b470bc16f421f08 Mon Sep 17 00:00:00 2001
From: Simon Svensson <sisve@devhost.se>
Date: Fri, 5 Apr 2019 09:39:14 +0200
Subject: [PATCH] Added support for multiple entity managers

---
 src/DoctrineQueueProvider.php |  2 +-
 src/Worker.php                | 33 ++++++++++++++++------------
 tests/WorkerTest.php          | 41 ++++++++++++++++++++++++++++++++++-
 3 files changed, 60 insertions(+), 16 deletions(-)

diff --git a/src/DoctrineQueueProvider.php b/src/DoctrineQueueProvider.php
index c77a596..4253507 100644
--- a/src/DoctrineQueueProvider.php
+++ b/src/DoctrineQueueProvider.php
@@ -47,7 +47,7 @@ protected function registerWorker()
 
         $this->app->singleton('safeQueue.worker', function ($app) {
             return new Worker($app['queue'], $app['events'],
-                $app['em'], $app['Illuminate\Contracts\Debug\ExceptionHandler']);
+                $app['Doctrine\Common\Persistence\ManagerRegistry'], $app['Illuminate\Contracts\Debug\ExceptionHandler']);
         });
     }
 
diff --git a/src/Worker.php b/src/Worker.php
index 29e356c..1c6da5d 100644
--- a/src/Worker.php
+++ b/src/Worker.php
@@ -3,6 +3,7 @@
 
 namespace MaxBrokman\SafeQueue;
 
+use Doctrine\Common\Persistence\ManagerRegistry;
 use Doctrine\ORM\EntityManager;
 use Exception;
 use Illuminate\Contracts\Debug\ExceptionHandler;
@@ -16,27 +17,27 @@
 /*final*/ class Worker extends IlluminateWorker
 {
     /**
-     * @var EntityManager
+     * @var ManagerRegistry
      */
-    private $entityManager;
+    protected $managerRegistry;
 
     /**
      * Worker constructor.
      *
      * @param QueueManager     $manager
      * @param Dispatcher       $events
-     * @param EntityManager    $entityManager
+     * @param ManagerRegistry    $managerRegistry
      * @param ExceptionHandler $exceptions
      */
     public function __construct(
         QueueManager $manager,
         Dispatcher $events,
-        EntityManager $entityManager,
+        ManagerRegistry $managerRegistry,
         ExceptionHandler $exceptions
     ) {
         parent::__construct($manager, $events, $exceptions);
 
-        $this->entityManager = $entityManager;
+        $this->managerRegistry = $managerRegistry;
     }
 
     /**
@@ -74,11 +75,11 @@ protected function runJob($job, $connectionName, WorkerOptions $options)
      */
     private function assertEntityManagerOpen()
     {
-        if ($this->entityManager->isOpen()) {
-            return;
+        foreach ($this->managerRegistry->getManagers() as $entityManager) {
+            if (!$entityManager->isOpen()) {
+                throw new EntityManagerClosedException;
+            }
         }
-
-        throw new EntityManagerClosedException;
     }
 
     /**
@@ -86,7 +87,9 @@ private function assertEntityManagerOpen()
      */
     private function assertEntityManagerClear()
     {
-        $this->entityManager->clear();
+        foreach ($this->managerRegistry->getManagers() as $entityManager) {
+            $entityManager->clear();
+        }
     }
 
     /**
@@ -96,11 +99,13 @@ private function assertEntityManagerClear()
      */
     private function assertGoodDatabaseConnection()
     {
-        $connection = $this->entityManager->getConnection();
+        foreach ($this->managerRegistry->getManagers() as $entityManager) {
+            $connection = $entityManager->getConnection();
 
-        if ($connection->ping() === false) {
-            $connection->close();
-            $connection->connect();
+            if ($connection->ping() === false) {
+                $connection->close();
+                $connection->connect();
+            }
         }
     }
 }
diff --git a/tests/WorkerTest.php b/tests/WorkerTest.php
index 050b29c..a421258 100644
--- a/tests/WorkerTest.php
+++ b/tests/WorkerTest.php
@@ -3,8 +3,10 @@
 
 namespace tests\MaxBrokman\SafeQueue;
 
+use Doctrine\Common\Persistence\ManagerRegistry;
 use Doctrine\DBAL\Connection;
 use Doctrine\ORM\EntityManager;
+use Doctrine\ORM\EntityManagerInterface;
 use Illuminate\Contracts\Cache\Repository;
 use Illuminate\Contracts\Debug\ExceptionHandler as Handler;
 use Illuminate\Contracts\Events\Dispatcher;
@@ -53,6 +55,11 @@ class WorkerTest extends \PHPUnit_Framework_TestCase
      */
     private $exceptions;
 
+    /**
+     * @var ManagerRegistry|m\MockInterface
+     */
+    private $managerRegistry;
+
     /**
      * @var Worker
      */
@@ -72,8 +79,9 @@ protected function setUp()
         $this->dbConnection  = m::mock(Connection::class);
         $this->cache         = m::mock(Repository::class);
         $this->exceptions    = m::mock(Handler::class);
+        $this->managerRegistry = m::mock(ManagerRegistry::class);
 
-        $this->worker = new Worker($this->queueManager, $this->dispatcher, $this->entityManager, $this->exceptions);
+        $this->worker = new Worker($this->queueManager, $this->dispatcher, $this->managerRegistry, $this->exceptions);
 
         $this->options = new WorkerOptions(0, 128, 0, 0, 0);
 
@@ -128,6 +136,37 @@ public function testChecksEmState()
         $this->dbConnection->shouldReceive('close')->once();
         $this->dbConnection->shouldReceive('connect')->once();
 
+        $this->managerRegistry->shouldReceive('getManagers')->andReturn(array($this->entityManager));
+
+        $this->worker->runNextJob('connection', 'queue', $this->options);
+    }
+
+    public function testMultipleEntityManagers() {
+        $job = m::mock(Job::class);
+        $job->shouldReceive('fire')->once();
+        $job->shouldIgnoreMissing();
+
+        $this->prepareToRunJob($job);
+
+        $this->entityManager->shouldReceive('isOpen')->once()->andReturn(true);
+        $this->entityManager->shouldReceive('clear')->once();
+
+        $this->dbConnection->shouldReceive('ping')->once()->andReturn(false);
+        $this->dbConnection->shouldReceive('close')->once();
+        $this->dbConnection->shouldReceive('connect')->once();
+
+        $secondConnection = m::mock(Connection::class);
+        $secondConnection->shouldReceive('ping')->once()->andReturn(false);
+        $secondConnection->shouldReceive('close')->once();
+        $secondConnection->shouldReceive('connect')->once();
+
+        $secondManager = m::mock(EntityManagerInterface::class);
+        $secondManager->shouldReceive('getConnection')->andReturn($secondConnection);
+        $secondManager->shouldReceive('isOpen')->once()->andReturn(true);
+        $secondManager->shouldReceive('clear')->once();
+
+        $this->managerRegistry->shouldReceive('getManagers')->andReturn(array($this->entityManager, $secondManager));
+
         $this->worker->runNextJob('connection', 'queue', $this->options);
     }
 }