io.gatling.core.action.builder.ActionBuilder Scala Examples

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 com.logimethods.smartmeter.inject

import akka.actor.{ActorRef, Props}
import io.gatling.core.Predef._
import io.gatling.core.action.builder.ActionBuilder

import com.logimethods.connector.gatling.to_nats._

import scala.concurrent.duration._
import java.util.Properties
import com.logimethods.smartmeter.generate._

class NatsStreamingInjection extends Simulation {

//  println("System properties: " + System.getenv())

try {
    val natsUrl = System.getenv("NATS_URI")
    val clusterID = System.getenv("NATS_CLUSTER_ID")

    var subject = System.getenv("GATLING_TO_NATS_SUBJECT")
    if (subject == null) {
      println("No Subject has been defined through the 'GATLING_TO_NATS_SUBJECT' Environment Variable!!!")
    } else {
      println("Will emit messages to " + subject)
      val natsProtocol = NatsStreamingProtocol(natsUrl, clusterID, subject)

      val usersPerSec = System.getenv("GATLING_USERS_PER_SEC").toDouble
      val duration = System.getenv("GATLING_DURATION").toInt
      val streamingDuration = System.getenv("STREAMING_DURATION").toInt
      val slot = System.getenv("TASK_SLOT").toInt
      val randomness = System.getenv("RANDOMNESS").toFloat
      val predictionLength = System.getenv("PREDICTION_LENGTH").toInt
      val timeRoot = System.getenv("TIME_ROOT").toInt
      TimeProvider.config = Some(timeRoot)

      val natsScn = scenario("smartmeter_"+slot).exec(
          NatsStreamingBuilder(new ConsumerInterpolatedVoltageProvider(slot, usersPerSec, streamingDuration,
                                                                       randomness, predictionLength)))
      setUp(
        natsScn.inject(constantUsersPerSec(usersPerSec) during (duration minute))
      ).protocols(natsProtocol)
    }
  } catch {
    case e: Exception => {
      println(e.toString())
      e.printStackTrace()
    }
  }
} 

package com.logimethods.smartmeter.inject

import akka.actor.{ActorRef, Props}
import io.gatling.core.Predef._
import io.gatling.core.action.builder.ActionBuilder

import com.logimethods.connector.gatling.to_nats._

import scala.concurrent.duration._
import java.util.Properties
import io.nats.client.Nats.PROP_URL

import com.logimethods.smartmeter.generate._

class NatsInjection extends Simulation {

//  println("System properties: " + System.getenv())

  try {
    val properties = new Properties()
    val natsUrl = System.getenv("NATS_URI")
    properties.setProperty(PROP_URL, natsUrl)

    val subject = System.getenv("GATLING_TO_NATS_SUBJECT")
    if (subject == null) {
      println("No Subject has been defined through the 'GATLING_TO_NATS_SUBJECT' Environment Variable!!!")
    } else {
      println("Will emit messages to " + subject)
      val natsProtocol = NatsProtocol(properties, subject)

      val usersPerSec = System.getenv("GATLING_USERS_PER_SEC").toDouble
      val duration = System.getenv("GATLING_DURATION").toInt
      val streamingDuration = System.getenv("STREAMING_DURATION").toInt
      val slot = System.getenv("TASK_SLOT").toInt
      val randomness = System.getenv("RANDOMNESS").toFloat
      val predictionLength = System.getenv("PREDICTION_LENGTH").toInt
      val timeRoot = System.getenv("TIME_ROOT").toInt
      TimeProvider.config = Some(timeRoot)

      val natsScn = scenario("smartmeter_"+slot).exec(
          NatsBuilder(new ConsumerInterpolatedVoltageProvider(slot, usersPerSec, streamingDuration,
                                                              randomness, predictionLength)))

      setUp(
        natsScn.inject(constantUsersPerSec(usersPerSec) during (duration minute))
      ).protocols(natsProtocol)
    }
  } catch {
    case e: Exception => {
      println(e.toString())
      e.printStackTrace()
    }
  }
} 

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.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.ringcentral.gatling.mongo

