Add return support

Author: TomasVotruba

rules-tests/CodeQuality/Rector/Expression/ConfiguredMockEntityToSetterObjectRector/Fixture/direct-return.php.inc

@@ -0,0 +1,35 @@
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Expression\ConfiguredMockEntityToSetterObjectRector\Fixture;
+
+use Rector\PHPUnit\Tests\CodeQuality\Rector\Expression\ConfiguredMockEntityToSetterObjectRector\Source\SomeEntityToBeConfigured;
+
+final class DirectReturn extends \PHPUnit\Framework\TestCase
+{
+    public function test()
+    {
+        return $this->createConfiguredMock(SomeEntityToBeConfigured::class, [
+            'getName'=> 'John',
+        ]);
+    }
+}
+
+?>
+-----
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Expression\ConfiguredMockEntityToSetterObjectRector\Fixture;
+
+use Rector\PHPUnit\Tests\CodeQuality\Rector\Expression\ConfiguredMockEntityToSetterObjectRector\Source\SomeEntityToBeConfigured;
+
+final class DirectReturn extends \PHPUnit\Framework\TestCase
+{
+    public function test()
+    {
+        $someEntityToBeConfigured = new \Rector\PHPUnit\Tests\CodeQuality\Rector\Expression\ConfiguredMockEntityToSetterObjectRector\Source\SomeEntityToBeConfigured();
+        $someEntityToBeConfigured->setName('John');
+        return $someEntityToBeConfigured;
+    }
+}
+
+?>

rules/CodeQuality/Rector/Expression/ConfiguredMockEntityToSetterObjectRector.php

@@ -4,16 +4,21 @@
 
 namespace Rector\PHPUnit\CodeQuality\Rector\Expression;
 
+use Webmozart\Assert\Assert;
+use Nette\Utils\Strings;
 use PhpParser\Node;
 use PhpParser\Node\Arg;
 use PhpParser\Node\Expr;
 use PhpParser\Node\Expr\Array_;
 use PhpParser\Node\Expr\Assign;
 use PhpParser\Node\Expr\MethodCall;
 use PhpParser\Node\Expr\New_;
+use PhpParser\Node\Expr\Variable;
 use PhpParser\Node\Name\FullyQualified;
+use PhpParser\Node\Stmt;
 use PhpParser\Node\Stmt\Class_;
 use PhpParser\Node\Stmt\Expression;
+use PhpParser\Node\Stmt\Return_;
 use PHPStan\Reflection\ReflectionProvider;
 use Rector\Doctrine\NodeAnalyzer\DoctrineEntityDetector;
 use Rector\PhpParser\AstResolver;
@@ -81,31 +86,37 @@ public function test()
      */
     public function getNodeTypes(): array
     {
-        return [Expression::class];
+        return [Expression::class, Return_::class];
     }
 
     /**
-     * @param Expression $node
-     * @return Expression[]|null
+     * @param Expression|Return_ $node
+     * @return Stmt[]|null
      */
     public function refactor(Node $node): ?array
     {
         if (! $this->testsNodeAnalyzer->isInTestClass($node)) {
             return null;
         }
 
-        if (! $node->expr instanceof Assign) {
-            return null;
-        }
+        $assign = null;
+        if ($node instanceof Return_) {
+            if ($node->expr instanceof MethodCall) {
+                $methodCall = $node->expr;
+            } else {
+                return null;
+            }
+        } elseif ($node->expr instanceof Assign) {
+            $assign = $node->expr;
+            if (! $assign->expr instanceof MethodCall) {
+                return null;
+            }
 
-        $assign = $node->expr;
-        if (! $assign->expr instanceof MethodCall) {
+            $methodCall = $assign->expr;
+        } else {
             return null;
         }
 
-        $objectVariable = $assign->var;
-
-        $methodCall = $assign->expr;
         if (! $this->isName($methodCall->name, 'createConfiguredMock')) {
             return null;
         }
@@ -126,11 +137,13 @@ public function refactor(Node $node): ?array
             return null;
         }
 
-        $assign->expr = new New_(new FullyQualified($doctrineClass));
+        if ($node instanceof Expression) {
+            Assert::isInstanceOf($assign, Assign::class);
+            return $this->createForAssign($doctrineClass, $assign, $definedGettersArg->value, $node);
+        }
 
-        $setterExpressions = $this->createEntitySetterExpressions($definedGettersArg->value, $objectVariable);
+        return $this->createForReturn($doctrineClass, $definedGettersArg->value, $node);
 
-        return array_merge([$node], $setterExpressions);
     }
 
     /**
@@ -191,4 +204,39 @@ private function matchDoctrineClassName(Expr $expr): string|null
 
         return $mockedClassValue;
     }
+
+    /**
+     * @return Stmt[]
+     */
+    private function createForReturn(string $doctrineClass, Array_ $array, Return_ $return): array
+    {
+        $shortClassName = Strings::after($doctrineClass, '\\', -1);
+        $objectVariable = new Variable(lcfirst((string) $shortClassName));
+
+        $new = new New_(new FullyQualified($doctrineClass));
+        $assign = new Assign($objectVariable, $new);
+
+        $setterExpressions = $this->createEntitySetterExpressions($array, $objectVariable);
+
+        $return->expr = $objectVariable;
+
+        return array_merge([new Expression($assign)], $setterExpressions, [$return]);
+    }
+
+    /**
+     * @return Stmt[]
+     */
+    private function createForAssign(
+        string $doctrineClass,
+        Assign $assign,
+        Array_ $definedGettersArray,
+        Expression $expression
+    ): array {
+        $assign->expr = new New_(new FullyQualified($doctrineClass));
+        $objectVariable = $assign->var;
+
+        $setterExpressions = $this->createEntitySetterExpressions($definedGettersArray, $objectVariable);
+
+        return array_merge([$expression], $setterExpressions);
+    }
 }