io.gatling.core.structure.ScenarioContext Scala Examples

package com.github.mnogu.gatling.mqtt.action

import com.github.mnogu.gatling.mqtt.protocol.{MqttComponents, MqttProtocol}
import com.github.mnogu.gatling.mqtt.request.builder.MqttAttributes
import io.gatling.core.action.Action
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.structure.ScenarioContext

class MqttRequestActionBuilder(mqttAttributes: MqttAttributes)
  extends ActionBuilder {

  override def build(
    ctx: ScenarioContext, next: Action
  ): Action = {
    import ctx._

    val mqttComponents : MqttComponents = protocolComponentsRegistry.components(MqttProtocol.MqttProtocolKey)
    
    new MqttRequestAction(
      mqttAttributes,
      coreComponents,
      mqttComponents.mqttProtocol,
      next
    )
  }
} 

package io.gatling.amqp.action

import akka.actor._
import io.gatling.amqp.config._
import io.gatling.amqp.data._
import io.gatling.amqp.request.builder._
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.structure.ScenarioContext

class AmqpActionBuilder(amqpRequestBuilder: AmqpRequestBuilder)(implicit amqp: AmqpProtocol) extends ActionBuilder {
  def build(system: ActorSystem, next: ActorRef, ctx: ScenarioContext): ActorRef = {
    amqpRequestBuilder.build match {
      case req: PublishRequest =>
        system.actorOf(AmqpPublishAction.props(req, next, amqp), actorName("AmqpPublishAction"))
      case req: ConsumeRequest =>
        system.actorOf(AmqpConsumeAction.props(req, next, amqp), actorName("AmqpConsumeAction"))
    }
  }
} 

package ch.tamedia.gatling

import ch.tamedia.gatling.actions.impl.{GrpcAsyncCallAction, GrpcSyncCallAction}
import ch.tamedia.gatling.actions.{GrpcAction, GrpcExecutableAction}
import ch.tamedia.gatling.grpc._
import io.gatling.core.action.Action
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.protocol.Protocols
import io.gatling.core.structure.ScenarioContext


case class GrpcActionBuilder(action: GrpcExecutableAction, checks: List[GrpcCheck]) extends ActionBuilder{

  def grpcProtocol(protocols: Protocols) = {
    protocols.protocol[GrpcProtocol].getOrElse(throw new UnsupportedOperationException("gRPC protocol wasn't registered"))
  }

  override def build(ctx: ScenarioContext, next: Action): Action = {
    import ctx._
    val statsEngine = coreComponents.statsEngine
    GrpcAction(action, checks, new GrpcProtocol, ctx.system, statsEngine, next)
//    functionName match {
//      case "async" => GrpcAction(GrpcAsyncCallAction(functionName, host, port, requestMessage = payload.get), checks, new GrpcProtocol, ctx.system, statsEngine, next)
//      case "sync"  => GrpcAction(GrpcSyncCallAction(functionName, host, port, requestMessage = payload.get), checks, new GrpcProtocol, ctx.system, statsEngine, next)
//      case       _ => throw new UnsupportedOperationException(s"Operation $functionName is not supported")
//    }
  }
} 

package ru.tinkoff.gatling.amqp.action

import io.gatling.core.action.Action
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.protocol.ProtocolComponentsRegistry
import io.gatling.core.structure.ScenarioContext
import ru.tinkoff.gatling.amqp.protocol.{AmqpComponents, AmqpProtocol}
import ru.tinkoff.gatling.amqp.request.AmqpAttributes

case class PublishBuilder(attributes: AmqpAttributes, configuration: GatlingConfiguration) extends ActionBuilder {

  private def components(protocolComponentsRegistry: ProtocolComponentsRegistry): AmqpComponents =
    protocolComponentsRegistry.components(AmqpProtocol.amqpProtocolKey)

  override def build(ctx: ScenarioContext, next: Action): Action = {
    import ctx._
    val amqpComponents = components(protocolComponentsRegistry)
    val statsEngine    = coreComponents.statsEngine

    new Publish(
      attributes,
      amqpComponents,
      statsEngine,
      coreComponents.clock,
      configuration,
      next,
      coreComponents.throttler,
      ctx.throttled
    )
  }
} 

