fix: copy headers when using express (#154)

lib/response.js

@@ -66,16 +66,7 @@ Response.prototype.setTimeout = function (msecs, callback) {
 Response.prototype.writeHead = function () {
   const result = http.ServerResponse.prototype.writeHead.apply(this, arguments)
 
-  this._lightMyRequest.headers = Object.assign({}, this.getHeaders())
-
-  // Add raw headers
-  ;['Date', 'Connection', 'Transfer-Encoding'].forEach((name) => {
-    const regex = new RegExp('\\r\\n' + name + ': ([^\\r]*)\\r\\n')
-    const field = this._header.match(regex)
-    if (field) {
-      this._lightMyRequest.headers[name.toLowerCase()] = field[1]
-    }
-  })
+  copyHeaders(this)
 
   return result
 }
@@ -116,6 +107,9 @@ Response.prototype.addTrailers = function (trailers) {
 }
 
 function generatePayload (response) {
+  if (response._lightMyRequest.headers === null) {
+    copyHeaders(response)
+  }
   // Prepare response object
   const res = {
     raw: {
@@ -159,4 +153,17 @@ function getNullSocket () {
   })
 }
 
+function copyHeaders (response) {
+  response._lightMyRequest.headers = Object.assign({}, response.getHeaders())
+
+  // Add raw headers
+  ;['Date', 'Connection', 'Transfer-Encoding'].forEach((name) => {
+    const regex = new RegExp('\\r\\n' + name + ': ([^\\r]*)\\r\\n')
+    const field = response._header.match(regex)
+    if (field) {
+      response._lightMyRequest.headers[name.toLowerCase()] = field[1]
+    }
+  })
+}
+
 module.exports = Response

package.json

@@ -15,6 +15,7 @@
     "@types/node": "^16.0.0",
     "benchmark": "^2.1.4",
     "end-of-stream": "^1.4.4",
+    "express": "^4.17.1",
     "form-auto-content": "^2.0.0",
     "form-data": "^4.0.0",
     "pre-commit": "^1.2.2",

test/test.js

@@ -10,6 +10,7 @@ const http = require('http')
 const { finished } = require('stream')
 const eos = require('end-of-stream')
 const semver = require('semver')
+const express = require('express')
 
 const inject = require('../index')
 const parseURL = require('../lib/parseURL')
@@ -1689,3 +1690,19 @@ test('multiple calls to req.destroy should not be called', (t) => {
     t.equal(res, null)
   })
 })
+
+test('passes headers when using an express app', (t) => {
+  t.plan(2)
+
+  const app = express()
+
+  app.get('/hello', (req, res) => {
+    res.setHeader('Some-Fancy-Header', 'a very cool value')
+    res.end()
+  })
+
+  inject(app, { method: 'GET', url: 'http://example.com:8080/hello' }, (err, res) => {
+    t.error(err)
+    t.equal(res.headers['some-fancy-header'], 'a very cool value')
+  })
+})