tests/test_ape_stage.py

import json
import sys
import tempfile
import unittest
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT))

from approach.ape_stage import extract_image_id, run_ape_stage


class ApeStageTests(unittest.TestCase):
    def test_extract_image_id_requires_documented_basename_shape(self):
        self.assertEqual(extract_image_id("123_4.jpg"), 123004)
        with self.assertRaises(ValueError):
            extract_image_id("123_4_view_a.png")
        with self.assertRaises(ValueError):
            extract_image_id("123_view_a.png")

    def test_multiple_records_are_enriched_and_checkpointed_as_valid_json(self):
        calls = []

        def fake_inference(**kwargs):
            calls.append(kwargs)
            category = kwargs["text_prompt"].split(":", 1)[0]
            return [
                {
                    "category_name": category,
                    "bbox": [1, 2, 3, 4],
                    "score": 0.9,
                }
            ]

        records = [
            {"question_id": 0, "text": {"objects": {"button": "round red"}}},
            {"question_id": 1, "text": {"objects": {"lever": "long silver"}}},
        ]
        questions = {
            0: {"image": "123_4.jpg"},
            1: {"image": "456_78.jpg"},
        }

        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            output_path = root / "predictions.json"
            results, errors = run_ape_stage(
                records=records,
                questions=questions,
                images_dir=root / "images",
                output_path=output_path,
                inference=fake_inference,
                inference_kwargs={"confidence_threshold": 0.15},
                resume=False,
            )

            on_disk = json.loads(output_path.read_text())

        self.assertEqual(errors, [])
        self.assertEqual(results, on_disk)
        self.assertEqual(len(results), 2)
        self.assertEqual(results[0]["image_id"], 123004)
        self.assertEqual(results[0]["category_id"], "button")
        self.assertEqual(results[1]["image_id"], 456078)
        self.assertEqual(results[1]["category_id"], "lever")
        self.assertEqual(calls[0]["confidence_threshold"], 0.15)

    def test_resume_does_not_duplicate_completed_images(self):
        calls = []

        def fake_inference(**kwargs):
            calls.append(kwargs)
            return [
                {
                    "category_name": "button",
                    "bbox": [1, 2, 3, 4],
                    "score": 0.9,
                }
            ]

        records = [{"question_id": 0, "text": {"objects": {"button": "red"}}}]
        questions = {0: {"image": "123_4.jpg"}}

        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            output_path = root / "predictions.json"
            first, _ = run_ape_stage(
                records,
                questions,
                root,
                output_path,
                fake_inference,
                resume=False,
            )
            second, _ = run_ape_stage(
                records,
                questions,
                root,
                output_path,
                fake_inference,
                resume=True,
            )

        self.assertEqual(first, second)
        self.assertEqual(len(calls), 1)

    def test_resume_tracks_completed_images_with_zero_detections(self):
        calls = []

        def empty_inference(**kwargs):
            calls.append(kwargs)
            return []

        records = [{"question_id": "0", "text": {"objects": {"button": "red"}}}]
        questions = {0: {"image": "123_4.jpg"}}

        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            output_path = root / "predictions.json"
            first, _ = run_ape_stage(
                records,
                questions,
                root,
                output_path,
                empty_inference,
                resume=False,
            )
            second, _ = run_ape_stage(
                records,
                questions,
                root,
                output_path,
                empty_inference,
                resume=True,
            )

            progress = json.loads((root / "predictions.progress.json").read_text())

        self.assertEqual(first, [])
        self.assertEqual(second, [])
        self.assertEqual(calls, [calls[0]])
        self.assertEqual(progress["completed_image_ids"], [123004])

    def test_errors_are_recorded_without_corrupting_predictions(self):
        def failing_inference(**kwargs):
            raise RuntimeError("synthetic detector failure")

        records = [{"question_id": 0, "text": {"objects": {"button": "red"}}}]
        questions = {0: {"image": "123_4.jpg"}}

        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            output_path = root / "predictions.json"
            error_path = root / "errors.json"
            results, errors = run_ape_stage(
                records,
                questions,
                root,
                output_path,
                failing_inference,
                error_path=error_path,
                resume=False,
            )

            self.assertEqual(json.loads(output_path.read_text()), [])
            self.assertEqual(json.loads(error_path.read_text()), errors)

        self.assertEqual(results, [])
        self.assertEqual(errors[0]["question_id"], 0)
        self.assertEqual(errors[0]["error_type"], "RuntimeError")
        self.assertNotIn("traceback", errors[0])


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