akka.stream.testkit.scaladsl.TestSource Scala Examples

package com.github.dnvriend.streams.sink

import akka.actor.{ Actor, ActorRef, Props }
import akka.stream.scaladsl.{ Sink, Source }
import akka.stream.testkit.TestPublisher
import akka.stream.testkit.scaladsl.TestSource
import akka.testkit.TestProbe
import com.github.dnvriend.streams.TestSpec
import scala.concurrent.duration._

import scala.reflect.ClassTag

// see: https://github.com/akka/akka/blob/4acc1cca6a27be0ff80f801de3640f91343dce94/akka-stream-tests/src/test/scala/akka/stream/scaladsl/ActorRefBackpressureSinkSpec.scala
object ActorRefWithAckTest {
  final val InitMessage = "start"
  final val CompleteMessage = "done"
  final val AckMessage = "ack"

  class Forwarder(ref: ActorRef) extends Actor {
    def receive = {
      case msg @ `InitMessage` ⇒
        sender() ! AckMessage
        ref forward msg
      case msg @ `CompleteMessage` ⇒
        ref forward msg
      case msg ⇒
        sender() ! AckMessage
        ref forward msg
    }
  }
}

class ActorRefWithAckTest extends TestSpec {
  import ActorRefWithAckTest._
  def createActor[A: ClassTag](testProbeRef: ActorRef): ActorRef =
    system.actorOf(Props(implicitly[ClassTag[A]].runtimeClass, testProbeRef))

  def withForwarder(xs: Int*)(f: TestProbe ⇒ Unit): Unit = {
    val tp = TestProbe()
    val ref = createActor[Forwarder](tp.ref)
    Source(xs.toList).runWith(Sink.actorRefWithAck(ref, InitMessage, AckMessage, CompleteMessage))
    try f(tp) finally killActors(ref)
  }

  def withTestPublisher[A](f: (TestPublisher.Probe[A], TestProbe, ActorRef) ⇒ Unit): Unit = {
    val tp = TestProbe()
    val ref = createActor[Forwarder](tp.ref)
    val pub: TestPublisher.Probe[A] = TestSource.probe[A].to(Sink.actorRefWithAck(ref, InitMessage, AckMessage, CompleteMessage)).run()
    try f(pub, tp, ref) finally killActors(ref)
  }

  it should "send the elements to the ActorRef" in {
    // which means that the forwarder actor that acts as a sink
    // will initially receive an InitMessage
    // next it will receive each `payload` element, here 1, 2 and 3,
    // finally the forwarder will receive the CompletedMessage, stating that
    // the producer completes the stream because there are no more elements (a finite stream)
    withForwarder(1, 2, 3) { tp ⇒
      tp.expectMsg(InitMessage)
      tp.expectMsg(1)
      tp.expectMsg(2)
      tp.expectMsg(3)
      tp.expectMsg(CompleteMessage)
      tp.expectNoMsg(100.millis)
    }
  }

  it should "send the elements to the ActorRef manually 1, 2 and 3" in {
    withTestPublisher[Int] { (pub, tp, _) ⇒
      pub.sendNext(1)
      tp.expectMsg(InitMessage)
      tp.expectMsg(1)

      pub.sendNext(2)
      tp.expectMsg(2)

      pub.sendNext(3)
      tp.expectMsg(3)

      pub.sendComplete()
      tp.expectMsg(CompleteMessage)
      tp.expectNoMsg(100.millis)
    }
  }

  it should "cancel stream when actor terminates" in {
    withTestPublisher[Int] { (pub, tp, ref) ⇒
      pub.sendNext(1)
      tp.expectMsg(InitMessage)
      tp.expectMsg(1)
      killActors(ref)
      pub.expectCancellation()
    }
  }
} 

package com.github.dnvriend.streams.sink

