diff --git a/lib/rules/use-storybook-expect.ts b/lib/rules/use-storybook-expect.ts
index 9415854..59dc553 100644
--- a/lib/rules/use-storybook-expect.ts
+++ b/lib/rules/use-storybook-expect.ts
@@ -3,14 +3,13 @@
  * @author Yann Braga
  */
 
-import { isPlayFunction } from '../utils'
-
 import { CategoryId } from '../utils/constants'
 import {
   isExpressionStatement,
   isCallExpression,
   isMemberExpression,
   isIdentifier,
+  isBlockStatement,
 } from '../utils/ast'
 
 import { createStorybookRule } from '../utils/create-storybook-rule'
@@ -67,6 +66,20 @@ export = createStorybookRule({
         node.specifiers.find((spec: any) => spec.imported.name === 'expect')
       )
     }
+
+    const checkExpectInvocations = (blockStatement) => {
+      if (!isBlockStatement(blockStatement)) {
+        return
+      }
+
+      const expressionBody = blockStatement.body || []
+      const expressionStatements = getExpressionStatements(expressionBody)
+      expressionStatements.forEach(({ expression }) => {
+        if (isExpect(expression)) {
+          expectInvocations.push(expression)
+        }
+      })
+    }
     //----------------------------------------------------------------------
     // Public
     //----------------------------------------------------------------------
@@ -80,20 +93,13 @@ export = createStorybookRule({
           isImportingFromStorybookExpect = true
         }
       },
-      AssignmentExpression(node: any) {
-        if (!isExpressionStatement(node.parent)) {
+      CallExpression(node) {
+        if (!isIdentifier(node.callee)) {
           return null
         }
 
-        if (isPlayFunction(node)) {
-          const { right } = node
-          const expressionBody = (right.body && right.body.body) || []
-          const expressionStatements = getExpressionStatements(expressionBody)
-          expressionStatements.forEach(({ expression }) => {
-            if (isExpect(expression)) {
-              expectInvocations.push(expression)
-            }
-          })
+        if (node.callee.name === 'expect') {
+          expectInvocations.push(node.callee)
         }
       },
       'Program:exit': function () {
diff --git a/tests/lib/rules/use-storybook-expect.ts b/tests/lib/rules/use-storybook-expect.ts
index eb7cd2e..fdff5df 100644
--- a/tests/lib/rules/use-storybook-expect.ts
+++ b/tests/lib/rules/use-storybook-expect.ts
@@ -18,41 +18,123 @@ import { AST_NODE_TYPES } from '@typescript-eslint/types'
 
 ruleTester.run('use-storybook-expect', rule, {
   valid: [
-    dedent(`
+    dedent`
       import { expect } from '@storybook/jest';
 
       Default.play = () => {
         expect(123).toEqual(123);
       }
-    `),
+    `,
+    dedent`
+      import { expect } from '@storybook/jest';
+
+      export const Basic = {
+        ...Default,
+        play: async (context) => {
+          expect(123).toEqual(123);
+        },
+      };
+    `,
   ],
 
   invalid: [
     {
-      code: dedent(`
+      code: dedent`
         Default.play = () => {
           expect(123).toEqual(123);
         }
-      `),
-      output: dedent(`
+      `,
+      output: dedent`
         import { expect } from '@storybook/jest';
         Default.play = () => {
           expect(123).toEqual(123);
         }
-      `),
+      `,
       errors: [
         {
           messageId: 'useExpectFromStorybook',
-          type: AST_NODE_TYPES.CallExpression,
+          type: AST_NODE_TYPES.Identifier,
           suggestions: [
             {
               messageId: 'updateImports',
-              output: dedent(`
+              output: dedent`
                 import { expect } from '@storybook/jest';
                 Default.play = () => {
                   expect(123).toEqual(123);
                 }
-              `),
+              `,
+            },
+          ],
+        },
+      ],
+    },
+    {
+      code: dedent`
+        const someInteraction = () => {
+          expect(123).toEqual(123);
+        }
+        Default.play = someInteraction
+      `,
+      output: dedent`
+        import { expect } from '@storybook/jest';
+        const someInteraction = () => {
+          expect(123).toEqual(123);
+        }
+        Default.play = someInteraction
+      `,
+      errors: [
+        {
+          messageId: 'useExpectFromStorybook',
+          type: AST_NODE_TYPES.Identifier,
+          suggestions: [
+            {
+              messageId: 'updateImports',
+              output: dedent`
+                import { expect } from '@storybook/jest';
+                const someInteraction = () => {
+                  expect(123).toEqual(123);
+                }
+                Default.play = someInteraction
+              `,
+            },
+          ],
+        },
+      ],
+    },
+    {
+      code: dedent`
+        export const Basic = {
+          ...Default,
+          play: async (context) => {
+            expect(123).toEqual(123);
+          },
+        };
+      `,
+      output: dedent`
+        import { expect } from '@storybook/jest';
+        export const Basic = {
+          ...Default,
+          play: async (context) => {
+            expect(123).toEqual(123);
+          },
+        };
+      `,
+      errors: [
+        {
+          messageId: 'useExpectFromStorybook',
+          type: AST_NODE_TYPES.Identifier,
+          suggestions: [
+            {
+              messageId: 'updateImports',
+              output: dedent`
+                import { expect } from '@storybook/jest';
+                export const Basic = {
+                  ...Default,
+                  play: async (context) => {
+                    expect(123).toEqual(123);
+                  },
+                };
+              `,
             },
           ],
         },