akka.stream.testkit.TestSubscriber Scala Examples

// Copyright (c) 2020 Digital Asset (Switzerland) GmbH and/or its affiliates. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

package com.daml.platform.apiserver.services.tracking

import akka.NotUsed
import akka.stream.OverflowStrategy
import akka.stream.scaladsl.{Keep, Source, SourceQueueWithComplete}
import akka.stream.testkit.TestSubscriber
import akka.stream.testkit.scaladsl.TestSink
import com.daml.ledger.api.testing.utils.{
  AkkaBeforeAndAfterAll,
  IsStatusException,
  TestingException
}
import com.daml.ledger.api.v1.command_service.SubmitAndWaitRequest
import com.daml.ledger.api.v1.commands.Commands
import com.daml.ledger.api.v1.completion.Completion
import com.daml.dec.DirectExecutionContext
import com.google.rpc.status.{Status => RpcStatus}
import io.grpc.Status
import org.scalatest.concurrent.ScalaFutures
import org.scalatest.{BeforeAndAfterEach, Matchers, Succeeded, WordSpec}

import scala.concurrent.ExecutionContext.Implicits.global

class TrackerImplTest
    extends WordSpec
    with Matchers
    with BeforeAndAfterEach
    with ScalaFutures
    with AkkaBeforeAndAfterAll {

  private var sut: Tracker = _
  private var consumer: TestSubscriber.Probe[NotUsed] = _
  private var queue: SourceQueueWithComplete[TrackerImpl.QueueInput] = _

  private def input(cid: Int) = SubmitAndWaitRequest(Some(Commands(commandId = cid.toString)))

  override protected def beforeEach(): Unit = {
    val (q, sink) = Source
      .queue[TrackerImpl.QueueInput](1, OverflowStrategy.dropNew)
      .map { in =>
        in.context.success(Completion(in.value.getCommands.commandId, Some(RpcStatus())))
        NotUsed
      }
      .toMat(TestSink.probe[NotUsed])(Keep.both)
      .run()
    queue = q
    sut = new TrackerImpl(q)
    consumer = sink
  }

  override protected def afterEach(): Unit = {
    consumer.cancel()
    queue.complete()
  }

  "Tracker Implementation" when {

    "input is submitted, and the queue is available" should {

      "work successfully" in {

        val resultF1 = sut.track(input(1))
        consumer.requestNext()
        val resultF = resultF1.flatMap(_ => sut.track(input(2)))(DirectExecutionContext)
        consumer.requestNext()
        whenReady(resultF)(_ => Succeeded)
      }
    }

    "input is submitted, and the queue is backpressuring" should {

      "return a RESOURCE_EXHAUSTED error" in {

        sut.track(input(1))
        whenReady(sut.track(input(2)).failed)(IsStatusException(Status.RESOURCE_EXHAUSTED))
      }
    }

    "input is submitted, and the queue has been completed" should {

      "return an ABORTED error" in {

        queue.complete()
        whenReady(sut.track(input(2)).failed)(IsStatusException(Status.ABORTED))
      }
    }

    "input is submitted, and the queue has failed" should {

      "return an ABORTED error" in {

        queue.fail(TestingException("The queue fails with this error."))
        whenReady(sut.track(input(2)).failed)(IsStatusException(Status.ABORTED))
      }
    }
  }
} 

package com.github.dnvriend

import akka.NotUsed
import akka.stream.scaladsl.Source
import akka.stream.testkit.TestSubscriber
import akka.stream.testkit.scaladsl.TestSink
import scala.collection.immutable.Seq
import scala.concurrent.duration._

trait AkkaStreamUtils { _: TestSpec ⇒
  implicit class SourceOps[A, M](src: Source[A, M]) {
    def testProbe(f: TestSubscriber.Probe[A] ⇒ Unit): Unit = {
      val tp = src.runWith(TestSink.probe(system))
      tp.within(10.seconds)(f(tp))
    }
  }

  def withIteratorSrc[T](start: Int = 0)(f: Source[Int, NotUsed] ⇒ Unit): Unit =
    f(Source.fromIterator(() ⇒ Iterator from start))

  def fromCollectionProbe[A](xs: Seq[A])(f: TestSubscriber.Probe[A] ⇒ Unit): Unit =
    f(Source(xs).runWith(TestSink.probe(system)))
} 

