package kb
// ========================================================================
// 多格式解析器单元测试
// ========================================================================
// 测试样本全部在内存里构造(docx 手工拼 zip、pdf 手工拼对象表、
// xlsx 用 excelize 生成),不依赖外部测试文件,CI 环境可直接跑
// ========================================================================
import (
"archive/zip"
"bytes"
"fmt"
"strings"
"testing"
"unicode/utf8"
"github.com/xuri/excelize/v2"
)
// buildMiniDocx 在内存构造一个最小可用的 docx(zip + word/document.xml)
func buildMiniDocx(t *testing.T, documentXML string) []byte {
t.Helper()
var buf bytes.Buffer
zw := zip.NewWriter(&buf)
w, err := zw.Create("word/document.xml")
if err != nil {
t.Fatalf("创建 zip 条目失败: %v", err)
}
if _, err := w.Write([]byte(documentXML)); err != nil {
t.Fatalf("写入 document.xml 失败: %v", err)
}
if err := zw.Close(); err != nil {
t.Fatalf("关闭 zip 失败: %v", err)
}
return buf.Bytes()
}
// TestExtractDocxText 验证 docx 抽取:段落分隔、换行、制表符、表格单元格
func TestExtractDocxText(t *testing.T) {
xml := `
感冒的中医辨证
风寒感冒:恶寒重发热轻无汗头痛
药名麻黄
`
data := buildMiniDocx(t, xml)
text, err := extractDocxText(data)
if err != nil {
t.Fatalf("extractDocxText 报错: %v", err)
}
for _, want := range []string{"感冒的中医辨证", "风寒感冒:\n恶寒重发热轻\t无汗头痛", "药名", "麻黄"} {
if !strings.Contains(text, want) {
t.Errorf("抽取文本缺少 %q,实际:\n%s", want, text)
}
}
// 段落之间必须有空行(ChunkPlainText 依赖双换行切段)
if !strings.Contains(text, "感冒的中医辨证\n\n") {
t.Errorf("段落后缺少空行分隔,实际:\n%s", text)
}
}
// TestExtractDocxText_InvalidZip 非 zip 内容要报友好错误而不是 panic
func TestExtractDocxText_InvalidZip(t *testing.T) {
if _, err := extractDocxText([]byte("这不是一个zip文件")); err == nil {
t.Fatal("非法 docx 应当报错")
}
}
// TestExtractHTMLText 验证 HTML 抽取:跳过 script/style、标题转 #、列表、表格
func TestExtractHTMLText(t *testing.T) {
htmlSrc := `
页面标题
中医基础理论
阴阳五行学说是中医的理论基础。
四诊
`
text, err := extractHTMLText([]byte(htmlSrc))
if err != nil {
t.Fatalf("extractHTMLText 报错: %v", err)
}
if strings.Contains(text, "alert") || strings.Contains(text, "color:red") || strings.Contains(text, "导航栏不要") {
t.Errorf("script/style/nav 内容未被过滤,实际:\n%s", text)
}
if !strings.Contains(text, "# 中医基础理论") || !strings.Contains(text, "## 四诊") {
t.Errorf("标题未转成 markdown 井号,实际:\n%s", text)
}
if !strings.Contains(text, "- 望诊") {
t.Errorf("列表项缺少 - 前缀,实际:\n%s", text)
}
// 走 ChunkMarkdown 后标题应进入 chunk.title
chunks := ChunkMarkdown(text, DefaultChunkOptions())
foundTitle := false
for _, c := range chunks {
if c.Title == "中医基础理论" || c.Title == "四诊" {
foundTitle = true
}
}
if !foundTitle {
t.Errorf("ChunkMarkdown 未把 HTML 标题切进 chunk.title,chunks=%+v", chunks)
}
}
// buildMiniPDF 在内存构造一个最小合法 PDF(单页 + Helvetica + 一行文本)
//
// 手工维护对象偏移量表(xref),这是 PDF 规范要求的最小骨架
func buildMiniPDF(text string) []byte {
content := fmt.Sprintf("BT /F1 12 Tf 72 720 Td (%s) Tj ET", text)
objs := []string{
"<< /Type /Catalog /Pages 2 0 R >>",
"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Contents 4 0 R /Resources << /Font << /F1 5 0 R >> >> >>",
fmt.Sprintf("<< /Length %d >>\nstream\n%s\nendstream", len(content), content),
"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
}
var buf bytes.Buffer
buf.WriteString("%PDF-1.4\n")
offsets := make([]int, len(objs)+1)
for i, o := range objs {
offsets[i+1] = buf.Len()
fmt.Fprintf(&buf, "%d 0 obj\n%s\nendobj\n", i+1, o)
}
xrefPos := buf.Len()
fmt.Fprintf(&buf, "xref\n0 %d\n", len(objs)+1)
buf.WriteString("0000000000 65535 f \n")
for i := 1; i <= len(objs); i++ {
fmt.Fprintf(&buf, "%010d 00000 n \n", offsets[i])
}
fmt.Fprintf(&buf, "trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF", len(objs)+1, xrefPos)
return buf.Bytes()
}
// TestExtractPDFText 验证文字版 PDF 能抽出文本
func TestExtractPDFText(t *testing.T) {
data := buildMiniPDF("Mahuang Decoction for cold")
text, err := extractPDFText(data)
if err != nil {
t.Fatalf("extractPDFText 报错: %v", err)
}
if !strings.Contains(text, "Mahuang") {
t.Errorf("PDF 文本抽取结果不含预期内容,实际: %q", text)
}
}
// TestExtractPDFText_Invalid 畸形 PDF 要报错而不是 panic
func TestExtractPDFText_Invalid(t *testing.T) {
if _, err := extractPDFText([]byte("%PDF-1.4 这是坏文件")); err == nil {
t.Fatal("畸形 PDF 应当报错")
}
}
// TestDecodeToUTF8_GBK 验证 GBK 字节流自动转码(「中医药」的 GBK 编码)
func TestDecodeToUTF8_GBK(t *testing.T) {
gbk := []byte{0xD6, 0xD0, 0xD2, 0xBD, 0xD2, 0xA9}
if got := decodeToUTF8(gbk); got != "中医药" {
t.Errorf("GBK 解码失败,got=%q", got)
}
// 合法 UTF-8 原样保留
if got := decodeToUTF8([]byte("中医药")); got != "中医药" {
t.Errorf("UTF-8 被误转码,got=%q", got)
}
// UTF-8 BOM 被剥掉
withBOM := append([]byte{0xEF, 0xBB, 0xBF}, []byte("中医药")...)
if got := decodeToUTF8(withBOM); got != "中医药" {
t.Errorf("UTF-8 BOM 未剥离,got=%q", got)
}
}
// TestParseCSV 验证 CSV 解析(含引号字段与不等列数容错)
func TestParseCSV(t *testing.T) {
csvData := []byte("药名,性味,功效\n麻黄,辛温,\"发汗解表,宣肺平喘\"\n桂枝,辛甘温\n")
chunks, raw, err := parseCSV(csvData, DefaultChunkOptions())
if err != nil {
t.Fatalf("parseCSV 报错: %v", err)
}
if len(chunks) == 0 {
t.Fatal("parseCSV 未产出 chunk")
}
if !strings.Contains(raw, "麻黄 | 辛温 | 发汗解表,宣肺平喘") {
t.Errorf("CSV 行未按「 | 」拼接,raw:\n%s", raw)
}
if !strings.Contains(raw, "桂枝 | 辛甘温") {
t.Errorf("不等列数的行解析失败,raw:\n%s", raw)
}
}
// TestParseExcel_MaxKB 三列带表头 → 保持 MaxKB 语义(行即 chunk)
func TestParseExcel_MaxKB(t *testing.T) {
f := excelize.NewFile()
sheet := f.GetSheetName(0)
_ = f.SetSheetRow(sheet, "A1", &[]interface{}{"分段标题", "分段内容", "问题"})
_ = f.SetSheetRow(sheet, "A2", &[]interface{}{"麻黄汤", "主治外感风寒表实证", "感冒怎么办;风寒感冒用什么方"})
buf, err := f.WriteToBuffer()
if err != nil {
t.Fatalf("生成 xlsx 失败: %v", err)
}
chunks, _, err := parseExcel(buf.Bytes(), DefaultChunkOptions())
if err != nil {
t.Fatalf("parseExcel 报错: %v", err)
}
if len(chunks) != 1 {
t.Fatalf("MaxKB 模式应产出 1 个 chunk,实际 %d 个", len(chunks))
}
if chunks[0].Title != "麻黄汤" || !strings.Contains(chunks[0].Content, "外感风寒") {
t.Errorf("MaxKB 三列语义解析错误: %+v", chunks[0])
}
if chunks[0].Meta == nil {
t.Errorf("related_questions 未写入 meta")
}
}
// TestParseExcel_Generic 四列以上无表头 → 通用表格模式(sheet 名为 chunk 标题)
func TestParseExcel_Generic(t *testing.T) {
f := excelize.NewFile()
sheet := f.GetSheetName(0)
_ = f.SetSheetName(sheet, "常用中药")
_ = f.SetSheetRow("常用中药", "A1", &[]interface{}{"麻黄", "辛温", "肺经", "发汗解表"})
_ = f.SetSheetRow("常用中药", "A2", &[]interface{}{"桂枝", "辛甘温", "心经", "温通经脉"})
buf, err := f.WriteToBuffer()
if err != nil {
t.Fatalf("生成 xlsx 失败: %v", err)
}
chunks, raw, err := parseExcel(buf.Bytes(), DefaultChunkOptions())
if err != nil {
t.Fatalf("parseExcel 报错: %v", err)
}
if len(chunks) == 0 {
t.Fatal("通用表格模式未产出 chunk")
}
if chunks[0].Title != "常用中药" {
t.Errorf("通用模式 chunk 标题应为 sheet 名,实际 %q", chunks[0].Title)
}
if !strings.Contains(raw, "麻黄 | 辛温 | 肺经 | 发汗解表") {
t.Errorf("通用表格行拼接错误,raw:\n%s", raw)
}
}
// TestTruncateRunes 验证查询截断不破坏 UTF-8 字符边界
func TestTruncateRunes(t *testing.T) {
if got := truncateRunes("麻黄汤治感冒", 3); got != "麻黄汤" {
t.Errorf("中文截断错误,got=%q", got)
}
if got := truncateRunes("短查询", 500); got != "短查询" {
t.Errorf("短于上限不应截断,got=%q", got)
}
if got := truncateRunes("abc", 0); got != "abc" {
t.Errorf("max<=0 应原样返回,got=%q", got)
}
// 截断结果必须仍是合法 UTF-8
long := strings.Repeat("风寒感冒", 200)
if cut := truncateRunes(long, 500); !utf8.ValidString(cut) || utf8.RuneCountInString(cut) != 500 {
t.Errorf("长文本截断错误,len=%d valid=%v", utf8.RuneCountInString(cut), utf8.ValidString(cut))
}
}
// TestParseFileFromBytes_Dispatch 验证扩展名分发与防御逻辑
func TestParseFileFromBytes_Dispatch(t *testing.T) {
// docx 走 docx 分支
docx := buildMiniDocx(t, `四君子汤补气健脾`)
pdoc, err := ParseFileFromBytes("方剂.docx", docx, DefaultChunkOptions())
if err != nil {
t.Fatalf("docx 分发失败: %v", err)
}
if pdoc.SourceType != "docx" || len(pdoc.Chunks) == 0 {
t.Errorf("docx 解析结果异常: type=%s chunks=%d", pdoc.SourceType, len(pdoc.Chunks))
}
// 老版 .doc 明确拒绝
if _, err := ParseFileFromBytes("旧文档.doc", []byte("x"), DefaultChunkOptions()); err == nil {
t.Error(".doc 应当被拒绝")
}
// 不支持的扩展名报错
if _, err := ParseFileFromBytes("图片.png", []byte("x"), DefaultChunkOptions()); err == nil {
t.Error("png 应当被拒绝")
}
// 空文本文件:解析成功但没有 chunk → 报「没有可导入的文本内容」
if _, err := ParseFileFromBytes("空.txt", []byte(" \n\n "), DefaultChunkOptions()); err == nil {
t.Error("空文件应当报错")
}
// html 走 html 分支且标题进 chunk.title
html := []byte(`温病条辨
太阴风温、温热、温疫、冬温,初起恶风寒者,桂枝汤主之。
`)
pdoc, err = ParseFileFromBytes("wenbing.html", html, DefaultChunkOptions())
if err != nil {
t.Fatalf("html 分发失败: %v", err)
}
if pdoc.SourceType != "html" || len(pdoc.Chunks) == 0 || pdoc.Chunks[0].Title != "温病条辨" {
t.Errorf("html 解析结果异常: type=%s chunks=%+v", pdoc.SourceType, pdoc.Chunks)
}
}