fn_registry/python/functions/infra/comfyui_ensure_server_test.py

"""Tests para comfyui_ensure_server.

Los tests no tocan systemd ni la red reales: inyectan un runner falso que
registra los comandos systemctl y se mockea el health check.
"""

import os
import subprocess

from . import comfyui_ensure_server as mod
from .comfyui_ensure_server import (
    _detect_lowvram,
    _render_unit,
    comfyui_ensure_server,
)


class FakeRunner:
    """Runner inyectable: respuestas programables por prefijo de comando."""

    def __init__(self, active_state="inactive"):
        self.calls = []
        self.active_state = active_state

    def __call__(self, cmd):
        self.calls.append(list(cmd))
        # nvidia-smi VRAM
        if cmd[:1] == ["nvidia-smi"]:
            return subprocess.CompletedProcess(cmd, 0, stdout="8192\n", stderr="")
        if cmd[:2] == ["systemctl", "--user"]:
            sub = cmd[2] if len(cmd) > 2 else ""
            if sub == "is-active":
                return subprocess.CompletedProcess(
                    cmd, 0, stdout=self.active_state + "\n", stderr=""
                )
            # daemon-reload, enable, start -> exito
            return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="")
        if cmd[:1] == ["ss"]:
            return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="")
        return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="")

    def ran(self, *needle):
        return any(call[: len(needle)] == list(needle) for call in self.calls)


def _fake_comfy_dir(tmp_path):
    """Crea un comfyui_dir falso con .venv/bin/python y main.py."""
    d = tmp_path / "ComfyUI"
    (d / ".venv" / "bin").mkdir(parents=True)
    (d / ".venv" / "bin" / "python").write_text("#!/bin/sh\n")
    (d / "main.py").write_text("# fake\n")
    return d


# --- helpers puros ---

def test_detect_lowvram_umbral_8gb():
    assert _detect_lowvram(8192) is True
    assert _detect_lowvram(8200) is True
    assert _detect_lowvram(8201) is False
    assert _detect_lowvram(24564) is False
    assert _detect_lowvram(None) is False


def test_render_unit_restart_always_y_wantedby():
    unit = _render_unit(
        "/x/.venv/bin/python", "/x/main.py", "/x", 8188, True, "ComfyUI test"
    )
    assert "Restart=always" in unit
    assert "on-failure" not in unit  # regla function_tags
    assert "WantedBy=default.target" in unit
    assert "ExecStart=/x/.venv/bin/python /x/main.py --port 8188 --lowvram" in unit
    assert "WorkingDirectory=/x" in unit


def test_render_unit_sin_lowvram():
    unit = _render_unit(
        "/x/.venv/bin/python", "/x/main.py", "/x", 9000, False, "ComfyUI test"
    )
    assert "--lowvram" not in unit
    assert "--port 9000" in unit


# --- orquestacion (runner falso + health mockeado) ---

def test_error_si_falta_venv(tmp_path, monkeypatch):
    monkeypatch.setenv("HOME", str(tmp_path))
    res = comfyui_ensure_server(
        comfyui_dir=str(tmp_path / "no_existe"), runner=FakeRunner()
    )
    assert res["ok"] is False
    assert "venv python no encontrado" in res["error"]


def test_idempotente_si_ya_activo_y_sano(tmp_path, monkeypatch):
    monkeypatch.setenv("HOME", str(tmp_path))
    monkeypatch.setattr(mod, "_health", lambda *a, **k: True)
    d = _fake_comfy_dir(tmp_path)
    runner = FakeRunner(active_state="active")

    # 1a llamada: instala el unit por primera vez (changed -> idempotent False),
    # pero como ya esta active+sano NO debe arrancar nada.
    res1 = comfyui_ensure_server(comfyui_dir=str(d), runner=runner)
    assert res1["ok"] is True
    assert res1["health"] is True
    assert res1["active"] == "active"
    assert res1["idempotent"] is False  # escribio el unit por primera vez
    assert not runner.ran("systemctl", "--user", "start", "comfyui")

    # 2a llamada: el unit ya existe identico -> no toca nada -> idempotent True.
    runner2 = FakeRunner(active_state="active")
    res2 = comfyui_ensure_server(comfyui_dir=str(d), runner=runner2)
    assert res2["ok"] is True
    assert res2["idempotent"] is True
    assert res2["reloaded"] is False  # no reescribio el unit
    assert not runner2.ran("systemctl", "--user", "start", "comfyui")


def test_arranque_fresco_escribe_unit_y_arranca(tmp_path, monkeypatch):
    monkeypatch.setenv("HOME", str(tmp_path))
    # health: False antes de arrancar, True despues
    estados = iter([False, True, True, True])
    monkeypatch.setattr(mod, "_health", lambda *a, **k: next(estados, True))
    d = _fake_comfy_dir(tmp_path)
    runner = FakeRunner(active_state="inactive")
    res = comfyui_ensure_server(comfyui_dir=str(d), runner=runner, health_timeout=5)
    assert res["ok"] is True
    assert res["health"] is True
    assert res["lowvram"] is True  # nvidia-smi falso devuelve 8192
    assert runner.ran("systemctl", "--user", "daemon-reload")
    assert runner.ran("systemctl", "--user", "enable", "comfyui")
    assert runner.ran("systemctl", "--user", "start", "comfyui")
    # el unit quedo escrito
    unit_path = os.path.join(
        str(tmp_path), ".config", "systemd", "user", "comfyui.service"
    )
    assert os.path.exists(unit_path)
    with open(unit_path) as f:
        assert "Restart=always" in f.read()


def test_lowvram_forzado_false_omite_flag(tmp_path, monkeypatch):
    monkeypatch.setenv("HOME", str(tmp_path))
    estados = iter([False, True])
    monkeypatch.setattr(mod, "_health", lambda *a, **k: next(estados, True))
    d = _fake_comfy_dir(tmp_path)
    runner = FakeRunner(active_state="inactive")
    res = comfyui_ensure_server(
        comfyui_dir=str(d), runner=runner, lowvram=False, health_timeout=5
    )
    assert res["lowvram"] is False
    unit_path = os.path.join(
        str(tmp_path), ".config", "systemd", "user", "comfyui.service"
    )
    with open(unit_path) as f:
        assert "--lowvram" not in f.read()