package ru.tinkoff.gatling.amqp.action

import io.gatling.core.action.Action
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.protocol.ProtocolComponentsRegistry
import io.gatling.core.structure.ScenarioContext
import ru.tinkoff.gatling.amqp.protocol.{AmqpComponents, AmqpProtocol}
import ru.tinkoff.gatling.amqp.request.{AmqpAttributes, AmqpExchange}

case class RequestReplyBuilder(attributes: AmqpAttributes,
                               replyDest: AmqpExchange,
                               setReplyTo: Boolean,
                               configuration: GatlingConfiguration)
    extends ActionBuilder {
  private def components(protocolComponentsRegistry: ProtocolComponentsRegistry): AmqpComponents =
    protocolComponentsRegistry.components(AmqpProtocol.amqpProtocolKey)

  override def build(ctx: ScenarioContext, next: Action): Action = {
    import ctx._
    val amqpComponents = components(protocolComponentsRegistry)
    val statsEngine    = coreComponents.statsEngine

    new RequestReply(
      attributes,
      replyDest,
      setReplyTo,
      None,
      amqpComponents,
      statsEngine,
      coreComponents.clock,
      next,
      coreComponents.throttler,
      ctx.throttled
    )

  }
} 

package org.apache.openwhisk.extension.whisk

import io.gatling.core.Predef._
import io.gatling.core.action.Action
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.session.Expression
import io.gatling.core.structure.ScenarioContext
import io.gatling.http.request.builder.{Http, HttpRequestBuilder}
import spray.json.DefaultJsonProtocol._
import spray.json._

case class OpenWhiskActionBuilderBase(requestName: Expression[String]) {

  implicit private val http = new Http(requestName)

  
  def invoke() = {
    OpenWhiskActionBuilder(http.post(path).queryParam("blocking", "true").basicAuth(uuid, key))
  }
}

case class OpenWhiskActionBuilder(http: HttpRequestBuilder) extends ActionBuilder {
  override def build(ctx: ScenarioContext, next: Action): Action = {
    http.build(ctx, next)
  }
} 

package io.github.gatling.sql.action

import io.github.gatling.sql.protocol.SqlProtocol
import io.github.gatling.sql.request.SqlAttributes
import io.gatling.commons.stats.{KO, OK}
import io.gatling.commons.util.ClockSingleton.nowMillis
import io.gatling.commons.validation.Validation
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.action.{Action, ExitableAction}
import io.gatling.core.protocol.ProtocolComponentsRegistry
import io.gatling.core.session.Session
import io.gatling.core.stats.StatsEngine
import io.gatling.core.stats.message.ResponseTimings
import io.gatling.core.structure.ScenarioContext
import io.gatling.core.util.NameGen
import java.sql.PreparedStatement

import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.Future

class SqlActionBuilder(attr: SqlAttributes) extends ActionBuilder with NameGen {

  private def components(protocolComponentsRegistry: ProtocolComponentsRegistry) =
    protocolComponentsRegistry.components(SqlProtocol.SqlProtocolKey)

  override def build(ctx: ScenarioContext, next: Action): Action = {
    import ctx._
    val statsEngine = coreComponents.statsEngine
    val sqlComponents = components(protocolComponentsRegistry)
    new SqlAction(genName(s"SQL: ${attr.tag}"), sqlComponents.sqlProtocol, statsEngine, next, attr)
  }

}

class SqlAction(val name: String, protocol: SqlProtocol, val statsEngine: StatsEngine, val next: Action,
                val attr: SqlAttributes) extends ExitableAction {

  def execute(session: Session): Unit = {
    val stmt: Validation[PreparedStatement] = attr.statement(session)

    stmt.onFailure(err => {
      statsEngine.logResponse(session, name, ResponseTimings(nowMillis, nowMillis), KO, None, Some("Error setting up statement: " + err), Nil)
      next ! session.markAsFailed
    })

    stmt.onSuccess({ stmt =>
      val start = nowMillis

      val result = Future {
        stmt.execute()
      }

      result.onFailure { case t =>
        statsEngine.reportUnbuildableRequest(session, name, t.getMessage)
      }

      result.onSuccess { case result =>

        val requestEndDate = nowMillis

        statsEngine.logResponse(
          session,
          name,
          ResponseTimings(startTimestamp = start, endTimestamp = requestEndDate),
          if (result) OK else KO,
          None,
          if (result) None else Some("Failed... TBD")
        )

        next ! session.markAsSucceeded
      }
    })
  }
} 

