MorphDoc/apps/server/tests/pdf-engine.test.ts

import { describe, expect, it, vi } from "vitest";
import type { Browser } from "playwright";
import { defaultExportConfig } from "@md-to-pdf/core";
import {
  ConcurrencyGate,
  PdfEngineClosedError,
  PdfEngineOverloadedError,
  PdfRenderTimeoutError,
  PlaywrightPdfGenerator,
  isAllowedPdfRequestUrl,
  type PdfRenderPayload
} from "../src/pdf-engine.js";

function payload(): PdfRenderPayload {
  return {
    articleHtml: '<article id="write"><h1>测试</h1></article>',
    fileName: "测试.md",
    metadata: {
      title: "测试",
      author: "",
      subject: "",
      keywords: [],
      language: "zh-CN"
    },
    features: [],
    themeCss: "#write { color: #333; }",
    exportConfig: defaultExportConfig
  };
}

function fakeBrowser(options: { hangOnGoto?: boolean } = {}) {
  const route = vi.fn();
  const closeContext = vi.fn(async () => undefined);
  const closeBrowser = vi.fn(async () => undefined);
  const pdfBuffer = Buffer.from("%PDF-test");
  const page = {
    setDefaultTimeout: vi.fn(),
    setDefaultNavigationTimeout: vi.fn(),
    emulateMedia: vi.fn(async () => undefined),
    goto: options.hangOnGoto
      ? vi.fn(() => new Promise(() => undefined))
      : vi.fn(async () => undefined),
    waitForFunction: vi.fn(async () => undefined),
    evaluate: vi.fn(async () => ({
      pageCount: 2,
      contentHeight: 2200,
      mermaidErrors: [],
      timings: {
        setupMs: 1,
        mermaidMs: 2,
        mermaidFitMs: 3,
        mermaidConversionMs: 3.5,
        resourceWaitMs: 4,
        paginationMs: 5,
        finalizeMs: 6,
        totalMs: 21
      }
    })),
    pdf: vi.fn(async () => pdfBuffer)
  };
  const context = {
    route,
    newPage: vi.fn(async () => page),
    close: closeContext
  };
  const browser = {
    newContext: vi.fn(async () => context),
    close: closeBrowser,
    on: vi.fn()
  } as unknown as Browser;

  return {
    browser,
    closeBrowser,
    closeContext,
    context,
    page,
    pdfBuffer,
    route
  };
}

describe("PDF 并发控制", () => {
  it("限制并发并拒绝超过排队上限的任务", async () => {
    const gate = new ConcurrencyGate(1, 1);
    let releaseFirst: (() => void) | undefined;
    const first = gate.run(
      () =>
        new Promise<void>((resolve) => {
          releaseFirst = resolve;
        })
    );
    const second = gate.run(async () => undefined);

    await expect(
      gate.run(async () => undefined)
    ).rejects.toBeInstanceOf(PdfEngineOverloadedError);

    releaseFirst?.();
    await first;
    await second;
  });
});

describe("PDF 网络策略", () => {
  it("只允许渲染同源和内嵌资源", () => {
    const origin = "http://127.0.0.1:5173";
    expect(
      isAllowedPdfRequestUrl(
        "http://127.0.0.1:5173/api/themes/a.css",
        origin
      )
    ).toBe(true);
    expect(isAllowedPdfRequestUrl("data:image/png;base64,AA", origin))
      .toBe(true);
    expect(isAllowedPdfRequestUrl("blob:http://127.0.0.1/id", origin))
      .toBe(true);
    expect(
      isAllowedPdfRequestUrl("https://example.com/font.woff2", origin)
    ).toBe(false);
  });
});

describe("Playwright PDF 生成器", () => {
  it("复用浏览器、隔离上下文并返回 PDF", async () => {
    const fake = fakeBrowser();
    const generator = new PlaywrightPdfGenerator({
      launchBrowser: async () => fake.browser,
      timeoutMs: 1_000
    });

    const first = await generator.generate(payload());
    const second = await generator.generate(payload());

    expect(first.pdf.toString()).toBe("%PDF-test");
    expect(first.pdf).toBe(fake.pdfBuffer);
    expect(first.pageCount).toBe(2);
    expect(first.timings).toEqual(
      expect.objectContaining({
        queueMs: expect.any(Number),
        browserMs: expect.any(Number),
        contextMs: expect.any(Number),
        navigationMs: expect.any(Number),
        documentRenderMs: expect.any(Number),
        runtimeSetupMs: 1,
        mermaidMs: 2,
        mermaidFitMs: 3,
        mermaidConversionMs: 3.5,
        resourceWaitMs: 4,
        paginationMs: 5,
        finalizeMs: 6,
        runtimeTotalMs: 21,
        pdfPrintMs: expect.any(Number),
        totalMs: expect.any(Number)
      })
    );
    expect(second.pdf.toString()).toBe("%PDF-test");
    expect(fake.browser.newContext).toHaveBeenCalledTimes(2);
    expect(fake.closeContext).toHaveBeenCalledTimes(2);
    expect(fake.page.pdf).toHaveBeenCalledWith(
      expect.objectContaining({
        pageRanges: "1-2",
        preferCSSPageSize: true,
        printBackground: true,
        scale: 1
      })
    );

    await generator.close();
    expect(fake.closeBrowser).toHaveBeenCalledTimes(1);
    await expect(generator.generate(payload())).rejects.toBeInstanceOf(
      PdfEngineClosedError
    );
  });

  it("可以在首个生成请求前预热浏览器", async () => {
    const fake = fakeBrowser();
    const launchBrowser = vi.fn(async () => fake.browser);
    const generator = new PlaywrightPdfGenerator({
      launchBrowser,
      timeoutMs: 1_000
    });

    await generator.warmup();
    await generator.warmup();
    expect(launchBrowser).toHaveBeenCalledTimes(1);
    expect(fake.browser.newContext).not.toHaveBeenCalled();

    await generator.generate(payload());
    expect(launchBrowser).toHaveBeenCalledTimes(1);
    await generator.close();
  });

  it("超时后关闭独立上下文", async () => {
    const fake = fakeBrowser({ hangOnGoto: true });
    const generator = new PlaywrightPdfGenerator({
      launchBrowser: async () => fake.browser,
      timeoutMs: 10
    });

    await expect(generator.generate(payload())).rejects.toBeInstanceOf(
      PdfRenderTimeoutError
    );
    expect(fake.closeContext).toHaveBeenCalledTimes(1);
    await generator.close();
  });
});