773 lines
28 KiB
Python
773 lines
28 KiB
Python
from __future__ import annotations
|
||
|
||
import io
|
||
import json
|
||
import xml.etree.ElementTree as ET
|
||
import zipfile
|
||
from datetime import datetime
|
||
|
||
import pytest
|
||
from docx import Document
|
||
from openpyxl import Workbook
|
||
from pptx import Presentation
|
||
from pptx.util import Inches
|
||
from pypdf import PdfWriter
|
||
|
||
from app.modules.data_process.algorithms import (
|
||
PdfPageText,
|
||
content_quality_flags,
|
||
desensitize_pii,
|
||
desensitize_structured_record,
|
||
detect_document_structure,
|
||
detect_pdf_document_noise,
|
||
detect_text_format,
|
||
extract_pdf_page_texts,
|
||
extract_structured_records,
|
||
generate_standard_records,
|
||
is_near_duplicate,
|
||
merge_short_blocks,
|
||
normalize_text,
|
||
parse_text_content,
|
||
preprocess_structured_records,
|
||
record_fingerprint,
|
||
remove_document_noise,
|
||
score_quality,
|
||
stable_split,
|
||
stable_split_assignments,
|
||
)
|
||
|
||
|
||
def _pdf_page_texts(*texts: str) -> tuple[PdfPageText, ...]:
|
||
pages: list[PdfPageText] = []
|
||
offset = 0
|
||
for page_number, text in enumerate(texts, start=1):
|
||
normalized = normalize_text(text)
|
||
if pages:
|
||
offset += 2
|
||
start = offset
|
||
offset += len(normalized)
|
||
pages.append(
|
||
PdfPageText(
|
||
page_number=page_number,
|
||
text=normalized,
|
||
source_start=start,
|
||
source_end=offset,
|
||
)
|
||
)
|
||
return tuple(pages)
|
||
|
||
|
||
def _minimal_pdf(text: str = "Hello PDF") -> bytes:
|
||
stream = f"BT /F1 12 Tf 72 720 Td ({text}) Tj ET".encode("ascii")
|
||
objects = [
|
||
b"<< /Type /Catalog /Pages 2 0 R >>",
|
||
b"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
|
||
(
|
||
b"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] "
|
||
b"/Resources << /Font << /F1 5 0 R >> >> /Contents 4 0 R >>"
|
||
),
|
||
b"<< /Length " + str(len(stream)).encode("ascii") + b" >>\nstream\n"
|
||
+ stream
|
||
+ b"\nendstream",
|
||
b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
|
||
]
|
||
result = bytearray(b"%PDF-1.4\n%\xe2\xe3\xcf\xd3\n")
|
||
offsets = [0]
|
||
for object_number, value in enumerate(objects, start=1):
|
||
offsets.append(len(result))
|
||
result.extend(f"{object_number} 0 obj\n".encode("ascii"))
|
||
result.extend(value)
|
||
result.extend(b"\nendobj\n")
|
||
xref_offset = len(result)
|
||
result.extend(f"xref\n0 {len(objects) + 1}\n".encode("ascii"))
|
||
result.extend(b"0000000000 65535 f \n")
|
||
for offset in offsets[1:]:
|
||
result.extend(f"{offset:010d} 00000 n \n".encode("ascii"))
|
||
result.extend(
|
||
(
|
||
f"trailer\n<< /Size {len(objects) + 1} /Root 1 0 R >>\n"
|
||
f"startxref\n{xref_offset}\n%%EOF\n"
|
||
).encode("ascii")
|
||
)
|
||
return bytes(result)
|
||
|
||
|
||
def _aes_encrypted_pdf(*, user_password: str) -> bytes:
|
||
writer = PdfWriter(clone_from=io.BytesIO(_minimal_pdf()))
|
||
writer.encrypt(
|
||
user_password=user_password,
|
||
owner_password="owner-secret",
|
||
algorithm="AES-256",
|
||
)
|
||
output = io.BytesIO()
|
||
writer.write(output)
|
||
return output.getvalue()
|
||
|
||
|
||
def _docx_bytes() -> bytes:
|
||
document = Document()
|
||
document.add_heading("服务说明", level=1)
|
||
document.add_paragraph("这是 DOCX 正文。")
|
||
table = document.add_table(rows=1, cols=2)
|
||
table.cell(0, 0).text = "字段"
|
||
table.cell(0, 1).text = "内容"
|
||
output = io.BytesIO()
|
||
document.save(output)
|
||
return output.getvalue()
|
||
|
||
|
||
def _xlsx_bytes() -> bytes:
|
||
workbook = Workbook()
|
||
worksheet = workbook.active
|
||
worksheet.title = "数据"
|
||
worksheet.append(["name", "score", "created_at"])
|
||
worksheet.append(["Alice", 95, datetime(2026, 7, 23, 10, 30)])
|
||
worksheet.append(["Bob", 88, datetime(2026, 7, 24, 9, 0)])
|
||
output = io.BytesIO()
|
||
workbook.save(output)
|
||
workbook.close()
|
||
return output.getvalue()
|
||
|
||
|
||
def _xlsx_with_worksheet_relationship(
|
||
raw: bytes,
|
||
target: str,
|
||
*,
|
||
target_mode: str | None = None,
|
||
) -> bytes:
|
||
member_name = "xl/_rels/workbook.xml.rels"
|
||
source = io.BytesIO(raw)
|
||
output = io.BytesIO()
|
||
with zipfile.ZipFile(source) as original, zipfile.ZipFile(output, "w") as rewritten:
|
||
for member in original.infolist():
|
||
content = original.read(member.filename)
|
||
if member.filename == member_name:
|
||
root = ET.fromstring(content)
|
||
worksheet_relationship = next(
|
||
element
|
||
for element in root
|
||
if element.attrib.get("Type", "").endswith("/worksheet")
|
||
)
|
||
worksheet_relationship.set("Target", target)
|
||
if target_mode is None:
|
||
worksheet_relationship.attrib.pop("TargetMode", None)
|
||
else:
|
||
worksheet_relationship.set("TargetMode", target_mode)
|
||
content = ET.tostring(root, encoding="utf-8", xml_declaration=True)
|
||
rewritten.writestr(member, content)
|
||
return output.getvalue()
|
||
|
||
|
||
def _pptx_bytes() -> bytes:
|
||
presentation = Presentation()
|
||
slide = presentation.slides.add_slide(presentation.slide_layouts[6])
|
||
text_box = slide.shapes.add_textbox(Inches(1), Inches(1), Inches(6), Inches(1))
|
||
text_box.text = "PPTX 页面正文"
|
||
output = io.BytesIO()
|
||
presentation.save(output)
|
||
return output.getvalue()
|
||
|
||
|
||
def test_parse_utf8_json_jsonl_csv_markdown_and_txt() -> None:
|
||
parsed_json = parse_text_content(
|
||
b'\xef\xbb\xbf{"data":[{"name":"\xe5\xbc\xa0\xe4\xb8\x89"}]}',
|
||
filename="records.json",
|
||
)
|
||
assert parsed_json.format == "json"
|
||
assert parsed_json.records == ({"name": "张三"},)
|
||
|
||
parsed_jsonl = parse_text_content('{"id":1}\n\n{"id":2}\n', filename="records.jsonl")
|
||
assert parsed_jsonl.format == "jsonl"
|
||
assert parsed_jsonl.records == ({"id": 1}, {"id": 2})
|
||
|
||
parsed_csv = parse_text_content("name,answer\r\nAlice,yes\r\nBob,no", filename="records.csv")
|
||
assert parsed_csv.format == "csv"
|
||
assert parsed_csv.text == "name,answer\nAlice,yes\nBob,no"
|
||
assert parsed_csv.records[1] == {"name": "Bob", "answer": "no"}
|
||
|
||
parsed_markdown = parse_text_content("# 标题\n\n正文", filename="README.md")
|
||
assert parsed_markdown.format == "markdown"
|
||
assert parsed_markdown.records == ()
|
||
|
||
parsed_txt = parse_text_content("普通文本", filename="note.txt")
|
||
assert parsed_txt.format == "txt"
|
||
assert parsed_txt.text == "普通文本"
|
||
|
||
|
||
def test_parse_pdf_docx_xlsx_and_pptx() -> None:
|
||
parsed_pdf = parse_text_content(_minimal_pdf(), filename="manual.pdf")
|
||
assert parsed_pdf.format == "pdf"
|
||
assert "Hello PDF" in parsed_pdf.text
|
||
assert parsed_pdf.records == ()
|
||
|
||
pdf_pages = extract_pdf_page_texts(_minimal_pdf())
|
||
assert len(pdf_pages) == 1
|
||
assert pdf_pages[0].page_number == 1
|
||
assert pdf_pages[0].text == "Hello PDF"
|
||
assert pdf_pages[0].source_start == 0
|
||
assert pdf_pages[0].source_end == len(parsed_pdf.text)
|
||
|
||
parsed_docx = parse_text_content(_docx_bytes(), filename="manual.docx")
|
||
assert parsed_docx.format == "docx"
|
||
assert "服务说明" in parsed_docx.text
|
||
assert "这是 DOCX 正文。" in parsed_docx.text
|
||
assert "字段\t内容" in parsed_docx.text
|
||
assert parsed_docx.records == ()
|
||
|
||
parsed_xlsx = parse_text_content(_xlsx_bytes(), filename="records.xlsx")
|
||
assert parsed_xlsx.format == "xlsx"
|
||
assert parsed_xlsx.records == (
|
||
{"name": "Alice", "score": 95, "created_at": "2026-07-23T10:30:00"},
|
||
{"name": "Bob", "score": 88, "created_at": "2026-07-24T09:00:00"},
|
||
)
|
||
assert json.loads(parsed_xlsx.text.splitlines()[0]) == parsed_xlsx.records[0]
|
||
|
||
parsed_pptx = parse_text_content(_pptx_bytes(), filename="slides.pptx")
|
||
assert parsed_pptx.format == "pptx"
|
||
assert parsed_pptx.text == "PPTX 页面正文"
|
||
assert parsed_pptx.records == ()
|
||
|
||
|
||
def test_pdf_document_noise_removes_headers_page_numbers_and_toc_safely() -> None:
|
||
pages = _pdf_page_texts(
|
||
"""
|
||
远光制度文件 文件编码 2024
|
||
秘密等级 商密【中】
|
||
第 1 页 共 5 页
|
||
正文第一页，关于适用范围的说明。
|
||
业务提示保留
|
||
第一页补充说明甲
|
||
第一页补充说明乙
|
||
第一页补充说明丙
|
||
""",
|
||
"""
|
||
远光制度文件 文件编码 2024
|
||
秘密等级 商密【中】
|
||
第 2 页 共 5 页
|
||
目 录
|
||
第一章 总则........3
|
||
第二章 报销申请........4
|
||
第三章 附则........5
|
||
""",
|
||
"""
|
||
远光制度文件 文件编码 2024
|
||
秘密等级 商密【中】
|
||
第 3 页 共 5 页
|
||
1.1 管理要求........6
|
||
1.2 审批职责 7
|
||
1.3 费用标准........8
|
||
1.4 例外处理........9
|
||
""",
|
||
"""
|
||
远光制度文件 文件编码 2024
|
||
秘密等级 商密【中】
|
||
第 4 页 共 5 页
|
||
正文中可以说“请参见第 3 页说明”，不应误删。
|
||
第 99 页 共 100 页
|
||
系统可用率........99.9%
|
||
业务提示保留
|
||
第四页补充说明甲
|
||
第四页补充说明乙
|
||
第四页补充说明丙
|
||
""",
|
||
"""
|
||
远光制度文件 文件编码 2024
|
||
秘密等级 商密【中】
|
||
第 5 页 共 5 页
|
||
本办法自发布之日起施行。
|
||
业务提示保留
|
||
第五页补充说明甲
|
||
第五页补充说明乙
|
||
第五页补充说明丙
|
||
""",
|
||
)
|
||
source = "\n\n".join(page.text for page in pages)
|
||
|
||
spans = detect_pdf_document_noise(pages)
|
||
cleaned = remove_document_noise(source, spans)
|
||
|
||
assert {span.kind for span in spans} == {
|
||
"page_number",
|
||
"repeated_margin",
|
||
"table_of_contents",
|
||
}
|
||
assert "远光制度文件" not in cleaned
|
||
assert "商密【中】" not in cleaned
|
||
assert "第 1 页 共 5 页" not in cleaned
|
||
assert "第一章 总则" not in cleaned
|
||
assert "1.2 审批职责 7" not in cleaned
|
||
assert "请参见第 3 页说明" in cleaned
|
||
assert "第 99 页 共 100 页" in cleaned
|
||
assert "系统可用率........99.9%" in cleaned
|
||
assert cleaned.count("业务提示保留") == 3
|
||
|
||
|
||
def test_pdf_document_noise_does_not_infer_repeated_margins_for_short_documents() -> None:
|
||
pages = _pdf_page_texts(
|
||
"公司内部文件\n正文 A",
|
||
"公司内部文件\n正文 B",
|
||
)
|
||
|
||
spans = detect_pdf_document_noise(pages)
|
||
|
||
assert not any(span.kind == "repeated_margin" for span in spans)
|
||
|
||
|
||
def test_xlsx_merged_multilevel_headers_are_flattened_without_losing_columns() -> None:
|
||
workbook = Workbook()
|
||
worksheet = workbook.active
|
||
worksheet.merge_cells("A1:A2")
|
||
worksheet.merge_cells("B1:C1")
|
||
worksheet["A1"] = "地区"
|
||
worksheet["B1"] = "销售"
|
||
worksheet["B2"] = "Q1"
|
||
worksheet["C2"] = "Q2"
|
||
worksheet.append(["华东", 100, 120])
|
||
output = io.BytesIO()
|
||
workbook.save(output)
|
||
workbook.close()
|
||
|
||
parsed = parse_text_content(output.getvalue(), filename="sales.xlsx")
|
||
assert parsed.records == ({"地区": "华东", "销售.Q1": 100, "销售.Q2": 120},)
|
||
|
||
|
||
def test_xlsx_header_inference_skips_more_than_eight_merged_report_titles() -> None:
|
||
workbook = Workbook()
|
||
worksheet = workbook.active
|
||
for row_number in range(1, 13):
|
||
worksheet.merge_cells(
|
||
start_row=row_number,
|
||
start_column=1,
|
||
end_row=row_number,
|
||
end_column=4,
|
||
)
|
||
worksheet.cell(row_number, 1, f"报表说明 {row_number}")
|
||
worksheet.append(["姓名", "部门", "得分", "日期"])
|
||
worksheet.append(["张三", "研发", 95, "2026-07-23"])
|
||
output = io.BytesIO()
|
||
workbook.save(output)
|
||
workbook.close()
|
||
|
||
parsed = parse_text_content(output.getvalue(), filename="report.xlsx")
|
||
assert parsed.records == (
|
||
{"姓名": "张三", "部门": "研发", "得分": 95, "日期": "2026-07-23"},
|
||
)
|
||
|
||
|
||
def test_xlsx_header_inference_ignores_continuous_body_merges() -> None:
|
||
workbook = Workbook()
|
||
worksheet = workbook.active
|
||
worksheet.append(["类别", "名称", "数量"])
|
||
worksheet.append(["水果", "苹果", 10])
|
||
worksheet.append([None, "香蕉", 12])
|
||
worksheet.append(["蔬菜", "白菜", 8])
|
||
worksheet.append([None, "萝卜", 9])
|
||
worksheet.merge_cells("A2:A3")
|
||
worksheet.merge_cells("A4:A5")
|
||
output = io.BytesIO()
|
||
workbook.save(output)
|
||
workbook.close()
|
||
|
||
parsed = parse_text_content(output.getvalue(), filename="inventory.xlsx")
|
||
assert parsed.records == (
|
||
{"类别": "水果", "名称": "苹果", "数量": 10},
|
||
{"类别": "", "名称": "香蕉", "数量": 12},
|
||
{"类别": "蔬菜", "名称": "白菜", "数量": 8},
|
||
{"类别": "", "名称": "萝卜", "数量": 9},
|
||
)
|
||
|
||
|
||
def test_xlsx_header_inference_supports_title_and_two_header_levels() -> None:
|
||
workbook = Workbook()
|
||
worksheet = workbook.active
|
||
worksheet.merge_cells("A1:C1")
|
||
worksheet["A1"] = "区域销售报表"
|
||
worksheet["A2"] = "统计日期"
|
||
worksheet["B2"] = "2026-07-23"
|
||
worksheet.merge_cells("A4:A5")
|
||
worksheet.merge_cells("B4:C4")
|
||
worksheet["A4"] = "地区"
|
||
worksheet["B4"] = "销售"
|
||
worksheet["B5"] = "Q1"
|
||
worksheet["C5"] = "Q2"
|
||
worksheet.append(["华南", 88, 92])
|
||
output = io.BytesIO()
|
||
workbook.save(output)
|
||
workbook.close()
|
||
|
||
parsed = parse_text_content(output.getvalue(), filename="two-level.xlsx")
|
||
assert parsed.records == (
|
||
{"地区": "华南", "销售.Q1": 88, "销售.Q2": 92},
|
||
)
|
||
|
||
|
||
def test_xlsx_header_inference_supports_title_and_three_header_levels() -> None:
|
||
workbook = Workbook()
|
||
worksheet = workbook.active
|
||
worksheet.merge_cells("A1:D1")
|
||
worksheet["A1"] = "年度销售分析报告"
|
||
worksheet["A2"] = "统计日期"
|
||
worksheet["B2"] = "2026-07-23"
|
||
worksheet.merge_cells("A4:A6")
|
||
worksheet.merge_cells("B4:D4")
|
||
worksheet.merge_cells("B5:C5")
|
||
worksheet.merge_cells("D5:D6")
|
||
worksheet["A4"] = "地区"
|
||
worksheet["B4"] = "销售"
|
||
worksheet["B5"] = "国内"
|
||
worksheet["D5"] = "海外"
|
||
worksheet["B6"] = "Q1"
|
||
worksheet["C6"] = "Q2"
|
||
worksheet.append(["华东", 100, 120, 80])
|
||
output = io.BytesIO()
|
||
workbook.save(output)
|
||
workbook.close()
|
||
|
||
parsed = parse_text_content(output.getvalue(), filename="three-level.xlsx")
|
||
assert parsed.records == (
|
||
{
|
||
"地区": "华东",
|
||
"销售.国内.Q1": 100,
|
||
"销售.国内.Q2": 120,
|
||
"销售.海外": 80,
|
||
},
|
||
)
|
||
|
||
|
||
def test_xlsx_header_inference_keeps_an_ordinary_single_header_row() -> None:
|
||
parsed = parse_text_content(_xlsx_bytes(), filename="ordinary.xlsx")
|
||
assert tuple(parsed.records[0]) == ("name", "score", "created_at")
|
||
assert len(parsed.records) == 2
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"target",
|
||
[
|
||
"./worksheets/../worksheets/sheet1.xml",
|
||
"./worksheets/%2e%2e/worksheets/sheet1.xml",
|
||
"../xl/worksheets/sheet1.xml",
|
||
"/xl/worksheets/./sheet1.xml",
|
||
],
|
||
)
|
||
def test_xlsx_worksheet_relationship_allows_safe_dot_segments(target: str) -> None:
|
||
raw = _xlsx_with_worksheet_relationship(_xlsx_bytes(), target)
|
||
parsed = parse_text_content(raw, filename="records.xlsx")
|
||
assert parsed.records[0]["name"] == "Alice"
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"target",
|
||
[
|
||
"../../outside.xml",
|
||
"worksheets\\sheet1.xml",
|
||
"%2e%2e/%2e%2e/outside.xml",
|
||
"%252e%252e/%252e%252e/outside.xml",
|
||
"https://example.com/sheet1.xml",
|
||
],
|
||
)
|
||
def test_xlsx_worksheet_relationship_rejects_path_traversal(target: str) -> None:
|
||
raw = _xlsx_with_worksheet_relationship(_xlsx_bytes(), target)
|
||
with pytest.raises(ValueError, match="unsafe worksheet path"):
|
||
parse_text_content(raw, filename="unsafe.xlsx")
|
||
|
||
|
||
def test_xlsx_worksheet_relationship_rejects_external_and_missing_targets() -> None:
|
||
external = _xlsx_with_worksheet_relationship(
|
||
_xlsx_bytes(),
|
||
"https://example.com/sheet1.xml",
|
||
target_mode="External",
|
||
)
|
||
with pytest.raises(ValueError, match="external relationship"):
|
||
parse_text_content(external, filename="external.xlsx")
|
||
|
||
missing = _xlsx_with_worksheet_relationship(
|
||
_xlsx_bytes(),
|
||
"worksheets/missing.xml",
|
||
)
|
||
with pytest.raises(ValueError, match="target does not exist"):
|
||
parse_text_content(missing, filename="missing.xlsx")
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("filename", "replacement"),
|
||
[
|
||
("legacy.doc", ".docx"),
|
||
("legacy.xls", ".xlsx"),
|
||
("legacy.ppt", ".pptx"),
|
||
],
|
||
)
|
||
def test_legacy_office_formats_require_conversion(filename: str, replacement: str) -> None:
|
||
with pytest.raises(ValueError, match=rf"convert the file to \{replacement}"):
|
||
parse_text_content(b"legacy", filename=filename)
|
||
|
||
|
||
def test_office_zip_bomb_and_invalid_pdf_are_rejected_before_parsing() -> None:
|
||
archive = io.BytesIO()
|
||
with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_DEFLATED) as package:
|
||
package.writestr("[Content_Types].xml", "<Types/>")
|
||
package.writestr("word/document.xml", b"A" * (2 * 1024 * 1024))
|
||
with pytest.raises(ValueError, match="unsafe compression ratio"):
|
||
parse_text_content(archive.getvalue(), filename="unsafe.docx")
|
||
|
||
active_xml = io.BytesIO()
|
||
with zipfile.ZipFile(active_xml, "w") as package:
|
||
package.writestr("[Content_Types].xml", "<Types/>")
|
||
package.writestr(
|
||
"word/document.xml",
|
||
'<!DOCTYPE document [<!ENTITY xxe SYSTEM "file:///etc/passwd">]><document/>',
|
||
)
|
||
with pytest.raises(ValueError, match="unsupported active XML"):
|
||
parse_text_content(active_xml.getvalue(), filename="active.docx")
|
||
|
||
with pytest.raises(ValueError, match="missing PDF header"):
|
||
parse_text_content(b"not a pdf", filename="broken.pdf")
|
||
|
||
blank_pdf = io.BytesIO()
|
||
blank_writer = PdfWriter()
|
||
blank_writer.add_blank_page(width=612, height=792)
|
||
blank_writer.write(blank_pdf)
|
||
with pytest.raises(ValueError, match="scanned PDF requires OCR"):
|
||
parse_text_content(blank_pdf.getvalue(), filename="scanned.pdf")
|
||
|
||
aes_pdf_without_open_password = parse_text_content(
|
||
_aes_encrypted_pdf(user_password=""),
|
||
filename="aes-no-password.pdf",
|
||
)
|
||
assert "Hello PDF" in aes_pdf_without_open_password.text
|
||
|
||
with pytest.raises(ValueError, match="password-protected PDF files are not supported"):
|
||
parse_text_content(
|
||
_aes_encrypted_pdf(user_password="secret"),
|
||
filename="aes-password.pdf",
|
||
)
|
||
|
||
|
||
def test_invalid_utf8_and_malformed_structured_content_fail_loudly() -> None:
|
||
with pytest.raises(ValueError, match="not valid UTF-8"):
|
||
parse_text_content(b"\xff\xfe", filename="broken.txt")
|
||
with pytest.raises(ValueError, match="invalid JSONL at line 2"):
|
||
extract_structured_records('{"id":1}\nnot-json', "jsonl")
|
||
with pytest.raises(ValueError, match="more fields"):
|
||
extract_structured_records("a,b\n1,2,3", "csv")
|
||
|
||
|
||
def test_detect_format_from_content_and_normalize() -> None:
|
||
assert detect_text_format(text='{"id":1}\n{"id":2}') == "jsonl"
|
||
assert detect_text_format(text="# Heading\ntext") == "markdown"
|
||
assert detect_text_format(text="a,b\n1,2") == "csv"
|
||
assert normalize_text("\ufeffＡＢＣ \r\n第二\x00行\u200b\t \r\n") == "ABC\n第二行"
|
||
|
||
|
||
def test_extract_json_scalar_and_nested_values_are_stable() -> None:
|
||
assert extract_structured_records("[1, true, null]", "json") == [
|
||
{"value": 1},
|
||
{"value": True},
|
||
{"value": None},
|
||
]
|
||
result = extract_structured_records(
|
||
json.dumps({"items": [{"text": " 内容 "}], "ignored": 1}, ensure_ascii=False),
|
||
"json",
|
||
)
|
||
assert result == [{"text": "内容"}]
|
||
|
||
|
||
def test_desensitize_pii_returns_masked_text_and_counts() -> None:
|
||
source = "邮箱 a.user+tag@example.com，手机 +86 13800138000，身份证 11010519491231002X。"
|
||
masked, counts = desensitize_pii(source)
|
||
assert masked == "邮箱 [EMAIL]，手机 [PHONE]，身份证 [ID_CARD]。"
|
||
assert counts == {"email": 1, "phone": 1, "id_card": 1, "total": 3}
|
||
|
||
|
||
def test_every_structured_preprocess_option_has_independent_behavior() -> None:
|
||
clean_source = [
|
||
{"id": "1", "name": "有效", "empty_column": ""},
|
||
{"id": "", "name": "缺少关键字段", "empty_column": ""},
|
||
{"id": "2", "name": "有效", "empty_column": ""},
|
||
]
|
||
assert preprocess_structured_records(clean_source, []) == clean_source
|
||
assert preprocess_structured_records(clean_source, ["clean_invalid"]) == [
|
||
{"id": "1", "name": "有效"},
|
||
{"id": "2", "name": "有效"},
|
||
]
|
||
|
||
nested = [{"id": 1, "profile": {"name": "张三", "level": 2}}]
|
||
assert "profile" in preprocess_structured_records(nested, [])[0]
|
||
assert preprocess_structured_records(nested, ["detect_structure"])[0] == {
|
||
"id": 1,
|
||
"profile.name": "张三",
|
||
"profile.level": 2,
|
||
}
|
||
|
||
duplicates = [
|
||
{"customer_id": "C-1", "value": "first"},
|
||
{"customer_id": "C-1", "value": "updated"},
|
||
{"customer_id": "", "value": "blank-one"},
|
||
{"customer_id": "", "value": "blank-two"},
|
||
]
|
||
assert len(preprocess_structured_records(duplicates, [])) == 4
|
||
deduplicated = preprocess_structured_records(duplicates, ["deduplicate"])
|
||
assert [record["value"] for record in deduplicated] == [
|
||
"first",
|
||
"blank-one",
|
||
"blank-two",
|
||
]
|
||
|
||
unnormalized = [{" User Name ": "ＡＢＣ\r\n第二行"}]
|
||
assert preprocess_structured_records(unnormalized, []) == unnormalized
|
||
assert preprocess_structured_records(unnormalized, ["normalize_format"]) == [
|
||
{"user_name": "ABC\n第二行"}
|
||
]
|
||
|
||
anomaly_source = [
|
||
{"id": 10_000 + index, "amount": amount, "text": "正常内容"}
|
||
for index, amount in enumerate((10, 10, 11, 11, 12, 12, 13, 1000))
|
||
]
|
||
assert len(preprocess_structured_records(anomaly_source, [])) == 8
|
||
filtered = preprocess_structured_records(anomaly_source, ["filter_anomaly"])
|
||
assert len(filtered) == 7
|
||
assert all(record["amount"] != 1000 for record in filtered)
|
||
assert max(record["id"] for record in filtered) > 10_000
|
||
|
||
sensitive = [{"姓名": "张三", "phone": "13800138000", "email": "a@b.com"}]
|
||
assert preprocess_structured_records(sensitive, []) == sensitive
|
||
masked = preprocess_structured_records(sensitive, ["desensitize"])[0]
|
||
assert masked == {"姓名": "[NAME]", "phone": "[PHONE]", "email": "[EMAIL]"}
|
||
|
||
|
||
def test_structured_desensitization_counts_and_document_helpers() -> None:
|
||
masked, counts = desensitize_structured_record(
|
||
{"联系人姓名": "李四", "说明": "邮箱 user@example.com，手机 13900139000"}
|
||
)
|
||
assert masked == {
|
||
"联系人姓名": "[NAME]",
|
||
"说明": "邮箱 [EMAIL]，手机 [PHONE]",
|
||
}
|
||
assert counts == {"email": 1, "phone": 1, "id_card": 0, "name": 1, "total": 3}
|
||
|
||
structure = detect_document_structure(
|
||
"# 第一章\n正文\n\n## 细节\n- 项目一\n- 项目二\n\n```python\nprint(1)\n```"
|
||
)
|
||
assert [heading.title for heading in structure.headings] == ["第一章", "细节"]
|
||
assert structure.list_block_count == 1
|
||
assert structure.code_block_count == 1
|
||
assert merge_short_blocks(["短一", "短二", "这是一段足够长的正文内容"], min_token_count=4)
|
||
assert "mojibake" in content_quality_flags("正常文字锟斤拷内容", min_chars=0, min_tokens=0)
|
||
assert is_near_duplicate(
|
||
"alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu nu xi omicron",
|
||
"alpha beta gamma, delta epsilon zeta eta theta iota kappa lambda mu nu xi omicron",
|
||
similarity_threshold=0.92,
|
||
max_hamming_distance=2,
|
||
)
|
||
|
||
|
||
def test_quality_scoring_covers_all_dimensions_and_duplicates() -> None:
|
||
valid = {
|
||
"instruction": "如何修改收货地址？",
|
||
"input": "订单尚未发货",
|
||
"output": "可以在订单详情页申请修改收货地址。",
|
||
}
|
||
source = "订单尚未发货时，可以在订单详情页申请修改收货地址。"
|
||
first_score = score_quality(valid, min_output_length=10, source_content=source)
|
||
assert first_score.is_valid
|
||
assert first_score.completeness == 100
|
||
assert first_score.length == 100
|
||
assert first_score.readability >= 90
|
||
assert first_score.relevance >= 70
|
||
assert first_score.duplicate == 100
|
||
|
||
duplicate_score = score_quality(valid, known_fingerprints={first_score.fingerprint})
|
||
assert duplicate_score.duplicate == 0
|
||
assert "duplicate_record" in duplicate_score.flags
|
||
|
||
unrelated_score = score_quality(
|
||
valid,
|
||
min_output_length=10,
|
||
source_content="量子计算使用量子比特处理信息。",
|
||
)
|
||
assert unrelated_score.relevance < first_score.relevance
|
||
assert "low_source_relevance" in unrelated_score.flags
|
||
|
||
invalid_score = score_quality({"instruction": "", "output": "短"}, min_output_length=10)
|
||
assert not invalid_score.is_valid
|
||
assert {"missing_instruction", "output_too_short"}.issubset(invalid_score.flags)
|
||
assert record_fingerprint(valid) == record_fingerprint(dict(reversed(list(valid.items()))))
|
||
|
||
|
||
def test_stable_split_is_reproducible_and_validates_ratios() -> None:
|
||
first = stable_split("record-42", seed="task-1")
|
||
assert stable_split("record-42", seed="task-1") == first
|
||
assert first in {"train", "validation", "test"}
|
||
assert stable_split("record-42", {"train": 100, "validation": 0, "test": 0}) == "train"
|
||
with pytest.raises(ValueError, match="sum to 100"):
|
||
stable_split("record", {"train": 80, "validation": 10, "test": 9})
|
||
|
||
|
||
def test_stable_split_assignments_use_exact_deterministic_quotas() -> None:
|
||
values = [f"record-{index}" for index in range(28)]
|
||
first = stable_split_assignments(values, seed="task-1")
|
||
second = stable_split_assignments(values, seed="task-1")
|
||
|
||
assert first == second
|
||
assert first.count("train") == 22
|
||
assert first.count("validation") == 3
|
||
assert first.count("test") == 3
|
||
|
||
|
||
def test_generate_standard_records_supports_json_qa_and_stable_variants() -> None:
|
||
previews = [
|
||
{
|
||
"id": "preview-json",
|
||
"edited_content": json.dumps(
|
||
{"instruction": "问题", "input": "上下文", "output": "答案"},
|
||
ensure_ascii=False,
|
||
),
|
||
},
|
||
{"id": "preview-qa", "editedContent": "问：如何操作？\n答：按步骤操作。"},
|
||
]
|
||
records = generate_standard_records(
|
||
previews,
|
||
qa_pairs_per_item=2,
|
||
semantic_enrichment=True,
|
||
split={"train": 100, "validation": 0, "test": 0},
|
||
split_seed="task-1",
|
||
)
|
||
assert len(records) == 4
|
||
assert records[0]["instruction"] == "问题"
|
||
assert records[0]["input"] == "上下文"
|
||
assert records[0]["output"] == "答案"
|
||
assert records[1]["instruction"].endswith("问题")
|
||
assert records[2]["instruction"] == "如何操作?"
|
||
assert records[2]["output"] == "按步骤操作。"
|
||
assert all(record["status"] == "valid" for record in records)
|
||
assert all(record["split"] == "train" for record in records)
|
||
assert records == generate_standard_records(
|
||
previews,
|
||
qa_pairs_per_item=2,
|
||
semantic_enrichment=True,
|
||
split={"train": 100, "validation": 0, "test": 0},
|
||
split_seed="task-1",
|
||
)
|
||
|
||
|
||
def test_generate_standard_records_supports_fifty_unique_semantic_variants() -> None:
|
||
records = generate_standard_records(
|
||
[{"id": "preview-50", "edited_content": "问：如何操作？\n答：按步骤操作。"}],
|
||
qa_pairs_per_item=50,
|
||
semantic_enrichment=True,
|
||
split={"train": 100, "validation": 0, "test": 0},
|
||
split_seed="task-50",
|
||
)
|
||
|
||
assert len(records) == 50
|
||
assert len({record["id"] for record in records}) == 50
|
||
assert len({record["instruction"] for record in records}) == 50
|
||
assert all(record["status"] == "valid" for record in records)
|
||
|
||
|
||
@pytest.mark.parametrize("qa_pairs_per_item", [0, 51])
|
||
def test_generate_standard_records_rejects_out_of_range_count(
|
||
qa_pairs_per_item: int,
|
||
) -> None:
|
||
with pytest.raises(ValueError, match=r"\[1, 50\]"):
|
||
generate_standard_records([], qa_pairs_per_item=qa_pairs_per_item)
|