akka.http.scaladsl.model.ws.Message Scala Examples

package com.wavesplatform.dex.api.ws.connection

import java.util.concurrent.ConcurrentLinkedQueue

import akka.Done
import akka.actor.{ActorRef, ActorSystem, Status}
import akka.http.scaladsl.Http
import akka.http.scaladsl.model.ws.{BinaryMessage, Message, TextMessage, WebSocketRequest}
import akka.stream.scaladsl.{Flow, Sink, Source}
import akka.stream.{CompletionStrategy, Materializer, OverflowStrategy}
import com.wavesplatform.dex.api.ws.protocol.{WsClientMessage, WsMessage, WsPingOrPong, WsServerMessage}
import com.wavesplatform.dex.domain.utils.ScorexLogging
import play.api.libs.json.Json

import scala.collection.JavaConverters._
import scala.concurrent.Future
import scala.concurrent.duration._
import scala.util.{Failure, Success, Try}

class WsConnection(uri: String, keepAlive: Boolean = true)(implicit system: ActorSystem, materializer: Materializer) extends ScorexLogging {

  log.info(s"""Connecting to Matcher WS API:
            |         URI = $uri
            |  Keep alive = $keepAlive""".stripMargin)

  import materializer.executionContext

  private val wsHandlerRef = system.actorOf(TestWsHandlerActor props keepAlive)

  protected def stringifyClientMessage(cm: WsClientMessage): TextMessage.Strict =
    WsMessage.toStrictTextMessage(cm)(WsClientMessage.wsClientMessageWrites)

  // From test to server
  private val source: Source[TextMessage.Strict, ActorRef] = {
    val completionMatcher: PartialFunction[Any, CompletionStrategy] = { case akka.actor.Status.Success(_) => CompletionStrategy.draining }
    val failureMatcher: PartialFunction[Any, Throwable]             = { case Status.Failure(cause)        => cause }

    Source
      .actorRef[WsClientMessage](completionMatcher, failureMatcher, 10, OverflowStrategy.fail)
      .map(stringifyClientMessage)
      .mapMaterializedValue { source =>
        wsHandlerRef.tell(TestWsHandlerActor.AssignSourceRef, source)
        source
      }
  }

  private val messagesBuffer: ConcurrentLinkedQueue[WsServerMessage] = new ConcurrentLinkedQueue[WsServerMessage]()

  // From server to test
  private val sink: Sink[Message, Future[Done]] = Sink.foreach {
    case tm: TextMessage =>
      for {
        strictText <- tm.toStrict(1.second).map(_.getStrictText)
        clientMessage <- {
          log.trace(s"Got $strictText")
          Try { Json.parse(strictText).as[WsServerMessage] } match {
            case Failure(exception) => Future.failed(exception)
            case Success(x) => {
              messagesBuffer.add(x)
              if (keepAlive) x match {
                case value: WsPingOrPong => wsHandlerRef ! value
                case _                   =>
              }
              Future.successful(x)
            }
          }
        }
      } yield clientMessage

    case bm: BinaryMessage =>
      bm.dataStream.runWith(Sink.ignore)
      Future.failed { new IllegalArgumentException("Binary messages are not supported") }
  }

  private val flow: Flow[Message, TextMessage.Strict, Future[Done]] = Flow.fromSinkAndSourceCoupled(sink, source).watchTermination() {
    case (_, f) =>
      f.onComplete {
        case Success(_) => log.info(s"WebSocket connection to $uri successfully closed")
        case Failure(e) => log.error(s"WebSocket connection to $uri closed with an error", e)
      }(materializer.executionContext)
      f
  }

  val (connectionResponse, closed) = Http().singleWebSocketRequest(WebSocketRequest(uri), flow)

  val connectionOpenedTs: Long                   = System.currentTimeMillis
  val connectionClosedTs: Future[Long]           = closed.map(_ => System.currentTimeMillis)
  val connectionLifetime: Future[FiniteDuration] = connectionClosedTs.map(cc => FiniteDuration(cc - connectionOpenedTs, MILLISECONDS))

  def messages: List[WsServerMessage] = messagesBuffer.iterator().asScala.toList
  def clearMessages(): Unit           = messagesBuffer.clear()

  def send(message: WsClientMessage): Unit = wsHandlerRef ! TestWsHandlerActor.SendToServer(message)

  def close(): Unit     = if (!isClosed) wsHandlerRef ! TestWsHandlerActor.CloseConnection
  def isClosed: Boolean = closed.isCompleted
} 

package com.wavesplatform.dex.load.ws

import akka.Done
import akka.actor.{ActorRef, ActorSystem, Status}
import akka.http.scaladsl.Http
import akka.http.scaladsl.model.ws.{BinaryMessage, Message, TextMessage, WebSocketRequest}
import akka.stream.scaladsl.{Flow, Sink, Source}
import akka.stream.{CompletionStrategy, Materializer, OverflowStrategy}
import com.wavesplatform.dex.api.ws.connection.TestWsHandlerActor
import com.wavesplatform.dex.api.ws.protocol.{WsClientMessage, WsMessage, WsServerMessage}
import com.wavesplatform.dex.domain.utils.ScorexLogging
import play.api.libs.json.Json

import scala.concurrent.Future
import scala.concurrent.duration.DurationInt
import scala.util.{Failure, Success, Try}

class WsConnection(uri: String, receive: WsServerMessage => Option[WsClientMessage])(implicit system: ActorSystem) extends ScorexLogging {

