From b77d90d71be4bfe755e4be70188439bf463510f1 Mon Sep 17 00:00:00 2001
From: Yevhen Sidelnyk <zsidelnik@gmail.com>
Date: Mon, 14 Oct 2024 20:00:56 +0300
Subject: [PATCH] fix: handled exception envelope stamps

---
 src/ExceptionHandlerMiddleware.php            | 18 ++++++++++++++----
 tests/Unit/ExceptionHandlerMiddlewareTest.php | 15 +++++++++++++++
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/ExceptionHandlerMiddleware.php b/src/ExceptionHandlerMiddleware.php
index eb1c8b6..7cce472 100644
--- a/src/ExceptionHandlerMiddleware.php
+++ b/src/ExceptionHandlerMiddleware.php
@@ -11,6 +11,7 @@
 use Symfony\Component\Messenger\Middleware\MiddlewareInterface;
 use Symfony\Component\Messenger\Middleware\StackInterface;
 use Symfony\Component\Messenger\Stamp\BusNameStamp;
+use Symfony\Component\Messenger\Stamp\HandledStamp;
 use Symfony\Component\Messenger\Stamp\HandlerArgumentsStamp;
 use Symfony\Component\String\ByteString;
 use Throwable;
@@ -33,11 +34,12 @@ public function handle(Envelope $envelope, StackInterface $stack): Envelope
             return $stack->next()->handle($envelope, $stack);
         } catch (Exception $exception) {
             $exceptionBusName = $this->getExceptionBusName($busName);
+            $handledExceptionEnvelope = $this->handleException($exceptionBusName, $exception, $message);
 
-            return $this->exceptionHandlerBus->dispatch(
-                Envelope::wrap($exception, [new BusNameStamp($exceptionBusName)]),
-                [new HandlerArgumentsStamp([$message])],
-            );
+            /** @var list<HandledStamp> $handledExceptionStamps */
+            $handledExceptionStamps = $handledExceptionEnvelope->all(HandledStamp::class);
+
+            return $envelope->with(...$handledExceptionStamps);
         }
     }
 
@@ -56,4 +58,12 @@ private function getExceptionBusName(ByteString $busName): string
     {
         return $busName->trimSuffix('.bus')->append('.exception.bus')->toString();
     }
+
+    private function handleException(string $exceptionBusName, Exception $exception, object $message): Envelope
+    {
+        return $this->exceptionHandlerBus->dispatch(
+            Envelope::wrap($exception, [new BusNameStamp($exceptionBusName)]),
+            [new HandlerArgumentsStamp([$message])],
+        );
+    }
 }
diff --git a/tests/Unit/ExceptionHandlerMiddlewareTest.php b/tests/Unit/ExceptionHandlerMiddlewareTest.php
index 4189d91..119e4db 100644
--- a/tests/Unit/ExceptionHandlerMiddlewareTest.php
+++ b/tests/Unit/ExceptionHandlerMiddlewareTest.php
@@ -8,11 +8,13 @@
 use PhPhD\ExceptionHandler\ExceptionHandlerMiddleware;
 use PHPUnit\Framework\MockObject\MockObject;
 use PHPUnit\Framework\TestCase;
+use RuntimeException;
 use stdClass;
 use Symfony\Component\Messenger\Envelope;
 use Symfony\Component\Messenger\MessageBus;
 use Symfony\Component\Messenger\Middleware\MiddlewareInterface;
 use Symfony\Component\Messenger\Middleware\StackMiddleware;
+use Symfony\Component\Messenger\Stamp\BusNameStamp;
 
 /**
  * @covers \PhPhD\ExceptionHandler\ExceptionHandlerMiddleware
@@ -45,4 +47,17 @@ public function testRequiresBusNameStamp(): void
 
         $this->middleware->handle($envelope, $this->stack);
     }
+
+    public function testOriginalEnvelopeStampsArePreserved(): void
+    {
+        $this->nextMiddleware->method('handle')->willThrowException(new RuntimeException());
+
+        $envelope = Envelope::wrap(new stdClass(), [new BusNameStamp('command.bus')]);
+
+        $resultEnvelope = $this->middleware->handle($envelope, $this->stack);
+
+        $busNameStamp = $resultEnvelope->last(BusNameStamp::class);
+        self::assertInstanceOf(BusNameStamp::class, $busNameStamp);
+        self::assertSame('command.bus', $busNameStamp->getBusName());
+    }
 }