package com.github.dnvriend.streams

import akka.NotUsed
import akka.actor.{ ActorRef, ActorSystem, PoisonPill }
import akka.event.{ Logging, LoggingAdapter }
import akka.stream.Materializer
import akka.stream.scaladsl.Source
import akka.stream.testkit.TestSubscriber
import akka.stream.testkit.scaladsl.TestSink
import akka.testkit.TestProbe
import akka.util.Timeout
import com.github.dnvriend.streams.util.ClasspathResources
import org.scalatest._
import org.scalatest.concurrent.{ Eventually, ScalaFutures }
import org.scalatestplus.play.guice.GuiceOneServerPerSuite
import play.api.inject.BindingKey
import play.api.libs.json.{ Format, Json }
import play.api.test.WsTestClient

import scala.collection.immutable._
import scala.concurrent.duration._
import scala.concurrent.{ ExecutionContext, Future }
import scala.reflect.ClassTag
import scala.util.Try

object Person {
  implicit val format: Format[Person] = Json.format[Person]
}

final case class Person(firstName: String, age: Int)

class TestSpec extends FlatSpec
    with Matchers
    with GivenWhenThen
    with OptionValues
    with TryValues
    with ScalaFutures
    with WsTestClient
    with BeforeAndAfterAll
    with BeforeAndAfterEach
    with Eventually
    with ClasspathResources
    with GuiceOneServerPerSuite {

  def getComponent[A: ClassTag] = app.injector.instanceOf[A]
  def getNamedComponent[A](name: String)(implicit ct: ClassTag[A]): A =
    app.injector.instanceOf[A](BindingKey(ct.runtimeClass.asInstanceOf[Class[A]]).qualifiedWith(name))

  // set the port number of the HTTP server
  override lazy val port: Int = 8081
  implicit val timeout: Timeout = 1.second
  implicit val pc: PatienceConfig = PatienceConfig(timeout = 30.seconds, interval = 300.millis)
  implicit val system: ActorSystem = getComponent[ActorSystem]
  implicit val ec: ExecutionContext = getComponent[ExecutionContext]
  implicit val mat: Materializer = getComponent[Materializer]
  val log: LoggingAdapter = Logging(system, this.getClass)

  // ================================== Supporting Operations ====================================
  def id: String = java.util.UUID.randomUUID().toString

  implicit class FutureToTry[T](f: Future[T]) {
    def toTry: Try[T] = Try(f.futureValue)
  }

  implicit class SourceOps[A](src: Source[A, NotUsed]) {
    def testProbe(f: TestSubscriber.Probe[A] ⇒ Unit): Unit =
      f(src.runWith(TestSink.probe(system)))
  }

  def withIterator[T](start: Int = 0)(f: Source[Int, NotUsed] ⇒ T): T =
    f(Source.fromIterator(() ⇒ Iterator from start))

  def fromCollection[A](xs: Iterable[A])(f: TestSubscriber.Probe[A] ⇒ Unit): Unit =
    f(Source(xs).runWith(TestSink.probe(system)))

  def killActors(refs: ActorRef*): Unit = {
    val tp = TestProbe()
    refs.foreach { ref ⇒
      tp watch ref
      tp.send(ref, PoisonPill)
      tp.expectTerminated(ref)
    }
  }
} 

package com.rbmhtechnology.eventuate.adapter.stream

import akka.actor.ActorSystem
import akka.pattern.AskTimeoutException
import akka.stream.scaladsl._
import akka.stream.testkit.TestSubscriber
import akka.stream.testkit.scaladsl.TestSink
import akka.stream._
import akka.testkit._

import com.rbmhtechnology.eventuate.ReplicationProtocol._
import com.rbmhtechnology.eventuate.VectorTime
import com.typesafe.config.ConfigFactory

import org.scalatest._

object ProgressSourceSpec {
  val SrcLogId = "A"

  val config = ConfigFactory.parseString("eventuate.log.read-timeout = 500ms")
}

class ProgressSourceSpec extends TestKit(ActorSystem("test", ProgressSourceSpec.config)) with WordSpecLike with Matchers with BeforeAndAfterAll with BeforeAndAfterEach {
  import ProgressSourceSpec._

  implicit val materializer: Materializer =
    ActorMaterializer()

  private var log: TestProbe = _
  private var snk: TestSubscriber.Probe[Long] = _