  import system.dispatcher
  private implicit val materializer = Materializer(system)
  private val wsHandlerRef          = system.actorOf(TestWsHandlerActor.props(keepAlive = true))

  log.info(s"Connecting to Matcher WS API: $uri")

  protected def stringifyClientMessage(cm: WsClientMessage): TextMessage.Strict =
    WsMessage.toStrictTextMessage(cm)(WsClientMessage.wsClientMessageWrites)

  // To server
  private val source: Source[TextMessage.Strict, ActorRef] = {
    val completionMatcher: PartialFunction[Any, CompletionStrategy] = { case akka.actor.Status.Success(_) => CompletionStrategy.draining }
    val failureMatcher: PartialFunction[Any, Throwable]             = { case Status.Failure(cause)        => cause }

    Source
      .actorRef[WsClientMessage](completionMatcher, failureMatcher, 10, OverflowStrategy.fail)
      .map(stringifyClientMessage)
      .mapMaterializedValue { source =>
        wsHandlerRef.tell(TestWsHandlerActor.AssignSourceRef, source)
        source
      }
  }

  // To client
  private val sink: Sink[Message, Future[Done]] = Sink.foreach {
    case tm: TextMessage => // TODO move to tests
      for {
        strictText <- tm.toStrict(1.second).map(_.getStrictText)
        clientMessage <- {
          log.trace(s"Got $strictText")
          Try { Json.parse(strictText).as[WsServerMessage] } match {
            case Failure(exception) => Future.failed(exception)
            case Success(x)         => Future.successful { receive(x).foreach(wsHandlerRef ! _) }
          }
        }
      } yield clientMessage

    case bm: BinaryMessage =>
      bm.dataStream.runWith(Sink.ignore)
      Future.failed { new IllegalArgumentException("Binary messages are not supported") }
  }

  private val flow: Flow[Message, TextMessage.Strict, Future[Done]] = Flow.fromSinkAndSourceCoupled(sink, source).watchTermination() {
    case (_, f) =>
      f.onComplete {
        case Success(_) => log.info(s"WebSocket connection to $uri successfully closed")
        case Failure(e) => log.error(s"WebSocket connection to $uri closed with an error", e)
      }(materializer.executionContext)
      f
  }

  val (connectionResponse, closed) = Http().singleWebSocketRequest(WebSocketRequest(uri), flow)

  def send(message: WsClientMessage): Unit = wsHandlerRef ! TestWsHandlerActor.SendToServer(message)

  def isClosed: Boolean = closed.isCompleted
  def close(): Future[Done] = {
    if (!isClosed) wsHandlerRef ! TestWsHandlerActor.CloseConnection
    closed
  }
} 

package de.nierbeck.floating.data.server

import akka.actor.{ActorRef, Props}
import akka.http.scaladsl.model.ws.{Message, TextMessage}
import akka.stream.FlowShape
import akka.stream.scaladsl.{Flow, GraphDSL, Merge, Source}
import de.nierbeck.floating.data.domain.Vehicle
import GraphDSL.Implicits._
import de.nierbeck.floating.data.server._
import de.nierbeck.floating.data.server.actors.websocket._


object Flows {

  def graphFlowWithStats(router: ActorRef): Flow[Message, Message, _] = {
    Flow.fromGraph(GraphDSL.create() { implicit builder =>


      // create an actor source
      val source = Source.actorPublisher[String](VehiclePublisher.props(router))

      // Graph elements we'll use
      val merge = builder.add(Merge[String](2))
      val filter = builder.add(Flow[String].filter(_ => false))

      // get BBox from request and send it to route, return nothing ...
      val mapMsgToString = builder.add(Flow[Message].map[String] {
        case TextMessage.Strict(msg) => {
          println(s"received message: $msg")
          if (msg.contains("close")) {
            router ! msg
          } else if (msg.contains("spark")) {
            router ! SPARK
          } else if (msg.contains("flink")) {
            router ! FLINK
          } else {
            val bbox = toBoundingBox(msg)
            println(s"transformedt to bbox: $bbox")
            router ! bbox
          }
          ""
        }
      })
      //outgoing message ...
      val mapStringToMsg = builder.add(Flow[String].map[Message](x => TextMessage.Strict(x)))

      //add source to flow
      val vehiclesSource = builder.add(source)

      // connect the graph
      mapMsgToString ~> filter ~> merge // this part of the merge will never provide msgs
      vehiclesSource ~> merge ~> mapStringToMsg

      // expose ports
      FlowShape(mapMsgToString.in, mapStringToMsg.out)
    })
  }

} 

package com.tristanpenman.chordial.demo

import akka.actor._
import akka.http.scaladsl.model.ws.{Message, TextMessage}
import akka.http.scaladsl.server.Directives._
import akka.http.scaladsl.server.Route
import akka.stream.scaladsl.{Flow, Source}
import ch.megard.akka.http.cors.scaladsl.CorsDirectives._

import scala.concurrent.ExecutionContext

final class WebSocketWorker(governor: ActorRef, eventsSource: Source[TextMessage, _])(implicit val ec: ExecutionContext)
    extends WebService {

  def route =
    // scalafmt: { indentOperator = spray }
    cors() {
      routes(governor) ~
      pathPrefix("eventstream")(getFromResourceDirectory("webapp")) ~
      handleWebSocketMessages(Flow[Message].take(0).prepend(eventsSource))
    }
}

