tests/test_pipeline_utils.py

import json
import tempfile
import unittest
from pathlib import Path

from PIL import Image

from approach.pipeline_utils import (
    enrich_ape_results,
    output_path_for_selection,
    parse_orienter_image_name,
    resolve_image_path,
    run_optional_reflection,
    select_jsonl_lines,
    write_json_atomic,
)


class PipelineUtilsTest(unittest.TestCase):
    def test_parse_orienter_image_name_is_strict(self):
        self.assertEqual(parse_orienter_image_name("123_4.jpg"), ("123", "4", 123004))
        with self.assertRaises(ValueError):
            parse_orienter_image_name("123_4_extra.jpg")
        with self.assertRaises(ValueError):
            parse_orienter_image_name("nested/123_4.jpg")
        with self.assertRaises(ValueError):
            parse_orienter_image_name("123_1000.jpg")

    def test_resolve_image_path_rejects_escape_from_images_directory(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            images_dir = Path(tmp_dir) / "images"
            images_dir.mkdir()
            self.assertEqual(
                resolve_image_path(images_dir, "nested/123_4.jpg"),
                (images_dir / "nested" / "123_4.jpg").resolve(),
            )
            with self.assertRaises(ValueError):
                resolve_image_path(images_dir, "../123_4.jpg")
            with self.assertRaises(ValueError):
                resolve_image_path(images_dir, str(Path(tmp_dir) / "123_4.jpg"))

    def test_select_jsonl_lines_supports_ranges_or_shards(self):
        lines = [f"{idx}\n" for idx in range(10)]
        self.assertEqual(select_jsonl_lines(lines, start_index=2, end_index=8), lines[2:8])
        self.assertEqual(
            select_jsonl_lines(lines, shard_index=1, num_shards=2),
            ["1\n", "3\n", "5\n", "7\n", "9\n"],
        )
        with self.assertRaises(ValueError):
            select_jsonl_lines(
                lines,
                start_index=2,
                end_index=8,
                shard_index=1,
                num_shards=2,
            )

    def test_output_path_for_selection_prevents_shard_overwrite(self):
        self.assertEqual(
            output_path_for_selection("predictions.json", shard_index=1, num_shards=4),
            "predictions.shard01-of-04.json",
        )
        self.assertEqual(
            output_path_for_selection("predictions.json", start_index=100, end_index=200),
            "predictions.rows100-200.json",
        )
        self.assertEqual(output_path_for_selection("predictions.json"), "predictions.json")

    def test_enrich_ape_results_keeps_added_metadata(self):
        enriched = enrich_ape_results(
            [{"bbox": [1, 2, 3, 4], "category_name": "button"}],
            "123_4.jpg",
            lambda image_name: 123004,
        )
        self.assertEqual(enriched[0]["image_id"], 123004)
        self.assertEqual(enriched[0]["category_id"], "button")
        self.assertEqual(enriched[0]["bbox"], [1, 2, 3, 4])

    def test_write_json_atomic_writes_one_valid_json_document(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            out_path = Path(tmp_dir) / "ape.json"
            write_json_atomic(str(out_path), [{"image_id": 1}, {"image_id": 2}])
            self.assertEqual(json.loads(out_path.read_text()), [{"image_id": 1}, {"image_id": 2}])

    def test_optional_reflection_is_disabled_by_default(self):
        self.assertIsNone(
            run_optional_reflection(
                "unused.png",
                [],
                detector=lambda candidates, previous: previous,
                enabled=False,
            )
        )

    def test_optional_reflection_feeds_back_to_detector(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            image_path = Path(tmp_dir) / "image.png"
            Image.new("RGB", (10, 10), "white").save(image_path)
            detector_calls = []

            def advisor(**kwargs):
                if detector_calls:
                    return {"verified": [0], "needs_refinement": [], "feedback": []}
                return {"verified": [], "needs_refinement": [0], "feedback": ["button label"]}

            def detector(candidates, previous):
                detector_calls.append(candidates)
                return [{"bbox": [2, 2, 3, 3], "category_name": "button"}]

            result = run_optional_reflection(
                str(image_path),
                [{"bbox": [1, 1, 3, 3], "category_name": "button"}],
                detector=detector,
                enabled=True,
                advisor=advisor,
            )

        self.assertEqual(detector_calls, [["button label"]])
        self.assertFalse(result["max_iterations_reached"])


if __name__ == "__main__":
    unittest.main()