YG-Soft/YG_FT
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(data-process): 完善文件解析与切分存储链路

Browse Source
This commit is contained in:
caoxiaozhu
2026-07-24 11:27:51 +08:00
parent 6d4bf85284
commit 33d0ed2e01
12 changed files with 4813 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files

546
backend/tests/test_data_process_algorithms.py
View File

@@ -1,23 +1,151 @@
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 (
chunk_unstructured,
content_quality_flags,
desensitize_pii,
desensitize_structured_record,
detect_document_structure,
detect_text_format,
estimate_token_count,
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,
score_quality,
stable_split,
)
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"}]}',
@@ -44,6 +172,269 @@ def test_parse_utf8_json_jsonl_csv_markdown_and_txt() -> None:
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_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")
@@ -80,7 +471,89 @@ def test_desensitize_pii_returns_masked_text_and_counts() -> None:
assert counts == {"email": 1, "phone": 1, "id_card": 1, "total": 3}
@pytest.mark.parametrize("method", ["semantic", "heading", "fixed", "custom"])
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,
)
@pytest.mark.parametrize("method", ["structure", "fixed", "custom"])
def test_chunk_methods_preserve_offsets_and_always_advance(method: str) -> None:
text = "# 第一章\n" + "甲。" * 18 + "\n# 第二章\n" + "乙。" * 18
kwargs = {"custom_delimiter": "\\n"} if method == "custom" else {}
@@ -99,8 +572,38 @@ def test_chunk_methods_preserve_offsets_and_always_advance(method: str) -> None:
assert all(chunk.start_line <= chunk.end_line for chunk in chunks)
def test_fixed_chunk_overlap_is_exact_when_chunks_are_large_enough() -> None:
text = " ".join(f"token{i}" for i in range(30))
def test_default_and_structure_chunking_split_headings_without_cross_section_overlap() -> None:
text = (
"# 第一章\n"
+ " ".join(f"alpha{i}" for i in range(18))
+ "\n# 第二章\n"
+ " ".join(f"beta{i}" for i in range(18))
)
normalized = normalize_text(text)
second_chapter_start = normalized.index("# 第二章")
kwargs = {"chunk_size": 10, "chunk_overlap": 3, "min_chunk_size": 4}
default_chunks = chunk_unstructured(text, **kwargs)
structure_chunks = chunk_unstructured(text, method="structure", **kwargs)
assert default_chunks == structure_chunks
assert len(structure_chunks) > 2
assert all(
chunk.content == normalized[chunk.start : chunk.end] for chunk in structure_chunks
)
assert all(
not (chunk.start < second_chapter_start < chunk.end) for chunk in structure_chunks
)
second_chapter_chunks = [
chunk for chunk in structure_chunks if chunk.start >= second_chapter_start
]
assert second_chapter_chunks[0].start == second_chapter_start
assert second_chapter_chunks[0].content.startswith("# 第二章")
def test_fixed_chunk_offsets_and_actual_token_overlap_are_exact() -> None:
text = " ".join(f"token{i}" for i in range(30))
normalized = normalize_text(text)
chunks = chunk_unstructured(
text,
method="fixed",
@@ -108,10 +611,16 @@ def test_fixed_chunk_overlap_is_exact_when_chunks_are_large_enough() -> None:
chunk_overlap=3,
min_chunk_size=4,
)
first_tokens = chunks[0].content.split()
second_tokens = chunks[1].content.split()
assert first_tokens[-3:] == second_tokens[:3]
assert chunks[0].token_count == 10
assert len(chunks) > 2
assert all(chunk.content == normalized[chunk.start : chunk.end] for chunk in chunks)
assert all(chunk.token_count == estimate_token_count(chunk.content) for chunk in chunks)
assert all(chunk.token_count == 10 for chunk in chunks[:-1])
for left, right in zip(chunks, chunks[1:]):
overlap_text = normalized[right.start : left.end]
assert right.start < left.end
assert estimate_token_count(overlap_text) == 3
assert left.content.endswith(overlap_text)
assert right.content.startswith(overlap_text)
def test_chunk_line_numbers_treat_newline_as_previous_line_boundary() -> None:
@@ -129,18 +638,7 @@ def test_chunk_line_numbers_treat_newline_as_previous_line_boundary() -> None:
assert chunks[1].start_line == 2
def test_heading_and_custom_boundaries_are_respected() -> None:
heading_text = "前言 " * 8 + "\n# 第二章\n" + "正文 " * 12
heading_chunks = chunk_unstructured(
heading_text,
method="heading",
chunk_size=20,
chunk_overlap=0,
min_chunk_size=4,
)
assert "# 第二章" not in heading_chunks[0].content
assert heading_chunks[1].content.startswith("#")
def test_custom_delimiter_is_preserved_as_the_chunk_boundary() -> None:
custom_chunks = chunk_unstructured(
"a b c d <CUT> e f g h i j",
method="custom",
@@ -172,6 +670,13 @@ def test_heading_and_custom_boundaries_are_respected() -> None:
)
def test_markdown_protected_blocks_are_not_split(field: str, block: str) -> None:
text = "前言。" * 15 + "\n" + block + "\n" + "结尾。" * 40
unprotected = chunk_unstructured(
text,
method="fixed",
chunk_size=40,
chunk_overlap=0,
min_chunk_size=10,
)
chunks = chunk_unstructured(
text,
method="fixed",
@@ -180,6 +685,7 @@ def test_markdown_protected_blocks_are_not_split(field: str, block: str) -> None
min_chunk_size=10,
**{field: True},
)
assert all(block not in chunk.content for chunk in unprotected)
assert any(block in chunk.content for chunk in chunks)
@@ -194,6 +700,8 @@ def test_markdown_protected_blocks_are_not_split(field: str, block: str) -> None
"cannot exceed",
),
({"method": "custom", "custom_delimiter": ""}, "custom_delimiter"),
({"method": "semantic"}, "unsupported chunk method"),
({"method": "heading"}, "unsupported chunk method"),
],
)
def test_chunk_configuration_validation(kwargs: dict[str, object], message: str) -> None:

808
backend/tests/test_data_process_api.py
View File

@@ -1,13 +1,23 @@
from __future__ import annotations
from copy import deepcopy
from io import BytesIO
from pathlib import Path
from typing import Any
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from openpyxl import Workbook
from app.api.v1.endpoints import data_process as data_process_endpoint
from app.api.v1.endpoints.data_process import router
from app.modules.data_process.algorithms import normalize_text
from app.modules.data_process.storage import (
DataProcessStorageError,
LocalDataProcessStorage,
get_data_process_storage,
)
from app.modules.data_process.store import InvalidStateError, NotFoundError, get_data_process_store
@@ -87,11 +97,26 @@ class FakeDataProcessStore:
def add_source_file(self, task_id: str, **payload: Any) -> dict[str, Any]:
self.get_task(task_id)
values = deepcopy(payload)
source_id = str(values.pop("id", None) or self._id("dpsf"))
storage_object_id = str(
values.pop("storage_object_id", None)
or f"db://data-process/{task_id}/{source_id}/v1"
)
raw_size = int(values.pop("raw_size"))
metadata = deepcopy(values.pop("metadata", {}))
metadata.setdefault(
"storage_backend",
"local" if storage_object_id.startswith("local://data-process/") else "database",
)
source = {
"id": self._id("dpsf"),
"id": source_id,
"task_id": task_id,
"version_no": 1,
**deepcopy(payload),
"storage_object_id": storage_object_id,
"size_bytes": raw_size,
"metadata": metadata,
**values,
}
self.sources[task_id].append(source)
self.tasks[task_id]["input_count"] += payload["record_count"]
@@ -163,14 +188,27 @@ class FakeDataProcessStore:
self.results[task_id] = []
def replace_preview_items(
self, task_id: str, items: list[dict[str, Any]]
self,
task_id: str,
items: list[dict[str, Any]],
*,
source_file_ids: list[str] | None = None,
) -> list[dict[str, Any]]:
self.previews[task_id] = [
created = [
{"id": self._id("dpp"), "task_id": task_id, **deepcopy(item)} for item in items
]
if source_file_ids is None:
self.previews[task_id] = created
else:
selected = set(source_file_ids)
self.previews[task_id] = [
item
for item in self.previews[task_id]
if item["source_file_id"] not in selected
] + created
self.results[task_id] = []
self.tasks[task_id]["progress"] = 20
return deepcopy(self.previews[task_id])
return deepcopy(created)
def list_preview_items(
self,
@@ -389,16 +427,59 @@ class FakeDataProcessStore:
return {"dataset": deepcopy(dataset), "created": True}
def make_client() -> tuple[TestClient, FakeDataProcessStore]:
def make_client(
tmp_path: Path,
) -> tuple[TestClient, FakeDataProcessStore, LocalDataProcessStorage]:
store = FakeDataProcessStore()
storage = LocalDataProcessStorage(tmp_path / "data-process")
app = FastAPI()
app.include_router(router, prefix="/modelTF")
app.dependency_overrides[get_data_process_store] = lambda: store
return TestClient(app), store
app.dependency_overrides[get_data_process_storage] = lambda: storage
return TestClient(app), store, storage
def test_data_process_full_contract_without_database() -> None:
client, store = make_client()
def _stored_files(storage: LocalDataProcessStorage) -> list[Path]:
return [path for path in storage.root.rglob("*") if path.is_file() or path.is_symlink()]
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() + 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())
result.extend(value)
result.extend(b"\nendobj\n")
xref_offset = len(result)
result.extend(f"xref\n0 {len(objects) + 1}\n".encode())
result.extend(b"0000000000 65535 f \n")
for offset in offsets[1:]:
result.extend(f"{offset:010d} 00000 n \n".encode())
result.extend(
(
f"trailer\n<< /Size {len(objects) + 1} /Root 1 0 R >>\n"
f"startxref\n{xref_offset}\n%%EOF\n"
).encode()
)
return bytes(result)
def test_data_process_full_contract_without_database(tmp_path: Path) -> None:
client, store, _ = make_client(tmp_path)
created = client.post(
"/modelTF/data-process",
json={
@@ -506,8 +587,143 @@ def test_data_process_full_contract_without_database() -> None:
)
def test_external_source_never_returns_fake_success() -> None:
client, _ = make_client()
def test_preview_build_replaces_only_selected_files_and_reports_file_counts(
tmp_path: Path,
) -> None:
client, store, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "逐文件预览", "process_type": "structured", "config": {}},
).json()["data"]["id"]
uploaded = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files=[
("files", ("first.jsonl", b'{"id":1}\n', "application/jsonl")),
(
"files",
("second.jsonl", b'{"id":2}\n{"id":3}\n', "application/jsonl"),
),
],
)
assert uploaded.status_code == 200
first_source, second_source = uploaded.json()["data"]["files"]
first_build = client.post(
f"/modelTF/data-process/{task_id}/preview/build",
json={"source_file_ids": [first_source["id"]]},
)
assert first_build.status_code == 200
first_data = first_build.json()["data"]
assert first_data["file_counts"] == {first_source["id"]: 1}
assert first_data["files"] == [
{
"source_file_id": first_source["id"],
"preview_count": 1,
"status": "completed",
}
]
first_item = first_data["items"][0]
edited = client.put(
f"/modelTF/data-process/{task_id}/preview/{first_item['id']}",
json={"edited_content": "人工确认后的第一文件预览"},
)
assert edited.status_code == 200
second_build = client.post(
f"/modelTF/data-process/{task_id}/preview/build",
json={"source_file_id": second_source["id"]},
)
assert second_build.status_code == 200
second_data = second_build.json()["data"]
assert second_data["file_counts"] == {second_source["id"]: 2}
assert second_data["files"] == [
{
"source_file_id": second_source["id"],
"preview_count": 2,
"status": "completed",
}
]
assert {item["source_file_id"] for item in store.previews[task_id]} == {
first_source["id"],
second_source["id"],
}
preserved_first = next(
item
for item in store.previews[task_id]
if item["source_file_id"] == first_source["id"]
)
assert preserved_first["id"] == first_item["id"]
assert preserved_first["edited_content"] == "人工确认后的第一文件预览"
previous_second_ids = {
item["id"]
for item in store.previews[task_id]
if item["source_file_id"] == second_source["id"]
}
next(
source
for source in store.sources[task_id]
if source["id"] == second_source["id"]
)["content"] = '{"id":4}\n'
rebuilt = client.post(
f"/modelTF/data-process/{task_id}/preview/build",
json={"source_file_ids": [second_source["id"]]},
)
assert rebuilt.status_code == 200
assert rebuilt.json()["data"]["file_counts"] == {second_source["id"]: 1}
current_second_ids = {
item["id"]
for item in store.previews[task_id]
if item["source_file_id"] == second_source["id"]
}
assert current_second_ids.isdisjoint(previous_second_ids)
assert len(current_second_ids) == 1
assert next(
item
for item in store.previews[task_id]
if item["source_file_id"] == first_source["id"]
)["id"] == first_item["id"]
def test_preview_build_rejects_unknown_and_cross_task_source_file_ids(
tmp_path: Path,
) -> None:
client, _, _ = make_client(tmp_path)
first_task_id = client.post(
"/modelTF/data-process",
json={"name": "归属任务一", "process_type": "structured", "config": {}},
).json()["data"]["id"]
second_task_id = client.post(
"/modelTF/data-process",
json={"name": "归属任务二", "process_type": "structured", "config": {}},
).json()["data"]["id"]
foreign_source = client.post(
f"/modelTF/data-process/{second_task_id}/source-files",
files={"files": ("foreign.jsonl", b'{"id":2}\n', "application/jsonl")},
).json()["data"]["files"][0]
unknown = client.post(
f"/modelTF/data-process/{first_task_id}/preview/build",
json={"source_file_ids": ["dpsf_not_found"]},
)
assert unknown.status_code == 404
foreign = client.post(
f"/modelTF/data-process/{first_task_id}/preview/build",
json={"source_file_id": foreign_source["id"]},
)
assert foreign.status_code == 404
ambiguous = client.post(
f"/modelTF/data-process/{first_task_id}/preview/build",
json={
"source_file_id": foreign_source["id"],
"source_file_ids": [foreign_source["id"]],
},
)
assert ambiguous.status_code == 422
def test_external_source_never_returns_fake_success(tmp_path: Path) -> None:
client, _, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "外部数据", "process_type": "external", "config": {}},
@@ -520,8 +736,8 @@ def test_external_source_never_returns_fake_success() -> None:
assert response.json()["detail"]["code"] == 501
def test_config_validation_and_stop_state() -> None:
client, store = make_client()
def test_config_validation_and_stop_state(tmp_path: Path) -> None:
client, store, _ = make_client(tmp_path)
invalid = client.post(
"/modelTF/data-process",
json={
@@ -537,6 +753,28 @@ def test_config_validation_and_stop_state() -> None:
)
assert invalid.status_code == 422
legacy_semantic = client.post(
"/modelTF/data-process",
json={
"name": "旧切分策略",
"process_type": "unstructured",
"config": {"chunk_method": "semantic"},
},
)
assert legacy_semantic.status_code == 422
assert "chunk_method" in legacy_semantic.text
missing_custom_delimiter = client.post(
"/modelTF/data-process",
json={
"name": "缺少自定义分隔符",
"process_type": "unstructured",
"config": {"chunk_method": "custom"},
},
)
assert missing_custom_delimiter.status_code == 422
assert "custom_delimiter" in missing_custom_delimiter.text
task_id = client.post(
"/modelTF/data-process",
json={"name": "可停止任务", "process_type": "structured", "config": {}},
@@ -547,11 +785,11 @@ def test_config_validation_and_stop_state() -> None:
assert stopped.json()["data"]["status"] == "stopped"
def test_upload_batch_is_atomic_and_empty_files_are_rejected() -> None:
client, store = make_client()
def test_upload_batch_is_atomic_and_empty_files_are_rejected(tmp_path: Path) -> None:
client, store, storage = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "批量上传", "process_type": "structured", "config": {}},
json={"name": "批量上传", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
duplicate_batch = client.post(
@@ -563,6 +801,18 @@ def test_upload_batch_is_atomic_and_empty_files_are_rejected() -> None:
)
assert duplicate_batch.status_code == 400
assert store.sources[task_id] == []
assert _stored_files(storage) == []
parse_failure = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files=[
("files", ("valid.txt", "先暂存的内容".encode(), "text/plain")),
("files", ("broken.txt", b"\xff", "text/plain")),
],
)
assert parse_failure.status_code == 400
assert store.sources[task_id] == []
assert _stored_files(storage) == []
empty = client.post(
f"/modelTF/data-process/{task_id}/source-files",
@@ -570,10 +820,45 @@ def test_upload_batch_is_atomic_and_empty_files_are_rejected() -> None:
)
assert empty.status_code == 400
assert store.sources[task_id] == []
assert _stored_files(storage) == []
def test_preprocess_deduplicates_and_quality_filter_removes_short_results() -> None:
client, _ = make_client()
def test_upload_preserves_store_error_when_storage_rollback_fails(
tmp_path: Path,
monkeypatch: Any,
) -> None:
client, store, storage = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "回滚异常", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
cleanup_attempts: list[str] = []
def fail_store(*_: Any, **__: Any) -> list[dict[str, Any]]:
raise ValueError("simulated database transaction failure")
def fail_cleanup(reference: str, **_: Any) -> bool:
cleanup_attempts.append(reference)
raise OSError("simulated storage cleanup failure")
monkeypatch.setattr(store, "add_source_files", fail_store)
monkeypatch.setattr(storage, "delete", fail_cleanup)
response = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("rollback.txt", b"rollback payload", "text/plain")},
)
assert response.status_code == 400
assert response.json()["detail"]["message"] == "simulated database transaction failure"
assert len(cleanup_attempts) == 1
assert store.sources[task_id] == []
def test_preprocess_deduplicates_and_quality_filter_removes_short_results(
tmp_path: Path,
) -> None:
client, _, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={
@@ -651,8 +936,8 @@ def test_stale_generation_worker_cannot_overwrite_new_run(monkeypatch: Any) -> N
assert store.tasks[task_id]["status"] == "running"
def test_result_status_cannot_be_forged_by_client() -> None:
client, _ = make_client()
def test_result_status_cannot_be_forged_by_client(tmp_path: Path) -> None:
client, _, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "状态保护", "process_type": "structured", "config": {}},
@@ -664,8 +949,8 @@ def test_result_status_cannot_be_forged_by_client() -> None:
assert response.status_code == 422
def test_start_rebuilds_preview_and_generates_in_one_request() -> None:
client, _ = make_client()
def test_start_rebuilds_preview_and_generates_in_one_request(tmp_path: Path) -> None:
client, _, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "一键处理", "process_type": "structured", "config": {}},
@@ -691,14 +976,487 @@ def test_start_rebuilds_preview_and_generates_in_one_request() -> None:
assert client.get(f"/modelTF/data-process/{task_id}/results").json()["data"]["total"] == 1
def test_unsupported_upload_format_returns_415() -> None:
client, _ = make_client()
def test_unsupported_upload_format_returns_415(tmp_path: Path) -> None:
client, _, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "格式限制", "process_type": "structured", "config": {}},
).json()["data"]["id"]
response = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("document.pdf", b"not a pdf", "application/pdf")},
files={"files": ("payload.exe", b"not supported", "application/octet-stream")},
)
assert response.status_code == 415
legacy = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("document.doc", b"legacy", "application/msword")},
)
assert legacy.status_code == 415
assert "convert the file to .docx" in legacy.json()["detail"]["message"]
def test_xlsx_upload_is_accepted_as_structured_records(tmp_path: Path) -> None:
client, store, storage = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "XLSX 上传", "process_type": "structured", "config": {}},
).json()["data"]["id"]
workbook = Workbook()
worksheet = workbook.active
worksheet.append(["question", "answer"])
worksheet.append(["问题一", "答案一"])
worksheet.append(["问题二", "答案二"])
output = BytesIO()
workbook.save(output)
workbook.close()
original_bytes = output.getvalue()
response = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={
"files": (
"records.xlsx",
original_bytes,
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
)
},
)
assert response.status_code == 200
source = response.json()["data"]["files"][0]
assert source["file_format"] == "xlsx"
assert source["record_count"] == 2
assert source["size_bytes"] == len(original_bytes)
assert source["storage_object_id"].startswith("local://data-process/")
assert str(storage.root) not in response.text
assert storage.read(source["storage_object_id"]) == original_bytes
stored_source = store.get_source_file(task_id, source["id"])
assert stored_source["id"] == source["id"]
assert stored_source["storage_object_id"] == source["storage_object_id"]
assert stored_source["metadata"]["storage_backend"] == "local"
assert stored_source["metadata"]["original_size_bytes"] == len(original_bytes)
assert '"question":"问题一"' in stored_source["content"]
content = client.get(
f"/modelTF/data-process/{task_id}/source-files/{source['id']}/content"
)
assert content.status_code == 200
assert '"answer":"答案二"' in content.json()["data"]["content"]
preview = client.post(f"/modelTF/data-process/{task_id}/preview/build")
assert preview.status_code == 200
assert preview.json()["data"]["total"] == 2
def test_pdf_raw_preview_streams_original_file_and_supports_ranges(tmp_path: Path) -> None:
client, store, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "PDF 原件预览", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
original_pdf = _minimal_pdf()
uploaded = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("说明 文档.pdf", original_pdf, "application/pdf")},
).json()["data"]["files"][0]
raw_url = f"/modelTF/data-process/{task_id}/source-files/{uploaded['id']}/raw"
full = client.get(raw_url)
assert full.status_code == 200
assert full.content == original_pdf
assert full.headers["content-type"] == "application/pdf"
assert full.headers["accept-ranges"] == "bytes"
assert full.headers["cache-control"] == "private, no-store"
assert full.headers["content-length"] == str(len(original_pdf))
assert full.headers["content-disposition"].startswith("inline;")
assert "%E8%AF%B4%E6%98%8E%20%E6%96%87%E6%A1%A3.pdf" in full.headers[
"content-disposition"
]
assert full.headers["etag"] == f'"{uploaded["checksum_sha256"]}"'
partial = client.get(raw_url, headers={"Range": "bytes=5-14"})
assert partial.status_code == 206
assert partial.content == original_pdf[5:15]
assert partial.headers["content-range"] == f"bytes 5-14/{len(original_pdf)}"
assert partial.headers["content-length"] == "10"
suffix = client.get(raw_url, headers={"Range": "bytes=-8"})
assert suffix.status_code == 206
assert suffix.content == original_pdf[-8:]
invalid = client.get(raw_url, headers={"Range": "bytes=0-1,4-5"})
assert invalid.status_code == 416
assert invalid.headers["content-range"] == f"bytes */{len(original_pdf)}"
pages_url = (
f"/modelTF/data-process/{task_id}/source-files/{uploaded['id']}/pdf-pages"
)
pages = client.get(pages_url)
assert pages.status_code == 200
assert pages.json()["data"] == {
"page_count": 1,
"pages": [
{
"page_number": 1,
"source_start": 0,
"source_end": len("Hello PDF"),
}
],
}
legacy_id = "dpsf_legacy_pdf"
store.add_source_file(
task_id,
id=legacy_id,
storage_object_id=f"db://data-process/{task_id}/{legacy_id}/v1",
name="legacy.pdf",
content="legacy extracted PDF text",
raw_size=len(original_pdf),
checksum_sha256="a" * 64,
file_format="pdf",
record_count=1,
metadata={"legacy": True},
)
legacy = client.get(
f"/modelTF/data-process/{task_id}/source-files/{legacy_id}/raw"
)
assert legacy.status_code == 410
legacy_pages = client.get(
f"/modelTF/data-process/{task_id}/source-files/{legacy_id}/pdf-pages"
)
assert legacy_pages.status_code == 410
def test_raw_inline_preview_rejects_non_pdf_source(tmp_path: Path) -> None:
client, _, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "非 PDF 原件", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
uploaded = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("notes.txt", b"plain source text", "text/plain")},
).json()["data"]["files"][0]
response = client.get(
f"/modelTF/data-process/{task_id}/source-files/{uploaded['id']}/raw"
)
assert response.status_code == 415
pages_response = client.get(
f"/modelTF/data-process/{task_id}/source-files/{uploaded['id']}/pdf-pages"
)
assert pages_response.status_code == 415
def test_delete_source_removes_owned_local_object_and_accepts_legacy_db_reference(
tmp_path: Path,
) -> None:
client, store, storage = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "删除原件", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
uploaded = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("原件.txt", "本地原始内容".encode(), "text/plain")},
).json()["data"]["files"][0]
reference = uploaded["storage_object_id"]
assert storage.read(reference) == "本地原始内容".encode()
deleted = client.delete(
f"/modelTF/data-process/{task_id}/source-files/{uploaded['id']}"
)
assert deleted.status_code == 200
assert deleted.json()["data"]["storage_cleanup_pending"] is False
with pytest.raises(DataProcessStorageError, match="does not exist"):
storage.read(reference)
legacy_id = "dpsf_legacy"
store.add_source_file(
task_id,
id=legacy_id,
storage_object_id=f"db://data-process/{task_id}/{legacy_id}/v1",
name="legacy.txt",
content="旧记录正文",
raw_size=len("旧记录正文".encode()),
checksum_sha256="a" * 64,
file_format="txt",
record_count=1,
metadata={"legacy": True},
)
legacy_deleted = client.delete(
f"/modelTF/data-process/{task_id}/source-files/{legacy_id}"
)
assert legacy_deleted.status_code == 200
assert legacy_deleted.json()["data"]["storage_cleanup_pending"] is False
def test_delete_reports_pending_cleanup_after_database_soft_delete(
tmp_path: Path,
monkeypatch: Any,
) -> None:
client, store, storage = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "待清理原件", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
uploaded = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("pending.txt", b"pending cleanup", "text/plain")},
).json()["data"]["files"][0]
def fail_cleanup(*_: Any, **__: Any) -> bool:
raise OSError("simulated storage failure")
monkeypatch.setattr(storage, "delete", fail_cleanup)
response = client.delete(
f"/modelTF/data-process/{task_id}/source-files/{uploaded['id']}"
)
assert response.status_code == 200
assert response.json()["data"]["storage_cleanup_pending"] is True
with pytest.raises(NotFoundError):
store.get_source_file(task_id, uploaded["id"])
assert storage.read(uploaded["storage_object_id"]) == b"pending cleanup"
def test_delete_rejects_polluted_reference_owned_by_another_source(tmp_path: Path) -> None:
client, store, storage = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "归属校验", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
uploaded = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={"files": ("safe.txt", b"owned content", "text/plain")},
).json()["data"]["files"][0]
target_reference = uploaded["storage_object_id"]
polluted_id = "dpsf_polluted"
store.add_source_file(
task_id,
id=polluted_id,
storage_object_id=target_reference,
name="polluted.txt",
content="polluted",
raw_size=8,
checksum_sha256="b" * 64,
file_format="txt",
record_count=1,
metadata={},
)
rejected = client.delete(
f"/modelTF/data-process/{task_id}/source-files/{polluted_id}"
)
assert rejected.status_code == 400
assert storage.read(target_reference) == b"owned content"
assert store.get_source_file(task_id, polluted_id)["id"] == polluted_id
def test_upload_format_must_match_process_type(tmp_path: Path) -> None:
client, _, _ = make_client(tmp_path)
structured_id = client.post(
"/modelTF/data-process",
json={"name": "结构化格式约束", "process_type": "structured", "config": {}},
).json()["data"]["id"]
structured_pdf = client.post(
f"/modelTF/data-process/{structured_id}/source-files",
files={"files": ("manual.pdf", b"not parsed", "application/pdf")},
)
assert structured_pdf.status_code == 415
unstructured_id = client.post(
"/modelTF/data-process",
json={"name": "非结构化格式约束", "process_type": "unstructured", "config": {}},
).json()["data"]["id"]
unstructured_xlsx = client.post(
f"/modelTF/data-process/{unstructured_id}/source-files",
files={"files": ("records.xlsx", b"not parsed", "application/octet-stream")},
)
assert unstructured_xlsx.status_code == 415
external_id = client.post(
"/modelTF/data-process",
json={"name": "外部数据格式约束", "process_type": "external", "config": {}},
).json()["data"]["id"]
external_upload = client.post(
f"/modelTF/data-process/{external_id}/source-files",
files={"files": ("records.jsonl", b'{"id":1}', "application/jsonl")},
)
assert external_upload.status_code == 409
def _preview_task(
content: str,
*,
options: list[str],
config: dict[str, Any] | None = None,
source_id: str = "source-1",
file_format: str = "txt",
) -> list[dict[str, Any]]:
task_config = {
"preprocess_options": options,
"chunk_size": 200,
"chunk_overlap": 20,
"min_chunk_size": 20,
**(config or {}),
}
return data_process_endpoint._build_preview_items(
{"process_type": "unstructured", "config": task_config},
[
{
"id": source_id,
"name": f"{source_id}.{file_format}",
"file_format": file_format,
"content": content,
}
],
)
def test_default_and_structure_preview_split_headings_without_cross_section_overlap() -> None:
content = (
"# 第一章\n"
+ " ".join(f"alpha{index}" for index in range(18))
+ "\n# 第二章\n"
+ " ".join(f"beta{index}" for index in range(18))
)
normalized = normalize_text(content)
second_chapter_start = normalized.index("# 第二章")
common_config = {"chunk_size": 10, "chunk_overlap": 3, "min_chunk_size": 4}
default_items = _preview_task(
content,
options=["preserve_context"],
config=common_config,
)
structure_items = _preview_task(
content,
options=["preserve_context"],
config={**common_config, "chunk_method": "structure"},
)
def snapshot(items: list[dict[str, Any]]) -> list[tuple[Any, ...]]:
return [
(
item["original_content"],
item["source_start"],
item["source_end"],
item["source_start_line"],
item["source_end_line"],
)
for item in items
]
assert snapshot(default_items) == snapshot(structure_items)
assert all(
item["original_content"]
== normalized[item["source_start"] : item["source_end"]]
for item in structure_items
)
assert all(
not (item["source_start"] < second_chapter_start < item["source_end"])
for item in structure_items
)
second_chapter_items = [
item for item in structure_items if item["source_start"] >= second_chapter_start
]
assert second_chapter_items[0]["source_start"] == second_chapter_start
assert second_chapter_items[0]["original_content"].startswith("# 第二章")
def test_every_unstructured_preprocess_option_changes_preview_behavior() -> None:
repeated = "@" * 120
assert len(_preview_task(repeated, options=[])) == 1
assert _preview_task(repeated, options=["clean_invalid_content"]) == []
structured_text = "# 第一章\n" + "甲。" * 30 + "\n# 第二章\n" + "乙。" * 30
detected = _preview_task(
structured_text,
options=["detect_document_structure"],
config={"chunk_method": "fixed", "chunk_size": 20, "min_chunk_size": 5},
)
undetected = _preview_task(
structured_text,
options=[],
config={"chunk_method": "fixed", "chunk_size": 20, "min_chunk_size": 5},
)
assert all("heading_path" in item["quality_score"] for item in detected)
assert {tuple(item["quality_score"]["heading_path"]) for item in detected} == {
("第一章",),
("第二章",),
}
assert all("heading_path" not in item["quality_score"] for item in undetected)
assert all(not ("第一章" in item["edited_content"] and "第二章" in item["edited_content"]) for item in detected)
short_lead = "a b. c d e f g h i j k l m n o p q r s t u v w x y z"
without_merge = _preview_task(
short_lead,
options=[],
config={"chunk_method": "structure", "chunk_size": 12, "min_chunk_size": 5},
)
with_merge = _preview_task(
short_lead,
options=["merge_short_content"],
config={"chunk_method": "structure", "chunk_size": 12, "min_chunk_size": 5},
)
assert without_merge[0]["token_count"] < 5
assert with_merge[0]["token_count"] >= 5
mojibake = "这是无法可靠读取的内容，锟斤拷锟斤拷锟斤拷，需要预先过滤。"
assert len(_preview_task(mojibake, options=[])) == 1
assert _preview_task(mojibake, options=["filter_low_quality"]) == []
first = "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu nu xi omicron pi rho sigma tau upsilon phi chi psi omega"
second = "alpha beta gamma, delta epsilon zeta eta theta iota kappa lambda mu nu xi omicron pi rho sigma tau upsilon phi chi psi omega"
sources = [
{"id": "near-1", "name": "one.txt", "file_format": "txt", "content": first},
{"id": "near-2", "name": "two.txt", "file_format": "txt", "content": second},
]
base_task = {
"process_type": "unstructured",
"config": {
"chunk_method": "fixed",
"chunk_size": 200,
"chunk_overlap": 0,
"min_chunk_size": 1,
"preprocess_options": [],
},
}
assert len(data_process_endpoint._build_preview_items(base_task, sources)) == 2
deduplicated_task = deepcopy(base_task)
deduplicated_task["config"]["preprocess_options"] = ["deduplicate_content"]
assert len(data_process_endpoint._build_preview_items(deduplicated_task, sources)) == 1
context_text = " ".join(f"token{index}" for index in range(45))
no_context = _preview_task(
context_text,
options=[],
config={"chunk_method": "fixed", "chunk_size": 20, "chunk_overlap": 5},
)
with_context = _preview_task(
context_text,
options=["preserve_context"],
config={"chunk_method": "fixed", "chunk_size": 20, "chunk_overlap": 5},
)
assert no_context[1]["source_start"] >= no_context[0]["source_end"]
assert with_context[1]["source_start"] < with_context[0]["source_end"]
sensitive = "联系人：张三，手机 13800138000，邮箱 user@example.com。"
plain = _preview_task(sensitive, options=[])[0]
masked = _preview_task(sensitive, options=["desensitize"])[0]
assert "张三" in plain["edited_content"]
assert "联系人:[NAME]" in masked["edited_content"]
assert "[PHONE]" in masked["edited_content"]
assert "[EMAIL]" in masked["edited_content"]
def test_stored_binary_document_text_is_not_reparsed_as_binary() -> None:
for file_format in ("pdf", "docx", "pptx"):
items = _preview_task(
f"{file_format.upper()} 已抽取正文，可直接进入切片处理。",
options=[],
file_format=file_format,
)
assert len(items) == 1
assert "已抽取正文" in items[0]["edited_content"]

