Expose assertAllStagesStopped

2018-06-15 18:16:58 +03:00 · 2018-06-15 18:16:58 +03:00 · 766944a592
commit 766944a592
parent 28746a4cfe
107 changed files with 263 additions and 85 deletions
--- a/akka-stream-testkit/src/main/scala/akka/stream/testkit/javadsl/StreamTestKit.scala
+++ b/akka-stream-testkit/src/main/scala/akka/stream/testkit/javadsl/StreamTestKit.scala
@ -0,0 +1,24 @@
+/*
+ * Copyright (C) 2018 Lightbend Inc. <https://www.lightbend.com>
+ */
+
+package akka.stream.testkit.javadsl
+
+import akka.stream.Materializer
+import akka.stream.impl.PhasedFusingActorMaterializer
+import akka.stream.testkit.scaladsl
+
+object StreamTestKit {
+
+  /**
+   * 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.
+   */
+  def assertAllStagesStopped(mat: Materializer): Unit =
+    mat match {
+      case impl: PhasedFusingActorMaterializer ⇒
+        scaladsl.StreamTestKit.assertNoChildren(impl.system, impl.supervisor)
+      case _ ⇒
+    }
+}
--- a/akka-stream-testkit/src/main/scala/akka/stream/testkit/scaladsl/StreamTestKit.scala
+++ b/akka-stream-testkit/src/main/scala/akka/stream/testkit/scaladsl/StreamTestKit.scala
@ -0,0 +1,62 @@
+/*
+ * Copyright (C) 2018 Lightbend Inc. <https://www.lightbend.com>
+ */
+
+package akka.stream.testkit.scaladsl
+
+import akka.actor.{ ActorRef, ActorSystem }
+import akka.annotation.InternalApi
+import akka.stream.Materializer
+import akka.stream.impl.{ PhasedFusingActorMaterializer, StreamSupervisor }
+import akka.testkit.TestProbe
+
+import scala.concurrent.duration._
+
+object StreamTestKit {
+
+  /**
+   * Asserts that after the given code block is ran, no stages are left over
+   * that were created by the given materializer.
+   *
+   * 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)
+        result
+      case _ ⇒ block
+    }
+
+  /** INTERNAL API */
+  @InternalApi private[testkit] def stopAllChildren(sys: ActorSystem, supervisor: ActorRef): Unit = {
+    val probe = TestProbe()(sys)
+    probe.send(supervisor, StreamSupervisor.StopChildren)
+    probe.expectMsg(StreamSupervisor.StoppedChildren)
+  }
+
+  /** INTERNAL API */
+  @InternalApi private[testkit] def assertNoChildren(sys: ActorSystem, supervisor: ActorRef): Unit = {
+    val probe = TestProbe()(sys)
+    probe.within(5.seconds) {
+      var children = Set.empty[ActorRef]
+      try probe.awaitAssert {
+        supervisor.tell(StreamSupervisor.GetChildren, probe.ref)
+        children = probe.expectMsgType[StreamSupervisor.Children].children
+        assert(
+          children.isEmpty,
+          s"expected no StreamSupervisor children, but got [${children.mkString(", ")}]")
+      }
+      catch {
+        case ex: Throwable ⇒
+          children.foreach(_ ! StreamSupervisor.PrintDebugDump)
+          throw ex
+      }
+    }
+  }
+
+}
+
--- a/akka-stream-testkit/src/main/scala/akka/stream/testkit/scaladsl/TestSink.scala
+++ b/akka-stream-testkit/src/main/scala/akka/stream/testkit/scaladsl/TestSink.scala
@ -10,6 +10,7 @@ import akka.stream._
 import akka.stream.scaladsl._
 import akka.stream.testkit.TestSubscriber.Probe
 import akka.stream.testkit._
+import akka.stream.testkit.StreamTestKit.ProbeSink

 /**
 * Factory methods for test sinks.
@ -20,6 +21,6 @@ object TestSink {
   * A Sink that materialized to a [[akka.stream.testkit.TestSubscriber.Probe]].
   */
  def probe[T](implicit system: ActorSystem): Sink[T, Probe[T]] =
-    Sink.fromGraph[T, TestSubscriber.Probe[T]](new StreamTestKit.ProbeSink(none, SinkShape(Inlet("ProbeSink.in"))))
+    Sink.fromGraph[T, TestSubscriber.Probe[T]](new ProbeSink(none, SinkShape(Inlet("ProbeSink.in"))))

 }
--- a/akka-stream-testkit/src/main/scala/akka/stream/testkit/scaladsl/TestSource.scala
+++ b/akka-stream-testkit/src/main/scala/akka/stream/testkit/scaladsl/TestSource.scala
@ -8,6 +8,7 @@ import akka.stream._
 import akka.stream.Attributes.none
 import akka.stream.scaladsl._
 import akka.stream.testkit._
+import akka.stream.testkit.StreamTestKit.ProbeSource

 import akka.actor.ActorSystem

@ -19,6 +20,6 @@ object TestSource {
  /**
   * A Source that materializes to a [[akka.stream.testkit.TestPublisher.Probe]].
   */
-  def probe[T](implicit system: ActorSystem) = Source.fromGraph[T, TestPublisher.Probe[T]](new StreamTestKit.ProbeSource(none, SourceShape(Outlet("ProbeSource.out"))))
+  def probe[T](implicit system: ActorSystem) = Source.fromGraph[T, TestPublisher.Probe[T]](new ProbeSource(none, SourceShape(Outlet("ProbeSource.out"))))

 }
--- a/akka-stream-testkit/src/test/scala/akka/stream/testkit/BaseTwoStreamsSetup.scala
+++ b/akka-stream-testkit/src/test/scala/akka/stream/testkit/BaseTwoStreamsSetup.scala
@ -9,7 +9,7 @@ import akka.stream.scaladsl._
 import org.reactivestreams.Publisher
 import scala.collection.immutable
 import scala.util.control.NoStackTrace