import akka.Done
import akka.actor.Actor.Receive
import akka.actor.{ ActorRef, Props }
import akka.event.LoggingReceive
import akka.stream.actor.ActorSubscriberMessage.{ OnComplete, OnError, OnNext }
import akka.stream.actor.{ ActorSubscriber, OneByOneRequestStrategy, RequestStrategy }
import akka.stream.scaladsl.{ Sink, Source }
import akka.stream.testkit.TestPublisher
import akka.stream.testkit.scaladsl.TestSource
import akka.testkit.TestProbe
import com.github.dnvriend.streams.TestSpec
import com.github.dnvriend.streams.sink.ActorSubscriberTest.TestActorSubscriber

import scala.concurrent.Future
import scala.reflect.ClassTag

object ActorSubscriberTest {
  final val OnNextMessage = "onNext"
  final val OnCompleteMessage = "onComplete"
  final val OnErrorMessage = "onError"

  class TestActorSubscriber(ref: ActorRef) extends ActorSubscriber {
    override protected val requestStrategy: RequestStrategy = OneByOneRequestStrategy
    override def receive: Receive = LoggingReceive {
      case OnNext(msg)    ⇒ ref ! OnNextMessage
      case OnComplete     ⇒ ref ! OnCompleteMessage
      case OnError(cause) ⇒ ref ! OnErrorMessage
    }
  }
}

//class ActorSubscriberTest extends TestSpec {
//  def withForwarder(xs: Int*)(f: TestProbe ⇒ Unit): Unit = {
//    val tp = TestProbe()
//    val ref = new TestActorSubscriber(tp.ref)
//    Source(xs.toList).to(Sink.actorSubscriber(Props())).mapMaterializedValue(_ ⇒ Future.successful[Done]).run()
//    try f(tp) finally killActors(ref)
//  }
//
//} 

package com.rbmhtechnology.eventuate.adapter.stream

import akka.actor._
import akka.stream._
import akka.stream.scaladsl.Keep
import akka.stream.testkit._
import akka.stream.testkit.scaladsl.{ TestSink, TestSource }
import akka.pattern
import akka.testkit._

import com.rbmhtechnology.eventuate.DurableEvent

import org.scalatest._

import scala.collection.immutable.Seq
import scala.concurrent._
import scala.concurrent.duration._
import scala.util.Random

class BatchWriteStageSpec extends TestKit(ActorSystem("test")) with WordSpecLike with Matchers with BeforeAndAfterAll with BeforeAndAfterEach {
  import BatchWriteStage.BatchWriter

  private val settings: DurableEventWriterSettings =
    new DurableEventWriterSettings(system.settings.config)

  implicit val materializer: Materializer =
    ActorMaterializer()

  private var src: TestPublisher.Probe[Seq[DurableEvent]] = _
  private var snk: TestSubscriber.Probe[Seq[DurableEvent]] = _

  override def beforeEach(): Unit = {
    val probes = TestSource.probe[Seq[DurableEvent]]
      .via(new BatchWriteStage(ec => writer(ec)))
      .toMat(TestSink.probe[Seq[DurableEvent]])(Keep.both)
      .run()

    src = probes._1
    snk = probes._2
  }

  override def afterEach(): Unit = {
    snk.cancel()
  }

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

  private def random: Int =
    Random.nextInt(100)

  private def writer(implicit ec: ExecutionContext): BatchWriter = events =>
    if (events.exists(_.payload == "boom")) Future(throw TestException)
    else pattern.after(random.millis, system.scheduler)(Future(events))

  "A BatchWriterStage" must {
    "write batches sequentially" in {
      val b1 = Seq("a", "b", "c").map(DurableEvent(_))
      val b2 = Seq("d", "e", "f").map(DurableEvent(_))
      val b3 = Seq("g", "h", "i").map(DurableEvent(_))

      snk.request(3)
      src.sendNext(b1)
      src.sendNext(b2)
      src.sendNext(b3)
      snk.expectNext() should be(b1)
      snk.expectNext() should be(b2)
      snk.expectNext() should be(b3)
    }
    "fail if the batch writer fails" in {
      val b = Seq("a", "boom", "c").map(DurableEvent(_))

      snk.request(3)
      src.sendNext(b)
      snk.expectError(TestException)
    }
  }
} 

