NianToB/tests/unit/knowledge-routes.test.ts

import { mkdtemp, readFile, rm, stat, writeFile } from 'node:fs/promises';
import crypto from 'node:crypto';
import { join } from 'node:path';
import { tmpdir } from 'node:os';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';

let userDataDir = '';

vi.mock('@electron/utils/paths', () => ({
  getDataDir: () => userDataDir,
}));

describe('knowledge routes', () => {
  beforeEach(async () => {
    vi.resetModules();
    userDataDir = await mkdtemp(join(tmpdir(), 'zhinian-knowledge-test-'));
  });

  afterEach(async () => {
    await rm(userDataDir, { recursive: true, force: true });
  });

  it('copies supported text files into the app knowledge directory', async () => {
    const sourceDir = await mkdtemp(join(tmpdir(), 'zhinian-knowledge-source-'));
    const sourcePath = join(sourceDir, 'faq.md');
    await writeFile(sourcePath, '# FAQ\n\nhello', 'utf8');

    const { importKnowledgeFiles } = await import('@electron/api/routes/knowledge');
    const result = await importKnowledgeFiles({
      workspaceId: 'workspace/test',
      filePaths: [sourcePath],
    });

    expect(result.rejected).toEqual([]);
    expect(result.documents).toHaveLength(1);
    expect(result.documents[0].name).toBe('faq.md');
    expect(result.documents[0].mimeType).toBe('text/markdown');
    expect(result.documents[0].storedPath).not.toBe(sourcePath);
    await expect(stat(result.documents[0].storedPath)).resolves.toBeTruthy();
    await expect(readFile(result.documents[0].storedPath, 'utf8')).resolves.toContain('hello');

    const registryPath = join(userDataDir, 'yinian', 'knowledge', 'workspace_test', 'registry.json');
    await expect(readFile(registryPath, 'utf8')).resolves.toContain('faq.md');

    await rm(sourceDir, { recursive: true, force: true });
  });

  it('builds chat context from the local backup after the source file is deleted', async () => {
    const sourceDir = await mkdtemp(join(tmpdir(), 'zhinian-knowledge-source-'));
    const sourcePath = join(sourceDir, 'policy.md');
    await writeFile(sourcePath, '# Policy\n\nRefund within 7 days.', 'utf8');

    const { importKnowledgeFiles, buildKnowledgeContext } = await import('@electron/api/routes/knowledge');
    const imported = await importKnowledgeFiles({
      workspaceId: 'workspace-context',
      filePaths: [sourcePath],
    });
    await rm(sourceDir, { recursive: true, force: true });

    const context = await buildKnowledgeContext({
      workspaceId: 'workspace-context',
      documentIds: [imported.documents[0].id],
    });

    expect(context.missing).toEqual([]);
    expect(context.documents[0].name).toBe('policy.md');
    expect(context.context).toContain('[知识库上下文]');
    expect(context.context).toContain('Refund within 7 days.');
  });

  it('builds chat context from extracted docx text', async () => {
    const sourceDir = await mkdtemp(join(tmpdir(), 'zhinian-knowledge-source-'));
    const sourcePath = join(sourceDir, 'handbook.docx');
    await writeMinimalDocx(sourcePath, 'Use the local backup for answers.');

    const { importKnowledgeFiles, buildKnowledgeContext } = await import('@electron/api/routes/knowledge');
    const imported = await importKnowledgeFiles({
      workspaceId: 'workspace-docx-context',
      filePaths: [sourcePath],
    });

    const context = await buildKnowledgeContext({
      workspaceId: 'workspace-docx-context',
      documentIds: [imported.documents[0].id],
    });

    expect(context.context).toContain('handbook.docx');
    expect(context.context).toContain('Use the local backup for answers.');

    await rm(sourceDir, { recursive: true, force: true });
  });

  it('rejects unsupported binary-looking files', async () => {
    const sourceDir = await mkdtemp(join(tmpdir(), 'zhinian-knowledge-source-'));
    const sourcePath = join(sourceDir, 'image.png');
    await writeFile(sourcePath, 'not really an image', 'utf8');

    const { importKnowledgeFiles } = await import('@electron/api/routes/knowledge');
    const result = await importKnowledgeFiles({
      workspaceId: 'workspace-test',
      filePaths: [sourcePath],
    });

    expect(result.documents).toEqual([]);
    expect(result.rejected).toEqual([{ filePath: sourcePath, reason: '仅支持文本类知识文件' }]);

    await rm(sourceDir, { recursive: true, force: true });
  });

  it('rejects missing files without changing the registry', async () => {
    const missingPath = join(tmpdir(), `missing-${crypto.randomUUID()}.md`);

    const { importKnowledgeFiles } = await import('@electron/api/routes/knowledge');
    const result = await importKnowledgeFiles({
      workspaceId: 'workspace-missing',
      filePaths: [missingPath],
    });

    expect(result.documents).toEqual([]);
    expect(result.rejected).toEqual([{ filePath: missingPath, reason: '文件不存在或不可读取' }]);

    const registryPath = join(userDataDir, 'yinian', 'knowledge', 'workspace-missing', 'registry.json');
    await expect(readFile(registryPath, 'utf8')).rejects.toThrow();
  });

  it('keeps newer imports at the top of the local registry', async () => {
    const sourceDir = await mkdtemp(join(tmpdir(), 'zhinian-knowledge-source-'));
    const oldPath = join(sourceDir, 'old.md');
    const newPath = join(sourceDir, 'new.md');
    await writeFile(oldPath, 'old knowledge', 'utf8');
    await writeFile(newPath, 'new knowledge', 'utf8');

    const { importKnowledgeFiles } = await import('@electron/api/routes/knowledge');
    await importKnowledgeFiles({ workspaceId: 'workspace-order', filePaths: [oldPath] });
    await importKnowledgeFiles({ workspaceId: 'workspace-order', filePaths: [newPath] });

    const registryPath = join(userDataDir, 'yinian', 'knowledge', 'workspace-order', 'registry.json');
    const registry = JSON.parse(await readFile(registryPath, 'utf8')) as Array<{ name: string }>;
    expect(registry.map((doc) => doc.name)).toEqual(['new.md', 'old.md']);

    await rm(sourceDir, { recursive: true, force: true });
  });

  it('copies docx files and extracts searchable text', async () => {
    const sourceDir = await mkdtemp(join(tmpdir(), 'zhinian-knowledge-source-'));
    const sourcePath = join(sourceDir, 'handbook.docx');
    await writeMinimalDocx(sourcePath, 'Welcome to Zhinian Handbook');

    const { importKnowledgeFiles } = await import('@electron/api/routes/knowledge');
    const result = await importKnowledgeFiles({
      workspaceId: 'workspace-docx',
      filePaths: [sourcePath],
    });

    expect(result.rejected).toEqual([]);
    expect(result.documents).toHaveLength(1);
    expect(result.documents[0].name).toBe('handbook.docx');
    expect(result.documents[0].mimeType).toBe('application/vnd.openxmlformats-officedocument.wordprocessingml.document');
    expect(result.documents[0].storedPath).not.toBe(sourcePath);
    expect(result.documents[0].textPath).toBeTruthy();
    await expect(readFile(result.documents[0].textPath!, 'utf8')).resolves.toContain('Welcome to Zhinian Handbook');

    await rm(sourceDir, { recursive: true, force: true });
  });

  it('rejects invalid docx files before adding them to the registry', async () => {
    const sourceDir = await mkdtemp(join(tmpdir(), 'zhinian-knowledge-source-'));
    const sourcePath = join(sourceDir, 'broken.docx');
    await writeFile(sourcePath, 'not a zip document', 'utf8');

    const { importKnowledgeFiles } = await import('@electron/api/routes/knowledge');
    const result = await importKnowledgeFiles({
      workspaceId: 'workspace-docx-broken',
      filePaths: [sourcePath],
    });

    expect(result.documents).toEqual([]);
    expect(result.rejected).toEqual([{ filePath: sourcePath, reason: 'Word 文档解析失败，请确认文件为 .docx 格式' }]);

    const registryPath = join(userDataDir, 'yinian', 'knowledge', 'workspace-docx-broken', 'registry.json');
    await expect(readFile(registryPath, 'utf8')).rejects.toThrow();

    await rm(sourceDir, { recursive: true, force: true });
  });
});

async function writeMinimalDocx(filePath: string, text: string) {
  const JSZip = (await import('jszip')).default;
  const zip = new JSZip();
  zip.file('[Content_Types].xml', `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
  <Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>
  <Default Extension="xml" ContentType="application/xml"/>
  <Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/>
</Types>`);
  zip.folder('_rels')?.file('.rels', `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
  <Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/>
</Relationships>`);
  zip.folder('word')?.file('document.xml', `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
  <w:body>
    <w:p><w:r><w:t>${escapeXml(text)}</w:t></w:r></w:p>
  </w:body>
</w:document>`);
  const buffer = await zip.generateAsync({ type: 'nodebuffer' });
  await writeFile(filePath, buffer);
}

function escapeXml(value: string): string {
  return value
    .replace(/&/g, '&amp;')
    .replace(/</g, '&lt;')
    .replace(/>/g, '&gt;')
    .replace(/"/g, '&quot;')
    .replace(/'/g, '&apos;');
}