Fix andThen with arrow function

Fixes #22

Author: simonihmig

__testfixtures__/acceptance/nested-in-and-then.input.js

@@ -26,3 +26,9 @@ test('visiting /twiddles', function(assert) {
 
   andThen(() => assert.ok(true));
 });
+
+test('visiting /twiddles', function(assert) {
+  andThen(() => {
+    click('.foo');
+  });
+});

__testfixtures__/acceptance/nested-in-and-then.output.js

@@ -14,4 +14,8 @@ test('visiting /twiddles', async function(assert) {
   assert.ok(true);
 
   assert.ok(true);
-});
+});
+
+test('visiting /twiddles', async function(assert) {
+  await click('.foo');
+});

lib/utils.js

@@ -428,7 +428,10 @@ function createImportStatement(j, source, imported, local) {
  */
 function findParentFunction(j, path) {
   while (path = path.parent) {
-    if (j.FunctionExpression.check(path.node) && !isAndThenCall(j, path) || j.FunctionDeclaration.check(path.node) || j.ArrowFunctionExpression.check(path.node)) {
+    if (
+      j.FunctionExpression.check(path.node) && !isAndThenCall(j, path)
+      || j.FunctionDeclaration.check(path.node)
+      || j.ArrowFunctionExpression.check(path.node) && !isAndThenCall(j, path)) {
       return path;
     }
   }
@@ -442,7 +445,8 @@ function findParentFunction(j, path) {
  */
 function isAndThenCall(j, path) {
   let parentNode = path.parent.node;
-  return j.FunctionExpression.check(path.node)
+  return (
+    j.FunctionExpression.check(path.node) || j.ArrowFunctionExpression.check(path.node))
     && j.CallExpression.check(parentNode)
     && j.Identifier.check(parentNode.callee)
     && parentNode.callee.name === 'andThen';