chore(mocha): refactor to use selenium-webdriver's mocha adapters (#…

…4013)

Closes #3985

lib/frameworks/mocha.js

@@ -1,5 +1,4 @@
 var q = require('q');
-var promise = require('selenium-webdriver').promise;
 
 /**
  * Execute the Runner's test cases through Mocha.
@@ -21,14 +20,56 @@ exports.run = function(runner, specs) {
   // wait until then to load mocha-webdriver adapters as well.
   mocha.suite.on('pre-require', function() {
     try {
-      global.after = wrapped(global.after);
-      global.afterEach = wrapped(global.afterEach);
-      global.before = wrapped(global.before);
-      global.beforeEach = wrapped(global.beforeEach);
-
-      global.it = wrapped(global.it);
-      global.it.only = wrapped(global.iit);
-      global.it.skip = wrapped(global.xit);
+      // We need to re-wrap all of the global functions, which `selenium-webdriver/testing` only
+      // does when it is required. So first we must remove it from the cache.
+      delete require.cache[require.resolve('selenium-webdriver/testing')];
+      var seleniumAdapter = require('selenium-webdriver/testing');
+
+      // Save unwrapped version
+      var unwrappedFns = {};
+      ['after', 'afterEach', 'before', 'beforeEach', 'it', 'xit', 'iit'].forEach(function(fnName) {
+        unwrappedFns[fnName] = global[fnName] || Mocha[fnName];
+      });
+
+      var wrapFn = function(seleniumWrappedFn, opt_fnName) {
+        // This does not work on functions that can be nested (e.g. `describe`)
+        return function() {
+          // Set globals to unwrapped version to avoid circular reference
+          var wrappedFns = {};
+          for (var fnName in unwrappedFns) {
+            wrappedFns[fnName] = global[fnName];
+            global[fnName] = unwrappedFns[fnName];
+          }
+
+          var args = arguments;
+          // Allow before/after hooks to use names
+          if (opt_fnName && (arguments.length > 1) && (seleniumWrappedFn.length < 2)) {
+            global[opt_fnName] = global[opt_fnName].bind(this, args[0]);
+            args = Array.prototype.slice.call(arguments, 1);
+          }
+
+          try {
+            seleniumWrappedFn.apply(this, args);
+          } finally {
+            // Restore wrapped version
+            for (fnName in wrappedFns) {
+              global[fnName] = wrappedFns[fnName];
+            }
+          }
+        };
+      };
+
+      // Wrap functions
+      global.after = wrapFn(seleniumAdapter.after, 'after');
+      global.afterEach = wrapFn(seleniumAdapter.afterEach, 'afterEach');
+      global.before = wrapFn(seleniumAdapter.before, 'before');
+      global.beforeEach = wrapFn(seleniumAdapter.beforeEach, 'beforeEach');
+
+      global.it = wrapFn(seleniumAdapter.it);
+      global.iit = wrapFn(seleniumAdapter.it.only);
+      global.xit = wrapFn(seleniumAdapter.xit);
+      global.it.only = wrapFn(seleniumAdapter.it.only);
+      global.it.skip = wrapFn(seleniumAdapter.it.skip);
     } catch (err) {
       deferred.reject(err);
     }
@@ -97,112 +138,3 @@ exports.run = function(runner, specs) {
 
   return deferred.promise;
 };
-
-
-
-var flow = (function() {
-  var initial = process.env['SELENIUM_PROMISE_MANAGER'];
-  try {
-    process.env['SELENIUM_PROMISE_MANAGER'] = '1';
-    return promise.controlFlow();
-  } finally {
-    if (initial === undefined) {
-      delete process.env['SELENIUM_PROMISE_MANAGER'];
-    } else {
-      process.env['SELENIUM_PROMISE_MANAGER'] = initial;
-    }
-  }
-})();
-
-/**
- * Wraps a function on Mocha's BDD interface so it runs inside a
- * webdriver.promise.ControlFlow and waits for the flow to complete before
- * continuing.
- * @param {!Function} globalFn The function to wrap.
- * @return {!Function} The new function.
- */
-function wrapped(globalFn) {
-  return function() {
-    if (arguments.length === 1) {
-      return globalFn(makeAsyncTestFn(arguments[0]));
-
-    } else if (arguments.length === 2) {
-      return globalFn(arguments[0], makeAsyncTestFn(arguments[1]));
-
-    } else {
-      throw Error('Invalid # arguments: ' + arguments.length);
-    }
-  };
-}
-
-/**
- * Wraps a function so that all passed arguments are ignored.
- * @param {!Function} fn The function to wrap.
- * @return {!Function} The wrapped function.
- */
-function seal(fn) {
-  return function() {
-    fn();
-  };
-}
-
-/**
- * Make a wrapper to invoke caller's test function, fn.  Run the test function
- * within a ControlFlow.
- *
- * Should preserve the semantics of Mocha's Runnable.prototype.run (See
- * https://github.com/mochajs/mocha/blob/master/lib/runnable.js#L192)
- *
- * @param {!Function} fn
- * @return {!Function}
- */
-function makeAsyncTestFn(fn) {
-  var isAsync = fn.length > 0;
-  var isGenerator = promise.isGenerator(fn);
-  if (isAsync && isGenerator) {
-    throw new TypeError(
-        'generator-based tests must not take a callback; for async testing,'
-            + ' return a promise (or yield on a promise)');
-  }
-
-  var ret = /** @type {function(this: mocha.Context)}*/ function(done) {
-    var self = this;
-    var runTest = function(resolve, reject) {
-      try {
-        if (self.isAsync) {
-          fn.call(self, function(err) { err ? reject(err) : resolve(); });
-        } else if (self.isGenerator) {
-          resolve(promise.consume(fn, self));
-        } else {
-          resolve(fn.call(self));
-        }
-      } catch (ex) {
-        reject(ex);
-      }
-    };
-
-    if (!promise.USE_PROMISE_MANAGER) {
-      new promise.Promise(runTest).then(seal(done), done);
-      return;
-    }
-
-    var runnable = this.runnable();
-    var mochaCallback = runnable.callback;
-    runnable.callback = function() {
-      flow.reset();
-      return mochaCallback.apply(this, arguments);
-    };
-
-    flow.execute(function controlFlowExecute() {
-      return new promise.Promise(function(fulfill, reject) {
-        return runTest(fulfill, reject);
-      }, flow);
-    }, runnable.fullTitle()).then(seal(done), done);
-  };
-
-  ret.toString = function() {
-    return fn.toString();
-  };
-
-  return ret;
-}