From 5ffd7f4ae8fe3c3fd3a609cc937cfc12c656fa55 Mon Sep 17 00:00:00 2001
From: Daniel Schep <dschep@gmail.com>
Date: Tue, 27 Nov 2018 09:48:03 -0500
Subject: [PATCH] Fix invoke local when using callback in nodejs

---
 .../fixture/asyncHandlerWithSuccess.js        |   4 +
 lib/plugins/aws/invokeLocal/index.js          | 103 +++++++++---------
 lib/plugins/aws/invokeLocal/index.test.js     |  18 +++
 3 files changed, 75 insertions(+), 50 deletions(-)

diff --git a/lib/plugins/aws/invokeLocal/fixture/asyncHandlerWithSuccess.js b/lib/plugins/aws/invokeLocal/fixture/asyncHandlerWithSuccess.js
index 188e7dd99..cc77b2457 100644
--- a/lib/plugins/aws/invokeLocal/fixture/asyncHandlerWithSuccess.js
+++ b/lib/plugins/aws/invokeLocal/fixture/asyncHandlerWithSuccess.js
@@ -22,6 +22,10 @@ module.exports.withMessageByCallback = (event, context, callback) => {
   return Promise.resolve();
 };
 
+module.exports.withMessageAndDelayByCallback = (event, context, callback) => {
+  setTimeout(() => callback(null, 'Succeed'), 1);
+};
+
 module.exports.withMessageByLambdaProxy = () =>
   Promise.resolve({
     statusCode: 200,
diff --git a/lib/plugins/aws/invokeLocal/index.js b/lib/plugins/aws/invokeLocal/index.js
index 1c0bd9e5d..83ee03228 100644
--- a/lib/plugins/aws/invokeLocal/index.js
+++ b/lib/plugins/aws/invokeLocal/index.js
@@ -312,60 +312,63 @@ class AwsInvokeLocal {
       this.serverless.cli.consoleLog(JSON.stringify(result, null, 4));
     }
 
-    const callback = (err, result) => {
-      if (!hasResponded) {
-        hasResponded = true;
-        if (err) {
-          handleError.call(this, err);
-        } else if (result) {
-          handleResult.call(this, result);
+    return new Promise((resolve) => {
+      const callback = (err, result) => {
+        if (!hasResponded) {
+          hasResponded = true;
+          if (err) {
+            handleError.call(this, err);
+          } else if (result) {
+            handleResult.call(this, result);
+          }
         }
+        resolve();
+      };
+
+      const startTime = new Date();
+      const timeout = Number(this.options.functionObj.timeout)
+        || Number(this.serverless.service.provider.timeout)
+        || 6;
+      let context = {
+        awsRequestId: 'id',
+        invokeid: 'id',
+        logGroupName: this.provider.naming.getLogGroupName(this.options.functionObj.name),
+        logStreamName: '2015/09/22/[HEAD]13370a84ca4ed8b77c427af260',
+        functionVersion: 'HEAD',
+        isDefaultFunctionVersion: true,
+
+        functionName: this.options.functionObj.name,
+        memoryLimitInMB: '1024',
+
+        succeed(result) {
+          return callback(null, result);
+        },
+        fail(error) {
+          return callback(error);
+        },
+        done(error, result) {
+          return callback(error, result);
+        },
+        getRemainingTimeInMillis() {
+          return Math.max((timeout * 1000) - ((new Date()).valueOf() - startTime.valueOf()), 0);
+        },
+      };
+
+      if (customContext) {
+        context = customContext;
       }
-    };
 
-    const startTime = new Date();
-    const timeout = Number(this.options.functionObj.timeout)
-      || Number(this.serverless.service.provider.timeout)
-      || 6;
-    let context = {
-      awsRequestId: 'id',
-      invokeid: 'id',
-      logGroupName: this.provider.naming.getLogGroupName(this.options.functionObj.name),
-      logStreamName: '2015/09/22/[HEAD]13370a84ca4ed8b77c427af260',
-      functionVersion: 'HEAD',
-      isDefaultFunctionVersion: true,
+      const maybeThennable = lambda(event, context, callback);
+      if (!_.isUndefined(maybeThennable) && _.isFunction(maybeThennable.then)) {
+        return maybeThennable
+          .then(
+            callback.bind(this, null),
+            callback.bind(this)
+          );
+      }
 
-      functionName: this.options.functionObj.name,
-      memoryLimitInMB: '1024',
-
-      succeed(result) {
-        return callback(null, result);
-      },
-      fail(error) {
-        return callback(error);
-      },
-      done(error, result) {
-        return callback(error, result);
-      },
-      getRemainingTimeInMillis() {
-        return Math.max((timeout * 1000) - ((new Date()).valueOf() - startTime.valueOf()), 0);
-      },
-    };
-
-    if (customContext) {
-      context = customContext;
-    }
-
-    const maybeThennable = lambda(event, context, callback);
-    if (!_.isUndefined(maybeThennable) && _.isFunction(maybeThennable.then)) {
-      return maybeThennable
-        .then(
-          callback.bind(this, null),
-          callback.bind(this)
-        );
-    }
-
-    return maybeThennable;
+      return maybeThennable;
+    });
   }
 }
 
diff --git a/lib/plugins/aws/invokeLocal/index.test.js b/lib/plugins/aws/invokeLocal/index.test.js
index 74c2cdda5..7e3deb1b9 100644
--- a/lib/plugins/aws/invokeLocal/index.test.js
+++ b/lib/plugins/aws/invokeLocal/index.test.js
@@ -625,6 +625,24 @@ describe('AwsInvokeLocal', () => {
       });
     });
 
+    describe("by callback method even if callback isn't called syncronously", () => {
+      it('should succeed once if succeed if by callback', () => {
+        awsInvokeLocal.serverless.config.servicePath = __dirname;
+
+        return awsInvokeLocal.invokeLocalNodeJs(
+          'fixture/asyncHandlerWithSuccess',
+          'withMessageAndDelayByCallback'
+        )
+          .then(() => {
+            expect(serverless.cli.consoleLog.lastCall.args[0]).to.contain('"Succeed"');
+            const calls = serverless.cli.consoleLog.getCalls().reduce((acc, call) => (
+              _.includes(call.args[0], 'Succeed') ? [call].concat(acc) : acc
+            ), []);
+            expect(calls.length).to.equal(1);
+          });
+      });
+    });
+
     describe('with Lambda Proxy with application/json response', () => {
       it('should succeed if succeed', () => {
         awsInvokeLocal.serverless.config.servicePath = __dirname;