sort out whilst's test args

3 files changed, 24 insertions, 18 deletions

diff --git a/lib/doWhilst.js b/lib/doWhilst.js
index a425999..163c6a9 100644
--- a/lib/doWhilst.js
+++ b/lib/doWhilst.js
@@ -1,8 +1,9 @@
-import whilst from './whilst';
+import noop from 'lodash/noop';
+import rest from 'lodash/rest';
 
 /**
  * The post-check version of [`whilst`]{@link module:ControlFlow.whilst}. To reflect the difference in
- * the order of operations, the arguments `test` and `fn` are switched.
+ * the order of operations, the arguments `test` and `iteratee` are switched.
  *
  * `doWhilst` is to `whilst` as `do while` is to `while` in plain JavaScript.
  *
@@ -12,20 +13,23 @@ import whilst from './whilst';
  * @method
  * @see [async.whilst]{@link module:ControlFlow.whilst}
  * @category Control Flow
- * @param {Function} fn - A function which is called each time `test` passes.
- * The function is passed a `callback(err)`, which must be called once it has
- * completed with an optional `err` argument. Invoked with (callback).
+ * @param {Function} iteratee - A function which is called each time `test`
+ * passes. The function is passed a `callback(err)`, which must be called once
+ * it has completed with an optional `err` argument. Invoked with (callback).
  * @param {Function} test - synchronous truth test to perform after each
- * execution of `fn`. Invoked with Invoked with the non-error callback results
- * of `fn`.
+ * execution of `iteratee`. Invoked with Invoked with the non-error callback
+ * results of `iteratee`.
  * @param {Function} [callback] - A callback which is called after the test
- * function has failed and repeated execution of `fn` has stopped. `callback`
- * will be passed an error and any arguments passed to the final `fn`'s
- * callback. Invoked with (err, [results]);
+ * function has failed and repeated execution of `iteratee` has stopped.
+ * `callback` will be passed an error and any arguments passed to the final
+ * `iteratee`'s callback. Invoked with (err, [results]);
  */
-export default function doWhilst(fn, test, callback) {
-    var calls = 0;
-    whilst(function() {
-        return ++calls <= 1 || test.apply(this, arguments);
-    }, fn, callback);
+export default function doWhilst(iteratee, test, callback) {
+    callback = callback || noop;
+    var next = rest(function(err, args) {
+        if (err) return callback(err);
+        if (test.apply(this, args)) return iteratee(next);
+        callback.apply(null, [null].concat(args));
+    });
+    iteratee(next);
 }
diff --git a/lib/whilst.js b/lib/whilst.js
index 3311f4c..e11a141 100644
--- a/lib/whilst.js
+++ b/lib/whilst.js
@@ -41,7 +41,7 @@ export default function whilst(test, iteratee, callback) {
     if (!test()) return callback(null);
     var next = rest(function(err, args) {
         if (err) return callback(err);
-        if (test.apply(this, args)) return iteratee(next);
+        if (test()) return iteratee(next);
         callback.apply(null, [null].concat(args));
     });
     iteratee(next);
diff --git a/mocha_test/whilst.js b/mocha_test/whilst.js
index a384916..e04c2b7 100644
--- a/mocha_test/whilst.js
+++ b/mocha_test/whilst.js
@@ -8,7 +8,8 @@ describe('whilst', function(){
 
         var count = 0;
         async.whilst(
-            function () {
+            function (c) {
+                expect(c).to.equal(undefined);
                 call_order.push(['test', count]);
                 return (count < 5);
             },
@@ -57,7 +58,8 @@ describe('whilst', function(){
                 count++;
                 cb(null, count);
             },
-            function () {
+            function (c) {
+                expect(c).to.equal(count);
                 call_order.push(['test', count]);
                 return (count < 5);
             },