pyflowx/tests/test_imaging.py

"""image_pipeline DAG 构造器测试.

验证:
* 生成的 Graph 拓扑正确 (任务数/依赖链).
* 实际执行后输出文件链完整.
* 未知操作抛 ValueError.
* convert 步骤改变扩展名.
* 自定义 output_dir 与 naming 模板.
"""

from __future__ import annotations

from pathlib import Path

import pytest

import pyflowx as px

pytest.importorskip("PIL")
from PIL import Image


def _make_test_image(path: Path, size: tuple[int, int] = (100, 80), color: str = "red") -> None:
    """创建测试图片."""
    Image.new("RGB", size, color).save(path)


class TestImagePipelineTopology:
    """image_pipeline 生成的 Graph 拓扑测试."""

    def test_single_step(self, tmp_path: Path) -> None:
        """单步流水线生成 1 个任务无依赖."""
        src = tmp_path / "src.png"
        _make_test_image(src, (100, 80))
        graph = px.image_pipeline(src, steps=[("resize", {"width": 50})])
        assert len(graph.all_specs()) == 1
        spec = next(iter(graph.all_specs().values()))
        assert spec.depends_on == ()

    def test_multi_step_chain(self, tmp_path: Path) -> None:
        """多步流水线形成链式依赖."""
        src = tmp_path / "src.png"
        _make_test_image(src, (100, 80))
        graph = px.image_pipeline(
            src,
            steps=[
                ("resize", {"width": 50}),
                ("watermark", {"text": "X"}),
                ("compress", {"quality": 80}),
            ],
        )
        specs = graph.all_specs()
        assert len(specs) == 3
        names = list(specs.keys())
        assert specs[names[0]].depends_on == ()
        assert specs[names[1]].depends_on == (names[0],)
        assert specs[names[2]].depends_on == (names[1],)

    def test_unknown_operation_raises(self, tmp_path: Path) -> None:
        """未知操作抛 ValueError."""
        src = tmp_path / "src.png"
        _make_test_image(src, (50, 50))
        with pytest.raises(ValueError, match="未知图片操作"):
            px.image_pipeline(src, steps=[("unknown_op", {})])


class TestImagePipelineExecution:
    """image_pipeline 实际执行测试."""

    def test_execute_resize_only(self, tmp_path: Path) -> None:
        """单步 resize 执行后输出文件存在且尺寸正确."""
        src = tmp_path / "src.png"
        _make_test_image(src, (100, 80))
        graph = px.image_pipeline(src, steps=[("resize", {"width": 50})])
        report = px.run(graph, strategy="sequential")
        assert report.success
        out = tmp_path / "src_resize.png"
        assert out.exists()
        with Image.open(out) as img:
            assert img.size[0] == 50

    def test_execute_chain(self, tmp_path: Path) -> None:
        """链式执行: resize → watermark → convert."""
        src = tmp_path / "src.png"
        _make_test_image(src, (200, 150))
        graph = px.image_pipeline(
            src,
            steps=[
                ("resize", {"width": 100}),
                ("watermark", {"text": "TEST"}),
                ("convert", {"format": "webp", "quality": 80}),
            ],
        )
        report = px.run(graph, strategy="sequential")
        assert report.success
        final_out = tmp_path / "src_resize_watermark_convert.webp"
        assert final_out.exists()
        with Image.open(final_out) as img:
            assert img.format == "WEBP"

    def test_execute_with_output_dir(self, tmp_path: Path) -> None:
        """自定义 output_dir 输出到指定目录."""
        src = tmp_path / "src.png"
        _make_test_image(src, (100, 80))
        out_dir = tmp_path / "output"
        graph = px.image_pipeline(
            src,
            steps=[("resize", {"width": 50})],
            output_dir=out_dir,
        )
        report = px.run(graph, strategy="sequential")
        assert report.success
        assert (out_dir / "src_resize.png").exists()

    def test_execute_convert_changes_extension(self, tmp_path: Path) -> None:
        """convert 步骤改变输出扩展名."""
        src = tmp_path / "src.png"
        _make_test_image(src, (50, 50))
        graph = px.image_pipeline(
            src,
            steps=[
                ("convert", {"format": "jpeg"}),
            ],
        )
        report = px.run(graph, strategy="sequential")
        assert report.success
        out = tmp_path / "src_convert.jpeg"
        assert out.exists()

    def test_execute_custom_naming(self, tmp_path: Path) -> None:
        """自定义 naming 模板."""
        src = tmp_path / "src.png"
        _make_test_image(src, (50, 50))
        graph = px.image_pipeline(
            src,
            steps=[("resize", {"width": 25})],
            naming="processed_{step}_{stem}{ext}",
        )
        report = px.run(graph, strategy="sequential")
        assert report.success
        out = tmp_path / "processed_resize_src.png"
        assert out.exists()

    def test_execute_flip_and_rotate(self, tmp_path: Path) -> None:
        """flip + rotate 组合执行."""
        src = tmp_path / "src.png"
        _make_test_image(src, (100, 80))
        graph = px.image_pipeline(
            src,
            steps=[
                ("flip", {"direction": "horizontal"}),
                ("rotate", {"degrees": 90, "expand": True}),
            ],
        )
        report = px.run(graph, strategy="sequential")
        assert report.success
        out = tmp_path / "src_flip_rotate.png"
        assert out.exists()