Web-components: Fix cleaning Lit Expression comments

app/web-components/src/client/docs/sourceDecorator.test.ts

@@ -1,13 +1,13 @@
 import { html } from 'lit-html';
 import { styleMap } from 'lit-html/directives/style-map';
 import { addons, useEffect } from '@storybook/addons';
-import type { StoryContext } from '@storybook/addons';
 import { SNIPPET_RENDERED } from '@storybook/docs-tools';
+import type { StoryContext, WebComponentsFramework } from '..';
 import { sourceDecorator } from './sourceDecorator';
 
 jest.mock('@storybook/addons');
 const mockedAddons = addons as jest.Mocked<typeof addons>;
-const mockedUseEffect = useEffect as jest.Mocked<typeof useEffect>;
+const mockedUseEffect = useEffect as jest.Mock;
 
 expect.addSnapshotSerializer({
   print: (val: any) => val,
@@ -16,16 +16,22 @@ expect.addSnapshotSerializer({
 
 const tick = () => new Promise((r) => setTimeout(r, 0));
 
-const makeContext = (name: string, parameters: any, args: any, extra?: object): StoryContext => ({
-  id: `lit-test--${name}`,
-  kind: 'js-text',
-  name,
-  parameters,
-  args,
-  argTypes: {},
-  globals: {},
-  ...extra,
-});
+const makeContext = (
+  name: string,
+  parameters: any,
+  args: any,
+  extra?: Partial<StoryContext<WebComponentsFramework>>
+) =>
+  ({
+    id: `lit-test--${name}`,
+    kind: 'js-text',
+    name,
+    parameters,
+    args,
+    argTypes: {},
+    globals: {},
+    ...extra,
+  } as StoryContext<WebComponentsFramework>);
 
 describe('sourceDecorator', () => {
   let mockChannel: { on: jest.Mock; emit?: jest.Mock };
@@ -106,4 +112,23 @@ describe('sourceDecorator', () => {
     sourceDecorator(storyFn, context);
     expect(transformSource).toHaveBeenCalledWith('<div>args story</div>', context);
   });
+
+  it('should clean lit expression comments', async () => {
+    const storyFn = (args: any) => html`<div>${args.slot}</div>`;
+    const context = makeContext(
+      'args',
+      { __isArgsStory: true },
+      { slot: 'some content' },
+      { originalStoryFn: storyFn }
+    );
+    // bind args to storyFn, as it's done in Storybook
+    const boundStoryFn = storyFn.bind(null, context.args);
+    sourceDecorator(boundStoryFn, context);
+    await tick();
+    expect(mockChannel.emit).toHaveBeenCalledWith(
+      SNIPPET_RENDERED,
+      'lit-test--args',
+      '<div>some content</div>'
+    );
+  });
 });

app/web-components/src/client/docs/sourceDecorator.ts

@@ -6,6 +6,9 @@ import { SNIPPET_RENDERED, SourceType } from '@storybook/docs-tools';
 
 import type { WebComponentsFramework } from '..';
 
+// Taken from https://github.com/lit/lit/blob/main/packages/lit-html/src/test/test-utils/strip-markers.ts
+const LIT_EXPRESSION_COMMENTS = /<!--\?lit\$[0-9]+\$-->|<!--\??-->/g;
+
 function skipSourceRender(context: StoryContext<WebComponentsFramework>) {
   const sourceParams = context?.parameters.docs?.source;
   const isArgsStory = context?.parameters.__isArgsStory;
@@ -44,7 +47,10 @@ export function sourceDecorator(
   if (!skipSourceRender(context)) {
     const container = window.document.createElement('div');
     render(story, container);
-    source = applyTransformSource(container.innerHTML.replace(/<!---->/g, ''), context);
+    source = applyTransformSource(
+      container.innerHTML.replace(LIT_EXPRESSION_COMMENTS, ''),
+      context
+    );
   }
 
   return story;