Implement custom CheckTypeStep

I am writing gatling scenarios testing a web socket that returns base64 text frames.

`

ws(“Ping WS”)
.sendText("${request}")
.check(wsAwait.within(1.second).until(1)
// .regex("${response}")
)

`

until or except return a CheckTypeStep that only give me 3 methods: regex, jsonPath, jsonpJsonPath: https://github.com/gatling/gatling/blob/master/gatling-http/src/main/scala/io/gatling/http/check/async/AsyncCheckSupport.scala#L47
I’d like to pimp that CheckTypeStep with a new method that takes 2-3 params: 1) a function string => X; 2) a predicate of X; 3) an optional predate of a collection of X.
There is no documentation anywhere showing how to write custom checks for web sockets and I am a noob in scala so I’m having a hard time trying to write this.

Any help would be greatly appreciated.

Alright I have somewhat of answer and implementation.
Turns out I had to rewrite the entire AsyncCheckDSL for need:

`

import akka.japi.Predicate
import project.Envelope
import io.gatling.commons.validation._
import io.gatling.core.session._
import io.gatling.http.check.async._
import io.gatling.core.check._
import io.gatling.core.check.extractor.{Extractor, SingleArity}
import org.asynchttpclient.util.Base64

import scala.concurrent.duration.FiniteDuration

object CustomAsyncCheckDSL {

// val myWsListen = new MyTimeoutStep(false)
val myWsAwait = new MyTimeoutStep(true)

class MyTimeoutStep(await: Boolean) {
def within(timeout: FiniteDuration) = new MyExpectationStep(await, timeout)
}

class MyExpectationStep(await: Boolean, timeout: FiniteDuration) {
def until(count: Int) = new MyCheckTypeStep(await, timeout, UntilCount(count))
def expect(count: Int) = new MyCheckTypeStep(await, timeout, ExpectedCount(count))
def expect(range: Range) = new MyCheckTypeStep(await, timeout, ExpectedRange(range))
}

class MyCheckTypeStep(await: Boolean, timeout: FiniteDuration, expectation: Expectation) {

def compareEnvelope(pred: Predicate[Envelope]) = {
new ValidatorCheckBuilder[AsyncCheck, String, Envelope, Boolean](
AsyncCheckBuilders.extender(await, timeout, expectation),
EnvelopePreparer.safeParse,
new CustomExtractor(pred).expressionSuccess
)
}
}

// inspired from HttpStatusCheckBuilder
class CustomExtractor(pred: Predicate[Envelope]) extends Extractor[Envelope, Boolean] with SingleArity {
override def name: String = “useless”
override def apply(prepared: Envelope) = {
println(“PREPARED:\n” + prepared)
pred.test(prepared) match {
case true => Some(true).success
case false => “Not matched”.failure
}
}
}

object EnvelopePreparer {
private val ErrorMapper: String => String = "Protobuf failed to parse Envelope: " + _
def safeParse(msg: String): Validation[Envelope] =
safely(ErrorMapper)(Envelope.parseFrom(Base64.decode(msg)).success)
}

}

`

Suggestions to improve this implementations are more than welcome.