From 2486f691738eff99c0ef17c0a0791c15e6d5aa34 Mon Sep 17 00:00:00 2001
From: Kelly Joseph Price <kellyjosephprice@gmail.com>
Date: Tue, 25 Feb 2025 12:58:00 -0800
Subject: [PATCH] fix: toc

---
 __tests__/browser/markdown.test.js         |  1 +
 __tests__/custom-components/index.test.tsx | 54 ++++++++++++++--------
 example/components.ts                      |  2 +-
 lib/run.tsx                                |  5 +-
 processor/plugin/toc.ts                    |  6 +--
 5 files changed, 40 insertions(+), 28 deletions(-)

diff --git a/__tests__/browser/markdown.test.js b/__tests__/browser/markdown.test.js
index 5ee938bc4..b350b3a1f 100644
--- a/__tests__/browser/markdown.test.js
+++ b/__tests__/browser/markdown.test.js
@@ -13,6 +13,7 @@ describe('visual regression tests', () => {
     const docs = [
       'callouts',
       'calloutTests',
+      'childTests',
       'codeBlocks',
       // skipping this because they sporadically failure with network timing
       // issues
diff --git a/__tests__/custom-components/index.test.tsx b/__tests__/custom-components/index.test.tsx
index 743158e49..d28fa83b5 100644
--- a/__tests__/custom-components/index.test.tsx
+++ b/__tests__/custom-components/index.test.tsx
@@ -3,32 +3,30 @@ import type { RMDXModule } from '../../types';
 import { render, screen } from '@testing-library/react';
 import React from 'react';
 
-import { execute } from '../helpers';
 import { compile, run } from '../../lib';
-import { RMDXModule } from '../../types';
+import { execute } from '../helpers';
 
 describe('Custom Components', () => {
-  let Example;
-  let Multiple;
-
-  beforeEach(async () => {
-    Example = await execute('It works!', {}, {}, { getDefault: false });
-    Multiple = await execute(
-      `
+  const Example = 'It works!';
+  const Multiple = `
 export const First = () => <div>First</div>;
 export const Second = () => <div>Second</div>;
-  `,
-      {},
-      {},
-      { getDefault: false },
-    );
-  });
+  `;
+  const Nesting = `
+export const WithChildren = ({ children }) => <div>{children}</div>;
+
+<div>{props.children}</div>
+`;
 
   it('renders custom components', async () => {
     const doc = `
 <Example />
     `;
-    const Page = (await execute(doc, undefined, { components: { Example } })) as RMDXModule['default'];
+    const Page = (await execute(
+      doc,
+      { components: { Example } },
+      { components: { Example } },
+    )) as RMDXModule['default'];
     render(<Page />);
 
     expect(screen.getByText('It works!')).toBeVisible();
@@ -40,17 +38,33 @@ export const Second = () => <div>Second</div>;
 
 <Second />
     `;
-    const Page = (await execute(doc, undefined, { components: { Multiple } })) as RMDXModule['default'];
+    const Page = (await execute(
+      doc,
+      { components: { Multiple } },
+      { components: { Multiple } },
+    )) as RMDXModule['default'];
     render(<Page />);
 
     expect(screen.getByText('First')).toBeVisible();
     expect(screen.getByText('Second')).toBeVisible();
   });
 
+  it('renders a nested exported custom component', async () => {
+    const doc = '<Nesting><WithChildren>Hello, Test User!</WithChildren></Nesting>';
+    const Page = (await execute(
+      doc,
+      { components: { Nesting } },
+      { components: { Nesting } },
+    )) as RMDXModule['default'];
+    render(<Page />);
+
+    expect(screen.getByText('Hello, Test User!')).toBeVisible();
+  });
+
   it('renders the default export of a custom component and passes through props', async () => {
-    const Test = (await run(await compile(`{props.attr}`))) as RMDXModule;
-    const doc = `<Test attr="Hello" />`;
-    const Page = await run(await compile(doc), { components: { Test } });
+    const Test = '{props.attr}';
+    const doc = '<Test attr="Hello" />';
+    const Page = await run(await compile(doc, { components: { Test } }), { components: { Test } });
     render(<Page.default />);
 
     expect(screen.getByText('Hello')).toBeVisible();
diff --git a/example/components.ts b/example/components.ts
index 8b9b3ec80..547619e5f 100644
--- a/example/components.ts
+++ b/example/components.ts
@@ -14,7 +14,7 @@ export const Test = ({ color = 'thistle' } = {}) => {
 export default Test;
   `,
   MultipleExports: `
-export const One = () => <div>"One"</div>;
+export const One = () => "One";
 
 export const Two = () => "Two";
   `,
diff --git a/lib/run.tsx b/lib/run.tsx
index 51cce06b1..e33d8d334 100644
--- a/lib/run.tsx
+++ b/lib/run.tsx
@@ -55,11 +55,10 @@ const run = async (string: string, _opts: RunOpts = {}) => {
   const { Fragment } = runtime;
   const { components = {}, terms, variables, baseUrl, imports = {}, ...opts } = _opts;
 
-  const tocsByTag = {};
+  const tocsByTag: Record<string, RMDXModule['toc']> = {};
   const promises: [string, RMDXModule][] = await Promise.all(
     Object.entries(components).map(async ([tag, body]) => {
       const code = await compile(body);
-      console.log(code);
       const mod = await run(code);
       return [tag, mod];
     }),
@@ -81,8 +80,6 @@ const run = async (string: string, _opts: RunOpts = {}) => {
     return memo;
   }, {});
 
-  console.log(string);
-
   const exec = (text: string, { useMDXComponents = makeUseMDXComponents(exportedComponents) }: RunOpts = {}) => {
     return mdxRun(text, {
       ...runtime,
diff --git a/processor/plugin/toc.ts b/processor/plugin/toc.ts
index d3f607dac..3eb85d0b4 100644
--- a/processor/plugin/toc.ts
+++ b/processor/plugin/toc.ts
@@ -1,4 +1,4 @@
-import type { CustomComponents, HastHeading, IndexableElements, TocList, TocListItem } from '../../types';
+import type { CustomComponents, HastHeading, IndexableElements, RMDXModule, TocList, TocListItem } from '../../types';
 import type { Root } from 'hast';
 import type { MdxjsEsm } from 'mdast-util-mdxjs-esm';
 import type { Transformer } from 'unified';
@@ -89,11 +89,11 @@ const tocToHast = (headings: HastHeading[] = []): TocList => {
   return ast;
 };
 
-export const tocHastToMdx = (toc: IndexableElements[], components: CustomComponents) => {
+export const tocHastToMdx = (toc: IndexableElements[], components: Record<string, RMDXModule['toc']>) => {
   const tree: Root = { type: 'root', children: toc };
 
   visit(tree, 'mdxJsxFlowElement', (node, index, parent) => {
-    const subToc = components[node.name].toc || [];
+    const subToc = components[node.name] || [];
     parent.children.splice(index, 1, ...subToc);
   });