package com.ringcentral.gatling.mongo.action

import com.ringcentral.gatling.mongo.command._
import com.ringcentral.gatling.mongo.protocol.{MongoComponents, MongoProtocolKey}
import io.gatling.core.action.Action
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.structure.ScenarioContext

class MongoActionBuilder(command: MongoCommand, configuration: GatlingConfiguration) extends ActionBuilder {

  protected def mongoComponents(ctx: ScenarioContext): MongoComponents = {
    ctx.protocolComponentsRegistry.components(MongoProtocolKey)
  }

  override def build(ctx: ScenarioContext, next: Action): Action = {
    val statsEngine = ctx.coreComponents.statsEngine
    val components = mongoComponents(ctx)
    val databaseContext = components.mongoContext
    command match {
      case rawCommand: MongoRawCommand => new MongoRawCommandAction(rawCommand, databaseContext.database, statsEngine, configuration, next)
      case countCommand: MongoCountCommand => new MongoCountAction(countCommand, databaseContext.database, statsEngine, configuration, next)
      case findCommand: MongoFindCommand => new MongoFindAction(findCommand, databaseContext.database, statsEngine, configuration, next)
      case removeCommand: MongoRemoveCommand => new MongoRemoveAction(removeCommand, databaseContext.database, statsEngine, configuration, next)
      case insertCommand: MongoInsertCommand => new MongoInsertAction(insertCommand, databaseContext.database, statsEngine, configuration, next)
      case updateCommand: MongoUpdateCommand => new MongoUpdateAction(updateCommand, databaseContext.database, statsEngine, configuration, next)
    }

  }
} 

package com.github.phisgr.gatling.grpc.action

import com.github.phisgr.gatling.grpc.HeaderPair
import com.github.phisgr.gatling.grpc.check.{GrpcCheck, ResponseExtract}
import com.github.phisgr.gatling.grpc.protocol.GrpcProtocol
import io.gatling.commons.validation.Success
import io.gatling.core.action.Action
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.check.{MultipleFindCheckBuilder, ValidatorCheckBuilder}
import io.gatling.core.session.{Expression, ExpressionSuccessWrapper, Session}
import io.gatling.core.structure.ScenarioContext
import io.grpc.{CallOptions, Metadata, MethodDescriptor}

import scala.collection.breakOut

case class GrpcCallActionBuilder[Req, Res](
  requestName: Expression[String],
  method: MethodDescriptor[Req, Res],
  payload: Expression[Req],
  callOptions: Expression[CallOptions] = CallOptions.DEFAULT.expressionSuccess,
  reversedHeaders: List[HeaderPair[_]] = Nil,
  checks: List[GrpcCheck[Res]] = Nil,
  protocolOverride: Option[GrpcProtocol] = None,
  isSilent: Boolean = false
) extends ActionBuilder {
  override def build(ctx: ScenarioContext, next: Action): Action = GrpcCallAction(this, ctx, next)

  def callOptions(callOptions: Expression[CallOptions]) = copy(
    callOptions = callOptions
  )

  
  def callOptions(callOptions: => CallOptions) = copy(
    callOptions = { _: Session => Success(callOptions) }
  )

  def header[T](key: Metadata.Key[T])(value: Expression[T]) = copy(
    reversedHeaders = HeaderPair(key, value) :: reversedHeaders
  )

  private def mapToList[T, U](s: Seq[T])(f: T => U) = s.map[U, List[U]](f)(breakOut)

  def check(checks: GrpcCheck[Res]*) = copy(
    checks = this.checks ::: checks.toList
  )

  // In fact they can be added to checks using .check
  // but the type Res cannot be inferred there
  def extract[X](
    f: Res => Option[X])(
    ts: (ValidatorCheckBuilder[ResponseExtract, Res, X] => GrpcCheck[Res])*
  ) = {
    val e = ResponseExtract.extract(f)
    copy(
      checks = checks ::: mapToList(ts)(_.apply(e))
    )
  }

  def exists[X](f: Res => Option[X]) = extract(f)(_.exists.build(ResponseExtract.materializer))

  def extractMultiple[X](
    f: Res => Option[Seq[X]])(
    ts: (MultipleFindCheckBuilder[ResponseExtract, Res, X] => GrpcCheck[Res])*
  ) = {
    val e = ResponseExtract.extractMultiple[Res, X](f)
    copy(
      checks = checks ::: mapToList(ts)(_.apply(e))
    )
  }

  def target(protocol: GrpcProtocol) = copy(protocolOverride = Some(protocol))

  def silent = copy(isSilent = true)

} 

