73 lines
2.9 KiB
Python
73 lines
2.9 KiB
Python
|
import tempfile
|
||
|
|
import unittest
|
||
|
|
from pathlib import Path
|
||
|
|
from zipfile import ZIP_DEFLATED, ZipFile
|
||
|
|
|
||
|
|
from pypdf import PdfWriter
|
||
|
|
from pypdf.generic import DecodedStreamObject, DictionaryObject, NameObject
|
||
|
|
|
||
|
|
from rag.document_parser import parse_document
|
||
|
|
|
||
|
|
|
||
|
|
class DocumentParserTests(unittest.TestCase):
|
||
|
|
def test_parses_txt_and_md_as_utf8_text(self):
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
root = Path(tmp)
|
||
|
|
txt = root / "notice.txt"
|
||
|
|
md = root / "notice.md"
|
||
|
|
txt.write_text("基金风险提示", encoding="utf-8")
|
||
|
|
md.write_text("# 产品说明\n\n开放式基金", encoding="utf-8")
|
||
|
|
|
||
|
|
self.assertEqual(parse_document(txt), "基金风险提示")
|
||
|
|
self.assertEqual(parse_document(md), "# 产品说明\n\n开放式基金")
|
||
|
|
|
||
|
|
def test_extracts_text_from_docx(self):
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
docx = Path(tmp) / "notice.docx"
|
||
|
|
document_xml = (
|
||
|
|
'<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
|
||
|
|
'<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">'
|
||
|
|
'<w:body><w:p><w:r><w:t>基金产品说明</w:t></w:r></w:p>'
|
||
|
|
'<w:p><w:r><w:t>风险揭示</w:t></w:r></w:p></w:body></w:document>'
|
||
|
|
)
|
||
|
|
with ZipFile(docx, "w", ZIP_DEFLATED) as archive:
|
||
|
|
archive.writestr("word/document.xml", document_xml)
|
||
|
|
|
||
|
|
self.assertEqual(parse_document(docx), "基金产品说明\n风险揭示")
|
||
|
|
|
||
|
|
def test_extracts_text_from_pdf(self):
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
pdf = Path(tmp) / "notice.pdf"
|
||
|
|
writer = PdfWriter()
|
||
|
|
page = writer.add_blank_page(width=612, height=792)
|
||
|
|
font = writer._add_object(
|
||
|
|
DictionaryObject(
|
||
|
|
{
|
||
|
|
NameObject("/Type"): NameObject("/Font"),
|
||
|
|
NameObject("/Subtype"): NameObject("/Type1"),
|
||
|
|
NameObject("/BaseFont"): NameObject("/Helvetica"),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
)
|
||
|
|
page[NameObject("/Resources")] = DictionaryObject(
|
||
|
|
{NameObject("/Font"): DictionaryObject({NameObject("/F1"): font})}
|
||
|
|
)
|
||
|
|
page[NameObject("/Contents")] = DecodedStreamObject()
|
||
|
|
page[NameObject("/Contents")].set_data(b"BT /F1 12 Tf 72 720 Td (Fund FAQ) Tj ET")
|
||
|
|
with pdf.open("wb") as output:
|
||
|
|
writer.write(output)
|
||
|
|
|
||
|
|
self.assertEqual(parse_document(pdf), "Fund FAQ")
|
||
|
|
|
||
|
|
def test_rejects_unsupported_extension(self):
|
||
|
|
with tempfile.TemporaryDirectory() as tmp:
|
||
|
|
path = Path(tmp) / "notice.xlsx"
|
||
|
|
path.write_bytes(b"not supported")
|
||
|
|
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
parse_document(path)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|