Support waitFor()-type modifiers in async arrow transform

lib/babel-plugin-transform-ember-concurrency-async-tasks.js

@@ -78,68 +78,108 @@ const TransformAsyncMethodsIntoGeneratorMethods = {
       }
 
       // Thus far, we've established that value is `myTask = task(...)`.
-      // Now we need to check if the last argument is an async ArrowFunctionExpress
+      // Now we need to check if the last argument is an async ArrowFunctionExpression,
+      // possibly wrapped in other modifier functions such as `waitFor()`
 
-      const maybeAsyncArrowPath = path.get(
+      // If there are modifier functions applied, this will capture the
+      // top-level one
+      let rootModifierPath;
+
+      let maybeAsyncArrowPath = path.get(
         `value.arguments.${value.arguments.length - 1}`
       );
-      if (!maybeAsyncArrowPath && !maybeAsyncArrowPath.node) {
-        return;
-      }
-      const maybeAsyncArrow = maybeAsyncArrowPath.node;
-      if (
-        maybeAsyncArrow &&
-        maybeAsyncArrow.type === 'ArrowFunctionExpression' &&
-        maybeAsyncArrow.async
-      ) {
-        convertFunctionExpressionIntoGenerator(
-          maybeAsyncArrowPath,
-          state,
-          factoryFunctionName
-        );
+      while (maybeAsyncArrowPath && maybeAsyncArrowPath.node) {
+        const maybeAsyncArrow = maybeAsyncArrowPath.node;
+
+        if (
+          maybeAsyncArrow.type === 'ArrowFunctionExpression' &&
+          maybeAsyncArrow.async
+        ) {
+          // It's an async arrow function, so convert it
+          convertFunctionExpressionIntoGenerator(
+            maybeAsyncArrowPath,
+            rootModifierPath,
+            state,
+            factoryFunctionName
+          );
+          break;
+        } else if (maybeAsyncArrow.type === 'CallExpression') {
+          // It's a call expression, so save it as the modifier functions root
+          // if we don't already have one and then traverse into it
+          rootModifierPath = rootModifierPath || maybeAsyncArrowPath;
+          maybeAsyncArrowPath = maybeAsyncArrowPath.get('arguments.0');
+        } else {
+          break;
+        }
       }
     }
   },
 };
 
 function convertFunctionExpressionIntoGenerator(
-  path,
+  taskFnPath,
+  rootModifierPath,
   state,
   factoryFunctionName
 ) {
-  if (path && path.node.async) {
-    if (isArrowFunctionExpression(path)) {
+  if (taskFnPath && taskFnPath.node.async) {
+    if (isArrowFunctionExpression(taskFnPath)) {
       // At this point we have something that looks like
       //
       //    foo = task(this?, {}?, async () => {})
       //
+      // or (if there are modifier functions applied)
+      //
+      //    foo = task(this?, {}?, modifier1(modifier2(async () => {})))
+      //
       // and we need to convert it to
       //
       //    foo = buildTask(contextFn, options | null, taskName, bufferPolicyName?)
       //
       // where conextFn is
       //
       //    () => ({ context: this, generator: function * () { ... } })
+      //
+      // or (if there are modifier functions applied)
+      //
+      //    () => ({ context: this, generator: modifier1(modifier2(function * () { ... } })))
+
+      // Before we start moving things around, let's save off the task()
+      // CallExpression path
+      const taskPath = (rootModifierPath || taskFnPath).parentPath;
 
-      // Replace the async arrow fn with a generator fn
-      let asyncArrowFnBody = path.node.body;
+      // Transform the async arrow task function into a generator function
+      // (we'll do the actual transformation of `await`s into `yield`s below)
+      let asyncArrowFnBody = taskFnPath.node.body;
       if (asyncArrowFnBody.type !== 'BlockStatement') {
         // Need to convert `async () => expr` with `async () => { return expr }`
         asyncArrowFnBody = blockStatement([returnStatement(asyncArrowFnBody)]);
       }
 
       const taskGeneratorFn = functionExpression(
-        path.node.id,
-        path.node.params,
+        taskFnPath.node.id,
+        taskFnPath.node.params,
         asyncArrowFnBody,
         true
       );
 
+      // Replace the async arrow task function with the generator function
+      // in-place in the tree (and update `taskFnPath` to point to the new,
+      // generator, task function)
+      taskFnPath = taskFnPath.replaceWith(taskGeneratorFn)[0];
+
       const contextFn = arrowFunctionExpression(
         [],
         objectExpression([
           objectProperty(identifier('context'), thisExpression()),
-          objectProperty(identifier('generator'), taskGeneratorFn),
+          objectProperty(
+            identifier('generator'),
+            // We've swapped out the task fn for a generator function, possibly
+            // inside some modifier functions. Now we want to move that whole
+            // tree, including any modifier functions, into this generator
+            // property.
+            (rootModifierPath || taskFnPath).node
+          ),
         ])
       );
 
@@ -152,15 +192,15 @@ function convertFunctionExpressionIntoGenerator(
         );
       }
 
-      const originalArgs = path.parentPath.node.arguments;
+      const originalArgs = taskPath.node.arguments;
 
       // task(this, async() => {}) was the original API, but we don't actually
       // need the `this` arg (we determine the `this` context from the contextFn async arrow fn)
       if (originalArgs[0] && originalArgs[0].type === 'ThisExpression') {
         originalArgs.shift();
       }
 
-      const taskName = extractTaskNameFromClassProperty(path);
+      const taskName = extractTaskNameFromClassProperty(taskPath);
       let optionsOrNull;
 
       // remaining args should either be [options, async () => {}] or [async () => {}]
@@ -192,7 +232,7 @@ function convertFunctionExpressionIntoGenerator(
         ]
       );
 
-      let newPath = path.parentPath.replaceWith(buildTaskCall)[0];
+      let newPath = taskPath.replaceWith(buildTaskCall)[0];
       newPath.traverse({
         FunctionExpression(path) {
           if (!path.node.generator) {
@@ -224,11 +264,11 @@ const TransformAwaitIntoYield = {
  * in this method we extract the name from the ClassProperty assignment so that we can pass it in
  * to the options hash when constructing the Task.
  *
- * @param {babel.NodePath<babel.types.ArrowFunctionExpression>} asyncArrowFnPath
+ * @param {babel.NodePath<babel.types.CallExpression>} taskPath
  * @returns {string | null}
  */
-function extractTaskNameFromClassProperty(asyncArrowFnPath) {
-  const maybeClassPropertyPath = asyncArrowFnPath.parentPath.parentPath;
+function extractTaskNameFromClassProperty(taskPath) {
+  const maybeClassPropertyPath = taskPath.parentPath;
   if (
     maybeClassPropertyPath &&
     maybeClassPropertyPath.node.type === 'ClassProperty'

tests/integration/async-arrow-task-test.js

@@ -2,7 +2,14 @@ import { module, test } from 'qunit';
 import { setupRenderingTest } from 'ember-qunit';
 // eslint-disable-next-line ember/no-computed-properties-in-native-classes
 import { computed, set } from '@ember/object';
-import { click, render, settled } from '@ember/test-helpers';
+import {
+  click,
+  getSettledState,
+  render,
+  settled,
+  waitUntil,
+} from '@ember/test-helpers';
+import { waitFor } from '@ember/test-waiters';
 import { hbs } from 'ember-cli-htmlbars';
 import {
   task,
@@ -142,6 +149,76 @@ module('Integration | async-arrow-task', function (hooks) {
     await finishTest(assert);
   });
 
+  test('modifiers', async function (assert) {
+    let modifier1Called = false;
+    let modifier2Called = false;
+
+    function modifier1(fn) {
+      return function* (...args) {
+        modifier1Called = true;
+        yield* fn(args);
+      };
+    }
+    function modifier2(fn) {
+      return function* (...args) {
+        modifier2Called = true;
+        yield* fn(args);
+      };
+    }
+
+    this.owner.register(
+      'component:test-async-arrow-task',
+      class extends TestComponent {
+        myTask = task(
+          this,
+          modifier1(
+            modifier2(async (arg) => {
+              return arg;
+            })
+          )
+        );
+      }
+    );
+
+    await render(hbs`<TestAsyncArrowTask />`);
+    await click('button#start');
+    assert.true(modifier1Called);
+    assert.true(modifier2Called);
+  });
+
+  test('waitFor modifier', async function (assert) {
+    assert.expect(9);
+
+    let { promise, resolve } = defer();
+
+    this.owner.register(
+      'component:test-async-arrow-task',
+      class extends TestComponent {
+        myTask = task(
+          this,
+          waitFor(async (arg) => {
+            set(this, 'resolved', await promise);
+            assert.strictEqual(this.myTask.name, 'myTask');
+            return arg;
+          })
+        );
+      }
+    );
+
+    await render(hbs`<TestAsyncArrowTask />`);
+    click('button#start');
+    await waitUntil(() => this.element.textContent.includes('Running!'));
+
+    assert.true(getSettledState().hasPendingWaiters);
+
+    resolve('Wow!');
+    await settled();
+
+    assert.false(getSettledState().hasPendingWaiters);
+
+    await finishTest(assert);
+  });
+
   test('dropTask and other shorthand tasks (with `this` arg)', async function (assert) {
     assert.expect(13);