Test coverage for stages using invokeWithFeedback #23953

This commit is contained in:
Johan Andrén 2017-12-04 16:34:16 +01:00 committed by GitHub
parent 891cf30348
commit a380d01337
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 30 additions and 20 deletions

View file

@ -336,8 +336,6 @@ import scala.collection.generic.CanBuildFrom
override def toString: String = "QueueSink"
override def createLogicAndMaterializedValue(inheritedAttributes: Attributes) = {
var logicCallback: AsyncCallback[Output[T]] = null
val stageLogic = new GraphStageLogic(shape) with InHandler with SinkQueueWithCancel[T] {
type Received[E] = Try[Option[E]]
@ -397,18 +395,17 @@ import scala.collection.generic.CanBuildFrom
override def onUpstreamFinish(): Unit = enqueueAndNotify(Success(None))
override def onUpstreamFailure(ex: Throwable): Unit = enqueueAndNotify(Failure(ex))
logicCallback = callback
setHandler(in, this)
// SinkQueueWithCancel impl
override def pull(): Future[Option[T]] = {
val p = Promise[Option[T]]
logicCallback.invokeWithFeedback(Pull(p))
callback.invokeWithFeedback(Pull(p))
.onFailure { case NonFatal(e) p.tryFailure(e) }(akka.dispatch.ExecutionContexts.sameThreadExecutionContext)
p.future
}
override def cancel(): Unit = {
logicCallback.invoke(QueueSink.Cancel)
callback.invoke(QueueSink.Cancel)
}
}