Repository: mcollina/autocannon
Pull request #304 (2020-11-19 07:57:57)

feat: Add support for initializing custom context object
Fixes https://github.com/mcollina/autocannon/issues/279

- Added `initialContext` option which is used to initialize context instead of empty object
- Covered with tests
- Updated readme.

mcollina: can you fix the conflicts?

Patch:
diff --git a/README.md b/README.md
index bfdee12..c308aeb 100644
--- a/README.md
+++ b/README.md
@@ -286,6 +286,7 @@ Start autocannon against the given target.
     * `body`: A `String` or a `Buffer` containing the body of the request. Insert one or more randomly generated IDs into the body by including `[<id>]` where the randomly generated ID should be inserted (Must also set idReplacement to true). This can be useful in soak testing POST endpoints where one or more fields must be unique. Leave undefined for an empty body. _OPTIONAL_ default: `undefined`.
     * `form`: A `String` or an `Object` containing the multipart/form-data options or a path to the JSON file containing them
     * `headers`: An `Object` containing the headers of the request. _OPTIONAL_ default: `{}`.
+    * `initialContext`: An object that you'd like to initialize your context with. Checkout [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 [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_
@@ -321,7 +322,7 @@ Each client will loop over the `requests` array, would it contain one or several
 While going through available requests, the client will maintain a `context`: an object you can use in `onResponse` and `setupRequest` functions, to store and read some contextual data.
 Please check the `request-context.js` file in samples.
 
-Note that `context` object will be cleared when restarting to the first available request, ensuring similar runs.
+Note that `context` object will be reset to `initialContext` (or `{}` it is not provided) when restarting to the first available request, ensuring similar runs.
 
 
 ### autocannon.track(instance[, opts])
diff --git a/lib/requestIterator.js b/lib/requestIterator.js
index d758b9f..20d0275 100644
--- a/lib/requestIterator.js
+++ b/lib/requestIterator.js
@@ -3,6 +3,7 @@
 const hyperid = require('hyperid')(true)
 const inherits = require('util').inherits
 const requestBuilder = require('./httpRequestBuilder')
+const clone = require('clone')
 
 function RequestIterator (opts) {
   if (!(this instanceof RequestIterator)) {
@@ -10,7 +11,8 @@ function RequestIterator (opts) {
   }
 
   this.resetted = false
-  this.context = {}
+  this.initialContext = opts.initialContext || {}
+  this.resetContext()
   this.reqDefaults = opts
   this.requestBuilder = requestBuilder(opts)
   this.setRequests(opts.requests)
@@ -18,12 +20,16 @@ function RequestIterator (opts) {
 
 inherits(RequestIterator, Object)
 
+RequestIterator.prototype.resetContext = function () {
+  this.context = clone(this.initialContext)
+}
+
 RequestIterator.prototype.nextRequest = function () {
   this.resetted = false
   ++this.currentRequestIndex
   // when looping over available request, clear context for a fresh start
   if (this.currentRequestIndex === this.requests.length) {
-    this.context = {}
+    this.resetContext()
     this.currentRequestIndex = 0
   }
   this.currentRequest = this.requests[this.currentRequestIndex]
diff --git a/lib/run.js b/lib/run.js
index 92a6a1b..97cde3e 100644
--- a/lib/run.js
+++ b/lib/run.js
@@ -186,6 +186,9 @@ function run (opts, tracker, cb) {
       if (!opts.connectionRate && opts.overallRate) {
         url.rate = distributeNums(opts.overallRate, i)
       }
+      if (opts.initialContext) {
+        url.initialContext = opts.initialContext
+      }
 
       if (opts.tlsOptions) {
         url.tlsOptions = opts.tlsOptions
diff --git a/samples/init-context.js b/samples/init-context.js
new file mode 100644
index 0000000..b1c27ca
--- /dev/null
+++ b/samples/init-context.js
@@ -0,0 +1,45 @@
+'use strict'
+
+const http = require('http')
+const autocannon = require('../autocannon')
+
+const server = http.createServer(handle)
+
+server.listen(0, startBench)
+
+function handle (req, res) {
+  const body = []
+  // this route handler simply returns whatever it gets from response
+  req
+    .on('data', chunk => body.push(chunk))
+    .on('end', () => res.end(Buffer.concat(body)))
+}
+
+function startBench () {
+  const url = 'http://localhost:' + server.address().port
+
+  autocannon({
+    url: url,
+    connections: 1,
+    amount: 1,
+    initialContext: { user: { firstName: 'Salman' } },
+    requests: [
+      {
+        // use data from context
+        method: 'PUT',
+        setupRequest: (req, context) => ({
+          ...req,
+          path: `/user/${context.user.id}`,
+          body: JSON.stringify({
+            ...context.user,
+            lastName: 'Mitha'
+          })
+        })
+      }
+    ]
+  }, finishedBench)
+
+  function finishedBench (err, res) {
+    console.log('finished bench', err, res)
+  }
+}


Test patch:
diff --git a/test/httpClient.test.js b/test/httpClient.test.js
index ba0c23d..0f8fe3c 100644
--- a/test/httpClient.test.js
+++ b/test/httpClient.test.js
@@ -174,7 +174,7 @@ test('client ignores falsy SNI servername', (t) => {
 test('client passes through tlsOptions to connect', (t) => {
   t.plan(4)
 
-  var opts = tlsServer.address()
+  const opts = tlsServer.address()
   opts.protocol = 'https:'
   opts.tlsOptions = {
     key: fs.readFileSync(path.join(__dirname, '/key.pem')),
diff --git a/test/requestIterator.test.js b/test/requestIterator.test.js
index de2a7b2..e27ef18 100644
--- a/test/requestIterator.test.js
+++ b/test/requestIterator.test.js
@@ -408,3 +408,30 @@ test('request iterator should return next request buffer', (t) => {
   const buffer = iterator.nextRequestBuffer()
   t.same(request2Res, buffer, 'request is okay')
 })
+
+test('request iterator should initialize context from options', (t) => {
+  t.plan(3)
+
+  const opts = server.address()
+  opts.initialContext = { foo: 'bar' }
+  opts.requests = [
+    {
+      setupRequest: (req, context) => {
+        t.deepEqual(context, { foo: 'bar' }, 'context should be initialized from opts')
+        context.baz = 'qux'
+        return req
+      }
+    },
+    {
+      setupRequest: (req, context) => {
+        t.deepEqual(context, { foo: 'bar', baz: 'qux' }, 'context should contain updated data')
+        return req
+      }
+    }
+  ]
+
+  const iterator = new RequestIterator(opts)
+  iterator.nextRequest()
+  // will reset and reinit context
+  iterator.nextRequest()
+})
diff --git a/test/run.test.js b/test/run.test.js
index a38b0e6..0ec38cd 100644
--- a/test/run.test.js
+++ b/test/run.test.js
@@ -245,6 +245,23 @@ test('run should callback with an error using expectBody and requests', (t) => {
   })
 })
 
+test('run should handle context correctly', (t) => {
+  t.plan(1)
+
+  initJob({
+    url: 'http://localhost:' + server.address().port,
+    connections: 1,
+    amount: 1,
+    initialContext: { init: 'context' },
+    requests: [{
+      setupRequest: (req, context) => {
+        t.deepEqual(context, { init: 'context' }, 'context should be initialized from opts')
+        return req
+      }
+    }]
+  })
+})
+
 test('run should allow users to enter timestrings to be used for duration', (t) => {
   t.plan(3)