package com.github.jeanadrien.gatling.mqtt.actions

import com.github.jeanadrien.gatling.mqtt.client.MqttQoS
import com.github.jeanadrien.gatling.mqtt.client.MqttQoS.MqttQoS
import com.softwaremill.quicklens._
import io.gatling.core.action.Action
import io.gatling.core.session._
import io.gatling.core.structure.ScenarioContext
import org.fusesource.mqtt.client.QoS


case class SubscribeActionBuilder(
    topic : Expression[String],
    qos   : MqttQoS = MqttQoS.AtMostOnce
) extends MqttActionBuilder {

    def qos(newQos : MqttQoS) : SubscribeActionBuilder = this.modify(_.qos).setTo(newQos)

    def qosAtMostOnce = qos(MqttQoS.AtMostOnce)

    def qosAtLeastOnce = qos(MqttQoS.AtLeastOnce)

    def qosExactlyOnce = qos(MqttQoS.ExactlyOnce)

    override def build(
        ctx : ScenarioContext, next : Action
    ) : Action = {
        new SubscribeAction(
            mqttComponents(ctx),
            ctx.coreComponents,
            topic,
            qos,
            next
        )
    }

} 

package com.github.jeanadrien.gatling.mqtt.actions

import com.github.jeanadrien.gatling.mqtt.client.MqttQoS
import com.github.jeanadrien.gatling.mqtt.client.MqttQoS.MqttQoS
import com.softwaremill.quicklens._
import io.gatling.core.action.Action
import io.gatling.core.session._
import io.gatling.core.structure.ScenarioContext


case class PublishActionBuilder(
    topic : Expression[String],
    payload : Expression[Array[Byte]],
    qos     : MqttQoS = MqttQoS.AtMostOnce,
    retain  : Boolean = false
) extends MqttActionBuilder {

    def qos(newQos : MqttQoS) : PublishActionBuilder = this.modify(_.qos).setTo(newQos)

    def qosAtMostOnce = qos(MqttQoS.AtMostOnce)

    def qosAtLeastOnce = qos(MqttQoS.AtLeastOnce)

    def qosExactlyOnce = qos(MqttQoS.ExactlyOnce)

    def retain(newRetain : Boolean) : PublishActionBuilder = this.modify(_.retain).setTo(newRetain)

    override def build(
        ctx : ScenarioContext, next : Action
    ) : Action = {
        new PublishAction(
            mqttComponents(ctx),
            ctx.coreComponents,
            topic,
            payload,
            qos,
            retain,
            next
        )
    }

} 

package com.github.jeanadrien.gatling.mqtt.actions

import com.github.jeanadrien.gatling.mqtt.client.MqttQoS.MqttQoS
import com.github.jeanadrien.gatling.mqtt.protocol.ConnectionSettings
import com.softwaremill.quicklens._
import io.gatling.core.action.Action
import io.gatling.core.session._
import io.gatling.core.structure.ScenarioContext
import org.fusesource.mqtt.client.QoS


