From 8598fc933aa227b7245ebeea50604dd3e69ff658 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Poirier=20Th=C3=A9or=C3=AAt?= <mpoiriert@gmail.com>
Date: Thu, 12 Dec 2024 23:20:46 -0500
Subject: [PATCH] [EntityMigrator] Integrate DoctrineEnvelopeEntityReference
 custom exception

---
 .../entity-migrator/Message/MigrateEntityCommand.php  | 11 ++++++-----
 .../Tests/Message/MigrateEntityCommandTest.php        |  7 ++++---
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/packages/entity-migrator/Message/MigrateEntityCommand.php b/packages/entity-migrator/Message/MigrateEntityCommand.php
index 65d84cd5..4a664587 100644
--- a/packages/entity-migrator/Message/MigrateEntityCommand.php
+++ b/packages/entity-migrator/Message/MigrateEntityCommand.php
@@ -3,22 +3,23 @@
 namespace Draw\Component\EntityMigrator\Message;
 
 use Draw\Component\EntityMigrator\Entity\EntityMigrationInterface;
+use Draw\Component\Messenger\DoctrineEnvelopeEntityReference\Exception\ObjectNotFoundException;
 use Draw\Component\Messenger\DoctrineEnvelopeEntityReference\Message\DoctrineReferenceAwareInterface;
-use Symfony\Component\Messenger\Exception\UnrecoverableMessageHandlingException;
+use Draw\Component\Messenger\DoctrineEnvelopeEntityReference\Stamp\PropertyReferenceStamp;
 
 class MigrateEntityCommand implements DoctrineReferenceAwareInterface
 {
-    private ?EntityMigrationInterface $entity;
+    private PropertyReferenceStamp|EntityMigrationInterface|null $entity;
 
     public function __construct(EntityMigrationInterface $entity)
     {
         $this->entity = $entity;
     }
 
-    public function getEntity(): ?EntityMigrationInterface
+    public function getEntity(): EntityMigrationInterface
     {
-        if (null === $this->entity) {
-            throw new UnrecoverableMessageHandlingException('Entity not found');
+        if (!$this->entity instanceof EntityMigrationInterface) {
+            throw new ObjectNotFoundException($this->entity?->getClass() ?? EntityMigrationInterface::class, $this->entity);
         }
 
         return $this->entity;
diff --git a/packages/entity-migrator/Tests/Message/MigrateEntityCommandTest.php b/packages/entity-migrator/Tests/Message/MigrateEntityCommandTest.php
index d9e51d00..dd3a3d65 100644
--- a/packages/entity-migrator/Tests/Message/MigrateEntityCommandTest.php
+++ b/packages/entity-migrator/Tests/Message/MigrateEntityCommandTest.php
@@ -5,9 +5,10 @@
 use Draw\Component\Core\Reflection\ReflectionAccessor;
 use Draw\Component\EntityMigrator\Entity\EntityMigrationInterface;
 use Draw\Component\EntityMigrator\Message\MigrateEntityCommand;
+use Draw\Component\Messenger\DoctrineEnvelopeEntityReference\Exception\ObjectNotFoundException;
+use Draw\Component\Messenger\DoctrineEnvelopeEntityReference\Stamp\PropertyReferenceStamp;
 use PHPUnit\Framework\Attributes\CoversClass;
 use PHPUnit\Framework\TestCase;
-use Symfony\Component\Messenger\Exception\UnrecoverableMessageHandlingException;
 
 /**
  * @internal
@@ -36,11 +37,11 @@ public function testGetEntity(): void
         ReflectionAccessor::setPropertyValue(
             $this->object,
             'entity',
-            null
+            $stamp = new PropertyReferenceStamp('entity', EntityMigrationInterface::class, ['id' => 1])
         );
 
         static::expectExceptionObject(
-            new UnrecoverableMessageHandlingException('Entity not found')
+            new ObjectNotFoundException(EntityMigrationInterface::class, $stamp)
         );
 
         $this->object->getEntity();