import com.ringcentral.gatling.mongo.action.MongoActionBuilder
import com.ringcentral.gatling.mongo.check.MongoCheckSupport
import com.ringcentral.gatling.mongo.command.{MongoCommandBuilder, MongoDslBuilder}
import com.ringcentral.gatling.mongo.feeder.MongoFeederSource
import com.ringcentral.gatling.mongo.protocol.{MongoProtocol, MongoProtocolFieldsBuilder, MongoProtocolUriBuilder}
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.feeder.RecordSeqFeederBuilder
import io.gatling.core.session.Expression
import play.api.libs.json.JsObject

import scala.concurrent.duration.{FiniteDuration, _}

trait MongoDsl extends MongoCheckSupport {

  def mongo(implicit configuration: GatlingConfiguration) = MongoProtocol

  def mongo(requestName: Expression[String])(implicit configuration: GatlingConfiguration) = new MongoDslBuilder(requestName, configuration)

  def mongoFeeder(url: String, collection: String, query: String, limit: Int = 100, batchSize: Int = 0, connectionTimeout: FiniteDuration = 5 seconds,
                  receiveTimeout: FiniteDuration = 30 seconds, postProcessor: JsObject => Map[String, Any] = MongoFeederSource.defaultPostProcessor): RecordSeqFeederBuilder[Any] =
    RecordSeqFeederBuilder(MongoFeederSource(url, collection, query, limit, batchSize, connectionTimeout, receiveTimeout, postProcessor))

  implicit def mongoProtocolUriBuilder2mongoProtocol(builder: MongoProtocolUriBuilder): MongoProtocol = builder.build()

  implicit def mongoProtocolBuilder2mongoProtocol(builder: MongoProtocolFieldsBuilder): MongoProtocol = builder.build()

  implicit def mongoCommandBuilder2ActionBuilder(commandBuilder: MongoCommandBuilder)(implicit configuration: GatlingConfiguration): ActionBuilder = {
    new MongoActionBuilder(commandBuilder.build(), configuration)
  }
} 

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 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 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 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 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 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 io.gatling.keycloak

import akka.actor.ActorDSL._
import akka.actor.ActorRef
import io.gatling.core.action.Interruptable
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.config.Protocols
import io.gatling.core.result.writer.DataWriterClient
import io.gatling.core.session.{Session, Expression}
import io.gatling.core.validation.Validation


class RefreshTokenActionBuilder(requestName: Expression[String]) extends ActionBuilder{
  override def build(next: ActorRef, protocols: Protocols): ActorRef = {
    actor(actorName("refresh-token"))(new RefreshTokenAction(requestName, next))
  }
}

class RefreshTokenAction(
                          requestName: Expression[String],
                          val next: ActorRef
                        ) extends Interruptable with ExitOnFailure with DataWriterClient {
  override def executeOrFail(session: Session): Validation[_] = {
    val requestAuth: MockRequestAuthenticator = session(MockRequestAuthenticator.KEY).as[MockRequestAuthenticator]
    Blocking(() =>
      Stopwatch(() => requestAuth.getKeycloakSecurityContext.refreshExpiredToken(false))
        .check(identity, _ => "Could not refresh token")
        .recordAndContinue(this, session, requestName(session).get)
    )
  }
} 

package io.gatling.keycloak

import javax.ws.rs.core.{HttpHeaders, Response}

import akka.actor.ActorDSL._
import akka.actor.ActorRef
import io.gatling.core.action.Interruptable
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.config.Protocols
import io.gatling.core.result.writer.DataWriterClient
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.validation.{Success, Validation}
import org.keycloak.admin.client.resource.RealmResource
import org.keycloak.representations.idm.{CredentialRepresentation, UserRepresentation}

