93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""tools/builtin_tools/pdf_reader.py 的单元测试。
|
|
|
|
用 PyMuPDF 现场生成含文本与内嵌图片的测试 PDF，验证文本结构化提取、
|
|
图片提取落地，以及异常分支（文件不存在 / 超过大小上限）。
|
|
|
|
需在装有 PyMuPDF 的 haocode 环境运行::
|
|
|
|
python -m unittest discover tests
|
|
"""
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
|
|
# 让 tests 目录能 import 到项目根目录下的 tools 包
|
|
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
|
if ROOT not in sys.path:
|
|
sys.path.insert(0, ROOT)
|
|
|
|
try:
|
|
import pymupdf
|
|
except ImportError: # 兼容旧版导入名
|
|
import fitz as pymupdf # type: ignore
|
|
|
|
from tools.builtin_tools.pdf_reader import (
|
|
extract_pdf_text,
|
|
extract_pdf_images,
|
|
)
|
|
|
|
|
|
def _make_png_bytes() -> bytes:
|
|
"""生成一张 8x8 红色小图的 PNG 字节流。"""
|
|
pix = pymupdf.Pixmap(pymupdf.csRGB, pymupdf.IRect(0, 0, 8, 8), 0)
|
|
pix.set_rect(pix.irect, (255, 0, 0))
|
|
data = pix.tobytes("png")
|
|
pix = None
|
|
return data
|
|
|
|
|
|
def _build_sample_pdf(path: str) -> None:
|
|
"""造一个 2 页 PDF：第 1 页含文本 + 图片，第 2 页仅文本。"""
|
|
doc = pymupdf.open()
|
|
p1 = doc.new_page()
|
|
p1.insert_text((72, 72), "Hello PDF page one")
|
|
p1.insert_image(pymupdf.Rect(72, 100, 172, 200), stream=_make_png_bytes())
|
|
p2 = doc.new_page()
|
|
p2.insert_text((72, 72), "Second page text here")
|
|
doc.save(path)
|
|
doc.close()
|
|
|
|
|
|
class PdfReaderTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.tmp.cleanup)
|
|
self.dir = self.tmp.name
|
|
self.pdf = os.path.join(self.dir, "sample.pdf")
|
|
_build_sample_pdf(self.pdf)
|
|
|
|
def test_extract_text_structure(self):
|
|
text, pages = extract_pdf_text(self.pdf)
|
|
self.assertEqual(pages, 2)
|
|
self.assertIn("[第 1 页]", text)
|
|
self.assertIn("[第 2 页]", text)
|
|
self.assertIn("Hello PDF page one", text)
|
|
self.assertIn("Second page text here", text)
|
|
|
|
def test_extract_images(self):
|
|
out_dir = os.path.join(self.dir, "imgs")
|
|
imgs = extract_pdf_images(self.pdf, out_dir)
|
|
self.assertGreaterEqual(len(imgs), 1)
|
|
im = imgs[0]
|
|
for key in ("page", "index", "abs_path", "mime", "size_kb", "width", "height"):
|
|
self.assertIn(key, im)
|
|
self.assertEqual(im["page"], 1)
|
|
self.assertEqual(im["mime"], "image/png")
|
|
self.assertTrue(os.path.isfile(im["abs_path"]))
|
|
self.assertGreater(os.path.getsize(im["abs_path"]), 0)
|
|
|
|
def test_missing_file(self):
|
|
with self.assertRaises(ValueError):
|
|
extract_pdf_text(os.path.join(self.dir, "nope.pdf"))
|
|
|
|
def test_oversize(self):
|
|
# 用一个极小的上限触发超大分支
|
|
with self.assertRaises(ValueError):
|
|
extract_pdf_text(self.pdf, max_bytes=10)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|