fix WC source decorator to handle document fragment correctly

storybookjs · May 11, 2023 · 7eb6f35 · 7eb6f35
1 parent 9029693
commit 7eb6f35
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 3 deletions.
diff --git a/code/renderers/web-components/src/docs/sourceDecorator.test.ts b/code/renderers/web-components/src/docs/sourceDecorator.test.ts
@@ -1,4 +1,4 @@
-import { html } from 'lit';
+import { html, render } from 'lit';
 import { styleMap } from 'lit/directives/style-map.js';
 import { addons, useEffect } from '@storybook/preview-api';
 import { SNIPPET_RENDERED } from '@storybook/docs-tools';
@@ -85,4 +85,34 @@ describe('sourceDecorator', () => {
       source: '<div>args story</div>',
     });
   });
+
+  it('should handle document fragment without removing its child nodes', async () => {
+    const storyFn = () =>
+      html`my
+        <div>args story</div>`;
+    const decoratedStoryFn = () => {
+      const fragment = document.createDocumentFragment();
+      render(storyFn(), fragment);
+      return fragment;
+    };
+    const context = makeContext('args', { __isArgsStory: true }, {});
+    const story = sourceDecorator(decoratedStoryFn, context);
+    await tick();
+    expect(mockChannel.emit).toHaveBeenCalledWith(SNIPPET_RENDERED, {
+      id: 'lit-test--args',
+      args: {},
+      source: `my
+        <div>args story</div>`,
+    });
+    expect(story).toMatchInlineSnapshot(`
+      <DocumentFragment>
+        <!---->
+        my
+              
+        <div>
+          args story
+        </div>
+      </DocumentFragment>
+    `);
+  });
 });
diff --git a/code/renderers/web-components/src/docs/sourceDecorator.ts b/code/renderers/web-components/src/docs/sourceDecorator.ts
@@ -40,7 +40,11 @@ export function sourceDecorator(
   });
   if (!skipSourceRender(context)) {
     const container = window.document.createElement('div');
-    render(renderedForSource, container);
+    if (renderedForSource instanceof DocumentFragment) {
+      render(renderedForSource.cloneNode(true), container);
+    } else {
+      render(renderedForSource, container);
+    }
     source = container.innerHTML.replace(LIT_EXPRESSION_COMMENTS, '');
   }
 

diff --git a/code/renderers/web-components/src/types.ts b/code/renderers/web-components/src/types.ts
@@ -1,7 +1,12 @@
 import type { StoryContext as StoryContextBase, WebRenderer } from '@storybook/types';
 import type { TemplateResult, SVGTemplateResult } from 'lit';
 
-export type StoryFnHtmlReturnType = string | Node | TemplateResult | SVGTemplateResult;
+export type StoryFnHtmlReturnType =
+  | string
+  | Node
+  | DocumentFragment
+  | TemplateResult
+  | SVGTemplateResult;
 
 export type StoryContext = StoryContextBase<WebComponentsRenderer>;