feat(babel-plugin-espower): call assertionVisitor with appropriate as…

…ync/generator flags

.gitignore

@@ -1,3 +1,4 @@
 /.idea
 /node_modules
 /npm-debug.log
+/test/fixtures/*/actual.js

lib/babel-estree-path.js

@@ -21,10 +21,24 @@ module.exports = function (babel) {
         }
     }
 
-    return function estreePath (nodePath) {
+    function findEnclosingFunction(nodePath) {
+        if (!nodePath) {
+            return null;
+        }
+        if (babel.types.isFunction(nodePath.node)) {
+            return nodePath.node;
+        }
+        return findEnclosingFunction(nodePath.parentPath);
+    }
+
+    function estreePath (nodePath) {
         var ancestors = [];
         traverseUp(ancestors, nodePath);
         ancestors.reverse();
         return ancestors;
-    };
+    }
+
+    estreePath.findEnclosingFunction = findEnclosingFunction;
+
+    return estreePath;
 };

lib/create-espower-visitor.js

@@ -30,7 +30,8 @@ function enterNodePath (nodePath, currentNode, parentNode, scope, file, matchers
                 sourceMap: file.opts.inputSourceMap
             }, options);
             treePath = esTreePath(nodePath);
-            assertionVisitor = new espower.AssertionVisitor(matcher, treePath, espowerOptions);
+            var enclosingFunction = esTreePath.findEnclosingFunction(nodePath);
+            assertionVisitor = new espower.AssertionVisitor(matcher, treePath, enclosingFunction, espowerOptions);
             assertionVisitor.enter(currentNode, parentNode);
             file.set('espowerAssertionVisitor', assertionVisitor);
         }

package.json

@@ -11,7 +11,7 @@
   "dependencies": {
     "array-find": "^1.0.0",
     "escallmatch": "^1.4.2",
-    "espower": "^1.0.6",
+    "espower": "^1.2.1",
     "xtend": "^4.0.0"
   },
   "devDependencies": {

test/fixtures/AwaitExpression/expected.js

@@ -0,0 +1,40 @@
+'use strict';
+
+function myAsync(a) {
+  return regeneratorRuntime.async(function myAsync$(context$1$0) {
+    while (1) switch (context$1$0.prev = context$1$0.next) {
+      case 0:
+        context$1$0.t0 = assert;
+        context$1$0.t1 = assert;
+        context$1$0.t2 = assert;
+        context$1$0.next = 5;
+        return regeneratorRuntime.awrap(a);
+
+      case 5:
+        context$1$0.t3 = context$1$0.sent;
+        context$1$0.t4 = context$1$0.t2._capt.call(context$1$0.t2, context$1$0.t3, 'arguments/0/left');
+        context$1$0.t5 = context$1$0.t4 === 3;
+        context$1$0.t6 = context$1$0.t1._capt.call(context$1$0.t1, context$1$0.t5, 'arguments/0');
+        context$1$0.t7 = {
+          async: true,
+          content: 'assert((await a) === 3)',
+          filepath: 'test/fixtures/AwaitExpression/fixture.js',
+          line: 4
+        };
+        context$1$0.t8 = context$1$0.t0._expr.call(context$1$0.t0, context$1$0.t6, context$1$0.t7);
+        assert(context$1$0.t8);
+
+      case 12:
+      case 'end':
+        return context$1$0.stop();
+    }
+  }, null, this);
+}
+
+function notAsync(a) {
+  assert(assert._expr(assert._capt(assert._capt(await(assert._capt(a, 'arguments/0/left/arguments/0')), 'arguments/0/left') === 3, 'arguments/0'), {
+    content: 'assert(await(a) === 3)',
+    filepath: 'test/fixtures/AwaitExpression/fixture.js',
+    line: 8
+  }));
+}

test/fixtures/AwaitExpression/fixture.js

@@ -0,0 +1,9 @@
+'use strict';
+
+async function myAsync(a){
+  assert((await (a)) === 3);
+}
+
+function notAsync(a){
+  assert((await (a)) === 3);
+}

test/fixtures/YieldExpression/expected.js

@@ -0,0 +1,41 @@
+'use strict';
+
+var marked0$0 = [gen].map(regeneratorRuntime.mark);
+function gen(a) {
+  return regeneratorRuntime.wrap(function gen$(context$1$0) {
+    while (1) switch (context$1$0.prev = context$1$0.next) {
+      case 0:
+        context$1$0.t0 = assert;
+        context$1$0.t1 = assert;
+        context$1$0.t2 = assert;
+        context$1$0.next = 5;
+        return a;
+
+      case 5:
+        context$1$0.t3 = context$1$0.sent;
+        context$1$0.t4 = context$1$0.t2._capt.call(context$1$0.t2, context$1$0.t3, 'arguments/0/left');
+        context$1$0.t5 = context$1$0.t4 === 3;
+        context$1$0.t6 = context$1$0.t1._capt.call(context$1$0.t1, context$1$0.t5, 'arguments/0');
+        context$1$0.t7 = {
+          content: 'assert((yield a) === 3)',
+          filepath: 'test/fixtures/YieldExpression/fixture.js',
+          generator: true,
+          line: 4
+        };
+        context$1$0.t8 = context$1$0.t0._expr.call(context$1$0.t0, context$1$0.t6, context$1$0.t7);
+        assert(context$1$0.t8);
+
+      case 12:
+      case 'end':
+        return context$1$0.stop();
+    }
+  }, marked0$0[0], this);
+}
+
+function notGen(a) {
+  assert(assert._expr(assert._capt(assert._capt(yield(assert._capt(a, 'arguments/0/left/arguments/0')), 'arguments/0/left') === 3, 'arguments/0'), {
+    content: 'assert(yield(a) === 3)',
+    filepath: 'test/fixtures/YieldExpression/fixture.js',
+    line: 8
+  }));
+}

test/fixtures/YieldExpression/fixture.js

@@ -0,0 +1,9 @@
+'use strict';
+
+function *gen(a){
+  assert((yield (a)) === 3);
+}
+
+function notGen(a){
+  assert((yield (a)) === 3);
+}

test/test.js

@@ -9,12 +9,16 @@ function testTransform (fixtureName, extraOptions) {
     it(fixtureName, function () {
         var fixtureFilepath = path.resolve(__dirname, 'fixtures', fixtureName, 'fixture.js');
         var expectedFilepath = path.resolve(__dirname, 'fixtures', fixtureName, 'expected.js');
+        var actualFilepath = path.resolve(__dirname, 'fixtures', fixtureName, 'actual.js');
         var result = babel.transformFileSync(fixtureFilepath, extend({
             plugins: ['../index']
         }, extraOptions));
-        var actual = result.code;
-        var expected = fs.readFileSync(expectedFilepath).toString();
-        assert.equal(actual + '\n', expected);
+        var actual = result.code + '\n';
+        var expected = fs.readFileSync(expectedFilepath, 'utf8');
+        if (actual != expected) {
+            fs.writeFileSync(actualFilepath, actual);
+        }
+        assert.equal(actual, expected);
     });
 }
 
@@ -40,6 +44,8 @@ describe('babel-plugin-espower', function () {
     testTransform('ClassExpression');
     testTransform('SpreadElement');
     testTransform('Property');
+    testTransform('YieldExpression');
+    testTransform('AwaitExpression');
     testTransform('inputSourceMap', {
         plugins: [
             createEspowerPlugin(babel, {