gpu-rent/tests/test_warm_up_skips.py

"""Warm re-up shortcuts: no git cascade, killer reuse, seed/ollama skips."""

from pathlib import Path

from gpu_rent.idle_killer import _remote_killer_reusable
from gpu_rent.provision import _models_manifest_fp, seed_civitai


def test_models_manifest_fp_stable(tmp_path):
    man = tmp_path / "models.yaml"
    man.write_text("models:\n  - kind: checkpoint\n", encoding="utf-8")

    class Cfg:
        models_manifest = man

    a = _models_manifest_fp(Cfg())
    b = _models_manifest_fp(Cfg())
    assert a == b
    assert len(a) == 16
    man.write_text("models:\n  - kind: lora\n", encoding="utf-8")
    assert _models_manifest_fp(Cfg()) != a


def test_seed_civitai_skips_when_cache_dests_present(tmp_path, monkeypatch):
    man = tmp_path / "models.yaml"
    man.write_text("models:\n  - kind: checkpoint\n", encoding="utf-8")
    from gpu_rent import provision

    class Entry:
        kind = "checkpoint"
        url = "https://huggingface.co/org/model/resolve/main/a.safetensors"
        version_id = None

    monkeypatch.setattr(provision, "parse_models", lambda _p: [Entry()])

    def fake_ssh(cfg, host, cmd, **kw):
        if ".gpu-rent-models-jobs.json" in cmd and "cat" in cmd:
            import json

            return json.dumps(
                {
                    "fp": provision._models_manifest_fp(cfg),
                    "dests": ["/mnt/swarm_data/Models/Stable-Diffusion/a.safetensors"],
                }
            )
        if "ALL_OK" in cmd or "test -f" in cmd:
            return "ALL_OK"
        return ""

    class Cfg:
        models_manifest = man
        civitai_api_token = "t"
        civitai_api_host = "https://civitai.com"
        hf_token = ""

    logs = []
    monkeypatch.setattr(provision, "run_ssh", fake_ssh)
    seed_civitai(Cfg(), "1.2.3.4", logs.append)
    assert any("skip" in x.lower() for x in logs)


def test_remote_killer_reusable_parses_ok(monkeypatch):
    class Cfg:
        pass

    def fake_ssh(cfg, host, cmd, **kw):
        return (
            "sid=abc-123\n"
            "want=abc-123\n"
            "creds=yes\n"
            "armed=yes\n"
            "rootcreds=yes\n"
            "timer=active\n"
        )

    from gpu_rent import idle_killer

    monkeypatch.setattr(idle_killer, "run_ssh", fake_ssh)
    assert _remote_killer_reusable(Cfg(), "h", "abc-123") is True
    assert _remote_killer_reusable(Cfg(), "h", "other") is False


def test_install_ollama_skips_restart_when_unit_unchanged():
    from importlib.resources import files

    text = files("gpu_rent.remote").joinpath("install_ollama.sh").read_text(encoding="utf-8")
    assert "cmp -s" in text
    assert "skip restart" in text


def test_cli_has_update_flag():
    from gpu_rent import cli

    src = Path(cli.__file__).read_text(encoding="utf-8")
    assert '"--update"' in src
    assert "force_update" in src


def test_autocomplete_merge_sets_is_installed():
    from gpu_rent import provision
    from gpu_rent.provision import _AUTOCOMPLETE_MERGE_PY, _ENSURE_INSTALLED_PY

    assert "IsInstalled: true" in _AUTOCOMPLETE_MERGE_PY
    assert "IsInstalled: true" in _ENSURE_INSTALLED_PY
    assert "comfy_on_disk" in _AUTOCOMPLETE_MERGE_PY
    assert "comfy_on_disk" in _ENSURE_INSTALLED_PY
    assert r"\1false" in _AUTOCOMPLETE_MERGE_PY
    assert r"\1false" in _ENSURE_INSTALLED_PY
    assert "GPU_RENT_COMFY_PRESENT" in _AUTOCOMPLETE_MERGE_PY
    assert "GPU_RENT_COMFY_PRESENT" in _ENSURE_INSTALLED_PY
    assert hasattr(provision, "ensure_settings_is_installed")


def test_ensure_installed_clears_flag_without_comfy(tmp_path):
    import os
    import re
    import subprocess
    import sys

    from gpu_rent.provision import _ENSURE_INSTALLED_PY

    settings = tmp_path / "Settings.fds"
    settings.write_text("IsInstalled: true\nTheme: modern_dark\n", encoding="utf-8")
    script = tmp_path / "ensure.py"
    script.write_text(_ENSURE_INSTALLED_PY, encoding="utf-8")
    env = os.environ.copy()
    env["GPU_RENT_SETTINGS_FDS"] = str(settings)
    env["GPU_RENT_COMFY_PRESENT"] = "0"
    out = subprocess.check_output([sys.executable, str(script)], env=env, text=True)
    assert "IsInstalled false" in out or "cleared" in out
    text = settings.read_text(encoding="utf-8")
    assert re.search(r"(?im)^\s*IsInstalled:\s*false\s*$", text)
    assert "Theme: modern_dark" in text


def test_ensure_installed_sets_true_when_comfy_present(tmp_path):
    import os
    import re
    import subprocess
    import sys

    from gpu_rent.provision import _ENSURE_INSTALLED_PY

    settings = tmp_path / "Settings.fds"
    settings.write_text("IsInstalled: false\n", encoding="utf-8")
    script = tmp_path / "ensure.py"
    script.write_text(_ENSURE_INSTALLED_PY, encoding="utf-8")
    env = os.environ.copy()
    env["GPU_RENT_SETTINGS_FDS"] = str(settings)
    env["GPU_RENT_COMFY_PRESENT"] = "1"
    subprocess.check_call([sys.executable, str(script)], env=env)
    text = settings.read_text(encoding="utf-8")
    assert re.search(r"(?im)^\s*IsInstalled:\s*true\s*$", text)


def test_install_comfy_patches_is_installed_on_skip():
    from importlib.resources import files

    text = files("gpu_rent.remote").joinpath("install_swarm_comfy.py").read_text(
        encoding="utf-8"
    )
    assert "ensure_settings_installed_flag" in text
    assert "leave /Install" in text
    assert "prepare_first_install" in text
    assert "already installed" in text
    assert "IsInstalled=false" in text
    assert "Open SwarmUI" not in text


def test_clear_settings_installed_flag(tmp_path, monkeypatch):
    from gpu_rent.remote import install_swarm_comfy as inst

    settings = tmp_path / "Settings.fds"
    settings.write_text("IsInstalled: true\nTheme: x\n", encoding="utf-8")
    monkeypatch.setattr(inst, "SETTINGS", settings)
    assert inst.clear_settings_installed_flag() is True
    text = settings.read_text(encoding="utf-8")
    assert "IsInstalled: false" in text
    assert "Theme: x" in text
    assert inst.clear_settings_installed_flag() is False