Files
pyflowx/tests/cli/test_pdftool.py
T
zhou fca6f17a5c chore: 升级项目适配 Python 3.10+ 并重构代码
1.  移除对 Python 3.8/3.9 的支持,更新 tox、pyproject.toml 配置
2.  替换 typing 导入为 collections.abc 标准库类型
3.  重构 CLI 系统从 argparse 迁移到 typer
4.  优化代码格式与类型注解,修复多处类型兼容问题
5.  更新依赖声明,移除 graphlib_backport 兼容包
2026-07-06 22:14:59 +08:00

490 lines
21 KiB
Python

"""Tests for cli.pdftool module."""
from __future__ import annotations
from pathlib import Path
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
from pyflowx.ops import pdftool as media
# ---------------------------------------------------------------------- #
# pdf_merge
# ---------------------------------------------------------------------- #
class TestPdfMerge:
"""Test pdf_merge function."""
def test_pdf_merge_files(self, tmp_path: Path) -> None:
"""Should merge PDF files."""
pytest.importorskip("pypdf")
input_files = [tmp_path / "input1.pdf", tmp_path / "input2.pdf"]
for f in input_files:
f.write_bytes(b"PDF content")
output_file = tmp_path / "merged.pdf"
with patch("pypdf.PdfReader"), patch("pypdf.PdfWriter") as mock_writer:
mock_writer_instance = MagicMock()
mock_writer.return_value = mock_writer_instance
media.pdf_merge(input_files, output_file)
assert mock_writer_instance.write.called
# ---------------------------------------------------------------------- #
# pdf_split
# ---------------------------------------------------------------------- #
class TestPdfSplit:
"""Test pdf_split function."""
def test_pdf_split_file(self, tmp_path: Path) -> None:
"""Should split PDF file."""
pytest.importorskip("pypdf")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_dir = tmp_path / "split"
with patch("pypdf.PdfReader") as mock_reader, patch("pypdf.PdfWriter"):
mock_reader_instance = MagicMock()
mock_reader.return_value = mock_reader_instance
mock_reader_instance.pages = [MagicMock()]
media.pdf_split(input_file, output_dir)
assert output_dir.exists()
# ---------------------------------------------------------------------- #
# pdf_compress
# ---------------------------------------------------------------------- #
class TestPdfCompress:
"""Test pdf_compress function."""
def test_pdf_compress_file(self, tmp_path: Path) -> None:
"""Should compress PDF file."""
pytest.importorskip("fitz")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_file = tmp_path / "compressed.pdf"
with patch("fitz.open") as mock_fitz_open:
mock_doc = MagicMock()
mock_fitz_open.return_value = mock_doc
# Mock save to actually create the file
def mock_save(*args: Any, **kwargs: Any):
output_file.write_bytes(b"Compressed PDF")
mock_doc.save = mock_save
media.pdf_compress(input_file, output_file)
assert output_file.exists()
# ---------------------------------------------------------------------- #
# pdf_extract_text
# ---------------------------------------------------------------------- #
class TestPdfExtractText:
"""Test pdf_extract_text function."""
def test_pdf_extract_text_file(self, tmp_path: Path) -> None:
"""Should extract text from PDF."""
pytest.importorskip("fitz")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_file = tmp_path / "output.txt"
with patch("fitz.open") as mock_fitz_open:
mock_doc = MagicMock()
mock_page = MagicMock()
mock_page.get_text.return_value = "Test text"
mock_doc.__iter__ = MagicMock(return_value=iter([mock_page]))
mock_fitz_open.return_value = mock_doc
media.pdf_extract_text(input_file, output_file)
assert output_file.exists()
# ---------------------------------------------------------------------- #
# pdf_extract_images
# ---------------------------------------------------------------------- #
class TestPdfExtractImages:
"""Test pdf_extract_images function."""
def test_pdf_extract_images_file(self, tmp_path: Path) -> None:
"""Should extract images from PDF."""
pytest.importorskip("fitz")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_dir = tmp_path / "images"
with patch("fitz.open") as mock_fitz_open:
mock_doc = MagicMock()
mock_page = MagicMock()
mock_page.get_images.return_value = [[0]]
mock_doc.__iter__ = MagicMock(return_value=iter([mock_page]))
mock_doc.extract_image.return_value = {"image": b"image data", "ext": "png"}
mock_fitz_open.return_value = mock_doc
media.pdf_extract_images(input_file, output_dir)
assert output_dir.exists()
# ---------------------------------------------------------------------- #
# pdf_add_watermark
# ---------------------------------------------------------------------- #
class TestPdfAddWatermark:
"""Test pdf_add_watermark function."""
def test_pdf_add_watermark_file(self, tmp_path: Path) -> None:
"""Should add watermark to PDF."""
pytest.importorskip("fitz")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_file = tmp_path / "watermarked.pdf"
with patch("fitz.open") as mock_fitz_open, patch("fitz.get_text_length") as mock_text_length:
mock_doc = MagicMock()
mock_page = MagicMock()
mock_page.rect = MagicMock(width=800, height=600)
mock_doc.__iter__ = MagicMock(return_value=iter([mock_page]))
mock_fitz_open.return_value = mock_doc
mock_text_length.return_value = 100
media.pdf_add_watermark(input_file, output_file)
assert mock_doc.save.called
# ---------------------------------------------------------------------- #
# pdf_rotate
# ---------------------------------------------------------------------- #
class TestPdfRotate:
"""Test pdf_rotate function."""
def test_pdf_rotate_file_90(self, tmp_path: Path) -> None:
"""Should rotate PDF by 90 degrees."""
pytest.importorskip("fitz")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_file = tmp_path / "rotated.pdf"
with patch("fitz.open") as mock_fitz_open:
mock_doc = MagicMock()
mock_page = MagicMock()
mock_doc.__iter__ = MagicMock(return_value=iter([mock_page]))
mock_fitz_open.return_value = mock_doc
media.pdf_rotate(input_file, output_file, rotation=90)
assert mock_doc.save.called
def test_pdf_rotate_file_180(self, tmp_path: Path) -> None:
"""Should rotate PDF by 180 degrees."""
pytest.importorskip("fitz")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_file = tmp_path / "rotated.pdf"
with patch("fitz.open") as mock_fitz_open:
mock_doc = MagicMock()
mock_page = MagicMock()
mock_doc.__iter__ = MagicMock(return_value=iter([mock_page]))
mock_fitz_open.return_value = mock_doc
media.pdf_rotate(input_file, output_file, rotation=180)
assert mock_doc.save.called
# ---------------------------------------------------------------------- #
# pdf_crop
# ---------------------------------------------------------------------- #
class TestPdfCrop:
"""Test pdf_crop function."""
def test_pdf_crop_file(self, tmp_path: Path) -> None:
"""Should crop PDF."""
pytest.importorskip("fitz")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_file = tmp_path / "cropped.pdf"
with patch("fitz.open") as mock_fitz_open, patch("fitz.Rect"):
mock_doc = MagicMock()
mock_page = MagicMock()
mock_page.rect = MagicMock(x0=0, y0=0, x1=800, y1=600)
mock_doc.__iter__ = MagicMock(return_value=iter([mock_page]))
mock_fitz_open.return_value = mock_doc
media.pdf_crop(input_file, output_file, margins=(10, 10, 10, 10))
assert mock_doc.save.called
# ---------------------------------------------------------------------- #
# pdf_info
# ---------------------------------------------------------------------- #
class TestPdfInfo:
"""Test pdf_info function."""
def test_pdf_info_file(self, tmp_path: Path) -> None:
"""Should show PDF info."""
pytest.importorskip("fitz")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
with patch("fitz.open") as mock_fitz_open:
mock_doc = MagicMock()
mock_doc.page_count = 10
mock_doc.metadata = {"title": "Test", "author": "Author"}
mock_fitz_open.return_value = mock_doc
media.pdf_info(input_file)
assert mock_fitz_open.called
# ---------------------------------------------------------------------- #
# pdf_ocr
# ---------------------------------------------------------------------- #
class TestPdfOcr:
"""Test pdf_ocr function."""
@pytest.mark.slow
def test_pdf_ocr_file(self, tmp_path: Path) -> None:
"""Should OCR PDF."""
pytest.importorskip("fitz")
pytest.importorskip("pytesseract")
pytest.importorskip("PIL")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_file = tmp_path / "ocr.pdf"
with (
patch("fitz.open") as mock_fitz_open,
patch("PIL.Image.frombytes"),
patch("pytesseract.image_to_string") as mock_ocr,
):
mock_doc = MagicMock()
mock_page = MagicMock()
mock_page.rect = MagicMock(width=800, height=600)
mock_doc.__iter__ = MagicMock(return_value=iter([mock_page]))
mock_fitz_open.return_value = mock_doc
mock_ocr.return_value = "OCR text"
media.pdf_ocr(input_file, output_file)
# Should complete OCR
# ---------------------------------------------------------------------- #
# pdf_repair
# ---------------------------------------------------------------------- #
class TestPdfRepair:
"""Test pdf_repair function."""
def test_pdf_repair_file(self, tmp_path: Path) -> None:
"""Should repair PDF."""
pytest.importorskip("fitz")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_file = tmp_path / "repaired.pdf"
with patch("fitz.open") as mock_fitz_open:
mock_doc = MagicMock()
mock_fitz_open.return_value = mock_doc
media.pdf_repair(input_file, output_file)
assert mock_doc.save.called
# ---------------------------------------------------------------------- #
# pdf_encrypt
# ---------------------------------------------------------------------- #
class TestPdfEncrypt:
"""Test pdf_encrypt function."""
def test_pdf_encrypt_file(self, tmp_path: Path) -> None:
"""Should encrypt PDF."""
pytest.importorskip("pypdf")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_file = tmp_path / "encrypted.pdf"
with patch("pypdf.PdfReader") as mock_reader, patch("pypdf.PdfWriter") as mock_writer:
mock_reader_instance = MagicMock()
mock_reader.return_value = mock_reader_instance
mock_reader_instance.pages = [MagicMock()]
mock_writer_instance = MagicMock()
mock_writer.return_value = mock_writer_instance
media.pdf_encrypt(input_file, output_file, "secret")
assert mock_writer_instance.encrypt.called
assert mock_writer_instance.write.called
def test_pdf_encrypt_not_installed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Should print message when pypdf not installed."""
monkeypatch.setattr(media, "HAS_PYPDF", False)
media.pdf_encrypt(tmp_path / "in.pdf", tmp_path / "out.pdf", "pw")
# ---------------------------------------------------------------------- #
# pdf_decrypt
# ---------------------------------------------------------------------- #
class TestPdfDecrypt:
"""Test pdf_decrypt function."""
def test_pdf_decrypt_encrypted(self, tmp_path: Path) -> None:
"""Should decrypt encrypted PDF."""
pytest.importorskip("pypdf")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_file = tmp_path / "decrypted.pdf"
with patch("pypdf.PdfReader") as mock_reader, patch("pypdf.PdfWriter") as mock_writer:
mock_reader_instance = MagicMock()
mock_reader.return_value = mock_reader_instance
mock_reader_instance.is_encrypted = True
mock_reader_instance.pages = [MagicMock()]
mock_writer_instance = MagicMock()
mock_writer.return_value = mock_writer_instance
media.pdf_decrypt(input_file, output_file, "secret")
assert mock_reader_instance.decrypt.called
assert mock_writer_instance.write.called
def test_pdf_decrypt_not_encrypted(self, tmp_path: Path) -> None:
"""Should handle non-encrypted PDF."""
pytest.importorskip("pypdf")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_file = tmp_path / "decrypted.pdf"
with patch("pypdf.PdfReader") as mock_reader, patch("pypdf.PdfWriter") as mock_writer:
mock_reader_instance = MagicMock()
mock_reader.return_value = mock_reader_instance
mock_reader_instance.is_encrypted = False
mock_reader_instance.pages = [MagicMock()]
mock_writer_instance = MagicMock()
mock_writer.return_value = mock_writer_instance
media.pdf_decrypt(input_file, output_file, "secret")
assert not mock_reader_instance.decrypt.called
def test_pdf_decrypt_not_installed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Should print message when pypdf not installed."""
monkeypatch.setattr(media, "HAS_PYPDF", False)
media.pdf_decrypt(tmp_path / "in.pdf", tmp_path / "out.pdf", "pw")
# ---------------------------------------------------------------------- #
# pdf_reorder
# ---------------------------------------------------------------------- #
class TestPdfReorder:
"""Test pdf_reorder function."""
def test_pdf_reorder_pages(self, tmp_path: Path) -> None:
"""Should reorder PDF pages."""
pytest.importorskip("pypdf")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_file = tmp_path / "reordered.pdf"
with patch("pypdf.PdfReader") as mock_reader, patch("pypdf.PdfWriter") as mock_writer:
mock_reader_instance = MagicMock()
mock_reader.return_value = mock_reader_instance
mock_reader_instance.pages = [MagicMock(), MagicMock(), MagicMock()]
mock_writer_instance = MagicMock()
mock_writer.return_value = mock_writer_instance
media.pdf_reorder(input_file, output_file, [2, 0, 1])
assert mock_writer_instance.add_page.call_count == 3
def test_pdf_reorder_out_of_range(self, tmp_path: Path) -> None:
"""Should skip out-of-range page numbers."""
pytest.importorskip("pypdf")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_file = tmp_path / "reordered.pdf"
with patch("pypdf.PdfReader") as mock_reader, patch("pypdf.PdfWriter") as mock_writer:
mock_reader_instance = MagicMock()
mock_reader.return_value = mock_reader_instance
mock_reader_instance.pages = [MagicMock()]
mock_writer_instance = MagicMock()
mock_writer.return_value = mock_writer_instance
media.pdf_reorder(input_file, output_file, [0, 5, -1])
assert mock_writer_instance.add_page.call_count == 1
def test_pdf_reorder_not_installed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Should print message when pypdf not installed."""
monkeypatch.setattr(media, "HAS_PYPDF", False)
media.pdf_reorder(tmp_path / "in.pdf", tmp_path / "out.pdf", [0])
# ---------------------------------------------------------------------- #
# pdf_to_images
# ---------------------------------------------------------------------- #
class TestPdfToImages:
"""Test pdf_to_images function."""
def test_pdf_to_images_convert(self, tmp_path: Path) -> None:
"""Should convert PDF to images."""
pytest.importorskip("fitz")
input_file = tmp_path / "input.pdf"
input_file.write_bytes(b"PDF content")
output_dir = tmp_path / "images"
with patch("fitz.open") as mock_fitz_open:
mock_doc = MagicMock()
mock_page = MagicMock()
mock_pixmap = MagicMock()
mock_page.get_pixmap.return_value = mock_pixmap
mock_doc.__iter__ = MagicMock(return_value=iter([mock_page]))
mock_fitz_open.return_value = mock_doc
media.pdf_to_images(input_file, output_dir, dpi=150)
assert output_dir.exists()
assert mock_pixmap.save.called
def test_pdf_to_images_not_installed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Should print message when PyMuPDF not installed."""
monkeypatch.setattr(media, "HAS_PYMUPDF", False)
media.pdf_to_images(tmp_path / "in.pdf", tmp_path / "out")
# ---------------------------------------------------------------------- #
# Not-installed branches for already-tested functions
# ---------------------------------------------------------------------- #
class TestNotInstalledBranches:
"""Test 'not installed' branches for PDF functions."""
def test_pdf_merge_not_installed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Should print message when pypdf not installed."""
monkeypatch.setattr(media, "HAS_PYPDF", False)
media.pdf_merge([tmp_path / "a.pdf"], tmp_path / "out.pdf")
def test_pdf_split_not_installed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Should print message when pypdf not installed."""
monkeypatch.setattr(media, "HAS_PYPDF", False)
media.pdf_split(tmp_path / "in.pdf", tmp_path / "out")
def test_pdf_compress_not_installed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Should print message when PyMuPDF not installed."""
monkeypatch.setattr(media, "HAS_PYMUPDF", False)
media.pdf_compress(tmp_path / "in.pdf", tmp_path / "out.pdf")
def test_pdf_extract_text_not_installed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Should print message when PyMuPDF not installed."""
monkeypatch.setattr(media, "HAS_PYMUPDF", False)
media.pdf_extract_text(tmp_path / "in.pdf", tmp_path / "out.txt")
def test_pdf_extract_images_not_installed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Should print message when PyMuPDF not installed."""
monkeypatch.setattr(media, "HAS_PYMUPDF", False)
media.pdf_extract_images(tmp_path / "in.pdf", tmp_path / "out")
def test_pdf_add_watermark_not_installed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Should print message when PyMuPDF not installed."""
monkeypatch.setattr(media, "HAS_PYMUPDF", False)
media.pdf_add_watermark(tmp_path / "in.pdf", tmp_path / "out.pdf")
def test_pdf_rotate_not_installed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Should print message when PyMuPDF not installed."""
monkeypatch.setattr(media, "HAS_PYMUPDF", False)
media.pdf_rotate(tmp_path / "in.pdf", tmp_path / "out.pdf")
def test_pdf_crop_not_installed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Should print message when PyMuPDF not installed."""
monkeypatch.setattr(media, "HAS_PYMUPDF", False)
media.pdf_crop(tmp_path / "in.pdf", tmp_path / "out.pdf", (1, 1, 1, 1))
def test_pdf_info_not_installed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Should print message when PyMuPDF not installed."""
monkeypatch.setattr(media, "HAS_PYMUPDF", False)
media.pdf_info(tmp_path / "in.pdf")
def test_pdf_repair_not_installed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
"""Should print message when PyMuPDF not installed."""
monkeypatch.setattr(media, "HAS_PYMUPDF", False)
media.pdf_repair(tmp_path / "in.pdf", tmp_path / "out.pdf")