Repository: mcollina/autocannon
Pull request #281 (2020-09-09 11:59:23)

feat: restart request iterator when setupRequest returns falsey
`setupRequest()` is very handy to perfom dynamic scenarios.
However, users do not necessarily realize they have to return a request object. If they don't, it'll crash the whole benchmark.

Besides a very common need is to skip the rest of the scenario.

This feature fixes both, by reset request iterator based on `setupRequest()`.



Patch:
diff --git a/README.md b/README.md
index 0e7e203..0048844 100644
--- a/README.md
+++ b/README.md
@@ -250,7 +250,7 @@ Start autocannon against the given target.
        * `headers`: When present, will override `opts.headers`. _OPTIONAL_.
        * `method`: When present, will override `opts.method`. _OPTIONAL_.
        * `path`: When present, will override `opts.path`. _OPTIONAL_.
-       * `setupRequest`: A `Function` you may provide to mutate the raw `request` object, e.g. `request.method = 'GET'`. It takes `request` (Object) and `context` (Object) parameters, and must return the modified request. _OPTIONAL_.
+       * `setupRequest`: A `Function` you may provide to mutate the raw `request` object, e.g. `request.method = 'GET'`. It takes `request` (Object) and `context` (Object) parameters, and must return the modified request. When it returns a falsey value, autocannon will restart from first request. _OPTIONAL_.
        * `onResponse`: A `Function` you may provide to process the received response. It takes `status` (Number), `body` (String) and `context` (Object) parameters. _OPTIONAL_.
     * `idReplacement`: A `Boolean` which enables the replacement of `[<id>]` tags within the request body with a randomly generated ID, allowing for unique fields to be sent with requests. Check out [an example of programmatic usage](./samples/using-id-replacement.js) can be found in the samples. _OPTIONAL_ default: `false`
     * `forever`: A `Boolean` which allows you to setup an instance of autocannon that restarts indefinitely after emiting results with the `done` event. Useful for efficiently restarting your instance. To stop running forever, you must cause a `SIGINT` or call the `.stop()` function on your instance. _OPTIONAL_ default: `false`
diff --git a/lib/httpRequestBuilder.js b/lib/httpRequestBuilder.js
index a4ee86d..80591f7 100644
--- a/lib/httpRequestBuilder.js
+++ b/lib/httpRequestBuilder.js
@@ -22,6 +22,7 @@ function requestBuilder (defaults) {
   // buildRequest takes an object, and turns it into a buffer representing the
   // http request.
   // second parameter is passed to setupRequest, when relevant
+  // will return null if setupRequest returns falsey result
   return function buildRequest (reqData, context) {
     // below is a hack to enable deep extending of the headers so the default
     // headers object isn't overwritten by accident
@@ -31,6 +32,9 @@ function requestBuilder (defaults) {
     reqData = Object.assign({}, defaults, reqData)
 
     reqData = reqData.setupRequest(reqData, context)
+    if (!reqData) {
+      return null
+    }
 
     // for some reason some tests fail with method === undefined
     // the reqData.method should be set to SOMETHING in this case
diff --git a/lib/requestIterator.js b/lib/requestIterator.js
index f647d19..9bebacb 100644
--- a/lib/requestIterator.js
+++ b/lib/requestIterator.js
@@ -27,7 +27,10 @@ RequestIterator.prototype.nextRequest = function () {
   this.currentRequest = this.requests[this.currentRequestIndex]
   // only builds if it has dynamic setup
   if (this.reqDefaults.idReplacement || typeof this.currentRequest.setupRequest === 'function') {
-    this.rebuildRequest()
+    if (!this.rebuildRequest()) {
+      // when setupRequest returned falsey, reset
+      this.currentRequestIndex = -1
+    }
   }
   return this.currentRequest
 }
@@ -73,12 +76,16 @@ RequestIterator.prototype.setRequest = function (newRequest) {
 }
 
 RequestIterator.prototype.rebuildRequest = function () {
+  let data
   if (this.currentRequest) {
-    const data = this.requestBuilder(this.currentRequest, this.context)
-    this.currentRequest.requestBuffer = this.reqDefaults.idReplacement
-      ? Buffer.from(data.toString().replace(/\[<id>\]/g, hyperid()))
-      : data
+    data = this.requestBuilder(this.currentRequest, this.context)
+    if (data) {
+      this.currentRequest.requestBuffer = this.reqDefaults.idReplacement
+        ? Buffer.from(data.toString().replace(/\[<id>\]/g, hyperid()))
+        : data
+    }
   }
+  return data
 }
 
 RequestIterator.prototype.recordBody = function (request, status, body) {


Test patch:
diff --git a/test/requestIterator.test.js b/test/requestIterator.test.js
index e598550..dca7f76 100644
--- a/test/requestIterator.test.js
+++ b/test/requestIterator.test.js
@@ -274,6 +274,51 @@ test('request iterator should properly mutate requests if a setupRequest functio
   t.same(iterator.currentRequest.requestBuffer, request6Res, 'request was okay')
 })
 
+test('request iterator should reset when setupRequest returns nothing', (t) => {
+  t.plan(6)
+
+  const opts = server.address()
+  opts.method = 'POST'
+
+  let i = 0
+
+  opts.requests = [
+    {
+      body: 'hello world',
+      setupRequest: req => {
+        req.body += i++
+        return i > 2 ? null : req
+      }
+    },
+    {
+      method: 'GET'
+    }
+  ]
+
+  const request1Res = `POST / HTTP/1.1\r\nHost: localhost:${server.address().port}\r\nConnection: keep-alive\r\nContent-Length: 12\r\n\r\nhello world0`
+  const request2Res = `GET / HTTP/1.1\r\nHost: localhost:${server.address().port}\r\nConnection: keep-alive\r\n\r\n`
+  const request3Res = `POST / HTTP/1.1\r\nHost: localhost:${server.address().port}\r\nConnection: keep-alive\r\nContent-Length: 12\r\n\r\nhello world1`
+
+  const iterator = new RequestIterator(opts)
+  // first POST, i becomes 1
+  t.same(iterator.currentRequest.requestBuffer.toString(), request1Res, 'request 1 was okay')
+  iterator.nextRequest()
+  // first GET, is is 1
+  t.same(iterator.currentRequest.requestBuffer.toString(), request2Res, 'request 2 was okay')
+  iterator.nextRequest()
+  // second POST, is becomes 2
+  t.same(iterator.currentRequest.requestBuffer.toString(), request3Res, 'request 3 was okay')
+  iterator.nextRequest()
+  // second GET, is is 2
+  t.same(iterator.currentRequest.requestBuffer.toString(), request2Res, 'request 4 was okay')
+  iterator.nextRequest()
+  // third POST: because is was 2, pipeline is reset. no GET, and we reuse latest POST
+  t.same(iterator.currentRequest.requestBuffer.toString(), request3Res, 'request 5 was okay')
+  iterator.nextRequest()
+  // first POST: is is higher than 2, we keep reseting pipeline, and reusing request body
+  t.same(iterator.currentRequest.requestBuffer.toString(), request3Res, 'request 6 was okay')
+})
+
 test('request iterator should maintain context while looping on requests', (t) => {
   t.plan(4)