  override def beforeEach(): Unit = {
    log = TestProbe()
    snk = Source.fromGraph(ProgressSource(SrcLogId, log.ref)).toMat(TestSink.probe[Long])(Keep.right).run()
  }

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

  "A ProgressSource" must {
    "complete after emitting a stored replication progress" in {
      snk.request(1)
      log.expectMsg(GetReplicationProgress(SrcLogId))
      log.sender() ! GetReplicationProgressSuccess(SrcLogId, 17, VectorTime.Zero)
      snk.expectNext(17)
      snk.expectComplete()
    }
    "fail if replication progress reading fails" in {
      snk.request(1)
      log.expectMsg(GetReplicationProgress(SrcLogId))
      log.sender() ! GetReplicationProgressFailure(TestException)
      snk.expectError(TestException)
    }
    "fail on replication progress reading timeout" in {
      snk.request(1)
      log.expectMsg(GetReplicationProgress(SrcLogId))
      snk.expectError should be(an[AskTimeoutException])
    }
  }
} 

package akka.stream.integration
package activemq

import akka.NotUsed
import akka.actor.ActorRef
import akka.stream.integration.PersonDomain.Person
import akka.stream.scaladsl.{ Flow, Keep }
import akka.stream.testkit.scaladsl.{ TestSink, TestSource }
import akka.stream.testkit.{ TestPublisher, TestSubscriber }
import akka.testkit.TestActor.AutoPilot
import akka.testkit.TestProbe
import JsonCamelMessageExtractor._
import JsonCamelMessageBuilder._

import scala.util.{ Failure, Success, Try }

  implicit def function1ToAutoPilot[S, T](f: S => T): AutoPilot = new AutoPilot {
    override def run(sender: ActorRef, msg: Any): AutoPilot = msg match {
      case s: S =>
        val tryT: Try[T] = Try(f(s))
        tryT match {
          case Success(t) =>
            sender ! t
            function1ToAutoPilot(f)
          case Failure(f) =>
            fail(s"Failed to apply supplied function to received message: $s", f)
        }
      case _ =>
        fail(s"Received message is not of the required type: $msg")
    }
  }
} 

package com.github.dnvriend

import akka.actor.{ ActorRef, ActorSystem, PoisonPill }
import akka.stream.Materializer
import akka.stream.scaladsl.Source
import akka.stream.testkit.TestSubscriber
import akka.stream.testkit.scaladsl.TestSink
import akka.testkit.TestProbe
import akka.util.Timeout
import org.scalatest._
import org.scalatest.concurrent.{ Eventually, ScalaFutures }
import org.scalatestplus.play.guice.GuiceOneServerPerSuite
import play.api.inject.BindingKey
import play.api.test.WsTestClient

import scala.concurrent.duration._
import scala.concurrent.{ ExecutionContext, Future }
import scala.reflect.ClassTag
import scala.util.Try

class TestSpec extends FlatSpec
    with Matchers
    with GivenWhenThen
    with OptionValues
    with TryValues
    with ScalaFutures
    with WsTestClient
    with BeforeAndAfterAll
    with BeforeAndAfterEach
    with Eventually
    with GuiceOneServerPerSuite {

  def getComponent[A: ClassTag] = app.injector.instanceOf[A]

  def getAnnotatedComponent[A](name: String)(implicit ct: ClassTag[A]): A =
    app.injector.instanceOf[A](BindingKey(ct.runtimeClass.asInstanceOf[Class[A]]).qualifiedWith(name))

  // set the port number of the HTTP server
  override lazy val port: Int = 8080
  implicit val timeout: Timeout = 10.seconds
  implicit val pc: PatienceConfig = PatienceConfig(timeout = 30.seconds, interval = 300.millis)
  implicit val system: ActorSystem = getComponent[ActorSystem]
  implicit val ec: ExecutionContext = getComponent[ExecutionContext]
  implicit val mat: Materializer = getComponent[Materializer]

  // ================================== Supporting Operations ====================================
  implicit class PimpedByteArray(self: Array[Byte]) {
    def getString: String = new String(self)
  }

  implicit class PimpedFuture[T](self: Future[T]) {
    def toTry: Try[T] = Try(self.futureValue)
  }

  implicit class SourceOps[A](src: Source[A, _]) {
    def testProbe(f: TestSubscriber.Probe[A] => Unit): Unit =
      f(src.runWith(TestSink.probe(system)))
  }

  def killActors(actors: ActorRef*): Unit = {
    val tp = TestProbe()
    actors.foreach { (actor: ActorRef) =>
      tp watch actor
      actor ! PoisonPill
      tp.expectTerminated(actor)
    }
  }

  override protected def beforeEach(): Unit = {
  }
} 

