/* * Copyright 2016 Dennis Vriend * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package com.github.dnvriend.streams.customstage import akka.stream.{ Outlet, Inlet, Attributes, FlowShape } import akka.stream.stage._ import akka.stream.testkit.scaladsl.TestSink import com.github.dnvriend.streams.TestSpec class Ex2CustomMapTest extends TestSpec { "CustomMapStage" should "be implemented with a PushPullStage" in { /** * A custom map stage that takes elements of type `A`, and converts the element to type 'B'. * The `CustomMapStage` does this by using the supplied function that converts `A` into `B` * * The PushPull stage has input and output ports. The `input` ports are the callback methods, * `onPush(elem,ctx)` and `onPull(ctx)`. The `output` ports are implemented as methods on * on the `Context` object that is supplied as a parameter on the event handler methods below. * * By calling exactly one "output port" method we wire up these four ports in various ways. * `Calling` an output port is called wireing, because the element that has been supplied by the akka-streams * runtime to the PushPullStage by calling the `onPush(in,ctx)` or the `onPull(ctx)` method is passed to exactly * one output port. * * The CustomMapStage calls `ctx.push()` from the `onPush()` event handler and it also calls * `ctx.pull()` from the `onPull` handler resulting in the conceptual wiring: * * +---------------------------------+ * | onPush(in,ctx) ctx.push(out) | * O-------------> f(in) -->---------O * | | * O-------------<------<------------O * | ctx.pull() onPull(ctx) | * +---------------------------------+ * * Map is a typical example of a one-to-one transformation of a stream, the element will be processed and * forwarded. * * Note: * A map is just a function from f: A => B, so we will extend the PushPullStage to create this map function */ class CustomMapStage[A, B](f: A ⇒ B) extends PushPullStage[A, B] { override def onPush(elem: A, ctx: Context[B]): SyncDirective = ctx.push(f(elem)) // transform the element and pushes it downstream when there is demand for it override def onPull(ctx: Context[B]): SyncDirective = ctx.pull() // request for more elements from upstream (other stages before us) } /** * To use the custom transformation stage, call `transform()` on a `Flow` or `Source` * which takes a factory function returning a Stage: `f: () => Stage` * * In the example below we use a TestProbe as the Source that generates demand and * does assertions. */ withIterator(1) { src ⇒ src.transform(() ⇒ new CustomMapStage(_ * 2)) .take(2) .runWith(TestSink.probe[Int]) .request(Int.MaxValue) .expectNext(2, 4) .expectComplete() } } it should "also be implemented using the PushStage" in { /** * When the stage just propagates the pull upwards to the `previous` stage, it is not necessary to override * the onPull handler at all. Such transformations are better of by extending the `PushStage`. The conceptual * mapping will still be the same. */ class CustomMapStage[A, B](f: A ⇒ B) extends PushStage[A, B] { override def onPush(elem: A, ctx: Context[B]): SyncDirective = ctx.push(f(elem)) // transform the element and pushes it downstream when there is demand for it } /** * To use the custom transformation stage, call `transform()` on a `Flow` or `Source` * which takes a factory function returning a Stage: `f: () => Stage` * * In the example below we use a TestProbe as the Source that generates demand and * does assertions. */ /** * To use the custom transformation stage, call `transform()` on a `Flow` or `Source` * which takes a factory function returning a Stage: `f: () => Stage` * * In the example below we use a TestProbe as the Source that generates demand and * does assertions. */ withIterator(1) { src ⇒ src.transform(() ⇒ new CustomMapStage(_ * 2)) .take(2) .runWith(TestSink.probe[Int]) .request(Int.MaxValue) .expectNext(2, 4) .expectComplete() } } it should "also be implemented as a GraphStage" in { class CustomMapStage[A, B](f: A ⇒ B) extends GraphStage[FlowShape[A, B]] { val in = Inlet[A]("Map.in") val out = Outlet[B]("Map.out") override def shape: FlowShape[A, B] = FlowShape.of(in, out) override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) { setHandler(in, new InHandler { override def onPush(): Unit = push(out, f(grab(in))) }) setHandler(out, new OutHandler { override def onPull(): Unit = pull(in) }) } } withIterator(1) { src ⇒ src.via(new CustomMapStage(_ * 2)) .take(2) .runWith(TestSink.probe[Int]) .request(Int.MaxValue) .expectNext(2, 4) .expectComplete() } } }