object WebSocketWorker {
  def apply(nodeRef: ActorRef, eventsSource: Source[TextMessage, _])(
      implicit ec: ExecutionContext): Route =
    new WebSocketWorker(nodeRef, eventsSource).route
} 

package com.pkinsky

import akka.http.scaladsl.Http
import akka.http.scaladsl.model.{HttpResponse, HttpRequest}
import akka.http.scaladsl.model.ws.{TextMessage, Message}
import akka.http.scaladsl.server.Directives._
import akka.http.scaladsl.server.PathMatchers.PathEnd
import akka.stream._
import akka.stream.scaladsl._
import com.softwaremill.react.kafka.ReactiveKafka
import play.api.libs.json.Json
import scala.concurrent.{Future, ExecutionContext}
import scala.concurrent.duration._
import scala.util.Success
import scala.language.postfixOps

object StreamingUpload extends App with AppContext {
  val kafkaPublisherGraph: RunnableGraph[SourceQueue[Event]] =
    Source.queue[Event](1024, OverflowStrategy.backpressure).to(kafka.publish[Event](eventTopic))

  val sourceQueue: SourceQueue[Event] = kafkaPublisherGraph.run

  val queueWriter: Sink[Event, Unit] =
    Flow[Event].mapAsync(1){ elem =>
      sourceQueue.offer(elem)
        .andThen{
          case Success(false) => println(s"failed to publish $elem to topic $eventTopic")
          case Success(true) => println(s"published $elem to topic $eventTopic")
        }
    }.to(Sink.ignore)

  val parseMessages: Flow[Message, Event, Unit] =
    Flow[Message]
      .collect{
        case TextMessage.Strict(t) =>
          val js = Json.parse(t)
          Json.fromJson[Event](js).get
      }

  val wsHandlerFlow: Flow[Message, Message, Unit] =
    Flow.fromSinkAndSource(
      sink = parseMessages.to(queueWriter),
      source = Source.maybe
    )

  val routes: Flow[HttpRequest, HttpResponse, Unit] =
      get {
        path(PathEnd) {
          getFromResource("test.html")
        } ~
          path("ws") {
            println("ws connection accepted")
            handleWebsocketMessages(wsHandlerFlow)
          }
      }

  Http().bindAndHandle(routes, "localhost", port).andThen{ case util.Failure(t) => println(s"error binding to localhost: $t")}

  println(s"listening on port $port, press ENTER to stop")

  awaitTermination()
}

object KafkaListener extends App with AppContext {
  val graph = kafka.consume[Event](eventTopic, "kafka_listener").toMat(Sink.foreach(println))(Keep.right)

  graph.run.onComplete(println)

  println(s"listening to Kafka topic $eventTopic, press ENTER to stop")

  awaitTermination()
} 

package chat

import java.net._

import scala.concurrent.ExecutionContext
import scala.concurrent.ExecutionContext.Implicits._
import akka.actor._
import akka.NotUsed
import akka.stream._
import akka.stream.scaladsl._
import akka.http.scaladsl.model.ws.{Message, TextMessage}
import akka.http.scaladsl.server.Directives._
import akka.http.scaladsl.model.RemoteAddress
import chat.admin.AdminService
import EventConstants._
//import chat.utils.HttpClient

class ChatService(env: String, admin: AdminService)
                 (implicit system: ActorSystem, mat: ActorMaterializer, dispatcher: ExecutionContext)
  extends WebService {

  private var chatSuper:  ActorRef  = null
  private var aggregator: ActorRef  = null
  private var httpClient: HttpClient= null

  private val servicePort = 8000
  private val serviceRoute= //<- adjustable depended on client url
    pathPrefix(IntNumber) {
      chatRoomID => {
        chatSuper ! CreateChatRoom(chatRoomID)
        extractClientIP {
          ip =>
            handleWebSocketMessages(newUser(chatRoomID, ip))
        }
      }
    }

  def regNode(port: Int) = {
    val localhost = InetAddress.getLocalHost
    val localIpAddress = localhost.getHostAddress

    environment.hostName = localIpAddress
    environment.port = port

    log.info(s"Server IP Address of System => ${localIpAddress}")
  }

  def incomingMessages(userActor: ActorRef): Sink[Message, NotUsed] = {
    Flow[Message].map {
      // transform websocket message to domain message
      case TextMessage.Strict(text) => UserActor.IncomingMessage(text)

      // PoisonPill asynchronously stops disconnected user actor
      //TODO : to deal with join, leave, text message types
    }.to(Sink.actorRef[UserActor.IncomingMessage](userActor, PoisonPill))
  }

  def outgoingMessages(userActor: ActorRef): Source[Message, NotUsed] = {
    Source.actorRef[UserActor.OutgoingMessage](20000, OverflowStrategy.fail)
      .mapMaterializedValue { outActor =>
        // give the user actor a way to send messages out
        userActor ! UserActor.Connected(outActor)
        NotUsed
      }.map(
      // transform domain message to web socket message
      (outMsg: UserActor.OutgoingMessage) => TextMessage(outMsg.text))
  }

  def newUser(chatRoomID: Int, ip: RemoteAddress): Flow[Message, Message, NotUsed] = {
    // new connection - new user actor
    val userActor = system.actorOf(Props(new UserActor(chatRoomID, chatSuper, ChatRooms.httpClient, ip)))

    // Set Sink & Source
    val incomingMsg = incomingMessages(userActor)
    val outgoingMsg = outgoingMessages(userActor)

    Flow.fromSinkAndSource(incomingMsg, outgoingMsg)
  }

  def start() {
    log.info( "ChatService staring ..." )

    environment.setEnvType(env)
    environment.version = system.settings.config.getString("akka.heimdallr-version" )

    chatSuper = system.actorOf(Props(classOf[ChatSupervisor], env), "cs")
    admin.setChatSupervisorActorRef(chatSuper)

    regNode(servicePort)
    serviceBind(this.getClass.getSimpleName, serviceRoute, servicePort)

    log.info(environment.HeimdallrLogo(system))
  }

  def stop(): Unit = {
    serviceUnbind(this.getClass.getSimpleName)
  }
} 

