voculinux/tests/test_model_ladder.py

"""Ticket #20: memory-aware model ladder, conditional backend order, retry.

All behavior is exercised with mocked probes and loaders: no GPU, no
network, no hardware. The threshold-derivation test pins the rung
thresholds to the provisional transcribe peaks (1132 MiB for q8_0, 846 MiB
for q5_0, ticket #19, not yet recorded in the repo) plus the 512 MiB
protocol headroom. Replacing them with a recorded measurement is a
deliberate change to this test.
"""

import os
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

from vocalinux.speech_recognition import model_ladder
from vocalinux.speech_recognition.model_ladder import (
    DEFAULT_RUNG,
    FLOOR_RUNG_ENGLISH,
    FLOOR_RUNG_MULTILINGUAL,
    HEADROOM_MIB,
    LOW_MEMORY_RUNG,
    PREFER_Q5_0_ABOVE_MIB,
    PREFER_Q8_0_ABOVE_MIB,
    cuda_toolkit_ready_from_probe,
    next_rung_down,
    preferred_backend_order,
    resolve_startup_rung,
    select_ladder_rung,
)
from vocalinux.speech_recognition.recognition_manager import SpeechRecognitionManager
from vocalinux.utils.whispercpp_model_info import WHISPERCPP_MODEL_INFO


def _size_mb(name: str) -> float:
    return float(WHISPERCPP_MODEL_INFO[name]["size_mb"])


class TestThresholdDerivation:
    """Thresholds match the baseline ticket's recorded recommendation."""

    #: Ticket #19 measured transcribe peaks (MiB) per rung.
    Q8_0_MEASURED_MIB: float = 1132.0
    Q5_0_MEASURED_MIB: float = 846.0

    def test_q8_threshold_is_measured_footprint_plus_headroom(self) -> None:
        assert PREFER_Q8_0_ABOVE_MIB == self.Q8_0_MEASURED_MIB + HEADROOM_MIB
        assert PREFER_Q8_0_ABOVE_MIB == 1644.0

    def test_q5_threshold_is_measured_footprint_plus_headroom(self) -> None:
        assert PREFER_Q5_0_ABOVE_MIB == self.Q5_0_MEASURED_MIB + HEADROOM_MIB
        assert PREFER_Q5_0_ABOVE_MIB == 1358.0

    def test_rungs_are_ordered(self) -> None:
        assert PREFER_Q8_0_ABOVE_MIB > PREFER_Q5_0_ABOVE_MIB
        assert _size_mb(DEFAULT_RUNG) > _size_mb(LOW_MEMORY_RUNG)


class TestSelectLadderRung:
    """Free memory picks the rung; the floor pairs with the language."""

    def test_plenty_of_memory_keeps_default(self) -> None:
        assert select_ladder_rung(True, PREFER_Q8_0_ABOVE_MIB) == DEFAULT_RUNG
        assert select_ladder_rung(False, 99999.0) == DEFAULT_RUNG

    def test_between_thresholds_steps_to_q5(self) -> None:
        assert select_ladder_rung(True, PREFER_Q8_0_ABOVE_MIB - 1) == LOW_MEMORY_RUNG
        assert select_ladder_rung(True, PREFER_Q5_0_ABOVE_MIB) == LOW_MEMORY_RUNG

    def test_below_q5_fits_the_floor_by_language(self) -> None:
        assert select_ladder_rung(True, PREFER_Q5_0_ABOVE_MIB - 1) == FLOOR_RUNG_ENGLISH
        assert select_ladder_rung(False, PREFER_Q5_0_ABOVE_MIB - 1) == FLOOR_RUNG_MULTILINGUAL
        assert select_ladder_rung(True, 0.0) == FLOOR_RUNG_ENGLISH
        assert select_ladder_rung(False, 0.0) == FLOOR_RUNG_MULTILINGUAL

    def test_recorded_boundaries_select_recorded_rungs(self) -> None:
        assert select_ladder_rung(True, 1644.0) == DEFAULT_RUNG
        assert select_ladder_rung(True, 1500.0) == LOW_MEMORY_RUNG
        assert select_ladder_rung(True, 1358.0) == LOW_MEMORY_RUNG
        assert select_ladder_rung(True, 1200.0) == FLOOR_RUNG_ENGLISH
        assert select_ladder_rung(False, 1200.0) == FLOOR_RUNG_MULTILINGUAL


