From cdf2bc24d41430d0e727e3d9a1ef802774065cc0 Mon Sep 17 00:00:00 2001
From: Marcin Kubala <marcin@mkubala.pl>
Date: Thu, 5 Jun 2014 23:30:11 +0200
Subject: [PATCH] +tes #15132 Add additional overloaded expectMsg to TestKit

---
 .../src/main/scala/akka/testkit/TestKit.scala | 16 +++++++--
 .../scala/akka/testkit/TestProbeSpec.scala    | 36 ++++++++++++++++++-
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/akka-testkit/src/main/scala/akka/testkit/TestKit.scala b/akka-testkit/src/main/scala/akka/testkit/TestKit.scala
index fb9338c1a0..ac6730a632 100644
--- a/akka-testkit/src/main/scala/akka/testkit/TestKit.scala
+++ b/akka-testkit/src/main/scala/akka/testkit/TestKit.scala
@@ -332,10 +332,20 @@ trait TestKitBase {
    */
   def expectMsg[T](max: FiniteDuration, obj: T): T = expectMsg_internal(max.dilated, obj)
 
-  private def expectMsg_internal[T](max: Duration, obj: T): T = {
+  /**
+   * Receive one message from the test actor and assert that it equals the
+   * given object. Wait time is bounded by the given duration, with an
+   * AssertionFailure being thrown in case of timeout.
+   *
+   * @return the received object
+   */
+  def expectMsg[T](max: FiniteDuration, hint: String, obj: T): T = expectMsg_internal(max.dilated, obj, Some(hint))
+
+  private def expectMsg_internal[T](max: Duration, obj: T, hint: Option[String] = None): T = {
     val o = receiveOne(max)
-    assert(o ne null, s"timeout ($max) during expectMsg while waiting for $obj")
-    assert(obj == o, s"expected $obj, found $o")
+    lazy val hintOrEmptyString = hint.map(": " + _).getOrElse("")
+    assert(o ne null, s"timeout ($max) during expectMsg while waiting for $obj" + hintOrEmptyString)
+    assert(obj == o, s"expected $obj, found $o" + hintOrEmptyString)
     o.asInstanceOf[T]
   }
 
diff --git a/akka-testkit/src/test/scala/akka/testkit/TestProbeSpec.scala b/akka-testkit/src/test/scala/akka/testkit/TestProbeSpec.scala
index 1553f90773..5db387be49 100644
--- a/akka-testkit/src/test/scala/akka/testkit/TestProbeSpec.scala
+++ b/akka-testkit/src/test/scala/akka/testkit/TestProbeSpec.scala
@@ -9,6 +9,7 @@ import akka.actor._
 import scala.concurrent.{ Future, Await }
 import scala.concurrent.duration._
 import akka.pattern.ask
+import scala.util.Try
 
 @org.junit.runner.RunWith(classOf[org.scalatest.junit.JUnitRunner])
 class TestProbeSpec extends AkkaSpec with DefaultTimeout {
@@ -39,7 +40,40 @@ class TestProbeSpec extends AkkaSpec with DefaultTimeout {
       probe1.send(probe2.ref, "hello")
       probe2.expectMsg(0 millis, "hello")
       probe2.lastMessage.sender ! "world"
-      probe1.expectMsg(0 millis, "world")
+      probe1.expectMsg(0 millis, "some hint here", "world")
+    }
+
+    def assertFailureMessageContains(expectedHint: String)(block: ⇒ Unit) {
+      Try {
+        block
+      } match {
+        case scala.util.Failure(e: AssertionError) ⇒
+          if (!(e.getMessage contains expectedHint))
+            fail(s"failure message did not contain hint! Was: ${e.getMessage}, expected to contain $expectedHint")
+        case scala.util.Failure(oth) ⇒
+          fail(s"expected AssertionError but got: $oth")
+        case scala.util.Success(result) ⇒
+          fail(s"expected failure but got: $result")
+      }
+    }
+
+    "throw AssertionError containing hint in its message if max await time is exceeded" in {
+      val probe = TestProbe()
+      val hint = "some hint"
+
+      assertFailureMessageContains(hint) {
+        probe.expectMsg(0 millis, hint, "hello")
+      }
+    }
+
+    "throw AssertionError containing hint in its message if received message doesn't match" in {
+      val probe = TestProbe()
+      val hint = "some hint"
+
+      assertFailureMessageContains(hint) {
+        probe.ref ! "hello"
+        probe.expectMsg(0 millis, hint, "bye")
+      }
     }
 
     "have an AutoPilot" in {