chore(promises): Wait for promises explicitly

See angular#68 for details
sjelin · Nov 22, 2016 · 9630ac7 · 9630ac7
1 parent e05cd3c
commit 9630ac7
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/index.js b/index.js
@@ -78,20 +78,26 @@ function wrapInControlFlow(flow, globalFn, fnName) {
 
         flow.execute(function controlFlowExecute() {
           return new webdriver.promise.Promise(function(fulfill, reject) {
+            function wrappedReject(err) {
+                var wrappedErr = new Error(err);
+                reject(wrappedErr);
+            }
             if (async) {
               // If testFn is async (it expects a done callback), resolve the promise of this
               // test whenever that callback says to.  Any promises returned from testFn are
               // ignored.
               var proxyDone = fulfill;
-              proxyDone.fail = function(err) {
-                var wrappedErr = new Error(err);
-                reject(wrappedErr);
-              };
+              proxyDone.fail = wrappedReject;
               testFn(proxyDone);
             } else {
               // Without a callback, testFn can return a promise, or it will
               // be assumed to have completed synchronously.
-              fulfill(testFn());
+              var ret = testFn();
+              if (webdriver.promise.isPromise(ret)) {
+                ret.then(fulfill, wrappedReject);
+              } else {
+                fulfill(ret);
+              }
             }
           }, flow);
         }, 'Run ' + fnName + description + ' in control flow').then(seal(done), function(err) {