gpu-rent/tests/test_hold_killer.py

"""Tests for hold parsing and idle-killer busy classification."""

from __future__ import annotations

import importlib.util
from datetime import datetime, timezone
from pathlib import Path

import pytest

from gpu_rent.errors import GpuRentError
from gpu_rent.hold import _parse_until


ROOT = Path(__file__).resolve().parents[1]
REMOTE_KILLER = ROOT / "src" / "gpu_rent" / "remote" / "idle_killer.py"


def _load_remote():
    spec = importlib.util.spec_from_file_location("idle_killer_remote", REMOTE_KILLER)
    assert spec and spec.loader
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


def test_parse_until_iso():
    ts = _parse_until("2030-01-01T00:00:00+00:00")
    assert ts == int(datetime(2030, 1, 1, tzinfo=timezone.utc).timestamp())


def test_parse_until_unix():
    assert _parse_until("1700000000") == 1700000000


def test_parse_until_bad():
    with pytest.raises(GpuRentError):
        _parse_until("not-a-date")


def test_classify_busy_from_status(monkeypatch):
    mod = _load_remote()

    class FakeResp:
        def __init__(self, payload):
            self._payload = payload

        def read(self):
            import json

            return json.dumps(self._payload).encode()

        def __enter__(self):
            return self

        def __exit__(self, *args):
            return False

    calls = {"n": 0}

    def fake_urlopen(req, timeout=0, context=None):
        calls["n"] += 1
        url = getattr(req, "full_url", None) or req.get_full_url()
        if "GetNewSession" in url:
            return FakeResp({"session_id": "abc"})
        return FakeResp(
            {
                "status": {"waiting_gens": 0, "live_gens": 0, "loading_models": 0},
                "backend_status": {"status": "idle"},
            }
        )

    monkeypatch.setattr(mod.urllib.request, "urlopen", fake_urlopen)
    busy, detail = mod.swarm_busy("http://127.0.0.1:7801")
    assert busy is False
    assert "idle" in detail


def test_classify_running_without_queue_not_busy(monkeypatch):
    """SwarmUI 'running' = ready; empty queue → idle-killer may stop GPU."""
    mod = _load_remote()

    class FakeResp:
        def __init__(self, payload):
            self._payload = payload

        def read(self):
            import json

            return json.dumps(self._payload).encode()

        def __enter__(self):
            return self

        def __exit__(self, *args):
            return False

    def fake_urlopen(req, timeout=0, context=None):
        url = getattr(req, "full_url", None) or req.get_full_url()
        if "GetNewSession" in url:
            return FakeResp({"session_id": "abc"})
        return FakeResp(
            {
                "status": {"waiting_gens": 0, "live_gens": 0, "loading_models": 0},
                "backend_status": {"status": "running", "any_loading": False},
            }
        )

    monkeypatch.setattr(mod.urllib.request, "urlopen", fake_urlopen)
    busy, detail = mod.swarm_busy("http://127.0.0.1:7801")
    assert busy is False
    assert "running" in detail


def test_classify_loading_is_busy(monkeypatch):
    mod = _load_remote()

    class FakeResp:
        def __init__(self, payload):
            self._payload = payload

        def read(self):
            import json

            return json.dumps(self._payload).encode()

        def __enter__(self):
            return self

        def __exit__(self, *args):
            return False

    def fake_urlopen(req, timeout=0, context=None):
        url = getattr(req, "full_url", None) or req.get_full_url()
        if "GetNewSession" in url:
            return FakeResp({"session_id": "abc"})
        return FakeResp(
            {
                "status": {"waiting_gens": 0, "live_gens": 0, "loading_models": 0},
                "backend_status": {"status": "loading", "any_loading": True},
            }
        )

    monkeypatch.setattr(mod.urllib.request, "urlopen", fake_urlopen)
    busy, detail = mod.swarm_busy("http://127.0.0.1:7801")
    assert busy is True
    assert "loading" in detail


def test_classify_busy_queue(monkeypatch):
    mod = _load_remote()

    class FakeResp:
        def __init__(self, payload):
            self._payload = payload

        def read(self):
            import json

            return json.dumps(self._payload).encode()

        def __enter__(self):
            return self

        def __exit__(self, *args):
            return False

    def fake_urlopen(req, timeout=0, context=None):
        url = getattr(req, "full_url", None) or req.get_full_url()
        if "GetNewSession" in url:
            return FakeResp({"session_id": "abc"})
        return FakeResp(
            {
                "status": {"waiting_gens": 2, "live_gens": 0, "loading_models": 0},
                "backend_status": {"status": "idle"},
            }
        )

    monkeypatch.setattr(mod.urllib.request, "urlopen", fake_urlopen)
    busy, detail = mod.swarm_busy("http://127.0.0.1:7801")
    assert busy is True
    assert "waiting=2" in detail


def test_swarm_unreachable_starts_busy_then_allows_idle(tmp_path, monkeypatch):
    mod = _load_remote()
    monkeypatch.setattr(mod, "DATA", tmp_path)
    monkeypatch.setattr(mod, "SWARM_DOWN_SINCE", tmp_path / ".gpu-rent-swarm-down-since")
    monkeypatch.setattr(mod, "SWARM_UNREACHABLE_IDLE_SEC", 100)

    def boom(*_a, **_k):
        raise mod.urllib.error.URLError("down")

    monkeypatch.setattr(mod.urllib.request, "urlopen", boom)

    busy1, d1 = mod.swarm_busy("http://127.0.0.1:7801")
    assert busy1 is True
    assert "clock start" in d1
    assert mod.SWARM_DOWN_SINCE.is_file()

    # Still within window
    since = mod.read_ts(mod.SWARM_DOWN_SINCE)
    assert since is not None
    mod.write_ts(mod.SWARM_DOWN_SINCE, since - 50)
    busy2, d2 = mod.swarm_busy("http://127.0.0.1:7801")
    assert busy2 is True
    assert "50s" in d2 or "/ 100s" in d2

    # Past window → not busy so idle clock can run
    mod.write_ts(mod.SWARM_DOWN_SINCE, since - 120)
    busy3, d3 = mod.swarm_busy("http://127.0.0.1:7801")
    assert busy3 is False
    assert "allow idle" in d3