feat(options): support assertBeforeCallbackName and assertAfterCallba…

…ckName
azu · Aug 24, 2019 · d2717eb · d2717eb
1 parent c542815
commit d2717eb
Show file tree

Hide file tree

Showing 12 changed files with 142 additions and 27 deletions.
diff --git a/src/ast-utils.ts b/src/ast-utils.ts
@@ -8,7 +8,10 @@ import {
     isIdentifier,
     isLiteral,
     isNullLiteral,
-    isDirective
+    isDirective,
+    callExpression,
+    identifier,
+    stringLiteral
 } from "@babel/types";
 import template from "@babel/template";
 
@@ -36,61 +39,114 @@ export const ERROR_COMMENT_PATTERN = /^([a-zA-Z]*?Error)/;
 export const PROMISE_RESOLVE_COMMENT_PATTERN = /^Resolve:\s*(.*?)\s*$/;
 export const PROMISE_REJECT_COMMENT_PATTERN = /^Reject:\s*(.*?)\s*$/;
 
-export function wrapAssert(actualNode: any, expectedNode: any): any {
+export interface wrapAssertOptions {
+    // callback name before assert
+    assertBeforeCallbackName?: string;
+    // callback name after assert
+    assertAfterCallbackName?: string;
+}
+
+export function wrapAssert(
+    {
+        actualNode,
+        expectedNode,
+        commentExpression,
+        id
+    }: { actualNode: any; expectedNode: any; commentExpression: string; id: string },
+    options: wrapAssertOptions
+): any {
     assert.notStrictEqual(typeof expectedNode, "undefined");
     const ACTUAL_NODE = actualNode;
     const EXPECTED_NODE = expectedNode;
+    const BEFORE_CALLBACK = options.assertBeforeCallbackName
+        ? callExpression(identifier(options.assertBeforeCallbackName), [stringLiteral(id)])
+        : undefined;
+    const AFTER_CALLBACK = options.assertAfterCallbackName
+        ? callExpression(identifier(options.assertAfterCallbackName), [stringLiteral(id)])
+        : undefined;
     if (isConsole(actualNode)) {
         const args = actualNode.arguments;
         const firstArgument = args[0];
-        return wrapAssert(firstArgument, expectedNode);
+        return wrapAssert(
+            {
+                actualNode: firstArgument,
+                expectedNode,
+                commentExpression,
+                id
+            },
+            options
+        );
     } else if (isIdentifier(expectedNode) && ERROR_COMMENT_PATTERN.test(expectedNode.name)) {
-        return template`assert.throws(function() {
+        return template`BEFORE_CALLBACK;assert.throws(function() {
                     ACTUAL_NODE
-               })`({
+               });AFTER_CALLBACK;`({
+            BEFORE_CALLBACK,
+            AFTER_CALLBACK,
             ACTUAL_NODE
         });
     } else if (expectedNode.type === "Resolve") {
         // getExpressionNodeFromCommentValue define the type
         const ARGS = isConsole(actualNode) ? actualNode.arguments[0] : actualNode;
         return template`Promise.resolve(ARGS).then(v => {
-            ${wrapAssert({ type: "Identifier", name: "v" }, expectedNode.node)}
+            ${wrapAssert(
+                {
+                    actualNode: { type: "Identifier", name: "v" },
+                    expectedNode: expectedNode.node,
+                    commentExpression,
+                    id
+                },
+                options
+            )}
             return v;
         });`({
             ARGS
         });
     } else if (expectedNode.type === "Reject") {
         const ARGS = isConsole(actualNode) ? actualNode.arguments[0] : actualNode;
-        return template`assert.rejects(ARGS)`({
+        return template`BEFORE_CALLBACK;assert.rejects(ARGS);AFTER_CALLBACK;`({
+            BEFORE_CALLBACK,
+            AFTER_CALLBACK,
             ARGS
         });
     } else if (isIdentifier(expectedNode) && expectedNode.name === "NaN") {
-        return template`assert.ok(isNaN(ACTUAL_NODE));`({
+        return template`BEFORE_CALLBACK;assert.ok(isNaN(ACTUAL_NODE));AFTER_CALLBACK;`({
+            BEFORE_CALLBACK,
+            AFTER_CALLBACK,
             ACTUAL_NODE
         });
     } else if (isNullLiteral(expectedNode)) {
-        return template`assert.strictEqual(ACTUAL_NODE, null)`({
+        return template`BEFORE_CALLBACK;assert.strictEqual(ACTUAL_NODE, null);AFTER_CALLBACK;`({
+            BEFORE_CALLBACK,
+            AFTER_CALLBACK,
             ACTUAL_NODE
         });
     } else if (isIdentifier(expectedNode) && expectedNode.name === "undefined") {
-        return template`assert.strictEqual(ACTUAL_NODE, undefined)`({
+        return template`BEFORE_CALLBACK;assert.strictEqual(ACTUAL_NODE, undefined);AFTER_CALLBACK`({
+            BEFORE_CALLBACK,
+            AFTER_CALLBACK,
             ACTUAL_NODE
         });
     } else if (isLiteral(expectedNode)) {
         // Handle Directive Prorogue as string literal
         if (isDirective(ACTUAL_NODE)) {
-            return template`assert.strictEqual(ACTUAL_NODE, EXPECTED_NODE)`({
+            return template`BEFORE_CALLBACK;assert.strictEqual(ACTUAL_NODE, EXPECTED_NODE);AFTER_CALLBACK;`({
+                BEFORE_CALLBACK,
+                AFTER_CALLBACK,
                 ACTUAL_NODE: (ACTUAL_NODE.value as any).extra.raw,
                 EXPECTED_NODE
             });
         } else {
-            return template`assert.strictEqual(ACTUAL_NODE, EXPECTED_NODE)`({
+            return template`BEFORE_CALLBACK;assert.strictEqual(ACTUAL_NODE, EXPECTED_NODE);AFTER_CALLBACK;`({
+                BEFORE_CALLBACK,
+                AFTER_CALLBACK,
                 ACTUAL_NODE,
                 EXPECTED_NODE
             });
         }
     } else {
-        return template`assert.deepStrictEqual(ACTUAL_NODE, EXPECTED_NODE)`({
+        return template`BEFORE_CALLBACK;assert.deepStrictEqual(ACTUAL_NODE, EXPECTED_NODE);AFTER_CALLBACK;`({
+            BEFORE_CALLBACK,
+            AFTER_CALLBACK,
             ACTUAL_NODE,
             EXPECTED_NODE
         });

diff --git a/src/comment-to-assert.ts b/src/comment-to-assert.ts
@@ -5,7 +5,8 @@ import {
     PROMISE_REJECT_COMMENT_PATTERN,
     PROMISE_RESOLVE_COMMENT_PATTERN,
     tryGetCodeFromComments,
-    wrapAssert
+    wrapAssert,
+    wrapAssertOptions
 } from "./ast-utils";
 import { transformFromAstSync } from "@babel/core";
 import { identifier, isExpressionStatement, File } from "@babel/types";
@@ -48,9 +49,9 @@ function getExpressionNodeFromCommentValue(string: string): { type: string } & {
     }
 }
 
-export interface toAssertFromSourceOptions {
+export type toAssertFromSourceOptions = {
     babel?: ParserOptions;
-}
+} & wrapAssertOptions;
 
 /**
  * transform code to asserted code
@@ -65,30 +66,43 @@ export function toAssertFromSource(code: string, options?: toAssertFromSourceOpt
     if (!ast) {
         throw new Error("Can not parse the code");
     }
-    const output = toAssertFromAST(ast);
+    const output = toAssertFromAST(ast, options);
     const babelFileResult = transformFromAstSync(output, code, { comments: true });
     if (!babelFileResult) {
         throw new Error("can not generate from ast: " + JSON.stringify(output));
     }
     return babelFileResult.code;
 }
 
-export interface toAssertFromASTOptions {}
-
 /**
  * transform AST to asserted AST.
  */
-export function toAssertFromAST(ast: File, _options: toAssertFromASTOptions = {}) {
+export function toAssertFromAST(ast: File, options: wrapAssertOptions = {}) {
     const replaceSet = new Set();
+    let id = 0;
     traverse(ast, {
         exit(path) {
             if (!replaceSet.has(path.node) && path.node.trailingComments) {
                 const commentExpression = tryGetCodeFromComments(path.node.trailingComments);
                 if (commentExpression) {
                     const commentExpressionNode = getExpressionNodeFromCommentValue(commentExpression);
                     const actualNode = isExpressionStatement(path.node) ? path.node.expression : path.node;
-                    const replacement = wrapAssert(actualNode, commentExpressionNode);
-                    path.replaceWith(replacement);
+                    const replacement = wrapAssert(
+                        {
+                            actualNode: actualNode,
+                            expectedNode: commentExpressionNode,
+                            commentExpression,
+                            id: String(`id:${id++}`)
+                        },
+                        options
+                    );
+                    if (Array.isArray(replacement)) {
+                        // prevent ∞ loop
+                        path.node.trailingComments = null;
+                        path.replaceWithMultiple(replacement);
+                    } else {
+                        path.replaceWith(replacement);
+                    }
                     replaceSet.add(path.node);
                 }
             }

diff --git a/test/snapshot-test.ts b/test/snapshot-test.ts
@@ -38,12 +38,23 @@ ${JSON.stringify(actual)}
             if (typeof actual !== "string") {
                 throw new Error("actual is not string");
             }
-            if (options.asyncCallbackName) {
+            if (options.assertAfterCallbackName && options.assertBeforeCallbackName) {
+                // finish after all called
+                let actualCallCount = 0;
+                const totalCountOfAssert = actual.split(options.assertBeforeCallbackName).length - 1;
                 vm.runInContext(
                     actual,
                     vm.createContext({
                         assert,
-                        done
+                        [options.assertBeforeCallbackName]: () => {
+                            // nope
+                        },
+                        [options.assertAfterCallbackName]: () => {
+                            actualCallCount++;
+                            if (actualCallCount === totalCountOfAssert) {
+                                done();
+                            }
+                        }
                     })
                 );
             } else {

diff --git a/test/snapshots/Promise.reject/output.js b/test/snapshots/Promise.reject/output.js
@@ -1 +1 @@
-assert.rejects(Promise.reject(new Error("message"))); // => Reject: "message"
+assert.rejects(Promise.reject(new Error("message"))); // => Reject: "message"
diff --git a/test/snapshots/Promise.resolve/output.js b/test/snapshots/Promise.resolve/output.js
@@ -1,4 +1,4 @@
 Promise.resolve(Promise.resolve(1)).then(v => {
   assert.strictEqual(v, 1);
   return v;
-}); // => Resolve: 1
+}); // => Resolve: 1
diff --git a/test/snapshots/callback.NaN/input.js b/test/snapshots/callback.NaN/input.js
@@ -0,0 +1 @@
+NaN; // => NaN
diff --git a/test/snapshots/callback.NaN/options.js b/test/snapshots/callback.NaN/options.js
@@ -0,0 +1,4 @@
+module.exports = {
+    assertBeforeCallbackName: "beforeCallback",
+    assertAfterCallbackName: "afterCallback"
+};
diff --git a/test/snapshots/callback.NaN/output.js b/test/snapshots/callback.NaN/output.js
@@ -0,0 +1,3 @@
+beforeCallback("id:0");
+assert.ok(isNaN(NaN));
+afterCallback("id:0");
diff --git a/test/snapshots/callback.multiple-assert/input.js b/test/snapshots/callback.multiple-assert/input.js
@@ -0,0 +1,4 @@
+1; // => 1
+"str"; // => "str"
+[1, 2, 3]; // => [1,2,3]
+Promise.resolve(1); // => Resolve: 1
diff --git a/test/snapshots/callback.multiple-assert/options.js b/test/snapshots/callback.multiple-assert/options.js
@@ -0,0 +1,4 @@
+module.exports = {
+    assertBeforeCallbackName: "beforeCallback",
+    assertAfterCallbackName: "afterCallback"
+};
diff --git a/test/snapshots/callback.multiple-assert/output.js b/test/snapshots/callback.multiple-assert/output.js
@@ -0,0 +1,18 @@
+beforeCallback("id:0");
+assert.strictEqual(1, 1);
+afterCallback("id:0");
+// => 1
+beforeCallback("id:1");
+assert.strictEqual("str", "str");
+afterCallback("id:1");
+// => "str"
+beforeCallback("id:2");
+assert.deepStrictEqual([1, 2, 3], [1, 2, 3]);
+afterCallback("id:2");
+// => [1,2,3]
+Promise.resolve(Promise.resolve(1)).then(v => {
+  beforeCallback("id:3");
+  assert.strictEqual(v, 1);
+  afterCallback("id:3");
+  return v;
+}); // => Resolve: 1
diff --git a/test/snapshots/console.log.Promise.resolve/output.js b/test/snapshots/console.log.Promise.resolve/output.js
@@ -1,4 +1,4 @@
 Promise.resolve(Promise.resolve(1)).then(v => {
   assert.strictEqual(v, 1);
   return v;
-}); // => Resolve: 1
+}); // => Resolve: 1