Handle ->shouldBeCalled()->willReturn() chained methods.

lib/Transunit/Pass/CallGlobalCollaboratorPass.php

@@ -5,6 +5,7 @@
 use PhpParser\Node;
 use PhpParser\NodeFinder;
 use Transunit\Pass;
+use Transunit\RootMethodCallExtractor;
 
 /**
  * ```
@@ -75,56 +76,43 @@ private function callGlobalCollaborators(Node\Stmt\Class_ $node): void
                 foreach ($stmts as $i => $stmt) {
                     $newStmts[] = $stmt;
 
-                    if (!$stmt instanceof Node\Stmt\Expression) {
+                    if (! $stmt instanceof Node\Stmt\Expression) {
                         continue;
                     }
 
-                    if (!$stmt->expr instanceof Node\Expr\MethodCall) {
+                    if (! $stmt->expr instanceof Node\Expr\MethodCall) {
                         continue;
                     }
 
-                    // ->willReturn()
-                    // ->shouldBeCalled()
-
-                    if (!$stmt->expr->var instanceof Node\Expr\MethodCall) {
+                    /** @see \Prophecy\Prophecy\MethodProphecy */
+                    if (! in_array($stmt->expr->name->toString(), [
+                        'shouldBeCalled',
+                        'shouldNotBeCalled',
+                        'shouldBeCalledTimes',
+                        'willReturn',
+                    ], true)) {
                         continue;
                     }
 
-                    // ->stubbedMethod()->willReturn()
-                    // ->mockedCall()->shouldBeCalled()
-
-                    if (!$stmt->expr->var->var instanceof Node\Expr\Variable) {
-                        continue;
-                    }
+                    $rootMethodCall = (new RootMethodCallExtractor())->locate($stmt);
 
-                    // $collaborator->stubbedMethod()->willReturn()
-                    // $collaborator->mockedCall()->shouldBeCalled()
+                    // At this point the method call is expected to be one of:
+                    // - $collaborator->stubbedMethod()->willReturn()
+                    // - $collaborator->mockedCall()->shouldBeCalled()
+                    // - $collaborator->mockedCall()->shouldBeCalled()->willReturn()
 
-                    if ($stmt->expr->var->var->name !== $param->var->name) {
+                    if ($rootMethodCall->var->name !== $param->var->name) {
                         continue;
                     }
 
                     // $this->collaborator->stubbedMethod()->willReturn()
 
-                    array_pop($newStmts);
-
-                    $newStmts[] = new Node\Stmt\Expression(
-                    // ->willReturn()
-                        new Node\Expr\MethodCall(
-                        // ->stubbedMethod()
-                            new Node\Expr\MethodCall(
-                            // ->collaborator
-                                new Node\Expr\PropertyFetch(
-                                    new Node\Expr\Variable('this'),
-                                    $param->var->name
-                                ),
-                                $stmt->expr->var->name,
-                                $stmt->expr->var->args
-                            ),
-                            $stmt->expr->name,
-                            $stmt->expr->args
-                        )
+                    $rootMethodCall->var = new Node\Expr\PropertyFetch(
+                        new Node\Expr\Variable('this'),
+                        $param->var->name
                     );
+
+                    $newStmts[] = $stmt;
                 }
 
                 $classMethodNode->stmts = $newStmts;

lib/Transunit/Pass/ExceptionAssertionPass.php

@@ -80,7 +80,7 @@ private function expectException(Node\Stmt\ClassMethod $node): void
             $subjectMethodName = $expression->expr->args[0]->value->value;
             $subjectMethodArgs = [];
 
-            if ($expression->expr->args[1]->value instanceof Node\Expr\Array_) {
+            if (count($expression->expr->args) > 1 && $expression->expr->args[1]->value instanceof Node\Expr\Array_) {
                 /** @var Node\ArrayItem $item */
                 foreach ($expression->expr->args[1]->value->items as $item) {
                     $subjectMethodArgs[] = new Node\Arg($item->value);

lib/Transunit/RootMethodCallExtractor.php

@@ -0,0 +1,26 @@
+<?php
+
+namespace Transunit;
+
+use PhpParser\Node\Expr\MethodCall;
+use PhpParser\Node\Stmt\Expression;
+
+class RootMethodCallExtractor
+{
+    public function locate(Expression $stmt): MethodCall
+    {
+        if (! $stmt->expr instanceof MethodCall) {
+            throw new \DomainException('Cannot locate root variable without a method call chain.');
+        }
+
+        $current = $stmt->expr;
+        $next = $stmt->expr->var;
+
+        while ($next instanceof MethodCall) {
+            $current = $next;
+            $next = $next->var;
+        }
+
+        return $current;
+    }
+}