Add unit tests for service functionality and model caching
- Implement tests for ben2 and da3 model caching to ensure models are reused for matching runtime keys. - Verify service state management, including client registration, job cancellation, and client expiration. - Test discovery file creation and removal based on service ownership. - Ensure frame extraction functionality works without models and produces expected results.
This commit is contained in:
252
tests/test_service.py
Normal file
252
tests/test_service.py
Normal file
@@ -0,0 +1,252 @@
|
||||
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" / "service.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 create_cache(self):
|
||||
cache = self.service.ModelCache.__new__(self.service.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_service_tracks_clients_and_rejects_shared_shutdown(self):
|
||||
fake_models = SimpleNamespace(
|
||||
status=lambda: {"ben2": "", "da3": ""},
|
||||
close=lambda: None,
|
||||
)
|
||||
with patch.object(self.service, "ModelCache", return_value=fake_models):
|
||||
state = self.service.ServiceState("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_service_and_model_cache_alive(self):
|
||||
fake_models = SimpleNamespace(
|
||||
status=lambda: {"ben2": "BEN2_BASE", "da3": ""},
|
||||
close=lambda: None,
|
||||
)
|
||||
with patch.object(self.service, "ModelCache", return_value=fake_models):
|
||||
state = self.service.ServiceState("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.service.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.service, "ModelCache", return_value=fake_models):
|
||||
state = self.service.ServiceState("token", "instance")
|
||||
state.register_client("client", "Blender")
|
||||
state.clients["client"]["last_seen"] -= (
|
||||
self.service.CLIENT_TIMEOUT_SECONDS + 1.0
|
||||
)
|
||||
|
||||
state.expire_clients()
|
||||
|
||||
self.assertEqual(state.snapshot()["clients"], 0)
|
||||
|
||||
def test_discovery_is_only_removed_by_matching_service(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = Path(directory) / "service.json"
|
||||
args = SimpleNamespace(
|
||||
discovery=str(path),
|
||||
port=51234,
|
||||
token="token",
|
||||
instance_id="owner",
|
||||
)
|
||||
with patch.object(self.service.os, "getpid", return_value=42):
|
||||
self.service.write_discovery(args)
|
||||
|
||||
payload = self.service.json.loads(path.read_text(encoding="utf-8"))
|
||||
self.assertEqual(payload["port"], 51234)
|
||||
self.assertEqual(payload["pid"], 42)
|
||||
self.service.remove_discovery(path, "other")
|
||||
self.assertTrue(path.is_file())
|
||||
self.service.remove_discovery(path, "owner")
|
||||
self.assertFalse(path.exists())
|
||||
|
||||
def test_service_dispatches_frame_extraction_without_models(self):
|
||||
fake_models = SimpleNamespace(
|
||||
status=lambda: {"ben2": "", "da3": ""},
|
||||
close=lambda: None,
|
||||
)
|
||||
with patch.object(self.service, "ModelCache", return_value=fake_models):
|
||||
state = self.service.ServiceState("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.service.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()
|
||||
Reference in New Issue
Block a user