Repository: mcollina/autocannon
Pull request #443 (2022-04-19 07:28:08)

feat: add `verifyBody` to customize the mismatch rules for response body
Hey bro. 
This PR is about customize the mismatch rules for response body.  And i added samples and unit tests. Please contact me if you have any questions. 
tks so much

GuanyuChen: It's my honor. Looking forward to the code being merged soon
GuanyuChen: > lgtm


GuanyuChen: I'm very sorry to bother you. Are there any other issues that need improvement? My project needs this feature very much 🙏🏻
mcollina: I'm currently on vacation for the next two weeks. I'll try to ship but can't guarantee.
GuanyuChen: > I'm currently on vacation for the next two weeks. I'll try to ship but can't guarantee.

get it. I'll patch it temporarily. have a good holiday! 

Patch:
diff --git a/.github/workflows/nodejs.yml b/.github/workflows/nodejs.yml
index 3736650..7ab7bab 100644
--- a/.github/workflows/nodejs.yml
+++ b/.github/workflows/nodejs.yml
@@ -5,6 +5,8 @@ on: [ push, pull_request ]
 jobs:
   test:
     runs-on: ${{ matrix.os }}
+    permissions:
+      contents: read
     strategy:
       matrix:
         os: [ubuntu-latest, windows-latest, macos-latest]
diff --git a/.gitignore b/.gitignore
index 933e16f..4727591 100644
--- a/.gitignore
+++ b/.gitignore
@@ -134,4 +134,5 @@ build/Release
 
 profile-*
 
