Callback implementation and tests for write and end

index.js

@@ -87,7 +87,7 @@ function compression (options) {
 
     // proxy
 
-    res.write = function write (chunk, encoding) {
+    res.write = function write (chunk, encoding, cb) {
       if (ended) {
         return false
       }
@@ -97,11 +97,11 @@ function compression (options) {
       }
 
       return stream
-        ? stream.write(toBuffer(chunk, encoding))
-        : _write.call(this, chunk, encoding)
+        ? stream.write(toBuffer(chunk, encoding), cb)
+        : _write.call(this, chunk, encoding, cb)
     }
 
-    res.end = function end (chunk, encoding) {
+    res.end = function end (chunk, encoding, cb) {
       if (ended) {
         return false
       }
@@ -116,7 +116,7 @@ function compression (options) {
       }
 
       if (!stream) {
-        return _end.call(this, chunk, encoding)
+        return _end.call(this, chunk, encoding, cb)
       }
 
       // mark ended
@@ -125,7 +125,7 @@ function compression (options) {
       // write Buffer for Node.js 0.8
       return chunk
         ? stream.end(toBuffer(chunk, encoding))
-        : stream.end()
+        : stream.end(null, null, cb)
     }
 
     res.on = function on (type, listener) {

test/compression.js

@@ -1257,6 +1257,68 @@ describe('compression()', function () {
       var socket = openSocketWithRequest(port)
     })
   })
+
+  describe('when callbacks are used', function () {
+    it('should call the passed callbacks in the order passed when compressing', function (done) {
+      var callbackOutput = []
+      var server = createServer(null, function (req, res) {
+        res.setHeader('Content-Type', 'text/plain')
+        res.write('Hello', null, function () {
+          callbackOutput.push(0)
+        })
+        res.write(' World', null, function () {
+          callbackOutput.push(1)
+        })
+        res.end(null, null, function () {
+          callbackOutput.push(2)
+        })
+      })
+
+      request(server)
+      .get('/')
+      .set('Accept-Encoding', 'gzip')
+      .expect('Content-Encoding', 'gzip')
+      .end(function (err) {
+        if (err) {
+          throw new Error(err)
+        }
+        assert.equal(callbackOutput.length, 3)
+        assert.deepEqual(callbackOutput, [0, 1, 2])
+        done()
+      })
+    })
+
+    it('should call the passed callbacks in the order passed when not compressing', function (done) {
+      var callbackOutput = []
+      var server = createServer(null, function (req, res) {
+        res.setHeader('Cache-Control', 'no-transform')
+        res.setHeader('Content-Type', 'text/plain')
+        res.write('hello,', null, function () {
+          callbackOutput.push(0)
+        })
+        res.write(' world', null, function () {
+          callbackOutput.push(1)
+        })
+        res.end(null, null, function () {
+          callbackOutput.push(2)
+        })
+      })
+
+      request(server)
+      .get('/')
+      .set('Accept-Encoding', 'gzip')
+      .expect('Cache-Control', 'no-transform')
+      .expect(shouldNotHaveHeader('Content-Encoding'))
+      .end(function (err) {
+        if (err) {
+          throw new Error(err)
+        }
+        assert.equal(callbackOutput.length, 3)
+        assert.deepEqual(callbackOutput, [0, 1, 2])
+        done()
+      })
+    })
+  })
 })
 
 function createServer (opts, fn) {