[Messenger] Fix TraceableBus with envelope

ogizanagi · sroze · commit 21e49d21d881 · 2018-05-09T15:46:48.000+01:00
diff --git a/src/Symfony/Component/Messenger/Tests/TraceableMessageBusTest.php b/src/Symfony/Component/Messenger/Tests/TraceableMessageBusTest.php
@@ -12,6 +12,7 @@
 namespace Symfony\Component\Messenger\Tests;
 
 use PHPUnit\Framework\TestCase;
+use Symfony\Component\Messenger\Envelope;
 use Symfony\Component\Messenger\MessageBusInterface;
 use Symfony\Component\Messenger\Tests\Fixtures\DummyMessage;
 use Symfony\Component\Messenger\TraceableMessageBus;
@@ -30,6 +31,18 @@ public function testItTracesResult()
         $this->assertSame(array(array('message' => $message, 'result' => $result)), $traceableBus->getDispatchedMessages());
     }
 
+    public function testItTracesResultWithEnvelope()
+    {
+        $envelope = Envelope::wrap($message = new DummyMessage('Hello'));
+
+        $bus = $this->getMockBuilder(MessageBusInterface::class)->getMock();
+        $bus->expects($this->once())->method('dispatch')->with($envelope)->willReturn($result = array('foo' => 'bar'));
+
+        $traceableBus = new TraceableMessageBus($bus);
+        $this->assertSame($result, $traceableBus->dispatch($envelope));
+        $this->assertSame(array(array('message' => $message, 'result' => $result)), $traceableBus->getDispatchedMessages());
+    }
+
     public function testItTracesExceptions()
     {
         $message = new DummyMessage('Hello');
diff --git a/src/Symfony/Component/Messenger/TraceableMessageBus.php b/src/Symfony/Component/Messenger/TraceableMessageBus.php
@@ -29,18 +29,20 @@ public function __construct(MessageBusInterface $decoratedBus)
      */
     public function dispatch($message)
     {
+        $messageToTrace = $message instanceof Envelope ? $message->getMessage() : $message;
+
         try {
             $result = $this->decoratedBus->dispatch($message);
 
             $this->dispatchedMessages[] = array(
-                'message' => $message,
+                'message' => $messageToTrace,
                 'result' => $result,
             );
 
             return $result;
         } catch (\Throwable $e) {
             $this->dispatchedMessages[] = array(
-                'message' => $message,
+                'message' => $messageToTrace,
                 'exception' => $e,
             );
 
