From a7124d3f42646fa284e039937d24033a2c223265 Mon Sep 17 00:00:00 2001
From: Thomas Prelot <tprelot@gmail.com>
Date: Wed, 7 Nov 2018 15:51:48 +0100
Subject: [PATCH] Add context services to locator

---
 .../DependencyInjection/TransportFactory.php  | 26 +++++++++++++------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/pkg/enqueue/Symfony/DependencyInjection/TransportFactory.php b/pkg/enqueue/Symfony/DependencyInjection/TransportFactory.php
index 19712ff8b..da360d711 100644
--- a/pkg/enqueue/Symfony/DependencyInjection/TransportFactory.php
+++ b/pkg/enqueue/Symfony/DependencyInjection/TransportFactory.php
@@ -156,6 +156,8 @@ public function buildContext(ContainerBuilder $container, array $config): void
             ->setFactory([new Reference($factoryId), 'createContext'])
         ;
 
+        $this->addServiceToLocator($container, 'context');
+
         if ('default' === $this->name) {
             $container->setAlias(Context::class, $this->format('context'));
         }
@@ -187,14 +189,8 @@ public function buildQueueConsumer(ContainerBuilder $container, array $config):
 
         $container->register($this->format('processor_registry'), ContainerProcessorRegistry::class);
 
-        $locatorId = 'enqueue.locator';
-        if ($container->hasDefinition($locatorId)) {
-            $locator = $container->getDefinition($locatorId);
-            $locator->replaceArgument(0, array_replace($locator->getArgument(0), [
-                $this->format('queue_consumer') => new Reference($this->format('queue_consumer')),
-                $this->format('processor_registry') => new Reference($this->format('processor_registry')),
-            ]));
-        }
+        $this->addServiceToLocator($container, 'queue_consumer');
+        $this->addServiceToLocator($container, 'processor_registry');
 
         if ('default' === $this->name) {
             $container->setAlias(QueueConsumerInterface::class, $this->format('queue_consumer'));
@@ -226,4 +222,18 @@ private function assertServiceExists(ContainerBuilder $container, string $servic
             throw new \InvalidArgumentException(sprintf('The service "%s" does not exist.', $serviceId));
         }
     }
+
+    private function addServiceToLocator(ContainerBuilder $container, string $serviceName): void
+    {
+        $locatorId = 'enqueue.locator';
+
+        if ($container->hasDefinition($locatorId)) {
+            $locator = $container->getDefinition($locatorId);
+
+            $map = $locator->getArgument(0);
+            $map[$this->format($serviceName)] = $this->reference($serviceName);
+
+            $locator->replaceArgument(0, $map);
+        }
+    }
 }