fix handling of expression functions (#3)

Author: hiddentao

package.json

@@ -3,7 +3,9 @@
   "version": "0.0.1",
   "description": "Babel plugin that adds the number of assertions found in each test with expect.assertions(n)",
   "main": "dist/index.js",
-  "files": ["dist"],
+  "files": [
+    "dist"
+  ],
   "scripts": {
     "build": "babel src -d dist --ignore *.test.js",
     "contributor": "all-contributors add",
@@ -22,7 +24,14 @@
     "type": "git",
     "url": "git+https://github.com/mattphillips/babel-jest-assertions.git"
   },
-  "keywords": ["babel", "jest", "plugin", "assertions", "expect", "test"],
+  "keywords": [
+    "babel",
+    "jest",
+    "plugin",
+    "assertions",
+    "expect",
+    "test"
+  ],
   "author": "Matt Phillips <[email protected]> (mattphillips.io)",
   "license": "MIT",
   "bugs": {
@@ -46,12 +55,20 @@
     "prettier": "^1.6.1"
   },
   "lint-staged": {
-    "*.js": ["yarn prettier", "git add"]
+    "*.js": [
+      "yarn prettier",
+      "git add"
+    ]
   },
   "jest": {
     "testEnvironment": "node",
-    "testPathIgnorePatterns": ["/node_modules/", "/fixtures/"],
-    "coveragePathIgnorePatterns": ["/node_modules/"],
+    "testPathIgnorePatterns": [
+      "/node_modules/",
+      "/fixtures/"
+    ],
+    "coveragePathIgnorePatterns": [
+      "/node_modules/"
+    ],
     "coverageThreshold": {
       "global": {
         "branches": 100,
@@ -60,5 +77,8 @@
         "statements": 100
       }
     }
+  },
+  "dependencies": {
+    "babel-generator": "^6.26.0"
   }
 }

src/__snapshots__/async.test.js.snap

@@ -270,3 +270,22 @@ describe('.add', () => {
 });
 "
 `;
+
+exports[`Handles expression functions 1`] = `
+"
+describe('.add', () => {
+  it('returns 1 when given 0 and 1', async () =>
+    expect(add(1, 0)).toEqual(1));
+});
+
+      ↓ ↓ ↓ ↓ ↓ ↓
+
+describe('.add', () => {
+  it('returns 1 when given 0 and 1', async () => {
+    expect.assertions(1);
+    expect.hasAssertions();
+    return expect(add(1, 0)).toEqual(1);
+  });
+});
+"
+`;

src/__snapshots__/sync.test.js.snap

@@ -130,7 +130,7 @@ describe('.add', () => {
 
     const a = 1; // expect(a).toEqual(1);
     const b = 2; /* expect(b).toEqual(2); */
-    
+
     expect(add(1, 0)).toEqual(1);
     /*
       expect(add(6, 1).toEqual(7));
@@ -163,3 +163,22 @@ describe('.add', () => {
 });
 "
 `;
+
+exports[`Handles expression functions 1`] = `
+"
+describe('.add', () => {
+  it('returns 1 when given 0 and 1', () =>
+    expect(add(1, 0)).toEqual(1));
+});
+
+      ↓ ↓ ↓ ↓ ↓ ↓
+
+describe('.add', () => {
+  it('returns 1 when given 0 and 1', () => {
+    expect.assertions(1);
+    expect.hasAssertions();
+    return expect(add(1, 0)).toEqual(1);
+  });
+});
+"
+`;

src/async.test.js

@@ -137,6 +137,15 @@ pluginTester({
         });
       });
       `
+    },
+    'Handles expression functions': {
+      snapshot: true,
+      code: `
+      describe('.add', () => {
+        it('returns 1 when given 0 and 1', async () =>
+          expect(add(1, 0)).toEqual(1));
+      });
+      `
     }
   }
 });

src/index.js

@@ -1,6 +1,8 @@
+const generate = require('babel-generator').default;
+
 const looksLike = require('./utils/looks-like');
 
-module.exports = function({ template, transformFromAst, types }) {
+module.exports = function({ template, types }) {
   return {
     name: 'babel-assertions',
     visitor: {
@@ -38,7 +40,17 @@ module.exports = function({ template, transformFromAst, types }) {
           return;
         }
 
-        const { code } = transformFromAst(types.Program([path.node]));
+        // get the test case body
+        let body = path.node.expression.arguments[1].body;
+
+        // if it's an expression, e.g: () => (expression)
+        if (types.isCallExpression(body)) {
+          // convert it into a block statement: () => { return (expression); }
+          body = path.node.expression.arguments[1].body = types.blockStatement([types.returnStatement(body)]);
+        }
+
+        // generate the code
+        const { code } = generate(body);
 
         // remove comments
         const normalisedCode = code.replace(/\/\/(.*)/g, '').replace(/\/\*([\s\S]*?)\*\//g, '');
@@ -47,7 +59,8 @@ module.exports = function({ template, transformFromAst, types }) {
         const containsExpectAssertions = normalisedCode.includes('expect.assertions(');
         const containsHasAssertions = normalisedCode.includes('expect.hasAssertions()');
 
-        const args = path.node.expression.arguments[1].body.body;
+        const args = body.body;
+
         if (!containsHasAssertions) {
           const hasAssertions = template('expect.hasAssertions();')();
           args.unshift(hasAssertions);

src/sync.test.js

@@ -80,14 +80,23 @@ pluginTester({
 
           const a = 1; // expect(a).toEqual(1);
           const b = 2; /* expect(b).toEqual(2); */
-          
+
           expect(add(1, 0)).toEqual(1);
           /*
             expect(add(6, 1).toEqual(7));
             */
         });
       });
       `
+    },
+    'Handles expression functions': {
+      snapshot: true,
+      code: `
+      describe('.add', () => {
+        it('returns 1 when given 0 and 1', () =>
+          expect(add(1, 0)).toEqual(1));
+      });
+      `
     }
   }
 });