Added support for multiple entity managers

maxbrokman · Apr 5, 2019 · d10bdcf · d10bdcf
1 parent 21744ad
commit d10bdcf
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 16 deletions.
diff --git 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
@@ -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,19 +75,21 @@ 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;
     }
 
     /**
      * To clear the em before doing any work.
      */
     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
@@ -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);
     }
 }