Gatling Integration Util

IntegrationUtils.scala

import java.util.concurrent.TimeUnit

import io.gatling.core.Predef._
import io.gatling.http.Predef._
import io.gatling.http.request.builder._
import io.gatling.core.session._
import io.gatling.http.check.HttpCheck
import io.gatling.http.request.StringBody
import io.gatling.core.structure.{PopulatedScenarioBuilder, ChainBuilder}
import io.gatling.http.request.builder.AbstractHttpRequestBuilder.toActionBuilder

class GatlingSim extends Simulation {

  class GatlingTests(int: IntegrationUtils) {
    import int._

    val gatTest = createTest("Is Gatling there?",
      step("Grab root", false, "/", headers = textHeader)
    )
  }

  val tests = new GatlingTests(new IntegrationUtils("http://gatling-tool.org", "https://gatling-tool.org"))

  val httpProtocol = http
    .acceptHeader("text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8")
    .acceptEncodingHeader("gzip, deflate")
    .acceptLanguageHeader("en-us,en;q=0.8")
    .disableCaching
    .disableResponseChunksDiscarding
    .extraInfoExtractor((protocol: String, status: Status, session: Session, req: Request, res: Response) => {
    val body = req.getMethod match {
      case x if x.equalsIgnoreCase("get") => ""
      case x if x.equalsIgnoreCase("post") => req.getStringData
      case x if x.equalsIgnoreCase("put") => req.getStringData
      case x if x.equalsIgnoreCase("patch") => req.getStringData
      case x if x.equalsIgnoreCase("delete") => ""
      case _ => ""
    }
    List(req.getUrl, body, res.statusCode, res.body.string())
  })

  setUp(tests.gatTest).protocols(httpProtocol)
}

class IntegrationUtils(val httpServer: String, val httpsServer: String) {

  val wxdLib = new tests.Wxd(this)

  val upsLib = new tests.Ups(this)

  val cmsLib = new tests.Cms(this)

  val wxCheckinLib = new tests.WxCheckin(this)

  val solrLib = new tests.Solr(this)

  val xLib = new tests.X(this)

  val textHeader = Map("Content-Type" -> "text/plain")
  val jsonHeader = Map("Content-Type" -> "application/json")
  
  type Step = HttpRequestWithParamsBuilder

  private def chainParam(r: Step, kv: (String, Expression[String])) = r.queryParam(kv._1, kv._2)
  def addParams(req: Step, params: Iterable[(String, Expression[String])]) = (req /: params)(chainParam)

  def addChecks(req: Step, stat: Int, checks: Iterable[HttpCheck]) = (req.check(status is stat) /: checks)(_.check(_))

  private def commonStep(
    title: String,
    secure: Boolean,
    url: String,
    method: String = "get",
    body: StringBody,
    stat: Int,
    headers: Map[String, String],
    params: List[(String, Expression[String])],
    checks: Iterable[HttpCheck]
  ): Step = {
    val server = if (secure) httpsServer; else httpServer
    val fullStepName = title + " (" + method + ")"
    val test = method match {
      case "get"|"GET" => http(fullStepName).httpRequestWithParams(method, server + url)
      case _ => http(fullStepName).httpRequestWithParams(method, server + url).body(body)
    }
    addChecks(addParams(test.headers(headers), params), stat, checks)
  }

  def step(
    title: String,
    secure: Boolean,
    url: String,
    method: String = "get",
    body: StringBody = new StringBody(""),
    stat: Int = 200,
    headers: Map[String, String] = jsonHeader,
    params: List[(String, Expression[String])] = Nil,
    checks: Iterable[HttpCheck] = Nil
  ): Step = {
    val p = if (secure) params; else ("api", ExpressionWrapper("key").expression) :: params
    commonStep(title, secure, url, method, body, stat, headers, p, checks)
  }

  private val JSONP_CALLBACK: (String, Expression[String]) = ("jsonp", "gatling_callback")
  def stepJsonP(
    title: String,
    secure: Boolean,
    url: String,
    method: String = "get",
    body: StringBody = new StringBody(""),
    stat: Int = 200,
    headers: Map[String, String] = jsonHeader,
    params: List[(String, Expression[String])] = Nil,
    checks: Iterable[HttpCheck] = Nil
  ): Step = {
    val pTemp = JSONP_CALLBACK :: params
    val p = method match {
      case "get" | "Get" | "GET" => pTemp
      case _ => (method, body.string) :: pTemp
    }
    val c: List[HttpCheck] = jsonpJsonPath("$.status").ofType[Int].is(stat) :: checks.toList
    commonStep(title, secure, url, "get", new StringBody(""), 200, headers, p, c)
  }

  def createTest(title: String, tests: Step*): PopulatedScenarioBuilder = createTest(title, tests.toList)
  def createTest(title: String, tests: Iterable[Step]) = {
    def chainIt(parent: ChainBuilder, http: Step) = {
      val next = parent.exec(http)
      http.commonAttributes.method match {
        case "get" | "Get" | "GET" => next
        case _ => next.pause("200", TimeUnit.MILLISECONDS)
      }
    }
    val chain = (exec(tests.head) /: tests.tail)(chainIt)
    scenario(title).group(title)(chain).inject(atOnceUsers(1))
  }
}