From ec15fd6cfdf4dda4158410da32e1db3f0302799a Mon Sep 17 00:00:00 2001
From: Roland <rk@rkuhn.info>
Date: Sun, 19 Feb 2012 00:09:04 +0100
Subject: [PATCH] add TestKit.setAutoPilot, see #1807

---
 akka-docs/scala/testing.rst                   | 23 +++++++++++++++++
 .../src/main/scala/akka/testkit/TestKit.scala | 22 +++++++++++++---
 .../scala/akka/testkit/TestProbeSpec.scala    | 25 +++++++++++++++++++
 3 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/akka-docs/scala/testing.rst b/akka-docs/scala/testing.rst
index db0f670c93..e99f6b6838 100644
--- a/akka-docs/scala/testing.rst
+++ b/akka-docs/scala/testing.rst
@@ -421,6 +421,14 @@ using a small example:
 .. includecode:: code/akka/docs/testkit/TestkitDocSpec.scala
    :include: imports-test-probe,my-double-echo,test-probe
 
+Here a the system under test is simulated by :class:`MyDoubleEcho`, which is
+supposed to mirror its input to two outputs. Attaching two test probes enables
+verification of the (simplistic) behavior. Another example would be two actors
+A and B which collaborate by A sending messages to B. In order to verify this
+message flow, a :class:`TestProbe` could be inserted as target of A, using the
+forwarding capabilities or auto-pilot described below to include a real B in
+the test setup.
+
 Probes may also be equipped with custom assertions to make your test code even
 more concise and clear:
 
@@ -455,6 +463,21 @@ network functioning:
 The ``dest`` actor will receive the same message invocation as if no test probe
 had intervened.
 
+Auto-Pilot
+^^^^^^^^^^
+
+Receiving messages in a queue for later inspection is nice, but in order to
+keep a test running and verify traces later you can also install an
+:class:`AutoPilot` in the participating test probes (actually in any
+:class:`TestKit`) which is invoked before enqueueing to the inspection queue.
+This code can be used to forward messages, e.g. in a chain ``A --> Probe -->
+B``, as long as a certain protocol is obeyed.
+
+.. includecode:: ../../akka-testkit/src/test/scala/akka/testkit/TestProbeSpec.scala#autopilot
+
+The :meth:`run` method must return the auto-pilot for the next message, wrapped
+in an :class:`Option`; setting it to :obj:`None` terminates the auto-pilot.
+
 Caution about Timing Assertions
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/akka-testkit/src/main/scala/akka/testkit/TestKit.scala b/akka-testkit/src/main/scala/akka/testkit/TestKit.scala
index bdfab36ede..965fc555e1 100644
--- a/akka-testkit/src/main/scala/akka/testkit/TestKit.scala
+++ b/akka-testkit/src/main/scala/akka/testkit/TestKit.scala
@@ -16,9 +16,14 @@ import akka.util.Timeout
 object TestActor {
   type Ignore = Option[PartialFunction[AnyRef, Boolean]]
 
+  trait AutoPilot {
+    def run(sender: ActorRef, msg: Any): Option[AutoPilot]
+  }
+
   case class SetIgnore(i: Ignore)
   case class Watch(ref: ActorRef)
   case class UnWatch(ref: ActorRef)
+  case class SetAutoPilot(ap: AutoPilot)
 
   trait Message {
     def msg: AnyRef
@@ -36,11 +41,15 @@ class TestActor(queue: BlockingDeque[TestActor.Message]) extends Actor {
 
   var ignore: Ignore = None
 
+  var autopilot: Option[AutoPilot] = None
+
   def receive = {
-    case SetIgnore(ign)   ⇒ ignore = ign
-    case x @ Watch(ref)   ⇒ context.watch(ref); queue.offerLast(RealMessage(x, self))
-    case x @ UnWatch(ref) ⇒ context.unwatch(ref); queue.offerLast(RealMessage(x, self))
+    case SetIgnore(ign)      ⇒ ignore = ign
+    case x @ Watch(ref)      ⇒ context.watch(ref); queue.offerLast(RealMessage(x, self))
+    case x @ UnWatch(ref)    ⇒ context.unwatch(ref); queue.offerLast(RealMessage(x, self))
+    case SetAutoPilot(pilot) ⇒ autopilot = Some(pilot)
     case x: AnyRef ⇒
+      autopilot = autopilot.flatMap(_.run(sender, x))
       val observe = ignore map (ignoreFunc ⇒ if (ignoreFunc isDefinedAt x) !ignoreFunc(x) else true) getOrElse true
       if (observe) queue.offerLast(RealMessage(x, sender))
   }
@@ -148,6 +157,13 @@ class TestKit(_system: ActorSystem) {
     expectMsg(msg)
   }
 
+  /**
+   * Install an AutoPilot to drive the testActor: the AutoPilot will be run
+   * for each received message and can be used to send or forward messages,
+   * etc. Each invocation must return the AutoPilot for the next round.
+   */
+  def setAutoPilot(pilot: TestActor.AutoPilot): Unit = testActor ! TestActor.SetAutoPilot(pilot)
+
   /**
    * Obtain current time (`System.nanoTime`) as Duration.
    */
diff --git a/akka-testkit/src/test/scala/akka/testkit/TestProbeSpec.scala b/akka-testkit/src/test/scala/akka/testkit/TestProbeSpec.scala
index 516bba7322..b272544141 100644
--- a/akka-testkit/src/test/scala/akka/testkit/TestProbeSpec.scala
+++ b/akka-testkit/src/test/scala/akka/testkit/TestProbeSpec.scala
@@ -40,6 +40,31 @@ class TestProbeSpec extends AkkaSpec with DefaultTimeout {
       probe1.expectMsg(0 millis, "world")
     }
 
+    "have an AutoPilot" in {
+      //#autopilot
+      val probe = TestProbe()
+      probe.setAutoPilot(new TestActor.AutoPilot {
+        def run(sender: ActorRef, msg: Any): Option[TestActor.AutoPilot] =
+          msg match {
+            case "stop" ⇒ None
+            case x      ⇒ testActor.tell(x, sender); Some(this)
+          }
+      })
+      //#autopilot
+      probe.ref ! "hallo"
+      probe.ref ! "welt"
+      probe.ref ! "stop"
+      expectMsg("hallo")
+      expectMsg("welt")
+      probe.expectMsg("hallo")
+      probe.expectMsg("welt")
+      probe.expectMsg("stop")
+      probe.ref ! "hallo"
+      probe.expectMsg("hallo")
+      testActor ! "end"
+      expectMsg("end") // verify that "hallo" did not get through
+    }
+
   }
 
 }