import importlib.util import sys import tempfile import unittest from pathlib import Path from types import ModuleType, SimpleNamespace from unittest.mock import patch PROJECT_ROOT = Path(__file__).resolve().parents[1] def load_ai_host(): specification = importlib.util.spec_from_file_location( "snapform_ai_host", PROJECT_ROOT / "src" / "snapform" / "ai_host.py", ) module = importlib.util.module_from_spec(specification) specification.loader.exec_module(module) return module class FakeCuda: class OutOfMemoryError(RuntimeError): pass @staticmethod def is_available(): return False @staticmethod def empty_cache(): return None class HostTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.host = load_ai_host() def create_cache(self): cache = self.host.ModelCache.__new__(self.host.ModelCache) cache.torch = SimpleNamespace(cuda=FakeCuda()) cache.ben2_model = None cache.ben2_key = None cache.da3_model = None cache.da3_key = None return cache def test_ben2_model_is_reused_for_matching_runtime_key(self): loads = [] class FakeModel: @classmethod def from_pretrained(cls, model_dir): loads.append(model_dir) return cls() def to(self, device): self.device = device return self def eval(self): return None ben2 = ModuleType("ben2") ben2.BEN_Base = FakeModel cache = self.create_cache() device = SimpleNamespace(type="cuda") with patch.dict(sys.modules, {"ben2": ben2}): first, _load_ms = cache.get_ben2("model", device) second, cached_load_ms = cache.get_ben2("model", device) self.assertIs(first, second) self.assertEqual(loads, ["model"]) self.assertEqual(cache.status()["ben2"], "model") self.assertEqual(cached_load_ms, 0.0) def test_da3_model_is_reused_for_matching_runtime_key(self): loads = [] class FakeModel: @classmethod def from_pretrained(cls, model_dir): loads.append(model_dir) return cls() def to(self, device): self.device = device return self def eval(self): return None package = ModuleType("depth_anything_3") api = ModuleType("depth_anything_3.api") api.DepthAnything3 = FakeModel cache = self.create_cache() device = SimpleNamespace(type="cuda") with patch.dict( sys.modules, { "depth_anything_3": package, "depth_anything_3.api": api, }, ): first = cache.get_da3("model", device) second = cache.get_da3("model", device) self.assertIs(first, second) self.assertEqual(loads, ["model"]) self.assertEqual(cache.status()["da3"], "model") def test_host_tracks_clients_and_rejects_shared_shutdown(self): fake_models = SimpleNamespace( status=lambda: {"ben2": "", "da3": ""}, close=lambda: None, ) with patch.object(self.host, "ModelCache", return_value=fake_models): state = self.host.HostState("token", "instance") state.register_client("first", "Blender A") state.register_client("second", "Blender B") self.assertEqual(state.snapshot()["clients"], 2) self.assertEqual(state.request_shutdown("first"), "shared") self.assertTrue( state.reserve("job", "first", "reconstruct", "status.json") ) self.assertFalse( state.reserve("other", "second", "reconstruct", "other.json") ) def test_cancelling_job_keeps_host_and_model_cache_alive(self): fake_models = SimpleNamespace( status=lambda: {"ben2": "BEN2_BASE", "da3": ""}, close=lambda: None, ) with patch.object(self.host, "ModelCache", return_value=fake_models): state = self.host.HostState("token", "instance") state.register_client("client", "Blender") with tempfile.TemporaryDirectory() as directory: status_path = Path(directory) / "status.json" state.reserve( "job", "client", "remove-background", str(status_path), ) active = state.cancel_job("job", "client") self.assertTrue(active["cancel_requested"]) self.assertTrue(state.is_job_cancelled("job")) state.run_job( { "job_id": "job", "command": "remove-background", "parameters": {"status": str(status_path)}, } ) status = self.host.json.loads( status_path.read_text(encoding="utf-8") ) self.assertEqual(status["state"], "cancelled") self.assertFalse(state.shutdown_event.is_set()) self.assertEqual(state.snapshot()["cached_models"]["ben2"], "BEN2_BASE") def test_expired_clients_are_removed(self): fake_models = SimpleNamespace( status=lambda: {"ben2": "", "da3": ""}, close=lambda: None, ) with patch.object(self.host, "ModelCache", return_value=fake_models): state = self.host.HostState("token", "instance") state.register_client("client", "Blender") state.clients["client"]["last_seen"] -= ( self.host.CLIENT_TIMEOUT_SECONDS + 1.0 ) state.expire_clients() self.assertEqual(state.snapshot()["clients"], 0) def test_discovery_is_only_removed_by_matching_host(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "host.json" args = SimpleNamespace( discovery=str(path), port=51234, token="token", instance_id="owner", ) with patch.object(self.host.os, "getpid", return_value=42): self.host.write_discovery(args) payload = self.host.json.loads(path.read_text(encoding="utf-8")) self.assertEqual(payload["port"], 51234) self.assertEqual(payload["pid"], 42) self.host.remove_discovery(path, "other") self.assertTrue(path.is_file()) self.host.remove_discovery(path, "owner") self.assertFalse(path.exists()) def test_host_dispatches_frame_extraction_without_models(self): fake_models = SimpleNamespace( status=lambda: {"ben2": "", "da3": ""}, close=lambda: None, ) with patch.object(self.host, "ModelCache", return_value=fake_models): state = self.host.HostState("token", "instance") state.register_client("client", "Blender") with tempfile.TemporaryDirectory() as directory: root = Path(directory) input_directory = root / "frames" input_directory.mkdir() (input_directory / "frame-0001.jpg").touch() (input_directory / "frame-0002.jpg").touch() result_directory = root / "result" status_path = root / "status.json" state.reserve( "job", "client", "extract-frames", str(status_path), ) state.run_job( { "job_id": "job", "command": "extract-frames", "parameters": { "input": str(input_directory), "output_dir": str(result_directory), "status": str(status_path), "video_frames": 16, "start_frame": 1, }, } ) status = self.host.json.loads( status_path.read_text(encoding="utf-8") ) self.assertTrue( (result_directory / "frame.0001.jpg").is_file() ) self.assertEqual(status["state"], "succeeded") if __name__ == "__main__": unittest.main()