diff --git a/lib/Transunit/Pass/LocalRevealPass.php b/lib/Transunit/Pass/LocalRevealPass.php
index 4e82688..158d608 100644
--- a/lib/Transunit/Pass/LocalRevealPass.php
+++ b/lib/Transunit/Pass/LocalRevealPass.php
@@ -6,9 +6,14 @@
 use PhpParser\NodeFinder;
 use PhpParser\NodeTraverser;
 use Transunit\Pass;
-use Transunit\Visitor\ParentConnectingVisitor;
-use Transunit\Visitor\RevealCollaboratorVisitor;
-
+use Transunit\Visitor;
+
+/**
+ * ```
+ * - $this->_testSubject->onKernelRequest($event);
+ * + $this->_testSubject->onKernelRequest($event->reveal());
+ * ```
+ */
 class LocalRevealPass implements Pass
 {
     public function find(NodeFinder $nodeFinder, $ast): array
@@ -48,8 +53,8 @@ private function reveal(Node\Stmt\ClassMethod $node): void
         }
 
         $subNodeTraverser = new NodeTraverser(
-            new ParentConnectingVisitor(),
-            new RevealCollaboratorVisitor($collabs)
+            new Visitor\ParentConnectingVisitor(),
+            new Visitor\RevealCollaboratorVisitor($collabs)
         );
 
         $node->stmts = $subNodeTraverser->traverse($node->stmts);
diff --git a/lib/Transunit/Pass/ReplaceCallsToGetWrappedObjectPass.php b/lib/Transunit/Pass/ReplaceCallsToGetWrappedObjectPass.php
index 4752409..48e81ce 100644
--- a/lib/Transunit/Pass/ReplaceCallsToGetWrappedObjectPass.php
+++ b/lib/Transunit/Pass/ReplaceCallsToGetWrappedObjectPass.php
@@ -6,6 +6,14 @@
 use PhpParser\NodeFinder;
 use Transunit\Pass;
 
