handle async functions in relevant utility methods

7 files changed, 80 insertions, 5 deletions

diff --git a/lib/ensureAsync.js b/lib/ensureAsync.js
index c60d07e..8582bba 100644
--- a/lib/ensureAsync.js
+++ b/lib/ensureAsync.js
@@ -1,5 +1,6 @@
 import setImmediate from './internal/setImmediate';
 import initialParams from './internal/initialParams';
+import { isAsync } from './internal/wrapAsync';
 
 /**
  * Wrap an async function and ensure it calls its callback on a later tick of
@@ -36,6 +37,7 @@ import initialParams from './internal/initialParams';
  * async.mapSeries(args, async.ensureAsync(sometimesAsync), done);
  */
 export default function ensureAsync(fn) {
+    if (isAsync(fn)) return fn;
     return initialParams(function (args, callback) {
         var sync = true;
         args.push(function () {
diff --git a/lib/forever.js b/lib/forever.js
index d1f3256..66bc472 100644
--- a/lib/forever.js
+++ b/lib/forever.js
@@ -35,7 +35,7 @@ import wrapAsync from './internal/wrapAsync';
  */
 export default function forever(fn, errback) {
     var done = onlyOnce(errback || noop);
-    var task = ensureAsync(wrapAsync(fn));
+    var task = wrapAsync(ensureAsync(fn));
 
     function next(err) {
         if (err) return done(err);
diff --git a/lib/internal/consoleFunc.js b/lib/internal/consoleFunc.js
index 6ab2f7b..35374d7 100644
--- a/lib/internal/consoleFunc.js
+++ b/lib/internal/consoleFunc.js
@@ -1,9 +1,10 @@
 import arrayEach from 'lodash/_arrayEach';
 import rest from './rest';
+import wrapAsync from './wrapAsync';
 
 export default function consoleFunc(name) {
     return rest(function (fn, args) {
-        fn.apply(null, args.concat(rest(function (err, args) {
+        wrapAsync(fn).apply(null, args.concat(rest(function (err, args) {
             if (typeof console === 'object') {
                 if (err) {
                     if (console.error) {
diff --git a/lib/memoize.js b/lib/memoize.js
index ec83620..d752d99 100644
--- a/lib/memoize.js
+++ b/lib/memoize.js
@@ -3,6 +3,7 @@ import rest from './internal/rest';
 
 import setImmediate from './internal/setImmediate';
 import initialParams from './internal/initialParams';
+import wrapAsync from './internal/wrapAsync';
 
 function has(obj, key) {
     return key in obj;
@@ -49,6 +50,7 @@ export default function memoize(fn, hasher) {
     var memo = Object.create(null);
     var queues = Object.create(null);
     hasher = hasher || identity;
+    var _fn = wrapAsync(fn);
     var memoized = initialParams(function memoized(args, callback) {
         var key = hasher.apply(null, args);
         if (has(memo, key)) {
@@ -59,7 +61,7 @@ export default function memoize(fn, hasher) {
             queues[key].push(callback);
         } else {
             queues[key] = [callback];
-            fn.apply(null, args.concat(rest(function(args) {
+            _fn.apply(null, args.concat(rest(function(args) {
                 memo[key] = args;
                 var q = queues[key];
                 delete queues[key];
diff --git a/lib/reflect.js b/lib/reflect.js
index cc8d6b6..0dbd397 100644
--- a/lib/reflect.js
+++ b/lib/reflect.js
@@ -1,5 +1,6 @@
 import initialParams from './internal/initialParams';
 import rest from './internal/rest';
+import wrapAsync from './internal/wrapAsync';
 
 /**
  * Wraps the function in another function that always returns data even when it
@@ -41,6 +42,7 @@ import rest from './internal/rest';
  * });
  */
 export default function reflect(fn) {
+    var _fn = wrapAsync(fn);
     return initialParams(function reflectOn(args, reflectCallback) {
         args.push(rest(function callback(err, cbArgs) {
             if (err) {
@@ -60,6 +62,6 @@ export default function reflect(fn) {
             }
         }));
 
-        return fn.apply(this, args);
+        return _fn.apply(this, args);
     });
 }
diff --git a/lib/timeout.js b/lib/timeout.js
index 93fe2e1..ad050e0 100644
--- a/lib/timeout.js
+++ b/lib/timeout.js
@@ -1,4 +1,5 @@
 import initialParams from './internal/initialParams';
+import wrapAsync from './internal/wrapAsync';
 
 /**
  * Sets a time limit on an asynchronous function. If the function does not call
@@ -64,10 +65,12 @@ export default function timeout(asyncFn, milliseconds, info) {
         originalCallback(error);
     }
 
+    var fn = wrapAsync(asyncFn);
+
     return initialParams(function (args, origCallback) {
         originalCallback = origCallback;
         // setup timer and call original function
         timer = setTimeout(timeoutCallback, milliseconds);
-        asyncFn.apply(null, args.concat(injectedCallback));
+        fn.apply(null, args.concat(injectedCallback));
     });
 }
diff --git a/mocha_test/es2017/asyncFunctions.js b/mocha_test/es2017/asyncFunctions.js
index dc22f62..60a1af3 100644
--- a/mocha_test/es2017/asyncFunctions.js
+++ b/mocha_test/es2017/asyncFunctions.js
@@ -594,4 +594,69 @@ module.exports = function () {
             done(err);
         })
     });
+
+    /**
+     * Utils
+     */
+
+    it('should handle async functions in dir', (done) => {
+        async.dir(async (val) => val, 'foo');
+        setTimeout(done);
+    });
+
+    it('should handle async functions in log', (done) => {
+        async.log(async (val) => val, 'foo');
+        setTimeout(done);
+    });
+
+    it('should handle async functions in ensureAsync', () => {
+        var fn = async.ensureAsync(asyncIdentity);
+        assert(fn === asyncIdentity);
+    });
+
+    it('should handle async functions in memoize', (done) => {
+        var fn = async.memoize(asyncIdentity);
+        fn(1, () => {
+            fn(1, done);
+        })
+    });
+
+    it('should handle async functions in reflect', (done) => {
+        var fn = async.reflect(asyncIdentity);
+        fn(1, (err, result) => {
+            expect(result).to.eql({value: 1});
+            done(err);
+        })
+    });
+
+    it('should handle async functions in reflect (error case)', (done) => {
+        var thrown;
+        var fn = async.reflect(async () => {
+            thrown = new Error('foo');
+            throw thrown;
+        });
+        fn(1, (err, result) => {
+            expect(result).to.eql({error: thrown});
+            done(err);
+        })
+    });
+
+    it('should handle async functions in timeout', (done) => {
+        var fn = async.timeout(asyncIdentity, 50);
+        fn(1, (err, result) => {
+            expect(result).to.eql(1);
+            done(err);
+        })
+    });
+
+    it('should handle async functions in timeout (error case)', (done) => {
+        var fn = async.timeout(async (val) => {
+            await new Promise((resolve) => setTimeout(resolve, 100));
+            return val;
+        }, 50);
+        fn(1, (err) => {
+            expect(err.message).to.match(/timed out/);
+            done();
+        })
+    });
 }