AnyImage/tests/test_ai_host.py

import importlib.util
import inspect
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"], "da3_model")

    def test_http_host_has_health_jobs_auth_and_single_task_guard(self):
        source = inspect.getsource(self.host.create_app)

        self.assertIn('@app.get("/health")', source)
        self.assertIn('@app.post("/jobs", status_code=202)', source)
        self.assertIn("x_snapform_token", source)
        self.assertIn("hmac.compare_digest", source)
        self.assertIn("parameters[\"status\"]", source)
        self.assertIn('@app.post("/clients", status_code=201)', source)
        self.assertIn('@app.post("/clients/{client_id}/heartbeat")', source)
        self.assertIn('@app.delete("/jobs/{job_id}", status_code=202)', source)
        self.assertIn('@app.post("/models/clear")', source)
        self.assertIn('@app.post("/shutdown", status_code=202)', source)
        self.assertNotIn("os._exit", source)

    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())


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