From 7eb6f35af40da41c145d05dbd193a720ccd883a6 Mon Sep 17 00:00:00 2001
From: Mikhail Bashkirov <bashmish@gmail.com>
Date: Thu, 11 May 2023 08:00:45 +0200
Subject: [PATCH] fix WC source decorator to handle document fragment correctly

---
 .../src/docs/sourceDecorator.test.ts          | 32 ++++++++++++++++++-
 .../src/docs/sourceDecorator.ts               |  6 +++-
 code/renderers/web-components/src/types.ts    |  7 +++-
 3 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/code/renderers/web-components/src/docs/sourceDecorator.test.ts b/code/renderers/web-components/src/docs/sourceDecorator.test.ts
index 57990d4c193d..d09a4136c7f5 100644
--- 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
index c5ccff5de92e..3a6f45c04001 100644
--- 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
index 33dfc412d751..32cbce64c690 100644
--- 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>;