superwork/vscode/webview-ui/src/codeReader/useMermaidDiagrams.test.tsx

// @vitest-environment happy-dom

import type { Root } from 'react-dom/client'
import { act } from 'react'
import { createRoot } from 'react-dom/client'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import { renderMarkdown } from './markdown'

/** 假的 mermaid：记下加载了几次、按哪个主题初始化、画了什么；源码里有 `坏` 就当解析失败。 */
const fake = vi.hoisted(() => ({
  imports: 0,
  theme: '',
  running: 0,
  mostRunning: 0,
  /** 每张图画之前等它，测并发时卡住 */
  gate: Promise.resolve(),
  initialize: vi.fn((config: { theme: string }) => {
    fake.theme = config.theme
  }),
  render: vi.fn(async (id: string, text: string) => {
    fake.running += 1
    fake.mostRunning = Math.max(fake.mostRunning, fake.running)
    await fake.gate
    fake.running -= 1
    if (text.includes('坏'))
      throw new Error('Parse error on line 2:\n坏\n^\nExpecting \'NODE\', got \'EOF\'')
    return { svg: `<svg id="${id}" data-theme="${fake.theme}"><text>${text.trim()}</text></svg>` }
  }),
}))

/** 加载一次 mermaid 记一次 */
function mermaidModule(): { default: Pick<typeof fake, 'initialize' | 'render'> } {
  fake.imports += 1
  return { default: { initialize: fake.initialize, render: fake.render } }
}

let root: Root
let page: HTMLDivElement

/** 挂上只用这个 hook 的组件；每个测试重新加载模块、重新登记假的 mermaid，mermaid 有没有被加载各算各的 */
async function mount(): Promise<void> {
  vi.resetModules()
  vi.doMock('mermaid', mermaidModule)
  const { useMermaidDiagrams } = await import('./useMermaidDiagrams')
  function Diagrams(): null {
    useMermaidDiagrams()
    return null
  }
  const host = document.createElement('div')
  document.body.append(host)
  root = createRoot(host)
  await act(async () => root.render(<Diagrams />))
}

/** 页面上加一段注释的 Markdown，返回它的容器 */
function addMarkdown(markdown: string): HTMLDivElement {
  const div = document.createElement('div')
  div.className = 'cr-markdown'
  div.innerHTML = renderMarkdown(markdown)
  page.append(div)
  return div
}

const flow = (text: string): string => `\`\`\`mermaid\nflowchart TD\n    ${text}\n\`\`\``

beforeEach(() => {
  vi.stubGlobal('IS_REACT_ACT_ENVIRONMENT', true)
  document.body.className = 'vscode-dark'
  page = document.createElement('div')
  document.body.append(page)
  fake.gate = Promise.resolve()
  fake.running = 0
  fake.mostRunning = 0
  fake.render.mockClear()
  fake.initialize.mockClear()
})

afterEach(async () => {
  await act(async () => root.unmount())
  document.body.replaceChildren()
  vi.unstubAllGlobals()
})

describe('useMermaidDiagrams', () => {
  it('页面上没有 mermaid 占位就不加载 mermaid', async () => {
    const before = fake.imports
    await mount()
    addMarkdown('```ts\nload()\n```\n\n说明')
    await new Promise(resolve => setTimeout(resolve, 150))
    expect(fake.imports).toBe(before)

    addMarkdown(flow('A["开始"] --> B["结束"]'))
    await vi.waitFor(() => expect(page.querySelector('.cr-mermaid svg')).not.toBeNull())
    expect(fake.imports).toBe(before + 1)
  })

  it('只画占位：代码块换成图，别的代码块不动；初始化不跑页面脚本、深色主题', async () => {
    await mount()
    const other = addMarkdown('```ts\nload()\n```')
    const diagram = addMarkdown(flow('A["读取会话"] --> B{"是本人的？"}'))
    await vi.waitFor(() => expect(diagram.querySelector('.cr-mermaid svg')).not.toBeNull())
    expect(diagram.querySelector('.cr-mermaid pre')).toBeNull()
    expect(diagram.querySelector('svg')!.dataset.theme).toBe('dark')
    expect(other.innerHTML).toBe('<pre><code class="language-ts">load()\n</code></pre>\n')
    expect(fake.render).toHaveBeenCalledTimes(1)
    expect(fake.render.mock.calls[0][1]).toBe('flowchart TD\n    A["读取会话"] --> B{"是本人的？"}\n')
    expect(fake.initialize).toHaveBeenCalledWith(expect.objectContaining({ startOnLoad: false, securityLevel: 'strict', theme: 'dark' }))
  })

  it('同样的源码、同样的主题只画一次：两处一起画，卡片重新挂上直接用画好的', async () => {
    await mount()
    const first = addMarkdown(flow('A --> B'))
    const second = addMarkdown(flow('A --> B'))
    await vi.waitFor(() => expect(second.querySelector('.cr-mermaid svg')).not.toBeNull())
    expect(first.querySelector('.cr-mermaid svg')).not.toBeNull()
    expect(fake.render).toHaveBeenCalledTimes(1)

    first.remove()
    second.remove()
    const again = addMarkdown(flow('A --> B'))
    await vi.waitFor(() => expect(again.querySelector('.cr-mermaid svg')).not.toBeNull())
    expect(fake.render).toHaveBeenCalledTimes(1)
  })

  it('主题变了按新主题重画；换回来用画过的', async () => {
    await mount()
    const diagram = addMarkdown(flow('A --> B'))
    await vi.waitFor(() => expect(diagram.querySelector('svg')?.dataset.theme).toBe('dark'))

    document.body.className = 'vscode-light'
    await vi.waitFor(() => expect(diagram.querySelector('svg')?.dataset.theme).toBe('default'))
    expect(fake.render).toHaveBeenCalledTimes(2)

    document.body.className = 'vscode-high-contrast'
    await vi.waitFor(() => expect(diagram.querySelector('svg')?.dataset.theme).toBe('dark'))
    expect(fake.render).toHaveBeenCalledTimes(2)
  })

  it('画不出来：留着代码块，下面写一行简短原因，悬停看全文', async () => {
    await mount()
    const diagram = addMarkdown(flow('坏 -->'))
    await vi.waitFor(() => expect(diagram.querySelector('.cr-mermaid-error')).not.toBeNull())
    expect(diagram.querySelector('.cr-mermaid pre code')!.textContent).toBe('flowchart TD\n    坏 -->\n')
    expect(diagram.querySelector('svg')).toBeNull()
    const error = diagram.querySelector<HTMLElement>('.cr-mermaid-error')!
    expect(error.textContent).toBe('图画不出来：Parse error on line 2:')
    expect(error.dataset.tooltip).toContain('Expecting \'NODE\', got \'EOF\'')
  })

  it('一次只画一张', async () => {
    let open = (): void => {}
    fake.gate = new Promise(resolve => open = resolve)
    await mount()
    const diagrams = [addMarkdown(flow('A --> B')), addMarkdown(flow('C --> D')), addMarkdown(flow('E --> F'))]
    await vi.waitFor(() => expect(fake.render).toHaveBeenCalledTimes(1))
    await new Promise(resolve => setTimeout(resolve, 100))
    expect(fake.render).toHaveBeenCalledTimes(1)
    open()
    await vi.waitFor(() => expect(diagrams.every(diagram => diagram.querySelector('svg'))).toBe(true))
    expect(fake.render).toHaveBeenCalledTimes(3)
    expect(fake.mostRunning).toBe(1)
  })
})