Add overridden duration timeout to StreamTestKit

stream-testkit/src/main/scala/org/apache/pekko/stream/testkit/javadsl/StreamTestKit.scala

@@ -19,6 +19,10 @@ import pekko.stream.{ Materializer, SystemMaterializer }
 import pekko.stream.impl.PhasedFusingActorMaterializer
 import pekko.stream.testkit.scaladsl
 
+import java.time.Duration
+import java.util.concurrent.TimeUnit
+import scala.concurrent.duration.FiniteDuration
+
 object StreamTestKit {
 
   /**
@@ -29,7 +33,21 @@ object StreamTestKit {
   def assertAllStagesStopped(mat: Materializer): Unit =
     mat match {
       case impl: PhasedFusingActorMaterializer =>
-        scaladsl.StreamTestKit.assertNoChildren(impl.system, impl.supervisor)
+        scaladsl.StreamTestKit.assertNoChildren(impl.system, impl.supervisor, None)
+      case _ =>
+    }
+
+  /**
+   * Assert that there are no stages running under a given materializer.
+   * Usually this assertion is run after a test-case to check that all of the
+   * stages have terminated successfully with an overridden duration that ignores
+   * `stream.testkit.all-stages-stopped-timeout`.
+   */
+  def assertAllStagesStopped(mat: Materializer, duration: Duration): Unit =
+    mat match {
+      case impl: PhasedFusingActorMaterializer =>
+        scaladsl.StreamTestKit.assertNoChildren(impl.system, impl.supervisor,
+          Some(FiniteDuration(duration.toMillis, TimeUnit.MILLISECONDS)))
       case _ =>
     }
 

stream-testkit/src/main/scala/org/apache/pekko/stream/testkit/scaladsl/StreamTestKit.scala

@@ -35,12 +35,30 @@ object StreamTestKit {
    * This assertion is useful to check that all of the stages have
    * terminated successfully.
    */
+  def assertAllStagesStopped[T](block: => T, overrideTimeout: FiniteDuration)(implicit materializer: Materializer): T =
+    materializer match {
+      case impl: PhasedFusingActorMaterializer =>
+        stopAllChildren(impl.system, impl.supervisor)
+        val result = block
+        assertNoChildren(impl.system, impl.supervisor, Some(overrideTimeout))
+        result
+      case _ => block
+    }
+
+  /**
+   * Asserts that after the given code block is ran, no stages are left over
+   * that were created by the given materializer with an overridden duration
+   * that ignores `stream.testkit.all-stages-stopped-timeout`.
+   *
+   * This assertion is useful to check that all of the stages have
+   * terminated successfully.
+   */
   def assertAllStagesStopped[T](block: => T)(implicit materializer: Materializer): T =
     materializer match {
       case impl: PhasedFusingActorMaterializer =>
         stopAllChildren(impl.system, impl.supervisor)
         val result = block
-        assertNoChildren(impl.system, impl.supervisor)
+        assertNoChildren(impl.system, impl.supervisor, None)
         result
       case _ => block
     }
@@ -53,10 +71,16 @@ object StreamTestKit {
   }
 
   /** INTERNAL API */
-  @InternalApi private[testkit] def assertNoChildren(sys: ActorSystem, supervisor: ActorRef): Unit = {
+  @InternalApi private[testkit] def assertNoChildren(sys: ActorSystem, supervisor: ActorRef,
+      overrideTimeout: Option[FiniteDuration]): Unit = {
     val probe = TestProbe()(sys)
-    val c = sys.settings.config.getConfig("pekko.stream.testkit")
-    val timeout = c.getDuration("all-stages-stopped-timeout", MILLISECONDS).millis
+    val timeout = overrideTimeout match {
+      case Some(value) => value
+      case None =>
+        val c = sys.settings.config.getConfig("pekko.stream.testkit")
+        c.getDuration("all-stages-stopped-timeout", MILLISECONDS).millis
+    }
+
     probe.within(timeout) {
       try probe.awaitAssert {
           supervisor.tell(StreamSupervisor.GetChildren, probe.ref)

stream-testkit/src/test/scala/org/apache/pekko/stream/testkit/StreamSpec.scala

@@ -31,6 +31,7 @@ import org.scalatest.Failed
 import com.typesafe.config.{ Config, ConfigFactory }
 
 abstract class StreamSpec(_system: ActorSystem) extends PekkoSpec(_system) {
+  private var _allStagesStoppedTimeout: Option[FiniteDuration] = None
   def this(config: Config) =
     this(
       ActorSystem(
@@ -43,6 +44,31 @@ abstract class StreamSpec(_system: ActorSystem) extends PekkoSpec(_system) {
 
   def this() = this(ActorSystem(TestKitUtils.testNameFromCallStack(classOf[StreamSpec], "".r), PekkoSpec.testConf))
 
+  def this(config: Config, overrideAllStagesStoppedTimeout: FiniteDuration) = {
+    this(config)
+    _allStagesStoppedTimeout = Some(overrideAllStagesStoppedTimeout)
+  }
+
+  def this(s: String, overrideAllStagesStoppedTimeout: FiniteDuration) = {
+    this(s)
+    _allStagesStoppedTimeout = Some(overrideAllStagesStoppedTimeout)
+  }
+
+  def this(configMap: Map[String, _], overrideAllStagesStoppedTimeout: FiniteDuration) = {
+    this(configMap)
+    _allStagesStoppedTimeout = Some(overrideAllStagesStoppedTimeout)
+  }
+
+  def this(overrideAllStagesStoppedTimeout: FiniteDuration) = {
+    this()
+    _allStagesStoppedTimeout = Some(overrideAllStagesStoppedTimeout)
+  }
+
+  def this(_system: ActorSystem, overrideAllStagesStoppedTimeout: FiniteDuration) = {
+    this(_system)
+    _allStagesStoppedTimeout = Some(overrideAllStagesStoppedTimeout)
+  }
+
   override def withFixture(test: NoArgTest) = {
     super.withFixture(test) match {
       case failed: Failed =>
@@ -73,7 +99,7 @@ abstract class StreamSpec(_system: ActorSystem) extends PekkoSpec(_system) {
           case impl: PhasedFusingActorMaterializer =>
             stopAllChildren(impl.system, impl.supervisor)
             val result = test.apply()
-            assertNoChildren(impl.system, impl.supervisor)
+            assertNoChildren(impl.system, impl.supervisor, _allStagesStoppedTimeout)
             result
           case _ => other
         }