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_ai_tasks_support_cooperative_cancellation(self): with self.assertRaises(self.worker.TaskCancelled): self.worker.check_cancelled(lambda: True) self.worker.check_cancelled(lambda: False) self.assertIn( "cancel_check=cancel_check", inspect.getsource(self.worker.reconstruct), ) self.assertIn( "cancel_check=cancel_check", inspect.getsource(self.worker.remove_background), ) def test_optional_background_removal_runs_before_da3_model_load(self): source = inspect.getsource(self.worker.reconstruct) self.assertLess( source.index("infer_foreground("), source.index("DepthAnything3.from_pretrained"), ) args = self.worker.parser().parse_args( [ "reconstruct", "--input", "input.png", "--model-dir", "da3", "--output-dir", "result", "--export-format", "mini_npz", "--status", "status.json", "--remove-background-model-dir", "ben2", "--process-res", "1024", "--process-res-method", "upper_bound_resize", "--conf-thresh-percentile", "40", "--num-max-points", "1000", "--show-cameras", "0", "--gs-views-interval", "1", "--video-frames", "16", "--device", "auto", "--use-ray-pose", "1", "--ref-view-strategy", "first", ] ) self.assertEqual(args.remove_background_model_dir, "ben2") 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_download_passes_optional_file_patterns(self): calls = [] fake_hugging_face = ModuleType("huggingface_hub") fake_hugging_face.snapshot_download = lambda **options: calls.append( options ) args = SimpleNamespace( repository="PramaLLC/BEN2", model_dir="model", status="status.json", allow_pattern=["config.json", "model.safetensors"], ) with tempfile.TemporaryDirectory() as directory, patch.dict( sys.modules, {"huggingface_hub": fake_hugging_face}, ), patch.object( self.worker, "create_download_progress_class", return_value=object, ): args.model_dir = str(Path(directory) / "model") args.status = str(Path(directory) / "status.json") self.worker.download_model(args) self.assertEqual( calls[0]["allow_patterns"], ["config.json", "model.safetensors"], ) def test_ben2_background_removal_writes_fast_transparent_png(self): inference_options = [] save_options = [] class FakeImage: def __enter__(self): return self def __exit__(self, *_args): return None def convert(self, mode): self.mode = mode return self def save(self, path, **options): save_options.append(options) Path(path).touch() class FakeModel: @classmethod def from_pretrained(cls, model_dir): self.assertEqual(model_dir, "model") return cls() def to(model_self, device): self.assertEqual(device.type, "cpu") return model_self def eval(self): return None def inference(self, image, **options): inference_options.append((image, options)) return FakeImage() fake_torch = ModuleType("torch") fake_torch.cuda = SimpleNamespace( is_available=lambda: False, synchronize=lambda: None, ) fake_torch.device = lambda device_type: SimpleNamespace(type=device_type) fake_torch.inference_mode = nullcontext fake_ben2 = ModuleType("ben2") fake_ben2.BEN_Base = FakeModel fake_pil = ModuleType("PIL") fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage()) with tempfile.TemporaryDirectory() as directory, patch.dict( sys.modules, { "torch": fake_torch, "ben2": fake_ben2, "PIL": fake_pil, }, ): root = Path(directory) input_path = root / "input.png" input_path.touch() args = SimpleNamespace( input=str(input_path), model_dir="model", output_dir=str(root / "result"), status=str(root / "status.json"), device="cpu", ) self.worker.remove_background(args) status = json.loads(Path(args.status).read_text(encoding="utf-8")) self.assertTrue((root / "result" / "foreground.png").is_file()) self.assertEqual( inference_options[0][1], {"refine_foreground": False}, ) self.assertEqual( save_options[0], {"format": "PNG", "compress_level": 1}, ) self.assertIn("Background removed in", status["message"]) 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()