From c632285ae22b58c0e23ccd39d10706eb7d9a3e94 Mon Sep 17 00:00:00 2001
From: Sammy Jelin <sjelin@gmail.com>
Date: Wed, 25 Jan 2017 14:44:20 -0800
Subject: [PATCH] chore(mocha): refactor to use `selenium-webdriver`'s mocha
 adapters

Closes https://github.com/angular/protractor/issues/3985
---
 lib/frameworks/mocha.js | 156 ++++++++++------------------------------
 1 file changed, 39 insertions(+), 117 deletions(-)

diff --git a/lib/frameworks/mocha.js b/lib/frameworks/mocha.js
index ea9bc65c5..a83c26d86 100644
--- a/lib/frameworks/mocha.js
+++ b/lib/frameworks/mocha.js
@@ -21,14 +21,45 @@ 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 seleniumMochaAdapters = require('selenium-webdriver/testing');
+
+      // Save unwrapped version
+      var unwrappedAdapters = {};
+      ['after', 'afterEach', 'before', 'beforeEach', 'it', 'xit', 'iit'].forEach(function(fnName) {
+        unwrappedAdapters = global[fnName] || Mocha[fnName];
+      });
+
+      // Does not work on functions that can be nested (e.g. `describe`)
+      function wrapFn(seleniumWrappedFn) {
+        // Set globals to unwrapped version to avoid circular reference
+        var wrappedAdapters = {};
+        for (var fnName in unwrappedAdapters) {
+          wrappedAdapters[fnName] = global[fnName];
+          global[fnName] = unwrappedAdapters[fnName];
+        }
+
+        seleniumWrappedFn.apply(this, arguments);
+
+        // Restore wrapped version
+        for (var fnName in wrappedAdapters) {
+          global[fnName] = wrappedAdapters[fnName];
+        }
+      }
+
+      // Wrap functions
+      global.after = wrapFn(seleniumMochaAdapters.after);
+      global.afterEach = wrapFn(seleniumMochaAdapters.afterEach);
+      global.before = wrapFn(seleniumMochaAdapters.before);
+      global.beforeEach = wrapFn(seleniumMochaAdapters.beforeEach);
+
+      global.it = wrapFn(seleniumMochaAdapters.it);
+      global.iit = wrapFn(seleniumMochaAdapters.it.only);
+      global.xit = wrapFn(seleniumMochaAdapters.xit);
+      global.it.only = wrapFn(seleniumMochaAdapters.it.only);
+      global.it.skip = wrapFn(seleniumMochaAdapters.it.skip);
     } catch (err) {
       deferred.reject(err);
     }
@@ -97,112 +128,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;
-}