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_service(): specification = importlib.util.spec_from_file_location( "snapform_service", PROJECT_ROOT / "src" / "snapform" / "ai_service" / "server.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 ServiceTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.service = load_service() def setUp(self): self.jobs_temp = tempfile.TemporaryDirectory() self.addCleanup(self.jobs_temp.cleanup) self.jobs_directory = Path(self.jobs_temp.name) / "jobs" def create_cache(self): cache = self.service.ModelSessionCache.__new__( self.service.ModelSessionCache ) cache.ben2_model = None cache.ben2_key = None cache.da3_model = None cache.da3_key = None cache.moge_model = None cache.moge_key = None return cache def test_ben2_model_is_reused_for_matching_runtime_key(self): loads = [] ben2 = ModuleType("ben2_onnx") ben2.create_session = lambda model_dir, device: ( loads.append((model_dir, device)) or object() ) cache = self.create_cache() device = SimpleNamespace(type="cuda") with patch.dict(sys.modules, {"ben2_onnx": 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", "cuda")]) self.assertEqual(cache.snapshot()["ben2"], "model") self.assertEqual(cached_load_ms, 0.0) def test_da3_model_is_reused_for_matching_runtime_key(self): loads = [] da3 = ModuleType("da3_onnx") da3.create_session = lambda model_dir, device: ( loads.append((model_dir, device)) or object() ) cache = self.create_cache() device = SimpleNamespace(type="cuda") with patch.dict( sys.modules, { "da3_onnx": da3, }, ): first = cache.get_da3("model", device) second = cache.get_da3("model", device) self.assertIs(first, second) self.assertEqual(loads, [("model", "cuda")]) self.assertEqual(cache.snapshot()["da3"], "model") def test_depth_model_cache_keeps_only_da3_or_moge(self): cache = self.create_cache() cache.ben2_model = object() cache.ben2_key = ("ben2", "cuda") cache.moge_model = object() cache.moge_key = ("moge", "cuda") device = SimpleNamespace(type="cuda") da3 = object() def load_da3(_model_dir, _device, key): cache.da3_model = da3 cache.da3_key = key return da3 with patch.object(cache, "_load_da3", side_effect=load_da3): self.assertIs(cache.get_da3("da3", device), da3) self.assertIsNone(cache.moge_model) self.assertIsNotNone(cache.ben2_model) moge = object() fake_moge = ModuleType("moge_onnx") fake_moge.create_session = lambda _model_dir, _device: moge with patch.dict(sys.modules, {"moge_onnx": fake_moge}): self.assertIs(cache.get_moge_onnx("moge", "cuda"), moge) self.assertIsNone(cache.da3_model) self.assertIsNotNone(cache.ben2_model) def test_dedicated_service_runs_only_one_job_at_a_time(self): fake_models = SimpleNamespace( snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}}, clear=lambda: None, close=lambda: None, ) with patch.object(self.service, "ModelManager", return_value=fake_models): server = self.service.create_server(jobs_directory=self.jobs_directory) parameters = {"model": "BASE"} first = server.submit("debug-da3", parameters, job_id="job") second = server.submit("debug-da3", parameters, job_id="other") self.assertIsNotNone(first) self.assertIsNone(second) self.assertTrue(server.snapshot("instance")["busy"]) self.assertNotIn("clients", server.snapshot("instance")) def test_service_registers_tasks_instead_of_model_names(self): server = self.service.create_server(jobs_directory=self.jobs_directory) self.assertEqual( set(server.handlers), { "download-model", "remove-background", "extract-frames", "generate-da3-depth", "generate-moge2-depth", "generate-da3-surface", "generate-moge2-surface", "debug-da3", "debug-moge2", "debug-ben2", }, ) def test_cancelling_job_keeps_service_and_model_cache_alive(self): fake_models = SimpleNamespace( snapshot=lambda: { "loaded": {"ben2": "BEN2_BASE", "da3": ""} }, clear=lambda: None, close=lambda: None, ) with patch.object(self.service, "ModelManager", return_value=fake_models): server = self.service.create_server(jobs_directory=self.jobs_directory) context = server.submit("remove-background", {}, job_id="job") self.assertTrue(context.request_cancel()) self.assertTrue(context.is_cancelled()) server.run(context) status = context.snapshot() self.assertEqual(status["state"], "cancelled") self.assertFalse(server.shutdown_event.is_set()) self.assertEqual( server.snapshot("instance")["resources"]["models"]["loaded"]["ben2"], "BEN2_BASE", ) def test_job_progress_is_kept_in_service_memory(self): server = self.service.JobServer("Test", jobs_directory=self.jobs_directory) @server.job("example") def example(context, _parameters): context.update(0.5, "Half way") context = server.submit("example", {}) server.run(context) self.assertRegex(context.job_id, r"^[0-9a-f]{32}$") self.assertEqual(context.directory.name, context.job_id) self.assertEqual(context.snapshot()["state"], "succeeded") self.assertEqual(context.snapshot()["progress"], 1.0) def test_job_server_owns_and_clears_shared_resources(self): events = [] resource = SimpleNamespace( snapshot=lambda: {"loaded": "model"}, clear=lambda: events.append("clear"), close=lambda: events.append("close"), ) server = self.service.JobServer("Test", jobs_directory=self.jobs_directory) server.add_resource("models", resource) @server.job("example") def example(context, _parameters): self.assertIs(context.resource("models"), resource) context = server.submit("example", {}, job_id="job") self.assertFalse(server.clear_resource("models")) server.run(context) self.assertTrue(server.clear_resource("models")) server.close() self.assertEqual( server.snapshot("instance")["resources"]["models"]["loaded"], "model", ) self.assertEqual(events, ["clear", "close"]) def test_service_dispatches_frame_extraction_without_models(self): fake_models = SimpleNamespace( snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}}, clear=lambda: None, close=lambda: None, ) with patch.object(self.service, "ModelManager", return_value=fake_models): server = self.service.create_server(jobs_directory=self.jobs_directory) 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() context = server.submit( "extract-frames", { "input": str(input_directory), "video_frames": 16, "start_frame": 1, }, job_id="job", ) server.run(context) status = context.snapshot() self.assertTrue( (context.directory / "frame.0001.jpg").is_file() ) self.assertEqual(status["state"], "succeeded") if __name__ == "__main__": unittest.main()