Fix return type of Query::getLockMode() (#11462)

… for DBAL 4

src/Query.php

@@ -635,9 +635,10 @@ public function setLockMode(LockMode|int $lockMode): self
     /**
      * Get the current lock mode for this query.
      *
-     * @return int|null The current lock mode of this query or NULL if no specific lock mode is set.
+     * @return LockMode|int|null The current lock mode of this query or NULL if no specific lock mode is set.
+     * @psalm-return LockMode::*|null
      */
-    public function getLockMode(): int|null
+    public function getLockMode(): LockMode|int|null
     {
         $lockMode = $this->getHint(self::HINT_LOCK_MODE);
 

tests/Tests/ORM/Query/QueryTest.php

@@ -12,9 +12,12 @@
 use Doctrine\DBAL\Connection;
 use Doctrine\DBAL\Driver;
 use Doctrine\DBAL\Driver\Result;
+use Doctrine\DBAL\LockMode;
 use Doctrine\DBAL\ParameterType;
 use Doctrine\DBAL\Platforms\AbstractPlatform;
 use Doctrine\DBAL\Types\Types;
+use Doctrine\ORM\EntityManagerInterface;
+use Doctrine\ORM\Query;
 use Doctrine\ORM\Query\Parameter;
 use Doctrine\ORM\Query\QueryException;
 use Doctrine\ORM\UnitOfWork;
@@ -39,8 +42,7 @@
 
 class QueryTest extends OrmTestCase
 {
-    /** @var EntityManagerMock */
-    protected $entityManager;
+    private EntityManagerMock $entityManager;
 
     protected function setUp(): void
     {
@@ -80,6 +82,30 @@ public function testSetParameters(): void
         self::assertEquals($parameters, $query->getParameters());
     }
 
+    /** @psalm-param LockMode::* $lockMode */
+    #[DataProvider('provideLockModes')]
+    public function testSetLockMode(LockMode|int $lockMode): void
+    {
+        $query = $this->entityManager->wrapInTransaction(static function (EntityManagerInterface $em) use ($lockMode): Query {
+            $query = $em->createQuery('select u from Doctrine\Tests\Models\CMS\CmsUser u where u.username = ?1');
+            $query->setLockMode($lockMode);
+
+            return $query;
+        });
+
+        self::assertSame($lockMode, $query->getLockMode());
+        self::assertSame($lockMode, $query->getHint(Query::HINT_LOCK_MODE));
+    }
+
+    /** @psalm-return list<array{LockMode::*}> */
+    public static function provideLockModes(): array
+    {
+        return [
+            [LockMode::PESSIMISTIC_READ],
+            [LockMode::PESSIMISTIC_WRITE],
+        ];
+    }
+
     public function testFree(): void
     {
         $query = $this->entityManager->createQuery('select u from Doctrine\Tests\Models\CMS\CmsUser u where u.username = ?1');