-import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.testkit.AkkaSpec

 abstract class BaseTwoStreamsSetup extends AkkaSpec {
--- a/akka-stream-testkit/src/test/scala/akka/stream/testkit/TestPublisherSubscriberSpec.scala
+++ b/akka-stream-testkit/src/test/scala/akka/stream/testkit/TestPublisherSubscriberSpec.scala
@ -7,7 +7,7 @@ package akka.stream.testkit
 import akka.stream.scaladsl.{ Sink, Source }
 import akka.stream.testkit.TestPublisher._
 import akka.stream.testkit.TestSubscriber._
-import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.{ ActorMaterializer, ActorMaterializerSettings }
 import org.reactivestreams.Subscription
 import akka.testkit.AkkaSpec
--- a/akka-stream-testkit/src/test/scala/akka/stream/testkit/Utils.scala
+++ b/akka-stream-testkit/src/test/scala/akka/stream/testkit/Utils.scala
@ -20,32 +20,6 @@ object Utils {

  case class TE(message: String) extends RuntimeException(message) with NoStackTrace

-  def assertAllStagesStopped[T](block: ⇒ T)(implicit materializer: Materializer): T =
-    materializer match {
-      case impl: PhasedFusingActorMaterializer ⇒
-        val probe = TestProbe()(impl.system)
-        probe.send(impl.supervisor, StreamSupervisor.StopChildren)
-        probe.expectMsg(StreamSupervisor.StoppedChildren)
-        val result = block
-        probe.within(5.seconds) {
-          var children = Set.empty[ActorRef]
-          try probe.awaitAssert {
-            impl.supervisor.tell(StreamSupervisor.GetChildren, probe.ref)
-            children = probe.expectMsgType[StreamSupervisor.Children].children
-            assert(
-              children.isEmpty,
-              s"expected no StreamSupervisor children, but got [${children.mkString(", ")}]")
-          }
-          catch {
-            case ex: Throwable ⇒
-              children.foreach(_ ! StreamSupervisor.PrintDebugDump)
-              throw ex
-          }
-        }
-        result
-      case _ ⇒ block
-    }
-
  def assertDispatcher(ref: ActorRef, dispatcher: String): Unit = ref match {
    case r: ActorRefWithCell ⇒
      if (r.underlying.props.dispatcher != dispatcher)
--- a/akka-stream-tests/src/test/java/akka/stream/StreamTest.java
+++ b/akka-stream-tests/src/test/java/akka/stream/StreamTest.java
@ -4,6 +4,9 @@

 package akka.stream;

+import akka.stream.testkit.javadsl.StreamTestKit;
+import org.junit.After;
+import org.junit.Before;
 import org.scalatest.junit.JUnitSuite;

 import akka.actor.ActorSystem;
@ -11,11 +14,23 @@ import akka.testkit.AkkaJUnitActorSystemResource;

 public abstract class StreamTest extends JUnitSuite {
    final protected ActorSystem system;
-    final protected ActorMaterializer materializer;
+    final private ActorMaterializerSettings settings;
+
+    protected ActorMaterializer materializer;

    protected StreamTest(AkkaJUnitActorSystemResource actorSystemResource) {
        system = actorSystemResource.getSystem();
-        ActorMaterializerSettings settings = ActorMaterializerSettings.create(system);
+        settings = ActorMaterializerSettings.create(system);
+    }
+
+    @Before
+    public void setUp() {
        materializer = ActorMaterializer.create(settings, system);
    }
+
+    @After
+    public void tearDown() {
+        StreamTestKit.assertAllStagesStopped(materializer);
+        materializer.shutdown();
+    }
 }
--- a/akka-stream-tests/src/test/java/akka/stream/javadsl/SourceTest.java
+++ b/akka-stream-tests/src/test/java/akka/stream/javadsl/SourceTest.java
@ -8,6 +8,7 @@ import akka.Done;
 import akka.NotUsed;
 import akka.actor.ActorRef;
 import akka.actor.Cancellable;
+import akka.actor.Status;
 import akka.japi.Pair;
 import akka.japi.function.*;
 import akka.japi.pf.PFBuilder;
@ -446,6 +447,7 @@ public class SourceTest extends StreamTest {
    probe.expectNoMessage(Duration.ofMillis(200));
    probe.expectMsgEquals("tick");
    probe.expectNoMessage(Duration.ofMillis(200));
+    cancellable.cancel();
  }

  @Test
@ -547,6 +549,7 @@ public class SourceTest extends StreamTest {
    probe.expectMsgEquals(1);
    ref.tell(2, ActorRef.noSender());
    probe.expectMsgEquals(2);
+    ref.tell(new Status.Success("ok"), ActorRef.noSender());
  }

  @Test
--- a/akka-stream-tests/src/test/java/akka/stream/javadsl/TcpTest.java
+++ b/akka-stream-tests/src/test/java/akka/stream/javadsl/TcpTest.java
@ -90,6 +90,8 @@ public class TcpTest extends StreamTest {
    for (int i = 0; i < testInput.size(); i ++) {
      assertEquals(testInput.get(i).head(), result[i]);
    }
+
+    b.unbind();
  }

  @Test
@ -111,6 +113,7 @@ public class TcpTest extends StreamTest {
          if (e.getCause() instanceof BindFailedException) {} // all good
          else throw new AssertionError("failed", e);
          // expected
+          b.unbind();
        } catch (Exception e) {
          throw new AssertionError("failed", e);
        }
--- a/akka-stream-tests/src/test/scala-jdk9-only/akka/stream/scaladsl/FlowPublisherSinkSpec.scala
+++ b/akka-stream-tests/src/test/scala-jdk9-only/akka/stream/scaladsl/FlowPublisherSinkSpec.scala
@ -8,6 +8,7 @@ import akka.stream.testkit.StreamSpec
 import akka.stream.{ ClosedShape, ActorMaterializer }

 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import scala.concurrent.duration._

 import scala.concurrent.Await
--- a/akka-stream-tests/src/test/scala/akka/stream/actor/ActorPublisherSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/actor/ActorPublisherSpec.scala
@ -9,6 +9,7 @@ import akka.stream.{ ClosedShape, ActorMaterializer, ActorMaterializerSettings,
 import akka.stream.scaladsl._
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.impl.ReactiveStreamsCompliance
 import akka.testkit.TestEvent.Mute
 import akka.testkit.{ EventFilter, ImplicitSender, TestProbe }
@ -410,7 +411,7 @@ class ActorPublisherSpec extends StreamSpec(ActorPublisherSpec.config) with Impl

    "be able to define a subscription-timeout, after which it should shut down" in {
      implicit val materializer = ActorMaterializer()
-      Utils.assertAllStagesStopped {
+      assertAllStagesStopped {
        val timeout = 150.millis
        val a = system.actorOf(timeoutingProps(testActor, timeout))
        val pub = ActorPublisher(a)
--- a/akka-stream-tests/src/test/scala/akka/stream/extra/FlowTimedSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/extra/FlowTimedSpec.scala
@ -9,6 +9,7 @@ import akka.stream.scaladsl.{ Source, Flow }
 import akka.stream.scaladsl.Sink
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.testkit.TestProbe
 import org.reactivestreams.{ Publisher, Subscriber }

--- a/akka-stream-tests/src/test/scala/akka/stream/impl/GraphStageLogicSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/impl/GraphStageLogicSpec.scala
@ -10,7 +10,7 @@ import akka.stream.testkit.StreamSpec
 import akka.stream._
 import akka.stream.scaladsl._
 import akka.stream.stage._
-import akka.stream.testkit.Utils.assertAllStagesStopped
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.scaladsl.TestSink
 import akka.stream.impl.fusing._
 import org.scalatest.concurrent.ScalaFutures
--- a/akka-stream-tests/src/test/scala/akka/stream/impl/TimeoutsSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/impl/TimeoutsSpec.scala
@ -9,6 +9,7 @@ import java.util.concurrent.TimeoutException
 import akka.Done
 import akka.stream.scaladsl._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.{ StreamSpec, TestPublisher, TestSubscriber }
 import akka.stream._
 import org.scalatest.{ Matchers, WordSpecLike }
--- a/akka-stream-tests/src/test/scala/akka/stream/impl/fusing/ActorGraphInterpreterSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/impl/fusing/ActorGraphInterpreterSpec.scala
@ -13,6 +13,7 @@ import akka.stream.impl.fusing.GraphStages.SimpleLinearGraphStage
 import akka.stream.scaladsl._
 import akka.stream.stage.{ GraphStage, GraphStageLogic, InHandler, OutHandler }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.{ StreamSpec, TestPublisher, TestSubscriber }
 import akka.testkit.{ EventFilter, TestLatch }

--- a/akka-stream-tests/src/test/scala/akka/stream/impl/fusing/GraphInterpreterPortsSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/impl/fusing/GraphInterpreterPortsSpec.scala
@ -6,6 +6,7 @@ package akka.stream.impl.fusing

 import akka.stream.testkit.StreamSpec
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 class GraphInterpreterPortsSpec extends StreamSpec with GraphInterpreterSpecKit {

--- a/akka-stream-tests/src/test/scala/akka/stream/impl/fusing/KeepGoingStageSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/impl/fusing/KeepGoingStageSpec.scala
@ -10,6 +10,7 @@ import akka.stream.testkit.StreamSpec
 import akka.stream.{ Attributes, Inlet, SinkShape, ActorMaterializer }
 import akka.stream.stage.{ InHandler, AsyncCallback, GraphStageLogic, GraphStageWithMaterializedValue }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 import scala.concurrent.{ Await, Promise, Future }
 import scala.concurrent.duration._
--- a/akka-stream-tests/src/test/scala/akka/stream/io/FileSinkSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/io/FileSinkSpec.scala
@ -16,6 +16,7 @@ import akka.stream.impl.StreamSupervisor.Children
 import akka.stream.scaladsl.{ FileIO, Sink, Source }
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream._
 import akka.util.ByteString
 import com.google.common.jimfs.{ Configuration, Jimfs }
--- a/akka-stream-tests/src/test/scala/akka/stream/io/FileSourceSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/io/FileSourceSpec.scala
@ -16,6 +16,7 @@ import akka.stream.impl.StreamSupervisor.Children
 import akka.stream.io.FileSourceSpec.Settings
 import akka.stream.scaladsl.{ FileIO, Keep, Sink }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import akka.stream.testkit.scaladsl.TestSink
 import akka.testkit.TestDuration
--- a/akka-stream-tests/src/test/scala/akka/stream/io/InputStreamSinkSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/io/InputStreamSinkSpec.scala
@ -15,6 +15,7 @@ import akka.stream.impl.io.InputStreamSinkStage
 import akka.stream.impl.{ PhasedFusingActorMaterializer, StreamSupervisor }
 import akka.stream.scaladsl.{ Keep, Source, StreamConverters }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.scaladsl.TestSource
 import akka.stream.testkit._
 import akka.testkit.TestProbe
--- a/akka-stream-tests/src/test/scala/akka/stream/io/InputStreamSourceSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/io/InputStreamSourceSpec.scala
@ -10,6 +10,7 @@ import java.util.concurrent.CountDownLatch
 import akka.stream.scaladsl.{ Sink, StreamConverters }
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.scaladsl.TestSink
 import akka.stream.{ ActorMaterializer, ActorMaterializerSettings }
 import akka.util.ByteString
--- a/akka-stream-tests/src/test/scala/akka/stream/io/OutputStreamSinkSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/io/OutputStreamSinkSpec.scala
@ -9,6 +9,7 @@ import java.io.OutputStream
 import akka.stream.scaladsl.{ Source, StreamConverters }
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.{ AbruptIOTerminationException, ActorMaterializer, ActorMaterializerSettings }
 import akka.testkit.TestProbe
 import akka.util.ByteString
--- a/akka-stream-tests/src/test/scala/akka/stream/io/OutputStreamSourceSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/io/OutputStreamSourceSpec.scala
@ -16,6 +16,7 @@ import akka.stream.impl.StreamSupervisor.Children
 import akka.stream.impl.io.OutputStreamSourceStage
 import akka.stream.scaladsl.{ Keep, Sink, StreamConverters }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import akka.stream.testkit.scaladsl.TestSink
 import akka.testkit.TestProbe
--- a/akka-stream-tests/src/test/scala/akka/stream/io/TcpSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/io/TcpSpec.scala
@ -15,6 +15,7 @@ import akka.stream._
 import akka.stream.scaladsl.Tcp.{ IncomingConnection, ServerBinding }
 import akka.stream.scaladsl.{ Flow, _ }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import akka.testkit.{ EventFilter, TestKit, TestLatch, TestProbe }
 import akka.testkit.SocketUtil.temporaryServerAddress
--- a/akka-stream-tests/src/test/scala/akka/stream/io/TlsSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/io/TlsSpec.scala
@ -24,6 +24,7 @@ import akka.stream.scaladsl._
 import akka.stream.stage._
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.util.ByteString
 import javax.net.ssl._

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/ActorRefBackpressureSinkSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/ActorRefBackpressureSinkSpec.scala
@ -9,6 +9,7 @@ import akka.actor.{ Actor, ActorRef, Props, Status }
 import akka.stream.ActorMaterializer
 import akka.stream.Attributes.inputBuffer
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import akka.stream.testkit.scaladsl._
 import akka.testkit.TestProbe
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/ActorRefSinkSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/ActorRefSinkSpec.scala
@ -7,6 +7,7 @@ package akka.stream.scaladsl
 import akka.stream.ActorMaterializer
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.scaladsl._
 import akka.actor.Actor
 import akka.actor.ActorRef
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/ActorRefSourceSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/ActorRefSourceSpec.scala
@ -9,6 +9,7 @@ import akka.stream.{ Attributes, ActorMaterializer, OverflowStrategy }
 import akka.stream.testkit._
 import akka.stream.testkit.scaladsl._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.actor.PoisonPill
 import akka.actor.Status

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/BidiFlowSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/BidiFlowSpec.scala
@ -7,6 +7,7 @@ package akka.stream.scaladsl
 import akka.NotUsed
 import akka.stream.testkit.StreamSpec
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.util.ByteString
 import akka.stream._
 import scala.concurrent.Await
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowAskSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowAskSpec.scala
@ -11,6 +11,7 @@ import akka.stream.ActorAttributes.supervisionStrategy
 import akka.stream.{ ActorAttributes, ActorMaterializer, Supervision }
 import akka.stream.Supervision.resumingDecider
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import akka.testkit.{ TestActors, TestProbe }

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowBufferSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowBufferSpec.scala
@ -11,6 +11,7 @@ import akka.stream.{ BufferOverflowException, ActorMaterializer, ActorMaterializ
 import akka.stream.testkit._
 import akka.stream.testkit.scaladsl._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 class FlowBufferSpec extends StreamSpec {

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowConcatAllSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowConcatAllSpec.scala
@ -12,6 +12,7 @@ import akka.stream.ActorMaterializer
 import akka.stream.ActorMaterializerSettings
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 class FlowConcatAllSpec extends StreamSpec {

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowConcatSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowConcatSpec.scala
@ -5,6 +5,7 @@
 package akka.stream.scaladsl

 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.scaladsl.TestSink
 import akka.stream.testkit.{ BaseTwoStreamsSetup, TestPublisher, TestSubscriber }
 import org.reactivestreams.Publisher
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowDelaySpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowDelaySpec.scala
@ -8,6 +8,7 @@ import akka.Done
 import akka.stream.Attributes._
 import akka.stream.OverflowStrategies.EmitEarly
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.scaladsl.TestSink
 import akka.stream.testkit.{ StreamSpec, TestPublisher, TestSubscriber }
 import akka.stream._
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowDetacherSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowDetacherSpec.scala
@ -8,7 +8,8 @@ import akka.stream._
 import scala.concurrent.Await
 import scala.concurrent.duration._
 import akka.stream.testkit.scaladsl.TestSink
-import akka.stream.testkit.{ StreamSpec, Utils }
+import akka.stream.testkit.StreamSpec
+import akka.stream.testkit.scaladsl.StreamTestKit._

 class FlowDetacherSpec extends StreamSpec {

@ -16,14 +17,14 @@ class FlowDetacherSpec extends StreamSpec {

  "A Detacher" must {

-    "pass through all elements" in Utils.assertAllStagesStopped {
+    "pass through all elements" in assertAllStagesStopped {
      Source(1 to 100)
        .detach
        .runWith(Sink.seq)
        .futureValue should ===(1 to 100)
    }

-    "pass through failure" in Utils.assertAllStagesStopped {
+    "pass through failure" in assertAllStagesStopped {
      val ex = new Exception("buh")
      val result = Source(1 to 100)
        .map(x ⇒ if (x == 50) throw ex else x)
@ -35,7 +36,7 @@ class FlowDetacherSpec extends StreamSpec {

    }

-    "emit the last element when completed without demand" in Utils.assertAllStagesStopped {
+    "emit the last element when completed without demand" in assertAllStagesStopped {
      Source.single(42)
        .detach
        .runWith(TestSink.probe)
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowDropWhileSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowDropWhileSpec.scala
@ -7,6 +7,7 @@ package akka.stream.scaladsl
 import akka.stream.ActorAttributes._
 import akka.stream.Supervision._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.ActorMaterializer
 import akka.stream.ActorMaterializerSettings
 import akka.stream.testkit._
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowFilterSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowFilterSpec.scala
@ -10,6 +10,7 @@ import java.util.concurrent.ThreadLocalRandom.{ current ⇒ random }
 import akka.stream.ActorAttributes._
 import akka.stream.Supervision._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.ActorMaterializer
 import akka.stream.ActorMaterializerSettings
 import akka.stream.testkit._
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowFlattenMergeSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowFlattenMergeSpec.scala
@ -7,7 +7,8 @@ package akka.stream.scaladsl
 import akka.NotUsed
 import akka.stream.stage.{ GraphStage, GraphStageLogic, OutHandler }
 import akka.stream._
-import akka.stream.testkit.Utils.{ TE, assertAllStagesStopped }
+import akka.stream.testkit.Utils.TE
+import akka.stream.testkit.scaladsl.StreamTestKit._

 import scala.concurrent._
 import scala.concurrent.duration._
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowFoldAsyncSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowFoldAsyncSpec.scala
@ -10,6 +10,7 @@ import akka.stream.ActorMaterializer
 import akka.stream.Supervision.{ restartingDecider, resumingDecider }
 import akka.stream.impl.ReactiveStreamsCompliance
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import org.scalatest.concurrent.PatienceConfiguration.Timeout

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowFoldSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowFoldSpec.scala
@ -11,6 +11,7 @@ import scala.concurrent.Await
 import scala.util.control.NoStackTrace
 import akka.stream.{ ActorAttributes, ActorMaterializer, Supervision }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 import scala.concurrent.duration._

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowForeachSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowForeachSpec.scala
@ -8,6 +8,7 @@ import scala.util.control.NoStackTrace
 import akka.stream.ActorMaterializer
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import scala.concurrent.Await
 import scala.concurrent.duration._

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowFromFutureSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowFromFutureSpec.scala
@ -11,6 +11,7 @@ import akka.stream.ActorMaterializer
 import akka.stream.ActorMaterializerSettings
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 class FlowFromFutureSpec extends StreamSpec {

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowGroupBySpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowGroupBySpec.scala
@ -20,6 +20,7 @@ import akka.stream.Supervision.resumingDecider
 import akka.stream.impl.fusing.GroupBy
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import org.reactivestreams.Publisher
 import org.scalatest.concurrent.PatienceConfiguration.Timeout
 import akka.stream.testkit.scaladsl.TestSource
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowGroupedWithinSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowGroupedWithinSpec.scala
@ -11,6 +11,7 @@ import java.util.concurrent.ThreadLocalRandom.{ current ⇒ random }
 import akka.stream.{ ActorMaterializer, ActorMaterializerSettings, ThrottleMode }
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 import akka.testkit.TimingTest
 import akka.util.ConstantFun
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowIdleInjectSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowIdleInjectSpec.scala
@ -5,7 +5,8 @@
 package akka.stream.scaladsl

 import akka.stream.{ ActorMaterializer, ActorMaterializerSettings }
-import akka.stream.testkit.{ StreamSpec, TestSubscriber, TestPublisher, Utils }
+import akka.stream.testkit.{ StreamSpec, TestSubscriber, TestPublisher }
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import scala.concurrent.Await
 import scala.concurrent.duration._

@ -18,13 +19,13 @@ class FlowIdleInjectSpec extends StreamSpec {

  "keepAlive" must {

-    "not emit additional elements if upstream is fast enough" in Utils.assertAllStagesStopped {
+    "not emit additional elements if upstream is fast enough" in assertAllStagesStopped {
      Await.result(
        Source(1 to 10).keepAlive(1.second, () ⇒ 0).grouped(1000).runWith(Sink.head),
        3.seconds) should ===(1 to 10)
    }

-    "emit elements periodically after silent periods" in Utils.assertAllStagesStopped {
+    "emit elements periodically after silent periods" in assertAllStagesStopped {
      val sourceWithIdleGap = Source(1 to 5) ++ Source(6 to 10).initialDelay(2.second)

      val result = Await.result(
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowInitialDelaySpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowInitialDelaySpec.scala
@ -6,7 +6,8 @@ package akka.stream.scaladsl

 import java.util.concurrent.TimeoutException
 import akka.stream.{ ActorMaterializer, ActorMaterializerSettings }
-import akka.stream.testkit.{ StreamSpec, Utils, TestSubscriber }
+import akka.stream.testkit.{ StreamSpec, TestSubscriber }
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import scala.concurrent.Await
 import scala.concurrent.duration._

@ -19,13 +20,13 @@ class FlowInitialDelaySpec extends StreamSpec {

  "Flow initialDelay" must {

-    "work with zero delay" in Utils.assertAllStagesStopped {
+    "work with zero delay" in assertAllStagesStopped {
      Await.result(
        Source(1 to 10).initialDelay(Duration.Zero).grouped(100).runWith(Sink.head),
        1.second) should ===(1 to 10)
    }

-    "delay elements by the specified time but not more" in Utils.assertAllStagesStopped {
+    "delay elements by the specified time but not more" in assertAllStagesStopped {
      a[TimeoutException] shouldBe thrownBy {
        Await.result(
          Source(1 to 10).initialDelay(2.seconds).initialTimeout(1.second).runWith(Sink.ignore),
@ -37,7 +38,7 @@ class FlowInitialDelaySpec extends StreamSpec {
        2.seconds)
    }

-    "properly ignore timer while backpressured" in Utils.assertAllStagesStopped {
+    "properly ignore timer while backpressured" in assertAllStagesStopped {
      val probe = TestSubscriber.probe[Int]()
      Source(1 to 10).initialDelay(0.5.second).runWith(Sink.fromSubscriber(probe))

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowInterleaveSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowInterleaveSpec.scala
@ -5,6 +5,7 @@
 package akka.stream.scaladsl

 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import org.reactivestreams.Publisher

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowIteratorSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowIteratorSpec.scala
@ -11,6 +11,7 @@ import akka.stream.ActorMaterializer
 import akka.stream.ActorMaterializerSettings
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.testkit.EventFilter

 class FlowIteratorSpec extends AbstractFlowIteratorSpec {
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowJoinSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowJoinSpec.scala
@ -7,6 +7,7 @@ package akka.stream.scaladsl
 import akka.stream.{ ActorMaterializer, ActorMaterializerSettings, FlowShape, OverflowStrategy }
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.scaladsl._
 import com.typesafe.config.ConfigFactory
 import org.scalatest.time._
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowKillSwitchSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowKillSwitchSpec.scala
@ -8,7 +8,8 @@ import akka.Done
 import akka.stream.testkit.StreamSpec
 import akka.stream.{ ActorMaterializer, ClosedShape, KillSwitches }
 import akka.stream.testkit.scaladsl.{ TestSink, TestSource }
-import akka.stream.testkit.Utils.{ TE, assertAllStagesStopped }
+import akka.stream.testkit.scaladsl.StreamTestKit._
+import akka.stream.testkit.Utils.TE

 import scala.concurrent.duration._

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowMapAsyncSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowMapAsyncSpec.scala
@ -12,6 +12,7 @@ import akka.stream.{ ActorAttributes, ActorMaterializer, Supervision }
 import akka.stream.Supervision.resumingDecider
 import akka.stream.impl.ReactiveStreamsCompliance
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import akka.stream.testkit.scaladsl.TestSink
 import akka.testkit.{ TestLatch, TestProbe }
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowMapAsyncUnorderedSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowMapAsyncUnorderedSpec.scala
@ -12,6 +12,7 @@ import akka.stream.ActorMaterializer
 import akka.stream.testkit._
 import akka.stream.testkit.scaladsl._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.testkit.TestLatch
 import akka.testkit.TestProbe
 import akka.stream.ActorAttributes.supervisionStrategy
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowMapConcatSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowMapConcatSpec.scala
@ -7,6 +7,7 @@ package akka.stream.scaladsl
 import akka.stream.testkit.scaladsl.TestSink
 import akka.stream.{ Supervision, ActorAttributes, ActorMaterializer, ActorMaterializerSettings }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import scala.util.control.NoStackTrace

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowMapErrorSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowMapErrorSpec.scala
@ -6,6 +6,7 @@ package akka.stream.scaladsl

 import akka.stream.testkit.StreamSpec
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.scaladsl.TestSink
 import akka.stream.{ ActorMaterializer, ActorMaterializerSettings }

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowMergeSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowMergeSpec.scala
@ -5,6 +5,7 @@
 package akka.stream.scaladsl

 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import org.reactivestreams.Publisher

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowOnCompleteSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowOnCompleteSpec.scala
@ -12,6 +12,7 @@ import akka.stream.ActorMaterializer
 import akka.stream.ActorMaterializerSettings
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.testkit.TestProbe

 class FlowOnCompleteSpec extends StreamSpec with ScriptedTest {
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowPrefixAndTailSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowPrefixAndTailSpec.scala
@ -11,6 +11,7 @@ import scala.util.control.NoStackTrace
 import akka.stream._
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 class FlowPrefixAndTailSpec extends StreamSpec {

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowRecoverSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowRecoverSpec.scala
@ -8,6 +8,7 @@ import akka.stream.testkit.StreamSpec
 import akka.stream.testkit.scaladsl.TestSink
 import akka.stream.{ ActorMaterializer, ActorMaterializerSettings }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 import scala.util.control.NoStackTrace

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowRecoverWithSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowRecoverWithSpec.scala
@ -9,6 +9,7 @@ import akka.stream.testkit.StreamSpec
 import akka.stream.testkit.scaladsl.TestSink
 import akka.stream._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 import scala.util.control.NoStackTrace

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowReduceSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowReduceSpec.scala
@ -10,6 +10,7 @@ import scala.concurrent.Await
 import scala.util.control.NoStackTrace
 import akka.stream.{ ActorAttributes, ActorMaterializer, Supervision }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 import scala.concurrent.duration._

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowScanSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowScanSpec.scala
@ -9,6 +9,7 @@ import akka.stream.testkit.StreamSpec
 import akka.stream.testkit.scaladsl.TestSink
 import akka.stream.{ ActorAttributes, ActorMaterializer, ActorMaterializerSettings, Supervision }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 import scala.collection.immutable
 import scala.concurrent.Await
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowSlidingSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowSlidingSpec.scala
@ -5,6 +5,7 @@
 package akka.stream.scaladsl

 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.{ ActorMaterializer, ActorMaterializerSettings }
 import akka.stream.testkit._
 import org.scalacheck.Gen
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowSpec.scala
@ -10,6 +10,7 @@ import akka.NotUsed
 import akka.actor._
 import akka.stream.impl._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import akka.stream._
 import akka.testkit.TestDuration
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowSplitAfterSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowSplitAfterSpec.scala
@ -10,6 +10,7 @@ import akka.stream.Supervision.resumingDecider
 import akka.stream.impl.SubscriptionTimeoutException
 import akka.stream.testkit.{ StreamSpec, TestPublisher, TestSubscriber }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import org.reactivestreams.Publisher
 import scala.concurrent.Await
 import scala.concurrent.duration._
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowSplitWhenSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowSplitWhenSpec.scala
@ -9,6 +9,7 @@ import akka.stream._
 import akka.stream.Supervision.resumingDecider
 import akka.stream.impl.SubscriptionTimeoutException
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import akka.stream.testkit.scaladsl.TestSink
 import org.reactivestreams.Publisher
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowTakeWhileSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowTakeWhileSpec.scala
@ -7,6 +7,7 @@ package akka.stream.scaladsl
 import akka.stream.ActorAttributes._
 import akka.stream.Supervision._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.ActorMaterializer
 import akka.stream.ActorMaterializerSettings
 import akka.stream.testkit._
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowTakeWithinSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowTakeWithinSpec.scala
@ -8,6 +8,7 @@ import scala.concurrent.duration._
 import akka.stream.ActorMaterializer
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 class FlowTakeWithinSpec extends StreamSpec {

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowThrottleSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowThrottleSpec.scala
@ -11,6 +11,7 @@ import akka.stream.ThrottleMode.{ Enforcing, Shaping }
 import akka.stream._
 import akka.stream.testkit._
 import akka.stream.testkit.scaladsl.TestSink
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.util.ByteString
 import scala.concurrent.Await
 import scala.concurrent.duration._
@ -27,7 +28,7 @@ class FlowThrottleSpec extends StreamSpec {
    ByteString(new Random().shuffle(0 to 255).take(length).map(_.toByte).toArray)

  "Throttle for single cost elements" must {
-    "work for the happy case" in Utils.assertAllStagesStopped {
+    "work for the happy case" in assertAllStagesStopped {
      //Source(1 to 5).throttle(1, 100.millis, 0, Shaping)
      Source(1 to 5).throttle(19, 1000.millis, -1, Shaping)
        .runWith(TestSink.probe[Int])
@ -36,7 +37,7 @@ class FlowThrottleSpec extends StreamSpec {
        .expectComplete()
    }

-    "accept very high rates" in Utils.assertAllStagesStopped {
+    "accept very high rates" in assertAllStagesStopped {
      Source(1 to 5).throttle(1, 1.nanos, 0, Shaping)
        .runWith(TestSink.probe[Int])
        .request(5)
@ -44,7 +45,7 @@ class FlowThrottleSpec extends StreamSpec {
        .expectComplete()
    }

-    "accept very low rates" in Utils.assertAllStagesStopped {
+    "accept very low rates" in assertAllStagesStopped {
      Source(1 to 5).throttle(1, 100.days, 1, Shaping)
        .runWith(TestSink.probe[Int])
        .request(5)
@ -53,7 +54,7 @@ class FlowThrottleSpec extends StreamSpec {
        .cancel() // We won't wait 100 days, sorry
    }

-    "work if there are two throttles in different streams" in Utils.assertAllStagesStopped {
+    "work if there are two throttles in different streams" in assertAllStagesStopped {
      val sharedThrottle = Flow[Int].throttle(1, 1.day, 1, Enforcing)

      // If there is accidental shared state then we would not be able to pass through the single element
@ -70,7 +71,7 @@ class FlowThrottleSpec extends StreamSpec {

    }

-    "emit single element per tick" in Utils.assertAllStagesStopped {
+    "emit single element per tick" in assertAllStagesStopped {
      val upstream = TestPublisher.probe[Int]()
      val downstream = TestSubscriber.probe[Int]()

@ -89,7 +90,7 @@ class FlowThrottleSpec extends StreamSpec {
      downstream.expectComplete()
    }

-    "not send downstream if upstream does not emit element" in Utils.assertAllStagesStopped {
+    "not send downstream if upstream does not emit element" in assertAllStagesStopped {
      val upstream = TestPublisher.probe[Int]()
      val downstream = TestSubscriber.probe[Int]()
      Source.fromPublisher(upstream).throttle(1, 300.millis, 0, Shaping).runWith(Sink.fromSubscriber(downstream))
@ -105,13 +106,13 @@ class FlowThrottleSpec extends StreamSpec {
      upstream.sendComplete()
    }

-    "cancel when downstream cancels" in Utils.assertAllStagesStopped {
+    "cancel when downstream cancels" in assertAllStagesStopped {
      val downstream = TestSubscriber.probe[Int]()
      Source(1 to 10).throttle(1, 300.millis, 0, Shaping).runWith(Sink.fromSubscriber(downstream))
      downstream.cancel()
    }

-    "send elements downstream as soon as time comes" in Utils.assertAllStagesStopped {
+    "send elements downstream as soon as time comes" in assertAllStagesStopped {
      val probe = Source(1 to 10).throttle(2, 750.millis, 0, Shaping).runWith(TestSink.probe[Int])
        .request(5)
      probe.receiveWithin(900.millis) should be(Seq(1, 2))
@ -122,7 +123,7 @@ class FlowThrottleSpec extends StreamSpec {
        .cancel()
    }

-    "burst according to its maximum if enough time passed" in Utils.assertAllStagesStopped {
+    "burst according to its maximum if enough time passed" in assertAllStagesStopped {
      val upstream = TestPublisher.probe[Int]()
      val downstream = TestSubscriber.probe[Int]()
      Source.fromPublisher(upstream).throttle(1, 200.millis, 5, Shaping).runWith(Sink.fromSubscriber(downstream))
@ -139,7 +140,7 @@ class FlowThrottleSpec extends StreamSpec {
      downstream.cancel()
    }

-    "burst some elements if have enough time" in Utils.assertAllStagesStopped {
+    "burst some elements if have enough time" in assertAllStagesStopped {
      val upstream = TestPublisher.probe[Int]()
      val downstream = TestSubscriber.probe[Int]()
      Source.fromPublisher(upstream).throttle(1, 200.millis, 5, Shaping).runWith(Sink.fromSubscriber(downstream))
@ -160,7 +161,7 @@ class FlowThrottleSpec extends StreamSpec {
      downstream.cancel()
    }

-    "throw exception when exceeding throughtput in enforced mode" in Utils.assertAllStagesStopped {
+    "throw exception when exceeding throughtput in enforced mode" in assertAllStagesStopped {
      Await.result(
        Source(1 to 5).throttle(1, 200.millis, 5, Enforcing).runWith(Sink.seq),
        2.seconds) should ===(1 to 5) // Burst is 5 so this will not fail
@ -172,7 +173,7 @@ class FlowThrottleSpec extends StreamSpec {
      }
    }

-    "properly combine shape and throttle modes" in Utils.assertAllStagesStopped {
+    "properly combine shape and throttle modes" in assertAllStagesStopped {
      Source(1 to 5).throttle(1, 100.millis, 5, Shaping)
        .throttle(1, 100.millis, 5, Enforcing)
        .runWith(TestSink.probe[Int])
@ -183,7 +184,7 @@ class FlowThrottleSpec extends StreamSpec {
  }

  "Throttle for various cost elements" must {
-    "work for happy case" in Utils.assertAllStagesStopped {
+    "work for happy case" in assertAllStagesStopped {
      Source(1 to 5).throttle(1, 100.millis, 0, (_) ⇒ 1, Shaping)
        .runWith(TestSink.probe[Int])
        .request(5)
@ -191,7 +192,7 @@ class FlowThrottleSpec extends StreamSpec {
        .expectComplete()
    }

-    "emit elements according to cost" in Utils.assertAllStagesStopped {
+    "emit elements according to cost" in assertAllStagesStopped {
      val list = (1 to 4).map(_ * 2).map(genByteString)
      Source(list).throttle(2, 200.millis, 0, _.length, Shaping)
        .runWith(TestSink.probe[ByteString])
@ -206,7 +207,7 @@ class FlowThrottleSpec extends StreamSpec {
        .expectComplete()
    }

-    "not send downstream if upstream does not emit element" in Utils.assertAllStagesStopped {
+    "not send downstream if upstream does not emit element" in assertAllStagesStopped {
      val upstream = TestPublisher.probe[Int]()
      val downstream = TestSubscriber.probe[Int]()
      Source.fromPublisher(upstream).throttle(2, 300.millis, 0, identity, Shaping).runWith(Sink.fromSubscriber(downstream))
@ -222,13 +223,13 @@ class FlowThrottleSpec extends StreamSpec {
      upstream.sendComplete()
    }

-    "cancel when downstream cancels" in Utils.assertAllStagesStopped {
+    "cancel when downstream cancels" in assertAllStagesStopped {
      val downstream = TestSubscriber.probe[Int]()
      Source(1 to 10).throttle(2, 200.millis, 0, identity, Shaping).runWith(Sink.fromSubscriber(downstream))
      downstream.cancel()
    }

-    "send elements downstream as soon as time comes" in Utils.assertAllStagesStopped {
+    "send elements downstream as soon as time comes" in assertAllStagesStopped {
      val probe = Source(1 to 10).throttle(4, 500.millis, 0, _ ⇒ 2, Shaping).runWith(TestSink.probe[Int])
        .request(5)
      probe.receiveWithin(600.millis) should be(Seq(1, 2))
@ -239,7 +240,7 @@ class FlowThrottleSpec extends StreamSpec {
        .cancel()
    }

-    "burst according to its maximum if enough time passed" in Utils.assertAllStagesStopped {
+    "burst according to its maximum if enough time passed" in assertAllStagesStopped {
      val upstream = TestPublisher.probe[Int]()
      val downstream = TestSubscriber.probe[Int]()
      Source.fromPublisher(upstream).throttle(2, 400.millis, 5, (_) ⇒ 1, Shaping).runWith(Sink.fromSubscriber(downstream))
@ -260,7 +261,7 @@ class FlowThrottleSpec extends StreamSpec {
      downstream.cancel()
    }

-    "burst some elements if have enough time" in Utils.assertAllStagesStopped {
+    "burst some elements if have enough time" in assertAllStagesStopped {
      val upstream = TestPublisher.probe[Int]()
      val downstream = TestSubscriber.probe[Int]()
      Source.fromPublisher(upstream).throttle(2, 400.millis, 5, (e) ⇒ if (e < 9) 1 else 20, Shaping).runWith(Sink.fromSubscriber(downstream))
@ -281,7 +282,7 @@ class FlowThrottleSpec extends StreamSpec {
      downstream.cancel()
    }

-    "throw exception when exceeding throughtput in enforced mode" in Utils.assertAllStagesStopped {
+    "throw exception when exceeding throughtput in enforced mode" in assertAllStagesStopped {
      Await.result(
        Source(1 to 4).throttle(2, 200.millis, 10, identity, Enforcing).runWith(Sink.seq),
        2.seconds) should ===(1 to 4) // Burst is 10 so this will not fail
@ -293,7 +294,7 @@ class FlowThrottleSpec extends StreamSpec {
      }
    }

-    "properly combine shape and enforce modes" in Utils.assertAllStagesStopped {
+    "properly combine shape and enforce modes" in assertAllStagesStopped {
      Source(1 to 5).throttle(2, 200.millis, 0, identity, Shaping)
        .throttle(1, 100.millis, 5, Enforcing)
        .runWith(TestSink.probe[Int])
@ -302,7 +303,7 @@ class FlowThrottleSpec extends StreamSpec {
        .expectComplete()
    }

-    "handle rate calculation function exception" in Utils.assertAllStagesStopped {
+    "handle rate calculation function exception" in assertAllStagesStopped {
      val ex = new RuntimeException with NoStackTrace
      Source(1 to 5).throttle(2, 200.millis, 0, (_) ⇒ { throw ex }, Shaping)
        .throttle(1, 100.millis, 5, Enforcing)
@ -311,7 +312,7 @@ class FlowThrottleSpec extends StreamSpec {
        .expectError(ex)
    }

-    "work for real scenario with automatic burst size" taggedAs TimingTest in Utils.assertAllStagesStopped {
+    "work for real scenario with automatic burst size" taggedAs TimingTest in assertAllStagesStopped {
      val startTime = System.nanoTime()
      val counter1 = new AtomicInteger
      val timestamp1 = new AtomicLong(System.nanoTime())
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowWatchSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowWatchSpec.scala
@ -7,6 +7,7 @@ package akka.stream.scaladsl
 import akka.actor.{ Actor, PoisonPill, Props }
 import akka.stream.ActorMaterializer
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import akka.testkit.TestActors

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowWatchTerminationSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowWatchTerminationSpec.scala
@ -9,6 +9,7 @@ import akka.pattern.pipe
 import akka.stream._
 import akka.stream.testkit.StreamSpec
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.scaladsl.{ TestSink, TestSource }

 import scala.util.control.NoStackTrace
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowWireTapSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowWireTapSpec.scala
@ -7,6 +7,7 @@ package akka.stream.scaladsl
 import akka.Done
 import akka.stream.ActorMaterializer
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._

 import scala.concurrent.duration._
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowZipSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowZipSpec.scala
@ -5,6 +5,7 @@
 package akka.stream.scaladsl

 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.{ BaseTwoStreamsSetup, TestSubscriber }
 import org.reactivestreams.Publisher

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowZipWithIndexSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FlowZipWithIndexSpec.scala
@ -5,6 +5,7 @@
 package akka.stream.scaladsl

 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.{ ActorMaterializer, ActorMaterializerSettings }
 import akka.stream.testkit.{ StreamSpec, TestSubscriber }

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FutureFlattenSourceSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/FutureFlattenSourceSpec.scala
@ -8,8 +8,9 @@ import java.util.concurrent.{ CompletableFuture, TimeUnit }

 import akka.stream._
 import akka.stream.stage.{ GraphStageLogic, GraphStageWithMaterializedValue }
-import akka.stream.testkit.Utils.{ TE, assertAllStagesStopped }
+import akka.stream.testkit.Utils.TE
 import akka.stream.testkit.{ StreamSpec, TestPublisher, TestSubscriber }
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.testkit.TestLatch

 import scala.concurrent.{ Await, Future, Promise }
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphBalanceSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphBalanceSpec.scala
@ -11,6 +11,7 @@ import akka.stream._
 import akka.stream.testkit._
 import akka.stream.testkit.scaladsl._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 class GraphBalanceSpec extends StreamSpec {

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphBroadcastSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphBroadcastSpec.scala
@ -11,6 +11,7 @@ import scala.concurrent.duration._
 import akka.stream._
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 class GraphBroadcastSpec extends StreamSpec {

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphConcatSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphConcatSpec.scala
@ -9,6 +9,7 @@ import scala.concurrent.{ Promise }
 import akka.stream._
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 class GraphConcatSpec extends TwoStreamsSetup {

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphMergeLatestSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphMergeLatestSpec.scala
@ -6,6 +6,7 @@ package akka.stream.scaladsl

 import akka.stream._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import akka.stream.testkit.scaladsl.TestSource

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphMergeSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphMergeSpec.scala
@ -11,6 +11,7 @@ import scala.concurrent.duration._

 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 class GraphMergeSpec extends TwoStreamsSetup {
  import GraphDSL.Implicits._
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphPartitionSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphPartitionSpec.scala
@ -7,6 +7,7 @@ package akka.stream.scaladsl
 import akka.stream.testkit._
 import akka.stream.{ OverflowStrategy, ActorMaterializer, ActorMaterializerSettings, ClosedShape }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import scala.concurrent.Await
 import scala.concurrent.duration._

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphStageTimersSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphStageTimersSpec.scala
@ -15,6 +15,7 @@ import scala.concurrent.duration._

 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 object GraphStageTimersSpec {
  case object TestSingleTimer
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphUnzipSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphUnzipSpec.scala
@ -8,6 +8,7 @@ import scala.concurrent.duration._
 import akka.stream.{ ClosedShape, OverflowStrategy, ActorMaterializerSettings, ActorMaterializer }
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 class GraphUnzipSpec extends StreamSpec {

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphUnzipWithSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphUnzipWithSpec.scala
@ -7,6 +7,7 @@ package akka.stream.scaladsl
 import akka.stream._
 import akka.stream.testkit.TestSubscriber.Probe
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import org.reactivestreams.Publisher
 import scala.concurrent.duration._
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphWireTapSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphWireTapSpec.scala
@ -6,6 +6,7 @@ package akka.stream.scaladsl

 import akka.stream._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import akka.stream.testkit.scaladsl.TestSink

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphZipNSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphZipNSpec.scala
@ -6,6 +6,7 @@ package akka.stream.scaladsl

 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream._

 import scala.concurrent.duration._
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphZipSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/GraphZipSpec.scala
@ -6,6 +6,7 @@ package akka.stream.scaladsl

 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream._

 import scala.concurrent.Await
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/HeadSinkSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/HeadSinkSpec.scala
@ -10,6 +10,7 @@ import scala.concurrent.duration._
 import akka.stream.{ AbruptTerminationException, ActorMaterializer, ActorMaterializerSettings }
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 class HeadSinkSpec extends StreamSpec with ScriptedTest {

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/HubSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/HubSpec.scala
@ -6,8 +6,9 @@ package akka.stream.scaladsl

 import akka.stream.{ ActorMaterializer, KillSwitches, ThrottleMode }
 import akka.stream.testkit.{ StreamSpec, TestPublisher, TestSubscriber }
-import akka.stream.testkit.Utils.{ TE, assertAllStagesStopped }
+import akka.stream.testkit.Utils.TE
 import akka.stream.testkit.scaladsl.{ TestSink, TestSource }
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.testkit.EventFilter

 import scala.collection.immutable
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/LastSinkSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/LastSinkSpec.scala
@ -11,6 +11,7 @@ import akka.stream.ActorMaterializer
 import akka.stream.ActorMaterializerSettings
 import akka.stream.testkit._
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._

 class LastSinkSpec extends StreamSpec with ScriptedTest {

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/LazilyAsyncSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/LazilyAsyncSpec.scala
@ -9,7 +9,7 @@ import java.util.concurrent.atomic.AtomicBoolean
 import akka.Done
 import akka.stream.ActorMaterializer
 import akka.stream.testkit.{ StreamSpec, TestSubscriber }
-import akka.stream.testkit.Utils.assertAllStagesStopped
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.testkit.DefaultTimeout
 import org.scalatest.concurrent.ScalaFutures

--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/LazyFlowSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/LazyFlowSpec.scala
@ -9,6 +9,7 @@ import akka.stream._
 import akka.stream.stage.{ GraphStageLogic, GraphStageWithMaterializedValue }
 import akka.stream.testkit.{ StreamSpec, TestPublisher }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.scaladsl.{ TestSink, TestSource }
 import akka.stream.testkit.{ StreamSpec, TestPublisher }
 import org.scalatest.concurrent.ScalaFutures
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/LazySinkSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/LazySinkSpec.scala
@ -14,6 +14,7 @@ import akka.stream.stage.{ GraphStage, GraphStageLogic }
 import akka.stream.testkit.{ StreamSpec, TestPublisher }
 import akka.stream.testkit.TestSubscriber.Probe
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.scaladsl.TestSink

 import scala.concurrent.{ Await, Future, Promise }
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/LazySourceSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/LazySourceSpec.scala
@ -9,7 +9,8 @@ import java.util.concurrent.atomic.AtomicBoolean
 import akka.Done
 import akka.stream.impl.LazySource
 import akka.stream.stage.{ GraphStage, GraphStageLogic }
-import akka.stream.testkit.Utils.{ TE, assertAllStagesStopped }
+import akka.stream.testkit.Utils.TE
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.{ StreamSpec, TestPublisher, TestSubscriber }
 import akka.stream.{ ActorMaterializer, Attributes, Outlet, SourceShape }
 import akka.testkit.DefaultTimeout
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/MaybeSourceSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/MaybeSourceSpec.scala
@ -5,7 +5,8 @@
 package akka.stream.scaladsl

 import akka.stream.{ AbruptStageTerminationException, ActorMaterializer }
-import akka.stream.testkit.{ StreamSpec, TestSubscriber, Utils }
+import akka.stream.testkit.{ StreamSpec, TestSubscriber }
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.testkit.DefaultTimeout

 import scala.concurrent.duration._
@ -17,7 +18,7 @@ class MaybeSourceSpec extends StreamSpec with DefaultTimeout {

  "The Maybe Source" must {

-    "complete materialized future with None when stream cancels" in Utils.assertAllStagesStopped {
+    "complete materialized future with None when stream cancels" in assertAllStagesStopped {
      val neverSource = Source.maybe[Int]
      val pubSink = Sink.asPublisher[Int](false)

@ -34,7 +35,7 @@ class MaybeSourceSpec extends StreamSpec with DefaultTimeout {
      f.future.futureValue shouldEqual None
    }

-    "allow external triggering of empty completion" in Utils.assertAllStagesStopped {
+    "allow external triggering of empty completion" in assertAllStagesStopped {
      val neverSource = Source.maybe[Int].filter(_ ⇒ false)
      val counterSink = Sink.fold[Int, Int](0) { (acc, _) ⇒ acc + 1 }

@ -46,7 +47,7 @@ class MaybeSourceSpec extends StreamSpec with DefaultTimeout {
      counterFuture.futureValue shouldEqual 0
    }

-    "allow external triggering of empty completion when there was no demand" in Utils.assertAllStagesStopped {
+    "allow external triggering of empty completion when there was no demand" in assertAllStagesStopped {
      val probe = TestSubscriber.probe[Int]()
      val promise = Source.maybe[Int].to(Sink.fromSubscriber(probe)).run()

@ -56,7 +57,7 @@ class MaybeSourceSpec extends StreamSpec with DefaultTimeout {
      probe.expectComplete()
    }

-    "allow external triggering of non-empty completion" in Utils.assertAllStagesStopped {
+    "allow external triggering of non-empty completion" in assertAllStagesStopped {
      val neverSource = Source.maybe[Int]
      val counterSink = Sink.head[Int]

@ -68,7 +69,7 @@ class MaybeSourceSpec extends StreamSpec with DefaultTimeout {
      counterFuture.futureValue shouldEqual 6
    }

-    "allow external triggering of onError" in Utils.assertAllStagesStopped {
+    "allow external triggering of onError" in assertAllStagesStopped {
      val neverSource = Source.maybe[Int]
      val counterSink = Sink.fold[Int, Int](0) { (acc, _) ⇒ acc + 1 }

@ -80,7 +81,7 @@ class MaybeSourceSpec extends StreamSpec with DefaultTimeout {
      counterFuture.failed.futureValue.getMessage should include("Boom")
    }

-    "complete materialized future when materializer is shutdown" in Utils.assertAllStagesStopped {
+    "complete materialized future when materializer is shutdown" in assertAllStagesStopped {
      val mat = ActorMaterializer()
      val neverSource = Source.maybe[Int]
      val pubSink = Sink.asPublisher[Int](false)
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/PublisherSinkSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/PublisherSinkSpec.scala
@ -8,6 +8,7 @@ import akka.stream.testkit.StreamSpec
 import akka.stream.{ ClosedShape, ActorMaterializer }

 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import scala.concurrent.duration._

 import scala.concurrent.Await
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/QueueSinkSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/QueueSinkSpec.scala
@ -9,6 +9,7 @@ import akka.pattern.pipe
 import akka.stream.Attributes.inputBuffer
 import akka.stream.{ ActorMaterializer, StreamDetachedException }
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._

 import scala.concurrent.{ Await, Promise }
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/QueueSourceSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/QueueSourceSpec.scala
@ -10,6 +10,7 @@ import akka.pattern.pipe
 import akka.stream._
 import akka.stream.impl.QueueSource
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit.scaladsl.TestSink
 import akka.stream.testkit.{ GraphStageMessages, StreamSpec, TestSourceStage, TestSubscriber }
 import akka.testkit.TestProbe
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/RestartSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/RestartSpec.scala
@ -8,8 +8,9 @@ import java.util.concurrent.atomic.AtomicInteger

 import akka.stream.scaladsl.RestartWithBackoffFlow.Delay
 import akka.stream.testkit.{ StreamSpec, TestPublisher, TestSubscriber }
-import akka.stream.testkit.Utils.{ TE, assertAllStagesStopped }
+import akka.stream.testkit.Utils.TE
 import akka.stream.testkit.scaladsl.{ TestSink, TestSource }
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.{ ActorMaterializer, Attributes, OverflowStrategy }
 import akka.testkit.{ DefaultTimeout, TestDuration }
 import akka.{ Done, NotUsed }
--- a/akka-stream-tests/src/test/scala/akka/stream/scaladsl/SinkAsJavaStreamSpec.scala
+++ b/akka-stream-tests/src/test/scala/akka/stream/scaladsl/SinkAsJavaStreamSpec.scala
@ -11,6 +11,7 @@ import akka.stream._
 import akka.stream.impl.{ PhasedFusingActorMaterializer, StreamSupervisor }
 import akka.stream.impl.StreamSupervisor.Children
 import akka.stream.testkit.Utils._
+import akka.stream.testkit.scaladsl.StreamTestKit._
 import akka.stream.testkit._
 import akka.stream.testkit.scaladsl.TestSource
 import akka.util.ByteString
@ -43,13 +44,13 @@ class SinkAsJavaStreamSpec extends StreamSpec(UnboundedMailboxConfig) {
      javaSource.count() should ===(0)
    }

-    "work with endless stream" in Utils.assertAllStagesStopped {
+    "work with endless stream" in assertAllStagesStopped {
      val javaSource = Source.repeat(1).runWith(StreamConverters.asJavaStream())
      javaSource.limit(10).count() should ===(10)
      javaSource.close()
    }

-    "allow overriding the dispatcher using Attributes" in Utils.assertAllStagesStopped {
+    "allow overriding the dispatcher using Attributes" in assertAllStagesStopped {
      val sys = ActorSystem("dispatcher-testing", UnboundedMailboxConfig)
      val materializer = ActorMaterializer()(sys)

@ -62,7 +63,7 @@ class SinkAsJavaStreamSpec extends StreamSpec(UnboundedMailboxConfig) {
      } finally shutdown(sys)
    }

-    "work in separate IO dispatcher" in Utils.assertAllStagesStopped {
+    "work in separate IO dispatcher" in assertAllStagesStopped {
      val sys = ActorSystem("dispatcher-testing", UnboundedMailboxConfig)
      val materializer = ActorMaterializer()(sys)

--- a/Show more
+++ b/Show more