242
backend/tests/test_data_process_storage.py Normal file
View File

@@ -0,0 +1,242 @@
from __future__ import annotations
from pathlib import Path, PurePosixPath
import pytest
from app.modules.data_process import storage as storage_module
from app.modules.data_process.storage import (
DataProcessStorageError,
LocalDataProcessStorage,
StagedSourceObject,
)
def _stage(
storage: LocalDataProcessStorage,
*,
batch_id: str = "batch-main",
task_id: str = "task-1",
source_file_id: str = "source-1",
version: int = 1,
name: str = "source.txt",
content: bytes = b"payload",
) -> StagedSourceObject:
return storage.stage_bytes(
batch_id=batch_id,
task_id=task_id,
source_file_id=source_file_id,
version=version,
name=name,
content=content,
)
def _create_symlink(link: Path, target: Path, *, target_is_directory: bool = False) -> None:
try:
link.symlink_to(target, target_is_directory=target_is_directory)
except (NotImplementedError, OSError) as exc:
pytest.skip(f"当前平台不支持创建测试所需的符号链接: {exc}")
def _assert_staging_empty(storage: LocalDataProcessStorage) -> None:
assert list((storage.root / ".staging").iterdir()) == []
def test_stage_publish_read_delete_roundtrip_with_unicode_filename(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
content = "第一行\n第二行，100% 完成".encode()
staged = _stage(
storage,
name="中文 数据 100%.csv",
content=content,
)
assert "%20" in staged.reference
assert "%25" in staged.reference
storage.publish([staged])
assert storage.read(staged.reference) == content
assert storage.delete(staged.reference) is True
assert storage.delete(staged.reference) is False
_assert_staging_empty(storage)
def test_db_reference_is_left_to_database_storage(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
assert storage.read("db://source-files/source-1") is None
assert storage.delete("db://source-files/source-1") is False
def test_owned_source_can_be_streamed_by_byte_range(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
content = b"0123456789abcdef"
staged = _stage(storage, content=content)
storage.publish([staged])
assert storage.file_size(
staged.reference,
expected_task_id="task-1",
expected_source_file_id="source-1",
) == len(content)
assert b"".join(storage.iter_bytes(
staged.reference,
expected_task_id="task-1",
expected_source_file_id="source-1",
expected_size=len(content),
start=4,
length=6,
chunk_size=2,
)) == b"456789"
with pytest.raises(DataProcessStorageError, match="owner mismatch"):
storage.file_size(
staged.reference,
expected_task_id="another-task",
expected_source_file_id="source-1",
)
with pytest.raises(DataProcessStorageError, match="does not match metadata"):
b"".join(storage.iter_bytes(
staged.reference,
expected_task_id="task-1",
expected_source_file_id="source-1",
expected_size=len(content) + 1,
))
@pytest.mark.parametrize(
"reference",
[
"local://data-process/../source-1/v1/file.txt",
"local://data-process/task-1/source-1/v1/file%2Fname.txt",
"local://data-process/task-1/source-1/v1/file.txt?download=1",
"local://data-process/task-1/source-1/v1/file.txt#fragment",
"https://data-process/task-1/source-1/v1/file.txt",
],
ids=[
"parent-traversal",
"percent-encoded-slash",
"query",
"fragment",
"wrong-scheme",
],
)
def test_unsafe_references_are_rejected(tmp_path: Path, reference: str) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
with pytest.raises(DataProcessStorageError):
storage.read(reference)
with pytest.raises(DataProcessStorageError):
storage.delete(reference)
def test_publish_rejects_intermediate_directory_symlink(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
outside = tmp_path / "outside"
outside.mkdir()
staged = _stage(storage, task_id="linked-task")
_create_symlink(
storage.root / "linked-task",
outside,
target_is_directory=True,
)
with pytest.raises(DataProcessStorageError, match="symlink|non-directory"):
storage.publish([staged])
assert list(outside.iterdir()) == []
_assert_staging_empty(storage)
def test_target_symlink_is_never_followed_or_deleted(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
staged = _stage(storage, task_id="task-link", source_file_id="source-link")
outside_file = tmp_path / "outside.txt"
outside_file.write_bytes(b"outside sentinel")
final_path = storage.root.joinpath(*staged._relative_path.parts)
final_path.parent.mkdir(parents=True)
_create_symlink(final_path, outside_file)
with pytest.raises(DataProcessStorageError, match="already exists"):
storage.publish([staged])
with pytest.raises(DataProcessStorageError, match="regular file"):
storage.read(staged.reference)
with pytest.raises(DataProcessStorageError, match="non-regular"):
storage.delete(staged.reference)
assert final_path.is_symlink()
assert outside_file.read_bytes() == b"outside sentinel"
_assert_staging_empty(storage)
def test_publish_rolls_back_first_object_when_second_target_collides(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
existing = _stage(
storage,
batch_id="batch-existing",
source_file_id="source-existing",
content=b"existing content",
)
storage.publish([existing])
first = _stage(
storage,
batch_id="batch-new",
source_file_id="source-new",
content=b"must be rolled back",
)
colliding_second = _stage(
storage,
batch_id="batch-new",
source_file_id="source-existing",
content=b"must not replace existing content",
)
with pytest.raises(DataProcessStorageError, match="already exists"):
storage.publish([first, colliding_second])
with pytest.raises(DataProcessStorageError, match="does not exist"):
storage.read(first.reference)
assert storage.read(existing.reference) == b"existing content"
_assert_staging_empty(storage)
def test_publish_rejects_manually_forged_staged_object(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
temporary_path = storage.root / ".staging" / "batch-forged" / "forged.tmp"
temporary_path.parent.mkdir()
temporary_path.write_bytes(b"forged content")
relative_path = PurePosixPath("task-forged", "source-forged", "v1", "forged.txt")
forged = StagedSourceObject(
reference="local://data-process/task-forged/source-forged/v1/forged.txt",
_temporary_path=temporary_path,
_relative_path=relative_path,
)
with pytest.raises(DataProcessStorageError, match="was not issued"):
storage.publish([forged])
with pytest.raises(DataProcessStorageError, match="was not issued"):
storage.discard([forged])
assert temporary_path.read_bytes() == b"forged content"
assert not storage.root.joinpath(*relative_path.parts).exists()
def test_relative_storage_configuration_is_anchored_to_backend_root(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
relative_configuration = Path("relative-storage") / tmp_path.name
backend_root = Path(storage_module.__file__).resolve().parents[3]
monkeypatch.chdir(tmp_path)
monkeypatch.setenv("DATA_PROCESS_STORAGE_DIR", str(relative_configuration))
storage_module.get_data_process_storage.cache_clear()
try:
configured_root = storage_module._configured_storage_root()
assert configured_root == backend_root / relative_configuration
assert not configured_root.exists()
finally:
storage_module.get_data_process_storage.cache_clear()

54
backend/tests/test_data_process_store.py Normal file
View File

@@ -0,0 +1,54 @@
from __future__ import annotations
import pytest
from app.modules.data_process.store import (
DataProcessStoreError,
_source_storage_descriptor,
)
def test_source_storage_descriptor_accepts_owned_local_and_legacy_db_references() -> None:
task_id = "dpt_task"
source_file_id = "dpsf_source"
local_reference = (
f"local://data-process/{task_id}/{source_file_id}/v1/source%20100%25.csv"
)
reference, metadata = _source_storage_descriptor(
{
"storage_object_id": local_reference,
"metadata": {"storage_backend": "spoofed", "content_type": "text/csv"},
},
task_id,
source_file_id,
)
assert reference == local_reference
assert metadata == {"storage_backend": "local", "content_type": "text/csv"}
legacy_reference, legacy_metadata = _source_storage_descriptor(
{"metadata": {"legacy": True}},
task_id,
source_file_id,
)
assert legacy_reference == f"db://data-process/{task_id}/{source_file_id}/v1"
assert legacy_metadata == {"storage_backend": "database", "legacy": True}
@pytest.mark.parametrize(
"reference",
[
"local://data-process/dpt_other/dpsf_source/v1/source.txt",
"db://data-process/dpt_task/dpsf_other/v1",
"/var/tmp/source.txt",
],
)
def test_source_storage_descriptor_rejects_unowned_or_unsupported_references(
reference: str,
) -> None:
with pytest.raises(DataProcessStoreError):
_source_storage_descriptor(
{"storage_object_id": reference},
"dpt_task",
"dpsf_source",
)