class TestResolveStartupRung:
    """A manual pin wins verbatim; only unpinned loads use the ladder."""

    def test_pinned_variant_never_overridden(self) -> None:
        assert resolve_startup_rung("tiny", True, 99999.0) == "tiny"
        assert resolve_startup_rung(LOW_MEMORY_RUNG, False, 0.0) == LOW_MEMORY_RUNG
        assert resolve_startup_rung(FLOOR_RUNG_ENGLISH, True, 99999.0) == FLOOR_RUNG_ENGLISH

    def test_unpinned_uses_ladder(self) -> None:
        assert resolve_startup_rung("", True, 99999.0) == DEFAULT_RUNG
        assert resolve_startup_rung("", False, 0.0) == FLOOR_RUNG_MULTILINGUAL


class TestNextRungDown:
    """One step down per failure; the floor fails loudly."""

    def test_default_steps_to_low_memory(self) -> None:
        assert next_rung_down(DEFAULT_RUNG, True) == LOW_MEMORY_RUNG
        assert next_rung_down(DEFAULT_RUNG, False) == LOW_MEMORY_RUNG

    def test_low_memory_steps_to_floor_by_language(self) -> None:
        assert next_rung_down(LOW_MEMORY_RUNG, True) == FLOOR_RUNG_ENGLISH
        assert next_rung_down(LOW_MEMORY_RUNG, False) == FLOOR_RUNG_MULTILINGUAL

    def test_floor_has_no_step_down(self) -> None:
        assert next_rung_down(FLOOR_RUNG_ENGLISH, True) is None
        assert next_rung_down(FLOOR_RUNG_MULTILINGUAL, False) is None

    def test_non_ladder_rungs_have_no_step_down(self) -> None:
        assert next_rung_down("tiny", True) is None
        assert next_rung_down("medium", False) is None


class TestCudaToolkitReadiness:
    """Driver alone is not a toolkit; arch must match too."""

    def test_complete_matching_toolkit_is_ready(self) -> None:
        assert cuda_toolkit_ready_from_probe(True, True, (8, 6)) is True
        assert cuda_toolkit_ready_from_probe(True, True, (7, 0)) is True

    def test_driver_without_toolkit_is_not_ready(self) -> None:
        assert cuda_toolkit_ready_from_probe(True, False, (8, 6)) is False

    def test_no_gpu_is_not_ready(self) -> None:
        assert cuda_toolkit_ready_from_probe(False, True, (8, 6)) is False
        assert cuda_toolkit_ready_from_probe(False, False, None) is False

    def test_old_architecture_is_not_ready(self) -> None:
        assert cuda_toolkit_ready_from_probe(True, True, (6, 1)) is False
        assert cuda_toolkit_ready_from_probe(True, True, None) is False


NVCC_VERSION_OUTPUT = "Cuda compilation tools, release 13.2, V13.2.51\n"
NVIDIA_SMI_OUTPUT = "3711, 8.6\n"


def _write_toolkit(
    root: Path,
    *,
    nvcc: bool = True,
    header: bool = True,
    libcudart: bool = True,
    target: str = "x86_64-linux",
) -> Path:
    """Lay out a CUDA toolkit tree shaped like the NVIDIA runfile install."""
    if nvcc:
        (root / "bin").mkdir(parents=True, exist_ok=True)
        binary = root / "bin" / "nvcc"
        binary.write_text("#!/bin/sh\n")
        binary.chmod(0o755)
    if header:
        (root / "include").mkdir(parents=True, exist_ok=True)
        (root / "include" / "cuda_runtime.h").write_text("/* header */\n")
    if libcudart:
        lib = root / "targets" / target / "lib"
        lib.mkdir(parents=True, exist_ok=True)
        (lib / "libcudart.so.13").write_text("")
        if target == "x86_64-linux":
            (root / "lib64").symlink_to(lib, target_is_directory=True)
    return root