package akkaviz.server

import akka.http.scaladsl.model.ws.{BinaryMessage, Message}
import akka.stream.scaladsl.Flow
import akka.util.ByteString
import akkaviz.protocol

trait ProtocolSerializationSupport {

  def protocolServerMessageToByteString: Flow[protocol.ApiServerMessage, ByteString, Any] = Flow[protocol.ApiServerMessage].map {
    msg => ByteString(protocol.IO.write(msg))
  }

  def websocketMessageToClientMessage: Flow[Message, protocol.ApiClientMessage, _] = Flow[Message].collect {
    case BinaryMessage.Strict(msg) =>
      protocol.IO.readClient(msg.asByteBuffer)
  }

}

object ProtocolSerializationSupport extends ProtocolSerializationSupport 

package akkaviz.server

import akka.actor.{ActorRef, ActorSystem, Kill, PoisonPill}
import akka.http.scaladsl.coding.Gzip
import akka.http.scaladsl.model._
import akka.http.scaladsl.model.ws.{BinaryMessage, Message}
import akka.http.scaladsl.server.Directives
import akka.stream.scaladsl._
import akka.stream.{Materializer, OverflowStrategy}
import akkaviz.config.Config
import akkaviz.events._
import akkaviz.events.types._
import akkaviz.persistence.{PersistenceSources, ReceivedRecord}
import akkaviz.protocol

import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.duration._

class Webservice(implicit materializer: Materializer, system: ActorSystem)
    extends Directives with SubscriptionSession with ReplSupport with AkkaHttpHelpers with ArchiveSupport
    with FrontendResourcesSupport with ProtocolSerializationSupport with BackendEventsMarshalling {

  def route: Flow[HttpRequest, HttpResponse, Any] = encodeResponseWith(Gzip) {
    get {
      path("stream") {
        handleWebSocketMessages(tracingEventsFlow.mapMaterializedValue(EventSystem.subscribe))
      }
    } ~
      archiveRouting ~
      replRouting ~
      frontendResourcesRouting
  }

  def tracingEventsFlow: Flow[Message, Message, ActorRef] = {
    val eventSrc = Source.actorRef[BackendEvent](Config.bufferSize, OverflowStrategy.dropNew)

    val wsIn = Flow[Message]
      .via(websocketMessageToClientMessage)
      .via(handleUserCommand)
      .scan(defaultSettings)(updateSettings)
      .expand(r => Iterator.continually(r))

    val out = wsIn.zipMat(eventSrc)((_, m) => m)
      .collect {
        case (settings, r: BackendEvent) if settings.eventAllowed(r) => r
      }.via(backendEventToProtocolFlow)
      .keepAlive(10.seconds, () => protocol.Ping)
      .via(protocolServerMessageToByteString)
      .map(BinaryMessage.Strict(_))

    out
  }

  private[this] val handleUserCommand: Flow[protocol.ApiClientMessage, ChangeSubscriptionSettings, _] = Flow[protocol.ApiClientMessage].mapConcat {
    case protocol.SetAllowedMessages(classNames) =>
      system.log.debug(s"Set allowed messages to $classNames")
      List(SetAllowedClasses(classNames))
    case protocol.ObserveActors(actors) =>
      system.log.debug(s"Set observed actors to $actors")
      List(SetActorEventFilter(actors))
    case protocol.SetReceiveDelay(duration) =>
      system.log.debug(s"Setting receive delay to $duration")
      EventSystem.setReceiveDelay(duration)
      Nil
    case protocol.SetEnabled(isEnabled) =>
      system.log.info(s"Setting EventSystem.setEnabled($isEnabled)")
      EventSystem.setEnabled(isEnabled)
      Nil
    case protocol.RefreshInternalState(actor) =>
      ActorSystems.refreshActorState(actor)
      Nil
    case protocol.PoisonPillActor(actor) =>
      ActorSystems.tell(actor, PoisonPill)
      Nil
    case protocol.KillActor(actor) =>
      ActorSystems.tell(actor, Kill)
      Nil
  }

  override def receivedOf(ref: String): Source[ReceivedRecord, _] = PersistenceSources.of(ref)

  override def receivedBetween(ref: String, ref2: String): Source[ReceivedRecord, _] = PersistenceSources.between(ref, ref2)

  override def isArchiveEnabled: Boolean = Config.enableArchive

} 

package example.akkawschat

import java.util.Date

import akka.actor.ActorSystem
import akka.http.scaladsl.model.ws.{ Message, TextMessage }

import scala.concurrent.duration._
import akka.http.scaladsl.server.Directives
import akka.stream.scaladsl.Flow
import upickle.default._
import shared.Protocol
import shared.Protocol._

import scala.util.Failure

