gpu-rent/tests/test_session.py

import pytest

from gpu_rent.config import load_config
from gpu_rent.errors import CloudError, GpuRentError
from gpu_rent.lock import SessionLock
from gpu_rent.paths import lock_path
from gpu_rent.session import cmd_stop, cmd_up
from gpu_rent.state import SessionState, save_state


class Server:
    def __init__(self, status="ACTIVE", server_id="s1"):
        self.id = server_id
        self.name = "gpu-rent"
        self.status = status
        self.flavor = {"id": "f1"}
        self.addresses = {
            "gpu-rent": [{"addr": "203.0.113.9", "OS-EXT-IPS:type": "floating"}]
        }


def _cfg(monkeypatch):
    monkeypatch.setenv("OS_AUTH_URL", "https://example.invalid/identity/v3")
    monkeypatch.setenv("OS_USER_DOMAIN_NAME", "999")
    monkeypatch.setenv("OS_USERNAME", "svc")
    monkeypatch.setenv("OS_PASSWORD", "secret")
    monkeypatch.setenv("OS_PROJECT_ID", "proj")
    monkeypatch.setenv("OS_REGION_NAME", "ru-7")
    monkeypatch.setenv("GPU_RENT_AZ", "ru-7a")
    return load_config(require_auth=True)


def _mock_bind(monkeypatch):
    monkeypatch.setattr(
        "gpu_rent.session.ensure_floating_ip",
        lambda conn, server, existing_id, existing_addr, log: ("203.0.113.9", "fip1"),
    )
    monkeypatch.setattr("gpu_rent.session.wait_ssh", lambda cfg, host, timeout=900.0, log=None: None)
    monkeypatch.setattr(
        "gpu_rent.session.run_ssh",
        lambda cfg, host, command, **kw: (
            "yes" if "gpu-rent-bootstrapped" in command else "inactive"
        ),
    )
    monkeypatch.setattr(
        "gpu_rent.session.run_bootstrap",
        lambda cfg, host, log, update=True, light=False: None,
    )
    monkeypatch.setattr("gpu_rent.session.provision_vm", lambda cfg, host, log, **kw: None)
    monkeypatch.setattr(
        "gpu_rent.session.ensure_swarm_comfy_installed",
        lambda cfg, host, log: None,
    )
    monkeypatch.setattr(
        "gpu_rent.session.seed_autocomplete",
        lambda cfg, host, log: False,
    )
    monkeypatch.setattr("gpu_rent.session.wait_backend_idle", lambda cfg, host, log, **kw: None)
    monkeypatch.setattr(
        "gpu_rent.session.verify_stack_on_vm",
        lambda cfg, host, log, **kw: [],
    )
    monkeypatch.setattr(
        "gpu_rent.session.verify_gpu_env",
        lambda cfg, host, log, **kw: [],
    )
    monkeypatch.setattr(
        "gpu_rent.session.tune_swarm_perf",
        lambda cfg, host, log: False,
    )
    monkeypatch.setattr(
        "gpu_rent.session.ensure_data_binds",
        lambda cfg, host, log, **kw: None,
    )
    monkeypatch.setattr(
        "gpu_rent.session.ensure_boot_snapshot",
        lambda conn, boot_volume_id, cfg, log: None,
    )
    monkeypatch.setattr("gpu_rent.session.notify_ready", lambda cfg, log: None)


def test_cmd_up_refuses_zero_gpu_quota(monkeypatch):
    monkeypatch.setattr("gpu_rent.session.connect", lambda cfg: object())
    monkeypatch.setattr("gpu_rent.session.compute_quotas", lambda conn: {"gpu": 0})
    with pytest.raises(CloudError, match="квота GPU"):
        cmd_up(_cfg(monkeypatch), yes=True)


def test_cmd_up_does_not_create_second_gpu(monkeypatch):
    created = []
    monkeypatch.setattr("gpu_rent.session.connect", lambda cfg: object())
    monkeypatch.setattr("gpu_rent.session.compute_quotas", lambda conn: {"gpu": 1})
    monkeypatch.setattr("gpu_rent.session.pick_existing_server", lambda conn: Server())
    _mock_bind(monkeypatch)
    monkeypatch.setattr(
        "gpu_rent.session.create_gpu_server",
        lambda *a, **k: created.append("created") or Server(),
    )
    save_state(
        SessionState(
            server_id="s1",
            floating_ip="203.0.113.9",
            bootstrapped=True,
            phase="ready_cloud",
        )
    )
    state = cmd_up(_cfg(monkeypatch), yes=True)
    assert created == []
    assert state.server_id == "s1"
    assert state.phase == "ready_cloud"
    assert state.floating_ip == "203.0.113.9"
    assert state.bootstrapped is True