def _probe_stub(nvcc_output: str | None):
    """Stand in for ``_run_probe``: nvidia-smi sees an RTX-class GPU."""

    def probe(argv, timeout=10):
        if argv[0] == "nvidia-smi":
            return NVIDIA_SMI_OUTPUT
        if argv[0] == "nvcc":
            return nvcc_output
        return None

    return probe


@pytest.fixture
def no_env_toolkit(tmp_path: Path):
    """Isolate detection from the host: no CUDA env vars, empty /usr/local."""
    usr_local = tmp_path / "usr-local"
    usr_local.mkdir()
    with (
        patch.dict(os.environ, {}, clear=True),
        patch.object(model_ladder, "_CUDA_USR_LOCAL_ROOT", usr_local),
        patch.object(model_ladder.platform, "machine", return_value="x86_64"),
    ):
        yield usr_local


class TestCudaToolkitDetection:
    """A toolkit counts from PATH, CUDA_HOME, CUDA_PATH or /usr/local/cuda*/."""

    def test_toolkit_found_via_path(self, no_env_toolkit: Path) -> None:
        with patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(NVCC_VERSION_OUTPUT)):
            assert model_ladder.cuda_toolkit_ready() is True

    @pytest.mark.parametrize("variable", ["CUDA_HOME", "CUDA_PATH"])
    def test_toolkit_found_via_cuda_env_root(
        self, tmp_path: Path, no_env_toolkit: Path, variable: str
    ) -> None:
        root = _write_toolkit(tmp_path / "opt" / "cuda")
        with (
            patch.dict(os.environ, {variable: str(root)}, clear=True),
            patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)),
        ):
            assert model_ladder.cuda_toolkit_ready() is True

    def test_toolkit_found_only_under_usr_local_cuda(self, no_env_toolkit: Path) -> None:
        """The real target: /usr/local/cuda -> cuda-13.2 with nvcc off PATH."""
        real = _write_toolkit(no_env_toolkit / "cuda-13.2")
        (no_env_toolkit / "cuda").symlink_to(real, target_is_directory=True)
        with patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)):
            assert model_ladder.cuda_toolkit_ready() is True

    def test_no_toolkit_anywhere_is_not_ready(self, no_env_toolkit: Path) -> None:
        with patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)):
            assert model_ladder.cuda_toolkit_ready() is False

    @pytest.mark.parametrize(
        "layout",
        [
            {"nvcc": False},
            {"header": False},
            {"libcudart": False},
            {"target": "sbsa-linux"},
        ],
        ids=["no-nvcc", "no-cuda-runtime-header", "no-libcudart", "other-architecture"],
    )
    def test_incomplete_toolkit_is_rejected(
        self, tmp_path: Path, no_env_toolkit: Path, layout: dict
    ) -> None:
        root = _write_toolkit(tmp_path / "partial", **layout)
        with (
            patch.dict(os.environ, {"CUDA_HOME": str(root)}, clear=True),
            patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)),
        ):
            assert model_ladder.cuda_toolkit_ready() is False

    def test_driver_without_toolkit_stays_not_ready_with_partial_tree(
        self, no_env_toolkit: Path
    ) -> None:
        """A /usr/local/cuda with headers only is not a toolkit."""
        _write_toolkit(no_env_toolkit / "cuda-13.2", nvcc=False, libcudart=False)
        with patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)):
            assert model_ladder.cuda_toolkit_ready() is False


class TestBuildAwareBackendOrder:
    """The order names only backends the installed build compiled in."""

    def test_vulkan_only_build_never_picks_cuda_even_with_toolkit(
        self, no_env_toolkit: Path
    ) -> None:
        _write_toolkit(no_env_toolkit / "cuda-13.2")
        (no_env_toolkit / "cuda").symlink_to(no_env_toolkit / "cuda-13.2")
        with patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)):
            assert model_ladder.cuda_toolkit_ready() is True
            assert preferred_backend_order("vulkan") == ("vulkan", "cpu")

    def test_cuda_build_picks_cuda_with_nvcc_off_path(self, no_env_toolkit: Path) -> None:
        _write_toolkit(no_env_toolkit / "cuda-13.2")
        (no_env_toolkit / "cuda").symlink_to(no_env_toolkit / "cuda-13.2")
        with patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)):
            assert preferred_backend_order("cuda") == ("cuda", "cpu")

    def test_cuda_build_without_toolkit_has_no_cuda_rung(self, no_env_toolkit: Path) -> None:
        with patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)):
            assert preferred_backend_order("cuda") == ("cpu",)

    def test_cpu_build_is_cpu_only(self) -> None:
        assert preferred_backend_order("cpu") == ("cpu",)


