ttulttul-ComfyUI-Modal/tests/test_ssh_worker.py

"""Tests for persistent SSH worker state and relay behavior."""

from __future__ import annotations

import io
import json
import os
from pathlib import Path
import socket
import subprocess
import sys
from types import SimpleNamespace
from typing import Any


def test_top_level_worker_entrypoint_reports_runtime_info() -> None:
    """The OCI entrypoint must import when the repository is a top-level path."""
    repo_root = Path(__file__).resolve().parents[1]
    environment = dict(os.environ)
    environment["PYTHONPATH"] = str(repo_root)
    environment["COMFY_MODAL_RUNTIME_FINGERPRINT"] = "test-fingerprint"

    completed = subprocess.run(
        [sys.executable, "-m", "remote.ssh_worker", "runtime-info"],
        cwd=repo_root,
        env=environment,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        check=True,
        text=True,
        timeout=15,
    )

    runtime_info = json.loads(completed.stdout)
    assert runtime_info["runtime_fingerprint"] == "test-fingerprint"
    assert runtime_info["protocol_version"] > 0


def test_worker_execution_state_registers_and_cancels(ssh_worker_module: Any) -> None:
    """Cancellation should target a stable invocation identity."""
    state = ssh_worker_module.WorkerExecutionState()

    cancellation = state.register("RIV_test")

    assert state.cancel("RIV_test") is True
    assert cancellation.is_set()
    state.unregister("RIV_test")
    assert state.cancel("RIV_test") is False


def test_worker_stage_profiles_streams_progress_and_result(
    ssh_worker_module: Any,
    llm_staging_module: Any,
    monkeypatch: Any,
    capsys: Any,
    tmp_path: Path,
) -> None:
    """The SSH staging command should expose machine-readable progress and metadata."""

    def stage(model_references: list[str], storage_root: Path, **kwargs: Any) -> Any:
        """Emit one progress update and return one immutable profile result."""
        assert model_references == ["owner/model"]
        assert storage_root == tmp_path
        kwargs["progress_callback"](
            SimpleNamespace(
                stage="download",
                message="Fetching files",
                value=1,
                maximum=2,
                unit="files",
                indeterminate=False,
                model_reference="owner/model",
            )
        )
        return [
            SimpleNamespace(
                to_dict=lambda: {
                    "requested_reference": "owner/model",
                    "profile_id": "hf-" + "b" * 64,
                    "revision": "8" * 40,
                }
            )
        ]

    monkeypatch.setattr(llm_staging_module, "resolve_and_stage_model_references", stage)

    results = ssh_worker_module.stage_profiles(["owner/model"], tmp_path)

    events = [json.loads(line) for line in capsys.readouterr().out.splitlines()]
    assert events[0]["kind"] == "progress"
    assert events[0]["max"] == 2
    assert events[0]["model_reference"] == "owner/model"
    assert events[1] == {"kind": "result", "results": results}


def test_worker_cancels_only_matching_owned_stager(
    ssh_worker_module: Any,
    monkeypatch: Any,
    tmp_path: Path,
) -> None:
    """The cancellation command should validate identity before killing a PID."""
    owner_id = "vast:RIV_test:owner"
    pid = 12345
    owner_path = ssh_worker_module._staging_owner_path(tmp_path, owner_id)
    owner_path.parent.mkdir(parents=True)
    owner_path.write_text(
        json.dumps(
            {
                "owner_id": owner_id,
                "pid": pid,
                "process_start": "start-tick",
            }
        ),
        encoding="utf-8",
    )
    lease_path = tmp_path / "llm_models" / "repo" / ".revision.download.lock"
    lease_path.parent.mkdir(parents=True)
    lease_path.write_text(
        json.dumps(
            {
                "owner_id": owner_id,
                "pid": pid,
                "process_start": "start-tick",
                "token": "lease-token",
            }
        ),
        encoding="utf-8",
    )
    starts = iter(("start-tick", None))
    killed: list[tuple[int, int]] = []
    monkeypatch.setattr(
        ssh_worker_module,
        "_linux_process_start",
        lambda _pid: next(starts),
    )
    monkeypatch.setattr(
        ssh_worker_module.Path,
        "read_bytes",
        lambda _path: (
            b"python\0-m\0remote.ssh_worker\0stage-profiles\0--owner-id\0"
            + owner_id.encode("utf-8")
            + b"\0"
        ),
    )
    monkeypatch.setattr(
        ssh_worker_module.os,
        "kill",
        lambda target_pid, sent_signal: killed.append((target_pid, sent_signal)),
    )

    assert ssh_worker_module.cancel_staging_process(tmp_path, owner_id) is True

    assert not owner_path.exists()
    assert not lease_path.exists()
    assert killed == [(pid, ssh_worker_module.signal.SIGTERM)]


def test_worker_request_relay_preserves_framed_binary_payload(
    ssh_worker_module: Any,
    remote_protocol_module: Any,
) -> None:
    """The docker-exec relay must preserve request and input frames byte-for-byte."""
    request = remote_protocol_module.encode_json_frame(
        remote_protocol_module.RemoteFrameKind.REQUEST,
        {"invocation_id": "RIV_test", "payload": {"payload_kind": "subgraph"}},
    )
    inputs = remote_protocol_module.encode_frame(
        remote_protocol_module.RemoteFrameKind.INPUTS,
        b"\x00\xfftensor-bytes",
    )
    left, right = socket.socketpair()
    try:
        ssh_worker_module._copy_request_frames(io.BytesIO(request + inputs), left)
        received = right.makefile("rb")
        try:
            assert remote_protocol_module.read_frame(received)[0] is remote_protocol_module.RemoteFrameKind.REQUEST
            assert remote_protocol_module.read_frame(received) == (
                remote_protocol_module.RemoteFrameKind.INPUTS,
                b"\x00\xfftensor-bytes",
            )
        finally:
            received.close()
    finally:
        left.close()
        right.close()