Repository: mcollina/autocannon
Pull request #366 (2021-05-05 12:13:20)

feat: added headers to the onResponse callback.

EladIsraeli: Resolves https://github.com/mcollina/autocannon/issues/361
EladIsraeli: Yeah sure, I added docs let me know if there is anything related to the docs needed here

Patch:
diff --git a/README.md b/README.md
index e3824c3..d436d8f 100644
--- a/README.md
+++ b/README.md
@@ -55,7 +55,7 @@ Available options:
   -W/--warmup
        Use a warm up interval before starting sampling.
        This enables startup processes to finish and traffic to normalize before sampling begins
-       use -c and -d sub args e.g. `--warmpup [ -c 1 -d 3 ]`
+       use -c and -d sub args e.g. `--warmup [ -c 1 -d 3 ]`
   --on-port
         Start the command listed after -- on the command line. When it starts listening on a port,
         start sending requests to that port. A URL is still required in order to send requests to
@@ -305,7 +305,7 @@ Start autocannon against the given target.
        * `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. When it returns a falsey value, autocannon will restart from first request. When using `workers`, you need to supply a file path that default exports a function instead (Check out [workers](#workers) section for more details) _OPTIONAL_
-       * `onResponse`: A `Function` you may provide to process the received response. It takes `status` (Number), `body` (String) and `context` (Object) parameters. When using `workers`, you need to supply a file path that default exports a function instead (Check out [workers](#workers) section for more details) _OPTIONAL_
+       * `onResponse`: A `Function` you may provide to process the received response. It takes `status` (Number), `body` (String) `context` (Object) parameters and `headers` (Key-Value Object). When using `workers`, you need to supply a file path that default exports a function instead (Check out [workers](#workers) section for more details) _OPTIONAL_
     * `har`: an `Object` of parsed [HAR](https://w3c.github.io/web-performance/specs/HAR/Overview.html) content. Autocannon will extra and use `entries.request`: `requests`, `method`, `form` and `body` options will be ignored. _NOTE_: you must ensure that entries are targeting the same domain as `url` option. _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/autocannon.js b/autocannon.js
index e3ed8e3..91db9a3 100755
--- a/autocannon.js
+++ b/autocannon.js
@@ -154,7 +154,7 @@ function parseArguments (argvs) {
   })
 
   if (argv.input) {
-    argv.body = fs.readFileSync(argv.input)
+    argv.body = fs.readFileSync(argv.input, 'utf8')
   }
 
   if (argv.headers) {
diff --git a/lib/httpClient.js b/lib/httpClient.js
index a5ee381..027d95f 100644
--- a/lib/httpClient.js
+++ b/lib/httpClient.js
@@ -95,7 +95,7 @@ function Client (opts) {
         return this.emit('mismatch', resp.body)
       }
 
-      this.requestIterator.recordBody(resp.req, resp.headers.statusCode, 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)
 
diff --git a/lib/requestIterator.js b/lib/requestIterator.js
index 0a1c4eb..1b1a561 100644
--- a/lib/requestIterator.js
+++ b/lib/requestIterator.js
@@ -5,6 +5,16 @@ const inherits = require('util').inherits
 const requestBuilder = require('./httpRequestBuilder')
 const clone = require('clone')
 
+const toHeaderKeyValue = (headersArray) => {
+  const headersKeyValue = {}
+
+  for (let i = 0; i < headersArray.length - 1; ++i) {
+    headersKeyValue[headersArray[i]] = headersArray[i + 1]
+  }
+
+  return headersKeyValue
+}
+
 function RequestIterator (opts) {
   if (!(this instanceof RequestIterator)) {
     return new RequestIterator(opts)
@@ -88,7 +98,7 @@ RequestIterator.prototype.rebuildRequest = function () {
   let data
   this.resetted = false
   if (this.currentRequest) {
-    this.currentRequest.headers = this.headers
+    this.currentRequest.headers = this.currentRequest.headers || this.headers
     data = this.requestBuilder(this.currentRequest, this.context)
     if (data) {
       this.currentRequest.requestBuffer = this.reqDefaults.idReplacement
@@ -106,9 +116,9 @@ RequestIterator.prototype.rebuildRequest = function () {
   return data
 }
 
-RequestIterator.prototype.recordBody = function (request, status, body) {
+RequestIterator.prototype.recordBody = function (request, status, body, headers) {
   if (request && typeof request.onResponse === 'function') {
-    request.onResponse(status, body, this.context)
+    request.onResponse(status, body, this.context, toHeaderKeyValue(headers || []))
   }
 }
 
diff --git a/package.json b/package.json
index f3e2b5e..3a71c1a 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "autocannon",
-  "version": "7.2.0",
+  "version": "7.3.0",
   "description": "Fast HTTP benchmarking tool written in Node.js",
   "main": "autocannon.js",
   "bin": {


Test patch:
diff --git a/test/argumentParsing.test.js b/test/argumentParsing.test.js
index 9131f2a..4ed1b13 100644
--- a/test/argumentParsing.test.js
+++ b/test/argumentParsing.test.js
@@ -2,6 +2,7 @@
 
 const test = require('tap').test
 const Autocannon = require('../autocannon')
+const fs = require('fs')
 
 test('parse argument', (t) => {
   t.plan(4)
@@ -158,3 +159,21 @@ test('parse argument with multiple url', (t) => {
   t.equal(args.url[0], 'http://localhost/foo/bar')
   t.equal(args.url[1], 'http://localhost/baz/qux')
 })
+
+test('parse argument with input file and multiple workers', (t) => {
+  t.plan(3)
+
+  const inputPath = 'help.txt'
+  const args = Autocannon.parseArguments([
+    '-m', 'POST',
+    '-w', '2',
+    '-a', 10,
+    '-i', inputPath,
+    '-H', 'Content-Type=application/json',
+    'http://localhost/foo/bar'
+  ])
+
+  t.equal(args.url, 'http://localhost/foo/bar')
+  t.equal(args.method, 'POST')
+  t.equal(args.body, fs.readFileSync(inputPath, 'utf8'))
+})
diff --git a/test/helper.js b/test/helper.js
index 4777f55..da7c750 100644
--- a/test/helper.js
+++ b/test/helper.js
@@ -24,7 +24,8 @@ function startServer (opts) {
   server.listen(opts.socketPath || 0)
 
   function handle (req, res) {
-    let { statusCode, body } = opts.responses[server.autocannonRequests] || {}
+    let { statusCode, body, headers } = opts.responses[server.autocannonRequests] || {}
+
     server.autocannonRequests++
 
     if (!statusCode) {
@@ -38,6 +39,10 @@ function startServer (opts) {
     res.statusCode = statusCode
     const reply = () => {
       const bodyToWrite = typeof body === 'function' ? body(req) : body
+
+      if (headers) {
+        res.writeHead(statusCode, headers)
+      }
       res.end(statusCode < 200 ? undefined : bodyToWrite)
     }
 
diff --git a/test/httpClient.test.js b/test/httpClient.test.js
index 5f33eac..e1d663f 100644
--- a/test/httpClient.test.js
+++ b/test/httpClient.test.js
@@ -287,6 +287,34 @@ test('client supports custom headers', (t) => {
   })
 })
 
+test('client supports custom headers in requests', (t) => {
+  t.plan(2)
+
+  const opts = server.address()
+  opts.requests = [
+    {
+      path: '/',
+      method: 'POST',
+      headers: {
+        'content-type': 'application/json'
+      },
+      body: JSON.stringify({
+        foo: 'example'
+      })
+    }
+  ]
+  const client = new Client(opts)
+
+  server.once('request', (req, res) => {
+    t.equal(req.headers['content-type'], 'application/json', 'custom header matches')
+  })
+
+  client.on('response', (statusCode, length) => {
+    t.equal(statusCode, 200, 'status code matches')
+    client.destroy()
+  })
+})
+
 test('client supports host custom header', (t) => {
   t.plan(2)
 
diff --git a/test/requestIterator.test.js b/test/requestIterator.test.js
index fac1860..1ca6a72 100644
--- a/test/requestIterator.test.js
+++ b/test/requestIterator.test.js
@@ -44,6 +44,35 @@ test('request iterator should create requests with overwritten defaults', (t) =>
     'request is okay')
 })
 
+test('request iterator should use headers from requests', (t) => {
+  t.plan(2)
+
+  const opts = server.address()
+  opts.requests = [
+    {
+      path: '/',
+      method: 'POST',
+      headers: {
+        'content-type': 'application/json'
+      },
+      body: JSON.stringify({ foo: 'bar' })
+    },
+    {
+      path: '/2',
+      method: 'POST',
+      headers: {
+        'content-type': 'text/html'
+      },
+      body: JSON.stringify({ foo: 'bar' })
+    }
+  ]
+  const iterator = new RequestIterator(opts)
+
+  t.same(iterator.currentRequest.headers['content-type'], 'application/json')
+  iterator.nextRequest()
+  t.same(iterator.currentRequest.headers['content-type'], 'text/html')
+})
+
 test('request iterator should create requests with overwritten defaults', (t) => {
   t.plan(3)
 
@@ -481,3 +510,25 @@ test('request iterator should use the same headers when set', (t) => {
   iterator.setRequest()
   t.same(iterator.currentRequest.requestBuffer, request6Res, iterator.currentRequest.requestBuffer.toString())
 })
+
+test('request iterator should invoke onResponse callback with headers in it', (t) => {
+  t.plan(4)
+
+  const opts = server.address()
+  opts.requests = [
+    {
+      onResponse: (status, body, context, headers) => {
+        t.same(status, 200)
+        t.same(body, 'ok')
+        t.same(context, {})
+        t.same(headers, { 'set-cookie': 123 })
+      }
+    }
+  ]
+
+  const iterator = new RequestIterator(opts)
+  iterator.recordBody(iterator.currentRequest, 200, 'ok', ['set-cookie', 123])
+  iterator.nextRequest()
+
+  t.end()
+})
diff --git a/test/run.test.js b/test/run.test.js
index 80a6d51..fa1c57e 100644
--- a/test/run.test.js
+++ b/test/run.test.js
@@ -917,3 +917,26 @@ test('The warmup should not pollute the main result set', (t) => {
     t.type(result.warmup, 'object')
   })
 })
+
+test('should get headers passed from server onResponse callback', async t => {
+  t.plan(3)
+  const server = helper.startServer({ responses: [{ statusCode: 200, body: 'ok', headers: { 'set-cookie': 123 } }] })
+
+  await initJob({
+    url: 'http://localhost:' + server.address().port,
+    connections: 1,
+    amount: 1,
+    requests: [
+      {
+        method: 'GET',
+        onResponse (status, body, context, headers) {
+          t.same(status, 200)
+          t.same(body, 'ok')
+          t.same(headers['set-cookie'], 123)
+        }
+      }
+    ]
+  })
+
+  t.end()
+})