class TestPreferredBackendOrder:
    """CUDA ranks first only with a complete toolkit."""

    def test_cuda_first_with_toolkit(self) -> None:
        with patch.object(model_ladder, "cuda_toolkit_ready", return_value=True):
            assert preferred_backend_order() == ("cuda", "vulkan", "cpu")

    def test_vulkan_then_cpu_without_toolkit(self) -> None:
        with patch.object(model_ladder, "cuda_toolkit_ready", return_value=False):
            assert preferred_backend_order() == ("vulkan", "cpu")


def _make_manager(**kwargs):
    """Create a whisper.cpp manager without running any initializer."""
    with (
        patch.object(SpeechRecognitionManager, "_init_vosk"),
        patch.object(SpeechRecognitionManager, "_init_whisper"),
        patch.object(SpeechRecognitionManager, "_init_whispercpp"),
    ):
        return SpeechRecognitionManager(
            engine="whisper_cpp",
            model_size="small",
            language="en-us",
            defer_download=True,
            **kwargs,
        )


class TestStepDownRetry:
    """The loader steps one rung down and retries once, then surfaces."""

    def test_first_failure_steps_down_and_retries_once(self) -> None:
        manager = _make_manager()
        manager.model_size = DEFAULT_RUNG
        calls = []

        def fail_once() -> None:
            calls.append(manager.model_size)
            if len(calls) == 1:
                raise RuntimeError("CUDA out of memory")

        with patch.object(SpeechRecognitionManager, "_init_whispercpp_once", side_effect=fail_once):
            manager._init_whispercpp()

        assert calls == [DEFAULT_RUNG, LOW_MEMORY_RUNG]
        assert manager.model_size == LOW_MEMORY_RUNG

    def test_second_failure_surfaces(self) -> None:
        manager = _make_manager()
        manager.model_size = DEFAULT_RUNG
        with patch.object(
            SpeechRecognitionManager, "_init_whispercpp_once", side_effect=RuntimeError("boom")
        ):
            try:
                manager._init_whispercpp()
            except RuntimeError:
                pass
            else:
                raise AssertionError("second failure must surface")
        assert manager.model_size == LOW_MEMORY_RUNG

    def test_manual_pin_never_steps_down(self) -> None:
        manager = _make_manager(whispercpp_variant_pinned=True)
        manager.model_size = DEFAULT_RUNG
        loader = MagicMock(side_effect=RuntimeError("CUDA out of memory"))
        with patch.object(SpeechRecognitionManager, "_init_whispercpp_once", loader):
            try:
                manager._init_whispercpp()
            except RuntimeError:
                pass
            else:
                raise AssertionError("manual pick must fail loudly")
        assert loader.call_count == 1
        assert manager.model_size == DEFAULT_RUNG

    def test_floor_rung_fails_loudly(self) -> None:
        manager = _make_manager()
        manager.model_size = FLOOR_RUNG_ENGLISH
        loader = MagicMock(side_effect=OSError("no space"))
        with patch.object(SpeechRecognitionManager, "_init_whispercpp_once", loader):
            try:
                manager._init_whispercpp()
            except OSError:
                pass
            else:
                raise AssertionError("floor failure must surface")
        assert loader.call_count == 1

    def test_non_ladder_model_does_not_step_down(self) -> None:
        manager = _make_manager()
        manager.model_size = "medium"
        loader = MagicMock(side_effect=RuntimeError("boom"))
        with patch.object(SpeechRecognitionManager, "_init_whispercpp_once", loader):
            try:
                manager._init_whispercpp()
            except RuntimeError:
                pass
            else:
                raise AssertionError("non-ladder failure must surface")
        assert loader.call_count == 1