package docs.scaladsl.mb

import com.lightbend.lagom.scaladsl.api.broker.Topic
import com.lightbend.lagom.scaladsl.server.LagomApplication
import com.lightbend.lagom.scaladsl.server.LagomApplicationContext
import com.lightbend.lagom.scaladsl.server.LagomServer
import com.lightbend.lagom.scaladsl.server.LocalServiceLocator
import com.lightbend.lagom.scaladsl.testkit.ServiceTest
import com.lightbend.lagom.scaladsl.testkit.TestTopicComponents
import play.api.libs.ws.ahc.AhcWSComponents
import org.scalatest.matchers.should.Matchers
import org.scalatest.wordspec.AsyncWordSpec
import akka.NotUsed
import akka.Done
import akka.stream.scaladsl.Source
import akka.stream.testkit.scaladsl.TestSink
import akka.stream.testkit.TestSubscriber
import akka.stream.testkit.TestSubscriber.Probe

abstract class PublishApplication(context: LagomApplicationContext)
    extends LagomApplication(context)
    with AhcWSComponents {
  override lazy val lagomServer = serverFor[service.PublishService](new service.PublishServiceImpl())
}

package service {
  import com.lightbend.lagom.scaladsl.api.Service
  import com.lightbend.lagom.scaladsl.broker.TopicProducer

  object PublishService {
    val TOPIC_NAME = "events"
  }

  trait PublishService extends Service {
    final override def descriptor = {
      import Service._
      named("brokerdocs")
        .withTopics(topic(PublishService.TOPIC_NAME, events))
        .withAutoAcl(true)
    }

    def events(): Topic[PubMessage]
  }

  case class PubMessage(message: String)

  object PubMessage {
    import play.api.libs.json.Format
    import play.api.libs.json.Json
    implicit val format: Format[PubMessage] = Json.format[PubMessage]
  }

  class PublishServiceImpl() extends PublishService {
    override def events(): Topic[PubMessage] =
      TopicProducer.singleStreamWithOffset { offset =>
        Source((1 to 10)).map(i => (PubMessage(s"msg $i"), offset))
      }
  }
}

class PublishServiceSpec extends AsyncWordSpec with Matchers {
  import service._

  //#topic-test-publishing-into-a-topic
  "The PublishService" should {
    "publish events on the topic" in ServiceTest.withServer(ServiceTest.defaultSetup) { ctx =>
      new PublishApplication(ctx) with LocalServiceLocator with TestTopicComponents
    } { server =>
      implicit val system = server.actorSystem
      implicit val mat    = server.materializer

      val client: PublishService = server.serviceClient.implement[PublishService]
      val source                 = client.events().subscribe.atMostOnceSource
      source
        .runWith(TestSink.probe[PubMessage])
        .request(1)
        .expectNext should ===(PubMessage("msg 1"))
    }
  }
  //#topic-test-publishing-into-a-topic
} 

package akka.persistence.couchbase.scaladsl

import akka.persistence.couchbase.TestActor
import akka.stream.testkit.TestSubscriber
import akka.stream.testkit.scaladsl.TestSink
import akka.testkit.TestProbe

import scala.concurrent.duration._

class PersistenceIdsQuerySpec extends AbstractCouchbaseSpec("PersistenceIdsQuerySpec") {
  "currentPersistenceIds" must {
    "work" in {
      val senderProbe = TestProbe()
      implicit val sender = senderProbe.ref
      val pa1 = system.actorOf(TestActor.props("p1"))
      pa1 ! "p1-evt-1"
      senderProbe.expectMsg("p1-evt-1-done")
      val pa2 = system.actorOf(TestActor.props("p2"))
      pa2 ! "p2-evt-1"
      senderProbe.expectMsg("p2-evt-1-done")

      awaitAssert(
        {
          val probe: TestSubscriber.Probe[String] = queries.currentPersistenceIds().runWith(TestSink.probe)

          probe.requestNext("p1")
          probe.requestNext("p2")
          probe.expectComplete()
        },
        readOurOwnWritesTimeout
      )
    }
  }