+/**
+ * ```
+ *      $finder->getIterator()->willReturn(new \ArrayIterator([
+ * -        $splFileInfo->getWrappedObject(),
+ * +        $splFileInfo->reveal(),
+ *      ]));
+ * ```
+ */
 class ReplaceCallsToGetWrappedObjectPass implements Pass
 {
     public function find(NodeFinder $nodeFinder, $ast): array
diff --git a/lib/Transunit/Pass/TestSubjectCallPass.php b/lib/Transunit/Pass/TestSubjectCallPass.php
new file mode 100644
index 0000000..da9f94f
--- /dev/null
+++ b/lib/Transunit/Pass/TestSubjectCallPass.php
@@ -0,0 +1,42 @@
+<?php
+
+namespace Transunit\Pass;
+
+use PhpParser\Node;
+use PhpParser\NodeFinder;
+use PhpParser\NodeTraverser;
+use Transunit\Pass;
+use Transunit\Visitor;
+
+/**
+ * ```
+ *       function it_contracts_out_agents(AgentRepository $agentRepository, EventDispatcher $eventDispatcher, Agent $agent47, ContractEvent $event)
+ *       {
+ *           $this->agentRepository->find(47)->willReturn($agent47);
+ *           $this->eventDispatcher->dispatch($event)->shouldBeCalled();
+ *
+ * -         $this->contractOut(47)->shouldReturn($agent47);
+ * +         $this->_testSubject->contractOut(47)->shouldReturn($agent47);
+ *       }
+ * ```
+ */
+class TestSubjectCallPass implements Pass
+{
+    public function find(NodeFinder $nodeFinder, $ast): array
+    {
+        return $nodeFinder->findInstanceOf($ast, Node\Stmt\ClassMethod::class);
+    }
+
+    public function rewrite(Node $node): void
+    {
+        if (!$node instanceof Node\Stmt\ClassMethod) {
+            return;
+        }
+
+        $subNodeTraverser = new NodeTraverser(
+            new Visitor\TestSubjectCallVisitor()
+        );
+
+        $node->stmts = $subNodeTraverser->traverse($node->stmts);
+    }
+}
diff --git a/lib/Transunit/Pass/TestSubjectCallsPass.php b/lib/Transunit/Pass/TestSubjectCallsPass.php
deleted file mode 100644
index 759a1f5..0000000
--- a/lib/Transunit/Pass/TestSubjectCallsPass.php
+++ /dev/null
@@ -1,87 +0,0 @@
-<?php
-
-namespace Transunit\Pass;
-
-use PhpParser\Node;
-use PhpParser\NodeFinder;
-use Transunit\Pass;
-
-/**
- * ```
- *       function it_contracts_out_agents(AgentRepository $agentRepository, EventDispatcher $eventDispatcher, Agent $agent47, ContractEvent $event)
- *       {
- *           $this->agentRepository->find(47)->willReturn($agent47);
- *           $this->eventDispatcher->dispatch($event)->shouldBeCalled();
- *
- * -         $this->contractOut(47)->shouldReturn($agent47);
- * +         $this->_testSubject->contractOut(47)->shouldReturn($agent47);
- *       }
- * ```
- */
-class TestSubjectCallsPass implements Pass
-{
-    public function find(NodeFinder $nodeFinder, $ast): array
-    {
-        return $nodeFinder->findInstanceOf($ast, Node\Stmt\ClassMethod::class);
-    }
-
-    public function rewrite(Node $node): void
-    {
-        if (!$node instanceof Node\Stmt\ClassMethod) {
-            return;
-        }
-
-        foreach ($node->stmts as $stmt) {
-            if (!$stmt instanceof Node\Stmt\Expression) {
-                continue;
-            }
-
-            $resolvedCall = $stmt->expr;
-
-            if ($resolvedCall instanceof Node\Expr\Assign) {
-                $resolvedCall = $resolvedCall->expr;
-            }
-
-            if (!$resolvedCall instanceof Node\Expr\MethodCall) {
-                continue;
-            }
-
-            if ($resolvedCall->var instanceof Node\Expr\MethodCall) {
-                $resolvedCall = $resolvedCall->var;
-            }
-
-            $this->callMethodOnTestSubject($resolvedCall);
-        }
-    }
-
-    /**
-     * $this->_testSubject->doSomething();
-     */
-    private function callMethodOnTestSubject(Node\Expr\MethodCall $stmt): void
-    {
-        if (!$stmt->var instanceof Node\Expr\Variable) {
-            return;
-        }
-
-        if ('this' !== $stmt->var->name) {
-            return;
-        }
-
-        if ('prophesize' === $stmt->name->toString()) {
-            return;
-        }
-
-        if ('beConstructedWith' === $stmt->name->toString()) {
-            return;
-        }
-
-        if ('beConstructedThrough' === $stmt->name->toString()) {
-            return;
-        }
-
-        $stmt->var = new Node\Expr\PropertyFetch(
-            new Node\Expr\Variable('this'),
-            '_testSubject'
-        );
-    }
-}
diff --git a/lib/Transunit/Transunit.php b/lib/Transunit/Transunit.php
index 07d535c..344328f 100644
--- a/lib/Transunit/Transunit.php
+++ b/lib/Transunit/Transunit.php
@@ -68,7 +68,7 @@ private static function processFile(string $path): string
             new Pass\AddTestMethodPrefixPass(),
             new Pass\TestSubjectInstantiationPass(),
             new Pass\GlobalRevealPass(),
-            new Pass\TestSubjectCallsPass(),
+            new Pass\TestSubjectCallPass(),
             new Pass\AssertionPass(), // run after CallTestSubjectPass.
             new Pass\ExceptionAssertionPass(),
             new Pass\DeclareGlobalCollaboratorPass(),
diff --git a/lib/Transunit/Visitor/TestSubjectCallVisitor.php b/lib/Transunit/Visitor/TestSubjectCallVisitor.php
new file mode 100644
index 0000000..0a260d5
--- /dev/null
+++ b/lib/Transunit/Visitor/TestSubjectCallVisitor.php
@@ -0,0 +1,43 @@
+<?php
+
+namespace Transunit\Visitor;
+
+use PhpParser\Node;
+use PhpParser\NodeVisitorAbstract;
+
+class TestSubjectCallVisitor extends NodeVisitorAbstract
+{
+    public function leaveNode(Node $node)
+    {
+        if (! $node instanceof Node\Expr\MethodCall) {
+            return $node;
+        }
+
+        if (!$node->var instanceof Node\Expr\Variable) {
+            return $node;
+        }
+
+        if ('this' !== $node->var->name) {
+            return $node;
+        }
+
+        if ('prophesize' === $node->name->toString()) {
+            return $node;
+        }
+
+        if ('beConstructedWith' === $node->name->toString()) {
+            return $node;
+        }
+
+        if ('beConstructedThrough' === $node->name->toString()) {
+            return $node;
+        }
+
+        $node->var = new Node\Expr\PropertyFetch(
+            new Node\Expr\Variable('this'),
+            '_testSubject'
+        );
+
+        return $node;
+    }
+}