make context-in-play-function detect the context name

storybookjs · Feb 4, 2023 · 27d9f2b · 27d9f2b
1 parent d8adb39
commit 27d9f2b
Show file tree

Hide file tree

Showing 3 changed files with 83 additions and 2 deletions.
diff --git a/lib/rules/context-in-play-function.ts b/lib/rules/context-in-play-function.ts
@@ -13,6 +13,9 @@ import {
   isTSNonNullExpression,
   isObjectExpression,
   isSpreadElement,
+  isArrowFunctionExpression,
+  isObjectPattern,
+  isRestElement,
 } from '../utils/ast'
 
 //------------------------------------------------------------------------------
@@ -66,6 +69,36 @@ export = createStorybookRule({
       return false
     }
 
+    const getParentParameterName = (node: TSESTree.Node): string | undefined => {
+      if (!isArrowFunctionExpression(node)) {
+        if (!node.parent) {
+          return undefined
+        }
+        return getParentParameterName(node.parent)
+      }
+      // No parameter found
+      if (node.params.length === 0) {
+        return undefined
+      }
+
+      if (node.params.length >= 1) {
+        const param = node.params[0]
+        if (isIdentifier(param)) {
+          return param.name
+        }
+        if (isObjectPattern(param)) {
+          const restElement = param.properties.find(isRestElement)
+          if (!restElement || !isIdentifier(restElement.argument)) {
+            // No rest element found
+            return undefined
+          }
+          return restElement.argument.name
+        }
+      }
+      return undefined
+
+    }
+
     // Expression passing an argument called context OR spreading a variable called context
     const isNotPassingContextCorrectly = (expr: TSESTree.CallExpression) => {
       const firstExpressionArgument = expr.arguments[0]
@@ -74,10 +107,16 @@ export = createStorybookRule({
         return true
       }
 
+      const contextVariableName = getParentParameterName(expr)
+
+      if (!contextVariableName) {
+        return true
+      }
+
       if (
         expr.arguments.length === 1 &&
         isIdentifier(firstExpressionArgument) &&
-        firstExpressionArgument.name === 'context'
+        firstExpressionArgument.name === contextVariableName
       ) {
         return false
       }
@@ -86,7 +125,9 @@ export = createStorybookRule({
         isObjectExpression(firstExpressionArgument) &&
         firstExpressionArgument.properties.some((prop) => {
           return (
-            isSpreadElement(prop) && isIdentifier(prop.argument) && prop.argument.name === 'context'
+            isSpreadElement(prop) &&
+            isIdentifier(prop.argument) &&
+            prop.argument.name === contextVariableName
           )
         })
       ) {

diff --git a/lib/utils/ast.ts b/lib/utils/ast.ts
@@ -30,6 +30,7 @@ export const isObjectExpression = isNodeOfType(AST_NODE_TYPES.ObjectExpression)
 export const isObjectPattern = isNodeOfType(AST_NODE_TYPES.ObjectPattern)
 export const isProperty = isNodeOfType(AST_NODE_TYPES.Property)
 export const isSpreadElement = isNodeOfType(AST_NODE_TYPES.SpreadElement)
+export const isRestElement = isNodeOfType(AST_NODE_TYPES.RestElement)
 export const isReturnStatement = isNodeOfType(AST_NODE_TYPES.ReturnStatement)
 export const isFunctionDeclaration = isNodeOfType(AST_NODE_TYPES.FunctionDeclaration)
 export const isFunctionExpression = isNodeOfType(AST_NODE_TYPES.FunctionExpression)

diff --git a/tests/lib/rules/context-in-play-function.test.ts b/tests/lib/rules/context-in-play-function.test.ts
@@ -39,6 +39,45 @@ ruleTester.run('context-in-play-function', rule, {
         }
       }
     `,
+    dedent`
+      export const SecondStory = Template.bind({})
+      SecondStory.play = async (ctx) => {
+        await FirstStory.play(ctx)
+      }
+    `,
+    dedent`
+      export const SecondStory = Template.bind({})
+      SecondStory.play = async ({ canvasElement, ...ctx }) => {
+        await FirstStory.play({ canvasElement, ...ctx })
+      }
+    `,
+    dedent`
+      export const SecondStory = Template.bind({})
+      SecondStory.play = async (ctx) => {
+        await FirstStory.play(ctx)
+      }
+    `,
+    dedent`
+      export const SecondStory = {
+        play: async (ctx) => {
+          await FirstStory.play(ctx)
+        }
+      }
+    `,
+    dedent`
+      export const SecondStory = {
+        play: async ({ ...ctx }) => {
+          await FirstStory.play({ ...ctx })
+        }
+      }
+    `,
+    dedent`
+      export const SecondStory = {
+        play: async ({ canvasElement, ...ctx }) => {
+          await FirstStory.play({ canvasElement, ...ctx })
+        }
+      }
+    `,
   ],
   invalid: [
     {