Gatling classes for testing an arbitrary async function

AsyncFunctionAction.scala

import io.gatling.commons.stats.{KO, OK}
import io.gatling.commons.util.TimeHelper._
import io.gatling.core.Predef._
import io.gatling.core.action.{Action, ExitableAction}
import io.gatling.core.action.builder.ActionBuilder
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 scala.concurrent.{ExecutionContext, Future}
import scala.util.{Failure, Success}

/**
 * For gatling 2.2.2.
 * Use with a function that takes a session, and returnes a Future representing whether the function succeeded.
 * ie:
 * val asyncSearch: Session => Future[Boolean] = ???
 * scenario("foo").exec(new AsyncFunctionActionBuilder("foo", asyncSearch))
 */

case class AsyncFunctionAction(
                                baseName: String,
                                override val next: Action,
                                override val statsEngine: StatsEngine,
                                f: (Session) => Future[Boolean])
                                (implicit val ec: ExecutionContext)
                                extends ExitableAction with NameGen {
  override val name = genName(baseName)

  override def execute(session: Session): Unit = {
    val start = nowMillis
    val functionAttempt = f(session)
    functionAttempt andThen {
      case Success(true) =>
        statsEngine.logResponse(session, name, ResponseTimings(start, nowMillis), OK, None, None, List())
      case Success(false) =>
        statsEngine.logResponse(session, name, ResponseTimings(start, nowMillis), KO, None, None, List())
      case Failure(e) =>
        statsEngine.logCrash(session, name, e.toString)
        statsEngine.logResponse(session, name, ResponseTimings(start, nowMillis), KO, None, None, List())
        session.markAsFailed
    } andThen {
      case _ => next.execute(session)
    }

  }
}

class AsyncFunctionActionBuilder(name: String, f: (Session) => Future[Boolean])(implicit val ec: ExecutionContext) extends ActionBuilder {
  override def build(ctx: ScenarioContext, next: Action): Action = {
    AsyncFunctionAction(name, next, ctx.coreComponents.statsEngine,f)(ec)
  }
}