class Webservice(implicit system: ActorSystem) extends Directives {
  val theChat = Chat.create()
  import system.dispatcher
  system.scheduler.scheduleAtFixedRate(15.second, 15.second) { () =>
    theChat.injectMessage(ChatMessage(sender = "clock", s"Bling! The time is ${new Date().toString}."))
  }

  def route =
    get {
      pathSingleSlash {
        getFromResource("web/index.html")
      } ~
        // Scala-JS puts them in the root of the resource directory per default,
        // so that's where we pick them up
        path("frontend-launcher.js")(getFromResource("frontend-launcher.js")) ~
        path("frontend-fastopt.js")(getFromResource("frontend-fastopt.js")) ~
        path("chat") {
          parameter("name") { name =>
            handleWebSocketMessages(websocketChatFlow(sender = name))
          }
        }
    } ~
      getFromResourceDirectory("web")

  def websocketChatFlow(sender: String): Flow[Message, Message, Any] =
    Flow[Message]
      .collect {
        case TextMessage.Strict(msg) => msg // unpack incoming WS text messages...
        // This will lose (ignore) messages not received in one chunk (which is
        // unlikely because chat messages are small) but absolutely possible
        // FIXME: We need to handle TextMessage.Streamed as well.
      }
      .via(theChat.chatFlow(sender)) // ... and route them through the chatFlow ...
      .map {
        case msg: Protocol.Message =>
          TextMessage.Strict(write(msg)) // ... pack outgoing messages into WS JSON messages ...
      }
      .via(reportErrorsFlow) // ... then log any processing errors on stdin

  def reportErrorsFlow[T]: Flow[T, T, Any] =
    Flow[T]
      .watchTermination()((_, f) => f.onComplete {
        case Failure(cause) =>
          println(s"WS stream failed with $cause")
        case _ => // ignore regular completion
      })
} 

package asura.core.actor.flow

import akka.NotUsed
import akka.actor.{ActorRef, PoisonPill}
import akka.http.scaladsl.model.ws.{Message, TextMessage}
import akka.stream.OverflowStrategy
import akka.stream.scaladsl.{Flow, Sink, Source}
import asura.common.actor.{ActorEvent, SenderMessage}
import asura.common.exceptions.InvalidStatusException
import asura.core.CoreConfig
import asura.core.util.JacksonSupport

import scala.concurrent.duration._

object WebSocketMessageHandler {

  val DEFAULT_BUFFER_SIZE = CoreConfig.DEFAULT_WS_ACTOR_BUFFER_SIZE
  val KEEP_ALIVE_INTERVAL = 2

  def newHandleFlow[T <: AnyRef](workActor: ActorRef, msgClass: Class[T]): Flow[Message, Message, NotUsed] = {
    val incomingMessages: Sink[Message, NotUsed] =
      Flow[Message].map {
        case TextMessage.Strict(text) => JacksonSupport.parse(text, msgClass)
        case _ => throw InvalidStatusException("Unsupported message type")
      }.to(Sink.actorRef[T](workActor, PoisonPill))
    val outgoingMessages: Source[Message, NotUsed] =
      Source.actorRef[ActorEvent](DEFAULT_BUFFER_SIZE, OverflowStrategy.dropHead)
        .mapMaterializedValue { outActor =>
          workActor ! SenderMessage(outActor)
          NotUsed
        }
        .map(result => TextMessage(JacksonSupport.stringify(result)))
        .keepAlive(KEEP_ALIVE_INTERVAL.seconds, () => TextMessage.Strict(""))
    Flow.fromSinkAndSource(incomingMessages, outgoingMessages)
  }

  def newHandleStringFlow[T <: AnyRef](workActor: ActorRef, msgClass: Class[T]): Flow[Message, Message, NotUsed] = {
    val incomingMessages: Sink[Message, NotUsed] =
      Flow[Message].map {
        case TextMessage.Strict(text) => JacksonSupport.parse(text, msgClass)
        case _ => throw InvalidStatusException("Unsupported message type")
      }.to(Sink.actorRef[T](workActor, PoisonPill))
    val outgoingMessages: Source[Message, NotUsed] =
      Source.actorRef[String](DEFAULT_BUFFER_SIZE, OverflowStrategy.dropHead)
        .mapMaterializedValue { outActor =>
          workActor ! SenderMessage(outActor)
          NotUsed
        }
        .map(result => TextMessage(result))
        .keepAlive(KEEP_ALIVE_INTERVAL.seconds, () => TextMessage.Strict(""))
    Flow.fromSinkAndSource(incomingMessages, outgoingMessages)
  }

  def stringToActorEventFlow[T <: AnyRef](workActor: ActorRef, msgClass: Class[T]): Flow[String, String, NotUsed] = {
    val incomingMessages: Sink[String, NotUsed] =
      Flow[String].map {
        case text: String => JacksonSupport.parse(text, msgClass)
      }.to(Sink.actorRef[T](workActor, PoisonPill))
    val outgoingMessages: Source[String, NotUsed] =
      Source.actorRef[ActorEvent](DEFAULT_BUFFER_SIZE, OverflowStrategy.dropHead)
        .mapMaterializedValue { outActor =>
          workActor ! SenderMessage(outActor)
          NotUsed
        }
        .map(result => JacksonSupport.stringify(result))
        .keepAlive(KEEP_ALIVE_INTERVAL.seconds, () => "")
    Flow.fromSinkAndSource(incomingMessages, outgoingMessages)
  }

