Hoist jest.mock with 2 arguments.

Summary:Comes from jestjs#796

How about throwing if second argument references outside variables? So we will be able to allow referencing later.

edit by cpojer:
bypass-lint internal eslint invariant-rule is silly
Closes jestjs#826

Differential Revision: D3096884

fb-gh-sync-id: 7fe7b24547f11046d7087e81f3c88d2f46097cba
fbshipit-source-id: 7fe7b24547f11046d7087e81f3c88d2f46097cba

Author: pvolok

CHANGELOG.md

@@ -12,6 +12,9 @@
 * Clear the terminal window when using `--watch`.
 * By default, `--watch` will now only runs tests related to changed files.
   `--watch=all` can be used to run all tests on file system changes.
+* Added the `jest.mock('moduleName', moduleFactory)` feature. `jest.mock` now
+  gets hoisted by default. `jest.doMock` was added to explicitly mock a module
+  without the hoisting feature of `babel-jest`.
 
 ## jest-cli 0.9.2, babel-jest 9.0.3
 

docs/API.md

@@ -13,9 +13,9 @@ next: troubleshooting
   - [`jest.currentTestPath()`](#jest-currenttestpath)
   - [`jest.disableAutomock()`](#jest-disableautomock)
   - [`jest.enableAutomock()`](#jest-enableautomock)
-  - [`jest.fn(implementation?)`](#jest-fn-implementation)
+  - [`jest.fn(?implementation)`](#jest-fn-implementation)
   - [`jest.genMockFromModule(moduleName)`](#jest-genmockfrommodule-modulename)
-  - [`jest.mock(moduleName)`](#jest-mock-modulename)
+  - [`jest.mock(moduleName, ?factory)`](#jest-mock-modulename-factory)
   - [`jest.runAllTicks()`](#jest-runallticks)
   - [`jest.runAllTimers()`](#jest-runalltimers)
   - [`jest.runOnlyPendingTimers()`](#jest-runonlypendingtimers)
@@ -136,7 +136,7 @@ Re-enables automatic mocking in the module loader.
 
 *Note: this method was previously called `autoMockOn`. When using `babel-jest`, calls to `enableAutomock` will automatically be hoisted to the top of the code block. Use `autoMockOn` if you want to explicitly avoid this behavior.*
 
-### `jest.fn(implementation?)`
+### `jest.fn(?implementation)`
 Returns a new, unused [mock function](#mock-functions). Optionally takes a mock
 implementation.
 
@@ -155,10 +155,27 @@ Given the name of a module, use the automatic mocking system to generate a mocke
 
 This is useful when you want to create a [manual mock](/jest/docs/manual-mocks.html) that extends the automatic mock's behavior.
 
-### `jest.mock(moduleName)`
+### `jest.mock(moduleName, ?factory)`
 Indicates that the module system should always return a mocked version of the specified module from `require()` (e.g. that it should never return the real module).
 
-This is normally useful under the circumstances where you have called [`jest.autoMockOff()`](#jest-automockoff), but still wish to specify that certain particular modules should be mocked by the module system.
+```js
+  jest.mock('moduleName');
+
+  const moduleName = require('moduleName'); // moduleName will be explicitly mocked
+```
+
+The second argument can be used to specify an explicit module factory that is being run instead of using Jest's automocking feature:
+
+```js
+  jest.mock('moduleName', () => {
+    return jest.fn(() => 42);
+  });
+
+  const moduleName = require('moduleName'); // This runs the function specified as second argument to `jest.mock`.
+  moduleName(); // Will return "42";
+```
+
+*Note: When using `babel-jest`, calls to `mock` will automatically be hoisted to the top of the code block. Use `doMock` if you want to explicitly avoid this behavior.*
 
 ### `jest.runAllTicks()`
 Exhausts the **micro**-task queue (usually interfaced in node via `process.nextTick`).
@@ -186,6 +203,8 @@ On occasion there are times where the automatically generated mock the module sy
 
 In these rare scenarios you can use this API to manually fill the slot in the module system's mock-module registry.
 
+*Note It is recommended to use [`jest.mock()`](#jest-mock-modulename-factory) instead. The `jest.mock` API's second argument is a module factory instead of the expected exported module object.*
+
 ### `jest.unmock(moduleName)`
 Indicates that the module system should never return a mocked version of the specified module from `require()` (e.g. that it should always return the real module).
 

packages/babel-plugin-jest-hoist/src/__test_modules__/e.js

@@ -0,0 +1,11 @@
+/**
+ * Copyright (c) 2014-present, Facebook, Inc. All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree. An additional grant
+ * of patent rights can be found in the PATENTS file in the same directory.
+ */
+
+'use strict';
+
+export default () => 'unmocked';

packages/babel-plugin-jest-hoist/src/__tests__/integration-automock-off-test.js

@@ -17,7 +17,7 @@ import b from '../__test_modules__/b';
 jest.disableAutomock();
 jest.mock('../__test_modules__/b');
 
-describe('babel-plugin-jest-unmock', () => {
+describe('babel-plugin-jest-hoist', () => {
   it('hoists disableAutomock call before imports', () => {
     expect(a._isMockFunction).toBe(undefined);
   });

packages/babel-plugin-jest-hoist/src/__tests__/integration-test.js

@@ -17,21 +17,38 @@ import a from '../__test_modules__/a';
 import b from '../__test_modules__/b';
 import c from '../__test_modules__/c';
 import d from '../__test_modules__/d';
+import e from '../__test_modules__/e';
 
 // These will all be hoisted above imports
 jest.unmock('react');
 jest.unmock('../__test_modules__/Unmocked');
 jest
   .unmock('../__test_modules__/c')
   .unmock('../__test_modules__/d');
+jest.mock('../__test_modules__/e', () => {
+  if (!global.CALLS) {
+    global.CALLS = 0;
+  }
+  global.CALLS++;
+
+  return {
+    _isMock: true,
+    fn: () => {
+      // The `jest.mock` transform will allow require, built-ins and globals.
+      const path = require('path');
+      const array = new Array(3);
+      array[0] = path.sep;
+      return jest.fn(() => array);
+    },
+  };
+});
 
 // These will not be hoisted
 jest.unmock('../__test_modules__/a').dontMock('../__test_modules__/b');
 jest.unmock('../__test_modules__/' + 'c');
 jest.dontMock('../__test_modules__/Mocked');
 
-
-describe('babel-plugin-jest-unmock', () => {
+describe('babel-plugin-jest-hoist', () => {
   it('hoists react unmock call before imports', () => {
     expect(typeof React).toEqual('object');
     expect(React.isValidElement.mock).toBe(undefined);
@@ -48,6 +65,27 @@ describe('babel-plugin-jest-unmock', () => {
     expect(d()).toEqual('unmocked');
   });
 
+  it('hoists mock call with 2 arguments', () => {
+    const path = require('path');
+
+    expect(e._isMock).toBe(true);
+
+    const mockFn = e.fn();
+    expect(mockFn()).toEqual([path.sep, undefined, undefined]);
+  });
+
+  it('only executes the module factories once', () => {
+    global.CALLS = 0;
+
+    require('../__test_modules__/e');
+    expect(global.CALLS).toEqual(1);
+
+    require('../__test_modules__/e');
+    expect(global.CALLS).toEqual(1);
+
+    delete global.CALLS;
+  });
+
   it('does not hoist dontMock calls before imports', () => {
     expect(Mocked._isMockFunction).toBe(true);
     expect((new Mocked()).isMocked).toEqual(undefined);

packages/babel-plugin-jest-hoist/src/index.js

@@ -8,21 +8,118 @@
 
 'use strict';
 
+function invariant(condition, message) {
+  if (!condition) {
+    throw new Error('babel-plugin-jest-hoist: ' + message);
+  }
+}
+
+// We allow `jest`, `require`, all default Node.js globals and all ES2015
+// built-ins to be used inside of a `jest.mock` factory.
+const WHITELISTED_IDENTIFIERS = {
+  jest: true,
+  require: true,
+  Infinity: true,
+  NaN: true,
+  undefined: true,
+  Object: true,
+  Function: true,
+  Boolean: true,
+  Symbol: true,
+  Error: true,
+  EvalError: true,
+  InternalError: true,
+  RangeError: true,
+  ReferenceError: true,
+  SyntaxError: true,
+  TypeError: true,
+  URIError: true,
+  Number: true,
+  Math: true,
+  Date: true,
+  String: true,
+  RegExp: true,
+  Array: true,
+  Int8Array: true,
+  Uint8Array: true,
+  Uint8ClampedArray: true,
+  Int16Array: true,
+  Uint16Array: true,
+  Int32Array: true,
+  Uint32Array: true,
+  Float32Array: true,
+  Float64Array: true,
+  Map: true,
+  Set: true,
+  WeakMap: true,
+  WeakSet: true,
+  ArrayBuffer: true,
+  DataView: true,
+  JSON: true,
+  Promise: true,
+  Generator: true,
+  GeneratorFunction: true,
+  Reflect: true,
+  Proxy: true,
+  Intl: true,
+  arguments: true,
+};
+Object.keys(global).forEach(name => WHITELISTED_IDENTIFIERS[name] = true);
+
 const JEST_GLOBAL = {name: 'jest'};
+const IDVisitor = {
+  ReferencedIdentifier(path) {
+    this.ids.add(path);
+  },
+};
 
 const FUNCTIONS = {
-  mock: {
-    checkArgs: args => args.length === 1 && args[0].isStringLiteral(),
-  },
-  unmock: {
-    checkArgs: args => args.length === 1 && args[0].isStringLiteral(),
-  },
-  disableAutomock: {
-    checkArgs: args => args.length === 0,
-  },
-  enableAutomock: {
-    checkArgs: args => args.length === 0,
+  mock: args => {
+    if (args.length === 1) {
+      return args[0].isStringLiteral();
+    } else if (args.length === 2) {
+      const moduleFactory = args[1];
+      invariant(
+        moduleFactory.isFunction(),
+        'The second argument of `jest.mock` must be a function.'
+      );
+
+      const ids = new Set();
+      const parentScope = moduleFactory.parentPath.scope;
+      moduleFactory.traverse(IDVisitor, {ids});
+      for (const id of ids) {
+        const name = id.node.name;
+        let found = false;
+        let scope = id.scope;
+
+        while (scope !== parentScope) {
+          if (scope.bindings[name]) {
+            found = true;
+            break;
+          }
+
+          scope = scope.parent;
+        }
+
+        if (!found) {
+          invariant(
+            scope.hasGlobal(name) && WHITELISTED_IDENTIFIERS[name],
+            'The second argument of `jest.mock()` is not allowed to ' +
+            'reference any outside variables.\n' +
+            'Invalid variable access: ' + name + '\n' +
+            'Whitelisted objects: ' +
+            Object.keys(WHITELISTED_IDENTIFIERS).join(', ') + '.'
+          );
+        }
+      }
+
+      return true;
+    }
+    return false;
   },
+  unmock: args => args.length === 1 && args[0].isStringLiteral(),
+  disableAutomock: args => args.length === 0,
+  enableAutomock: args => args.length === 0,
 };
 
 module.exports = babel => {
@@ -37,7 +134,7 @@ module.exports = babel => {
     return (
       property.isIdentifier() &&
       FUNCTIONS[property.node.name] &&
-      FUNCTIONS[property.node.name].checkArgs(expr.get('arguments')) &&
+      FUNCTIONS[property.node.name](expr.get('arguments')) &&
       (
         object.isIdentifier(JEST_GLOBAL) ||
         (callee.isMemberExpression() && shouldHoistExpression(object))