case class AddUserAttributes(
  requestName: Expression[String],
  realm: Expression[RealmResource],
  username: Expression[String],
  enabled: Expression[Boolean] = _ => Success(true),
  firstName: Option[Expression[String]] = None,
  lastName: Option[Expression[String]] = None,
  email: Option[Expression[String]] = None,
  password: Option[Expression[String]] = None,
  passwordTemporary: Option[Expression[Boolean]] = None,
  save: Option[(Session, String) => Session] = None
) {}


case class AddUserActionBuilder(
                                 attributes: AddUserAttributes
                               ) extends ActionBuilder {
  def newInstance(attributes: AddUserAttributes) = new AddUserActionBuilder(attributes)

  def enabled(enabled: Expression[Boolean]) = newInstance(attributes.copy(enabled = enabled))
  def password(password: Expression[String], temporary: Expression[Boolean]) =
    newInstance(attributes.copy(password = Some(password), passwordTemporary = Some(temporary)))
  def firstName(firstName: Expression[String]) = newInstance(attributes.copy(firstName = Some(firstName)))
  def lastName(lastName: Expression[String]) = newInstance(attributes.copy(lastName = Some(lastName)))
  def email(email: Expression[String]) = newInstance(attributes.copy(email = Some(email)))
  def saveWith[T](save: (Session, String) => Session) = newInstance(attributes.copy(save = Some(save)))

  def saveAs(name: String) = saveWith((session, id) => session.set(name, id))

  override def build(next: ActorRef, protocols: Protocols): ActorRef = {
    actor(actorName("add-user"))(new AddUserAction(attributes, next))
  }
}

class AddUserAction(
                     attributes: AddUserAttributes,
                     val next: ActorRef
                   ) extends Interruptable with ExitOnFailure with DataWriterClient {
  override def executeOrFail(session: Session): Validation[_] = {
    val user = new UserRepresentation
    user.setUsername(attributes.username(session).get)
    user.setEnabled(attributes.enabled(session).get)
    attributes.firstName.map(fn => user.setFirstName(fn(session).get))
    attributes.lastName.map(ln => user.setLastName(ln(session).get))
    attributes.email.map(e => user.setEmail(e(session).get))
    attributes.realm(session).map(realm =>
      Blocking(() =>
        Stopwatch(() => realm.users.create(user))
          .check(response => response.getStatus == 201, response => {
            val status = response.getStatusInfo.toString
            response.close()
            status
          })
          .recordAndStopOnFailure(this, session, attributes.requestName(session).get + ".create-user")
          .onSuccess(response => {
            val id = getUserId(response)
            response.close()
            val newSession = attributes.save.map(s => s(session, id)).getOrElse(session)

            attributes.password.map(password => {
              val credentials = new CredentialRepresentation
              credentials.setType(CredentialRepresentation.PASSWORD)
              credentials.setValue(password(newSession).get)
              attributes.passwordTemporary.map(a => credentials.setTemporary(a(newSession).get))
              Stopwatch(() => realm.users.get(id).resetPassword(credentials))
                .recordAndContinue(this, newSession, attributes.requestName(session).get + ".reset-password")
            }).getOrElse(next ! newSession)
          })
      )
    )
  }

  def getUserId(u: Response): String = {
    val location = u.getHeaderString(HttpHeaders.LOCATION)
    val lastSlash = location.lastIndexOf('/');
    if (lastSlash < 0) null else location.substring(lastSlash + 1)
  }
} 

package io.gatling.keycloak

import java.text.SimpleDateFormat
import java.util.{Date, Collections}

import akka.actor.ActorDSL.actor
import akka.actor.ActorRef
import io.gatling.core.action.Interruptable
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.config.Protocols
import io.gatling.core.result.writer.DataWriterClient
import io.gatling.core.session._
import io.gatling.core.validation._
import org.jboss.logging.Logger
import org.keycloak.adapters.spi.AuthOutcome
import org.keycloak.adapters.KeycloakDeploymentBuilder
import org.keycloak.adapters.spi.HttpFacade.Cookie
import org.keycloak.common.enums.SslRequired
import org.keycloak.representations.adapters.config.AdapterConfig