  def stringToActorEventFlow[T <: AnyRef](workActor: ActorRef): Flow[String, String, NotUsed] = {
    val incomingMessages: Sink[String, NotUsed] =
      Flow[String].to(Sink.actorRef[String](workActor, PoisonPill))
    val outgoingMessages: Source[String, NotUsed] =
      Source.actorRef[ActorEvent](DEFAULT_BUFFER_SIZE, OverflowStrategy.dropHead)
        .mapMaterializedValue { outActor =>
          workActor ! SenderMessage(outActor)
          NotUsed
        }
        .map(result => JacksonSupport.stringify(result))
        .keepAlive(KEEP_ALIVE_INTERVAL.seconds, () => "")
    Flow.fromSinkAndSource(incomingMessages, outgoingMessages)
  }
} 

package sttp.client.akkahttp

import akka.actor.ActorSystem
import akka.event.LoggingAdapter
import akka.http.scaladsl.model.ws.{Message, WebSocketRequest, WebSocketUpgradeResponse}
import akka.http.scaladsl.model.{HttpRequest, HttpResponse}
import akka.http.scaladsl.server.{ExceptionHandler, RejectionHandler, Route, RoutingLog}
import akka.http.scaladsl.settings.{ClientConnectionSettings, ConnectionPoolSettings, ParserSettings, RoutingSettings}
import akka.http.scaladsl.{Http, HttpsConnectionContext}
import akka.stream.Materializer
import akka.stream.scaladsl.Flow

import scala.concurrent.{ExecutionContext, ExecutionContextExecutor, Future}

trait AkkaHttpClient {
  def singleRequest(
      request: HttpRequest,
      settings: ConnectionPoolSettings
  ): Future[HttpResponse]

  def singleWebsocketRequest[WS_RESULT](
      request: WebSocketRequest,
      clientFlow: Flow[Message, Message, WS_RESULT],
      settings: ClientConnectionSettings
  )(implicit ec: ExecutionContext, mat: Materializer): Future[(WebSocketUpgradeResponse, WS_RESULT)]
}

object AkkaHttpClient {
  def default(
      system: ActorSystem,
      connectionContext: Option[HttpsConnectionContext],
      customLog: Option[LoggingAdapter]
  ): AkkaHttpClient =
    new AkkaHttpClient {
      private val http = Http()(system)

      override def singleRequest(
          request: HttpRequest,
          settings: ConnectionPoolSettings
      ): Future[HttpResponse] = {
        http.singleRequest(
          request,
          connectionContext.getOrElse(http.defaultClientHttpsContext),
          settings,
          customLog.getOrElse(system.log)
        )
      }

      override def singleWebsocketRequest[WS_RESULT](
          request: WebSocketRequest,
          clientFlow: Flow[Message, Message, WS_RESULT],
          settings: ClientConnectionSettings
      )(implicit ec: ExecutionContext, mat: Materializer): Future[(WebSocketUpgradeResponse, WS_RESULT)] = {
        val (wsResponse, wsResult) = http.singleWebSocketRequest(
          request,
          clientFlow,
          connectionContext.getOrElse(http.defaultClientHttpsContext),
          None,
          settings,
          customLog.getOrElse(system.log)
        )
        wsResponse.map((_, wsResult))
      }
    }

  def stubFromAsyncHandler(run: HttpRequest => Future[HttpResponse]): AkkaHttpClient =
    new AkkaHttpClient {
      def singleRequest(request: HttpRequest, settings: ConnectionPoolSettings): Future[HttpResponse] =
        run(request)

      override def singleWebsocketRequest[WS_RESULT](
          request: WebSocketRequest,
          clientFlow: Flow[Message, Message, WS_RESULT],
          settings: ClientConnectionSettings
      )(implicit ec: ExecutionContext, mat: Materializer): Future[(WebSocketUpgradeResponse, WS_RESULT)] =
        Future.failed(new RuntimeException("Websockets are not supported"))
    }

  def stubFromRoute(route: Route)(implicit
      routingSettings: RoutingSettings,
      parserSettings: ParserSettings,
      materializer: Materializer,
      routingLog: RoutingLog,
      executionContext: ExecutionContextExecutor = null,
      rejectionHandler: RejectionHandler = RejectionHandler.default,
      exceptionHandler: ExceptionHandler = null
  ): AkkaHttpClient = stubFromAsyncHandler(Route.asyncHandler(route))
} 

package sttp.client.akkahttp

import java.util.concurrent.ConcurrentLinkedQueue

import akka.Done
import akka.http.scaladsl.model.ws.{Message, TextMessage}
import akka.stream.Materializer
import akka.stream.scaladsl._
import org.scalatest.BeforeAndAfterAll
import org.scalatest.concurrent.{Eventually, IntegrationPatience}
import sttp.client._

import scala.collection.JavaConverters._
import scala.concurrent.duration._
import scala.concurrent.{ExecutionContext, Future, Promise}
import scala.util.Success
import org.scalatest.flatspec.AsyncFlatSpec
import org.scalatest.matchers.should.Matchers
import sttp.client.testing.HttpTest.wsEndpoint

