gpu-rent/tests/test_tunnel_watch.py

import sys
import time

import pytest

from gpu_rent.tunnel import decide_watch, poll_ctrl_d, run_tunnel, tunnel_forwards


def test_decide_ok_active():
    d = decide_watch("ACTIVE", tunnel_alive=True)
    assert d.kind == "ok"


def test_decide_reconnect_when_tunnel_dead():
    d = decide_watch("ACTIVE", tunnel_alive=False)
    assert d.kind == "reconnect"


def test_decide_unshelve_expired():
    for st in ("EXPIRED", "SHELVED", "SHELVED_OFFLOADED"):
        d = decide_watch(st, tunnel_alive=True)
        assert d.kind == "unshelve", st


def test_decide_exit_error():
    d = decide_watch("ERROR", tunnel_alive=True)
    assert d.kind == "exit"


def test_decide_exit_missing():
    d = decide_watch(None, tunnel_alive=False)
    assert d.kind == "exit"


def test_decide_soft_fail_keeps_tunnel():
    d = decide_watch("SOFT_FAIL", tunnel_alive=True)
    assert d.kind == "ok"
    assert "soft-fail" in d.detail


def test_tunnel_forwards_swarm_only():
    class Cfg:
        swarmui_local_port = 17801
        llm_runtime = "none"
        ollama_local_port = 17811
        enable_swarmui = True

    assert tunnel_forwards(Cfg()) == [(17801, 7801)]


def test_tunnel_forwards_prefers_cfg_over_stale_notes():
    """tunnel_forwards uses cfg only — notes must not add Ollama."""
    class Cfg:
        swarmui_local_port = 17801
        llm_runtime = "none"
        ollama_local_port = 17811
        enable_swarmui = True

    assert tunnel_forwards(Cfg()) == [(17801, 7801)]


def test_resolve_llm_uses_cfg_only(monkeypatch):
    from gpu_rent.access_card import resolve_llm_runtime

    class Cfg:
        llm_runtime = "none"

    monkeypatch.setattr(
        "gpu_rent.access_card.load_state",
        lambda: type("S", (), {"notes": {"llm_runtime": "ollama"}})(),
    )
    # Stale notes must not override live cfg=none
    assert resolve_llm_runtime(Cfg()) == "none"

    class Cfg2:
        llm_runtime = "ollama"

    assert resolve_llm_runtime(Cfg2()) == "ollama"


def test_poll_ctrl_d_skips_when_not_tty(monkeypatch):
    monkeypatch.setattr("gpu_rent.tunnel.sys.stdin.isatty", lambda: False)
    t0 = time.time()
    assert poll_ctrl_d(0.02) is False
    assert time.time() - t0 < 0.5


def test_poll_ctrl_d_windows_eot(monkeypatch):
    class Msvcrt:
        def kbhit(self) -> bool:
            return True

        def getch(self) -> bytes:
            return b"\x04"

    monkeypatch.setattr("gpu_rent.tunnel.sys.stdin.isatty", lambda: True)
    monkeypatch.setattr("gpu_rent.tunnel.os.name", "nt")
    monkeypatch.setitem(sys.modules, "msvcrt", Msvcrt())
    assert poll_ctrl_d(0.2) is True


def test_poll_ctrl_d_windows_ctrl_c(monkeypatch):
    class Msvcrt:
        def kbhit(self) -> bool:
            return True

        def getch(self) -> bytes:
            return b"\x03"

    monkeypatch.setattr("gpu_rent.tunnel.sys.stdin.isatty", lambda: True)
    monkeypatch.setattr("gpu_rent.tunnel.os.name", "nt")
    monkeypatch.setitem(sys.modules, "msvcrt", Msvcrt())
    with pytest.raises(KeyboardInterrupt):
        poll_ctrl_d(0.2)


class _Fwd:
    is_active = True

    def stop(self) -> None:
        pass


class _State:
    notes: dict = {}
    phase = ""
    floating_ip = "1.2.3.4"


class _Cfg:
    ssh_user = "u"
    ssh_private_key_path = "k"
    swarmui_local_port = 17801
    ollama_local_port = 17811
    llm_runtime = "none"
    enable_swarmui = True


def _stub_tunnel(monkeypatch) -> None:
    st = _State()
    monkeypatch.setattr("gpu_rent.tunnel._ssh_tunnel_forwarder", lambda: object)
    monkeypatch.setattr("gpu_rent.tunnel._start_forwarder", lambda *a, **k: _Fwd())
    monkeypatch.setattr("gpu_rent.ready.verify_stack_local", lambda *a, **k: [])
    monkeypatch.setattr("gpu_rent.llm_runtime.maybe_warmup_ollama_local", lambda *a, **k: None)
    monkeypatch.setattr("gpu_rent.tunnel.load_state", lambda: st)
    monkeypatch.setattr("gpu_rent.tunnel.save_state", lambda s: None)
    monkeypatch.setattr("gpu_rent.access_card.print_access_card", lambda *a, **k: None)
    monkeypatch.setattr("gpu_rent.vm_logs.print_log_digest", lambda *a, **k: None)
    monkeypatch.setattr("gpu_rent.local_watchdog.watchdog_installed", lambda: False)
    monkeypatch.setattr("gpu_rent.local_watchdog.clear_lease", lambda: None)


def test_run_tunnel_ctrl_d_stops_gpu(monkeypatch):
    _stub_tunnel(monkeypatch)
    logs: list[str] = []
    stopped: list[bool] = []
    run_tunnel(
        _Cfg(),
        "1.2.3.4",
        log=logs.append,
        stop_gpu=lambda: stopped.append(True),
        session_end_poll=lambda _t: True,
        poll_seconds=999,
    )
    assert stopped == [True]
    assert any("GPU остановлен" in x for x in logs)


def test_run_tunnel_ctrl_c_stops_gpu(monkeypatch):
    _stub_tunnel(monkeypatch)
    logs: list[str] = []
    stopped: list[bool] = []

    def boom() -> None:
        raise KeyboardInterrupt

    run_tunnel(
        _Cfg(),
        "1.2.3.4",
        log=logs.append,
        wait=boom,
        stop_gpu=lambda: stopped.append(True),
    )
    assert stopped == [True]
    assert any("GPU остановлен" in x for x in logs)
    assert not any("GPU жив" in x for x in logs)