From c10af529eb7db36dbfe983ba9ee2296971c9a67f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Tamarelle?= <jerome@tamarelle.net>
Date: Fri, 30 Aug 2024 08:19:32 +0200
Subject: [PATCH] Make manager registry resettable

---
 psalm-baseline.xml            | 13 +++---
 src/ManagerRegistry.php       | 39 +++++++++++++++++-
 tests/ManagerRegistryTest.php | 74 +++++++++++++++++++++++++++++++++++
 3 files changed, 118 insertions(+), 8 deletions(-)
 create mode 100644 tests/ManagerRegistryTest.php

diff --git a/psalm-baseline.xml b/psalm-baseline.xml
index 8532853d..349db7da 100644
--- a/psalm-baseline.xml
+++ b/psalm-baseline.xml
@@ -1,5 +1,5 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<files psalm-version="5.24.0@462c80e31c34e58cc4f750c656be3927e80e550e">
+<files psalm-version="5.25.0@01a8eb06b9e9cc6cfb6a320bf9fb14331919d505">
   <file src="config/value_resolver.php">
     <UndefinedClass>
       <code><![CDATA[ExpressionLanguage]]></code>
@@ -21,15 +21,14 @@
       <code><![CDATA[children]]></code>
     </UndefinedMethod>
   </file>
-  <file src="tests/DocumentValueResolverFunctionalTest.php">
-    <InvalidArgument>
-      <code><![CDATA[[DocumentValueResolverController::class, 'showUserByDefault']]]></code>
-      <code><![CDATA[[DocumentValueResolverController::class, 'showUserWithMapping']]]></code>
-    </InvalidArgument>
-  </file>
   <file src="tests/Form/Type/GuesserTestType.php">
     <MissingTemplateParam>
       <code><![CDATA[GuesserTestType]]></code>
     </MissingTemplateParam>
   </file>
+  <file src="tests/ManagerRegistryTest.php">
+    <RedundantCondition>
+      <code><![CDATA[assertSame]]></code>
+    </RedundantCondition>
+  </file>
 </files>
diff --git a/src/ManagerRegistry.php b/src/ManagerRegistry.php
index 41d5a420..2ec26834 100644
--- a/src/ManagerRegistry.php
+++ b/src/ManagerRegistry.php
@@ -6,12 +6,16 @@
 
 use Doctrine\ODM\MongoDB\DocumentManager;
 use Doctrine\ODM\MongoDB\MongoDBException;
+use ProxyManager\Proxy\LazyLoadingInterface;
 use Psr\Container\ContainerInterface;
 use Symfony\Bridge\Doctrine\ManagerRegistry as BaseManagerRegistry;
+use Symfony\Component\VarExporter\LazyObjectInterface;
+use Symfony\Contracts\Service\ResetInterface;
 
 use function array_keys;
+use function assert;
 
-class ManagerRegistry extends BaseManagerRegistry
+class ManagerRegistry extends BaseManagerRegistry implements ResetInterface
 {
     public function __construct(string $name, array $connections, array $managers, string $defaultConnection, string $defaultManager, string $proxyInterfaceName, ?ContainerInterface $container = null)
     {
@@ -42,4 +46,37 @@ public function getAliasNamespace(string $alias): string
 
         throw MongoDBException::unknownDocumentNamespace($alias);
     }
+
+    /**
+     * Clears all document managers.
+     */
+    public function reset(): void
+    {
+        foreach ($this->getManagerNames() as $managerName => $serviceId) {
+            $this->resetOrClearManager($managerName, $serviceId);
+        }
+    }
+
+    private function resetOrClearManager(string $managerName, string $serviceId): void
+    {
+        if (! $this->container->initialized($serviceId)) {
+            return;
+        }
+
+        $manager = $this->container->get($serviceId);
+
+        if ($manager instanceof LazyLoadingInterface || $manager instanceof LazyObjectInterface) {
+            $this->resetManager($managerName);
+
+            return;
+        }
+
+        assert($manager instanceof DocumentManager);
+
+        if (! $manager->isOpen()) {
+            return;
+        }
+
+        $manager->clear();
+    }
 }
diff --git a/tests/ManagerRegistryTest.php b/tests/ManagerRegistryTest.php
new file mode 100644
index 00000000..783ef259
--- /dev/null
+++ b/tests/ManagerRegistryTest.php
@@ -0,0 +1,74 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Doctrine\Bundle\MongoDBBundle\Tests;
+
+use Doctrine\Bundle\MongoDBBundle\ManagerRegistry;
+use Doctrine\ODM\MongoDB\DocumentManager;
+use Doctrine\Persistence\ObjectManager;
+use Symfony\Component\DependencyInjection\Container;
+use Symfony\Component\DependencyInjection\ContainerBuilder;
+use Symfony\Component\DependencyInjection\Dumper\PhpDumper;
+
+class ManagerRegistryTest extends TestCase
+{
+    public function testReset(): void
+    {
+        $container = new ContainerBuilder();
+        $container->register('manager.default', DocumentManagerStub::class)
+            ->setPublic(true);
+        $container->register('manager.lazy', DocumentManagerStub::class)
+            ->setPublic(true)
+            ->setLazy(true)
+            ->addTag('proxy', ['interface' => ObjectManager::class]);
+        $container->compile();
+
+        /** @var class-string<Container> $containerClass */
+        $containerClass = 'MongoDBManagerRepositoryTestResetContainer';
+        $dumper         = new PhpDumper($container);
+        eval('?' . '>' . $dumper->dump(['class' => $containerClass]));
+
+        $container  = new $containerClass();
+        $repository = new ManagerRegistry('MongoDB', [], [
+            'default' => 'manager.default',
+            'lazy' => 'manager.lazy',
+        ], '', '', '', $container);
+
+        DocumentManagerStub::$clearCount = 0;
+
+        $repository->reset();
+
+        // Service was not initialized, so reset should not be called
+        $this->assertSame(0, DocumentManagerStub::$clearCount);
+
+        // The lazy service is reinitialized instead of being cleared
+        $container->get('manager.lazy')->flush();
+        $repository->reset();
+        $this->assertSame(0, DocumentManagerStub::$clearCount);
+
+        // The default service is cleared when initialized
+        $container->get('manager.default')->flush();
+        $repository->reset();
+        $this->assertSame(1, DocumentManagerStub::$clearCount);
+    }
+}
+
+class DocumentManagerStub extends DocumentManager
+{
+    public static int $clearCount;
+
+    public function __construct()
+    {
+    }
+
+    /** {@inheritDoc} */
+    public function clear($objectName = null): void
+    {
+        self::$clearCount++;
+    }
+
+    public function flush(array $options = []): void
+    {
+    }
+}