Make sure that await interaction works with @storybook/test

Author: kasperpeulen

lib/rules/await-interactions.ts

@@ -132,7 +132,8 @@ export = createStorybookRule({
 
     const isUserEventFromStorybookImported = (node: TSESTree.ImportDeclaration) => {
       return (
-        node.source.value === '@storybook/testing-library' &&
+        (node.source.value === '@storybook/testing-library' ||
+          node.source.value === '@storybook/test') &&
         node.specifiers.find(
           (spec) =>
             isImportSpecifier(spec) &&
@@ -144,7 +145,7 @@ export = createStorybookRule({
 
     const isExpectFromStorybookImported = (node: TSESTree.ImportDeclaration) => {
       return (
-        node.source.value === '@storybook/jest' &&
+        (node.source.value === '@storybook/jest' || node.source.value === '@storybook/test') &&
         node.specifiers.find(
           (spec) => isImportSpecifier(spec) && spec.imported.name === 'expect'
         ) !== undefined

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

@@ -146,6 +146,61 @@ ruleTester.run('await-interactions', rule, {
         },
       ],
     },
+    {
+      code: dedent`
+        import { expect, findByText } from '@storybook/test'
+        WithModalOpen.play = async ({ args }) => {
+          // should complain
+          expect(args.onClick).toHaveBeenCalled()
+          const element = findByText(canvasElement, 'asdf')
+        }
+      `,
+      output: dedent`
+        import { expect, findByText } from '@storybook/test'
+        WithModalOpen.play = async ({ args }) => {
+          // should complain
+          await expect(args.onClick).toHaveBeenCalled()
+          const element = await findByText(canvasElement, 'asdf')
+        }
+      `,
+      only: true,
+      errors: [
+        {
+          messageId: 'interactionShouldBeAwaited',
+          data: { method: 'toHaveBeenCalled' },
+          suggestions: [
+            {
+              messageId: 'fixSuggestion',
+              output: dedent`
+                import { expect, findByText } from '@storybook/test'
+                WithModalOpen.play = async ({ args }) => {
+                  // should complain
+                  await expect(args.onClick).toHaveBeenCalled()
+                  const element = findByText(canvasElement, 'asdf')
+                }
+              `,
+            },
+          ],
+        },
+        {
+          messageId: 'interactionShouldBeAwaited',
+          data: { method: 'findByText' },
+          suggestions: [
+            {
+              messageId: 'fixSuggestion',
+              output: dedent`
+                import { expect, findByText } from '@storybook/test'
+                WithModalOpen.play = async ({ args }) => {
+                  // should complain
+                  expect(args.onClick).toHaveBeenCalled()
+                  const element = await findByText(canvasElement, 'asdf')
+                }
+              `,
+            },
+          ],
+        },
+      ],
+    },
     {
       code: dedent`
         WithModalOpen.play = ({ canvasElement }) => {