mockResponse.end(): added callback triggering once end() logic has run.

lib/mockResponse.js

@@ -376,20 +376,55 @@ function createResponse(options) {
         }
 
     };
+
+    /**
+     *  Function: getEndArguments
+     *
+     *  Utility function that parses and names parameters for the various
+     *  mockResponse.end() signatures. Reference:
+     *  https://nodejs.org/api/http.html#http_response_end_data_encoding_callback
+     *  
+     */
+    function getEndArguments(args) {
+        var data, encoding, callback;
+        if (args[0]) {
+            if (typeof args[0] === 'function') {
+                callback = args[0];
+            } else {
+                data = args[0];
+            }
+        }
+        if (args[1]) {
+            var type = typeof args[1];
+            if (type === 'function') {
+                callback = args[1];
+            } else if (type === 'string' || args[1] instanceof String){
+                encoding = args[1];
+            }
+        }
+        if (args[2] && typeof args[2] === 'function') {
+            callback = args[2];
+        }
+        return { data: data, encoding: encoding, callback: callback };
+    }
 
     /**
      *  Function: end
      *
      *  The 'end' function from node's HTTP API that finishes
      *  the connection request. This must be called.
+     *  
+     *  Signature: response.end([data[, encoding]][, callback])
      *
-     * Parameters:
+     *  Parameters:
      *
-     *  data - Optional data to return. Must be a string. Appended
-     *         to previous calls to <send>.
+     *  data - Optional data to return. Must be a string or Buffer instance. 
+     *         Appended to previous calls to <send>.
      *  encoding - Optional encoding value.
+     *  callback - Optional callback function, called once the logic has run
+     *
      */
-    mockResponse.end = function(data, encoding) {
+    mockResponse.end = function() {
         if (_endCalled) {
             // Do not emit this event twice.
             return;
@@ -400,13 +435,15 @@ function createResponse(options) {
         mockResponse.headersSent = true;
 
         _endCalled = true;
-
-        if (data) {
-            if (data instanceof Buffer) {
-                _chunks.push(data);
-                _size += data.length;
+
+        var args = getEndArguments(arguments);
+
+        if (args.data) {
+            if (args.data instanceof Buffer) {
+                _chunks.push(args.data);
+                _size += args.data.length;
             } else {
-                _data += data;
+                _data += args.data;
             }
         }
 
@@ -426,15 +463,19 @@ function createResponse(options) {
             }
         }
 
-        if (encoding) {
-            _encoding = encoding;
+        if (args.encoding) {
+            _encoding = args.encoding;
         }
-
+        
         mockResponse.emit('end');
         mockResponse.writableFinished = true; //Reference: https://nodejs.org/docs/latest-v12.x/api/http.html#http_request_writablefinished
         mockResponse.emit('finish');
-
+
+        if (args.callback) {
+            args.callback();
+        }
     };
+
 
     /**
      * Function: vary

test/lib/mockResponse.spec.js

@@ -1221,6 +1221,31 @@ describe('mockResponse', function () {
 
       it('should inherit from Node OutogingMessage.end()');
 
+      it('triggers callback provided as 1st argument', function() {
+        var callback = sinon.spy();
+        response.end(callback);
+
+        expect(callback).to.have.been.called;
+      });
+
+      it('triggers callback provided as 2nd argument', function() {
+        var payload = 'random-text';
+        var callback = sinon.spy();
+        response.end(Buffer.from(payload), callback);
+
+        expect(callback).to.have.been.called;
+        expect(response._getBuffer().toString()).to.equal(payload);
+      });
+
+      it('triggers callback provided as 3rd argument', function() {
+        var payload = 'random-text';
+        var callback = sinon.spy();
+        response.end(Buffer.from(payload), 'utf8', callback);
+
+        expect(callback).to.have.been.called;
+        expect(response._getBuffer().toString()).to.equal(payload);
+      });
+
     });
 
   });