pwielgolaski
4/17/2017 - 10:01 AM

Simulation with group checker

Simulation with group checker

package io.gatling.test

import io.gatling.app.Gatling
import io.gatling.commons.util.ClockSingleton.nowMillis
import io.gatling.commons.validation.Failure
import io.gatling.core.Predef._
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.action.{Action, ExitableAction}
import io.gatling.core.session.{Expression, Session}
import io.gatling.core.stats.StatsEngine
import io.gatling.core.stats.message.ResponseTimings
import io.gatling.core.structure.{ChainBuilder, ScenarioContext}
import io.gatling.core.util.NameGen
import io.gatling.http.Predef._


trait GatlingRunner {
  def main(args: Array[String]): Unit = {
    val simulation = getClass.getName().stripSuffix("$")
    Gatling.main(args ++ Array("--simulation", simulation))
  }
}

object MySimulation extends GatlingRunner {}

class CheckerActionBuilder(nameFunction: Expression[String], checkFunction: Expression[Session]) extends ActionBuilder {
  override def build(ctx: ScenarioContext, next: Action): Action = {
    new CheckerAction(nameFunction, checkFunction, ctx.coreComponents.statsEngine, next)
  }
}

class CheckerAction(nameFunction: Expression[String], checkFunction: Expression[Session], val statsEngine: StatsEngine, val next: Action)
  extends ExitableAction with NameGen {

  override val name: String = genName("checker")

  override def execute(session: Session): Unit = recover(session) {
    val state = checkFunction(session)
    state.onFailure(msg =>
      executeFailureNext(session.markAsFailed, nameFunction, msg, next)
    )
    state.onSuccess(s => next ! s)
  }

  private def executeFailureNext(session: Session, nameFunction: Expression[String], message: String, next: Action) = {
    nameFunction(session).map(name => {
      val timings = ResponseTimings(nowMillis, nowMillis)
      statsEngine.logResponse(session, name, timings, session.status, None, Some(message))
      next ! session.logGroupRequest(timings.responseTime, session.status)
    })
  }
}

class MySimulation extends Simulation {

  val httpConf = http.baseURL("http://httpbin.org")

  val incIndex: Expression[Session] = { s =>
    val oldValue = s("index").as[Int]
    s.set("index", oldValue + 1)
  }

  val checkState: Expression[Boolean] = { s =>
    val page = s.get("page").asOption[String].map(_.toInt).getOrElse(0)
    page < 5
  }

  def checker(name: Expression[String], checkFunction: Expression[Session]) = new CheckerActionBuilder(name, checkFunction)

  def repeatUntilWithCheck(name: Expression[String], condition: Expression[Boolean], statusCheck: Expression[Session])(chain: ChainBuilder) = {
    group(name)(asLongAs(condition)(chain)
      .exec(checker("checker", statusCheck))
    )
  }

  val scn = scenario("BasicSimulation")
    .exec(s => s.set("index", 1))
    .exec(repeatUntilWithCheck("repeat until", checkState, statusCheck = _ => Failure("Missing information"))(
      exec(http("try").get("/get?page=${index}").check(jsonPath("$.args.page").find.saveAs("page"))).exec(incIndex)))

  setUp(
    scn.inject(atOnceUsers(1))
  ).protocols(httpConf)

}