From 368e4ba5dd3aa02cc0492d72c77fcc9b45c3b065 Mon Sep 17 00:00:00 2001
From: Tomas Votruba <tomas.vot@gmail.com>
Date: Tue, 30 Jun 2026 13:33:40 +0200
Subject: [PATCH] [CodeQuality] Skip DecorateWillReturnMapWithExpectsMockRector
 on non-test methods

---
 .../Fixture/skip_private_helper.php.inc       | 16 +++++
 .../skip_protected_non_test_method.php.inc    | 14 +++++
 ...rateWillReturnMapWithExpectsMockRector.php | 62 ++++++++++++++-----
 3 files changed, 76 insertions(+), 16 deletions(-)
 create mode 100644 rules-tests/CodeQuality/Rector/Expression/DecorateWillReturnMapWithExpectsMockRector/Fixture/skip_private_helper.php.inc
 create mode 100644 rules-tests/CodeQuality/Rector/Expression/DecorateWillReturnMapWithExpectsMockRector/Fixture/skip_protected_non_test_method.php.inc

diff --git a/rules-tests/CodeQuality/Rector/Expression/DecorateWillReturnMapWithExpectsMockRector/Fixture/skip_private_helper.php.inc b/rules-tests/CodeQuality/Rector/Expression/DecorateWillReturnMapWithExpectsMockRector/Fixture/skip_private_helper.php.inc
new file mode 100644
index 000000000..263305aa7
--- /dev/null
+++ b/rules-tests/CodeQuality/Rector/Expression/DecorateWillReturnMapWithExpectsMockRector/Fixture/skip_private_helper.php.inc
@@ -0,0 +1,16 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Expression\DecorateWillReturnMapWithExpectsMockRector\Fixture;
+
+final class SkipPrivateHelper extends \PHPUnit\Framework\TestCase
+{
+    private function createSomeMock()
+    {
+        $someMock = $this->createMock(\stdClass::class);
+        $someMock->method('some')->willReturnMap([1, 2]);
+
+        return $someMock;
+    }
+}
diff --git a/rules-tests/CodeQuality/Rector/Expression/DecorateWillReturnMapWithExpectsMockRector/Fixture/skip_protected_non_test_method.php.inc b/rules-tests/CodeQuality/Rector/Expression/DecorateWillReturnMapWithExpectsMockRector/Fixture/skip_protected_non_test_method.php.inc
new file mode 100644
index 000000000..d4d500c51
--- /dev/null
+++ b/rules-tests/CodeQuality/Rector/Expression/DecorateWillReturnMapWithExpectsMockRector/Fixture/skip_protected_non_test_method.php.inc
@@ -0,0 +1,14 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Expression\DecorateWillReturnMapWithExpectsMockRector\Fixture;
+
+final class SkipProtectedNonTestMethod extends \PHPUnit\Framework\TestCase
+{
+    protected function prepareSomeMock()
+    {
+        $someMock = $this->createMock(\stdClass::class);
+        $someMock->method('some')->willReturnMap([1, 2]);
+    }
+}
diff --git a/rules/CodeQuality/Rector/Expression/DecorateWillReturnMapWithExpectsMockRector.php b/rules/CodeQuality/Rector/Expression/DecorateWillReturnMapWithExpectsMockRector.php
index bb0840619..d804ad361 100644
--- a/rules/CodeQuality/Rector/Expression/DecorateWillReturnMapWithExpectsMockRector.php
+++ b/rules/CodeQuality/Rector/Expression/DecorateWillReturnMapWithExpectsMockRector.php
@@ -11,10 +11,12 @@
 use PhpParser\Node\Expr\Variable;
 use PhpParser\Node\Identifier;
 use PhpParser\Node\Scalar\Int_;
+use PhpParser\Node\Stmt\ClassMethod;
 use PhpParser\Node\Stmt\Expression;
 use PHPStan\Type\ObjectType;
-use Rector\PHPStan\ScopeFetcher;
+use Rector\PhpParser\Node\BetterNodeFinder;
 use Rector\PHPUnit\Enum\PHPUnitClassName;
+use Rector\PHPUnit\NodeAnalyzer\TestsNodeAnalyzer;
 use Rector\Rector\AbstractRector;
 use Rector\ValueObject\MethodName;
 use Symplify\RuleDocGenerator\ValueObject\CodeSample\CodeSample;
@@ -25,6 +27,12 @@
  */
 final class DecorateWillReturnMapWithExpectsMockRector extends AbstractRector
 {
+    public function __construct(
+        private readonly TestsNodeAnalyzer $testsNodeAnalyzer,
+        private readonly BetterNodeFinder $betterNodeFinder,
+    ) {
+    }
+
     public function getRuleDefinition(): RuleDefinition
     {
         return new RuleDefinition('Decorate willReturnMap() calls with expects on the mock object if missing', [
@@ -78,54 +86,76 @@ protected function setUp(): void
 
     public function getNodeTypes(): array
     {
-        return [Expression::class];
+        return [ClassMethod::class];
     }
 
     /**
-     * @param Expression $node
-     * @return Expression|null
+     * @param ClassMethod $node
+     * @return ClassMethod|null
      */
     public function refactor(Node $node)
     {
-        if (! $node->expr instanceof MethodCall) {
+        // allowed as can be flexible
+        if ($this->isName($node, MethodName::SET_UP)) {
             return null;
         }
 
-        $methodCall = $node->expr;
-        if (! $this->isName($methodCall->name, 'willReturnMap')) {
+        // only decorate inside test methods, skip helpers and other non-test methods
+        if (! $this->testsNodeAnalyzer->isTestClassMethod($node)) {
             return null;
         }
 
-        $scope = ScopeFetcher::fetch($node);
+        $hasChanged = false;
 
-        // allowed as can be flexible
-        if ($scope->getFunctionName() === MethodName::SET_UP) {
+        /** @var Expression[] $expressions */
+        $expressions = $this->betterNodeFinder->findInstancesOf($node, [Expression::class]);
+        foreach ($expressions as $expression) {
+            if ($this->refactorExpression($expression)) {
+                $hasChanged = true;
+            }
+        }
+
+        if (! $hasChanged) {
             return null;
         }
 
+        return $node;
+    }
+
+    private function refactorExpression(Expression $expression): bool
+    {
+        if (! $expression->expr instanceof MethodCall) {
+            return false;
+        }
+
+        $methodCall = $expression->expr;
+        if (! $this->isName($methodCall->name, 'willReturnMap')) {
+            return false;
+        }
+
         $topmostCall = $this->resolveTopmostCall($methodCall);
 
         // already covered
         if ($this->isName($topmostCall->name, 'expects')) {
-            return null;
+            return false;
         }
 
         if (! $this->isObjectType($topmostCall->var, new ObjectType(PHPUnitClassName::MOCK_OBJECT))) {
-            return null;
+            return false;
         }
 
         if ($methodCall->isFirstClassCallable()) {
-            return null;
+            return false;
         }
 
         // count values in will map arg
         $willReturnMapArg = $methodCall->getArgs()[0] ?? null;
         if (! $willReturnMapArg instanceof Arg) {
-            return null;
+            return false;
         }
 
         if (! $willReturnMapArg->value instanceof Array_) {
-            return null;
+            return false;
         }
 
         $array = $willReturnMapArg->value;
@@ -141,7 +171,7 @@ public function refactor(Node $node)
             ))]
         );
 
-        return $node;
+        return true;
     }
 
     private function resolveTopmostCall(MethodCall $methodCall): MethodCall