62 lines
1.7 KiB
Python
62 lines
1.7 KiB
Python
"""
|
|
Excel / PDF 抽纯文本：不猜字段。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import os
|
|
import sys
|
|
import unittest
|
|
|
|
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
|
if ROOT not in sys.path:
|
|
sys.path.insert(0, ROOT)
|
|
|
|
from agent.policy.attachment_extract import allowed_attachment, extract_attachment_text
|
|
|
|
|
|
def _xlsx_bytes() -> bytes:
|
|
from openpyxl import Workbook
|
|
|
|
wb = Workbook()
|
|
ws = wb.active
|
|
ws.title = "询价"
|
|
ws["A1"] = "海运"
|
|
ws["B1"] = "上海到洛杉矶"
|
|
ws["A2"] = "1x40HQ"
|
|
buf = io.BytesIO()
|
|
wb.save(buf)
|
|
return buf.getvalue()
|
|
|
|
|
|
class ExtractAttachmentTests(unittest.TestCase):
|
|
def test_allow_only_excel_pdf(self) -> None:
|
|
self.assertTrue(allowed_attachment("a.XLSX"))
|
|
self.assertTrue(allowed_attachment("b.xls"))
|
|
self.assertTrue(allowed_attachment("c.pdf"))
|
|
self.assertFalse(allowed_attachment("d.png"))
|
|
self.assertFalse(allowed_attachment("e.docx"))
|
|
self.assertFalse(allowed_attachment("f.zip"))
|
|
|
|
def test_reject_image_bytes(self) -> None:
|
|
got = extract_attachment_text("shot.png", b"\x89PNG")
|
|
self.assertFalse(got["ok"])
|
|
self.assertEqual(got["error"], "only_xlsx_xls_pdf")
|
|
|
|
def test_xlsx_cells_become_text(self) -> None:
|
|
got = extract_attachment_text("sea.xlsx", _xlsx_bytes())
|
|
self.assertTrue(got["ok"])
|
|
self.assertIn("海运", got["text"])
|
|
self.assertIn("上海到洛杉矶", got["text"])
|
|
self.assertIn("1x40HQ", got["text"])
|
|
|
|
def test_empty_bytes_fail(self) -> None:
|
|
got = extract_attachment_text("a.xlsx", b"")
|
|
self.assertFalse(got["ok"])
|
|
self.assertTrue(got["empty"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|