Fix exception for return andThen()

Author: simonihmig

__testfixtures__/acceptance/andthen.input.js

@@ -0,0 +1,27 @@
+import { test } from 'qunit';
+import moduleForAcceptance from '../../tests/helpers/module-for-acceptance';
+
+moduleForAcceptance('click');
+
+test('visiting /foo', function(assert) {
+  visit('/foo');
+  andThen(function() {
+  });
+
+  andThen(function() {
+    assert.ok(true);
+  });
+});
+
+test('visiting /foo', function(assert) {
+  visit('/foo');
+  return andThen(function() {
+  });
+});
+
+test('visiting /foo', function(assert) {
+  visit('/foo');
+  return andThen(function() {
+    assert.ok(true);
+  });
+});

__testfixtures__/acceptance/andthen.output.js

@@ -0,0 +1,19 @@
+import { visit } from '@ember/test-helpers';
+import { test } from 'qunit';
+import moduleForAcceptance from '../../tests/helpers/module-for-acceptance';
+
+moduleForAcceptance('click');
+
+test('visiting /foo', async function(assert) {
+  await visit('/foo');
+  assert.ok(true);
+});
+
+test('visiting /foo', async function(assert) {
+  await visit('/foo');
+});
+
+test('visiting /foo', async function(assert) {
+  await visit('/foo');
+  assert.ok(true);
+});

lib/utils.js

@@ -96,8 +96,6 @@ function isJQuerySelectExpression(j, node, path) {
   return false;
 }
 
-
-
 /**
  * Check if `node` is a `find(selector)` expression
  *
@@ -130,15 +128,15 @@ function migrateSelector(j, selector) {
 
     // Transform tail-eq selector to a find all
     if (tailEqRegex.test(string)) {
-      let [, query, eqIndex ] = string.match(/(.+?):eq\((\d+)\)$/);
+      let [, query, eqIndex] = string.match(/(.+?):eq\((\d+)\)$/);
       selector.value = query;
       if (eqIndex === '0') {
         // findAll('*')[0] === find('*')
         addImportStatement(['find']);
         return createFindExpression(j, [selector]);
       } else {
         addImportStatement(['findAll']);
-        return createArraySubscriptExpression(j, createFindAllExpression(j, [ selector ]), +eqIndex);
+        return createArraySubscriptExpression(j, createFindAllExpression(j, [selector]), +eqIndex);
       }
     }
   }
@@ -248,11 +246,11 @@ function createQuerySelectorAllExpression(j, args, fileRoot) {
 
 function isQuerySelectorAllCallExpression(j, node) {
   return j.CallExpression.check(node)
-  && j.MemberExpression.check(node.callee)
-  && j.MemberExpression.check(node.callee.object)
-  // && j.ThisExpression.check(node.callee.object.object)
-  && node.callee.object.property.name === 'element'
-  && node.callee.property.name === 'querySelectorAll';
+    && j.MemberExpression.check(node.callee)
+    && j.MemberExpression.check(node.callee.object)
+    // && j.ThisExpression.check(node.callee.object.object)
+    && node.callee.object.property.name === 'element'
+    && node.callee.property.name === 'querySelectorAll';
 }
 
 /**
@@ -332,7 +330,7 @@ function createBlurExpression(j, selector) {
   return j.awaitExpression(
     j.callExpression(
       j.identifier('blur'),
-     [migrateSelector(j, selector)]
+      [migrateSelector(j, selector)]
     )
   );
 }
@@ -475,13 +473,20 @@ function dropAndThen(j, root) {
     .find(j.CallExpression, { callee: { name: 'andThen' } })
     .map((path) => path.parent)
     .replaceWith(({ node }) => {
-      let body = node.expression.arguments[0].body;
-
-      if (j.CallExpression.check(body)) {
-        return j.expressionStatement(body);
+      let body;
+      if (j.ExpressionStatement.check(node)) {
+        body = node.expression.arguments[0].body;
+      } else if (j.ReturnStatement.check(node)) {
+        body = node.argument.arguments[0].body;
       }
 
-      return body.body;
+      if (body) {
+        if (j.CallExpression.check(body)) {
+          return j.expressionStatement(body);
+        }
+
+        return body.body;
+      }
     });
 
   if (replacements.length > 0) {
@@ -491,7 +496,7 @@ function dropAndThen(j, root) {
 
 function makeAwait(j, collection) {
   collection
-    .filter(({ parent: { node }}) => !j.AwaitExpression.check(node))
+    .filter(({ parent: { node } }) => !j.AwaitExpression.check(node))
     .replaceWith(({ node }) => j.awaitExpression(node))
     .forEach((path) => {