case class ConnectActionBuilder(
    connectionSettings : ConnectionSettings = ConnectionSettings(
        clientId = None, // default : Random
        cleanSession = None, // default : true
        userName = None,
        password = None,
        willTopic = None,
        willMessage = None,
        willQos = None,
        willRetain = None
    )
) extends MqttActionBuilder {

    def clientId(clientId : Expression[String]) = this.modify(_.connectionSettings.clientId).setTo(Some(clientId))

    def cleanSession(cleanSession : Boolean) = this.modify(_.connectionSettings.cleanSession).setTo(Some(cleanSession))

    def userName(userName : Expression[String]) = this.modify(_.connectionSettings.userName).setTo(Some(userName))

    def password(password : Expression[String]) = this.modify(_.connectionSettings.password).setTo(Some(password))

    def willTopic(willTopic : Expression[String]) = this.modify(_.connectionSettings.willTopic).setTo(Some(willTopic))

    def willMessage(willMessage : Expression[String]) = this.modify(_.connectionSettings.willMessage)
        .setTo(Some(willMessage))

    def willQos(willQos : MqttQoS) = this.modify(_.connectionSettings.willQos).setTo(Some(willQos))

    def willRetain(willRetain : Boolean) = this.modify(_.connectionSettings.willRetain).setTo(Some(willRetain))

    override def build(
        ctx : ScenarioContext, next : Action
    ) : Action = {
        new ConnectAction(
            mqttComponents(ctx),
            ctx.coreComponents,
            connectionSettings,
            next
        )
    }

} 

package com.github.jeanadrien.gatling.mqtt.actions

import com.softwaremill.quicklens._
import io.gatling.core.action.Action
import io.gatling.core.structure.ScenarioContext

import scala.concurrent.duration._


case class WaitForMessagesActionBuilder(
    timeout : FiniteDuration = 60 seconds
) extends MqttActionBuilder {

    def timeout(duration : FiniteDuration) : WaitForMessagesActionBuilder = this.modify(_.timeout).setTo(duration)

    override def build(
        ctx : ScenarioContext, next : Action
    ) : Action = {
        new WaitForMessagesAction(
            mqttComponents(ctx),
            ctx.coreComponents,
            timeout,
            next
        )
    }

} 

package com.github.jeanadrien.gatling.mqtt.actions

import com.github.jeanadrien.gatling.mqtt.client.MqttQoS
import com.github.jeanadrien.gatling.mqtt.client.MqttQoS.MqttQoS
import com.softwaremill.quicklens._
import io.gatling.core.action.Action
import io.gatling.core.session._
import io.gatling.core.structure.ScenarioContext
import org.fusesource.mqtt.client.QoS

import scala.concurrent.duration._


case class PublishAndWaitActionBuilder(
    topic : Expression[String],
    payload         : Expression[Array[Byte]],
    payloadFeedback : Array[Byte] => Array[Byte] => Boolean = PayloadComparison.sameBytesContent,
    qos             : MqttQoS = MqttQoS.AtMostOnce,
    retain          : Boolean = false,
    timeout         : FiniteDuration = 60 seconds
) extends MqttActionBuilder {

    def qos(newQos : MqttQoS) : PublishAndWaitActionBuilder = this.modify(_.qos).setTo(newQos)

    def qosAtMostOnce = qos(MqttQoS.AtMostOnce)

    def qosAtLeastOnce = qos(MqttQoS.AtLeastOnce)

    def qosExactlyOnce = qos(MqttQoS.ExactlyOnce)

    def retain(newRetain : Boolean) : PublishAndWaitActionBuilder = this.modify(_.retain).setTo(newRetain)

    def payloadFeedback(fn : Array[Byte] => Array[Byte] => Boolean) : PublishAndWaitActionBuilder = this
        .modify(_.payloadFeedback).setTo(fn)

    def timeout(duration : FiniteDuration) : PublishAndWaitActionBuilder = this.modify(_.timeout).setTo(duration)

    override def build(
        ctx : ScenarioContext, next : Action
    ) : Action = {
        new PublishAndWaitAction(
            mqttComponents(ctx),
            ctx.coreComponents,
            topic,
            payload,
            payloadFeedback,
            qos,
            retain,
            timeout,
            next
        )
    }

}