tests/test_runtime.py

"""Accelerator cleanup behavior for regular and ZeroGPU processes."""

from __future__ import annotations

from types import SimpleNamespace

from core import runtime


class FakeCuda:
    """Record cleanup calls without importing or initializing CUDA."""

    def __init__(self, initialized: bool) -> None:
        self.initialized = initialized
        self.empty_cache_calls = 0
        self.ipc_collect_calls = 0

    def is_initialized(self) -> bool:
        return self.initialized

    def empty_cache(self) -> None:
        self.empty_cache_calls += 1

    def ipc_collect(self) -> None:
        self.ipc_collect_calls += 1


def test_cleanup_does_not_initialize_cuda(monkeypatch) -> None:
    cuda = FakeCuda(initialized=False)
    monkeypatch.setattr(runtime, "import_torch", lambda: SimpleNamespace(cuda=cuda))

    runtime.cleanup_memory()

    assert cuda.empty_cache_calls == 0
    assert cuda.ipc_collect_calls == 0


def test_cleanup_releases_an_existing_cuda_context(monkeypatch) -> None:
    cuda = FakeCuda(initialized=True)
    monkeypatch.setattr(runtime, "import_torch", lambda: SimpleNamespace(cuda=cuda))

    runtime.cleanup_memory()

    assert cuda.empty_cache_calls == 1
    assert cuda.ipc_collect_calls == 1


def test_cleanup_skips_an_unavailable_mps_backend(monkeypatch) -> None:
    cuda = FakeCuda(initialized=False)
    calls = 0

    def empty_cache() -> None:
        nonlocal calls
        calls += 1

    mps = SimpleNamespace(empty_cache=empty_cache)
    mps_backend = SimpleNamespace(is_available=lambda: False)
    torch = SimpleNamespace(
        cuda=cuda,
        mps=mps,
        backends=SimpleNamespace(mps=mps_backend),
    )
    monkeypatch.setattr(runtime, "import_torch", lambda: torch)

    runtime.cleanup_memory()

    assert calls == 0


def test_cleanup_errors_never_escape(monkeypatch) -> None:
    cuda = FakeCuda(initialized=True)

    def fail_cleanup() -> None:
        raise RuntimeError("cleanup failed")

    cuda.empty_cache = fail_cleanup  # type: ignore[method-assign]
    monkeypatch.setattr(runtime, "import_torch", lambda: SimpleNamespace(cuda=cuda))

    runtime.cleanup_memory()