From c9d7dd1e249fc1f309f8d688d22fe725cb469921 Mon Sep 17 00:00:00 2001
From: Valeri Karpov <val@karpov.io>
Date: Wed, 16 May 2018 11:32:04 -0400
Subject: [PATCH] feat: add option to check wrapped function return value for
 promises

---
 index.js | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/index.js b/index.js
index 2d04cfb..d66ded2 100644
--- a/index.js
+++ b/index.js
@@ -260,6 +260,7 @@ Kareem.prototype.wrap = function(name, fn, context, args, options) {
   const _this = this;
 
   options = options || {};
+  const checkForPromise = options.checkForPromise;
 
   this.execPre(name, context, args, function(error) {
     if (error) {
@@ -273,7 +274,24 @@ Kareem.prototype.wrap = function(name, fn, context, args, options) {
     }
 
     const end = (typeof lastArg === 'function' ? args.length - 1 : args.length);
-    fn.apply(context, args.slice(0, end).concat(_cb));
+    const numParameters = fn.length;
+    const ret = fn.apply(context, args.slice(0, end).concat(_cb));
+
+    if (checkForPromise) {
+      if (ret != null && typeof ret.then === 'function') {
+        // Thenable, use it
+        return ret.then(
+          res => _cb(null, res),
+          err => _cb(err)
+        );
+      }
+
+      // If `fn()` doesn't have a callback argument and doesn't return a
+      // promise, assume it is sync
+      if (numParameters < end + 1) {
+        return _cb(null, ret);
+      }
+    }
 
     function _cb() {
       const args = arguments;