From 19769e2a87f611d4accadef01141e61cd96b140a Mon Sep 17 00:00:00 2001
From: Patrik Nordwall <patrik.nordwall@gmail.com>
Date: Fri, 21 Dec 2018 15:08:47 +0100
Subject: [PATCH] Cleanup Typed TestProbe implementation

* use a consistent api/impl technique
* _internal methods are not dilating
* add expectTerminated with default timeout
* FiniteDuration in awaitAssert
---
 .../typed/internal/TestProbeImpl.scala        | 162 +++++++++++++-----
 .../testkit/typed/javadsl/TestProbe.scala     | 145 +++++++---------
 .../testkit/typed/scaladsl/TestProbe.scala    |  88 +++++-----
 .../testkit/typed/scaladsl/package.scala      |   2 +-
 4 files changed, 226 insertions(+), 171 deletions(-)

diff --git a/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/internal/TestProbeImpl.scala b/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/internal/TestProbeImpl.scala
index 7fe9417f55..15325e9603 100644
--- a/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/internal/TestProbeImpl.scala
+++ b/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/internal/TestProbeImpl.scala
@@ -4,25 +4,37 @@
 
 package akka.actor.testkit.typed.internal
 
+import java.time.{ Duration ⇒ JDuration }
+import java.util.concurrent.BlockingDeque
+import java.util.concurrent.LinkedBlockingDeque
 import java.util.concurrent.atomic.AtomicInteger
-import java.util.concurrent.{ BlockingDeque, LinkedBlockingDeque }
 import java.util.function.Supplier
+import java.util.{ List ⇒ JList }
 
-import akka.actor.typed.scaladsl.Behaviors
-import akka.actor.typed.{ ActorRef, ActorSystem, Behavior, Terminated }
-import akka.annotation.InternalApi
-import akka.actor.testkit.typed.javadsl.{ TestProbe ⇒ JavaTestProbe }
-import akka.actor.testkit.typed.scaladsl.{ TestDuration, TestProbe ⇒ ScalaTestProbe }
-import akka.actor.testkit.typed.{ FishingOutcome, TestKitSettings }
-import akka.util.PrettyDuration._
-import akka.util.{ BoxedType, Timeout }
-import akka.util.JavaDurationConverters._
 import scala.annotation.tailrec
+import scala.collection.JavaConverters._
 import scala.collection.immutable
 import scala.concurrent.Await
 import scala.concurrent.duration._
+import scala.reflect.ClassTag
 import scala.util.control.NonFatal
 
