Fix await interactions use cases

lib/rules/await-interactions.ts

@@ -7,7 +7,15 @@ import type { CallExpression, Identifier, Node } from '@typescript-eslint/types/
 
 import { createStorybookRule } from '../utils/create-storybook-rule'
 import { CategoryId } from '../utils/constants'
-import { isMemberExpression, isIdentifier, isAwaitExpression, isCallExpression } from '../utils/ast'
+import {
+  isMemberExpression,
+  isIdentifier,
+  isAwaitExpression,
+  isCallExpression,
+  isArrowFunctionExpression,
+  isReturnStatement,
+  isTSNonNullExpression,
+} from '../utils/ast'
 
 //------------------------------------------------------------------------------
 // Rule Definition
@@ -47,12 +55,19 @@ export = createStorybookRule({
       'waitForElement',
       'waitForDomChange',
       'userEvent',
+      'play',
     ]
 
     const getMethodThatShouldBeAwaited = (expr: CallExpression) => {
       const shouldAwait = (name: any) => {
         return FUNCTIONS_TO_BE_AWAITED.includes(name) || name.startsWith('findBy')
       }
+
+      // When an expression is a return value it doesn't need to be awaited
+      if (isArrowFunctionExpression(expr.parent) || isReturnStatement(expr.parent)) {
+        return null
+      }
+
       if (
         isMemberExpression(expr.callee) &&
         isIdentifier(expr.callee.object) &&
@@ -61,6 +76,15 @@ export = createStorybookRule({
         return expr.callee.object
       }
 
+      if (
+        isTSNonNullExpression(expr.callee) &&
+        isMemberExpression(expr.callee.expression) &&
+        isIdentifier(expr.callee.expression.property) &&
+        shouldAwait(expr.callee.expression.property.name)
+      ) {
+        return expr.callee.expression.property
+      }
+
       if (
         isMemberExpression(expr.callee) &&
         isIdentifier(expr.callee.property) &&

lib/utils/ast.ts

@@ -35,3 +35,4 @@ export const isProgram = isNodeOfType(AST_NODE_TYPES.Program)
 export const isTSTypeAliasDeclaration = isNodeOfType(AST_NODE_TYPES.TSTypeAliasDeclaration)
 export const isTSInterfaceDeclaration = isNodeOfType(AST_NODE_TYPES.TSInterfaceDeclaration)
 export const isTSAsExpression = isNodeOfType(AST_NODE_TYPES.TSAsExpression)
+export const isTSNonNullExpression = isNodeOfType(AST_NODE_TYPES.TSNonNullExpression)

tests/lib/rules/await-interactions.test.ts

@@ -48,6 +48,33 @@ ruleTester.run('await-interactions', rule, {
       }
     `,
     'await expect(foo).toBe(bar)',
+    dedent`
+      Basic.play = async () => {
+        await waitForElementToBeRemoved(() => canvas.findByText('Loading...'))
+        await waitForElementToBeRemoved(() => userEvent.hover(canvas.getByTestId('password-error-info')))
+      }
+    `,
+    // // @TODO: https://github.com/storybookjs/eslint-plugin-storybook/issues/47
+    // dedent`
+    //   import { userEvent } from '../utils'
+    //   Basic.play = async () => {
+    //     userEvent.click(canvas.getByRole('button'))
+    //   }
+    // `,
+    // // @TODO: https://github.com/storybookjs/eslint-plugin-storybook/issues/28
+    // dedent`
+    //   Block.parameters = {
+    //     async puppeteerTest(page) {
+    //       const element = await page.$('[data-test-block]');
+    //       await element.hover();
+    //       const textContent = await element.getProperty('textContent');
+    //       const text = await textContent.jsonValue();
+    //       expect(text).toBe('I am hovered');
+    //     },
+    //   };
+    // `,
+    'Basic.play = async () => userEvent.click(button)',
+    'Basic.play = async () => { return userEvent.click(button) }',
   ],
   invalid: [
     {
@@ -236,5 +263,33 @@ ruleTester.run('await-interactions', rule, {
         },
       ],
     },
+    {
+      code: dedent`
+        export const ThirdStory = {
+          play: async (context) => {
+            FirstStory.play(context)
+            SecondStory.play!(context)
+          }
+        }
+      `,
+      output: dedent`
+        export const ThirdStory = {
+          play: async (context) => {
+            await FirstStory.play(context)
+            await SecondStory.play!(context)
+          }
+        }
+      `,
+      errors: [
+        {
+          messageId: 'interactionShouldBeAwaited',
+          data: { method: 'play' },
+        },
+        {
+          messageId: 'interactionShouldBeAwaited',
+          data: { method: 'play' },
+        },
+      ],
+    },
   ],
 })