import scala.collection.JavaConverters._

case class AuthorizeAttributes(
  requestName: Expression[String],
  uri: Expression[String],
  cookies: Expression[List[Cookie]],
  sslRequired: SslRequired = SslRequired.EXTERNAL,
  resource: String = null,
  password: String = null,
  realm: String = null,
  realmKey: String = null,
  authServerUrl: Expression[String] = _ => Failure("no server url")
) {
  def toAdapterConfig(session: Session) = {
    val adapterConfig = new AdapterConfig
    adapterConfig.setSslRequired(sslRequired.toString)
    adapterConfig.setResource(resource)
    adapterConfig.setCredentials(Collections.singletonMap("secret", password))
    adapterConfig.setRealm(realm)
    adapterConfig.setRealmKey(realmKey)
    adapterConfig.setAuthServerUrl(authServerUrl(session).get)
    adapterConfig
  }
}

class AuthorizeActionBuilder(attributes: AuthorizeAttributes) extends ActionBuilder {
  def newInstance(attributes: AuthorizeAttributes) = new AuthorizeActionBuilder(attributes)

  def sslRequired(sslRequired: SslRequired) = newInstance(attributes.copy(sslRequired = sslRequired))
  def resource(resource: String) = newInstance(attributes.copy(resource = resource))
  def clientCredentials(password: String) = newInstance(attributes.copy(password = password))
  def realm(realm: String) = newInstance(attributes.copy(realm = realm))
  def realmKey(realmKey: String) = newInstance(attributes.copy(realmKey = realmKey))
  def authServerUrl(authServerUrl: Expression[String]) = newInstance(attributes.copy(authServerUrl = authServerUrl))

  override def build(next: ActorRef, protocols: Protocols): ActorRef = {
    actor(actorName("authorize"))(new AuthorizeAction(attributes, next))
  }
}

object AuthorizeAction {
  val logger = Logger.getLogger(classOf[AuthorizeAction])
}

class AuthorizeAction(
                       attributes: AuthorizeAttributes,
                       val next: ActorRef
                     ) extends Interruptable with ExitOnFailure with DataWriterClient {
  override def executeOrFail(session: Session): Validation[_] = {
    val facade = new MockHttpFacade()
    val deployment = KeycloakDeploymentBuilder.build(attributes.toAdapterConfig(session));
    facade.request.setURI(attributes.uri(session).get);
    facade.request.setCookies(attributes.cookies(session).get.map(c => (c.getName, c)).toMap.asJava)
    var nextSession = session
    val requestAuth: MockRequestAuthenticator = session(MockRequestAuthenticator.KEY).asOption[MockRequestAuthenticator] match {
      case Some(ra) => ra
      case None =>
        val tmp = new MockRequestAuthenticator(facade, deployment, new MockTokenStore, -1, session.userId)
        nextSession = session.set(MockRequestAuthenticator.KEY, tmp)
        tmp
    }

    Blocking(() => {
      AuthorizeAction.logger.debugf("%s: Authenticating %s%n", new SimpleDateFormat("HH:mm:ss,SSS").format(new Date()).asInstanceOf[Any], session("username").as[Any], Unit)
      Stopwatch(() => requestAuth.authenticate())
        .check(result => result == AuthOutcome.AUTHENTICATED, result => {
          AuthorizeAction.logger.warnf("%s: Failed auth %s%n", new SimpleDateFormat("HH:mm:ss,SSS").format(new Date()).asInstanceOf[Any], session("username").as[Any], Unit)
          result.toString
        })
        .recordAndContinue(AuthorizeAction.this, nextSession, attributes.requestName(session).get)
    })
  }
} 

package io.gatling.keycloak

