306 lines
11 KiB
Go
306 lines
11 KiB
Go
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 := `<?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>感冒的中医辨证</w:t></w:r></w:p>
|
||
<w:p><w:r><w:t>风寒感冒：</w:t><w:br/><w:t>恶寒重发热轻</w:t><w:tab/><w:t>无汗头痛</w:t></w:r></w:p>
|
||
<w:tbl>
|
||
<w:tr><w:tc><w:p><w:r><w:t>药名</w:t></w:r></w:p></w:tc><w:tc><w:p><w:r><w:t>麻黄</w:t></w:r></w:p></w:tc></w:tr>
|
||
</w:tbl>
|
||
</w:body>
|
||
</w:document>`
|
||
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 := `<!DOCTYPE html>
|
||
<html><head><title>页面标题</title><style>.a{color:red}</style><script>alert(1)</script></head>
|
||
<body>
|
||
<nav>导航栏不要</nav>
|
||
<h1>中医基础理论</h1>
|
||
<p>阴阳五行学说是中医的理论基础。</p>
|
||
<h2>四诊</h2>
|
||
<ul><li>望诊</li><li>闻诊</li></ul>
|
||
<table><tr><td>寒证</td><td>热证</td></tr></table>
|
||
</body></html>`
|
||
|
||
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, `<w:document xmlns:w="x"><w:body><w:p><w:r><w:t>四君子汤补气健脾</w:t></w:r></w:p></w:body></w:document>`)
|
||
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(`<html><body><h1>温病条辨</h1><p>太阴风温、温热、温疫、冬温，初起恶风寒者，桂枝汤主之。</p></body></html>`)
|
||
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)
|
||
}
|
||
}
|