  "live persistenceIds" must {
    "show new persistence ids" in {
      val senderProbe = TestProbe()
      implicit val sender = senderProbe.ref

      val queryProbe: TestSubscriber.Probe[String] =
        queries.persistenceIds().runWith(TestSink.probe)

      queryProbe.request(10)

      val pa3 = system.actorOf(TestActor.props("p3"))
      pa3 ! "p3-evt-1"
      senderProbe.expectMsg("p3-evt-1-done")

      awaitAssert({
        queryProbe.expectNext("p3")
      }, 5.seconds)

      val pa4 = system.actorOf(TestActor.props("p4"))
      pa4 ! "p4-evt-1"
      senderProbe.expectMsg("p4-evt-1-done")

      // we shouldn't see p3 again
      queryProbe.expectNext("p4")
      // also not after p4 (it could come out of order)
      queryProbe.expectNoMessage(noMsgTimeout)
    }
  }
} 

package com.github.krasserm.ases

import akka.stream.ActorMaterializer
import akka.stream.scaladsl.{Flow, Keep}
import akka.stream.testkit.scaladsl.{TestSink, TestSource}
import akka.stream.testkit.{TestPublisher, TestSubscriber}
import akka.testkit.TestKit
import org.scalatest.{BeforeAndAfterAll, Suite}

import scala.collection.immutable.Seq

trait StreamSpec extends BeforeAndAfterAll { this: TestKit with Suite =>
  implicit val materializer = ActorMaterializer()

  val emitterId = "emitter"

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

  def probes[I, O, M](flow: Flow[I, O, M]): (TestPublisher.Probe[I], TestSubscriber.Probe[O]) =
    TestSource.probe[I].viaMat(flow)(Keep.left).toMat(TestSink.probe[O])(Keep.both).run()

  def durables[A](emitted: Seq[Emitted[A]], offset: Int = 0): Seq[Durable[A]] =
    emitted.zipWithIndex.map { case (e, i) => e.durable(i + offset) }
} 

package akka.persistence.cassandra

import akka.persistence.cassandra.query.TestActor
import akka.persistence.cassandra.query._
import akka.persistence.journal.Tagged
import akka.persistence.query.NoOffset
import akka.stream.testkit.TestSubscriber
import akka.stream.testkit.scaladsl.TestSink

import scala.collection.immutable
import scala.concurrent.Future

class EventsByTagStressSpec extends CassandraSpec(s"""
    akka.persistence.cassandra {
      events-by-tag {
        max-message-batch-size = 25
      }
    }
  """) {

  implicit val ec = system.dispatcher

  val writers = 10
  val readers = 20
  val messages = 5000

  "EventsByTag" must {

    "work under load" in {
      val pas = (0 until writers).map { i =>
        system.actorOf(TestActor.props(s"pid$i"))
      }

      val eventsByTagQueries: immutable.Seq[(Int, TestSubscriber.Probe[(String, Int)])] = (0 until readers).map { i =>
        val probe = queryJournal
          .eventsByTag("all", NoOffset)
          .map(i => {
            (i.persistenceId, i.event.asInstanceOf[Int])
          })
          .runWith(TestSink.probe)
        (i, probe)
      }

      system.log.info("Started events by tag queries")

      val writes: Future[Unit] = Future {
        system.log.info("Sending messages")
        (0 until messages).foreach { i =>
          pas.foreach(ref => {
            ref ! Tagged(i, Set("all"))
            expectMsg(s"$i-done")
          })
        }
        system.log.info("Sent messages")
      }
      writes.onComplete(result => system.log.info("{}", result))

      system.log.info("Reading messages")
      var latestValues: Map[(Int, String), Int] = Map.empty.withDefault(_ => -1)
      (0 until messages).foreach { _ =>
        (0 until writers).foreach { _ =>
          eventsByTagQueries.foreach {
            case (probeNr, probe) =>
              // should be in order per persistence id per probe
              val (pid, msg) = probe.requestNext()
              latestValues((probeNr, pid)) shouldEqual (msg - 1)
              latestValues += (probeNr, pid) -> msg
          }
        }
      }
      system.log.info("Received all messages {}", latestValues)
    }

  }
}