package me.snov.akka.sqs.shape

import akka.Done
import akka.stream.scaladsl.{Keep, Sink}
import akka.stream.testkit.scaladsl.TestSource
import com.amazonaws.handlers.AsyncHandler
import com.amazonaws.services.sqs.model._
import me.snov.akka.sqs._
import me.snov.akka.sqs.client.SqsClient
import org.mockito.Mockito._
import org.mockito.ArgumentMatchers._
import org.mockito.invocation.InvocationOnMock
import org.mockito.stubbing.Answer
import org.scalatest.mockito.MockitoSugar.mock
import org.scalatest.{FlatSpec, Matchers}

import scala.concurrent.Await
import scala.concurrent.duration._

class SqsAckSinkShapeSpec extends FlatSpec with Matchers with DefaultTestContext {
  it should "delete messages on Ack" in {

    val sqsClient = mock[SqsClient]
    when(sqsClient.deleteAsync(any(), any())).thenAnswer(
      new Answer[Object] {
        override def answer(invocation: InvocationOnMock): Object = {
          val receiptHandle = invocation.getArgument[String](0)
          val callback = invocation.getArgument[AsyncHandler[DeleteMessageRequest, DeleteMessageResult]](1)
          callback.onSuccess(
            new DeleteMessageRequest().withReceiptHandle(receiptHandle),
            new DeleteMessageResult
          )
          None
        }
      }
    )

    val (probe, future) = TestSource.probe[MessageActionPair]
      .toMat(Sink.fromGraph(SqsAckSinkShape(sqsClient)))(Keep.both)
      .run()

    probe
      .sendNext((new Message().withReceiptHandle("123"), Ack()))
      .sendComplete()

    Await.result(future, 1.second) shouldBe Done
    verify(sqsClient, times(1)).deleteAsync(any(), any())
  }

  it should "requeue messages on RequeueWithDelay" in {

    val sqsClient = mock[SqsClient]
    when(sqsClient.sendWithDelayAsync(any[String], any[Int], any())).thenAnswer(
      new Answer[Object] {
        override def answer(invocation: InvocationOnMock): Object = {
          val body = invocation.getArgument[String](0)
          val delay = invocation.getArgument[Int](1)
          val callback = invocation.getArgument[AsyncHandler[SendMessageRequest, SendMessageResult]](2)
          callback.onSuccess(
            new SendMessageRequest().withMessageBody(body).withDelaySeconds(delay),
            new SendMessageResult().withMessageId("12345")
          )
          None
        }
      }
    )

    val (probe, future) = TestSource.probe[MessageActionPair]
      .toMat(Sink.fromGraph(SqsAckSinkShape(sqsClient)))(Keep.both)
      .run()

    probe
      .sendNext((new Message().withBody("foo"), RequeueWithDelay(9)))
      .sendComplete()

    Await.result(future, 100.second) shouldBe Done
    verify(sqsClient, times(1)).sendWithDelayAsync(any(), any(), any())
  }
} 

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.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 com.qubit.pubsub.akka

import akka.NotUsed
import akka.actor.ActorSystem
import akka.stream.scaladsl.{Keep, Sink}
import akka.stream.testkit.scaladsl.TestSource
import akka.stream.{ActorMaterializer, Attributes, Graph, SinkShape}
import com.google.common.base.Charsets
import com.qubit.pubsub.PubSubIntegrationTest
import com.qubit.pubsub.akka.attributes.{
  PubSubClientAttribute,
  PubSubStageBufferSizeAttribute
}
import com.qubit.pubsub.client.PubSubMessage
import org.scalatest.{BeforeAndAfterAll, FunSuite, Matchers}

import scala.concurrent.Await
import scala.concurrent.duration._
import scala.util.Try

