AnyImage/tests/test_worker.py

import importlib.util
import inspect
import json
import sys
import tempfile
import unittest
from contextlib import nullcontext
from pathlib import Path
from types import ModuleType, SimpleNamespace
from unittest.mock import patch


PROJECT_ROOT = Path(__file__).resolve().parents[1]


def load_worker():
    specification = importlib.util.spec_from_file_location(
        "snapform_worker",
        PROJECT_ROOT / "src" / "snapform" / "worker.py",
    )
    module = importlib.util.module_from_spec(specification)
    specification.loader.exec_module(module)
    return module


class WorkerTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.worker = load_worker()

    def test_collect_images_is_sorted_and_filters_suffixes(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            (root / "b.PNG").touch()
            (root / "a.jpg").touch()
            (root / "notes.txt").touch()

            images = self.worker.collect_images(root)

        self.assertEqual([path.name for path in images], ["a.jpg", "b.PNG"])

    def test_reconstruct_reports_status_before_loading_torch(self):
        source = inspect.getsource(self.worker.reconstruct)

        self.assertLess(
            source.index('"Starting AI runtime"'),
            source.index("import torch"),
        )

    def test_prepares_a_single_image_file(self):
        with tempfile.TemporaryDirectory() as directory:
            image_path = Path(directory) / "input.png"
            image_path.touch()

            images, temporary_directory = self.worker.prepare_input(
                image_path,
                16,
                directory,
            )

        self.assertEqual(images, [image_path])
        self.assertIsNone(temporary_directory)

    def test_rejects_unsupported_input_file(self):
        with tempfile.TemporaryDirectory() as directory:
            input_path = Path(directory) / "input.txt"
            input_path.touch()

            with self.assertRaisesRegex(RuntimeError, "Unsupported input type"):
                self.worker.prepare_input(input_path, 16, directory)

    def test_uses_official_da3_export_parameters_for_all_four_formats(self):
        inference_options = []
        export_calls = []

        class FakeModel:
            @classmethod
            def from_pretrained(cls, model_dir):
                if model_dir != "model":
                    raise AssertionError(f"Unexpected model directory: {model_dir}")
                return cls()

            def to(self, device):
                if device.type != "cpu":
                    raise AssertionError(f"Unexpected device: {device.type}")
                return self

            def eval(self):
                return None

            def inference(self, images, **options):
                inference_options.append((images, options))
                return SimpleNamespace()

        def fake_export(prediction, export_format, export_dir, **options):
            export_calls.append((prediction, export_format, export_dir, options))
            output_directory = Path(export_dir)
            if export_format == "glb":
                (output_directory / "scene.glb").touch()
            elif export_format == "gs_ply":
                ply_directory = output_directory / "gs_ply"
                ply_directory.mkdir()
                (ply_directory / "0000.ply").touch()
            elif export_format == "depth_vis":
                depth_directory = output_directory / "depth_vis"
                depth_directory.mkdir()
                (depth_directory / "0000.jpg").touch()
            else:
                mini_directory = output_directory / "exports" / "mini_npz"
                mini_directory.mkdir(parents=True)
                (mini_directory / "results.npz").touch()

        fake_torch = ModuleType("torch")
        fake_torch.device = lambda device_type: SimpleNamespace(
            type=device_type,
        )
        fake_torch.inference_mode = nullcontext
        fake_api = ModuleType("depth_anything_3.api")
        fake_api.DepthAnything3 = FakeModel
        fake_export_module = ModuleType("depth_anything_3.utils.export")
        fake_export_module.export = fake_export
        fake_utils = ModuleType("depth_anything_3.utils")
        fake_utils.__path__ = []
        fake_package = ModuleType("depth_anything_3")
        fake_package.__path__ = []

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            input_path = root / "input.png"
            input_path.touch()
            with patch.dict(
                sys.modules,
                {
                    "torch": fake_torch,
                    "depth_anything_3": fake_package,
                    "depth_anything_3.api": fake_api,
                    "depth_anything_3.utils": fake_utils,
                    "depth_anything_3.utils.export": fake_export_module,
                },
            ):
                for export_format in self.worker.EXPORT_FORMATS:
                    args = SimpleNamespace(
                        input=str(input_path),
                        video_frames=16,
                        output_dir=str(root / export_format),
                        export_format=export_format,
                        status=str(root / f"{export_format}.json"),
                        device="cpu",
                        model_dir="model",
                        process_res=518,
                        process_res_method="lower_bound_resize",
                        use_ray_pose=1,
                        ref_view_strategy="middle",
                        conf_thresh_percentile=35.0,
                        num_max_points=12345,
                        show_cameras=0,
                        gs_views_interval=2,
                    )
                    self.worker.reconstruct(args)
                    status = json.loads(
                        Path(args.status).read_text(encoding="utf-8")
                    )
                    self.assertEqual(status["progress"], 1.0)
                    self.assertIn(export_format, status["message"])

        self.assertEqual(len(inference_options), 4)
        self.assertEqual(len(export_calls), 4)
        for images, options in inference_options:
            self.assertEqual(images, [str(input_path)])
            self.assertEqual(options["process_res"], 518)
            self.assertEqual(
                options["process_res_method"],
                "lower_bound_resize",
            )
            self.assertTrue(options["use_ray_pose"])
            self.assertEqual(options["ref_view_strategy"], "middle")
        for _prediction, export_format, _export_dir, options in export_calls:
            if export_format == "glb":
                self.assertEqual(
                    options,
                    {
                        "glb": {
                            "conf_thresh_percentile": 35.0,
                            "num_max_points": 12345,
                            "show_cameras": False,
                            "export_depth_vis": False,
                        }
                    },
                )
            elif export_format == "gs_ply":
                self.assertEqual(
                    options,
                    {"gs_ply": {"gs_views_interval": 2}},
                )
            else:
                self.assertEqual(options, {})

        infer_gs_values = [options["infer_gs"] for _, options in inference_options]
        self.assertEqual(infer_gs_values, [False, True, False, False])


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