+import akka.actor.testkit.typed.FishingOutcome
+import akka.actor.testkit.typed.TestKitSettings
+import akka.actor.testkit.typed.javadsl.{ TestProbe ⇒ JavaTestProbe }
+import akka.actor.testkit.typed.scaladsl.TestDuration
+import akka.actor.testkit.typed.scaladsl.{ TestProbe ⇒ ScalaTestProbe }
+import akka.actor.typed.ActorRef
+import akka.actor.typed.ActorSystem
+import akka.actor.typed.Behavior
+import akka.actor.typed.Terminated
+import akka.actor.typed.scaladsl.Behaviors
+import akka.annotation.InternalApi
+import akka.util.BoxedType
+import akka.util.JavaDurationConverters._
+import akka.util.PrettyDuration._
+import akka.util.Timeout
+
 @InternalApi
 private[akka] object TestProbeImpl {
   private val testActorId = new AtomicInteger(0)
@@ -53,7 +65,7 @@ private[akka] object TestProbeImpl {
 private[akka] final class TestProbeImpl[M](name: String, system: ActorSystem[_]) extends JavaTestProbe[M] with ScalaTestProbe[M] {
 
   import TestProbeImpl._
-  protected implicit val settings = TestKitSettings(system)
+  protected implicit val settings: TestKitSettings = TestKitSettings(system)
   private val queue = new LinkedBlockingDeque[M]
   private val terminations = new LinkedBlockingDeque[Terminated]
 
@@ -65,8 +77,6 @@ private[akka] final class TestProbeImpl[M](name: String, system: ActorSystem[_])
    */
   private var lastWasNoMessage = false
 
-  private var lastMessage: Option[M] = None
-
   private val testActor: ActorRef[M] = {
     // FIXME arbitrary timeout?
     implicit val timeout: Timeout = Timeout(3.seconds)
@@ -78,14 +88,14 @@ private[akka] final class TestProbeImpl[M](name: String, system: ActorSystem[_])
 
   override def remainingOrDefault: FiniteDuration = remainingOr(settings.SingleExpectDefaultTimeout.dilated)
 
-  override def getRemainingOrDefault: java.time.Duration = remainingOrDefault.asJava
+  override def getRemainingOrDefault: JDuration = remainingOrDefault.asJava
 
   override def remaining: FiniteDuration = end match {
     case f: FiniteDuration ⇒ f - now
     case _                 ⇒ assertFail("`remaining` may not be called outside of `within`")
   }
 
-  override def getRemaining: java.time.Duration = remaining.asJava
+  override def getRemaining: JDuration = remaining.asJava
 
   override def remainingOr(duration: FiniteDuration): FiniteDuration = end match {
     case x if x eq Duration.Undefined ⇒ duration
@@ -93,33 +103,38 @@ private[akka] final class TestProbeImpl[M](name: String, system: ActorSystem[_])
     case f: FiniteDuration            ⇒ f - now
   }
 
-  override def getRemainingOr(duration: java.time.Duration): java.time.Duration =
+  override def getRemainingOr(duration: JDuration): JDuration =
     remainingOr(duration.asScala).asJava
 
-  private def remainingOrDilated(max: Duration): FiniteDuration = max match {
-    case x if x eq Duration.Undefined ⇒ remainingOrDefault
-    case x if !x.isFinite             ⇒ throw new IllegalArgumentException("max duration cannot be infinite")
-    case f: FiniteDuration            ⇒ f.dilated
-  }
+  override def within[T](min: FiniteDuration, max: FiniteDuration)(f: ⇒ T): T =
+    within_internal(min, max.dilated, f)
 
-  override protected def within_internal[T](min: FiniteDuration, max: FiniteDuration, f: ⇒ T): T = {
-    val _max = max.dilated
+  override def within[T](max: FiniteDuration)(f: ⇒ T): T =
+    within_internal(Duration.Zero, max.dilated, f)
+
+  override def within[T](min: JDuration, max: JDuration)(f: Supplier[T]): T =
+    within_internal(min.asScala, max.asScala.dilated, f.get())
+
+  def within[T](max: JDuration)(f: Supplier[T]): T =
+    within_internal(Duration.Zero, max.asScala.dilated, f.get())
+
+  private def within_internal[T](min: FiniteDuration, max: FiniteDuration, f: ⇒ T): T = {
     val start = now
     val rem = if (end == Duration.Undefined) Duration.Inf else end - start
     assert(rem >= min, s"required min time $min not possible, only ${rem.pretty} left")
 
     lastWasNoMessage = false
 
-    val max_diff = _max min rem
-    val prev_end = end
-    end = start + max_diff
+    val maxDiff = max min rem
+    val prevEnd = end
+    end = start + maxDiff
 
-    val ret = try f finally end = prev_end
+    val ret = try f finally end = prevEnd
 
     val diff = now - start
     assert(min <= diff, s"block took ${diff.pretty}, should at least have been $min")
     if (!lastWasNoMessage) {
-      assert(diff <= max_diff, s"block took ${diff.pretty}, exceeding ${max_diff.pretty}")
+      assert(diff <= maxDiff, s"block took ${diff.pretty}, exceeding ${maxDiff.pretty}")
     }
 
     ret
@@ -129,13 +144,13 @@ private[akka] final class TestProbeImpl[M](name: String, system: ActorSystem[_])
 
   override def expectMessage[T <: M](max: FiniteDuration, obj: T): T = expectMessage_internal(max.dilated, obj)
 
-  override def expectMessage[T <: M](max: java.time.Duration, obj: T): T =
+  override def expectMessage[T <: M](max: JDuration, obj: T): T =
     expectMessage(max.asScala, obj)
 
   override def expectMessage[T <: M](max: FiniteDuration, hint: String, obj: T): T =
     expectMessage_internal(max.dilated, obj, Some(hint))
 
-  override def expectMessage[T <: M](max: java.time.Duration, hint: String, obj: T): T =
+  override def expectMessage[T <: M](max: JDuration, hint: String, obj: T): T =
     expectMessage(max.asScala, hint, obj)
 
   private def expectMessage_internal[T <: M](max: Duration, obj: T, hint: Option[String] = None): T = {
@@ -150,7 +165,7 @@ private[akka] final class TestProbeImpl[M](name: String, system: ActorSystem[_])
 
   override def receiveOne(): M = receiveOne(remainingOrDefault)
 
-  override def receiveOne(max: java.time.Duration): M = receiveOne(max.asScala)
+  override def receiveOne(max: JDuration): M = receiveOne(max.asScala)
 
   def receiveOne(max: FiniteDuration): M =
     receiveOne_internal(max.dilated).
@@ -173,14 +188,13 @@ private[akka] final class TestProbeImpl[M](name: String, system: ActorSystem[_])
       }
     )
     lastWasNoMessage = false
-    lastMessage = message
     message
   }
 
   override def expectNoMessage(max: FiniteDuration): Unit =
     expectNoMessage_internal(max)
 
-  override def expectNoMessage(max: java.time.Duration): Unit =
+  override def expectNoMessage(max: JDuration): Unit =
     expectNoMessage(max.asScala)
 
   override def expectNoMessage(): Unit =
@@ -194,7 +208,19 @@ private[akka] final class TestProbeImpl[M](name: String, system: ActorSystem[_])
     }
   }
 
-  override protected def expectMessageClass_internal[C](max: FiniteDuration, c: Class[C]): C = {
+  override def expectMessageType[T <: M](implicit t: ClassTag[T]): T =
+    expectMessageClass_internal(remainingOrDefault, t.runtimeClass.asInstanceOf[Class[T]])
+
+  override def expectMessageType[T <: M](max: FiniteDuration)(implicit t: ClassTag[T]): T =
+    expectMessageClass_internal(max.dilated, t.runtimeClass.asInstanceOf[Class[T]])
+
+  override def expectMessageClass[T <: M](clazz: Class[T]): T =
+    expectMessageClass_internal(getRemainingOrDefault.asScala, clazz)
+
+  override def expectMessageClass[T <: M](clazz: Class[T], max: JDuration): T =
+    expectMessageClass_internal(max.asScala.dilated, clazz)
+
+  private def expectMessageClass_internal[C](max: FiniteDuration, c: Class[C]): C = {
     val o = receiveOne_internal(max)
     val bt = BoxedType(c)
     o match {
@@ -204,7 +230,19 @@ private[akka] final class TestProbeImpl[M](name: String, system: ActorSystem[_])
     }
   }
 
-  override protected def receiveN_internal(n: Int, max: FiniteDuration): immutable.Seq[M] = {
+  override def receiveN(n: Int): immutable.Seq[M] =
+    receiveN_internal(n, remainingOrDefault)
+
+  override def receiveN(n: Int, max: FiniteDuration): immutable.Seq[M] =
+    receiveN_internal(n, max.dilated)
+
+  override def receiveMessages(n: Int): JList[M] =
+    receiveN_internal(n, getRemainingOrDefault.asScala).asJava
+
+  override def receiveMessages(n: Int, max: JDuration): JList[M] =
+    receiveN_internal(n, max.asScala.dilated).asJava
+
+  private def receiveN_internal(n: Int, max: FiniteDuration): immutable.Seq[M] = {
     val stop = max + now
     for (x ← 1 to n) yield {
       val timeout = stop - now
@@ -216,7 +254,19 @@ private[akka] final class TestProbeImpl[M](name: String, system: ActorSystem[_])
     }
   }
 
-  override protected def fishForMessage_internal(max: FiniteDuration, hint: String, fisher: M ⇒ FishingOutcome): List[M] = {
+  override def fishForMessage(max: FiniteDuration, hint: String)(fisher: M ⇒ FishingOutcome): immutable.Seq[M] =
+    fishForMessage_internal(max.dilated, hint, fisher)
+
+  override def fishForMessage(max: FiniteDuration)(fisher: M ⇒ FishingOutcome): immutable.Seq[M] =
+    fishForMessage(max, "")(fisher)
+
+  override def fishForMessage(max: JDuration, fisher: java.util.function.Function[M, FishingOutcome]): JList[M] =
+    fishForMessage(max, "", fisher)
+
+  override def fishForMessage(max: JDuration, hint: String, fisher: java.util.function.Function[M, FishingOutcome]): JList[M] =
+    fishForMessage_internal(max.asScala.dilated, hint, fisher.apply).asJava
+
+  private def fishForMessage_internal(max: FiniteDuration, hint: String, fisher: M ⇒ FishingOutcome): List[M] = {
     @tailrec def loop(timeout: FiniteDuration, seen: List[M]): List[M] = {
       val start = System.nanoTime()
       val maybeMsg = receiveOne_internal(timeout)
@@ -243,10 +293,19 @@ private[akka] final class TestProbeImpl[M](name: String, system: ActorSystem[_])
       }
     }
 
-    loop(max.dilated, Nil)
+    loop(max, Nil)
   }
 
-  override def expectTerminated[U](actorRef: ActorRef[U], max: FiniteDuration): Unit = {
+  override def expectTerminated[U](actorRef: ActorRef[U], max: FiniteDuration): Unit =
+    expectTerminated_internal(actorRef, max.dilated)
+
+  override def expectTerminated[U](actorRef: ActorRef[U]): Unit =
+    expectTerminated_internal(actorRef, remainingOrDefault)
+
+  override def expectTerminated[U](actorRef: ActorRef[U], max: JDuration): Unit =
+    expectTerminated_internal(actorRef, max.asScala.dilated)
+
+  private def expectTerminated_internal[U](actorRef: ActorRef[U], max: FiniteDuration): Unit = {
     testActor.asInstanceOf[ActorRef[AnyRef]] ! WatchActor(actorRef)
     val message =
       if (max == Duration.Zero) {
@@ -260,15 +319,26 @@ private[akka] final class TestProbeImpl[M](name: String, system: ActorSystem[_])
     assert(message.ref == actorRef, s"expected [${actorRef.path}] to stop, but saw [${message.ref.path}] stop")
   }
 
-  override def expectTerminated[U](actorRef: ActorRef[U], max: java.time.Duration): Unit =
-    expectTerminated(actorRef, max.asScala)
+  override def awaitAssert[A](a: ⇒ A, max: FiniteDuration, interval: FiniteDuration): A =
+    awaitAssert_internal(a, max.dilated, interval)
 
-  override def awaitAssert[A](max: java.time.Duration, interval: java.time.Duration, supplier: Supplier[A]): A =
-    awaitAssert(supplier.get(), if (max == java.time.Duration.ZERO) Duration.Undefined else max.asScala, interval.asScala)
+  override def awaitAssert[A](a: ⇒ A, max: FiniteDuration): A =
+    awaitAssert_internal(a, max.dilated, 100.millis)
 
-  override def awaitAssert[A](a: ⇒ A, max: Duration = Duration.Undefined, interval: Duration = 100.millis): A = {
-    val _max = remainingOrDilated(max)
-    val stop = now + _max
+  override def awaitAssert[A](a: ⇒ A): A =
+    awaitAssert_internal(a, remainingOrDefault, 100.millis)
+
+  override def awaitAssert[A](max: JDuration, interval: JDuration, supplier: Supplier[A]): A =
+    awaitAssert_internal(supplier.get(), max.asScala.dilated, interval.asScala)
+
+  def awaitAssert[A](max: JDuration, supplier: Supplier[A]): A =
+    awaitAssert(max, JDuration.ofMillis(100), supplier)
+
+  def awaitAssert[A](supplier: Supplier[A]): A =
+    awaitAssert(getRemainingOrDefault, supplier)
+
+  private def awaitAssert_internal[A](a: ⇒ A, max: FiniteDuration, interval: FiniteDuration): A = {
+    val stop = now + max
 
     @tailrec
     def poll(t: Duration): A = {
@@ -294,7 +364,7 @@ private[akka] final class TestProbeImpl[M](name: String, system: ActorSystem[_])
       }
     }
 
-    poll(_max min interval)
+    poll(max min interval)
   }
 
   /**
diff --git a/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/javadsl/TestProbe.scala b/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/javadsl/TestProbe.scala
index 5513841e5f..727d098003 100644
--- a/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/javadsl/TestProbe.scala
+++ b/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/javadsl/TestProbe.scala
@@ -8,19 +8,14 @@ import java.time.Duration
 import java.util.function.Supplier
 import java.util.{ List ⇒ JList }
 
-import akka.actor.typed.{ ActorRef, ActorSystem }
-import akka.annotation.DoNotInherit
-import akka.annotation.InternalApi
+import akka.actor.testkit.typed.FishingOutcome
+import akka.actor.testkit.typed.TestKitSettings
 import akka.actor.testkit.typed.internal.TestProbeImpl
-import akka.actor.testkit.typed.{ FishingOutcome, TestKitSettings }
-import akka.actor.testkit.typed.scaladsl.TestDuration
-import akka.util.JavaDurationConverters._
+import akka.actor.typed.ActorRef
+import akka.actor.typed.ActorSystem
+import akka.annotation.DoNotInherit
 import akka.util.unused
 
-import scala.collection.immutable
-import scala.collection.JavaConverters._
-import scala.concurrent.duration.FiniteDuration
-
 object FishingOutcomes {
   /**
    * Consume this message and continue with the next
@@ -108,8 +103,8 @@ abstract class TestProbe[M] {
    * take maximum wait times are available in a version which implicitly uses
    * the remaining time governed by the innermost enclosing `within` block.
    *
-   * Note that the timeout is scaled using Duration.dilated, which uses the
-   * configuration entry "akka.actor.testkit.typed.timefactor", while the min Duration is not.
+   * Note that the max timeout is scaled using the configuration entry "akka.actor.testkit.typed.timefactor",
+   * while the min Duration is not.
    *
    * {{{
    * val ret = within(50 millis) {
@@ -118,19 +113,12 @@ abstract class TestProbe[M] {
    * }
    * }}}
    */
-  def within[T](min: Duration, max: Duration)(f: Supplier[T]): T =
-    within_internal(min.asScala, max.asScala, f.get())
+  def within[T](min: Duration, max: Duration)(f: Supplier[T]): T
 
   /**
    * Same as calling `within(0 seconds, max)(f)`.
    */
-  def within[T](max: Duration)(f: Supplier[T]): T =
-    within_internal(scala.concurrent.duration.Duration.Zero, max.asScala, f.get())
-
-  /**
-   * INTERNAL API
-   */
-  @InternalApi protected def within_internal[T](min: FiniteDuration, max: FiniteDuration, f: ⇒ T): T
+  def within[T](max: Duration)(f: Supplier[T]): T
 
   /**
    * Same as `expectMessage(remainingOrDefault, obj)`, but using the
@@ -163,65 +151,23 @@ abstract class TestProbe[M] {
   def expectNoMessage(max: Duration): Unit
 
   /**
-   * Assert that no message is received. Waits for the default period configured as `akka.actor.testkit.typed.expect-no-message-default`
-   * That value is dilated.
+   * Assert that no message is received. Waits for the default period configured as `akka.actor.testkit.typed.expect-no-message-default`.
+   * That timeout is scaled using the configuration entry "akka.actor.testkit.typed.timefactor".
    */
   def expectNoMessage(): Unit
 
-  /**
-   * Expect the given actor to be stopped or stop within the given timeout or
-   * throw an [[AssertionError]].
-   */
-  def expectTerminated[U](actorRef: ActorRef[U], max: Duration): Unit
-
-  /**
-   * Evaluate the given assert every `interval` until it does not throw an exception and return the
-   * result.
-   *
-   * If the `max` timeout expires the last exception is thrown.
-   *
-   * Note that the timeout is scaled using Duration.dilated, which uses the configuration entry "akka.test.timefactor".
-   */
-  def awaitAssert[A](max: Duration, interval: Duration, supplier: Supplier[A]): A
-
-  /**
-   * Evaluate the given assert every 100 milliseconds until it does not throw an exception and return the
-   * result.
-   *
-   * If the `max` timeout expires the last exception is thrown.
-   *
-   * Note that the timeout is scaled using Duration.dilated, which uses the configuration entry "akka.test.timefactor".
-   */
-  def awaitAssert[A](max: Duration, supplier: Supplier[A]): A =
-    awaitAssert(max, Duration.ofMillis(100), supplier)
-
-  /**
-   * Evaluate the given assert every 100 milliseconds until it does not throw an exception and return the
-   * result. A max time is taken it from the innermost enclosing `within` block.
-   */
-  def awaitAssert[A](supplier: Supplier[A]): A =
-    awaitAssert(Duration.ZERO, supplier)
-
-  // FIXME awaitAssert(Procedure): Unit would be nice for java people to not have to return null
-
   /**
    * Same as `expectMessageType(clazz, remainingOrDefault)`,but using the
    * default timeout as deadline.
    */
-  def expectMessageClass[T <: M](clazz: Class[T]): T =
-    expectMessageClass_internal(getRemainingOrDefault.asScala, clazz)
+  def expectMessageClass[T <: M](clazz: Class[T]): T
 
   /**
    * Wait for a message of type M and return it when it arrives, or fail if the `max` timeout is hit.
-   * The timeout is dilated.
+   *
+   * Note that the timeout is scaled using the configuration entry "akka.actor.testkit.typed.timefactor".
    */
-  def expectMessageClass[T <: M](clazz: Class[T], max: Duration): T =
-    expectMessageClass_internal(max.asScala.dilated, clazz)
-
-  /**
-   * INTERNAL API
-   */
-  @InternalApi protected def expectMessageClass_internal[C](max: FiniteDuration, c: Class[C]): C
+  def expectMessageClass[T <: M](clazz: Class[T], max: Duration): T
 
   /**
    * Receive one message of type `M` within the default timeout as deadline.
@@ -237,17 +183,14 @@ abstract class TestProbe[M] {
   /**
    * Same as `receiveMessages(n, remaining)` but using the default timeout as deadline.
    */
-  def receiveMessages(n: Int): JList[M] = receiveN_internal(n, getRemainingOrDefault.asScala).asJava
+  def receiveMessages(n: Int): JList[M]
 
   /**
    * Receive `n` messages in a row before the given deadline.
+   *
+   * Note that the timeout is scaled using the configuration entry "akka.actor.testkit.typed.timefactor".
    */
-  def receiveMessages(n: Int, max: Duration): JList[M] = receiveN_internal(n, max.asScala.dilated).asJava
-
-  /**
-   * INTERNAL API
-   */
-  @InternalApi protected def receiveN_internal(n: Int, max: FiniteDuration): immutable.Seq[M]
+  def receiveMessages(n: Int, max: Duration): JList[M]
 
   /**
    * Java API: Allows for flexible matching of multiple messages within a timeout, the fisher function is fed each incoming
@@ -262,23 +205,57 @@ abstract class TestProbe[M] {
    * is decorated with some fishing details and the test is failed (making it convenient to use this method with a
    * partial function).
    *
-   * @param max Max total time without the fisher function returning `CompleteFishing` before failing
-   *            The timeout is dilated.
+   * @param max Max total time without the fisher function returning `CompleteFishing` before failing.
+   *            The timeout is scaled using the configuration entry "akka.actor.testkit.typed.timefactor".
    * @return The messages accepted in the order they arrived
    */
-  def fishForMessage(max: Duration, fisher: java.util.function.Function[M, FishingOutcome]): java.util.List[M] =
-    fishForMessage(max, "", fisher)
+  def fishForMessage(max: Duration, fisher: java.util.function.Function[M, FishingOutcome]): java.util.List[M]
 
   /**
    * Same as the other `fishForMessage` but includes the provided hint in all error messages
    */
-  def fishForMessage(max: Duration, hint: String, fisher: java.util.function.Function[M, FishingOutcome]): java.util.List[M] =
-    fishForMessage_internal(max.asScala, hint, fisher.apply).asJava
+  def fishForMessage(max: Duration, hint: String, fisher: java.util.function.Function[M, FishingOutcome]): java.util.List[M]
 
   /**
-   * INTERNAL API
+   * Expect the given actor to be stopped or stop within the given timeout or
+   * throw an [[AssertionError]].
+   *
+   * Note that the timeout is scaled using the configuration entry "akka.actor.testkit.typed.timefactor".
    */
-  @InternalApi protected def fishForMessage_internal(max: FiniteDuration, hint: String, fisher: M ⇒ FishingOutcome): List[M]
+  def expectTerminated[U](actorRef: ActorRef[U], max: Duration): Unit
+
+  /**
+   * Expect the given actor to be stopped or stop within the default timeout.
+   */
+  def expectTerminated[U](actorRef: ActorRef[U]): Unit
+
+  /**
+   * Evaluate the given assert every `interval` until it does not throw an exception and return the
+   * result.
+   *
+   * If the `max` timeout expires the last exception is thrown.
+   *
+   * Note that the timeout is scaled using the configuration entry "akka.actor.testkit.typed.timefactor".
+   */
+  def awaitAssert[A](max: Duration, interval: Duration, supplier: Supplier[A]): A
+
+  /**
+   * Evaluate the given assert every 100 milliseconds until it does not throw an exception and return the
+   * result.
+   *
+   * If the `max` timeout expires the last exception is thrown.
+   *
+   * Note that the timeout is scaled using the configuration entry "akka.actor.testkit.typed.timefactor".
+   */
+  def awaitAssert[A](max: Duration, supplier: Supplier[A]): A
+
+  /**
+   * Evaluate the given assert every 100 milliseconds until it does not throw an exception and return the
+   * result. A max time is taken it from the innermost enclosing `within` block.
+   */
+  def awaitAssert[A](supplier: Supplier[A]): A
+
+  // FIXME awaitAssert(Procedure): Unit would be nice for java people to not have to return null
 
   /**
    * Stops the [[TestProbe.getRef]], which is useful when testing watch and termination.
diff --git a/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/scaladsl/TestProbe.scala b/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/scaladsl/TestProbe.scala
index 7d2acee0c6..33129e60a0 100644
--- a/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/scaladsl/TestProbe.scala
+++ b/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/scaladsl/TestProbe.scala
@@ -4,14 +4,16 @@
 
 package akka.actor.testkit.typed.scaladsl
 
-import akka.actor.typed.{ ActorRef, ActorSystem }
-import akka.annotation.DoNotInherit
-import akka.actor.testkit.typed.internal.TestProbeImpl
-import akka.actor.testkit.typed.{ FishingOutcome, TestKitSettings }
-
+import scala.collection.immutable
 import scala.concurrent.duration._
 import scala.reflect.ClassTag
-import scala.collection.immutable
+
+import akka.actor.testkit.typed.FishingOutcome
+import akka.actor.testkit.typed.TestKitSettings
+import akka.actor.testkit.typed.internal.TestProbeImpl
+import akka.actor.typed.ActorRef
+import akka.actor.typed.ActorSystem
+import akka.annotation.DoNotInherit
 
 object FishingOutcomes {
   /**
@@ -84,8 +86,8 @@ object TestProbe {
    * take maximum wait times are available in a version which implicitly uses
    * the remaining time governed by the innermost enclosing `within` block.
    *
-   * Note that the timeout is scaled using Duration.dilated, which uses the
-   * configuration entry "akka.actor.testkit.typed.timefactor", while the min Duration is not.
+   * Note that the max timeout is scaled using the configuration entry "akka.actor.testkit.typed.timefactor",
+   * while the min Duration is not.
    *
    * {{{
    * val ret = within(50 millis) {
@@ -94,15 +96,12 @@ object TestProbe {
    * }
    * }}}
    */
-  def within[T](min: FiniteDuration, max: FiniteDuration)(f: ⇒ T): T =
-    within_internal(min, max, f)
+  def within[T](min: FiniteDuration, max: FiniteDuration)(f: ⇒ T): T
+
   /**
    * Same as calling `within(0 seconds, max)(f)`.
    */
-  def within[T](max: FiniteDuration)(f: ⇒ T): T =
-    within_internal(Duration.Zero, max, f)
-
-  protected def within_internal[T](min: FiniteDuration, max: FiniteDuration, f: ⇒ T): T
+  def within[T](max: FiniteDuration)(f: ⇒ T): T
 
   /**
    * Same as `expectMessage(remainingOrDefault, obj)`, but using the default timeout as deadline.
@@ -134,24 +133,20 @@ object TestProbe {
   def expectNoMessage(max: FiniteDuration): Unit
 
   /**
-   * Assert that no message is received. Waits for the default period configured as `akka.actor.testkit.typed.expect-no-message-default`
-   * That value is dilated.
+   * Assert that no message is received. Waits for the default period configured as `akka.actor.testkit.typed.expect-no-message-default`.
+   * That timeout is scaled using the configuration entry "akka.actor.testkit.typed.timefactor".
    */
   def expectNoMessage(): Unit
 
   /**
    * Same as `expectMessageType[T](remainingOrDefault)`, but using the default timeout as deadline.
    */
-  def expectMessageType[T <: M](implicit t: ClassTag[T]): T =
-    expectMessageClass_internal(remainingOrDefault, t.runtimeClass.asInstanceOf[Class[T]])
+  def expectMessageType[T <: M](implicit t: ClassTag[T]): T
 
   /**
    * Expect a message of type T to arrive within `max` or fail. `max` is dilated.
    */
-  def expectMessageType[T <: M](max: FiniteDuration)(implicit t: ClassTag[T]): T =
-    expectMessageClass_internal(max.dilated, t.runtimeClass.asInstanceOf[Class[T]])
-
-  protected def expectMessageClass_internal[C](max: FiniteDuration, c: Class[C]): C
+  def expectMessageType[T <: M](max: FiniteDuration)(implicit t: ClassTag[T]): T
 
   /**
    * Receive one message of type `M` within the default timeout as deadline.
@@ -167,14 +162,14 @@ object TestProbe {
   /**
    * Same as `receiveN(n, remaining)` but using the default timeout as deadline.
    */
-  def receiveN(n: Int): immutable.Seq[M] = receiveN_internal(n, remainingOrDefault)
+  def receiveN(n: Int): immutable.Seq[M]
 
   /**
    * Receive `n` messages in a row before the given deadline.
+   *
+   * Note that the timeout is scaled using the configuration entry "akka.actor.testkit.typed.timefactor".
    */
-  def receiveN(n: Int, max: FiniteDuration): immutable.Seq[M] = receiveN_internal(n, max.dilated)
-
-  protected def receiveN_internal(n: Int, max: FiniteDuration): immutable.Seq[M]
+  def receiveN(n: Int, max: FiniteDuration): immutable.Seq[M]
 
   /**
    * Allows for flexible matching of multiple messages within a timeout, the fisher function is fed each incoming
@@ -190,20 +185,16 @@ object TestProbe {
    * is decorated with some fishing details and the test is failed (making it convenient to use this method with a
    * partial function).
    *
-   * @param max Max total time without the fisher function returning `CompleteFishing` before failing
-   *            The timeout is dilated.
+   * @param max Max total time without the fisher function returning `CompleteFishing` before failing.
+   *            The timeout is scaled using the configuration entry "akka.actor.testkit.typed.timefactor".
    * @return The messages accepted in the order they arrived
    */
-  def fishForMessage(max: FiniteDuration, hint: String)(fisher: M ⇒ FishingOutcome): immutable.Seq[M] =
-    fishForMessage_internal(max, hint, fisher)
+  def fishForMessage(max: FiniteDuration, hint: String)(fisher: M ⇒ FishingOutcome): immutable.Seq[M]
 
   /**
    * Same as the other `fishForMessage` but with no hint
    */
-  def fishForMessage(max: FiniteDuration)(fisher: M ⇒ FishingOutcome): immutable.Seq[M] =
-    fishForMessage(max, "")(fisher)
-
-  protected def fishForMessage_internal(max: FiniteDuration, hint: String, fisher: M ⇒ FishingOutcome): immutable.Seq[M]
+  def fishForMessage(max: FiniteDuration)(fisher: M ⇒ FishingOutcome): immutable.Seq[M]
 
   /**
    * Expect the given actor to be stopped or stop within the given timeout or
@@ -211,19 +202,36 @@ object TestProbe {
    */
   def expectTerminated[U](actorRef: ActorRef[U], max: FiniteDuration): Unit
 
+  /**
+   * Expect the given actor to be stopped or stop within the default timeout.
+   */
+  def expectTerminated[U](actorRef: ActorRef[U]): Unit
+
   /**
    * Evaluate the given assert every `interval` until it does not throw an exception and return the
    * result.
    *
    * If the `max` timeout expires the last exception is thrown.
    *
-   * If no timeout is given, take it from the innermost enclosing `within`
-   * block.
-   *
-   * Note that the timeout is scaled using Duration.dilated,
-   * which uses the configuration entry "akka.test.timefactor".
+   * Note that the timeout is scaled using the configuration entry "akka.actor.testkit.typed.timefactor".
    */
-  def awaitAssert[A](a: ⇒ A, max: Duration = Duration.Undefined, interval: Duration = 100.millis): A
+  def awaitAssert[A](a: ⇒ A, max: FiniteDuration, interval: FiniteDuration): A
+
+  /**
+   * Evaluate the given assert every 100 ms until it does not throw an exception and return the
+   * result.
+   *
+   * If the `max` timeout expires the last exception is thrown.
+   */
+  def awaitAssert[A](a: ⇒ A, max: FiniteDuration): A
+
+  /**
+   * Evaluate the given assert every 100 ms until it does not throw an exception and return the
+   * result.
+   *
+   * If the default timeout expires the last exception is thrown.
+   */
+  def awaitAssert[A](a: ⇒ A): A
 
   /**
    * Stops the [[TestProbe.ref]], which is useful when testing watch and termination.
diff --git a/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/scaladsl/package.scala b/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/scaladsl/package.scala
index a0d869c72a..cd1a878a09 100644
--- a/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/scaladsl/package.scala
+++ b/akka-actor-testkit-typed/src/main/scala/akka/actor/testkit/typed/scaladsl/package.scala
@@ -10,7 +10,7 @@ package object scaladsl {
 
   /**
    * Scala API. Scale timeouts (durations) during tests with the configured
-   * 'akka.test.timefactor'.
+   * 'akka.actor.testkit.typed.timefactor'.
    * Implicit class providing `dilated` method.
    *
    * {{{