From facc4f78139a40369ff46b9fe0233b13bd7a8756 Mon Sep 17 00:00:00 2001
From: Michael Shilman <michael@lab80.co>
Date: Thu, 5 May 2022 09:06:37 +0800
Subject: [PATCH] Fix default export postprocessing

---
 .vscode/settings.json |  3 ++
 src/mdx2.test.ts      | 65 +++++++++++++++++++++++++++++++++++++++++++
 src/mdx2.ts           |  6 ++--
 3 files changed, 71 insertions(+), 3 deletions(-)
 create mode 100644 .vscode/settings.json

diff --git a/.vscode/settings.json b/.vscode/settings.json
new file mode 100644
index 0000000..66d2c07
--- /dev/null
+++ b/.vscode/settings.json
@@ -0,0 +1,3 @@
+{
+  "deepscan.enable": true
+}
\ No newline at end of file
diff --git a/src/mdx2.test.ts b/src/mdx2.test.ts
index 7042765..c4f7706 100644
--- a/src/mdx2.test.ts
+++ b/src/mdx2.test.ts
@@ -46,6 +46,71 @@ describe('mdx2', () => {
       const mdxStoryNameToKey = { foo: 'foo' };
     `);
   });
+
+  it('full snapshot', () => {
+    const input = dedent`
+      # hello
+
+      <Meta title="foobar" />
+
+      world {2 + 1}
+
+      <Story name="foo">bar</Story>
+    `;
+    // @ts-ignore
+    expect(compileSync(input)).toMatchInlineSnapshot(`
+/*@jsxRuntime automatic @jsxImportSource react*/
+import { assertIsFn, AddContext } from "@storybook/addon-docs";
+import {Fragment as _Fragment, jsx as _jsx, jsxs as _jsxs} from "react/jsx-runtime";
+function MDXContent(props = {}) {
+  const {wrapper: MDXLayout} = props.components || ({});
+  return MDXLayout ? _jsx(MDXLayout, Object.assign({}, props, {
+    children: _jsx(_createMdxContent, {})
+  })) : _createMdxContent();
+  function _createMdxContent() {
+    const _components = Object.assign({
+      h1: "h1",
+      p: "p"
+    }, props.components), {Meta, Story} = _components;
+    if (!Meta) _missingMdxReference("Meta", true);
+    if (!Story) _missingMdxReference("Story", true);
+    return _jsxs(_Fragment, {
+      children: [_jsx(_components.h1, {
+        children: "hello"
+      }), "\\n", _jsx(Meta, {
+        title: "foobar"
+      }), "\\n", _jsxs(_components.p, {
+        children: ["world ", 2 + 1]
+      }), "\\n", _jsx(Story, {
+        name: "foo",
+        children: "bar"
+      })]
+    });
+  }
+}
+function _missingMdxReference(id, component) {
+  throw new Error("Expected " + (component ? "component" : "object") + " \`" + id + "\` to be defined: you likely forgot to import, pass, or provide it.");
+}
+// =========
+export const foo = () => (
+          "bar"
+        );
+foo.storyName = 'foo';
+foo.parameters = { storySource: { source: '\\"bar\\"' } };
+
+const componentMeta = { title: 'foobar', includeStories: ["foo"],  };
+
+const mdxStoryNameToKey = {"foo":"foo"};
+
+componentMeta.parameters = componentMeta.parameters || {};
+componentMeta.parameters.docs = {
+  ...(componentMeta.parameters.docs || {}),
+  page: () => <AddContext mdxStoryNameToKey={mdxStoryNameToKey} mdxComponentAnnotations={componentMeta}><MDXContent /></AddContext>,
+};
+
+export default componentMeta;
+`);
+  });
 });
 
 describe('docs-mdx-compiler-plugin', () => {
diff --git a/src/mdx2.ts b/src/mdx2.ts
index 8c13a16..a3330a7 100644
--- a/src/mdx2.ts
+++ b/src/mdx2.ts
@@ -115,14 +115,14 @@ export const plugin = (store: any) => (root: any) => {
 
 export const postprocess = (code: string, extractedExports: string) => {
   const lines = code.toString().trim().split('\n');
-  const last = lines.pop();
+
+  // /*@jsxRuntime automatic @jsxImportSource react*/
   const first = lines.shift();
 
   return [
     first,
     'import { assertIsFn, AddContext } from "@storybook/addon-docs";',
-    ...lines,
-       last,
+    ...lines.filter((line) => !line.match(/^export default/)),
     SEPARATOR,
     extractedExports,
   ].join('\n');