class PubSubSinkIT
    extends FunSuite
    with Matchers
    with BeforeAndAfterAll
    with PubSubIntegrationTest {

  implicit val actorSystem = ActorSystem("pubsub-stream-test")
  implicit val materializer = ActorMaterializer()

  override def testName = "pubsubsink"

  override def beforeAll(): Unit = {
    Await.ready(client.createTopic(testTopic), timeout)
    Await
      .ready(client.createSubscription(testSubscription, testTopic), timeout)
  }

  override def afterAll(): Unit = {
    actorSystem.terminate()
    Await.ready(client.deleteSubscription(testSubscription), timeout)
    Await.ready(client.deleteTopic(testTopic), timeout)
  }

  test("PubSubSink success") {
    val sinkGraph: Graph[SinkShape[PubSubMessage], NotUsed] =
      new PubSubSink(testTopic, 1.second)
    val sinkAttributes = Attributes(
      List(PubSubClientAttribute(client), PubSubStageBufferSizeAttribute(30)))
    val pubsubSink = Sink.fromGraph(sinkGraph).withAttributes(sinkAttributes)

    val (pub, _) = TestSource
      .probe[Array[Byte]]
      .map(PubSubMessage(_))
      .toMat(pubsubSink)(Keep.both)
      .run()

    Range(0, 100)
      .map(i => s"xxx$i".getBytes(Charsets.UTF_8))
      .foreach(pub.sendNext)
    pub.sendComplete()

    // wait for buffers to flush
    Try(Thread.sleep(1000))

    val output = Await.result(client.pull(testSubscription, 100), timeout)
    client.ack(testSubscription, output.map(m => m.ackId))

    output should not be (null)
    output should have size (100)
    output
      .map(m => new String(m.payload.payload, Charsets.UTF_8))
      .forall(_.startsWith("xxx")) should be(true)
  }
} 

package cmwell.util.streams.test

import akka.stream._
import akka.stream.scaladsl.{GraphDSL, RunnableGraph, Source}
import akka.stream.testkit.scaladsl.{TestSink, TestSource}
import akka.stream.testkit.TestPublisher.{Probe => SrcProbe}
import akka.stream.testkit.TestSubscriber.{Probe => SnkProbe}
import cmwell.util.stream.MapInitAndLast
import scala.concurrent.duration.DurationInt

class MapInitAndLastTests extends StreamSpec {

  def generateGraph[In](): (SrcProbe[In],SnkProbe[(In,Boolean)]) = {
    val src = TestSource.probe[In]
    val snk = TestSink.probe[(In,Boolean)]
    RunnableGraph.fromGraph(GraphDSL.create(src, snk)((a, b) => (a, b)) {
      implicit b => {
        (s1, s2) => {
          import GraphDSL.Implicits._

          val mial = b.add(new MapInitAndLast[In, (In,Boolean)](_ -> false, _ -> true))

          s1 ~> mial ~> s2

          ClosedShape
        }
      }
    }).run()
  }

  describe("MapInitAndLast Stage"){
    it("should buffer a single element"){
      val (src,snk) = generateGraph[Int]()
      snk.request(99)
      src.sendNext(1)
      snk.expectNoMessage(300.millis)
      src.sendComplete()
      snk.expectNext((1,true))
      snk.expectComplete()
    }

    it("should treat last element differently") {
      val (src,snk) = generateGraph[Int]()
      snk.request(99)
      src.sendNext(1)
      snk.expectNoMessage(300.millis)
      src.sendNext(2)
      snk.expectNext((1,false))
      src.sendNext(3)
      snk.expectNext((2,false))
      src.sendComplete()
      snk.expectNext((3,true))
      snk.expectComplete()
    }

    it("should propagate back-pressure"){
      val (src,snk) = generateGraph[Int]()
      snk.ensureSubscription()
      src.sendNext(1)
      snk.expectNoMessage(300.millis)
      src.sendNext(1)
      snk.expectNoMessage(300.millis)
      src.sendComplete()
      snk.expectNoMessage(300.millis)
      snk.request(1)
      snk.expectNext((1,false))
      snk.request(1)
      snk.expectNext((1,true))
      snk.expectComplete()
    }
  }
} 

