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

lib/mockResponse.js

@@ -382,14 +382,18 @@ function createResponse(options) {
      *
      *  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 +404,15 @@ function createResponse(options) {
         mockResponse.headersSent = true;
 
         _endCalled = true;
-
-        if (data) {
-            if (data instanceof Buffer) {
-                _chunks.push(data);
-                _size += data.length;
+
+        var args = mockResponse.end.handleSignatures(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 +432,47 @@ 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: end.handleSignatures
+     *
+     *  Utility function that parses and names mockResponse.end various 
+     *  function signatures. Reference:
+     *  https://nodejs.org/api/http.html#http_response_end_data_encoding_callback
+     *  
+     */
+    mockResponse.end.handleSignatures = function (args) {
+        var data, encoding, callback;
+        if (args[0]) {
+            if (typeof args[0] === "function") {
+                callback = args[0];
+            } else {
+                data = args[0];
+            }
+        }
+        if (args[1]) {
+            if (typeof args[1] === "function") {
+                callback = args[1];
+            } else if (typeof args[1] == "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: vary

test/lib/mockResponse.spec.js

@@ -1221,6 +1221,30 @@ describe('mockResponse', function () {
 
       it('should inherit from Node OutogingMessage.end()');
 
+      it("triggers callback provided as 1st argument", function() {
+        var buff = Buffer.from('random-text');
+        var called = false;
+        function cb() { called = true; }
+        response.end(cb);
+        expect(called).to.eql(true);
+      });
+
+      it("triggers callback provided as 2nd argument", function() {
+        var buff = Buffer.from('random-text');
+        var called = false;
+        function cb() { called = true; }
+        response.end(buff, cb);
+        expect(called).to.eql(true);
+      });
+
+      it("triggers callback provided as 3rd argument", function() {
+        var buff = Buffer.from('random-text');
+        var called = false;
+        function cb() { called = true; }
+        response.end(buff, "utf8", cb);
+        expect(called).to.eql(true);
+      });
+
     });
 
   });