From 1f485cdc52aa37d1acd34199f654fa46da8eeec8 Mon Sep 17 00:00:00 2001
From: Tomas Votruba <tomas.vot@gmail.com>
Date: Fri, 25 Apr 2025 14:35:38 +0200
Subject: [PATCH] [code-quality] Add TypeWillReturnCallableArrowFunctionRector

---
 config/sets/phpunit-code-quality.php          |   5 +-
 .../Fixture/assigned_final_in_setup.php.inc   |  53 ++++
 .../Fixture/fill_known_param_type.php.inc     |  37 +++
 .../Fixture/include_matches.php.inc           |  39 +++
 .../Source/SomeFinalMockedClass.php           |  13 +
 .../Source/SomeMockedClass.php                |  12 +
 ...lReturnCallableArrowFunctionRectorTest.php |  28 ++
 .../config/configured_rule.php                |  10 +
 .../SetUpAssignedMockTypesResolver.php        |  84 ++++++
 ...eWillReturnCallableArrowFunctionRector.php | 285 ++++++++++++++++++
 ...MethodParametersAndReturnTypesResolver.php |  76 +++++
 .../ValueObject/ParamTypesAndReturnType.php   |  32 ++
 12 files changed, 673 insertions(+), 1 deletion(-)
 create mode 100644 rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/assigned_final_in_setup.php.inc
 create mode 100644 rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/fill_known_param_type.php.inc
 create mode 100644 rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/include_matches.php.inc
 create mode 100644 rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Source/SomeFinalMockedClass.php
 create mode 100644 rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Source/SomeMockedClass.php
 create mode 100644 rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/TypeWillReturnCallableArrowFunctionRectorTest.php
 create mode 100644 rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/config/configured_rule.php
 create mode 100644 rules/CodeQuality/NodeAnalyser/SetUpAssignedMockTypesResolver.php
 create mode 100644 rules/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector.php
 create mode 100644 rules/CodeQuality/Reflection/MethodParametersAndReturnTypesResolver.php
 create mode 100644 rules/CodeQuality/ValueObject/ParamTypesAndReturnType.php

diff --git a/config/sets/phpunit-code-quality.php b/config/sets/phpunit-code-quality.php
index e5d7c38c..540b41d8 100644
--- a/config/sets/phpunit-code-quality.php
+++ b/config/sets/phpunit-code-quality.php
@@ -9,6 +9,7 @@
 use Rector\PHPUnit\CodeQuality\Rector\Class_\RemoveDataProviderParamKeysRector;
 use Rector\PHPUnit\CodeQuality\Rector\Class_\SingleMockPropertyTypeRector;
 use Rector\PHPUnit\CodeQuality\Rector\Class_\TestWithToDataProviderRector;
+use Rector\PHPUnit\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector;
 use Rector\PHPUnit\CodeQuality\Rector\Class_\YieldDataProviderRector;
 use Rector\PHPUnit\CodeQuality\Rector\ClassMethod\AddInstanceofAssertForNullableInstanceRector;
 use Rector\PHPUnit\CodeQuality\Rector\ClassMethod\DataProviderArrayItemsNewLinedRector;
@@ -60,10 +61,12 @@
         NarrowSingleWillReturnCallbackRector::class,
         SingleWithConsecutiveToWithRector::class,
 
+        // type declarations
+        TypeWillReturnCallableArrowFunctionRector::class,
+
         NarrowUnusedSetUpDefinedPropertyRector::class,
 
         // specific asserts
-
         AssertCompareOnCountableWithMethodToAssertCountRector::class,
         AssertComparisonToSpecificMethodRector::class,
         AssertNotOperatorRector::class,