class AkkaHttpWebsocketTest
    extends AsyncFlatSpec
    with Matchers
    with BeforeAndAfterAll
    with Eventually
    with IntegrationPatience {
  implicit val ec: ExecutionContext = scala.concurrent.ExecutionContext.global
  implicit val backend: SttpBackend[Future, Nothing, Flow[Message, Message, *]] = AkkaHttpBackend()

  it should "send and receive ten messages" in {
    val received = new ConcurrentLinkedQueue[String]()

    val sink: Sink[Message, Future[Done]] = collectionSink(received)

    val n = 10
    val source: Source[Message, Promise[Option[Message]]] =
      Source((1 to n).map(i => TextMessage(s"test$i"))).concatMat(Source.maybe[Message])(Keep.right)

    val flow: Flow[Message, Message, (Future[Done], Promise[Option[Message]])] =
      Flow.fromSinkAndSourceMat(sink, source)(Keep.both)

    basicRequest.get(uri"$wsEndpoint/ws/echo").openWebsocket(flow).flatMap { r =>
      eventually {
        received.asScala.toList shouldBe (1 to n).map(i => s"echo: test$i").toList
      }

      r.result._2.complete(Success(None)) // the source should now complete
      r.result._1.map(_ => succeed) // the future should be completed once the stream completes (and the ws closes)
    }
  }

  it should "receive two messages" in {
    val received = new ConcurrentLinkedQueue[String]()
    val sink: Sink[Message, Future[Done]] = collectionSink(received)
    val source: Source[Message, Promise[Option[Message]]] = Source.maybe[Message]

    val flow: Flow[Message, Message, Promise[Option[Message]]] =
      Flow.fromSinkAndSourceMat(sink, source)(Keep.right)

    basicRequest.get(uri"$wsEndpoint/ws/send_and_wait").openWebsocket(flow).flatMap { r =>
      eventually {
        received.asScala.toList shouldBe List("test10", "test20")
      }
      r.result.success(None) // closing
      succeed
    }
  }

  it should "error if the endpoint is not a websocket" in {
    basicRequest.get(uri"$wsEndpoint/echo").openWebsocket(Flow.apply[Message]).failed.map { t =>
      t shouldBe a[NotAWebsocketException]
    }
  }

  def collectionSink(queue: ConcurrentLinkedQueue[String]): Sink[Message, Future[Done]] =
    Sink
      .setup[Message, Future[Done]] { (_materializer, _) =>
        Flow[Message]
        // mapping with parallelism 1 so that messages don't get reordered
          .mapAsync(1) {
            case m: TextMessage =>
              implicit val materializer: Materializer = _materializer
              m.toStrict(1.second).map(Some(_))
            case _ => Future.successful(None)
          }
          .collect {
            case Some(TextMessage.Strict(text)) => text
          }
          .toMat(Sink.foreach(queue.add))(Keep.right)
      }
      .mapMaterializedValue(_.flatMap(identity))

  override protected def afterAll(): Unit = {
    backend.close()
    super.afterAll()
  }
} 

package sttp.client.akkahttp

import akka.actor.ActorSystem
import akka.http.scaladsl.model.ws.{Message, TextMessage}
import akka.stream.OverflowStrategy
import akka.stream.scaladsl.{Flow, Keep, Sink, Source}
import org.scalatest.BeforeAndAfterAll
import org.scalatest.concurrent.ScalaFutures
import org.scalatest.flatspec.AnyFlatSpec
import org.scalatest.matchers.should.Matchers
import sttp.client._
import sttp.model.Headers

import scala.concurrent.{Await, Future}
import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.duration._

class SttpBackendStubAkkaTests extends AnyFlatSpec with Matchers with ScalaFutures with BeforeAndAfterAll {

  implicit val system: ActorSystem = ActorSystem()

  override protected def afterAll(): Unit = {
    Await.result(system.terminate().map(_ => ()), 5.seconds)
  }

  "backend stub" should "cycle through responses using a single sent request" in {
    // given
    implicit val backend = AkkaHttpBackend.stub
      .whenRequestMatches(_ => true)
      .thenRespondCyclic("a", "b", "c")

    // when
    def r = basicRequest.get(uri"http://example.org/a/b/c").send().futureValue

    // then
    r.body shouldBe Right("a")
    r.body shouldBe Right("b")
    r.body shouldBe Right("c")
    r.body shouldBe Right("a")
  }

  it should "use given flow as web socket handler" in {
    // This test is an example how can we test client flow.
    // We check behavior of client when connected to echo server.
    // Client responsibility was to send two messages to the server and collect received messages.
    val useHandler: Flow[Message, Message, Future[Seq[Message]]] => Future[Seq[Message]] = clientFlow => {
      val ((outQueue, clientReceivedMessages), inQueue) = Source
        .queue(1, OverflowStrategy.fail)
        .viaMat(clientFlow)(Keep.both)
        .toMat(Sink.queue())(Keep.both)
        .run()

      def echoMsg(): Future[Unit] =
        inQueue.pull().flatMap {
          case None =>
            echoMsg()
          case Some(msg) =>
            outQueue.offer(TextMessage(s"echo: " + msg.asTextMessage.getStrictText)).map(_ => ())
        }

      (for {
        _ <- outQueue.offer(TextMessage("Hi!"))
        _ <- echoMsg()
        _ <- echoMsg()
        _ = outQueue.complete()
        _ <- outQueue.watchCompletion()
      } yield ()).flatMap(_ => clientReceivedMessages)
    }

    val clientFlow: Flow[Message, Message, Future[Seq[Message]]] = {
      Flow.fromSinkAndSourceMat(
        Sink.seq[Message],
        Source((1 to 2).map(i => TextMessage(s"test$i")))
      )(Keep.left)
    }

    implicit val b = AkkaHttpBackend.stub
      .whenRequestMatches(_ => true)
      .thenHandleOpenWebSocket(Headers(List.empty), useHandler)

    val receivedMessages = basicRequest
      .get(uri"wss://echo.websocket.org")
      .openWebsocket(clientFlow)
      .flatMap(_.result)
      .futureValue
      .toList

    receivedMessages shouldBe List("Hi!", "echo: test1", "echo: test2").map(TextMessage(_))
  }
} 