package cmwell.tools.data.utils.chunkers

import akka.stream.scaladsl.Keep
import akka.stream.testkit.scaladsl.{TestSink, TestSource}
import akka.util.ByteString
import cmwell.tools.data.helpers.BaseStreamSpec

import scala.concurrent.duration._

class GroupSpecAutoFusingOn  extends { val autoFusing = true  } with GroupChunkerSpec
class GroupSpecAutoFusingOff extends { val autoFusing = false } with GroupChunkerSpec

trait GroupChunkerSpec extends BaseStreamSpec {
  "GroupChunker" should "emit elements when new group has arrived" in {
    val (pub, sub) = TestSource.probe[String]
      .map(x => ByteString(x.toString))
      .via(GroupChunker(b => ByteString(b.size), 2.seconds)) // group byte-strings by size
      .map(_.map(_.utf8String))
      .toMat(TestSink.probe[Seq[String]])(Keep.both)
      .run()

    sub.request(100)
    pub.sendNext("hello")
    pub.sendNext("world")
    pub.sendNext("nba")
    pub.sendNext("ibm")
    pub.sendNext("what")
    pub.sendNext("is")
    pub.sendNext("life")
    pub.sendComplete()
    sub.expectNext(Seq("hello", "world"))
    sub.expectNext(Seq("nba", "ibm"))
    sub.expectNext(Seq("what"))
    sub.expectNext(Seq("is"))
    sub.expectNext(Seq("life"))
    sub.expectComplete()
  }

  it should "emit elements when time threshold has reached" in {
    val (pub, sub) = TestSource.probe[String]
      .map(x => ByteString(x.toString))
      .via(GroupChunker(b => ByteString(b.size), 2.seconds)) // group byte-strings by size
      .map(_.map(_.utf8String))
      .toMat(TestSink.probe[Seq[String]])(Keep.both)
      .run()

    sub.request(4)

    pub.sendNext("one")
    sub.expectNext(Seq("one"))

    pub.sendNext("two")
    sub.expectNext(Seq("two"))

    pub.sendNext("four")
    pub.sendNext("five")
    pub.sendComplete()
    sub.expectNext(Seq("four","five"))
    sub.expectComplete()
  }
} 

package org.squbs.stream

import akka.actor.ActorSystem
import akka.stream.ActorMaterializer
import akka.stream.testkit.scaladsl.TestSource
import akka.testkit.TestKit
import com.typesafe.config.{Config, ConfigFactory}
import org.scalatest._
import org.squbs.lifecycle.GracefulStop
import org.squbs.unicomplex._

import scala.concurrent.duration._

object UnicomplexActorPublisherJSpec {
  val myConfig: Config = ConfigFactory.parseString(
    """
      | squbs.actorsystem-name = UnicomplexActorPublisherJSpec
    """.stripMargin)
  val boot = UnicomplexBoot(myConfig).createUsing((name, config) => ActorSystem(name, config))
    .scanResources("/")
    .initExtensions
    .start()
}

final class UnicomplexActorPublisherJSpec extends TestKit(UnicomplexActorPublisherJSpec.boot.actorSystem)
    with FlatSpecLike with Matchers with BeforeAndAfterAll {
  implicit val materializer = ActorMaterializer()
  val duration = 10.second

  val in = TestSource.probe[String]

  // expose probe port(s)
  val mat = new UnicomplexActorPublisherJ(system).runnableGraph()
  val (pub, sub) = mat.toScala
  val (pubIn, pubTrigger) = pub.toScala

  override def afterAll(): Unit = {
    Unicomplex(system).uniActor ! GracefulStop
  }

  "UnicomplexTriggerJ" should "activate flow by unicomplex" in {
    // send 2 elements to in
    pubIn.sendNext("1")
    pubIn.sendNext("2")
    sub.request(2)
    sub.expectNext(duration, "1")
    sub.expectNext("2")

    // re-send Active to unicomplex trigger, flow continues
    sub.request(2)
    sub.expectNoMessage(remainingOrDefault)
    pubTrigger ! SystemState
    pubIn.sendNext("3")
    pubIn.sendNext("4")
    sub.expectNext("3", "4")
  }
} 

