awaitable retry

7 files changed, 63 insertions, 10 deletions

diff --git a/lib/internal/awaitify.js b/lib/internal/awaitify.js
index cad271b..57628de 100644
--- a/lib/internal/awaitify.js
+++ b/lib/internal/awaitify.js
@@ -1,7 +1,7 @@
 // conditionally promisify a function.
 // only return a promise if a callback is omitted
-export default function awaitify (asyncFn, arity) {
-    if (arity == null) throw new Error('arity is undefined')
+export default function awaitify (asyncFn, arity = asyncFn.length) {
+    if (!arity) throw new Error('arity is undefined')
     function awaitable (...args) {
         if (typeof args[arity - 1] === 'function') {
             return asyncFn.apply(this, args)
diff --git a/lib/internal/initialParams.js b/lib/internal/initialParams.js
index 4014785..9ea253a 100644
--- a/lib/internal/initialParams.js
+++ b/lib/internal/initialParams.js
@@ -1,6 +1,6 @@
 export default function (fn) {
     return function (...args/*, callback*/) {
         var callback = args.pop();
-        fn.call(this, args, callback);
+        return fn.call(this, args, callback);
     };
 }
diff --git a/lib/retry.js b/lib/retry.js
index 2a353dc..14624e1 100644
--- a/lib/retry.js
+++ b/lib/retry.js
@@ -1,5 +1,6 @@
 import noop from './internal/noop';
 import wrapAsync from './internal/wrapAsync';
+import { promiseCallback, PROMISE_SYMBOL } from './internal/promiseCallback';
 
 function constant(value) {
     return function () {
@@ -39,6 +40,7 @@ function constant(value) {
  * task has succeeded, or after the final failed attempt. It receives the `err`
  * and `result` arguments of the last attempt at completing the `task`. Invoked
  * with (err, results).
+ * @returns {Promise} a promise if no callback provided
  *
  * @example
  *
@@ -101,11 +103,11 @@ export default function retry(opts, task, callback) {
     };
 
     if (arguments.length < 3 && typeof opts === 'function') {
-        callback = task || noop;
+        callback = task || promiseCallback();
         task = opts;
     } else {
         parseTimes(options, opts);
-        callback = callback || noop;
+        callback = callback || promiseCallback();
     }
 
     if (typeof task !== 'function') {
@@ -129,6 +131,7 @@ export default function retry(opts, task, callback) {
     }
 
     retryAttempt();
+    return callback[PROMISE_SYMBOL]
 }
 
 function parseTimes(acc, t) {
diff --git a/lib/retryable.js b/lib/retryable.js
index 9422eed..ad99801 100644
--- a/lib/retryable.js
+++ b/lib/retryable.js
@@ -1,6 +1,7 @@
 import retry from './retry';
 import initialParams from './internal/initialParams';
-import wrapAsync from './internal/wrapAsync';
+import {default as wrapAsync, isAsync} from './internal/wrapAsync';
+import { promiseCallback, PROMISE_SYMBOL } from './internal/promiseCallback';
 
 /**
  * A close relative of [`retry`]{@link module:ControlFlow.retry}.  This method
@@ -14,7 +15,8 @@ import wrapAsync from './internal/wrapAsync';
  * @see [async.retry]{@link module:ControlFlow.retry}
  * @category Control Flow
  * @param {Object|number} [opts = {times: 5, interval: 0}| 5] - optional
- * options, exactly the same as from `retry`
+ * options, exactly the same as from `retry`, except for a `opts.arity` that
+ * is the arity of the `task` function, defaulting to `task.length`
  * @param {AsyncFunction} task - the asynchronous function to wrap.
  * This function will be passed any arguments passed to the returned wrapper.
  * Invoked with (...args, callback).
@@ -30,13 +32,21 @@ import wrapAsync from './internal/wrapAsync';
  *     })]
  * }, callback);
  */
-export default function (opts, task) {
+export default function retryable (opts, task) {
     if (!task) {
         task = opts;
         opts = null;
     }
+    let arity = (opts && opts.arity) || task.length
+    if (isAsync(task)) {
+        arity += 1
+    }
     var _task = wrapAsync(task);
     return initialParams((args, callback) => {
+        if (args.length < arity - 1 || callback == null) {
+            args.push(callback)
+            callback = promiseCallback()
+        }
         function taskFn(cb) {
             _task(...args, cb);
         }
@@ -44,5 +54,6 @@ export default function (opts, task) {
         if (opts) retry(opts, taskFn, callback);
         else retry(taskFn, callback);
 
+        return callback[PROMISE_SYMBOL]
     });
 }
diff --git a/lib/transform.js b/lib/transform.js
index bfc2aac..786f5d3 100644
--- a/lib/transform.js
+++ b/lib/transform.js
@@ -1,7 +1,7 @@
 import eachOf from './eachOf';
 import once from './internal/once';
 import wrapAsync from './internal/wrapAsync';
-import { promiseCallback, PROMISE_SYMBOL } from './internal/promiseCallback'
+import { promiseCallback, PROMISE_SYMBOL } from './internal/promiseCallback';
 
 /**
  * A relative of `reduce`.  Takes an Object or Array, and iterates over each
diff --git a/test/es2017/awaitableFunctions.js b/test/es2017/awaitableFunctions.js
index 5a8df07..add5a33 100644
--- a/test/es2017/awaitableFunctions.js
+++ b/test/es2017/awaitableFunctions.js
@@ -496,4 +496,43 @@ module.exports = function () {
         ], 2)
         expect(calls).to.eql([1, 1, 1, 1])
     });
+
+    it('should return a Promise: retryable', async () => {
+        expect (async.retryable.name).to.contain('retryable')
+        let counter = 0
+        const calls = []
+        const fn = async.retryable(async (a, b) => {
+            calls.push(a, b)
+            counter++
+            if (counter < 3) throw new Error()
+        })
+        const promise = fn(1, 2)
+        expect(promise.then).to.be.a('function')
+        await promise
+        expect(calls).to.eql([1, 2, 1, 2, 1, 2])
+    });
+    it('should return a Promise: retryable (arity 0)', async () => {
+        expect (async.retryable.name).to.contain('retryable')
+        let counter = 0
+        const calls = []
+        const fn = async.retryable({times: 5}, async () => {
+            calls.push(0)
+            counter++
+            if (counter < 3) throw new Error()
+        })
+        await fn()
+        expect(calls).to.eql([0, 0, 0])
+    });
+
+    it('should return a Promise: retry', async () => {
+        expect (async.retry.name).to.contain('retry')
+        let counter = 0
+        const calls = []
+        await async.retry(async () => {
+            calls.push(counter)
+            counter++
+            if (counter < 3) throw new Error()
+        })
+        expect(calls).to.eql([0, 1, 2])
+    });
 };
diff --git a/test/retry.js b/test/retry.js
index 7369f8a..e7bc36a 100644
--- a/test/retry.js
+++ b/test/retry.js
@@ -115,7 +115,7 @@ describe("retry", () => {
         async.retry((cb) => {
             calls++;
             cb("fail");
-        });
+        }).catch(() => {});
         setTimeout(() => {
             expect(calls).to.equal(5);
             done();