import akka.actor.ActorRef
import akka.actor.ActorDSL._
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.action.Interruptable
import io.gatling.core.config.Protocols
import io.gatling.core.result.writer.DataWriterClient
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.validation.Validation
import org.keycloak.admin.client.resource.RealmResource
import org.keycloak.representations.idm.UserRepresentation

import scala.collection.JavaConversions._

case class FindUserAttributes(
  requestName: Expression[String],
  realm: Expression[RealmResource],
  username: Option[Expression[String]] = None,
  firstName: Option[Expression[String]] = None,
  lastName: Option[Expression[String]] = None,
  email: Option[Expression[String]] = None,
  firstResult: Option[Expression[Integer]] = None,
  maxResults: Option[Expression[Integer]] = None,
  use: Option[(Session, List[UserRepresentation]) => Session] = None
) {}

object FindUserBuilder {
  implicit def toActionBuilder(builder: FindUserBuilder) = new FindUserActionBuilder(builder.attributes)
}

class FindUserBuilder(private val attributes: FindUserAttributes) {
  def username(username: Expression[String]) = new FindUserBuilder(attributes.copy(username = Some(username)))
  def firstName(firstName: Expression[String]) = new FindUserBuilder(attributes.copy(firstName = Some(firstName)))
  def lastName(lastName: Expression[String]) = new FindUserBuilder(attributes.copy(lastName = Some(lastName)))
  def email(email: Expression[String]) = new FindUserBuilder(attributes.copy(email = Some(email)))
  def firstResult(firstResult: Expression[Integer]) = new FindUserBuilder(attributes.copy(firstResult = Some(firstResult)))
  def maxResults(maxResults: Expression[Integer]) = new FindUserBuilder(attributes.copy(maxResults = Some(maxResults)))
  def use(use: (Session, List[UserRepresentation]) => Session) = new FindUserActionBuilder(attributes.copy(use = Some(use)))
}

class FindUserActionBuilder(attributes: FindUserAttributes) extends ActionBuilder {
  override def build(next: ActorRef, protocols: Protocols): ActorRef =
    actor(actorName("find-user"))(new FindUserAction(attributes, next))
}

class FindUserAction(
                      attributes: FindUserAttributes,
                      val next: ActorRef
                    ) extends Interruptable with ExitOnFailure with DataWriterClient {
  override def executeOrFail(session: Session): Validation[_] =
    attributes.realm(session).flatMap(realm =>
      Blocking(() =>
        Stopwatch(() => {
          realm.users().search(
            attributes.username.map(a => a(session).get).orNull,
            attributes.firstName.map(a => a(session).get).orNull,
            attributes.lastName.map(a => a(session).get).orNull,
            attributes.email.map(a => a(session).get).orNull,
            attributes.firstResult.map(a => a(session).get).orNull,
            attributes.maxResults.map(a => a(session).get).orNull)
        })
          .recordAndContinue(this, session, attributes.requestName(session).get, users => {
            attributes.use.map(use => use(session, users.toList)).getOrElse(session)
          })
      )
    )
} 

package io.gatling.keycloak

import akka.actor.ActorDSL._
import akka.actor.ActorRef
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.config.Protocols
import io.gatling.core.session.Expression
import org.keycloak.admin.client.resource.RealmResource


case class Admin(requestName: Expression[String]) {
  def addUser(realm: Expression[RealmResource], username: Expression[String]) =
    new AddUserActionBuilder(new AddUserAttributes(requestName, realm, username))

  def removeUser(realm: Expression[RealmResource], userId: Expression[String]) =
    new ActionBuilder {
      override def build(next: ActorRef, protocols: Protocols): ActorRef =
        actor(actorName("remove-user"))(new RemoveUserAction(requestName, realm, userId, next))
    }

  def findUser(realm: Expression[RealmResource]) =
    new FindUserBuilder(new FindUserAttributes(requestName, realm))

}