package org.squbs.stream

import akka.actor.ActorSystem
import akka.stream.ActorMaterializer
import akka.stream.scaladsl.Keep
import akka.stream.testkit.scaladsl.{TestSink, TestSource}
import akka.testkit.TestKit
import com.typesafe.config.{Config, ConfigFactory}
import org.scalatest._
import org.squbs.lifecycle.GracefulStop
import org.squbs.unicomplex._

import scala.concurrent.duration._

object UnicomplexActorPublisherSpec {
  val myConfig: Config = ConfigFactory.parseString(
    """
      | squbs.actorsystem-name = UnicomplexActorPublisherSpec
    """.stripMargin)
  val boot = UnicomplexBoot(myConfig).createUsing((name, config) => ActorSystem(name, config))
    .scanResources("/")
    .initExtensions
    .start()
}

final class UnicomplexActorPublisherSpec extends TestKit(UnicomplexActorPublisherSpec.boot.actorSystem)
    with FlatSpecLike with Matchers with BeforeAndAfterAll {

  implicit val materializer = ActorMaterializer()
  val duration = 10.second

  val in = TestSource.probe[String]

  // expose probe port(s)
  val ((pubIn, pubTrigger), sub) = LifecycleManaged().source(in).toMat(TestSink.probe[String](system))(Keep.both).run()

  override def afterAll(): Unit = {
    Unicomplex(system).uniActor ! GracefulStop
  }

  "UnicomplexTrigger" should "activate flow by unicomplex" in {
    // send 2 elements to in
    pubIn.sendNext("1")
    pubIn.sendNext("2")
    sub.request(2)
    sub.expectNext(duration, "1")
    sub.expectNext("2")

    // re-send Active to unicomplex trigger, flow continues
    sub.request(2)
    sub.expectNoMessage(remainingOrDefault)
    pubTrigger ! SystemState
    pubIn.sendNext("3")
    pubIn.sendNext("4")
    sub.expectNext("3", "4")
  }
} 

package uk.gov.hmrc.akka.xml

import akka.util.ByteString
import org.scalatest.{FlatSpec, Matchers}
import scala.concurrent.duration._
import akka.actor.ActorSystem
import akka.stream.ActorMaterializer
import akka.stream.scaladsl._
import akka.stream.testkit.scaladsl.TestSink
import akka.stream.testkit.scaladsl.TestSource

class ParsingStageSpec extends FlatSpec {


  def createStream(instructions:Seq[XMLInstruction],validationMaxSize: Option[Int] = None) = {
    val as = ActorSystem("CompleteChunkSpec")
    val am = ActorMaterializer()(as)
    val source = TestSource.probe[ParsingData](as)
    val sink = TestSink.probe[(ByteString, Set[XMLElement])](as)
    val chunk = ParsingStage.parser(instructions)

    //source.via(chunk).alsoTo(Sink.foreach(a => println(">> " + a._1.decodeString("UTF-8") + " | " + a._2))).toMat(sink)(Keep.both).run()(am)
    source.via(chunk).toMat(sink)(Keep.both).run()(am)
  }

  def getEmptyResult(in:String):(ByteString, Set[XMLElement]) = {
    (ByteString(in), Set.empty[XMLElement])
  }