diff --git a/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/assigned_final_in_setup.php.inc b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/assigned_final_in_setup.php.inc
new file mode 100644
index 00000000..78ac2462
--- /dev/null
+++ b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/assigned_final_in_setup.php.inc
@@ -0,0 +1,53 @@
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Fixture;
+
+use Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Source\SomeFinalMockedClass;
+use PHPUnit\Framework\MockObject\MockObject;
+use PHPUnit\Framework\TestCase;
+
+final class AssignedFinalInSetup extends TestCase
+{
+    private MockObject $someFinalMockedClass;
+
+    protected function setUp(): void
+    {
+        $this->someFinalMockedClass = $this->createMock(SomeFinalMockedClass::class);
+    }
+
+    public function test($value): void
+    {
+        $this->someFinalMockedClass
+            ->method('anotherMethod')
+            ->willReturnCallback(fn ($age) => $value);
+    }
+}
+
+?>
+-----
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Fixture;
+
+use Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Source\SomeFinalMockedClass;
+use PHPUnit\Framework\MockObject\MockObject;
+use PHPUnit\Framework\TestCase;
+
+final class AssignedFinalInSetup extends TestCase
+{
+    private MockObject $someFinalMockedClass;
+
+    protected function setUp(): void
+    {
+        $this->someFinalMockedClass = $this->createMock(SomeFinalMockedClass::class);
+    }
+
+    public function test($value): void
+    {
+        $this->someFinalMockedClass
+            ->method('anotherMethod')
+            ->willReturnCallback(fn (int $age): float => $value);
+    }
+}
+
+?>
diff --git a/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/fill_known_param_type.php.inc b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/fill_known_param_type.php.inc
new file mode 100644
index 00000000..e6573845
--- /dev/null
+++ b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/fill_known_param_type.php.inc
@@ -0,0 +1,37 @@
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Fixture;
+
+use PHPUnit\Framework\TestCase;
+use Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Source\SomeMockedClass;
+
+final class FillKnownParamType extends TestCase
+{
+    public function test($value): void
+    {
+        $this->createMock(SomeMockedClass::class)
+            ->method('someMethod')
+            ->willReturnCallback(fn ($name) => $value);
+    }
+}
+
+?>
+-----
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Fixture;
+
+use PHPUnit\Framework\TestCase;
+use Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Source\SomeMockedClass;
+
+final class FillKnownParamType extends TestCase
+{
+    public function test($value): void
+    {
+        $this->createMock(SomeMockedClass::class)
+            ->method('someMethod')
+            ->willReturnCallback(fn (string $name): int => $value);
+    }
+}
+
+?>
diff --git a/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/include_matches.php.inc b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/include_matches.php.inc
new file mode 100644
index 00000000..8ac5d894
--- /dev/null
+++ b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/include_matches.php.inc
@@ -0,0 +1,39 @@
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Fixture;
+
+use PHPUnit\Framework\TestCase;
+use Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Source\SomeMockedClass;
+
+final class IncludeMatcher extends TestCase
+{
+    public function test($value): void
+    {
+        $this->createMock(SomeMockedClass::class)
+            ->expects($this->any())
+            ->method('someMethod')
+            ->willReturnCallback(fn ($name) => $value);
+    }
+}
+
+?>
+-----
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Fixture;
+
+use PHPUnit\Framework\TestCase;
+use Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Source\SomeMockedClass;
+
+final class IncludeMatcher extends TestCase
+{
+    public function test($value): void
+    {
+        $this->createMock(SomeMockedClass::class)
+            ->expects($this->any())
+            ->method('someMethod')
+            ->willReturnCallback(fn (string $name): int => $value);
+    }
+}
+
+?>
diff --git a/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Source/SomeFinalMockedClass.php b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Source/SomeFinalMockedClass.php
new file mode 100644
index 00000000..bafcfedf
--- /dev/null
+++ b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Source/SomeFinalMockedClass.php
@@ -0,0 +1,13 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Source;
+
+final class SomeFinalMockedClass
+{
+    public function anotherMethod(int $age): float
+    {
+        return 25.55;
+    }
+}
diff --git a/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Source/SomeMockedClass.php b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Source/SomeMockedClass.php
new file mode 100644
index 00000000..456fea95
--- /dev/null
+++ b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Source/SomeMockedClass.php
@@ -0,0 +1,12 @@
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Source;
+
+// non final on purpose so PHPStan can analyze it
+class SomeMockedClass
+{
+    public function someMethod(string $name): int
+    {
+        return 100;
+    }
+}
diff --git a/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/TypeWillReturnCallableArrowFunctionRectorTest.php b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/TypeWillReturnCallableArrowFunctionRectorTest.php
new file mode 100644
index 00000000..acedc785
--- /dev/null
+++ b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/TypeWillReturnCallableArrowFunctionRectorTest.php
@@ -0,0 +1,28 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector;
+
+use Iterator;
+use PHPUnit\Framework\Attributes\DataProvider;
+use Rector\Testing\PHPUnit\AbstractRectorTestCase;
+
+final class TypeWillReturnCallableArrowFunctionRectorTest extends AbstractRectorTestCase
+{
+    #[DataProvider('provideData')]
+    public function test(string $filePath): void
+    {
+        $this->doTestFile($filePath);
+    }
+
+    public static function provideData(): Iterator
+    {
+        return self::yieldFilesFromDirectory(__DIR__ . '/Fixture');
+    }
+
+    public function provideConfigFilePath(): string
+    {
+        return __DIR__ . '/config/configured_rule.php';
+    }
+}
diff --git a/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/config/configured_rule.php b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/config/configured_rule.php
new file mode 100644
index 00000000..588e3031
--- /dev/null
+++ b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/config/configured_rule.php
@@ -0,0 +1,10 @@
+<?php
+
+declare(strict_types=1);
+
+use Rector\Config\RectorConfig;
+use Rector\PHPUnit\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector;
+
+return static function (RectorConfig $rectorConfig): void {
+    $rectorConfig->rule(TypeWillReturnCallableArrowFunctionRector::class);
+};
diff --git a/rules/CodeQuality/NodeAnalyser/SetUpAssignedMockTypesResolver.php b/rules/CodeQuality/NodeAnalyser/SetUpAssignedMockTypesResolver.php
new file mode 100644
index 00000000..c1742f52
--- /dev/null
+++ b/rules/CodeQuality/NodeAnalyser/SetUpAssignedMockTypesResolver.php
@@ -0,0 +1,84 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Rector\PHPUnit\CodeQuality\NodeAnalyser;
+
+use PhpParser\Node\Expr\Assign;
+use PhpParser\Node\Expr\ClassConstFetch;
+use PhpParser\Node\Expr\MethodCall;
+use PhpParser\Node\Expr\PropertyFetch;
+use PhpParser\Node\Expr\Variable;
+use PhpParser\Node\Stmt\Class_;
+use PhpParser\Node\Stmt\ClassMethod;
+use PhpParser\Node\Stmt\Expression;
+use Rector\NodeNameResolver\NodeNameResolver;
+use Rector\ValueObject\MethodName;
+use Webmozart\Assert\Assert;
+
+final readonly class SetUpAssignedMockTypesResolver
+{
+    public function __construct(
+        private NodeNameResolver $nodeNameResolver
+    ) {
+    }
+
+    /**
+     * @return array<string, string>
+     */
+    public function resolveFromClass(Class_ $class): array
+    {
+        $setUpClassMethod = $class->getMethod(MethodName::SET_UP);
+        if (! $setUpClassMethod instanceof ClassMethod) {
+            return [];
+        }
+
+        $propertyNameToMockedTypes = [];
+        foreach ((array) $setUpClassMethod->stmts as $stmt) {
+            if (! $stmt instanceof Expression) {
+                continue;
+            }
+
+            if (! $stmt->expr instanceof Assign) {
+                continue;
+            }
+
+            $assign = $stmt->expr;
+            if (! $assign->expr instanceof MethodCall) {
+                continue;
+            }
+
+            if (! $this->nodeNameResolver->isNames($assign->expr->name, ['createMock', 'getMockBuilder'])) {
+                continue;
+            }
+
+            if (! $assign->var instanceof PropertyFetch && ! $assign->var instanceof Variable) {
+                continue;
+            }
+
+            $mockedClassNameExpr = $assign->expr->getArgs()[0]
+                ->value;
+            if (! $mockedClassNameExpr instanceof ClassConstFetch) {
+                continue;
+            }
+
+            $propertyOrVariableName = $this->resolvePropertyOrVariableName($assign->var);
+            $mockedClass = $this->nodeNameResolver->getName($mockedClassNameExpr->class);
+
+            Assert::string($mockedClass);
+
+            $propertyNameToMockedTypes[$propertyOrVariableName] = $mockedClass;
+        }
+
+        return $propertyNameToMockedTypes;
+    }
+
+    private function resolvePropertyOrVariableName(PropertyFetch|Variable $propertyFetchOrVariable): ?string
+    {
+        if ($propertyFetchOrVariable instanceof Variable) {
+            return $this->nodeNameResolver->getName($propertyFetchOrVariable);
+        }
+
+        return $this->nodeNameResolver->getName($propertyFetchOrVariable->name);
+    }
+}
diff --git a/rules/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector.php b/rules/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector.php
new file mode 100644
index 00000000..52fca877
--- /dev/null
+++ b/rules/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector.php
@@ -0,0 +1,285 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Rector\PHPUnit\CodeQuality\Rector\Class_;
+
+use PhpParser\Node;
+use PhpParser\Node\Expr;
+use PhpParser\Node\Expr\ArrowFunction;
+use PhpParser\Node\Expr\Closure;
+use PhpParser\Node\Expr\MethodCall;
+use PhpParser\Node\Expr\PropertyFetch;
+use PhpParser\Node\Expr\Variable;
+use PhpParser\Node\Scalar\String_;
+use PhpParser\Node\Stmt\Class_;
+use PHPStan\Type\IntersectionType;
+use PHPStan\Type\MixedType;
+use PHPStan\Type\NeverType;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\Type;
+use PHPUnit\Framework\MockObject\Builder\InvocationMocker;
+use Rector\Enum\ClassName;
+use Rector\PHPStanStaticTypeMapper\Enum\TypeKind;
+use Rector\PHPUnit\CodeQuality\NodeAnalyser\SetUpAssignedMockTypesResolver;
+use Rector\PHPUnit\CodeQuality\Reflection\MethodParametersAndReturnTypesResolver;
+use Rector\PHPUnit\CodeQuality\ValueObject\ParamTypesAndReturnType;
+use Rector\PHPUnit\NodeAnalyzer\TestsNodeAnalyzer;
+use Rector\Rector\AbstractRector;
+use Rector\StaticTypeMapper\StaticTypeMapper;
+use Symplify\RuleDocGenerator\ValueObject\CodeSample\CodeSample;
+use Symplify\RuleDocGenerator\ValueObject\RuleDefinition;
+
+/**
+ * @see \Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\TypeWillReturnCallableArrowFunctionRectorTest
+ */
+final class TypeWillReturnCallableArrowFunctionRector extends AbstractRector
+{
+    /**
+     * @var string
+     */
+    private const WILL_RETURN_CALLBACK = 'willReturnCallback';
+
+    public function __construct(
+        private readonly TestsNodeAnalyzer $testsNodeAnalyzer,
+        private readonly StaticTypeMapper $staticTypeMapper,
+        private readonly SetUpAssignedMockTypesResolver $setUpAssignedMockTypesResolver,
+        private readonly MethodParametersAndReturnTypesResolver $methodParametersAndReturnTypesResolver,
+    ) {
+    }
+
+    public function getRuleDefinition(): RuleDefinition
+    {
+        return new RuleDefinition(
+            'Decorate callbacks and arrow functions in willReturnCallback() with known param/return types based on reflection method',
+            [
+                new CodeSample(
+                    <<<'CODE_SAMPLE'
+use PHPUnit\Framework\TestCase;
+
+final class SomeTest extends TestCase
+{
+    public function testSomething()
+    {
+        $this->createMock(SomeClass::class)
+            ->method('someMethod')
+            ->willReturnCallback(function ($arg) {
+                return $arg;
+            });
+    }
+}
+
+final class SomeClass
+{
+    public function someMethod(string $arg): string
+    {
+        return $arg . ' !';
+    }
+}
+CODE_SAMPLE
+                    ,
+                    <<<'CODE_SAMPLE'
+use PHPUnit\Framework\TestCase;
+
+final class SomeTest extends TestCase
+{
+    public function testSomething()
+    {
+        $this->createMock(SomeClass::class)
+            ->method('someMethod')
+            ->willReturnCallback(
+                function (string $arg): string {
+                    return $arg;
+                }
+            );
+    }
+}
+CODE_SAMPLE
+                ),
+            ]
+        );
+    }
+
+    /**
+     * @return array<class-string<Node>>
+     */
+    public function getNodeTypes(): array
+    {
+        return [Class_::class];
+    }
+
+    /**
+     * @param Class_ $node
+     */
+    public function refactor(Node $node): ?Class_
+    {
+        if (! $this->testsNodeAnalyzer->isInTestClass($node)) {
+            return null;
+        }
+
+        $hasChanged = false;
+
+        $propertyNameToMockedTypes = $this->setUpAssignedMockTypesResolver->resolveFromClass($node);
+
+        $this->traverseNodesWithCallable($node->getMethods(), function (Node $node) use (
+            &$hasChanged,
+            $propertyNameToMockedTypes
+        ) {
+            if (! $node instanceof MethodCall || $node->isFirstClassCallable()) {
+                return null;
+            }
+
+            if (! $this->isName($node->name, self::WILL_RETURN_CALLBACK)) {
+                return null;
+            }
+
+            $innerArg = $node->getArgs()[0]
+                ->value;
+            if (! $innerArg instanceof ArrowFunction && ! $innerArg instanceof Closure) {
+                return null;
+            }
+
+            if (! $node->var instanceof MethodCall) {
+                return null;
+            }
+
+            $parentMethodCall = $node->var;
+            if (! $this->isName($parentMethodCall->name, 'method')) {
+                return null;
+            }
+
+            $methodNameExpr = $parentMethodCall->getArgs()[0]
+                ->value;
+            if (! $methodNameExpr instanceof String_) {
+                return null;
+            }
+
+            $methodName = $methodNameExpr->value;
+            $callerType = $this->getType($parentMethodCall->var);
+
+            if ($callerType instanceof ObjectType && $callerType->getClassName() === InvocationMocker::class) {
+                $parentMethodCall = $parentMethodCall->var;
+
+                if ($parentMethodCall instanceof MethodCall) {
+                    $callerType = $this->getType($parentMethodCall->var);
+                }
+            }
+
+            $callerType = $this->fallbackMockedObjectInSetUp(
+                $callerType,
+                $parentMethodCall,
+                $propertyNameToMockedTypes
+            );
+
+            // we need mocks
+            if (! $callerType instanceof IntersectionType) {
+                return null;
+            }
+
+            $hasChanged = false;
+
+            $parameterTypesAndReturnType = $this->methodParametersAndReturnTypesResolver->resolveFromReflection(
+                $callerType,
+                $methodName
+            );
+
+            if (! $parameterTypesAndReturnType instanceof ParamTypesAndReturnType) {
+                return null;
+            }
+
+            foreach ($innerArg->params as $key => $param) {
+                // avoid typing variadic parameters
+                if ($param->variadic) {
+                    continue;
+                }
+
+                // already filled, lets skip it
+                if ($param->type instanceof Node) {
+                    continue;
+                }
+
+                $nativeParameterType = $parameterTypesAndReturnType->getParamTypes()[$key] ?? null;
+                // we need specific non-mixed type
+                if ($nativeParameterType === null) {
+                    continue;
+                }
+
+                if ($nativeParameterType instanceof MixedType) {
+                    continue;
+                }
+
+                $parameterTypeNode = $this->staticTypeMapper->mapPHPStanTypeToPhpParserNode(
+                    $nativeParameterType,
+                    TypeKind::PARAM
+                );
+
+                if (! $parameterTypeNode instanceof Node) {
+                    continue;
+                }
+
+                $param->type = $parameterTypeNode;
+                $hasChanged = true;
+            }
+
+            if (! $innerArg->returnType instanceof Node) {
+                $returnTypeNode = $this->staticTypeMapper->mapPHPStanTypeToPhpParserNode(
+                    $parameterTypesAndReturnType->getReturnType(),
+                    TypeKind::RETURN
+                );
+
+                if ($returnTypeNode instanceof Node) {
+                    $innerArg->returnType = $returnTypeNode;
+                    $hasChanged = true;
+                }
+            }
+
+            $hasChanged = true;
+        });
+
+        if (! $hasChanged) {
+            return null;
+        }
+
+        return $node;
+    }
+
+    /**
+     * @param array<string, string> $propertyNameToMockedTypes
+     */
+    private function fallbackMockedObjectInSetUp(
+        Type $callerType,
+        Expr $expr,
+        array $propertyNameToMockedTypes
+    ): mixed {
+        if (! $callerType instanceof ObjectType && ! $callerType instanceof NeverType) {
+            return $callerType;
+        }
+
+        if (! $expr instanceof MethodCall) {
+            return $callerType;
+        }
+
+        if ($callerType instanceof ObjectType && $callerType->getClassName() !== ClassName::MOCK_OBJECT) {
+            return $callerType;
+        }
+
+        // type is missing, because of "final" keyword on mocked class
+        // resolve from constructor instead
+        if (! $expr->var instanceof PropertyFetch && ! $expr->var instanceof Variable) {
+            return $callerType;
+        }
+
+        if ($expr->var instanceof Variable) {
+            $propertyOrVariableName = $this->getName($expr->var);
+        } else {
+            $propertyOrVariableName = $this->getName($expr->var->name);
+        }
+
+        if (isset($propertyNameToMockedTypes[$propertyOrVariableName])) {
+            $mockedType = $propertyNameToMockedTypes[$propertyOrVariableName];
+            return new IntersectionType([$callerType, new ObjectType($mockedType)]);
+        }
+
+        return $callerType;
+    }
+}
diff --git a/rules/CodeQuality/Reflection/MethodParametersAndReturnTypesResolver.php b/rules/CodeQuality/Reflection/MethodParametersAndReturnTypesResolver.php
new file mode 100644
index 00000000..57d8f66c
--- /dev/null
+++ b/rules/CodeQuality/Reflection/MethodParametersAndReturnTypesResolver.php
@@ -0,0 +1,76 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Rector\PHPUnit\CodeQuality\Reflection;
+
+use PHPStan\Reflection\ClassReflection;
+use PHPStan\Reflection\ExtendedMethodReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
+use PHPStan\Type\IntersectionType;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\Type;
+use Rector\Enum\ClassName;
+use Rector\PHPUnit\CodeQuality\ValueObject\ParamTypesAndReturnType;
+
+final class MethodParametersAndReturnTypesResolver
+{
+    public function resolveFromReflection(
+        IntersectionType $intersectionType,
+        string $methodName
+    ): ?ParamTypesAndReturnType {
+        foreach ($intersectionType->getTypes() as $intersectionedType) {
+            if (! $intersectionedType instanceof ObjectType) {
+                continue;
+            }
+
+            if ($intersectionedType->getClassName() === ClassName::MOCK_OBJECT) {
+                continue;
+            }
+
+            $classReflection = $intersectionedType->getClassReflection();
+            if (! $classReflection instanceof ClassReflection) {
+                continue;
+            }
+
+            if (! $classReflection->hasNativeMethod($methodName)) {
+                continue;
+            }
+
+            $mockedMethodReflection = $classReflection->getNativeMethod($methodName);
+
+            $parameterTypes = $this->resolveParameterTypes($mockedMethodReflection);
+            $returnType = $this->resolveReturnType($mockedMethodReflection);
+
+            return new ParamTypesAndReturnType($parameterTypes, $returnType);
+        }
+
+        return null;
+    }
+
+    /**
+     * @return Type[]
+     */
+    private function resolveParameterTypes(ExtendedMethodReflection $extendedMethodReflection): array
+    {
+        $extendedParametersAcceptor = ParametersAcceptorSelector::combineAcceptors(
+            $extendedMethodReflection->getVariants()
+        );
+
+        $parameterTypes = [];
+        foreach ($extendedParametersAcceptor->getParameters() as $parameterReflection) {
+            $parameterTypes[] = $parameterReflection->getType();
+        }
+
+        return $parameterTypes;
+    }
+
+    private function resolveReturnType(ExtendedMethodReflection $extendedMethodReflection): Type
+    {
+        $extendedParametersAcceptor = ParametersAcceptorSelector::combineAcceptors(
+            $extendedMethodReflection->getVariants()
+        );
+
+        return $extendedParametersAcceptor->getReturnType();
+    }
+}
diff --git a/rules/CodeQuality/ValueObject/ParamTypesAndReturnType.php b/rules/CodeQuality/ValueObject/ParamTypesAndReturnType.php
new file mode 100644
index 00000000..d4acbb8c
--- /dev/null
+++ b/rules/CodeQuality/ValueObject/ParamTypesAndReturnType.php
@@ -0,0 +1,32 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Rector\PHPUnit\CodeQuality\ValueObject;
+
+use PHPStan\Type\Type;
+
+final readonly class ParamTypesAndReturnType
+{
+    /**
+     * @param Type[] $paramTypes
+     */
+    public function __construct(
+        private array $paramTypes,
+        private ?Type $returnType
+    ) {
+    }
+
+    /**
+     * @return Type[]
+     */
+    public function getParamTypes(): array
+    {
+        return $this->paramTypes;
+    }
+
+    public function getReturnType(): ?Type
+    {
+        return $this->returnType;
+    }
+}