refactor(prefer-expect-assertions): use ESQuery selectors (#170)

jest-community · Oct 2, 2018 · be29fba · be29fba
1 parent 6fa003b
commit be29fba
Showing 1 changed file with 9 additions and 28 deletions.
diff --git a/rules/prefer-expect-assertions.js b/rules/prefer-expect-assertions.js
@@ -31,14 +31,6 @@ const isExpectAssertionsOrHasAssertionsCall = expression => {
   }
 };
 
-const isTestOrItFunction = node => {
-  return (
-    node.type === 'CallExpression' &&
-    node.callee &&
-    (node.callee.name === 'it' || node.callee.name === 'test')
-  );
-};
-
 const getFunctionFirstLine = functionBody => {
   return functionBody[0] && functionBody[0].expression;
 };
@@ -47,14 +39,6 @@ const isFirstLineExprStmt = functionBody => {
   return functionBody[0] && functionBody[0].type === 'ExpressionStatement';
 };
 
-const getTestFunctionBody = node => {
-  try {
-    return node.arguments[1].body.body;
-  } catch (e) {
-    return undefined;
-  }
-};
-
 const reportMsg = (context, node) => {
   context.report({
     message: ruleMsg,
@@ -70,18 +54,15 @@ module.exports = {
   },
   create(context) {
     return {
-      CallExpression(node) {
-        if (isTestOrItFunction(node)) {
-          const testFuncBody = getTestFunctionBody(node);
-          if (testFuncBody) {
-            if (!isFirstLineExprStmt(testFuncBody)) {
-              reportMsg(context, node);
-            } else {
-              const testFuncFirstLine = getFunctionFirstLine(testFuncBody);
-              if (!isExpectAssertionsOrHasAssertionsCall(testFuncFirstLine)) {
-                reportMsg(context, node);
-              }
-            }
+      'CallExpression[callee.name=/^it|test$/][arguments.1.body.body]'(node) {
+        const testFuncBody = node.arguments[1].body.body;
+
+        if (!isFirstLineExprStmt(testFuncBody)) {
+          reportMsg(context, node);
+        } else {
+          const testFuncFirstLine = getFunctionFirstLine(testFuncBody);
+          if (!isExpectAssertionsOrHasAssertionsCall(testFuncFirstLine)) {
+            reportMsg(context, node);
           }
         }
       },