def test_cmd_up_unshelves_expired(monkeypatch):
    unshelved = []
    created = []
    monkeypatch.setattr("gpu_rent.session.connect", lambda cfg: object())
    monkeypatch.setattr("gpu_rent.session.compute_quotas", lambda conn: {"gpu": 1})
    monkeypatch.setattr(
        "gpu_rent.session.pick_existing_server", lambda conn: Server(status="EXPIRED")
    )
    monkeypatch.setattr(
        "gpu_rent.session.unshelve",
        lambda conn, server, log: unshelved.append(server.id) or Server(status="ACTIVE"),
    )
    _mock_bind(monkeypatch)
    monkeypatch.setattr(
        "gpu_rent.session.create_gpu_server",
        lambda *a, **k: created.append("created"),
    )
    state = cmd_up(_cfg(monkeypatch), yes=True)
    assert unshelved == ["s1"]
    assert created == []
    assert state.unshelved_at
    assert state.phase == "ready_cloud"
    assert state.bootstrapped is True


def test_cmd_stop_deletes_compute_keeps_disks(monkeypatch):
    deleted = []
    monkeypatch.setattr(
        "gpu_rent.session.delete_server",
        lambda conn, server, log: deleted.append(server.id),
    )
    monkeypatch.setattr(
        "gpu_rent.session.delete_floating_ip",
        lambda conn, fip_id, address, log: deleted.append("fip"),
    )
    save_state(
        SessionState(
            server_id="s1",
            boot_volume_id="b1",
            data_volume_id="d1",
            floating_ip="1.1.1.1",
            bootstrapped=True,
        )
    )

    class Conn:
        class compute:
            @staticmethod
            def get_server(sid):
                return Server(server_id=sid)

    monkeypatch.setattr("gpu_rent.session.connect", lambda cfg: Conn())
    monkeypatch.setattr("gpu_rent.session.pick_existing_server", lambda conn: Server())
    state = cmd_stop(_cfg(monkeypatch))
    assert "s1" in deleted
    assert state.phase == "idle"
    assert state.server_id is None
    assert state.bootstrapped is False
    assert state.boot_volume_id == "b1"
    assert state.data_volume_id == "d1"


def test_lock_busy(monkeypatch):
    lock_path().parent.mkdir(parents=True, exist_ok=True)
    lock_path().write_text("1", encoding="utf-8")
    monkeypatch.setattr("gpu_rent.lock._pid_alive", lambda pid: True)
    monkeypatch.setattr("gpu_rent.lock.os.getpid", lambda: 99)
    with pytest.raises(GpuRentError, match="уже работает"):
        SessionLock().__enter__()


def test_bind_access_keeps_provision_llm_notes(monkeypatch):
    """In-memory SessionState must not resurrect llm_error after provision_llm saved."""
    from gpu_rent.ready import ServiceCheck
    from gpu_rent.state import load_state

    def fake_provision(cfg, host, log, **kw):
        st = load_state()
        st.notes = dict(st.notes or {})
        st.notes["llm_runtime"] = "ollama"
        st.notes.pop("llm_error", None)
        save_state(st)

    monkeypatch.setattr("gpu_rent.session.connect", lambda cfg: object())
    monkeypatch.setattr("gpu_rent.session.compute_quotas", lambda conn: {"gpu": 1})
    monkeypatch.setattr("gpu_rent.session.pick_existing_server", lambda conn: Server())
    _mock_bind(monkeypatch)
    monkeypatch.setattr("gpu_rent.session.provision_vm", fake_provision)
    monkeypatch.setattr(
        "gpu_rent.session.verify_stack_on_vm",
        lambda cfg, host, log, **kw: [
            ServiceCheck("swarmui", True, "HTTP 200"),
            ServiceCheck("ollama", True, "1 models (qwen)"),
        ],
    )
    monkeypatch.setattr(
        "gpu_rent.session.create_gpu_server",
        lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not create")),
    )
    save_state(
        SessionState(
            server_id="s1",
            floating_ip="203.0.113.9",
            bootstrapped=True,
            phase="ready_cloud",
            notes={
                "llm_error": "Ollama /api/tags без моделей из ollama-models.yaml",
                "llm_runtime": "none",
            },
        )
    )
    state = cmd_up(_cfg(monkeypatch), yes=True)
    assert state.notes.get("llm_error") is None
    assert state.notes.get("llm_runtime") == "ollama"
    disk = load_state()
    assert disk.notes.get("llm_error") is None
    assert disk.notes.get("llm_runtime") == "ollama"