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,19 @@
 
 namespace Rector\PHPUnit\CodeQuality\Rector\Expression;
 
+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\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 +84,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 Node\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 +135,26 @@ public function refactor(Node $node): ?array
             return null;
         }
 
-        $assign->expr = new New_(new FullyQualified($doctrineClass));
+        if ($assign instanceof Assign) {
+            $assign->expr = new New_(new FullyQualified($doctrineClass));
+            $objectVariable = $assign->var;
+
+            $setterExpressions = $this->createEntitySetterExpressions($definedGettersArg->value, $objectVariable);
+
+            return array_merge([$node], $setterExpressions);
+        }
+
+        $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($definedGettersArg->value, $objectVariable);
 
-        return array_merge([$node], $setterExpressions);
+        $node->expr = $objectVariable;
+        return array_merge([new Expression($assign)], $setterExpressions, [$node]);
+
     }
 
     /**