gpu-rent/tests/test_wait_backend_idle.py

"""Tests for wait_backend_idle READY semantics and install helpers."""

from gpu_rent.ready import _REMOTE_POLL


def test_remote_poll_empty_is_busy_not_ready():
    assert 'bstat == "empty"' in _REMOTE_POLL
    assert "BUSY backend=empty" in _REMOTE_POLL


def test_remote_poll_running_is_ready():
    """SwarmUI: running = healthy ready."""
    assert 'bstat == "running"' in _REMOTE_POLL
    assert "READY backend=running" in _REMOTE_POLL


def test_remote_poll_idle_suspended_is_ready():
    """Suspended idle backends are installed — ready for up (wake on gen)."""
    assert "READY backend=idle" in _REMOTE_POLL
    assert "BUSY backend=idle" not in _REMOTE_POLL


def test_remote_poll_loading_is_busy():
    assert 'bstat in ("loading", "some_loading")' in _REMOTE_POLL
    assert "Comfy стартует" in _REMOTE_POLL


def test_install_swarm_comfy_script_payload():
    from importlib.resources import files

    text = files("gpu_rent.remote").joinpath("install_swarm_comfy.py").read_text(
        encoding="utf-8"
    )
    assert "InstallConfirmWS" in text
    assert '"backend": "comfyui"' in text
    assert '"models": "none"' in text
    assert "modern_dark" in text
    assert "backends present (idle/suspended)" in text
    assert ".gpu-rent-comfy-installing" in text
    assert "RestartBackends" in text
    assert 'bstat == "errored"' in text
    assert "sanitize_backends_fds" in text or "recover_empty_backends" in text
    assert "AddNewBackend" in text
    assert "EditBackend" in text
    assert "configure_comfy_backend" in text
    assert "ensure_dlbackend_bind" in text
    assert "DeleteBackend" in text or "delete_backend" in text
    assert "absolute StartScript" in text or "reconfigure errored" in text
    assert "/mnt/swarm_data" in text or "DATA /" in text
    assert 'bstat in ("empty", "disabled", "all_disabled", "unknown")' in text


def test_wait_backend_idle_fail_fast_on_long_loading(monkeypatch):
    from gpu_rent.errors import CloudError
    from gpu_rent import ready

    class Cfg:
        pass

    def fake_ssh(*a, **k):
        cmd = a[2] if len(a) > 2 else ""
        if "journalctl" in str(cmd):
            return "loading models…"
        return "BUSY backend=loading (Comfy стартует)"

    diag_calls = {"n": 0}

    def fake_diag(*a, **k):
        diag_calls["n"] += 1
        return "DIAG ok"

    monkeypatch.setattr(ready, "run_ssh", fake_ssh)
    monkeypatch.setattr(ready, "collect_swarm_diagnostics", fake_diag)
    try:
        ready.wait_backend_idle(
            Cfg(),
            "1.2.3.4",
            [].append,
            timeout=600.0,
            poll_every=0.01,
            loading_fail_sec=0.05,
        )
        assert False, "expected CloudError"
    except CloudError as exc:
        assert "loading" in str(exc).lower()
    assert diag_calls["n"] == 1


def test_install_script_uses_none_frontend():
    from importlib.resources import files

    text = files("gpu_rent.remote").joinpath("install_swarm_comfy.py").read_text(
        encoding="utf-8"
    )
    assert '"FrontendVersion": "None"' in text
    assert "stop_for_remount" in text or "перед remount" in text


def test_wait_backend_idle_fail_fast_on_disabled(monkeypatch):
    from gpu_rent.errors import CloudError
    from gpu_rent import ready

    class Cfg:
        pass

    def fake_ssh(*a, **k):
        return "BUSY backend=disabled"

    diag_calls = {"n": 0}

    def fake_diag(*a, **k):
        diag_calls["n"] += 1
        return "DIAG ok"

    monkeypatch.setattr(ready, "run_ssh", fake_ssh)
    monkeypatch.setattr(ready, "collect_swarm_diagnostics", fake_diag)
    try:
        ready.wait_backend_idle(
            Cfg(),
            "1.2.3.4",
            [].append,
            timeout=600.0,
            poll_every=0.01,
            disabled_fail_sec=0.05,
        )
        assert False, "expected CloudError"
    except CloudError as exc:
        assert "disabled" in str(exc).lower()
        assert "startscript" in str(exc).lower()
    assert diag_calls["n"] == 1


def test_swarm_diag_script_covers_api_and_journal():
    from importlib.resources import files

    text = files("gpu_rent.remote").joinpath("swarm_diag.py").read_text(encoding="utf-8")
    assert "ListBackends" in text
    assert "journalctl" in text
    assert ".gpu-rent-last-diag.txt" in text
    assert "nvidia-smi" in text


def test_verify_gpu_env_fail_fast_empty_dlbackend(monkeypatch):
    import json

    from gpu_rent.errors import CloudError
    from gpu_rent.ready import verify_gpu_env
    import gpu_rent.ssh_ops as ssh_ops

    class Cfg:
        enable_swarmui = True

    payload = {
        "ok": False,
        "checks": [
            {"name": "nvidia-smi", "required": True, "ok": True, "detail": "ok"},
            {"name": "cuda", "required": True, "ok": True, "detail": "ok"},
            {
                "name": "torch",
                "required": True,
                "ok": False,
                "detail": "ComfyUI venv python не найден (dlbackend пуст — SwarmUI Install не прогоняли)",
            },
        ],
    }
    calls = {"n": 0}

    def fake(*a, **k):
        calls["n"] += 1
        return json.dumps(payload)

    monkeypatch.setattr(ssh_ops, "run_python", fake)
    logs: list[str] = []
    try:
        verify_gpu_env(Cfg(), "1.2.3.4", logs.append, timeout=600.0, poll_every=0.1)
        assert False, "expected CloudError"
    except CloudError as exc:
        assert "fail-fast" in str(exc).lower() or "torch" in str(exc).lower()
    assert calls["n"] == 1


def test_verify_gpu_env_fail_fast_torch_no_cuda(monkeypatch):
    import json

    from gpu_rent.errors import CloudError
    from gpu_rent.ready import verify_gpu_env
    import gpu_rent.ssh_ops as ssh_ops

    class Cfg:
        enable_swarmui = True

    payload = {
        "ok": False,
        "checks": [
            {"name": "nvidia-smi", "required": True, "ok": True, "detail": "ok"},
            {"name": "cuda", "required": True, "ok": True, "detail": "ok"},
            {
                "name": "torch",
                "required": True,
                "ok": False,
                "detail": "torch=2.0 cuda=None available=false (CPU wheel / без cuda — не заживёт само)",
            },
        ],
    }
    calls = {"n": 0}

    def fake(*a, **k):
        calls["n"] += 1
        return json.dumps(payload)

    monkeypatch.setattr(ssh_ops, "run_python", fake)
    try:
        verify_gpu_env(Cfg(), "1.2.3.4", [].append, timeout=600.0, poll_every=0.1)
        assert False, "expected CloudError"
    except CloudError:
        pass
    assert calls["n"] == 1