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: 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.wait_backend_idle", 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)
    monkeypatch.setattr("gpu_rent.session.print_mcp_snippet", 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__()