class TestStartupRungSelection:
    """Unpinned defaults resolve through the ladder; explicit picks do not."""

    def _manager_with(self, saved_size, pinned_variant, language="en-us"):
        from vocalinux.ui.config_manager import ConfigManager

        manager = ConfigManager.__new__(ConfigManager)
        manager.config = {
            "speech_recognition": {
                "whisper_cpp_model_size": saved_size,
                "whisper_cpp_model_variant": pinned_variant,
                "language": language,
            }
        }
        return manager

    def test_unpinned_default_selects_rung_by_memory(self) -> None:
        manager = self._manager_with("tiny", "", "en-us")
        with patch(
            "vocalinux.speech_recognition.model_ladder.free_memory_mib", return_value=99999.0
        ):
            with patch(
                "vocalinux.utils.whispercpp_model_info.is_model_downloaded", return_value=False
            ):
                assert manager.get_model_size_for_engine("whisper_cpp") == DEFAULT_RUNG

    def test_unpinned_default_low_memory_selects_q5(self) -> None:
        manager = self._manager_with("tiny", "", "en-us")
        with patch(
            "vocalinux.speech_recognition.model_ladder.free_memory_mib",
            return_value=PREFER_Q5_0_ABOVE_MIB,
        ):
            with patch(
                "vocalinux.utils.whispercpp_model_info.is_model_downloaded", return_value=False
            ):
                assert manager.get_model_size_for_engine("whisper_cpp") == LOW_MEMORY_RUNG

    def test_unpinned_default_tight_memory_selects_floor(self) -> None:
        manager = self._manager_with("tiny", "", "pl")
        with patch("vocalinux.speech_recognition.model_ladder.free_memory_mib", return_value=100.0):
            with patch(
                "vocalinux.utils.whispercpp_model_info.is_model_downloaded", return_value=False
            ):
                assert manager.get_model_size_for_engine("whisper_cpp") == FLOOR_RUNG_MULTILINGUAL

    def test_pinned_variant_never_uses_ladder(self) -> None:
        manager = self._manager_with("tiny", LOW_MEMORY_RUNG, "en-us")
        with patch(
            "vocalinux.speech_recognition.model_ladder.free_memory_mib", return_value=99999.0
        ):
            assert manager.get_model_size_for_engine("whisper_cpp") == LOW_MEMORY_RUNG

    def test_explicit_size_keeps_legacy_derivation(self) -> None:
        manager = self._manager_with("medium", "", "en-us")
        with patch(
            "vocalinux.speech_recognition.model_ladder.free_memory_mib", return_value=99999.0
        ):
            with patch(
                "vocalinux.utils.whispercpp_model_info.is_model_downloaded", return_value=False
            ):
                assert manager.get_model_size_for_engine("whisper_cpp") == "medium.en"


class TestFreeMemoryForBackend:
    """The memory reading follows the installed build (spec #16, ticket #20)."""

    def test_vulkan_build_reads_discrete_vram_when_present(self) -> None:
        with (
            patch.object(model_ladder, "free_vram_mib", return_value=2600.0),
            patch.object(model_ladder, "free_system_ram_mib", return_value=9000.0),
        ):
            assert model_ladder.free_memory_for_backend("vulkan") == (2600.0, "VRAM")

    def test_cuda_build_reads_discrete_vram_when_present(self) -> None:
        with patch.object(model_ladder, "free_vram_mib", return_value=1500.0):
            assert model_ladder.free_memory_for_backend("cuda") == (1500.0, "VRAM")

    def test_gpu_build_without_an_nvidia_reading_reads_ram(self) -> None:
        """Integrated graphics: no discrete VRAM reading, so system RAM decides."""
        with (
            patch.object(model_ladder, "free_vram_mib", return_value=None),
            patch.object(model_ladder, "free_system_ram_mib", return_value=6000.0),
        ):
            assert model_ladder.free_memory_for_backend("vulkan") == (6000.0, "RAM")

    def test_cpu_build_never_reads_vram(self) -> None:
        vram = MagicMock(return_value=2600.0)
        with (
            patch.object(model_ladder, "free_vram_mib", vram),
            patch.object(model_ladder, "free_system_ram_mib", return_value=1200.0),
        ):
            assert model_ladder.free_memory_for_backend("cpu") == (1200.0, "RAM")
        vram.assert_not_called()