-.devcontainer
\ No newline at end of file
+.devcontainer
+.history
\ No newline at end of file
diff --git a/README.md b/README.md
index f666d78..44d6ec5 100644
--- a/README.md
+++ b/README.md
@@ -263,6 +263,7 @@ autocannon({
   // ...
   workers: 4,
   setupClient: '/full/path/to/setup-client.js',
+  verifyBody: '/full/path/to/verify-body.js'
   requests: [
     {
       // ...
@@ -299,6 +300,7 @@ Start autocannon against the given target.
     * `headers`: An `Object` containing the headers of the request. _OPTIONAL_ default: `{}`.
     * `initialContext`: An object that you'd like to initialize your context with. Check out [an example of initializing context](./samples/init-context.js). _OPTIONAL_
     * `setupClient`: A `Function` which will be passed the `Client` object for each connection to be made. This can be used to customise each individual connection headers and body using the API shown below. The changes you make to the client in this function will take precedence over the default `body` and `headers` you pass in here. There is an example of this in the samples folder. _OPTIONAL_ default: `function noop () {}`. When using `workers`, you need to supply a file path that default exports a function instead (Check out the [workers](#workers) section for more details).
+    * `verifyBody`: A `Function` which will be passed the `Client` object for each connection to be made. This can be used to customise the validation rules for response body. This function will take precedence over the `expectBody` you pass in here. There is an example of this in the samples folder. When using `workers`, you need to supply a file path that default exports a function instead (Check out the [workers](#workers) section for more details).
     * `maxConnectionRequests`: A `Number` stating the max requests to make per connection. `amount` takes precedence if both are set. _OPTIONAL_
     * `maxOverallRequests`: A `Number` stating the max requests to make overall. Can't be less than `connections`. `maxConnectionRequests` takes precedence if both are set. _OPTIONAL_
     * `connectionRate`: A `Number` stating the rate of requests to make per second from each individual connection. No rate limiting by default. _OPTIONAL_
diff --git a/lib/httpClient.js b/lib/httpClient.js
index 95dc77e..9ce4a81 100644
--- a/lib/httpClient.js
+++ b/lib/httpClient.js
@@ -91,15 +91,18 @@ function Client (opts) {
       return this._resetConnection()
     }
     if (!this.destroyed) {
-      if (this.opts.expectBody && this.opts.expectBody !== resp.body) {
-        return this.emit('mismatch', resp.body)
-      }
-
       this.requestIterator.recordBody(resp.req, resp.headers.statusCode, resp.body, resp.headers.headers)
 
       this.emit('response', resp.headers.statusCode, resp.bytes, resp.duration, this.rate)
 
       this._doRequest()
+
+      const isFn = typeof this.opts.verifyBody === 'function'
+      if (isFn && !this.opts.verifyBody(resp.body)) {
+        return this.emit('mismatch', resp.body)
+      } else if (!isFn && this.opts.expectBody && this.opts.expectBody !== resp.body) {
+        return this.emit('mismatch', resp.body)
+      }
     }
   }
 
diff --git a/lib/run.js b/lib/run.js
index 7e0ea2f..5d914fd 100644
--- a/lib/run.js
+++ b/lib/run.js
@@ -65,6 +65,7 @@ function run (opts, tracker, cb) {
     url.body = opts.form ? opts.form.getBuffer() : opts.body
     url.headers = opts.form ? Object.assign({}, opts.headers, opts.form.getHeaders()) : opts.headers
     url.setupClient = opts.setupClient
+    url.verifyBody = opts.verifyBody
     url.timeout = opts.timeout
     url.origin = `${url.protocol}//${url.host}`
     // only keep requests for that origin, or default to requests from options
diff --git a/lib/validate.js b/lib/validate.js
index 642f728..d069ef5 100644
--- a/lib/validate.js
+++ b/lib/validate.js
@@ -25,6 +25,7 @@ function safeRequire (path) {
 
 function defaultOpts (opts) {
   const setupClient = opts.workers ? opts.setupClient : safeRequire(opts.setupClient)
+  const verifyBody = opts.workers ? opts.verifyBody : safeRequire(opts.verifyBody)
 
   const requests = opts.requests
     ? opts.requests.map((r) => {
@@ -43,6 +44,7 @@ function defaultOpts (opts) {
     ...defaultOptions,
     ...opts,
     ...(setupClient ? { setupClient } : undefined),
+    ...(verifyBody ? { verifyBody } : undefined),
     ...(requests ? { requests } : undefined)
   }
 }
@@ -78,6 +80,10 @@ module.exports = function validateOpts (opts, cbPassedIn) {
     return new Error('Invalid option setupClient, please provide a function (or file path when in workers mode)')
   }
 
+  if (!isValidFn(opts.verifyBody)) {
+    return new Error('Invalid option verifyBody, please provide a function (or file path when in workers mode)')
+  }
+
   if (!checkURL(opts.url) && !opts.socketPath) {
     return new Error('url or socketPath option required')
   }
diff --git a/lib/worker.js b/lib/worker.js
index 9c8953c..637eef5 100644
--- a/lib/worker.js
+++ b/lib/worker.js
@@ -30,6 +30,7 @@ function runTracker (opts, cb) {
     ...opts,
     ...(opts.form ? { form: multipart(opts.form) } : undefined),
     ...(opts.setupClient ? { setupClient: require(opts.setupClient) } : undefined),
+    ...(opts.verifyBody ? { verifyBody: require(opts.verifyBody) } : undefined),
     requests: opts.requests
       ? opts.requests.map(r => ({
           ...r,
diff --git a/samples/customise-verifyBody-workers.js b/samples/customise-verifyBody-workers.js
new file mode 100644
index 0000000..6d76ed3
--- /dev/null
+++ b/samples/customise-verifyBody-workers.js
@@ -0,0 +1,29 @@
+'use strict'
+
+const http = require('http')
+const path = require('path')
+const autocannon = require('../autocannon')
+
+const server = http.createServer(handle)
+
+server.listen(0, startBench)
+
+function handle (req, res) {
+  res.end('hello world')
+}
+
+function startBench () {
+  const url = 'http://localhost:' + server.address().port
+
+  autocannon({
+    url: url,
+    connections: 1000,
+    duration: 10,
+    workers: 2,
+    verifyBody: path.join(__dirname, 'helpers', 'verify-body')
+  }, finishedBench)
+
+  function finishedBench (err, res) {
+    console.log('finished bench', err, res)
+  }
+}
diff --git a/samples/customise-verifyBody.js b/samples/customise-verifyBody.js
new file mode 100644
index 0000000..d5c5230
--- /dev/null
+++ b/samples/customise-verifyBody.js
@@ -0,0 +1,31 @@
+'use strict'
+
+const http = require('http')
+const autocannon = require('../autocannon')
+
+const server = http.createServer(handle)
+
+server.listen(0, startBench)
+
+function handle (req, res) {
+  res.end('hello world')
+}
+
+function startBench () {
+  const url = 'http://localhost:' + server.address().port
+
+  autocannon({
+    url: url,
+    connections: 1000,
+    duration: 10,
+    verifyBody: verifyBody
+  }, finishedBench)
+
+  function verifyBody (body) {
+    return body.indexOf('<html>') > -1
+  }
+
+  function finishedBench (err, res) {
+    console.log('finished bench', err, res)
+  }
+}
diff --git a/samples/helpers/verify-body.js b/samples/helpers/verify-body.js
new file mode 100644
index 0000000..ae9eab8
--- /dev/null
+++ b/samples/helpers/verify-body.js
@@ -0,0 +1,5 @@
+'use strict'
+
+module.exports = (body) => {
+  return true
+}


Test patch:
diff --git a/test/run.test.js b/test/run.test.js
index ad33d8f..8154d45 100644
--- a/test/run.test.js
+++ b/test/run.test.js
@@ -327,6 +327,42 @@ test('run should produce 0 mismatches with expectBody set and matches', (t) => {
   })
 })
 
+test('run should produce count of mismatches with verifyBody set', (t) => {
+  t.plan(2)
+
+  initJob({
+    url: 'http://localhost:' + server.address().port,
+    verifyBody: function () {
+      return false
+    },
+    maxOverallRequests: 10,
+    timeout: 100
+  }, function (err, result) {
+    t.error(err)
+    t.equal(result.mismatches, 10)
+    t.end()
+  })
+})
+
+test('run should produce 0 mismatches with verifyBody set and return true', (t) => {
+  t.plan(2)
+
+  const responseBody = 'hello dave'
+  const server = helper.startServer({ body: responseBody })
+
+  initJob({
+    url: 'http://localhost:' + server.address().port,
+    verifyBody: function (body) {
+      return body.indexOf('hello') > -1
+    },
+    maxOverallRequests: 10
+  }, function (err, result) {
+    t.error(err)
+    t.equal(result.mismatches, 0)
+    t.end()
+  })
+})
+
 test('run should accept a unix socket/windows pipe', (t) => {
   t.plan(11)
 
diff --git a/test/utils/verify-body.js b/test/utils/verify-body.js
new file mode 100644
index 0000000..e1cdc8f
--- /dev/null
+++ b/test/utils/verify-body.js
@@ -0,0 +1,3 @@
+module.exports = (body) => {
+  return false
+}
diff --git a/test/workers.test.js b/test/workers.test.js
index 8f0671b..6fcdec3 100644
--- a/test/workers.test.js
+++ b/test/workers.test.js
@@ -136,6 +136,36 @@ test('setupRequest and onResponse work with workers', { skip: !hasWorkerSupport
   })
 })
 
+test('verifyBody work with workers', { skip: !hasWorkerSupport }, (t) => {
+  const server = http.createServer((req, res) => {
+    // it's not easy to assert things within setupRequest and onResponse
+    // when in workers mode. So, we set something in onResponse and use in the
+    // next Request and make sure it exist or we return 404.
+    if (req.method === 'GET' && req.url !== '/test-123?some=thing&bar=baz') {
+      res.statusCode = 404
+      res.end('NOT OK')
+      return
+    }
+
+    res.end('OK')
+  })
+  server.listen(0)
+  server.unref()
+
+  initJob({
+    url: 'http://localhost:' + server.address().port,
+    connections: 2,
+    amount: 4,
+    workers: 1,
+    verifyBody: path.join(__dirname, './utils/verify-body')
+  }, function (err, result) {
+    t.error(err)
+
+    t.equal(4, result.mismatches, 'should have 4 mismatches requests')
+    t.end()
+  })
+})
+
 test('setupClient works with workers', { skip: !hasWorkerSupport }, (t) => {
   const server = http.createServer((req, res) => {
     if (req.headers.custom !== 'my-header') {