package io.vamp.http_api.ws

import java.util.UUID

import akka.actor.PoisonPill
import akka.http.scaladsl.model.HttpRequest
import akka.http.scaladsl.model.ws.{ Message, TextMessage }
import akka.http.scaladsl.server.Route
import akka.stream._
import akka.stream.scaladsl.{ Flow, Sink, Source }
import akka.util.Timeout
import io.vamp.common.akka.IoC._
import io.vamp.common.http.{ HttpApiDirectives, HttpApiHandlers, TerminateFlowStage }
import io.vamp.common.{ Config, Namespace }
import io.vamp.http_api.ws.WebSocketActor.{ SessionClosed, SessionEvent, SessionOpened, SessionRequest }
import io.vamp.http_api.{ AbstractRoute, LogDirective }

import scala.concurrent.Future

trait WebSocketRoute extends AbstractRoute with WebSocketMarshaller with HttpApiHandlers {
  this: HttpApiDirectives with LogDirective ⇒

  implicit def materializer: Materializer

  private lazy val limit = Config.int("vamp.http-api.websocket.stream-limit")

  protected def websocketApiHandler(implicit namespace: Namespace, timeout: Timeout): Route

  def websocketRoutes(implicit namespace: Namespace, timeout: Timeout) = {
    pathEndOrSingleSlash {
      get {
        extractRequest { request ⇒
          handleWebSocketMessages {
            websocket(request)
          }
        }
      }
    }
  }

  protected def filterWebSocketOutput(message: AnyRef)(implicit namespace: Namespace, timeout: Timeout): Future[Boolean] = Future.successful(true)

  private def apiHandler(implicit namespace: Namespace, timeout: Timeout) = Route.asyncHandler(log {
    websocketApiHandler
  })

  private def websocket(origin: HttpRequest)(implicit namespace: Namespace, timeout: Timeout): Flow[AnyRef, Message, Any] = {
    val id = UUID.randomUUID()

    val in = Flow[AnyRef].collect {
      case TextMessage.Strict(message)  ⇒ Future.successful(message)
      case TextMessage.Streamed(stream) ⇒ stream.limit(limit()).completionTimeout(timeout.duration).runFold("")(_ + _)
    }.mapAsync(parallelism = 3)(identity)
      .mapConcat(unmarshall)
      .map(SessionRequest(apiHandler, id, origin, _))
      .to(Sink.actorRef[SessionEvent](actorFor[WebSocketActor], SessionClosed(id)))

    val out = Source.actorRef[AnyRef](16, OverflowStrategy.dropHead)
      .mapMaterializedValue(actorFor[WebSocketActor] ! SessionOpened(id, _))
      .via(new TerminateFlowStage[AnyRef](_ == PoisonPill))
      .mapAsync(parallelism = 3)(message ⇒ filterWebSocketOutput(message).map(f ⇒ f → message))
      .collect { case (true, m) ⇒ m }
      .map(message ⇒ TextMessage.Strict(marshall(message)))

    Flow.fromSinkAndSource(in, out)
  }
} 

package chat

import akka.NotUsed
import akka.actor._
import akka.http.scaladsl._
import akka.http.scaladsl.model.ws.{Message, TextMessage}
import akka.http.scaladsl.server.Directives._
import akka.stream._
import akka.stream.scaladsl._

import scala.concurrent.duration._
import scala.concurrent.Await
import scala.io.StdIn

object Server {
  def main(args: Array[String]): Unit = {
    implicit val system = ActorSystem()
    implicit val materializer = ActorMaterializer()

    val chatRoom = system.actorOf(Props(new ChatRoom), "chat")

    def newUser(): Flow[Message, Message, NotUsed] = {
      // new connection - new user actor
      val userActor = system.actorOf(Props(new User(chatRoom)))

      val incomingMessages: Sink[Message, NotUsed] =
        Flow[Message].map {
          // transform websocket message to domain message
          case TextMessage.Strict(text) => User.IncomingMessage(text)
        }.to(Sink.actorRef[User.IncomingMessage](userActor, PoisonPill))

      val outgoingMessages: Source[Message, NotUsed] =
        Source.actorRef[User.OutgoingMessage](10, OverflowStrategy.fail)
        .mapMaterializedValue { outActor =>
          // give the user actor a way to send messages out
          userActor ! User.Connected(outActor)
          NotUsed
        }.map(
          // transform domain message to web socket message
          (outMsg: User.OutgoingMessage) => TextMessage(outMsg.text))

      // then combine both to a flow
      Flow.fromSinkAndSource(incomingMessages, outgoingMessages)
    }

    val route =
      path("chat") {
        get {
          handleWebSocketMessages(newUser())
        }
      }

    val binding = Await.result(Http().bindAndHandle(route, "127.0.0.1", 8080), 3.seconds)


    // the rest of the sample code will go here
    println("Started server at 127.0.0.1:8080, press enter to kill server")
    StdIn.readLine()
    system.terminate()
  }
}