fix: ensure story wrapper is resolved correctly

mgadewoll · mgadewoll · commit b7ff818a914f · 2024-06-24T17:30:39.000+02:00
- uses Story wrapper for ComboBox story as example
diff --git a/packages/eui/.storybook/addons/code-snippet/decorators/render_jsx.tsx b/packages/eui/.storybook/addons/code-snippet/decorators/render_jsx.tsx
@@ -164,8 +164,6 @@ export const renderJsx = (
         : // @ts-expect-error (Converted from ts-ignore)
           reactElementToJSXString.default;
 
-    const shouldResolveChildren =
-      context?.parameters?.codeSnippet?.resolveChildren === true;
     const shouldResolveStoryElementOnly =
       context?.parameters?.codeSnippet?.resolveStoryElementOnly === true;
 
@@ -175,15 +173,17 @@ export const renderJsx = (
       // manual flag to remove an outer story wrapper and resolve its children instead
       // useful when complex custom stories are build where the actual story component is
       // not the outer component but part of a composition within another wrapper
-      if (shouldResolveChildren) {
-        node = getResolvedStoryChild(child, context);
-        // resolves the story element only and removes any wrapper or siblings
-      } else if (shouldResolveStoryElementOnly) {
+      if (shouldResolveStoryElementOnly) {
         const storyNode = getStoryComponent(child, context);
 
         if (storyNode) {
           node = storyNode;
         }
+        // manual flag to remove an outer story wrapper and resolve its children instead
+        // useful when complex custom stories are build where the actual story component is
+        // not the outer component but part of a composition within another wrapper
+      } else if (isStoryWrapper(child, context)) {
+        node = getResolvedStoryChild(child, context);
       } else {
         // removes outer wrapper components but leaves:
         // - stateful wrappers (kept and renamed later via displayName)
diff --git a/packages/eui/.storybook/addons/code-snippet/decorators/utils.ts b/packages/eui/.storybook/addons/code-snippet/decorators/utils.ts
@@ -15,6 +15,7 @@ import {
   ComponentType,
   ExoticComponent,
   isValidElement,
+  Component,
 } from 'react';
 import { Args, StoryContext } from '@storybook/react';
 import * as prettier from 'prettier';
@@ -56,8 +57,11 @@ export const getElementDisplayName = (
   node: ReactElement
 ): string | undefined => {
   let displayName;
+  const isClassComponent = node instanceof Component;
 
-  if (typeof node.type === 'function' || typeof node.type === 'object') {
+  if (isClassComponent) {
+    displayName = node.constructor.name;
+  } else if (typeof node.type === 'function' || typeof node.type === 'object') {
     const component = node.type as FunctionComponent;
     displayName = component.displayName ?? component.name ?? undefined;
   }
@@ -92,10 +96,10 @@ export const getEmotionComponentDisplayName = (
       : emotionTypeName ?? emotionLabelData?.displayName;
 
     // remove internal component underscore markings
-    return replacementName ? replacementName.replace('_', '') : displayName;
+    return replacementName ?? displayName;
   }
 
-  return displayName;
+  return displayName ? displayName.replace('_', '') : displayName;
 };
 
 export const getStoryComponentDisplayName = (
@@ -127,15 +131,18 @@ export const isStoryComponent = (
 ): boolean => {
   if (!context) return false;
 
-  const displayName = getEmotionComponentDisplayName(node)?.replace(/^_/, '');
+  const isClassComponent = node instanceof Component;
+  const displayName = isClassComponent
+    ? node.constructor?.name
+    : getEmotionComponentDisplayName(node)?.replace(/^_/, '');
   const isCurrentStory = displayName
     ? displayName === context?.component?.displayName
     : false;
 
   return isCurrentStory;
 };
 
-const isStoryWrapper = (node: ReactElement, context: StoryContext) => {
+export const isStoryWrapper = (node: ReactElement, context: StoryContext) => {
   const displayName = getEmotionComponentDisplayName(node);
   const isStoryWrapper =
     (typeof displayName === 'string' && displayName.startsWith('Story')) ||
@@ -210,9 +217,15 @@ export const getResolvedStoryChild = (
   child: ReactElement,
   context: StoryContext
 ) => {
-  return child.type && typeof child.type === 'function'
-    ? (child.type as (args: Args) => ReactElement)(context?.args)
-    : child;
+  if (!child.type) return child;
+  if (typeof child.type !== 'function') return child;
+
+  const isClassComponent = child.type.prototype instanceof Component;
+  const resolvedChild = isClassComponent
+    ? child
+    : (child.type as (args: Args) => ReactElement)(context?.args);
+
+  return resolvedChild;
 };
 
 /**
@@ -245,9 +258,15 @@ export const getStoryComponent = (
           // skip non-ReactElement children
           if (!isValidElement(child)) continue;
 
+          const displayName = getEmotionComponentDisplayName(child);
           // Story wrappers need to be resolved first to ensure the right data
           const resolvedChild = getResolvedStoryChild(child, context);
-          resolveChildren(resolvedChild);
+          const resolvedDisplayName =
+            getEmotionComponentDisplayName(resolvedChild);
+
+          if (resolvedDisplayName !== displayName) {
+            resolveChildren(resolvedChild);
+          }
         }
       } else if (
         // CASE: story wrapper; no children
@@ -260,15 +279,15 @@ export const getStoryComponent = (
         const resolvedDisplayName =
           getEmotionComponentDisplayName(resolvedChild);
 
-        if (resolvedDisplayName && resolvedDisplayName !== displayName) {
+        if (resolvedDisplayName !== displayName) {
           resolveChildren(resolvedChild);
         }
       } else if (
         // CASE: single child element
         childNode.props?.children &&
         !Array.isArray(childNode.props?.children)
       ) {
-        resolveChildren(childNode.props?.children);
+        storyNode = childNode.props.children;
       }
     }
   };
@@ -302,9 +321,11 @@ export const getDefaultPropsfromDocgenInfo = (
 
   if (!storyComponent) return undefined;
 
-  let propsInfo = isEmotionComponent(storyComponent)
-    ? storyComponent.props?.[EMOTION_TYPE_KEY]?.__docgenInfo.props
-    : storyComponent.type?.__docgenInfo?.props;
+  let propsInfo =
+    isEmotionComponent(storyComponent) &&
+    typeof storyComponent.props?.[EMOTION_TYPE_KEY] !== 'string'
+      ? storyComponent.props?.[EMOTION_TYPE_KEY]?.__docgenInfo.props
+      : storyComponent.type?.__docgenInfo?.props;
 
   const args = context.args;
 
diff --git a/packages/eui/src/components/combo_box/combo_box.stories.tsx b/packages/eui/src/components/combo_box/combo_box.stories.tsx
@@ -133,7 +133,12 @@ export const WithTooltip: Story = {
 hideStorybookControls(WithTooltip, ['onChange']);
 
 export const CustomMatcher: Story = {
-  render: (args) => <StatefulCustomMatcher {...args} />,
+  parameters: {
+    codeSnippet: {
+      resolveStoryElementOnly: true,
+    },
+  },
+  render: (args) => <StoryCustomMatcher {...args} />,
 };
 
 export const Groups: Story = {
@@ -237,7 +242,7 @@ const StatefulComboBox = ({
   );
 };
 
-const StatefulCustomMatcher = ({
+const StoryCustomMatcher = ({
   singleSelection,
   onChange,
   ...args
