tests/test_run_ape.py

import json
import tempfile
import unittest
from pathlib import Path

from approach.run_ape import (
    _resolve_ape_checkpoint,
    _verify_ape_checkpoint,
    build_parser,
    run,
)


class RunApeTests(unittest.TestCase):
    def test_custom_checkpoint_requires_explicit_trust(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            checkpoint = Path(tmpdir) / "custom.pth"
            checkpoint.write_bytes(b"not-the-released-checkpoint")
            with self.assertRaises(ValueError):
                _verify_ape_checkpoint(checkpoint)
            _verify_ape_checkpoint(checkpoint, trust_custom_checkpoint=True)

    def test_checkpoint_resolver_accepts_repo_relative_path(self):
        repo_relative = Path("approach/ovod/APE/ape_d_model_final.pth")
        resolved = _resolve_ape_checkpoint(Path("approach/ovod/APE"), repo_relative)

        self.assertEqual(resolved, (Path.cwd() / repo_relative).resolve())

    def test_configurable_runner_selects_shard_and_writes_isolated_output(self):
        calls = []

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

        with tempfile.TemporaryDirectory() as tmpdir:
            root = Path(tmpdir)
            questions = root / "questions.jsonl"
            candidates = root / "candidates.jsonl"
            output = root / "predictions.json"
            questions.write_text(
                "\n".join(
                    [
                        json.dumps({"question_id": 0, "image": "123_4.jpg"}),
                        json.dumps({"question_id": 1, "image": "456_7.jpg"}),
                    ]
                )
                + "\n"
            )
            candidates.write_text(
                "\n".join(
                    [
                        json.dumps(
                            {
                                "question_id": 0,
                                "text": '{"objects": {"button": "round red"}}',
                            }
                        ),
                        json.dumps(
                            {
                                "question_id": 1,
                                "text": {"objects": {"lever": "long silver"}},
                            }
                        ),
                    ]
                )
                + "\n"
            )
            args = build_parser().parse_args(
                [
                    "--questions",
                    str(questions),
                    "--candidates",
                    str(candidates),
                    "--images-dir",
                    str(root / "images"),
                    "--output",
                    str(output),
                    "--shard-index",
                    "0",
                    "--num-shards",
                    "2",
                ]
            )

            report = run(args, inference=inference)
            shard_output = root / "predictions.shard00-of-02.json"

            self.assertTrue(shard_output.exists())
            self.assertTrue((root / "predictions.shard00-of-02.progress.json").exists())
            self.assertEqual(json.loads(shard_output.read_text())[0]["image_id"], 123004)

        self.assertEqual(report["records_selected"], 1)
        self.assertEqual(len(calls), 1)
        self.assertEqual(calls[0]["confidence_threshold"], 0.15)


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