  it should "Extract XMLInstruction from a xml broken into pieces (even xml tags are broken up)" in {
    val idHeader = XMLExtract(List("xml","header","id"))
    val aaHeader = XMLExtract(List("xml","header","aa"))

    //This is our entire test xml: <xml><header><id>Joska</id><aa>Pista</aa><bb>Miska</bb><cc/><dd/></header></xml>  //The test xml

    val (pub,sub) = createStream(Seq(idHeader,aaHeader))
    sub.request(10)
    pub.sendNext(ParsingData(ByteString("<xml><hea"),Set.empty,5))
    sub.expectNext(getEmptyResult("<xml><hea"))
    pub.sendNext(ParsingData(ByteString("der><id>Jo"),Set.empty,19))
    sub.expectNext(getEmptyResult("der><id>Jo"))
    pub.sendNext(ParsingData(ByteString("ska</i"),Set.empty,26))
    sub.expectNext(getEmptyResult("ska</i"))
    pub.sendNext(ParsingData(ByteString("d><aa>Pista</a"),Set.empty,32))
    sub.expectNext((ByteString("d><aa>Pista</a"),Set(XMLElement(List("xml", "header", "id"),Map(),Some("Joska")))))
    pub.sendNext(ParsingData(ByteString("a><bb>Mis"),Set.empty,38))
    sub.expectNext((ByteString("a><bb>Mis"),Set(XMLElement(List("xml", "header", "aa"),Map(),Some("Pista")))))
    pub.sendNext(ParsingData(ByteString("ka</bb><cc/"),Set.empty,50))
    sub.expectNext(getEmptyResult("ka</bb><cc/"))
    pub.sendNext(ParsingData(ByteString("><dd"),Set.empty,57))
    sub.expectNext(getEmptyResult("><dd"))
    pub.sendNext(ParsingData(ByteString("/></header></xml>"),Set.empty,57))
    sub.expectNext(getEmptyResult("/></header></xml>"))
    pub.sendComplete()
  }

} 

package uk.gov.hmrc.akka.xml

import akka.util.ByteString
import org.scalatest.{FlatSpec, Matchers}
import scala.concurrent.duration._
import akka.actor.ActorSystem
import akka.stream.ActorMaterializer
import akka.stream.scaladsl._
import akka.stream.testkit.scaladsl.TestSink
import akka.stream.testkit.scaladsl.TestSource

class CompleteChunkSpec extends FlatSpec {

  def createStream() = {
    val as = ActorSystem("CompleteChunkSpec")
    val am = ActorMaterializer()(as)
    val source = TestSource.probe[ByteString](as)
    val sink = TestSink.probe[ParsingData](as)
    val chunk = CompleteChunkStage.parser()

    //source.map(a => {println("<< " + a.decodeString("UTF-8"));a}).via(chunk).alsoTo(Sink.foreach(a => println(">> " + a))).toMat(sink)(Keep.both).run()(am)  //Use for debugging
    source.via(chunk).toMat(sink)(Keep.both).run()(am)
  }


  it should "only let whole xml tags through" in {
    //This is our entire test xml: <xml><header><id>Joska</id><aa>Pista</aa><bb>Miska</bb></header></xml>
    val (pub,sub) = createStream()
    sub.request(20)
    pub.sendNext(ByteString("<xml><hea"))
    sub.expectNext(ParsingData(ByteString("<xml>"), Set.empty, 5))
    pub.sendNext(ByteString("der><id>Jo"))
    sub.expectNext(ParsingData(ByteString("<header><id>Jo"), Set.empty, 19))
    pub.sendNext(ByteString("ska</i"))
    sub.expectNext(ParsingData(ByteString("ska"), Set.empty, 22))
    pub.sendNext(ByteString("d><aa>Pista</a"))
    sub.expectNext(ParsingData(ByteString("</id><aa>Pista"), Set(), 36))
    pub.sendNext(ByteString("a><bb>Mis"))
    sub.expectNext(ParsingData(ByteString("</aa><bb>Mis"), Set(), 48))
    pub.sendNext(ByteString("ka</bb></he"))
    sub.expectNext(ParsingData(ByteString("ka</bb>"), Set(), 55))
    pub.sendNext(ByteString("ader></xml>"))
    sub.expectNext(ParsingData(ByteString("</header></xml>"), Set(), 70))
    pub.sendComplete()
    sub.expectNext(ParsingData(ByteString.empty,  Set(XMLElement(List(),Map("Stream Size" -> "70"), Some("Stream Size"))), 70))
    sub.expectComplete()
  }




}