Akka Streams Testing: TestKit, Probes, and Stream Verification

Akka Streams Testing: TestKit, Probes, and Stream Verification

Akka Streams introduces a different testing challenge from regular actors. Streams are lazy, composable, and backpressure-aware — properties that make them powerful but also make naive testing approaches fail. The akka-stream-testkit module provides the right primitives for verifying stream behavior without races or flakiness.

Setup

Add the testkit to your dependencies:

// build.sbt
libraryDependencies ++= Seq(
  "com.typesafe.akka" %% "akka-stream"         % "2.8.0",
  "com.typesafe.akka" %% "akka-stream-testkit"  % "2.8.0" % Test,
  "com.typesafe.akka" %% "akka-testkit"         % "2.8.0" % Test,
  "org.scalatest"     %% "scalatest"             % "3.2.15" % Test
)

Base test structure:

import akka.actor.ActorSystem
import akka.stream.scaladsl._
import akka.stream.testkit.scaladsl._
import akka.testkit.TestKit
import org.scalatest.BeforeAndAfterAll
import org.scalatest.wordspec.AnyWordSpecLike

class MyStreamSpec
    extends TestKit(ActorSystem("MyStreamSpec"))
    with AnyWordSpecLike
    with BeforeAndAfterAll {

  override def afterAll(): Unit = {
    TestKit.shutdownActorSystem(system)
  }

  // tests go here
}

Testing Simple Flows with runWith

For straightforward transformations, runWith plus Sink.seq or Sink.head is the simplest approach:

"A filtering flow" should {
  "pass only even numbers" in {
    val flow = Flow[Int].filter(_ % 2 == 0)

    val result = Source(1 to 10)
      .via(flow)
      .runWith(Sink.seq)

    // runWith returns a Future — use Await or ScalaFutures
    import scala.concurrent.Await
    import scala.concurrent.duration._

    val elements = Await.result(result, 3.seconds)
    assert(elements == Seq(2, 4, 6, 8, 10))
  }
}

For cleaner async handling, mix in ScalaFutures:

import org.scalatest.concurrent.ScalaFutures

class MyStreamSpec
    extends TestKit(ActorSystem("test"))
    with AnyWordSpecLike
    with ScalaFutures
    with BeforeAndAfterAll {

  implicit val patience: PatienceConfig =
    PatienceConfig(timeout = 5.seconds, interval = 50.millis)

  "A mapping flow" should {
    "double each element" in {
      val result = Source(1 to 5)
        .map(_ * 2)
        .runWith(Sink.seq)

      result.futureValue shouldBe Seq(2, 4, 6, 8, 10)
    }
  }
}

TestSource and TestSink Probes

The TestSource and TestSink probes give you fine-grained control over demand and element injection:

import akka.stream.testkit.scaladsl.{TestSource, TestSink}

"A rate-limiting flow" should {
  "emit elements on demand" in {
    val (pub, sub) = TestSource[Int]()
      .via(Flow[Int].map(_ * 10))
      .toMat(TestSink[Int]())(Keep.both)
      .run()

    sub.request(3)          // request 3 elements
    pub.sendNext(1)
    pub.sendNext(2)
    pub.sendNext(3)

    sub.expectNext(10)
    sub.expectNext(20)
    sub.expectNext(30)

    pub.sendComplete()
    sub.expectComplete()
  }
}

Testing Backpressure

Backpressure verification is one of the hardest things to test in stream code. TestSink lets you explicitly control demand:

"A buffered flow" should {
  "not emit without demand" in {
    val (pub, sub) = TestSource[String]()
      .via(Flow[String].buffer(5, OverflowStrategy.backpressure))
      .toMat(TestSink[String]())(Keep.both)
      .run()

    // Send elements without requesting — should buffer, not emit
    pub.sendNext("a")
    pub.sendNext("b")
    pub.sendNext("c")

    sub.expectNoMessage(100.millis)  // nothing emitted yet

    sub.request(2)
    sub.expectNext("a")
    sub.expectNext("b")
    sub.expectNoMessage(50.millis)   // "c" still buffered

    sub.request(1)
    sub.expectNext("c")
  }
}

Testing Error Handling

"A resilient flow" should {
  "recover from upstream errors" in {
    val (pub, sub) = TestSource[Int]()
      .via(
        Flow[Int]
          .map { n =>
            if (n == 3) throw new RuntimeException("bad element")
            else n * 2
          }
          .recover { case _: RuntimeException => -1 }
      )
      .toMat(TestSink[Int]())(Keep.both)
      .run()

    sub.request(5)
    pub.sendNext(1)
    pub.sendNext(2)
    pub.sendNext(3)  // triggers error
    pub.sendNext(4)

    sub.expectNext(2)
    sub.expectNext(4)
    sub.expectNext(-1)  // recovered value
    sub.expectNext(8)
  }
}

Testing Fan-Out and Fan-In

Broadcast and merge graphs need multi-probe setups:

"A broadcast flow" should {
  "send elements to all sinks" in {
    val (pub, sub1, sub2) = RunnableGraph.fromGraph(GraphDSL.createGraph(
      TestSource[Int](),
      TestSink[Int](),
      TestSink[Int]()
    )((p, s1, s2) => (p, s1, s2)) { implicit b =>
      (src, sink1, sink2) =>
        import GraphDSL.Implicits._

        val bcast = b.add(Broadcast[Int](2))
        src ~> bcast ~> sink1
                bcast ~> sink2
        ClosedShape
    }).run()

    sub1.request(3)
    sub2.request(3)

    pub.sendNext(1)
    pub.sendNext(2)
    pub.sendNext(3)

    sub1.expectNext(1, 2, 3)
    sub2.expectNext(1, 2, 3)

    pub.sendComplete()
    sub1.expectComplete()
    sub2.expectComplete()
  }
}

Testing Stateful Stages

Custom GraphStage implementations require testing internal state transitions:

// A custom stage that counts elements and emits a summary on completion
class CountingStage[T] extends GraphStage[FlowShape[T, String]] {
  val in = Inlet[T]("CountingStage.in")
  val out = Outlet[String]("CountingStage.out")
  override val shape = FlowShape(in, out)

  override def createLogic(inheritedAttributes: Attributes) =
    new GraphStageLogic(shape) {
      private var count = 0

      setHandler(in, new InHandler {
        override def onPush(): Unit = {
          count += 1
          pull(in)
        }
        override def onUpstreamFinish(): Unit = {
          emit(out, s"Total: $count")
          completeStage()
        }
      })

      setHandler(out, new OutHandler {
        override def onPull(): Unit = pull(in)
      })
    }
}

"CountingStage" should {
  "emit total count on completion" in {
    val result = Source(1 to 100)
      .via(new CountingStage[Int])
      .runWith(Sink.head)

    result.futureValue shouldBe "Total: 100"
  }

  "handle empty stream" in {
    val result = Source.empty[Int]
      .via(new CountingStage[Int])
      .runWith(Sink.headOption)

    result.futureValue shouldBe Some("Total: 0")
  }
}

Materializer and ActorSystem Lifecycle

A common source of test flakiness is sharing a materializer across tests with different lifecycles. Keep it simple:

// Shared system for all tests in the file — shut down in afterAll
class MyStreamSpec extends TestKit(ActorSystem("test"))
    with AnyWordSpecLike
    with BeforeAndAfterAll {

  override def afterAll(): Unit = {
    TestKit.shutdownActorSystem(system)
  }
}

// Or use AkkaSpec from Akka's test utilities for automatic cleanup

Avoid creating a new ActorSystem per test — it's expensive. The pattern above (one system per spec class) is the right balance.

Integration: Streams + External Systems

When streams interact with databases or HTTP clients, use a test double at the boundary:

// Production code
object DatabaseWriter {
  def flow(repo: UserRepo): Flow[User, WriteResult, NotUsed] =
    Flow[User].mapAsync(4)(repo.save)
}

// Test with a mock repo
"DatabaseWriter" should {
  "write all users" in {
    val written = scala.collection.mutable.ListBuffer[User]()

    val mockRepo = new UserRepo {
      override def save(user: User): Future[WriteResult] = {
        written += user
        Future.successful(WriteResult(user.id, success = true))
      }
    }

    val users = List(User("alice"), User("bob"), User("carol"))

    val results = Source(users)
      .via(DatabaseWriter.flow(mockRepo))
      .runWith(Sink.seq)

    results.futureValue should have size 3
    written.map(_.name) should contain allOf("alice", "bob", "carol")
  }
}

Timing and Throttle Testing

Testing time-based operators like throttle and groupedWithin requires care:

"A throttled stream" should {
  "not exceed the rate limit" in {
    val startTime = System.currentTimeMillis()

    val result = Source(1 to 5)
      .throttle(2, 1.second)  // 2 elements per second
      .runWith(Sink.seq)

    result.futureValue(PatienceConfig(10.seconds)) shouldBe Seq(1, 2, 3, 4, 5)

    val elapsed = System.currentTimeMillis() - startTime
    // 5 elements at 2/sec takes at least 2 seconds
    assert(elapsed >= 2000, s"Expected >=2000ms but got ${elapsed}ms")
  }
}

For CI environments, avoid hard timing assertions when possible — use probe-based demand control instead, which doesn't depend on wall clock time.

Summary

Akka Streams testing follows a clear hierarchy of complexity. Start with runWith(Sink.seq).futureValue for simple flows. Reach for TestSource and TestSink probes when you need demand control or backpressure verification. Use graph DSL with multiple probes for fan-out/fan-in scenarios. Keep one ActorSystem per spec class to avoid startup overhead.

The key insight is that most stream bugs are backpressure bugs — elements emitted without downstream demand, buffers overflowing, or stages not propagating completion. The probe-based approach makes these visible because you control demand explicitly rather than letting the runtime consume elements as fast as possible.

Read more

Start now free