ttulttul-ComfyUI-Modal/tests/test_sync_engine.py

"""Tests for asset syncing and custom_nodes archiving."""

from __future__ import annotations

from pathlib import Path
from typing import Any


def test_sync_file_deduplicates_by_hash(
    settings_module: Any,
    sync_engine_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """Repeated file syncs should reuse the same remote path and marker."""
    monkeypatch.setattr(sync_engine_module, "modal", None)
    asset_path = tmp_path / "model.safetensors"
    asset_path.write_bytes(b"model-bytes")

    settings = settings_module.ModalSyncSettings(
        app_name="app",
        auto_deploy=True,
        allow_ephemeral_fallback=False,
        enable_memory_snapshot=True,
        enable_gpu_memory_snapshot=False,
        execution_mode="local",
        sync_custom_nodes=False,
        volume_name="volume",
        route_path="/modal/queue_prompt",
        marker_property="is_modal_remote",
        local_storage_root=tmp_path / "storage",
        remote_storage_root="/storage",
        custom_nodes_archive_name="custom_nodes_bundle.zip",
        comfyui_root=None,
        custom_nodes_dir=None,
    )

    engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
    first = engine.sync_file(asset_path)
    second = engine.sync_file(asset_path)

    assert first.remote_path == second.remote_path
    assert first.sha256 == second.sha256
    assert first.uploaded is True
    assert second.uploaded is False
    assert (settings.local_storage_root / first.remote_path.lstrip("/")).exists()
    assert (settings.local_storage_root / f"hashes/{first.sha256}.done").exists()


def test_sync_custom_nodes_directory_creates_archive(
    settings_module: Any,
    sync_engine_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """The sync engine should archive and mirror a custom_nodes directory."""
    monkeypatch.setattr(sync_engine_module, "modal", None)
    custom_nodes_dir = tmp_path / "custom_nodes"
    package_dir = custom_nodes_dir / "example"
    package_dir.mkdir(parents=True)
    (package_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8")

    settings = settings_module.ModalSyncSettings(
        app_name="app",
        auto_deploy=True,
        allow_ephemeral_fallback=False,
        enable_memory_snapshot=True,
        enable_gpu_memory_snapshot=False,
        execution_mode="remote",
        sync_custom_nodes=True,
        volume_name="volume",
        route_path="/modal/queue_prompt",
        marker_property="is_modal_remote",
        local_storage_root=tmp_path / "storage",
        remote_storage_root="/storage",
        custom_nodes_archive_name="custom_nodes_bundle.zip",
        comfyui_root=None,
        custom_nodes_dir=custom_nodes_dir,
    )

    engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
    bundle = engine.sync_custom_nodes_directory()

    assert bundle is not None
    assert bundle.remote_path.startswith("/custom_nodes/")
    assert bundle.uploaded is True
    assert (settings.local_storage_root / bundle.remote_path.lstrip("/")).exists()


def test_hash_directory_ignores_virtualenv_and_bytecode_artifacts(
    settings_module: Any,
    sync_engine_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """Ignored directories and file suffixes should not affect custom_nodes hashing."""
    monkeypatch.setattr(sync_engine_module, "modal", None)
    custom_nodes_dir = tmp_path / "custom_nodes"
    package_dir = custom_nodes_dir / "example"
    package_dir.mkdir(parents=True)
    (package_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8")

    settings = settings_module.ModalSyncSettings(
        app_name="app",
        auto_deploy=True,
        allow_ephemeral_fallback=False,
        enable_memory_snapshot=True,
        enable_gpu_memory_snapshot=False,
        execution_mode="remote",
        sync_custom_nodes=True,
        volume_name="volume",
        route_path="/modal/queue_prompt",
        marker_property="is_modal_remote",
        local_storage_root=tmp_path / "storage",
        remote_storage_root="/storage",
        custom_nodes_archive_name="custom_nodes_bundle.zip",
        comfyui_root=None,
        custom_nodes_dir=custom_nodes_dir,
    )

    engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
    baseline_hash = engine._hash_directory(custom_nodes_dir)

    ignored_venv_dir = custom_nodes_dir / ".venv" / "lib"
    ignored_venv_dir.mkdir(parents=True)
    (ignored_venv_dir / "ignored.py").write_text("print('ignored')\n", encoding="utf-8")
    ignored_cache_dir = package_dir / "__pycache__"
    ignored_cache_dir.mkdir()
    (ignored_cache_dir / "example.cpython-312.pyc").write_bytes(b"bytecode")
    (package_dir / "ignored.pyc").write_bytes(b"bytecode")
    (package_dir / "ignored.log").write_text("log noise\n", encoding="utf-8")

    ignored_hash = engine._hash_directory(custom_nodes_dir)

    assert ignored_hash == baseline_hash


def test_sync_file_reuses_cached_hash_for_unchanged_file(
    settings_module: Any,
    sync_engine_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """Unchanged files should reuse the persisted file digest instead of re-reading the payload."""
    monkeypatch.setattr(sync_engine_module, "modal", None)
    asset_path = tmp_path / "model.safetensors"
    asset_path.write_bytes(b"model-bytes")

    settings = settings_module.ModalSyncSettings(
        app_name="app",
        auto_deploy=True,
        allow_ephemeral_fallback=False,
        enable_memory_snapshot=True,
        enable_gpu_memory_snapshot=False,
        execution_mode="local",
        sync_custom_nodes=False,
        volume_name="volume",
        route_path="/modal/queue_prompt",
        marker_property="is_modal_remote",
        local_storage_root=tmp_path / "storage",
        remote_storage_root="/storage",
        custom_nodes_archive_name="custom_nodes_bundle.zip",
        comfyui_root=None,
        custom_nodes_dir=None,
    )

    first_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
    first_sha = first_engine._hash_file(asset_path)

    second_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)

    def fail_open(*args: Any, **kwargs: Any) -> Any:
        raise AssertionError("Expected cached hash lookup to avoid reopening the unchanged file.")

    monkeypatch.setattr(Path, "open", fail_open)
    second_sha = second_engine._hash_file(asset_path)

    assert second_sha == first_sha


def test_sync_custom_nodes_directory_reuses_cached_archive(
    settings_module: Any,
    sync_engine_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """An unchanged custom_nodes tree should reuse its digest-keyed local archive."""
    monkeypatch.setattr(sync_engine_module, "modal", None)
    custom_nodes_dir = tmp_path / "custom_nodes"
    package_dir = custom_nodes_dir / "example"
    package_dir.mkdir(parents=True)
    (package_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8")

    settings = settings_module.ModalSyncSettings(
        app_name="app",
        auto_deploy=True,
        allow_ephemeral_fallback=False,
        enable_memory_snapshot=True,
        enable_gpu_memory_snapshot=False,
        execution_mode="remote",
        sync_custom_nodes=True,
        volume_name="volume",
        route_path="/modal/queue_prompt",
        marker_property="is_modal_remote",
        local_storage_root=tmp_path / "storage",
        remote_storage_root="/storage",
        custom_nodes_archive_name="custom_nodes_bundle.zip",
        comfyui_root=None,
        custom_nodes_dir=custom_nodes_dir,
    )

    first_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
    first_bundle = first_engine.sync_custom_nodes_directory()
    assert first_bundle is not None

    class NeverExistsVolume:
        """Volume double that forces archive reuse by reporting no remote markers."""

        def exists(self, remote_path: str) -> bool:
            return False

        def put_file(self, local_path: Path, remote_path: str) -> None:
            return None

        def put_bytes(self, payload: bytes, remote_path: str) -> None:
            return None

    second_engine = sync_engine_module.ModalAssetSyncEngine(
        volume=NeverExistsVolume(),
        settings=settings,
    )

    def fail_create_archive(*args: Any, **kwargs: Any) -> Any:
        raise AssertionError("Expected cached custom_nodes archive to be reused.")

    monkeypatch.setattr(second_engine, "_create_archive", fail_create_archive)
    second_bundle = second_engine.sync_custom_nodes_directory()

    assert second_bundle is not None
    assert second_bundle.sha256 == first_bundle.sha256
    assert second_bundle.uploaded is True
    assert second_engine._cached_custom_nodes_archive_path(first_bundle.sha256).exists()


def test_sync_file_backfills_marker_when_remote_payload_already_exists(
    settings_module: Any,
    sync_engine_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """An existing deterministic remote asset should be reused even when its marker is missing."""
    monkeypatch.setattr(sync_engine_module, "modal", None)
    asset_path = tmp_path / "model.safetensors"
    asset_path.write_bytes(b"model-bytes")

    settings = settings_module.ModalSyncSettings(
        app_name="app",
        auto_deploy=True,
        allow_ephemeral_fallback=False,
        enable_memory_snapshot=True,
        enable_gpu_memory_snapshot=False,
        execution_mode="local",
        sync_custom_nodes=False,
        volume_name="volume",
        route_path="/modal/queue_prompt",
        marker_property="is_modal_remote",
        local_storage_root=tmp_path / "storage",
        remote_storage_root="/storage",
        custom_nodes_archive_name="custom_nodes_bundle.zip",
        comfyui_root=None,
        custom_nodes_dir=None,
    )

    class ExistingRemoteFileVolume:
        """Volume double with a missing marker but a pre-existing deterministic payload path."""

        def __init__(self) -> None:
            """Initialize captured writes."""
            self.put_file_calls: list[tuple[Path, str]] = []
            self.put_bytes_calls: list[tuple[bytes, str]] = []

        def exists(self, remote_path: str) -> bool:
            return remote_path.startswith("/assets/") and remote_path.endswith("_model.safetensors")

        def put_file(self, local_path: Path, remote_path: str) -> None:
            self.put_file_calls.append((local_path, remote_path))

        def put_bytes(self, payload: bytes, remote_path: str) -> None:
            self.put_bytes_calls.append((payload, remote_path))

    volume = ExistingRemoteFileVolume()
    engine = sync_engine_module.ModalAssetSyncEngine(volume=volume, settings=settings)

    synced_asset = engine.sync_file(asset_path)

    assert synced_asset.uploaded is False
    assert volume.put_file_calls == []
    assert len(volume.put_bytes_calls) == 1
    assert volume.put_bytes_calls[0][1] == f"/hashes/{synced_asset.sha256}.done"


def test_sync_custom_nodes_directory_backfills_marker_when_remote_bundle_already_exists(
    settings_module: Any,
    sync_engine_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """An existing hash-named remote custom_nodes bundle should be reused without rebuilding or reuploading."""
    monkeypatch.setattr(sync_engine_module, "modal", None)
    custom_nodes_dir = tmp_path / "custom_nodes"
    package_dir = custom_nodes_dir / "example"
    package_dir.mkdir(parents=True)
    (package_dir / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n", encoding="utf-8")

    settings = settings_module.ModalSyncSettings(
        app_name="app",
        auto_deploy=True,
        allow_ephemeral_fallback=False,
        enable_memory_snapshot=True,
        enable_gpu_memory_snapshot=False,
        execution_mode="remote",
        sync_custom_nodes=True,
        volume_name="volume",
        route_path="/modal/queue_prompt",
        marker_property="is_modal_remote",
        local_storage_root=tmp_path / "storage",
        remote_storage_root="/storage",
        custom_nodes_archive_name="custom_nodes_bundle.zip",
        comfyui_root=None,
        custom_nodes_dir=custom_nodes_dir,
    )

    engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
    directory_hash = engine._hash_directory(custom_nodes_dir)
    remote_path = f"/custom_nodes/{directory_hash}_{settings.custom_nodes_archive_name}"
    marker_path = f"/hashes/custom_nodes_{directory_hash}.done"

    class ExistingRemoteBundleVolume:
        """Volume double with a missing marker but a pre-existing deterministic custom_nodes bundle path."""

        def __init__(self) -> None:
            """Initialize captured writes."""
            self.put_file_calls: list[tuple[Path, str]] = []
            self.put_bytes_calls: list[tuple[bytes, str]] = []

        def exists(self, candidate_path: str) -> bool:
            return candidate_path == remote_path

        def put_file(self, local_path: Path, candidate_path: str) -> None:
            self.put_file_calls.append((local_path, candidate_path))

        def put_bytes(self, payload: bytes, candidate_path: str) -> None:
            self.put_bytes_calls.append((payload, candidate_path))

    volume = ExistingRemoteBundleVolume()
    engine = sync_engine_module.ModalAssetSyncEngine(volume=volume, settings=settings)

    def fail_create_archive(*args: Any, **kwargs: Any) -> Any:
        raise AssertionError("Expected the existing hash-named remote bundle to be reused.")

    monkeypatch.setattr(engine, "_create_archive", fail_create_archive)
    bundle = engine.sync_custom_nodes_directory()

    assert bundle is not None
    assert bundle.sha256 == directory_hash
    assert bundle.remote_path == remote_path
    assert bundle.uploaded is False
    assert volume.put_file_calls == []
    assert len(volume.put_bytes_calls) == 1
    assert volume.put_bytes_calls[0][1] == marker_path


def test_remote_mode_uses_modal_volume_backend_when_sdk_is_available(
    settings_module: Any,
    sync_engine_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """Remote execution mode should upload into a real Modal volume when the SDK is available."""

    class FakeBatch:
        """Capture files uploaded through Modal batch_upload."""

        def __init__(self, uploads: list[tuple[Any, str]]) -> None:
            """Store uploaded file references."""
            self.uploads = uploads

        def __enter__(self) -> "FakeBatch":
            """Return the active batch context."""
            return self

        def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
            """Close the fake batch context."""
            return None

        def put_file(self, source: Any, remote_path: str) -> None:
            """Record the uploaded file or bytes payload."""
            self.uploads.append((source, remote_path))

    class FakeVolume:
        """Minimal Modal volume double."""

        def __init__(self) -> None:
            """Initialize the fake upload state."""
            self.paths: set[str] = set()
            self.uploads: list[tuple[Any, str]] = []

        def listdir(self, remote_path: str, recursive: bool = False) -> list[str]:
            """Return a non-empty listing when the path has already been uploaded."""
            return [remote_path] if remote_path in self.paths else []

        def batch_upload(self) -> FakeBatch:
            """Return a fake batch uploader."""
            return FakeBatch(self.uploads)

    fake_volume = FakeVolume()

    class FakeModal:
        """Minimal modal SDK double that returns a stable volume handle."""

        class Volume:
            """Namespace for volume lookups."""

            @staticmethod
            def from_name(name: str, create_if_missing: bool = False) -> FakeVolume:
                """Return the fake volume for any lookup."""
                return fake_volume

    monkeypatch.setattr(sync_engine_module, "modal", FakeModal)

    settings = settings_module.ModalSyncSettings(
        app_name="app",
        auto_deploy=True,
        allow_ephemeral_fallback=False,
        enable_memory_snapshot=True,
        enable_gpu_memory_snapshot=False,
        execution_mode="remote",
        sync_custom_nodes=True,
        volume_name="volume",
        route_path="/modal/queue_prompt",
        marker_property="is_modal_remote",
        local_storage_root=tmp_path / "storage",
        remote_storage_root="/storage",
        custom_nodes_archive_name="custom_nodes_bundle.zip",
        comfyui_root=None,
        custom_nodes_dir=None,
    )

    engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
    assert type(engine.volume).__name__ == "ModalVolumeBackend"

    asset_path = tmp_path / "encoder.safetensors"
    asset_path.write_bytes(b"weights")
    synced = engine.sync_file(asset_path)

    uploaded_remote_paths = [remote_path for _, remote_path in fake_volume.uploads]
    assert synced.remote_path in uploaded_remote_paths
    assert f"/hashes/{synced.sha256}.done" in uploaded_remote_paths


def test_modal_volume_backend_treats_missing_path_as_cache_miss(
    sync_engine_module: Any,
    monkeypatch: Any,
) -> None:
    """A missing Modal volume path should behave like a normal absent marker file."""

    class FakeNotFoundError(Exception):
        """Stand-in for modal.exception.NotFoundError."""

    class FakeVolume:
        """Minimal Modal volume double that raises on missing listdir."""

        def listdir(self, remote_path: str, recursive: bool = False) -> list[str]:
            """Simulate Modal's missing-path behavior."""
            raise FakeNotFoundError(remote_path)

    class FakeModal:
        """Minimal modal SDK double exposing the exception namespace."""

        exception = type("FakeExceptionNamespace", (), {"NotFoundError": FakeNotFoundError})

        class Volume:
            """Namespace for volume lookups."""

            @staticmethod
            def from_name(name: str, create_if_missing: bool = False) -> FakeVolume:
                """Return a fake volume that always reports missing paths."""
                return FakeVolume()

    monkeypatch.setattr(sync_engine_module, "modal", FakeModal)
    backend = sync_engine_module.ModalVolumeBackend("volume")

    assert backend.exists("/hashes/missing.done") is False


def test_modal_volume_backend_caches_exists_results_and_uploaded_paths(
    sync_engine_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """Repeated existence checks for the same path should avoid repeated Modal metadata calls."""

    class FakeBatch:
        """Capture one uploaded path."""

        def __init__(self, volume: "FakeVolume") -> None:
            """Store the backing volume."""
            self.volume = volume

        def __enter__(self) -> "FakeBatch":
            """Return the active batch context."""
            return self

        def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
            """Close the fake batch context."""
            return None

        def put_file(self, source: Any, remote_path: str) -> None:
            """Record the uploaded path as present."""
            self.volume.paths.add(remote_path)

    class FakeVolume:
        """Minimal Modal volume double with listdir accounting."""

        def __init__(self) -> None:
            """Initialize fake storage state."""
            self.paths: set[str] = set()
            self.listdir_calls = 0

        def listdir(self, remote_path: str, recursive: bool = False) -> list[str]:
            """Return a listing for known paths while counting calls."""
            self.listdir_calls += 1
            return [remote_path] if remote_path in self.paths else []

        def batch_upload(self) -> FakeBatch:
            """Return a fake uploader."""
            return FakeBatch(self)

    fake_volume = FakeVolume()

    class FakeModal:
        """Minimal modal SDK double exposing a fake volume."""

        exception = type(
            "FakeExceptionNamespace",
            (),
            {
                "NotFoundError": FileNotFoundError,
            },
        )

        class Volume:
            """Namespace for volume lookups."""

            @staticmethod
            def from_name(name: str, create_if_missing: bool = False) -> FakeVolume:
                """Return the fake volume for any lookup."""
                return fake_volume

    monkeypatch.setattr(sync_engine_module, "modal", FakeModal)
    backend = sync_engine_module.ModalVolumeBackend("volume")

    assert backend.exists("/hashes/present.done") is False
    assert backend.exists("/hashes/present.done") is False
    assert fake_volume.listdir_calls == 1

    local_path = tmp_path / "marker.txt"
    local_path.write_text("marker", encoding="utf-8")
    backend.put_file(local_path, "/hashes/present.done")

    assert backend.exists("/hashes/present.done") is True
    assert fake_volume.listdir_calls == 1


def test_modal_volume_backend_retries_rate_limited_calls(
    sync_engine_module: Any,
    monkeypatch: Any,
) -> None:
    """Transient Modal rate limiting should back off and retry instead of failing immediately."""

    class FakeResourceExhaustedError(Exception):
        """Stand-in for modal.exception.ResourceExhaustedError."""

    class FakeVolume:
        """Minimal Modal volume double that rate limits once."""

        def __init__(self) -> None:
            """Initialize the listdir attempt counter."""
            self.listdir_calls = 0

        def listdir(self, remote_path: str, recursive: bool = False) -> list[str]:
            """Raise once, then succeed."""
            self.listdir_calls += 1
            if self.listdir_calls == 1:
                raise FakeResourceExhaustedError("rate limited")
            return [remote_path]

    fake_volume = FakeVolume()

    class FakeModal:
        """Minimal modal SDK double exposing retryable error types."""

        exception = type(
            "FakeExceptionNamespace",
            (),
            {
                "NotFoundError": FileNotFoundError,
                "ResourceExhaustedError": FakeResourceExhaustedError,
            },
        )

        class Volume:
            """Namespace for volume lookups."""

            @staticmethod
            def from_name(name: str, create_if_missing: bool = False) -> FakeVolume:
                """Return the fake volume for any lookup."""
                return fake_volume

    sleep_calls: list[float] = []
    monkeypatch.setattr(sync_engine_module, "modal", FakeModal)
    monkeypatch.setattr(sync_engine_module.time, "sleep", lambda seconds: sleep_calls.append(seconds))
    backend = sync_engine_module.ModalVolumeBackend("volume")

    assert backend.exists("/hashes/present.done") is True
    assert fake_volume.listdir_calls == 2
    assert sleep_calls == [0.25]