From 3cd80a2e028aaac27e8cdf79397dfc8448b27827 Mon Sep 17 00:00:00 2001
From: Suleiman Dibirov <idsulik@gmail.com>
Date: Wed, 11 Dec 2024 14:10:22 +0200
Subject: [PATCH] fix: Fix LazyServiceFactory to handle exception in
 constructor

Signed-off-by: Suleiman Dibirov <idsulik@gmail.com>
---
 src/Proxy/LazyServiceFactory.php      |  9 ++++-
 test/Proxy/LazyServiceFactoryTest.php | 49 +++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/src/Proxy/LazyServiceFactory.php b/src/Proxy/LazyServiceFactory.php
index 56ccfe4a..73f91331 100644
--- a/src/Proxy/LazyServiceFactory.php
+++ b/src/Proxy/LazyServiceFactory.php
@@ -10,6 +10,7 @@
 use ProxyManager\Proxy\LazyLoadingInterface;
 use ProxyManager\Proxy\VirtualProxyInterface;
 use Psr\Container\ContainerInterface;
+use Throwable;
 
 use function sprintf;
 
@@ -42,8 +43,14 @@ public function __invoke(
     ): VirtualProxyInterface {
         if (isset($this->servicesMap[$name])) {
             $initializer = static function (&$wrappedInstance, LazyLoadingInterface $proxy) use ($callback): bool {
+                $initializer = $proxy->getProxyInitializer();
                 $proxy->setProxyInitializer(null);
-                $wrappedInstance = $callback();
+                try {
+                    $wrappedInstance = $callback();
+                } catch (Throwable $e) {
+                    $proxy->setProxyInitializer($initializer);
+                    throw $e;
+                }
 
                 return true;
             };
diff --git a/test/Proxy/LazyServiceFactoryTest.php b/test/Proxy/LazyServiceFactoryTest.php
index 582f135f..e648d0cf 100644
--- a/test/Proxy/LazyServiceFactoryTest.php
+++ b/test/Proxy/LazyServiceFactoryTest.php
@@ -15,6 +15,7 @@
 use ProxyManager\Proxy\LazyLoadingInterface;
 use ProxyManager\Proxy\VirtualProxyInterface;
 use Psr\Container\ContainerInterface;
+use RuntimeException;
 
 #[CoversClass(LazyServiceFactory::class)]
 final class LazyServiceFactoryTest extends TestCase
@@ -94,4 +95,52 @@ static function ($className, $initializer) use ($expectedService, $proxy): MockO
 
         self::assertSame($expectedService, $result, 'service created not match the expected');
     }
+
+    public function testHandlesExceptionInInitializer(): void
+    {
+        $callback = function (): void {
+            throw new RuntimeException('Test exception');
+        };
+
+        $proxy       = $this->createMock(LazyLoadingInterface::class);
+        $initializer = static fn(): bool => true;
+
+        $proxy->expects(self::once())
+            ->method('getProxyInitializer')
+            ->willReturn($initializer);
+
+        $proxy
+            ->expects(self::exactly(2))
+            ->method('setProxyInitializer')
+            ->willReturnMap(
+                [
+                    [null],
+                    [$initializer],
+                ]
+            );
+
+        $expectedService = $this->createMock(VirtualProxyInterface::class);
+
+        $this->proxyFactory
+            ->expects(self::once())
+            ->method('createProxy')
+            ->willReturnCallback(
+                /** @psalm-suppress UnusedVariable */
+                static function (string $className, callable $initializer) use ($expectedService, $proxy): MockObject {
+                    self::assertEquals('FooClass', $className);
+
+                    try {
+                        $wrappedInstance = null;
+                        $initializer($wrappedInstance, $proxy);
+                        self::fail('Expected exception was not thrown');
+                    } catch (RuntimeException) {
+                    }
+
+                    return $expectedService;
+                }
+            );
+
+        $result = $this->factory->__invoke($this->container, 'fooService', $callback);
+        self::assertSame($expectedService, $result);
+    }
 }