from __future__ import annotations import shutil from pathlib import Path import pytest from imagepipeline.core.pipeline import Pipeline from imagepipeline.core.resume import ( expected_output_filenames, materialize_external_outputs, step_outputs_complete, ) from imagepipeline.core.step import StepDefinition from imagepipeline.modules.imagemagick_grayscale import ImageMagickGrayscale from imagepipeline.modules.registry import get_module from imagepipeline.modules.rembg import RembgModule from imagepipeline.utils.gmic import finalize_gmic_output, split_gmic_command from tests.conftest import make_png class TestGmicCommandSplit: def test_splits_command_and_arguments(self) -> None: assert split_gmic_command("-gcd_stereo_img 0,0,2.028,1,1.714,3.06,4,1,0") == [ "-gcd_stereo_img", "0,0,2.028,1,1.714,3.06,4,1,0", ] def test_preserves_quoted_empty_argument(self) -> None: parts = split_gmic_command('-fx_custom_gradient 0,0,0,"",1,0') assert parts == ["-fx_custom_gradient", "0,0,0,,1,0"] class TestFinalizeGmicOutput: def test_keeps_frame_000001_and_removes_000000(self, tmp_path: Path) -> None: output_dir = tmp_path / "out" output_dir.mkdir() intended = output_dir / "photo.png" frame_000000 = output_dir / "photo_000000.png" frame_000001 = output_dir / "photo_000001.png" frame_000000.write_bytes(b"discard") frame_000001.write_bytes(b"keep") result = finalize_gmic_output(output_dir, intended) assert result == intended assert intended.read_bytes() == b"keep" assert not frame_000000.exists() assert not frame_000001.exists() def test_leaves_single_output_unchanged(self, tmp_path: Path) -> None: output_dir = tmp_path / "out" output_dir.mkdir() intended = output_dir / "photo.png" intended.write_bytes(b"single") result = finalize_gmic_output(output_dir, intended) assert result == intended assert intended.read_bytes() == b"single" def test_renames_only_000000_when_000001_missing(self, tmp_path: Path) -> None: output_dir = tmp_path / "out" output_dir.mkdir() intended = output_dir / "photo.png" frame_000000 = output_dir / "photo_000000.png" frame_000000.write_bytes(b"only") result = finalize_gmic_output(output_dir, intended) assert result == intended assert intended.read_bytes() == b"only" assert not frame_000000.exists() class TestExpectedOutputFilenames: def test_rembg_maps_jpg_inputs_to_png_outputs(self, tmp_path: Path) -> None: jpg = tmp_path / "photo.jpg" jpg.write_bytes(b"jpeg") step = StepDefinition( step_id="rembg_01", module_name="rembg", module=RembgModule, input_refs=["input"], params={}, output_dir_name="rembg_01", ) names = expected_output_filenames( step, matched_groups=[[jpg]], input_paths=[jpg], params=RembgModule.validate_module_params({}), ) assert names == ["photo.png"] def test_rembg_resume_detects_existing_png_outputs(self, tmp_path: Path) -> None: output_dir = tmp_path / "rembg_01" output_dir.mkdir() png = output_dir / "photo.png" make_png(png) jpg = tmp_path / "input" / "photo.jpg" jpg.parent.mkdir() jpg.write_bytes(b"jpeg") step = StepDefinition( step_id="rembg_01", module_name="rembg", module=RembgModule, input_refs=["input"], params={}, output_dir_name="rembg_01", ) params = RembgModule.validate_module_params({}) expected = expected_output_filenames( step, matched_groups=[[jpg]], input_paths=[jpg], params=params, ) assert step_outputs_complete([output_dir / name for name in expected]) class TestPipelineResume: @pytest.mark.skipif(not shutil.which("magick"), reason="ImageMagick not installed") def test_continue_skips_completed_steps(self, input_dir: Path, output_base: Path, capsys) -> None: with Pipeline( name="resume_test", input_dir=input_dir, output_base=output_base, verbose=True, ) as p: first = p.step("imagemagick_grayscale", inputs="input") p.step("imagemagick_grayscale", inputs=first) root = p.run() capsys.readouterr() with Pipeline( name="resume_test", input_dir=input_dir, output_base=output_base, verbose=True, continue_from=root, ) as p: step_a = p.step("imagemagick_grayscale", inputs="input") p.step("imagemagick_grayscale", inputs=step_a) resumed_root = p.run() assert resumed_root == root output = capsys.readouterr().out assert "Skipped step imagemagick_grayscale_01" in output assert "Skipped step imagemagick_grayscale_02" in output @pytest.mark.skipif(not shutil.which("magick"), reason="ImageMagick not installed") def test_existing_outputs_reuse_external_folder( self, input_dir: Path, output_base: Path, tmp_path: Path, capsys ) -> None: external = tmp_path / "external_rembg" external.mkdir() for src in input_dir.iterdir(): if src.is_file(): make_png(external / src.name, width=4, height=4, rgb=(10, 20, 30)) with Pipeline( name="external_test", input_dir=input_dir, output_base=output_base, verbose=True, existing_outputs={"imagemagick_grayscale_01": external}, ) as p: reused = p.step("imagemagick_grayscale", inputs="input") p.step("imagemagick_grayscale", inputs=reused) root = p.run() output = capsys.readouterr().out assert "Reused external output for imagemagick_grayscale_01" in output assert (root / "imagemagick_grayscale_01" / "photo_a.png").exists() @pytest.mark.skipif(not shutil.which("magick"), reason="ImageMagick not installed") def test_existing_outputs_with_custom_step_id( self, input_dir: Path, output_base: Path, tmp_path: Path, capsys ) -> None: external = tmp_path / "external_gray" external.mkdir() for src in input_dir.iterdir(): if src.is_file(): make_png(external / src.name, width=4, height=4, rgb=(10, 20, 30)) with Pipeline( name="custom_id_external", input_dir=input_dir, output_base=output_base, verbose=True, existing_outputs={"input_bokeh": external}, ) as p: reused = p.step( "imagemagick_grayscale", inputs="input", step_id="input_bokeh" ) p.step("imagemagick_grayscale", inputs=reused) root = p.run() output = capsys.readouterr().out assert "Reused external output for input_bokeh" in output assert (root / "input_bokeh" / "photo_a.png").exists() assert not (root / "imagemagick_grayscale_01").exists() class TestMaterializeExternal: def test_links_files_by_stem(self, tmp_path: Path) -> None: external = tmp_path / "external" external.mkdir() make_png(external / "photo_a.png") output_dir = tmp_path / "out" step = StepDefinition( step_id="imagemagick_grayscale_01", module_name="imagemagick_grayscale", module=ImageMagickGrayscale, input_refs=["input"], params={}, output_dir_name="imagemagick_grayscale_01", ) input_paths = [tmp_path / "photo_a.png"] paths = materialize_external_outputs( external, output_dir, step, matched_groups=[[path] for path in input_paths], input_paths=input_paths, params={}, ) assert len(paths) == 1 assert paths[0].name == "photo_a.png" assert paths[0].is_symlink() def test_step_outputs_complete(self, tmp_path: Path) -> None: output_dir = tmp_path / "done" output_dir.mkdir() make_png(output_dir / "photo.png") assert step_outputs_complete([output_dir / "photo.png"]) assert not step_outputs_complete([output_dir / "missing.png"])