unsloth/tests/studio/install/test_rocm_support.py

"""Tests for AMD ROCm support across install pathways.

Verifies that ROCm detection and installation logic works correctly
WITHOUT breaking existing CUDA, CPU, macOS, and Windows pathways.
All tests use mocks -- no AMD hardware required.
"""

import importlib.util
import json
import os
import subprocess
import sys
from pathlib import Path
from unittest.mock import MagicMock, patch, PropertyMock

import pytest


# ── Load modules under test ──────────────────────────────────────────────────

PACKAGE_ROOT = Path(__file__).resolve().parents[3]

# install_llama_prebuilt.py
_PREBUILT_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
_PREBUILT_SPEC = importlib.util.spec_from_file_location(
    "studio_install_llama_prebuilt", _PREBUILT_PATH
)
assert _PREBUILT_SPEC is not None and _PREBUILT_SPEC.loader is not None
prebuilt_mod = importlib.util.module_from_spec(_PREBUILT_SPEC)
sys.modules[_PREBUILT_SPEC.name] = prebuilt_mod
_PREBUILT_SPEC.loader.exec_module(prebuilt_mod)

HostInfo = prebuilt_mod.HostInfo
AssetChoice = prebuilt_mod.AssetChoice
PrebuiltFallback = prebuilt_mod.PrebuiltFallback
resolve_upstream_asset_choice = prebuilt_mod.resolve_upstream_asset_choice
runtime_patterns_for_choice = prebuilt_mod.runtime_patterns_for_choice
_apply_host_overrides = prebuilt_mod._apply_host_overrides
_normalize_forwarded_gfx = prebuilt_mod._normalize_forwarded_gfx

# install_python_stack.py
_STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py"
_STACK_SPEC = importlib.util.spec_from_file_location("studio_install_python_stack", _STACK_PATH)
assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None
stack_mod = importlib.util.module_from_spec(_STACK_SPEC)
sys.modules[_STACK_SPEC.name] = stack_mod
_STACK_SPEC.loader.exec_module(stack_mod)

_detect_rocm_version = stack_mod._detect_rocm_version
_ensure_rocm_torch = stack_mod._ensure_rocm_torch
_has_rocm_gpu = stack_mod._has_rocm_gpu
_has_usable_nvidia_gpu = stack_mod._has_usable_nvidia_gpu
_ROCM_TORCH_INDEX = stack_mod._ROCM_TORCH_INDEX
_windows_rocm_index_url = stack_mod._windows_rocm_index_url
_detect_windows_gfx_arch = stack_mod._detect_windows_gfx_arch
_install_bnb_windows_rocm = stack_mod._install_bnb_windows_rocm


def _extract_sh_function_body(source: str, name: str) -> str:
    """Return the body of a shell function from `source` by brace matching.

    Used by structural tests that need to assert ordering of helper
    calls inside a specific function rather than across the whole
    install.sh file.
    """
    needle = f"{name}() {{"
    start = source.find(needle)
    if start < 0:
        return ""
    depth = 0
    i = start + len(needle) - 1  # land on the opening brace
    n = len(source)
    while i < n:
        ch = source[i]
        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                return source[start : i + 1]
        i += 1
    return source[start:]


# ── Helper: build HostInfo for different scenarios ──────────────────────────


def nvidia_host(**overrides) -> HostInfo:
    """NVIDIA Linux x86_64 host."""
    defaults = dict(
        system = "Linux",
        machine = "x86_64",
        is_windows = False,
        is_linux = True,
        is_macos = False,
        is_x86_64 = True,
        is_arm64 = False,
        nvidia_smi = "/usr/bin/nvidia-smi",
        driver_cuda_version = (12, 6),
        compute_caps = ["89"],
        visible_cuda_devices = None,
        has_physical_nvidia = True,
        has_usable_nvidia = True,
        has_rocm = False,
    )
    defaults.update(overrides)
    return HostInfo(**defaults)


def rocm_host(**overrides) -> HostInfo:
    """AMD ROCm Linux x86_64 host (no NVIDIA)."""
    defaults = dict(
        system = "Linux",
        machine = "x86_64",
        is_windows = False,
        is_linux = True,
        is_macos = False,
        is_x86_64 = True,
        is_arm64 = False,
        nvidia_smi = None,
        driver_cuda_version = None,
        compute_caps = [],
        visible_cuda_devices = None,
        has_physical_nvidia = False,
        has_usable_nvidia = False,
        has_rocm = True,
    )
    defaults.update(overrides)
    return HostInfo(**defaults)


def cpu_host(**overrides) -> HostInfo:
    """CPU-only Linux x86_64 host."""
    defaults = dict(
        system = "Linux",
        machine = "x86_64",
        is_windows = False,
        is_linux = True,
        is_macos = False,
        is_x86_64 = True,
        is_arm64 = False,
        nvidia_smi = None,
        driver_cuda_version = None,
        compute_caps = [],
        visible_cuda_devices = None,
        has_physical_nvidia = False,
        has_usable_nvidia = False,
        has_rocm = False,
    )
    defaults.update(overrides)
    return HostInfo(**defaults)


def macos_host(**overrides) -> HostInfo:
    """macOS arm64 host."""
    defaults = dict(
        system = "Darwin",
        machine = "arm64",
        is_windows = False,
        is_linux = False,
        is_macos = True,
        is_x86_64 = False,
        is_arm64 = True,
        nvidia_smi = None,
        driver_cuda_version = None,
        compute_caps = [],
        visible_cuda_devices = None,
        has_physical_nvidia = False,
        has_usable_nvidia = False,
        has_rocm = False,
    )
    defaults.update(overrides)
    return HostInfo(**defaults)


def windows_host(**overrides) -> HostInfo:
    """Windows x86_64 host."""
    defaults = dict(
        system = "Windows",
        machine = "amd64",
        is_windows = True,
        is_linux = False,
        is_macos = False,
        is_x86_64 = True,
        is_arm64 = False,
        nvidia_smi = None,
        driver_cuda_version = None,
        compute_caps = [],
        visible_cuda_devices = None,
        has_physical_nvidia = False,
        has_usable_nvidia = False,
        has_rocm = False,
    )
    defaults.update(overrides)
    return HostInfo(**defaults)


def windows_rocm_host(**overrides) -> HostInfo:
    """Windows x86_64 host with ROCm."""
    defaults = dict(
        system = "Windows",
        machine = "amd64",
        is_windows = True,
        is_linux = False,
        is_macos = False,
        is_x86_64 = True,
        is_arm64 = False,
        nvidia_smi = None,
        driver_cuda_version = None,
        compute_caps = [],
        visible_cuda_devices = None,
        has_physical_nvidia = False,
        has_usable_nvidia = False,
        has_rocm = True,
    )
    defaults.update(overrides)
    return HostInfo(**defaults)


# ── Upstream asset fixture ───────────────────────────────────────────────────

LLAMA_TAG = "b8508"

UPSTREAM_ASSETS = {
    f"llama-{LLAMA_TAG}-bin-ubuntu-x64.tar.gz": f"https://example.com/{LLAMA_TAG}-linux-cpu.tar.gz",
    f"llama-{LLAMA_TAG}-bin-ubuntu-rocm-7.2-x64.tar.gz": f"https://example.com/{LLAMA_TAG}-linux-rocm.tar.gz",
    f"llama-{LLAMA_TAG}-bin-win-cpu-x64.zip": f"https://example.com/{LLAMA_TAG}-win-cpu.zip",
    f"llama-{LLAMA_TAG}-bin-win-cuda-12.4-x64.zip": f"https://example.com/{LLAMA_TAG}-win-cuda.zip",
    f"llama-{LLAMA_TAG}-bin-win-hip-radeon-x64.zip": f"https://example.com/{LLAMA_TAG}-win-hip.zip",
    f"llama-{LLAMA_TAG}-bin-macos-arm64.tar.gz": f"https://example.com/{LLAMA_TAG}-macos-arm64.tar.gz",
    f"llama-{LLAMA_TAG}-bin-macos-x64.tar.gz": f"https://example.com/{LLAMA_TAG}-macos-x64.tar.gz",
}


# TEST: install_llama_prebuilt.py -- resolve_upstream_asset_choice


class TestResolveUpstreamAssetChoice:
    """Verify that the asset selection logic picks the right binary for each platform."""

    @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS)
    def test_nvidia_linux_gets_cpu_asset(self, mock_assets):
        """NVIDIA host should NOT hit the ROCm path -- gets CPU asset (CUDA handled elsewhere)."""
        host = nvidia_host()
        choice = resolve_upstream_asset_choice(host, LLAMA_TAG)
        assert choice.install_kind == "linux-cpu"
        assert "ubuntu-x64" in choice.name
        assert "rocm" not in choice.name

    @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS)
    def test_rocm_linux_gets_rocm_prebuilt(self, mock_assets):
        """AMD ROCm Linux host should get the ROCm prebuilt."""
        host = rocm_host()
        choice = resolve_upstream_asset_choice(host, LLAMA_TAG)
        assert choice.install_kind == "linux-rocm"
        assert "rocm" in choice.name

    @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS)
    def test_cpu_linux_gets_cpu_asset(self, mock_assets):
        """CPU-only Linux host should get CPU asset."""
        host = cpu_host()
        choice = resolve_upstream_asset_choice(host, LLAMA_TAG)
        assert choice.install_kind == "linux-cpu"
        assert "ubuntu-x64" in choice.name

    @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS)
    def test_macos_arm64_gets_macos_asset(self, mock_assets):
        """macOS arm64 host should get macOS asset."""
        host = macos_host()
        choice = resolve_upstream_asset_choice(host, LLAMA_TAG)
        assert choice.install_kind == "macos-arm64"
        assert "macos-arm64" in choice.name

    @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS)
    def test_windows_cpu_gets_cpu_asset(self, mock_assets):
        """Windows CPU-only host should get Windows CPU asset."""
        host = windows_host()
        choice = resolve_upstream_asset_choice(host, LLAMA_TAG)
        assert choice.install_kind == "windows-cpu"
        assert "win-cpu" in choice.name

    @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS)
    def test_windows_rocm_gets_hip_asset(self, mock_assets):
        """Windows ROCm host should get Windows HIP asset."""
        host = windows_rocm_host()
        choice = resolve_upstream_asset_choice(host, LLAMA_TAG)
        assert choice.install_kind == "windows-hip"
        assert "hip" in choice.name

    @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS)
    def test_mixed_nvidia_rocm_prefers_nvidia(self, mock_assets):
        """Host with both NVIDIA and ROCm should use NVIDIA (CPU path here, CUDA elsewhere)."""
        host = nvidia_host(has_rocm = True)
        choice = resolve_upstream_asset_choice(host, LLAMA_TAG)
        # NVIDIA hosts go through the normal path (CUDA handled by resolve_linux_cuda_choice)
        assert choice.install_kind == "linux-cpu"
        assert "rocm" not in choice.name

    @patch.object(prebuilt_mod, "github_release_assets")
    def test_rocm_linux_no_prebuilt_falls_back(self, mock_assets):
        """AMD ROCm host should fall back to source build when no ROCm prebuilt exists."""
        # Remove the ROCm asset from available assets
        assets_without_rocm = {k: v for k, v in UPSTREAM_ASSETS.items() if "rocm" not in k}
        mock_assets.return_value = assets_without_rocm
        host = rocm_host()
        with pytest.raises(PrebuiltFallback, match = "ROCm detected"):
            resolve_upstream_asset_choice(host, LLAMA_TAG)

    @patch.object(prebuilt_mod, "github_release_assets")
    def test_windows_rocm_no_hip_falls_to_cpu(self, mock_assets):
        """Windows+ROCm with HIP prebuilt missing should fall through to CPU."""
        assets_no_hip = {k: v for k, v in UPSTREAM_ASSETS.items() if "hip" not in k}
        mock_assets.return_value = assets_no_hip
        host = windows_rocm_host()
        choice = resolve_upstream_asset_choice(host, LLAMA_TAG)
        assert choice.install_kind == "windows-cpu"

    @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS)
    def test_macos_rocm_impossible_has_rocm_false(self, mock_assets):
        """macOS host should never have has_rocm=True in practice; verify it gets macOS asset."""
        host = macos_host(has_rocm = True)
        choice = resolve_upstream_asset_choice(host, LLAMA_TAG)
        assert choice.install_kind == "macos-arm64"

    @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS)
    def test_linux_aarch64_rocm_gets_prebuilt_fallback(self, mock_assets):
        """Linux aarch64 with ROCm -- no x86_64 match, should raise PrebuiltFallback."""
        host = rocm_host(machine = "aarch64", is_x86_64 = False, is_arm64 = True)
        with pytest.raises(PrebuiltFallback):
            resolve_upstream_asset_choice(host, LLAMA_TAG)


# TEST: install_llama_prebuilt.py -- runtime_patterns_for_choice


class TestRuntimePatterns:
    """Verify runtime file patterns for all install kinds."""

    def test_linux_cpu_patterns(self):
        choice = AssetChoice(
            repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-cpu"
        )
        patterns = runtime_patterns_for_choice(choice)
        assert "llama-server" in patterns
        assert "llama-quantize" in patterns
        # Broad lib*.so* covers libllama, libggml, libmtmd, libggml-cpu-*,
        # plus the libllama-<binary>-impl.so split that ggml-org/llama.cpp
        # #23462 introduced between b9279 and b9283.
        assert "lib*.so*" in patterns

    def test_linux_cuda_patterns(self):
        choice = AssetChoice(
            repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-cuda"
        )
        patterns = runtime_patterns_for_choice(choice)
        # libggml-cuda.so is matched by lib*.so* now.
        assert "lib*.so*" in patterns

    def test_linux_rocm_patterns(self):
        choice = AssetChoice(
            repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-rocm"
        )
        patterns = runtime_patterns_for_choice(choice)
        # libggml-hip.so is matched by lib*.so* now.
        assert "lib*.so*" in patterns
        assert "llama-server" in patterns

    def test_windows_hip_patterns(self):
        choice = AssetChoice(
            repo = "",
            tag = "",
            name = "",
            url = "",
            source_label = "",
            install_kind = "windows-hip",
        )
        patterns = runtime_patterns_for_choice(choice)
        # Narrowed from "*.exe" to the two binaries Studio actually
        # invokes, mirroring the Linux/macOS pattern style.
        assert "llama-server.exe" in patterns
        assert "llama-quantize.exe" in patterns
        assert "*.dll" in patterns

    def test_macos_patterns(self):
        choice = AssetChoice(
            repo = "",
            tag = "",
            name = "",
            url = "",
            source_label = "",
            install_kind = "macos-arm64",
        )
        patterns = runtime_patterns_for_choice(choice)
        assert "lib*.dylib" in patterns


# TEST: install_llama_prebuilt.py -- HostInfo.has_rocm field


class TestHostInfoRocm:
    """Verify has_rocm field does not affect other HostInfo behavior."""

    def test_has_rocm_default_false(self):
        host = HostInfo(
            system = "Linux",
            machine = "x86_64",
            is_windows = False,
            is_linux = True,
            is_macos = False,
            is_x86_64 = True,
            is_arm64 = False,
            nvidia_smi = None,
            driver_cuda_version = None,
            compute_caps = [],
            visible_cuda_devices = None,
            has_physical_nvidia = False,
            has_usable_nvidia = False,
        )
        assert host.has_rocm is False

    def test_has_rocm_explicit_true(self):
        host = rocm_host()
        assert host.has_rocm is True

    def test_nvidia_host_no_rocm(self):
        host = nvidia_host()
        assert host.has_rocm is False
        assert host.has_usable_nvidia is True

    def test_detect_host_has_rocm_detection_logic(self):
        """detect_host() should have ROCm GPU detection logic."""
        import inspect

        source = inspect.getsource(prebuilt_mod.detect_host)
        # Must probe for actual GPU, not just tool presence
        assert "rocminfo" in source or "amd-smi" in source

    def test_detect_host_windows_rocm_detection(self):
        """detect_host() source should have Windows-specific ROCm GPU detection."""
        import inspect

        source = inspect.getsource(prebuilt_mod.detect_host)
        assert "hipinfo" in source or "amd-smi" in source


# TEST: install_python_stack.py -- _detect_rocm_version


class TestDetectRocmVersion:
    """Verify ROCm version detection from various sources."""

    def test_no_rocm_returns_none(self, tmp_path):
        """No ROCm installed should return None."""
        with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}):
            with patch("shutil.which", return_value = None):
                result = _detect_rocm_version()
                assert result is None

    def test_version_from_file(self, tmp_path):
        """Reads version from /opt/rocm/.info/version."""
        info_dir = tmp_path / ".info"
        info_dir.mkdir()
        (info_dir / "version").write_text("7.1.0-12345\n")
        with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}):
            result = _detect_rocm_version()
            assert result == (7, 1)

    def test_version_62(self, tmp_path):
        """Reads ROCm 6.2 version."""
        info_dir = tmp_path / ".info"
        info_dir.mkdir()
        (info_dir / "version").write_text("6.2.0\n")
        with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}):
            result = _detect_rocm_version()
            assert result == (6, 2)

    def test_hipconfig_fallback(self, tmp_path):
        """Falls back to hipconfig --version when file not found."""
        with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}):
            mock_result = MagicMock()
            mock_result.returncode = 0
            mock_result.stdout = b"6.3.21234.2\n"
            with patch("shutil.which", return_value = "/usr/bin/hipconfig"):
                with patch("subprocess.run", return_value = mock_result):
                    result = _detect_rocm_version()
                    assert result == (6, 3)

    def test_dpkg_fallback_without_hipconfig(self, tmp_path):
        """dpkg rocm-core fallback works when amd-smi and hipconfig are absent
        (regression: a shadowing local re import raised UnboundLocalError)."""

        def which(cmd):
            return "/usr/bin/dpkg-query" if cmd == "dpkg-query" else None

        mock_result = MagicMock()
        mock_result.returncode = 0
        mock_result.stdout = "1:6.3.0-1\n"
        with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}):
            with patch("shutil.which", side_effect = which):
                with patch("subprocess.run", return_value = mock_result):
                    assert _detect_rocm_version() == (6, 3)

    def test_empty_version_file(self, tmp_path):
        """Empty version file should return None."""
        info_dir = tmp_path / ".info"
        info_dir.mkdir()
        (info_dir / "version").write_text("")
        with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}):
            with patch("shutil.which", return_value = None):
                result = _detect_rocm_version()
                assert result is None

    def test_version_with_epoch_prefix(self, tmp_path):
        """Debian epoch prefix (2:6.2.0) -- version file has no epoch, so should parse."""
        info_dir = tmp_path / ".info"
        info_dir.mkdir()
        # Version files don't typically have epoch prefix, but lib/rocm_version might
        (info_dir / "version").write_text("6.2.0\n")
        with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}):
            result = _detect_rocm_version()
            assert result == (6, 2)

    def test_multiple_version_sources_first_wins(self, tmp_path):
        """When both .info/version and lib/rocm_version exist, first found wins."""
        info_dir = tmp_path / ".info"
        info_dir.mkdir()
        (info_dir / "version").write_text("7.1.0\n")
        lib_dir = tmp_path / "lib"
        lib_dir.mkdir()
        (lib_dir / "rocm_version").write_text("6.3.0\n")
        with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}):
            result = _detect_rocm_version()
            assert result == (7, 1)  # .info/version checked first

    def test_hipconfig_multiline_output(self, tmp_path):
        """hipconfig with multi-line output -- should use first line."""
        with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}):
            mock_result = MagicMock()
            mock_result.returncode = 0
            mock_result.stdout = b"6.3.21234.2\nSome extra info\n"
            with patch("shutil.which", return_value = "/usr/bin/hipconfig"):
                with patch("subprocess.run", return_value = mock_result):
                    result = _detect_rocm_version()
                    assert result == (6, 3)

    def test_hipconfig_timeout(self, tmp_path):
        """hipconfig that times out should return None."""
        with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}):
            with patch("shutil.which", return_value = "/usr/bin/hipconfig"):
                with patch(
                    "subprocess.run",
                    side_effect = subprocess.TimeoutExpired("hipconfig", 5),
                ):
                    result = _detect_rocm_version()
                    assert result is None


# TEST: install_python_stack.py -- _ensure_rocm_torch


class TestEnsureRocmTorch:
    """Verify ROCm torch reinstall logic."""

    @patch.object(stack_mod, "pip_install")
    @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
    def test_no_rocm_skips(self, mock_nvidia, mock_pip):
        """No ROCm toolchain should skip entirely."""
        # _detect_windows_gfx_arch pinned to None: on a real AMD test host its
        # WMI name fallback would otherwise answer and defeat the "no ROCm
        # anywhere" premise of this test.
        with patch.object(stack_mod, "_detect_windows_gfx_arch", return_value = None):
            with patch("os.path.isdir", return_value = False):
                with patch("shutil.which", return_value = None):
                    _ensure_rocm_torch()
        mock_pip.assert_not_called()

    @patch.object(stack_mod, "pip_install")
    @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
    @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
    @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
    def test_torch_already_has_cuda_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
        """If torch already has CUDA, should skip ROCm reinstall."""
        mock_probe = MagicMock()
        mock_probe.returncode = 0
        mock_probe.stdout = b"12.6\n"  # CUDA version string
        with patch("os.path.isdir", return_value = True):
            with patch("subprocess.run", return_value = mock_probe):
                _ensure_rocm_torch()
        mock_pip.assert_not_called()

    @patch.object(stack_mod, "pip_install")
    @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
    @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
    @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
    def test_torch_already_has_hip_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
        """If torch already has HIP, should skip ROCm reinstall."""
        mock_probe = MagicMock()
        mock_probe.returncode = 0
        mock_probe.stdout = b"7.1.12345\n"  # HIP version string
        with patch("os.path.isdir", return_value = True):
            with patch("subprocess.run", return_value = mock_probe):
                _ensure_rocm_torch()
        mock_pip.assert_not_called()

    @patch.object(stack_mod, "IS_WINDOWS", False)
    @patch.object(stack_mod, "pip_install_try", return_value = True)
    @patch.object(stack_mod, "pip_install")
    @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
    @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
    @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
    def test_cpu_torch_gets_rocm_reinstall(
        self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
    ):
        """CPU-only torch on ROCm host should trigger reinstall."""
        mock_probe = MagicMock()
        mock_probe.returncode = 0
        mock_probe.stdout = b"\n"  # empty = no GPU backend
        with patch("os.path.isdir", return_value = True):
            with patch("subprocess.run", return_value = mock_probe):
                _ensure_rocm_torch()
        # Should install torch via pip_install and bitsandbytes via pip_install_try.
        assert mock_pip.call_count == 1
        assert "rocm7.1" in str(mock_pip.call_args_list[0])
        assert mock_pip_try.call_count >= 1
        assert "bitsandbytes" in str(mock_pip_try.call_args_list[0])
        assert mock_pip_try.call_args.kwargs["force_pip"] is True

    @patch.object(stack_mod, "IS_WINDOWS", False)
    @patch.object(stack_mod, "pip_install")
    @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
    @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
    @patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 3))
    def test_rocm_63_selects_correct_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
        """ROCm 6.3 should select rocm6.3 tag."""
        mock_probe = MagicMock()
        mock_probe.returncode = 0
        mock_probe.stdout = b"\n"
        with patch("os.path.isdir", return_value = True):
            with patch("subprocess.run", return_value = mock_probe):
                _ensure_rocm_torch()
        torch_call = mock_pip.call_args_list[0]
        assert "rocm6.3" in str(torch_call)

    @patch.object(stack_mod, "pip_install")
    @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
    @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
    @patch.object(stack_mod, "_detect_rocm_version", return_value = (5, 0))
    def test_old_rocm_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
        """ROCm version too old (below 6.0) should skip."""
        mock_probe = MagicMock()
        mock_probe.returncode = 0
        mock_probe.stdout = b"\n"
        with patch("os.path.isdir", return_value = True):
            with patch("subprocess.run", return_value = mock_probe):
                _ensure_rocm_torch()
        mock_pip.assert_not_called()

    @patch.object(stack_mod, "IS_WINDOWS", False)
    @patch.object(stack_mod, "pip_install")
    @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
    @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
    @patch.object(stack_mod, "_detect_rocm_version", return_value = None)
    def test_version_unreadable_prints_warning(
        self, mock_ver, mock_gpu, mock_nvidia, mock_pip, capsys
    ):
        """ROCm detected but version unreadable should print warning and skip."""
        with patch("os.path.isdir", return_value = True):
            _ensure_rocm_torch()
        mock_pip.assert_not_called()
        captured = capsys.readouterr()
        assert "unreadable" in captured.out

    @patch.object(stack_mod, "IS_WINDOWS", False)
    @patch.object(stack_mod, "pip_install")
    @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
    @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
    @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
    def test_rocm_72_selects_72_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
        """ROCm 7.2 should select rocm7.2 tag (now in mapping with torch 2.11.0)."""
        mock_probe = MagicMock()
        mock_probe.returncode = 0
        mock_probe.stdout = b"\n"
        with patch("os.path.isdir", return_value = True):
            with patch("subprocess.run", return_value = mock_probe):
                _ensure_rocm_torch()
        torch_call = mock_pip.call_args_list[0]
        assert "rocm7.2" in str(torch_call)

    @patch.object(stack_mod, "IS_WINDOWS", False)
    @patch.object(stack_mod, "pip_install_try", return_value = True)
    @patch.object(stack_mod, "pip_install")
    @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
    @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
    @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
    def test_probe_timeout_triggers_reinstall(
        self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
    ):
        """Probe subprocess timeout should not crash; should proceed to reinstall."""
        with patch("os.path.isdir", return_value = True):
            with patch("subprocess.run", side_effect = subprocess.TimeoutExpired("python", 30)):
                _ensure_rocm_torch()
        # If probe times out, the function should treat torch as unusable and reinstall
        # both torch (via pip_install) and bitsandbytes (via pip_install_try).
        assert mock_pip.call_count == 1
        assert "rocm7.1" in str(mock_pip.call_args_list[0])
        assert mock_pip_try.call_count >= 1
        assert mock_pip_try.call_args.kwargs["force_pip"] is True

    @patch.object(stack_mod, "pip_install")
    @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
    @patch.object(stack_mod, "_has_rocm_gpu", return_value = False)
    def test_no_gpu_with_rocm_tools_skips(self, mock_gpu, mock_nvidia, mock_pip):
        """ROCm tools present but no actual AMD GPU should skip entirely."""
        # Pin the Windows arch probe to None: on a real AMD test host the WMI
        # name fallback would otherwise answer and defeat the "no actual GPU"
        # premise (the Linux path under test uses _has_rocm_gpu, mocked False).
        with patch.object(stack_mod, "_detect_windows_gfx_arch", return_value = None):
            with patch("os.path.isdir", return_value = True):
                _ensure_rocm_torch()
        mock_pip.assert_not_called()

    @patch.object(stack_mod, "pip_install")
    @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
    @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True)
    def test_torch_backend_cuda_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip):
        """UNSLOTH_TORCH_BACKEND=cuda must short-circuit before any GPU probe."""
        with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cuda"}):
            # Reload _TORCH_BACKEND from the patched environment.
            with patch.object(stack_mod, "_TORCH_BACKEND", "cuda"):
                _ensure_rocm_torch()
        mock_pip.assert_not_called()

    @patch.object(stack_mod, "pip_install")
    @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
    @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True)
    def test_torch_backend_cpu_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip):
        """UNSLOTH_TORCH_BACKEND=cpu must short-circuit before any GPU probe."""
        with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cpu"}):
            with patch.object(stack_mod, "_TORCH_BACKEND", "cpu"):
                _ensure_rocm_torch()
        mock_pip.assert_not_called()


# TEST: install_python_stack.py -- _has_rocm_gpu KFD sysfs vendor_id guard


class TestHasRocmGpuKfdVendorGuard:
    """Verify that the KFD sysfs fallback rejects non-AMD (NVIDIA) KFD nodes.

    These tests are source-level: they verify the regex and logic present in
    the _has_rocm_gpu implementation rather than running the sysfs traversal
    (which requires Linux path conventions).
    """

    def _src(self) -> str:
        """Return the source of _has_rocm_gpu from install_python_stack.py."""
        import inspect
        return inspect.getsource(stack_mod._has_rocm_gpu)

    def test_vendor_id_check_present(self):
        """_has_rocm_gpu sysfs fallback must check vendor_id 4098 (AMD 0x1002)."""
        src = self._src()
        assert "vendor_id" in src, (
            "_has_rocm_gpu KFD sysfs fallback must read the properties file "
            "to check vendor_id and exclude NVIDIA KFD nodes"
        )
        assert "4098" in src, (
            "_has_rocm_gpu must require AMD vendor_id 4098 (0x1002) in the "
            "KFD node properties to avoid false positives on NVIDIA systems"
        )

    def test_vendor_regex_pattern_anchored(self):
        """The vendor_id regex must use a word boundary to avoid partial matches."""
        import re as _re

        src = self._src()
        # The pattern should have a word boundary before and after the number
        # so "vendor_id 41098" doesn't match "vendor_id 4098".
        assert (
            _re.search(r"\\b.*vendor_id.*\\b", src) or "\\bvendor_id" in src
        ), "_has_rocm_gpu vendor_id check should use word boundary anchors"

    def test_sysfs_fallback_guarded_by_non_win32(self):
        """KFD sysfs fallback must be Linux-only (guarded by sys.platform != 'win32')."""
        src = self._src()
        assert "win32" in src, "_has_rocm_gpu sysfs fallback must be guarded by sys.platform check"

    def test_cpu_node_excluded(self):
        """gpu_id == '0' must be excluded (CPU topology nodes)."""
        src = self._src()
        assert (
            '!= "0"' in src or "== '0'" in src or "!= '0'" in src or '"0"' in src
        ), "_has_rocm_gpu must skip gpu_id 0 nodes (CPU nodes)"

    def test_install_sh_has_vendor_check(self):
        """_has_amd_rocm_gpu in install.sh sysfs fallback must also check vendor_id 4098."""
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        func_start = source.find("_has_amd_rocm_gpu()")
        func_end = source.find("\n}", func_start)
        func_body = source[func_start:func_end]
        assert "vendor_id" in func_body, "_has_amd_rocm_gpu sysfs fallback must check vendor_id"
        assert "4098" in func_body, "_has_amd_rocm_gpu must require AMD vendor_id 4098 (0x1002)"

    def test_has_rocm_gpu_returns_false_when_nvidia_present(self):
        """_has_rocm_gpu must return False immediately when _has_usable_nvidia_gpu is True.

        This is the primary guard: even if rocminfo, amd-smi, or KFD sysfs
        produce a false positive, an NVIDIA GPU always wins.
        """
        with patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True):
            with patch("shutil.which", return_value = "/usr/bin/rocminfo"):
                # Simulate rocminfo claiming an AMD GPU is present
                mock_result = MagicMock()
                mock_result.returncode = 0
                mock_result.stdout = "Name: gfx1100\n"
                with patch("subprocess.run", return_value = mock_result):
                    assert not stack_mod._has_rocm_gpu(), (
                        "_has_rocm_gpu must return False when NVIDIA GPU is detected, "
                        "regardless of what rocminfo reports"
                    )

    def test_install_sh_has_rocm_gpu_nvidia_guard(self):
        """_has_amd_rocm_gpu in install.sh must call _has_usable_nvidia_gpu and return 1 if true."""
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        func_start = source.find("_has_amd_rocm_gpu()")
        func_end = source.find("\n}", func_start)
        func_body = source[func_start:func_end]
        assert (
            "_has_usable_nvidia_gpu" in func_body
        ), "_has_amd_rocm_gpu must call _has_usable_nvidia_gpu to block NVIDIA hosts"
        assert (
            "return 1" in func_body
        ), "_has_amd_rocm_gpu must return 1 (false) when NVIDIA GPU is detected"

    def test_has_usable_nvidia_gpu_proc_fallback_present(self):
        """`_has_usable_nvidia_gpu` must have a /proc/driver/nvidia fallback."""
        import inspect

        src = inspect.getsource(stack_mod._has_usable_nvidia_gpu)
        assert "/proc/driver/nvidia" in src, (
            "_has_usable_nvidia_gpu must fall back to /proc/driver/nvidia/gpus when "
            "nvidia-smi subprocess fails, to handle PATH gaps and driver init races"
        )

    def test_install_sh_has_usable_nvidia_gpu_proc_fallback(self):
        """_has_usable_nvidia_gpu in install.sh must also have a /proc/driver/nvidia fallback."""
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        func_start = source.find("_has_usable_nvidia_gpu()")
        func_end = source.find("\n}", func_start)
        func_body = source[func_start:func_end]
        assert "/proc/driver/nvidia" in func_body, (
            "_has_usable_nvidia_gpu in install.sh must fall back to "
            "/proc/driver/nvidia/gpus when nvidia-smi fails"
        )


# TEST: install_python_stack.py -- _ROCM_TORCH_INDEX mapping


class TestRocmTorchIndex:
    """Verify the ROCm version -> torch index tag mapping."""

    def test_mapping_is_sorted_descending(self):
        """Keys should be in descending order for the next() iteration to work."""
        keys = list(_ROCM_TORCH_INDEX.keys())
        assert keys == sorted(keys, reverse = True)

    def test_rocm_72_in_mapping(self):
        """ROCm 7.2 should be in the active mapping (torch 2.11.0 now supported)."""
        assert (7, 2) in _ROCM_TORCH_INDEX
        assert _ROCM_TORCH_INDEX[(7, 2)] == "rocm7.2"

    def test_rocm_71_maps_correctly(self):
        assert _ROCM_TORCH_INDEX[(7, 1)] == "rocm7.1"

    def test_rocm_63_maps_correctly(self):
        assert _ROCM_TORCH_INDEX[(6, 3)] == "rocm6.3"

    def test_rocm_60_maps_correctly(self):
        assert _ROCM_TORCH_INDEX[(6, 0)] == "rocm6.0"

    def test_all_tags_use_download_pytorch(self):
        """All tags should be for download.pytorch.org, not repo.radeon.com."""
        for tag in _ROCM_TORCH_INDEX.values():
            assert tag.startswith("rocm")
            assert "radeon" not in tag

    def test_newer_rocm_selects_best_match(self):
        """ROCm 7.2 (now in map) should select rocm7.2 directly."""
        ver = (7, 2)
        tag = next(
            (
                t
                for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True)
                if ver >= (maj, mn)
            ),
            None,
        )
        assert tag == "rocm7.2"

    def test_rocm_64_selects_64(self):
        ver = (6, 4)
        tag = next(
            (
                t
                for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True)
                if ver >= (maj, mn)
            ),
            None,
        )
        assert tag == "rocm6.4"


# TEST: hardware.py -- IS_ROCM flag and detect_hardware


class TestHardwareRocmFlag:
    """Verify IS_ROCM flag behavior without importing the full hardware module."""

    def test_hardware_py_has_is_rocm(self):
        """hardware.py should define IS_ROCM."""
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        assert "IS_ROCM: bool" in source and "False" in source

    def test_hardware_py_sets_is_rocm_on_hip(self):
        """detect_hardware() should set IS_ROCM when torch.version.hip is set."""
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        assert 'torch.version, "hip"' in source or "torch.version.hip" in source

    def test_hardware_py_still_returns_cuda_for_rocm(self):
        """DeviceType should remain CUDA even on ROCm -- no DeviceType.ROCM."""
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        # Ensure ROCM is NOT a DeviceType member
        enum_section = source.split("class DeviceType")[1].split("\n\n")[0]
        assert "ROCM" not in enum_section

    def test_hardware_py_has_rocm_in_package_versions(self):
        """get_package_versions() should include 'rocm' key."""
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        assert '"rocm"' in source

    def test_hardware_py_device_type_cuda_references_intact(self):
        """All existing DeviceType.CUDA references should still be present."""
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        # Key functions that must still reference DeviceType.CUDA
        assert "DeviceType.CUDA" in source
        assert "DEVICE = DeviceType.CUDA" in source

    def test_is_rocm_exported_from_init(self):
        """IS_ROCM should be exported from hardware __init__.py."""
        init_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py"
        source = init_path.read_text(encoding = "utf-8")
        assert "IS_ROCM" in source

    def test_is_rocm_in_all_list(self):
        """IS_ROCM should be in __all__ list in __init__.py."""
        init_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py"
        source = init_path.read_text(encoding = "utf-8")
        # Extract __all__ section
        assert '"IS_ROCM"' in source

    def test_get_package_versions_returns_rocm_key(self):
        """get_package_versions() source should return both 'cuda' and 'rocm' keys."""
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        # Find the get_package_versions function body
        func_start = source.find("def get_package_versions")
        func_body = source[func_start : source.find("\ndef ", func_start + 1)]
        assert '"cuda"' in func_body
        assert '"rocm"' in func_body

    def test_distributed_stubs_cover_is_torchelastic_launched(self):
        """_determine_attention_impl_for_gpu_estimate must stub is_torchelastic_launched.

        resolve_attention_implementation calls is_torchelastic_launched() on
        Windows ROCm where torch.distributed ships without that helper, causing
        a warning: 'module torch.distributed has no attribute is_torchelastic_launched'.
        """
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        assert "is_torchelastic_launched" in source

    def test_distributed_stubs_cover_core_helpers(self):
        """_determine_attention_impl_for_gpu_estimate must stub the four core distributed helpers."""
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        for attr in ("is_initialized", "is_available", "get_rank", "get_world_size"):
            assert attr in source, f"distributed stub for '{attr}' missing from hardware.py"


# TEST: tokenizer_utils.py -- error message


class TestTokenizerErrorMessage:
    """Verify the AMD error message is updated."""

    def test_no_old_amd_message(self):
        """Old 'We do not support AMD' message should be gone."""
        tu_path = PACKAGE_ROOT / "unsloth" / "tokenizer_utils.py"
        source = tu_path.read_text(encoding = "utf-8")
        assert "We do not support AMD" not in source

    def test_new_message_has_docs_link(self):
        """New message should point to Unsloth AMD docs."""
        tu_path = PACKAGE_ROOT / "unsloth" / "tokenizer_utils.py"
        source = tu_path.read_text(encoding = "utf-8")
        assert "docs.unsloth.ai" in source or "No GPU detected" in source


# TEST: install.sh -- structural checks


class TestInstallShStructure:
    """Verify install.sh structural properties without running it."""

    def test_no_here_strings(self):
        """install.sh must not use the bash-only `<<<` here-string operator.

        `<<<` inside a quoted literal (e.g. a marker in a printf) is just data,
        not a here-string, so strip quoted spans first: this still catches a
        real `cmd <<< word` (outside quotes) without false-positiving on data.
        """
        import re

        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        # <<< is bash-only; breaks dash
        for i, line in enumerate(source.splitlines(), 1):
            stripped = line.lstrip()
            if stripped.startswith("#"):
                continue
            # Remove quoted string literals so `<<<` inside them is ignored;
            # a genuine here-string operator lives outside any quotes.
            unquoted = re.sub(r"'[^']*'", "", line)
            unquoted = re.sub(r'"[^"]*"', "", unquoted)
            assert "<<<" not in unquoted, f"install.sh:{i} uses non-POSIX <<< here-string"

    def test_rocm_detection_present(self):
        """install.sh should have ROCm detection in get_torch_index_url."""
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        assert "amd-smi" in source
        assert "rocm" in source.lower()

    def test_cuda_precedence(self):
        """ROCm detection should only run when nvidia-smi is absent.

        install.sh defines _has_amd_rocm_gpu and _has_usable_nvidia_gpu
        helpers near each other (file-position order has no semantic
        meaning), so check the runtime ordering inside
        get_torch_index_url instead: NVIDIA branch runs first and the
        AMD/ROCm branch only fires inside the `if [ -z "$_smi" ]`
        block.
        """
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        body = _extract_sh_function_body(source, "get_torch_index_url")
        nvidia_call = body.find("_has_usable_nvidia_gpu")
        # Gate changed from [ -z "$_smi" ] to [ "$_nvidia_detected" -eq 0 ] to
        # handle proc-only NVIDIA hosts where nvidia-smi is absent but _has_usable_nvidia_gpu
        # returns true via /proc/driver/nvidia/gpus.
        no_nvidia_branch = body.find('if [ "$_nvidia_detected" -eq 0 ]')
        if no_nvidia_branch < 0:
            no_nvidia_branch = body.find('if [ -z "$_smi" ]')
        rocm_call = body.find("_has_amd_rocm_gpu")
        assert nvidia_call >= 0, "get_torch_index_url should call _has_usable_nvidia_gpu"
        assert no_nvidia_branch >= 0, "get_torch_index_url should gate ROCm on no-nvidia branch"
        assert (
            rocm_call > no_nvidia_branch
        ), "ROCm detection should sit inside the 'no NVIDIA' branch"
        assert (
            nvidia_call < no_nvidia_branch
        ), "NVIDIA detection should run before the no-NVIDIA branch"

    def test_bitsandbytes_amd_install(self):
        """install.sh should install bitsandbytes for AMD when ROCm detected."""
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        assert "bitsandbytes" in source
        assert "rocm*)" in source  # case pattern for ROCm URLs

    def test_cpu_hint_mentions_amd(self):
        """CPU-only hint should mention AMD ROCm."""
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        assert "ROCm" in source

    def test_rocm72_supported_future_capped(self):
        """ROCm 7.2 should pass through directly; 7.3+ falls back to rocm7.2."""
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        assert 'echo "$_base/rocm7.2"' in source  # fallback for unknown future versions
        # Allowlisted versions should pass through directly
        assert "rocm6.*" in source
        assert "rocm7.0" in source
        assert "rocm7.1" in source
        assert "rocm7.2" in source

    def test_rocm_tag_validation_guard_exists(self):
        """install.sh should validate _rocm_tag with a case guard."""
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        assert "rocm[1-9]*.[0-9]*)" in source
        assert '_rocm_tag=""' in source  # rejection path

    def test_dpkg_epoch_handling(self):
        """install.sh should strip Debian epoch prefix from dpkg-query output."""
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        assert "sed 's/^[0-9]*://' " in source or "sed 's/^[0-9]*://'" in source

    def test_no_double_bracket_in_rocm_block(self):
        """ROCm detection block should not use [[ ]] (bash-only, not POSIX).
        Note: [[:space:]], [[:digit:]] etc. are valid POSIX character classes, not bash [[ ]]."""
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        func_start = source.find("get_torch_index_url()")
        func_end = source.find("\n}", func_start)
        func_body = source[func_start:func_end]
        import re

        for i, line in enumerate(func_body.splitlines(), 1):
            stripped = line.lstrip()
            if stripped.startswith("#"):
                continue
            # Remove POSIX character classes [[:foo:]] before checking for [[ ]]
            cleaned = re.sub(r"\[\[:[a-z]+:\]\]", "", line)
            assert "[[" not in cleaned, f"get_torch_index_url line {i} uses non-POSIX [["

    def test_no_arithmetic_expansion_in_rocm_block(self):
        """ROCm detection block should not use (( )) (bash-only)."""
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        func_start = source.find("get_torch_index_url()")
        func_end = source.find("\n}", func_start)
        func_body = source[func_start:func_end]
        for i, line in enumerate(func_body.splitlines(), 1):
            stripped = line.lstrip()
            if stripped.startswith("#"):
                continue
            assert (
                "((" not in line or "))" not in line or "$(()" in line
            ), f"get_torch_index_url line {i} may use non-POSIX (( ))"

    def test_macos_returns_cpu_before_rocm_check(self):
        """macOS should return CPU immediately (before any ROCm check)."""
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        func_start = source.find("get_torch_index_url()")
        func_body = source[func_start:]
        darwin_pos = func_body.find("Darwin")
        rocm_pos = func_body.find("amd-smi")
        assert darwin_pos < rocm_pos, "macOS check should come before ROCm detection"

    def test_unsloth_torch_backend_exported_after_get_torch_index_url(self):
        """install.sh must export UNSLOTH_TORCH_BACKEND after TORCH_INDEX_URL is set.

        This lets install_python_stack.py skip ROCm torch operations on CUDA
        and CPU hosts without re-running GPU detection in a subprocess.
        """
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        torch_url_pos = source.find("TORCH_INDEX_URL=$(get_torch_index_url)")
        backend_pos = source.find("UNSLOTH_TORCH_BACKEND")
        assert backend_pos > 0, "UNSLOTH_TORCH_BACKEND must be set in install.sh"
        assert (
            backend_pos > torch_url_pos
        ), "UNSLOTH_TORCH_BACKEND must be set AFTER TORCH_INDEX_URL is resolved"
        # Verify all three cases are covered
        assert '"cuda"' in source[backend_pos : backend_pos + 500]
        assert '"rocm"' in source[backend_pos : backend_pos + 500]
        assert '"cpu"' in source[backend_pos : backend_pos + 500]
        # Must be exported so subprocesses (setup.sh, install_python_stack.py) see it
        assert "export UNSLOTH_TORCH_BACKEND" in source

    def test_kfd_sysfs_amd_vendor_check_in_has_amd_rocm_gpu(self):
        """_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098.

        NVIDIA open kernel module (560+) registers KFD nodes with vendor_id
        4318 (0x10DE). Without the vendor check, _has_amd_rocm_gpu returns 0
        (true) on NVIDIA-only hosts that have the nvidia-open driver, causing
        get_torch_index_url to select a ROCm wheel index.
        """
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        func_start = source.find("_has_amd_rocm_gpu()")
        func_end = source.find("\n}", func_start)
        func_body = source[func_start:func_end]
        assert (
            "vendor_id" in func_body
        ), "_has_amd_rocm_gpu sysfs fallback must check vendor_id to exclude NVIDIA KFD nodes"
        assert (
            "4098" in func_body
        ), "_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (0x1002)"

    def test_kfd_awk_resets_state_per_file(self):
        """KFD sysfs awk must reset gpu/amd state per file (FNR==1).

        Without the reset, a Ryzen+NVIDIA host where node 0 is an AMD CPU
        agent (vendor_id 4098, gpu_id 0) and node 1 is an NVIDIA GPU
        (gpu_id > 0, vendor_id 4318) can produce a false positive: node 0
        sets amd=1, node 1 sets gpu=1, and the combined state triggers found=1
        before vendor_id 4318 is seen on node 1.
        """
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        func_start = source.find("_has_amd_rocm_gpu()")
        func_end = source.find("\n}", func_start)
        func_body = source[func_start:func_end]
        assert "FNR==1" in func_body, (
            "_has_amd_rocm_gpu KFD awk must reset state per file with FNR==1 "
            "to avoid false positives on Ryzen+NVIDIA hosts with multiple KFD nodes"
        )

    def test_get_torch_index_url_uses_nvidia_detected_flag(self):
        """get_torch_index_url must track NVIDIA detection independently of _smi.

        When _has_usable_nvidia_gpu returns true via /proc/driver/nvidia fallback
        but nvidia-smi is not on PATH, _smi stays empty. Without a separate
        _nvidia_detected flag, the function falls into the AMD/CPU branch even
        though NVIDIA was confirmed, silently installing CPU wheels instead of CUDA.
        """
        sh_path = PACKAGE_ROOT / "install.sh"
        source = sh_path.read_text(encoding = "utf-8")
        func_start = source.find("get_torch_index_url()")
        func_end = source.find("\n}", func_start)
        func_body = source[func_start:func_end]
        assert "_nvidia_detected" in func_body, (
            "get_torch_index_url must use a _nvidia_detected flag (separate from "
            "_smi) so that proc-only NVIDIA detection still selects CUDA wheels"
        )
        # The AMD/ROCm branch must be gated on _nvidia_detected being 0, not on
        # _smi being empty.
        assert (
            '_nvidia_detected" -eq 0' in func_body or "_nvidia_detected" in func_body
        ), "get_torch_index_url AMD branch must be skipped when _nvidia_detected=1"


# TEST: Live regression on current host (NVIDIA B200 expected)


class TestLiveRegression:
    """Live checks that run on the actual host -- skip if no NVIDIA GPU."""

    def test_get_torch_index_url_returns_cuda_on_nvidia(self):
        """On an NVIDIA machine, get_torch_index_url should return a CUDA URL."""
        import shutil

        if not shutil.which("nvidia-smi"):
            pytest.skip("No nvidia-smi available")
        # Skip if nvidia-smi exists but does not actually list a GPU on this
        # host (containers occasionally ship the binary without a driver).
        check = subprocess.run(
            [
                "bash",
                "-c",
                "nvidia-smi -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{f=1} END{exit !f}'",
            ],
            capture_output = True,
        )
        if check.returncode != 0:
            pytest.skip("nvidia-smi is on PATH but no GPU is listed")

        sh_path = PACKAGE_ROOT / "install.sh"
        # get_torch_index_url calls _has_usable_nvidia_gpu and
        # _has_amd_rocm_gpu, so all three function definitions must be
        # in scope when we eval the extract.
        extract_cmd = (
            f"sed -n '/^_has_amd_rocm_gpu()/,/^}}$/p; "
            f"/^_has_usable_nvidia_gpu()/,/^}}$/p; "
            f"/^get_torch_index_url()/,/^}}$/p' '{sh_path}'"
        )
        result = subprocess.run(
            ["bash", "-c", f'eval "$({extract_cmd})"; get_torch_index_url'],
            capture_output = True,
            text = True,
            timeout = 30,
        )
        if result.returncode != 0:
            pytest.skip("Could not extract get_torch_index_url for live test")
        url = result.stdout.strip()
        assert "cu1" in url or "cuda" in url.lower(), f"Expected CUDA URL, got: {url}"


# TEST: worker.py -- ROCm Mamba/SSM source build path

# Load worker.py module
_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py"
_EXPORT_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "export" / "worker.py"
# The torchao Windows-ROCm stub was de-duplicated out of the export/training
# workers into a shared module; both workers now call into it.
_TORCHAO_STUB_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "_torchao_stub.py"
# The wheel-probe subprocess was hoisted out of worker.py into wheel_utils
# during the wheel-resolver refactor; the probe script literal lives there.
_WHEEL_UTILS_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "wheel_utils.py"


class TestWorkerRocmMambaSsm:
    """Verify worker.py Mamba/SSM install logic on ROCm."""

    def test_probe_returns_hip_version_field(self):
        """The wheel probe should include hip_version, and worker.py should
        consume it."""
        assert "hip_version" in _WHEEL_UTILS_PATH.read_text(encoding = "utf-8")
        assert "hip_version" in _WORKER_PATH.read_text(encoding = "utf-8")

    def test_probe_script_has_getattr_hip(self):
        """Probe script should use getattr for torch.version.hip (safe on CUDA)."""
        source = _WHEEL_UTILS_PATH.read_text(encoding = "utf-8")
        assert "getattr(torch.version, 'hip', None)" in source

    def test_direct_wheel_url_returns_none_without_cuda_major(self, monkeypatch):
        """_direct_wheel_url should return None when cuda_major is empty (ROCm)."""
        # Load module for function access
        _worker_spec = importlib.util.spec_from_file_location("test_worker", _WORKER_PATH)
        assert _worker_spec is not None and _worker_spec.loader is not None
        worker_mod = importlib.util.module_from_spec(_worker_spec)

        # Stub worker.py's imports via monkeypatch so the fakes (notably a
        # non-package "utils") are undone and don't break later tests that
        # import the real utils.* package.
        loggers_mock = MagicMock()
        loggers_mock.get_logger = MagicMock(return_value = MagicMock())
        monkeypatch.setitem(sys.modules, "structlog", MagicMock())
        monkeypatch.setitem(sys.modules, "loggers", loggers_mock)
        monkeypatch.setitem(sys.modules, "utils", MagicMock())
        monkeypatch.setitem(sys.modules, "utils.hardware", MagicMock())

        try:
            _worker_spec.loader.exec_module(worker_mod)
        except Exception:
            pytest.skip("Could not load worker module in test environment")

        env_rocm = {
            "python_tag": "cp312",
            "torch_mm": "2.6",
            "cuda_major": "",
            "hip_version": "7.1.12345",
            "cxx11abi": "TRUE",
        }
        result = worker_mod._direct_wheel_url(
            filename_prefix = "causal_conv1d",
            package_version = "1.6.1",
            release_tag = "v1.6.1.post4",
            release_base_url = "https://github.com/Dao-AILab/causal-conv1d/releases/download",
            env = env_rocm,
        )
        assert result is None

    def test_hipcc_check_exists_in_source(self):
        """worker.py should check for hipcc before ROCm source builds."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        assert "hipcc" in source

    def test_rocm_source_build_status_message(self):
        """worker.py should send a specific status for ROCm source compilation."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        assert "Compiling" in source and "from source for ROCm" in source

    def test_rocm_build_failure_message(self):
        """worker.py should send a clear error on ROCm build failure."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        assert "Failed to compile" in source and "for ROCm" in source

    def test_timeout_on_install(self):
        """worker.py should have a timeout on pip install subprocess."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        assert "TimeoutExpired" in source
        assert "timeout" in source


# TEST: amd.py -- AMD GPU monitoring


class TestAmdGpuMonitoring:
    """Verify amd.py module structure and mock behavior."""

    def test_amd_py_exists(self):
        """amd.py should exist in the hardware directory."""
        amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
        assert amd_path.exists()

    def test_amd_py_has_required_functions(self):
        """amd.py should export the same function signatures as nvidia.py."""
        amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
        source = amd_path.read_text(encoding = "utf-8")
        assert "def get_physical_gpu_count" in source
        assert "def get_primary_gpu_utilization" in source
        assert "def get_visible_gpu_utilization" in source

    def test_amd_smi_json_parsing(self, monkeypatch):
        """Verify _extract_gpu_metrics parses amd-smi JSON correctly."""
        amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
        _amd_spec = importlib.util.spec_from_file_location("test_amd", amd_path)
        assert _amd_spec is not None and _amd_spec.loader is not None
        amd_mod = importlib.util.module_from_spec(_amd_spec)

        loggers_mock = MagicMock()
        loggers_mock.get_logger = MagicMock(return_value = MagicMock())
        monkeypatch.setitem(sys.modules, "loggers", loggers_mock)

        try:
            _amd_spec.loader.exec_module(amd_mod)
        except Exception:
            pytest.skip("Could not load amd module in test environment")

        # Simulate amd-smi metric JSON output
        gpu_data = {
            "usage": {"gfx_activity": "85"},
            "temperature": {"edge": "72"},
            "power": {
                "current_socket_power": "200.5",
                "power_cap": "300",
            },
            "vram": {
                "vram_used": 8192,  # MB
                "vram_total": 16384,  # MB
            },
        }
        metrics = amd_mod._extract_gpu_metrics(gpu_data)
        assert metrics["gpu_utilization_pct"] == 85.0
        assert metrics["temperature_c"] == 72.0
        assert metrics["power_draw_w"] == 200.5
        assert metrics["power_limit_w"] == 300.0
        assert metrics["vram_used_gb"] == round(8192 / 1024, 2)
        assert metrics["vram_total_gb"] == round(16384 / 1024, 2)
        assert metrics["vram_utilization_pct"] is not None
        assert metrics["power_utilization_pct"] is not None

    def test_amd_primary_gpu_with_mock(self, monkeypatch):
        """get_primary_gpu_utilization returns correct dict with mocked amd-smi."""
        amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
        _amd_spec = importlib.util.spec_from_file_location("test_amd2", amd_path)
        assert _amd_spec is not None and _amd_spec.loader is not None
        amd_mod = importlib.util.module_from_spec(_amd_spec)

        loggers_mock = MagicMock()
        loggers_mock.get_logger = MagicMock(return_value = MagicMock())
        monkeypatch.setitem(sys.modules, "loggers", loggers_mock)

        try:
            _amd_spec.loader.exec_module(amd_mod)
        except Exception:
            pytest.skip("Could not load amd module")

        # _first_visible_amd_gpu_id() short-circuits to None when any of
        # HIP / ROCR / CUDA_VISIBLE_DEVICES is "" or "-1". CI often sets
        # CUDA_VISIBLE_DEVICES="", so the test must not inherit that.
        for var in (
            "HIP_VISIBLE_DEVICES",
            "ROCR_VISIBLE_DEVICES",
            "CUDA_VISIBLE_DEVICES",
        ):
            monkeypatch.delenv(var, raising = False)

        # amd-smi is gated off on Windows w/o a HIP SDK; this test mocks it as
        # available, so opt in so the gate allows it on every platform.
        monkeypatch.setenv("UNSLOTH_ENABLE_AMD_SMI", "1")

        mock_json = json.dumps(
            [
                {
                    "usage": {"gfx_activity": "50"},
                    "temperature": {"edge": "65"},
                    "power": {"current_socket_power": "150", "power_cap": "250"},
                    "vram": {"vram_used": 4096, "vram_total": 16384},
                }
            ]
        )
        mock_result = MagicMock()
        mock_result.returncode = 0
        mock_result.stdout = mock_json

        # The premise is "amd-smi exists and answers": the absence guard
        # which()-checks before spawning, so hosts without a real amd-smi
        # (Linux CI, driver-only Windows) need which mocked too.
        with patch.object(amd_mod.shutil, "which", return_value = "/usr/bin/amd-smi"):
            with patch.object(subprocess, "run", return_value = mock_result):
                result = amd_mod.get_primary_gpu_utilization()
        assert result["available"] is True
        assert result["gpu_utilization_pct"] == 50.0
        assert result["temperature_c"] == 65.0

    def test_amd_smi_not_found_returns_unavailable(self, monkeypatch):
        """get_primary_gpu_utilization returns available=False when amd-smi is missing."""
        amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
        _amd_spec = importlib.util.spec_from_file_location("test_amd3", amd_path)
        assert _amd_spec is not None and _amd_spec.loader is not None
        amd_mod = importlib.util.module_from_spec(_amd_spec)

        loggers_mock = MagicMock()
        loggers_mock.get_logger = MagicMock(return_value = MagicMock())
        monkeypatch.setitem(sys.modules, "loggers", loggers_mock)

        try:
            _amd_spec.loader.exec_module(amd_mod)
        except Exception:
            pytest.skip("Could not load amd module")

        # Opt in so the call reaches subprocess.run (gated off on Windows w/o a
        # HIP SDK); testing the OSError handling here.
        with (
            patch.dict(os.environ, {"UNSLOTH_ENABLE_AMD_SMI": "1"}),
            patch.object(subprocess, "run", side_effect = OSError("amd-smi not found")),
        ):
            result = amd_mod.get_primary_gpu_utilization()
        assert result["available"] is False

    def test_amd_timeout_returns_unavailable(self, monkeypatch):
        """get_primary_gpu_utilization handles timeout gracefully."""
        amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
        _amd_spec = importlib.util.spec_from_file_location("test_amd4", amd_path)
        assert _amd_spec is not None and _amd_spec.loader is not None
        amd_mod = importlib.util.module_from_spec(_amd_spec)

        loggers_mock = MagicMock()
        loggers_mock.get_logger = MagicMock(return_value = MagicMock())
        monkeypatch.setitem(sys.modules, "loggers", loggers_mock)

        try:
            _amd_spec.loader.exec_module(amd_mod)
        except Exception:
            pytest.skip("Could not load amd module")

        # Opt in so the call reaches subprocess.run (gated off on Windows w/o a
        # HIP SDK); testing the timeout handling here.
        with (
            patch.dict(os.environ, {"UNSLOTH_ENABLE_AMD_SMI": "1"}),
            patch.object(
                subprocess,
                "run",
                side_effect = subprocess.TimeoutExpired("amd-smi", 5),
            ),
        ):
            result = amd_mod.get_primary_gpu_utilization()
        assert result["available"] is False


# TEST: hardware.py -- IS_ROCM branching to amd.py


class TestHardwareAmdBranching:
    """Verify hardware.py branches to amd.py when IS_ROCM is True."""

    def test_hardware_imports_amd_module(self):
        """hardware.py should import from amd module when IS_ROCM."""
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        assert "from . import amd" in source

    def test_hardware_branches_on_is_rocm_for_utilization(self):
        """get_gpu_utilization should dispatch to amd.py via _smi_query
        when IS_ROCM, and the dispatcher itself must check IS_ROCM and
        import the amd backend."""
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        func_start = source.find("def get_gpu_utilization")
        func_body = source[func_start : source.find("\ndef ", func_start + 1)]
        assert '_smi_query("get_primary_gpu_utilization"' in func_body
        smi = source[
            source.find("def _smi_query") : source.find("\ndef ", source.find("def _smi_query") + 1)
        ]
        assert "IS_ROCM" in smi
        assert "from . import amd" in smi

    def test_hardware_branches_on_is_rocm_for_visible(self):
        """get_visible_gpu_utilization should dispatch to amd.py via
        _smi_query when IS_ROCM."""
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        func_start = source.find("def get_visible_gpu_utilization")
        func_body = source[func_start : source.find("\ndef ", func_start + 1)]
        # The dispatcher call may wrap onto multiple lines; allow whitespace
        # between the open paren and the literal func name argument.
        import re as _re

        assert _re.search(r'_smi_query\(\s*"get_visible_gpu_utilization"', func_body)
        smi = source[
            source.find("def _smi_query") : source.find("\ndef ", source.find("def _smi_query") + 1)
        ]
        assert "IS_ROCM" in smi
        assert "from . import amd" in smi

    def test_hardware_branches_on_is_rocm_for_physical_count(self):
        """get_physical_gpu_count should try amd.py when IS_ROCM."""
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        func_start = source.find("def get_physical_gpu_count")
        func_body = source[func_start : source.find("\ndef ", func_start + 1)]
        assert "IS_ROCM" in func_body
        assert "from . import amd" in func_body


# TEST: hardware.py -- apply_gpu_ids ROCm fallback (issue #5180)


class TestApplyGpuIdsRocmFallback:
    """Verify apply_gpu_ids sets HIP_VISIBLE_DEVICES on ROCm hosts even when
    IS_ROCM is still False (worker subprocess before detect_hardware runs)."""

    def test_apply_gpu_ids_falls_back_to_torch_version_hip(self):
        """apply_gpu_ids should probe torch.version.hip when IS_ROCM is False and no ROCm env vars are set."""
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        func_start = source.find("def apply_gpu_ids")
        func_body = source[func_start : source.find("\ndef ", func_start + 1)]
        assert 'getattr(_torch.version, "hip", None)' in func_body

    def test_apply_gpu_ids_sets_hip_and_rocr_visible_devices(self):
        """apply_gpu_ids should set both HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES on ROCm."""
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        func_start = source.find("def apply_gpu_ids")
        func_body = source[func_start : source.find("\ndef ", func_start + 1)]
        assert 'os.environ["HIP_VISIBLE_DEVICES"] = value' in func_body
        assert 'os.environ["ROCR_VISIBLE_DEVICES"] = value' in func_body

    def test_apply_gpu_ids_rocm_fallback_is_guarded_by_try_except(self):
        """torch import in apply_gpu_ids must be wrapped in try/except so a missing torch never crashes."""
        hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
        source = hw_path.read_text(encoding = "utf-8")
        func_start = source.find("def apply_gpu_ids")
        func_body = source[func_start : source.find("\ndef ", func_start + 1)]
        assert "import torch as _torch" in func_body
        assert "except Exception" in func_body


# TEST: install_python_stack.py -- Windows AMD warning


class TestWindowsRocmWarning:
    """Verify Windows AMD GPU detection and warning message."""

    def test_windows_amd_warning_in_source(self):
        """install_python_stack.py should warn Windows AMD users."""
        source = _STACK_PATH.read_text(encoding = "utf-8")
        assert "AMD GPU detected" in source

    def test_windows_amd_warning_checks_hipinfo_or_amdsmi(self):
        """Warning should check for hipinfo or amd-smi."""
        source = _STACK_PATH.read_text(encoding = "utf-8")
        assert "hipinfo" in source
        assert "amd-smi" in source

    def test_windows_amd_warning_has_docs_link(self):
        """Warning should include AMD docs link."""
        source = _STACK_PATH.read_text(encoding = "utf-8")
        assert "docs.unsloth.ai/get-started/install-and-update/amd" in source


# TEST: unsloth/kernels/utils.py -- is_rdna() expansion


class TestIsRdnaExpansion:
    """Verify is_rdna() covers RDNA2, RDNA3, RDNA3.5, RDNA4 architectures."""

    def test_is_rdna_source_has_rdna2(self):
        """is_rdna() should include RDNA2 architectures."""
        utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
        source = utils_path.read_text(encoding = "utf-8")
        func_start = source.find("def is_rdna()")
        func_body = source[func_start : source.find("\ndef ", func_start + 1)]
        assert "gfx1030" in func_body
        assert "gfx1031" in func_body
        assert "gfx1032" in func_body
        assert "gfx1033" in func_body
        assert "gfx1034" in func_body
        assert "gfx1035" in func_body
        assert "gfx1036" in func_body

    def test_is_rdna_source_has_rdna3(self):
        """is_rdna() should include RDNA3 architectures."""
        utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
        source = utils_path.read_text(encoding = "utf-8")
        func_start = source.find("def is_rdna()")
        func_body = source[func_start : source.find("\ndef ", func_start + 1)]
        assert "gfx1100" in func_body
        assert "gfx1101" in func_body
        assert "gfx1102" in func_body
        assert "gfx1103" in func_body

    def test_is_rdna_source_has_rdna35(self):
        """is_rdna() should include RDNA3.5 architectures."""
        utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
        source = utils_path.read_text(encoding = "utf-8")
        func_start = source.find("def is_rdna()")
        func_body = source[func_start : source.find("\ndef ", func_start + 1)]
        assert "gfx1150" in func_body
        assert "gfx1151" in func_body
        assert "gfx1152" in func_body

    def test_is_rdna_source_has_rdna4(self):
        """is_rdna() should include RDNA4 architectures."""
        utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
        source = utils_path.read_text(encoding = "utf-8")
        func_start = source.find("def is_rdna()")
        func_body = source[func_start : source.find("\ndef ", func_start + 1)]
        assert "gfx1200" in func_body
        assert "gfx1201" in func_body

    def test_is_cdna_not_changed(self):
        """is_cdna() should remain unchanged (no RDNA architectures added)."""
        utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py"
        source = utils_path.read_text(encoding = "utf-8")
        func_start = source.find("def is_cdna()")
        func_body = source[func_start : source.find("\ndef ", func_start + 1)]
        assert "gfx940" in func_body
        assert "gfx941" in func_body
        assert "gfx942" in func_body
        assert "gfx950" in func_body
        # RDNA architectures should NOT be in is_cdna
        assert "gfx1030" not in func_body
        assert "gfx1100" not in func_body


# TEST: install_python_stack.py -- _windows_rocm_index_url arch mapping


class TestWindowsRocmIndexUrl:
    """Verify GPU arch → AMD pip index URL mapping."""

    def test_gfx1200_maps_to_gfx120x_all(self):
        url = stack_mod._windows_rocm_index_url("gfx1200")
        assert url is not None
        assert "gfx120X-all" in url

    def test_gfx1201_maps_to_gfx120x_all(self):
        url = stack_mod._windows_rocm_index_url("gfx1201")
        assert url is not None
        assert "gfx120X-all" in url

    def test_gfx1151_maps_to_gfx1151(self):
        url = stack_mod._windows_rocm_index_url("gfx1151")
        assert url is not None
        assert "gfx1151" in url

    def test_gfx1150_maps_to_gfx1150(self):
        url = stack_mod._windows_rocm_index_url("gfx1150")
        assert url is not None
        assert "gfx1150" in url

    def test_gfx1100_maps_to_gfx110x_all(self):
        url = stack_mod._windows_rocm_index_url("gfx1100")
        assert url is not None
        assert "gfx110X-all" in url

    def test_unknown_arch_returns_none(self):
        assert stack_mod._windows_rocm_index_url("gfx9999") is None

    def test_none_arch_returns_none(self):
        assert stack_mod._windows_rocm_index_url(None) is None

    def test_url_ends_with_slash(self):
        """AMD pip index URLs must end with / for --index-url compatibility."""
        url = stack_mod._windows_rocm_index_url("gfx1200")
        assert url is not None
        assert url.endswith("/")

    def test_base_url_uses_repo_amd_com_by_default(self):
        url = stack_mod._windows_rocm_index_url("gfx1200")
        assert url is not None
        assert "repo.amd.com" in url

    def test_mirror_env_var_overrides_base(self, monkeypatch):
        monkeypatch.setenv("UNSLOTH_ROCM_WINDOWS_MIRROR", "https://my-mirror.example.com/rocm/whl")
        # Reload module-level constant by calling helper directly
        url = stack_mod._windows_rocm_index_url("gfx1200")
        # The env var is read at module load time for _ROCM_WINDOWS_INDEX_BASE,
        # so just verify the helper itself doesn't error.
        assert url is not None


# TEST: install_python_stack.py -- _detect_windows_gfx_arch


class TestDetectWindowsGfxArch:
    """Verify hipinfo parsing for GPU arch detection on Windows."""

    def test_returns_none_when_hipinfo_not_on_path(self):
        # Also neutralise the venv-hipInfo and WMI-name fallbacks: this test
        # pins "no probe source available -> None", and the suite may run on a
        # real AMD host where WMI would legitimately answer.
        with patch("shutil.which", return_value = None):
            with patch("os.path.isfile", return_value = False):
                with patch("subprocess.run", side_effect = FileNotFoundError):
                    result = stack_mod._detect_windows_gfx_arch()
        assert result is None

    def test_parses_gcnarchname_from_hipinfo_output(self):
        mock_result = MagicMock()
        mock_result.returncode = 0
        mock_result.stdout = b"gcnArchName : gfx1200\nsome other line\n"
        with patch("shutil.which", return_value = "/usr/bin/hipinfo"):
            with patch("subprocess.run", return_value = mock_result):
                result = stack_mod._detect_windows_gfx_arch()
        assert result == "gfx1200"

    def test_returns_none_on_nonzero_returncode(self):
        mock_result = MagicMock()
        mock_result.returncode = 1
        mock_result.stdout = b"gcnArchName : gfx1200\n"
        with patch("shutil.which", return_value = "/usr/bin/hipinfo"):
            with patch("subprocess.run", return_value = mock_result):
                result = stack_mod._detect_windows_gfx_arch()
        assert result is None

    def test_returns_none_when_no_gcnarchname_in_output(self):
        # hipinfo answers but without a gcnArchName line. Route only the
        # hipinfo/amd-smi probes to the mock; the WMI fallback must get
        # nothing (FileNotFoundError) -- otherwise the mocked device name
        # would legitimately resolve via the name table.
        mock_result = MagicMock()
        mock_result.returncode = 0
        mock_result.stdout = b"deviceName : SomeUnknownDevice\n"

        def _run(cmd, **kwargs):
            if cmd and "powershell" in str(cmd[0]).lower():
                raise FileNotFoundError(cmd[0])
            return mock_result

        with patch("shutil.which", return_value = "/usr/bin/hipinfo"):
            with patch("subprocess.run", side_effect = _run):
                result = stack_mod._detect_windows_gfx_arch()
        assert result is None

    def test_returns_none_on_timeout(self):
        with patch("shutil.which", return_value = "/usr/bin/hipinfo"):
            with patch(
                "subprocess.run",
                side_effect = subprocess.TimeoutExpired("hipinfo", 10),
            ):
                result = stack_mod._detect_windows_gfx_arch()
        assert result is None

    def test_strips_whitespace_from_arch(self):
        mock_result = MagicMock()
        mock_result.returncode = 0
        mock_result.stdout = b"  gcnArchName :   gfx1201  \n"
        with patch("shutil.which", return_value = "/usr/bin/hipinfo"):
            with patch("subprocess.run", return_value = mock_result):
                result = stack_mod._detect_windows_gfx_arch()
        assert result == "gfx1201"


# TEST: install_python_stack.py -- GPU-name / WMI fallback (no amd-smi, no hipinfo)


class TestGfxArchNameFallback:
    """amd-smi does not exist on Windows (neither Adrenalin consistently nor
    the HIP SDK ship a CLI) and driver-only hosts lack hipinfo too. The
    detection chain must still resolve the arch from the GPU marketing name
    (WMI), mirroring setup.ps1's $nameArchTable."""

    @pytest.mark.parametrize(
        "name, expected",
        [
            ("AMD Radeon(TM) 8060S Graphics", "gfx1151"),
            ("AMD Ryzen AI MAX+ 395 w/ Radeon 8060S", "gfx1151"),
            ("AMD Radeon(TM) 890M", "gfx1150"),
            ("AMD Ryzen AI 9 HX 370 w/ Radeon 890M", "gfx1150"),
            ("AMD Radeon RX 9070 XT", "gfx1201"),
            ("AMD Radeon RX 9070", "gfx1200"),
            ("AMD Radeon RX 7700S", "gfx1102"),  # (?!S) lookahead must not hit gfx1100
            ("AMD Radeon RX 7700 XT", "gfx1100"),
            ("AMD Radeon(TM) 780M", "gfx1103"),
            ("NVIDIA GeForce RTX 4090", None),
            ("Microsoft Basic Display Adapter", None),
            ("", None),
        ],
    )
    def test_name_to_arch_mapping(self, name, expected):
        assert stack_mod._gfx_arch_from_gpu_name(name) == expected

    def test_wmi_fallback_resolves_arch_without_any_tools(self):
        """hipinfo absent everywhere + amd-smi absent -> WMI name fallback."""
        ps_result = MagicMock()
        ps_result.returncode = 0
        ps_result.stdout = b"AMD Radeon(TM) 8060S Graphics\r\nMicrosoft Basic Display Adapter\r\n"

        def _run(cmd, **kwargs):
            if cmd and "powershell.exe" in str(cmd[0]).lower():
                return ps_result
            raise FileNotFoundError(cmd[0])

        with patch.dict(os.environ, {}, clear = False):
            for _v in (
                "HIP_PATH",
                "ROCM_PATH",
                "UNSLOTH_ROCM_GFX_ARCH",
                "UNSLOTH_ENABLE_AMD_SMI",
            ):
                os.environ.pop(_v, None)
            with patch("shutil.which", return_value = None):
                with patch("os.path.isfile", return_value = False):
                    with patch("subprocess.run", side_effect = _run):
                        result = stack_mod._detect_windows_gfx_arch()
        assert result == "gfx1151"

    def test_wmi_fallback_returns_none_for_non_amd_hosts(self):
        ps_result = MagicMock()
        ps_result.returncode = 0
        ps_result.stdout = b"NVIDIA GeForce RTX 4090\r\n"

        def _run(cmd, **kwargs):
            if cmd and "powershell.exe" in str(cmd[0]).lower():
                return ps_result
            raise FileNotFoundError(cmd[0])

        with patch.dict(os.environ, {}, clear = False):
            for _v in ("HIP_PATH", "ROCM_PATH", "UNSLOTH_ROCM_GFX_ARCH"):
                os.environ.pop(_v, None)
            with patch("shutil.which", return_value = None):
                with patch("os.path.isfile", return_value = False):
                    with patch("subprocess.run", side_effect = _run):
                        result = stack_mod._detect_windows_gfx_arch()
        assert result is None

    def test_stack_probes_venv_hipinfo(self):
        """The venv Scripts dir hipInfo.exe (shipped by AMD torch wheels) must
        be a probe candidate so `studio update` works on driver-only hosts."""
        source = _STACK_PATH.read_text(encoding = "utf-8")
        assert 'os.path.join(os.path.dirname(sys.executable), "hipInfo.exe")' in source

    def test_prebuilt_resolve_exe_probes_venv_dir(self):
        """install_llama_prebuilt's _resolve_exe must include the venv Scripts
        candidate for the same driver-only standalone-rerun scenario."""
        source = _PREBUILT_PATH.read_text(encoding = "utf-8")
        assert "_venv_candidate" in source

    def test_runtime_monitor_guards_amd_smi_absence(self):
        """amd.py must which()-check amd-smi before spawning so absence
        disables the poller in one step (no FileNotFoundError strikes)."""
        amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
        source = amd_path.read_text(encoding = "utf-8")
        assert 'shutil.which("amd-smi") is None' in source


# TEST: install_python_stack.py -- _install_bnb_windows_rocm


class TestInstallBnbWindowsRocm:
    """Verify AMD Windows BNB wheel install helper."""

    @pytest.fixture(autouse = True)
    def _isolate_sitecustomize_persistence(self, monkeypatch, request):
        """Keep helper tests from writing to the active interpreter site-packages."""
        if request.node.name.startswith("test_persist"):
            return
        monkeypatch.setattr(
            stack_mod,
            "_persist_bnb_rocm_version",
            lambda version: True,
        )

    def test_calls_pip_install_try_with_win_amd64_url(self):
        """Should call pip_install_try with the win_amd64 wheel URL via plain pip."""
        with patch.object(stack_mod, "pip_install_try", return_value = True) as mock_pip:
            stack_mod._install_bnb_windows_rocm()
        assert mock_pip.call_count == 1
        call_args = str(mock_pip.call_args_list[0])
        assert "bitsandbytes" in call_args
        assert "win_amd64" in call_args
        # Must force plain pip (uv mangles the bitsandbytes wheel) -- see
        # https://unsloth.ai/docs/get-started/install/amd/amd-hackathon
        assert mock_pip.call_args.kwargs.get("force_pip") is True

    def test_forces_plain_pip_not_uv(self):
        """The bnb wheel must be installed with plain pip, never uv."""
        with patch.object(stack_mod, "pip_install_try", return_value = True) as mock_pip:
            stack_mod._install_bnb_windows_rocm()
        assert mock_pip.call_args.kwargs.get("force_pip") is True

    def test_does_not_touch_uv_skip_env_var(self):
        """The UV_SKIP_WHEEL_FILENAME_CHECK hack is gone; the env must be untouched."""
        observed = {}

        def _capture(*args, **kwargs):
            observed["during"] = os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK")
            return True

        with patch.dict(os.environ, {}, clear = False):
            os.environ.pop("UV_SKIP_WHEEL_FILENAME_CHECK", None)
            with patch.object(stack_mod, "pip_install_try", side_effect = _capture):
                stack_mod._install_bnb_windows_rocm()
            assert observed.get("during") is None
            assert "UV_SKIP_WHEEL_FILENAME_CHECK" not in os.environ

    def test_returns_false_on_pip_failure(self):
        """A failed pip_install_try must surface as a False return, not BNB_ROCM_VERSION."""
        with patch.dict(os.environ, {}, clear = False):
            os.environ.pop("BNB_ROCM_VERSION", None)
            with patch.object(stack_mod, "pip_install_try", return_value = False):
                result = stack_mod._install_bnb_windows_rocm()
            assert result is False
            assert "BNB_ROCM_VERSION" not in os.environ

    def test_no_op_when_win_amd64_url_missing(self):
        """Should be silent no-op if win_amd64 key absent from _BNB_ROCM_PRERELEASE_URLS."""
        with patch.object(stack_mod, "_BNB_ROCM_PRERELEASE_URLS", {}):
            with patch.object(stack_mod, "pip_install_try") as mock_pip:
                stack_mod._install_bnb_windows_rocm()
        mock_pip.assert_not_called()

    def test_sets_bnb_rocm_version_from_detected_dll(self):
        """BNB_ROCM_VERSION is set from the DLL detected after install."""
        with patch.dict(os.environ, {}, clear = False):
            os.environ.pop("BNB_ROCM_VERSION", None)
            os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None)
            with patch.object(stack_mod, "pip_install_try", return_value = True):
                with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "72"):
                    stack_mod._install_bnb_windows_rocm()
            assert os.environ.get("BNB_ROCM_VERSION") == "72"

    def test_sets_bnb_rocm_version_from_newer_dll(self):
        """If AMD ships a newer DLL (e.g. rocm713.dll), that version is used."""
        with patch.dict(os.environ, {}, clear = False):
            os.environ.pop("BNB_ROCM_VERSION", None)
            os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None)
            with patch.object(stack_mod, "pip_install_try", return_value = True):
                with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "713"):
                    stack_mod._install_bnb_windows_rocm()
            assert os.environ.get("BNB_ROCM_VERSION") == "713"

    def test_falls_back_to_72_when_detection_fails(self):
        """Falls back to '72' when DLL detection returns None."""
        with patch.dict(os.environ, {}, clear = False):
            os.environ.pop("BNB_ROCM_VERSION", None)
            os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None)
            with patch.object(stack_mod, "pip_install_try", return_value = True):
                with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = None):
                    stack_mod._install_bnb_windows_rocm()
            assert os.environ.get("BNB_ROCM_VERSION") == "72"

    def test_does_not_override_existing_bnb_rocm_version(self):
        """An explicit BNB_ROCM_VERSION in the caller's env must not be clobbered."""
        with patch.dict(os.environ, {"BNB_ROCM_VERSION": "60"}):
            os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None)
            with patch.object(stack_mod, "pip_install_try", return_value = True):
                stack_mod._install_bnb_windows_rocm()
            assert os.environ.get("BNB_ROCM_VERSION") == "60"

    def test_does_not_persist_existing_bnb_rocm_version(self):
        """A caller override must not become the venv's managed default."""
        with patch.dict(os.environ, {"BNB_ROCM_VERSION": "60"}):
            os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None)
            with patch.object(stack_mod, "pip_install_try", return_value = True):
                with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver") as mock_detect:
                    with patch.object(
                        stack_mod, "_persist_bnb_rocm_version", return_value = True
                    ) as mock_persist:
                        stack_mod._install_bnb_windows_rocm()

            assert os.environ.get("BNB_ROCM_VERSION") == "60"
            mock_detect.assert_not_called()
            mock_persist.assert_not_called()

    def test_redetects_when_bnb_rocm_version_came_from_sitecustomize(self):
        """Persisted defaults should not mask a newer DLL suffix after reinstall."""
        with patch.dict(
            os.environ,
            {
                "BNB_ROCM_VERSION": "72",
                stack_mod._BNB_ROCM_VERSION_SOURCE_ENV: (
                    stack_mod._BNB_ROCM_VERSION_SOURCE_SITECUSTOMIZE
                ),
            },
        ):
            with patch.object(stack_mod, "pip_install_try", return_value = True):
                with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "713"):
                    with patch.object(
                        stack_mod, "_persist_bnb_rocm_version", return_value = True
                    ) as mock_persist:
                        stack_mod._install_bnb_windows_rocm()

            assert os.environ.get("BNB_ROCM_VERSION") == "713"
            assert (
                os.environ.get(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV)
                == stack_mod._BNB_ROCM_VERSION_SOURCE_DETECTED
            )
            mock_persist.assert_called_once_with("713")

    def test_persists_bnb_rocm_version_for_direct_venv_python(self, tmp_path):
        """BNB_ROCM_VERSION must apply to a fresh Python process in the venv."""
        site_packages = tmp_path / "site-packages"

        with patch.dict(os.environ, {}, clear = False):
            os.environ.pop("BNB_ROCM_VERSION", None)
            os.environ.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None)
            with patch.object(stack_mod, "pip_install_try", return_value = True):
                with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "72"):
                    with patch.object(
                        stack_mod.sysconfig, "get_path", return_value = str(site_packages)
                    ):
                        stack_mod._install_bnb_windows_rocm()

        sitecustomize = site_packages / "sitecustomize.py"
        source = sitecustomize.read_text(encoding = "utf-8")
        assert "BNB_ROCM_VERSION" in source
        assert stack_mod._BNB_ROCM_VERSION_SOURCE_ENV in source
        assert "'72'" in source

        probe_env = os.environ.copy()
        probe_env.pop("BNB_ROCM_VERSION", None)
        probe_env.pop(stack_mod._BNB_ROCM_VERSION_SOURCE_ENV, None)
        probe_env["PYTHONPATH"] = str(site_packages)
        result = subprocess.run(
            [
                sys.executable,
                "-c",
                (
                    "import os; "
                    "print(os.environ.get('BNB_ROCM_VERSION', ''), "
                    "os.environ.get('UNSLOTH_BNB_ROCM_VERSION_SOURCE', ''))"
                ),
            ],
            env = probe_env,
            stdout = subprocess.PIPE,
            stderr = subprocess.PIPE,
            text = True,
            check = True,
        )
        assert result.stdout.strip() == "72 sitecustomize"

    def test_persist_bnb_rocm_version_replaces_existing_managed_block(self, tmp_path):
        """Updating sitecustomize.py must not duplicate the managed BNB block."""
        site_packages = tmp_path / "site-packages"
        site_packages.mkdir()
        sitecustomize = site_packages / "sitecustomize.py"
        sitecustomize.write_text(
            "EXISTING = True\n"
            "# BEGIN Unsloth BNB_ROCM_VERSION\n"
            "import os as _unsloth_os\n"
            "_unsloth_os.environ.setdefault('BNB_ROCM_VERSION', '72')\n"
            "# END Unsloth BNB_ROCM_VERSION\n",
            encoding = "utf-8",
        )

        with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)):
            assert stack_mod._persist_bnb_rocm_version("713") is True

        source = sitecustomize.read_text(encoding = "utf-8")
        assert source.count("# BEGIN Unsloth BNB_ROCM_VERSION") == 1
        assert "EXISTING = True" in source
        assert "'713'" in source
        assert "'72'" not in source

    def test_persist_bnb_rocm_version_handles_non_utf8_sitecustomize(self, tmp_path):
        """A legacy non-UTF-8 sitecustomize.py should not abort installation."""
        site_packages = tmp_path / "site-packages"
        site_packages.mkdir()
        sitecustomize = site_packages / "sitecustomize.py"
        sitecustomize.write_bytes(b"\xff\xfe\x00")

        with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)):
            assert stack_mod._persist_bnb_rocm_version("72") is False

    def test_persist_bnb_rocm_version_repairs_truncated_block(self, tmp_path):
        """A managed block missing its END marker is replaced, not duplicated."""
        site_packages = tmp_path / "site-packages"
        site_packages.mkdir()
        sitecustomize = site_packages / "sitecustomize.py"
        sitecustomize.write_text(
            "EXISTING = True\n"
            "# BEGIN Unsloth BNB_ROCM_VERSION\n"
            "import os as _unsloth_os\n"
            "_unsloth_os.environ.setdefault('BNB_ROCM_VERSION', '72')\n",
            encoding = "utf-8",
        )

        with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)):
            assert stack_mod._persist_bnb_rocm_version("713") is True

        source = sitecustomize.read_text(encoding = "utf-8")
        assert source.count("# BEGIN Unsloth BNB_ROCM_VERSION") == 1
        assert source.count("# END Unsloth BNB_ROCM_VERSION") == 1
        assert "EXISTING = True" in source
        assert "'713'" in source
        assert "'72'" not in source

    def test_persist_bnb_rocm_version_dedupes_duplicate_blocks(self, tmp_path):
        """Multiple managed blocks collapse to one while preserving user content."""
        site_packages = tmp_path / "site-packages"
        site_packages.mkdir()
        sitecustomize = site_packages / "sitecustomize.py"
        block = (
            "# BEGIN Unsloth BNB_ROCM_VERSION\n"
            "import os as _unsloth_os\n"
            "_unsloth_os.environ.setdefault('BNB_ROCM_VERSION', '72')\n"
            "# END Unsloth BNB_ROCM_VERSION\n"
        )
        sitecustomize.write_text(block + "USER_MID = 1\n" + block, encoding = "utf-8")

        with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)):
            assert stack_mod._persist_bnb_rocm_version("713") is True

        source = sitecustomize.read_text(encoding = "utf-8")
        assert source.count("# BEGIN Unsloth BNB_ROCM_VERSION") == 1
        assert source.count("# END Unsloth BNB_ROCM_VERSION") == 1
        assert "USER_MID = 1" in source
        assert "'713'" in source
        assert "'72'" not in source

    def test_persist_bnb_rocm_version_atomic_no_leftover_tmp(self, tmp_path):
        """The write-then-rename path must not leave its temp file behind."""
        site_packages = tmp_path / "site-packages"
        site_packages.mkdir()

        with patch.object(stack_mod.sysconfig, "get_path", return_value = str(site_packages)):
            assert stack_mod._persist_bnb_rocm_version("72") is True

        leftovers = [p.name for p in site_packages.iterdir() if "unsloth-tmp" in p.name]
        assert leftovers == []
        assert (site_packages / "sitecustomize.py").exists()


class TestRuntimeBnbRocmSourceGuards:
    """Runtime entrypoints redetect managed defaults but keep caller overrides."""

    _MAIN_PATH = PACKAGE_ROOT / "studio" / "backend" / "main.py"
    _TRAINING_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py"

    def test_main_gate_redetects_persisted_default(self):
        source = self._MAIN_PATH.read_text(encoding = "utf-8")
        assert 'os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize"' in source
        assert 'os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"' in source

    def test_worker_gate_redetects_persisted_default(self):
        source = self._TRAINING_WORKER_PATH.read_text(encoding = "utf-8")
        assert 'os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize"' in source
        assert 'os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"' in source

    def test_fallback_prefers_seeded_value_over_hardcoded_72(self):
        """A failed redetect must not downgrade a persisted suffix to '72'."""
        for path in (self._MAIN_PATH, self._TRAINING_WORKER_PATH):
            source = path.read_text(encoding = "utf-8")
            assert 'os.environ.get("BNB_ROCM_VERSION") or "72"' in source, path.name


class TestDetectBnbRocmDllVer:
    """Unit tests for _detect_bnb_rocm_dll_ver()."""

    def test_returns_none_when_bnb_not_installed(self):
        """Returns None if bitsandbytes is not importable."""
        import importlib.util
        with patch.object(importlib.util, "find_spec", return_value = None):
            assert stack_mod._detect_bnb_rocm_dll_ver() is None

    def test_detects_rocm72_dll(self, tmp_path):
        """Returns '72' when libbitsandbytes_rocm72.dll is present."""
        (tmp_path / "libbitsandbytes_rocm72.dll").write_text("")
        mock_spec = MagicMock()
        mock_spec.submodule_search_locations = [str(tmp_path)]
        import importlib.util

        with patch.object(importlib.util, "find_spec", return_value = mock_spec):
            assert stack_mod._detect_bnb_rocm_dll_ver() == "72"

    def test_detects_rocm713_dll(self, tmp_path):
        """Returns '713' when libbitsandbytes_rocm713.dll is present."""
        (tmp_path / "libbitsandbytes_rocm713.dll").write_text("")
        mock_spec = MagicMock()
        mock_spec.submodule_search_locations = [str(tmp_path)]
        import importlib.util

        with patch.object(importlib.util, "find_spec", return_value = mock_spec):
            assert stack_mod._detect_bnb_rocm_dll_ver() == "713"

    def test_returns_none_when_only_cuda_dlls(self, tmp_path):
        """Returns None when only CUDA DLLs are present (no ROCm DLL)."""
        (tmp_path / "libbitsandbytes_cuda121.dll").write_text("")
        mock_spec = MagicMock()
        mock_spec.submodule_search_locations = [str(tmp_path)]
        import importlib.util

        with patch.object(importlib.util, "find_spec", return_value = mock_spec):
            assert stack_mod._detect_bnb_rocm_dll_ver() is None

    def test_picks_highest_suffix_when_multiple_dlls(self, tmp_path):
        """Returns the highest numeric suffix when multiple ROCm DLL variants exist.

        Filesystem glob order is not guaranteed, so the function must not stop
        at the first match — it must always return the highest one.
        """
        (tmp_path / "libbitsandbytes_rocm72.dll").write_text("")
        (tmp_path / "libbitsandbytes_rocm713.dll").write_text("")
        mock_spec = MagicMock()
        mock_spec.submodule_search_locations = [str(tmp_path)]
        import importlib.util

        with patch.object(importlib.util, "find_spec", return_value = mock_spec):
            assert stack_mod._detect_bnb_rocm_dll_ver() == "713"


# TEST: install_python_stack.py -- UNSLOTH_ROCM_TORCH_INSTALLED early-return path


class TestRocmTorchInstalledEnvVar:
    """Verify UNSLOTH_ROCM_TORCH_INSTALLED=1 skips main install but still installs BNB."""

    @staticmethod
    def _ok_torch_probe(*a, **kw):
        # subprocess.run probe returns 0 when torch imports as ROCm
        rv = MagicMock()
        rv.returncode = 0
        return rv

    @patch.object(stack_mod, "_install_bnb_windows_rocm")
    @patch.object(stack_mod, "pip_install")
    def test_env_var_skips_main_pip_install(self, mock_pip, mock_bnb):
        """UNSLOTH_ROCM_TORCH_INSTALLED=1 should not trigger torch pip_install."""
        with (
            patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}),
            patch.object(stack_mod.subprocess, "run", side_effect = self._ok_torch_probe),
        ):
            stack_mod._ensure_rocm_torch()
        mock_pip.assert_not_called()

    @patch.object(stack_mod, "_install_bnb_windows_rocm")
    @patch.object(stack_mod, "pip_install")
    def test_env_var_calls_bnb_install(self, mock_pip, mock_bnb):
        """UNSLOTH_ROCM_TORCH_INSTALLED=1 should still call _install_bnb_windows_rocm."""
        with (
            patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}),
            patch.object(stack_mod.subprocess, "run", side_effect = self._ok_torch_probe),
        ):
            stack_mod._ensure_rocm_torch()
        mock_bnb.assert_called_once()

    @patch.object(stack_mod, "_install_bnb_windows_rocm")
    @patch.object(stack_mod, "pip_install")
    def test_env_var_sets_rocm_windows_flag(self, mock_pip, mock_bnb):
        """UNSLOTH_ROCM_TORCH_INSTALLED=1 should set _rocm_windows_torch_installed."""
        stack_mod._rocm_windows_torch_installed = False
        with (
            patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}),
            patch.object(stack_mod.subprocess, "run", side_effect = self._ok_torch_probe),
        ):
            stack_mod._ensure_rocm_torch()
        assert stack_mod._rocm_windows_torch_installed is True

    @patch.object(stack_mod, "_install_bnb_windows_rocm")
    @patch.object(stack_mod, "pip_install")
    def test_env_var_falls_through_when_torch_missing(self, mock_pip, mock_bnb):
        """If the venv was wiped between runs, the stale env-var must not suppress reinstall."""
        stack_mod._rocm_windows_torch_installed = False

        def _bad_probe(*a, **kw):
            rv = MagicMock()
            rv.returncode = 1
            return rv

        with (
            patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}),
            patch.object(stack_mod.subprocess, "run", side_effect = _bad_probe),
            patch.object(stack_mod, "IS_WINDOWS", False),
            patch.object(stack_mod, "IS_MACOS", True),
        ):
            stack_mod._ensure_rocm_torch()
        # macOS branch is the next exit -- but the point is the early-return did NOT fire.
        mock_bnb.assert_not_called()


# TEST: worker.py -- Windows ROCm patches (source-level checks)


class TestWorkerWindowsRocmPatches:
    """Verify worker.py contains the required Windows ROCm runtime patches."""

    def test_grouped_mm_dispatch_patch_present(self):
        """worker.py must register a _grouped_mm CUDA dispatch override."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        assert '_gm_lib.impl("_grouped_mm"' in source

    def test_grouped_mm_patch_targets_cuda_dispatch_key(self):
        """The dispatch override must target the CUDA key (not CompositeImplicitAutograd)."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        assert '"_grouped_mm", _grouped_mm_safe_impl, "CUDA"' in source

    def test_grouped_mm_lib_kept_alive(self):
        """The Library object must be stored to prevent GC clearing the registration."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        assert "_WINDOWS_ROCM_GROUPED_MM_LIB" in source

    def test_grouped_mm_handles_offs_grouped_case(self):
        """_grouped_mm fallback must handle the grouped (offs!=None) variant."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        assert "offs_list" in source
        assert "offs.tolist()" in source

    def test_worker_calls_shared_torchao_stub(self):
        """worker.py must invoke the shared torchao stub entrypoint."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        assert "install_torchao_windows_rocm_stub()" in source

    def test_export_worker_calls_shared_torchao_stub(self):
        """export/worker.py must invoke the same shared torchao stub entrypoint."""
        source = _EXPORT_WORKER_PATH.read_text(encoding = "utf-8")
        assert "install_torchao_windows_rocm_stub()" in source

    def test_torchao_stub_uses_stub_type_meta(self):
        """Torchao stub must use _StubTypeMeta so isinstance() returns False not TypeError."""
        source = _TORCHAO_STUB_PATH.read_text(encoding = "utf-8")
        assert "_StubTypeMeta" in source

    def test_stub_type_meta_has_instancecheck(self):
        """_StubTypeMeta must define __instancecheck__ returning False."""
        source = _TORCHAO_STUB_PATH.read_text(encoding = "utf-8")
        assert "__instancecheck__" in source

    def test_stub_subpackage_finder_registered(self):
        """_StubSubpackageFinder must be appended to sys.meta_path."""
        source = _TORCHAO_STUB_PATH.read_text(encoding = "utf-8")
        assert "sys.meta_path.append(_StubSubpackageFinder())" in source

    def test_torchao_key_submodules_pre_stubbed(self):
        """Key torchao submodules (dtypes, quantization) must be pre-stubbed."""
        source = _TORCHAO_STUB_PATH.read_text(encoding = "utf-8")
        assert "torchao.dtypes" in source
        assert "torchao.quantization" in source

    def test_torchdynamo_disabled_on_windows_rocm(self):
        """worker.py should disable dynamo on Windows ROCm as belt-and-suspenders."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        assert "TORCHDYNAMO_DISABLE" in source

    def test_bnb_rocm_version_set_on_windows_rocm(self):
        """worker.py must set BNB_ROCM_VERSION in the Windows ROCm section.

        BNB auto-detects HIP version from torch.version.hip, which can mismatch
        the DLL suffix in the AMD prerelease wheel.  The worker must detect the
        actual DLL suffix and override BNB's auto-detection before ML imports.
        """
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        # Env var must be set
        assert "BNB_ROCM_VERSION" in source
        # Detection helper must be used
        assert "_detect_bnb_rocm_dll_ver" in source or "libbitsandbytes_rocm" in source
        # "72" must appear as the safe fallback
        assert '"72"' in source or "'72'" in source

    def test_bnb_rocm_version_set_before_ml_imports(self):
        """BNB_ROCM_VERSION must appear in section 1f, before section 2 ML imports."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        idx_bnb = source.find("BNB_ROCM_VERSION")
        # Use the specific section-2 marker that appears in the worker process
        # entry-point function (not the trainer helper which has its own "# ── 2.").
        idx_sec2 = source.find("# ── 2. Now import ML libraries")
        assert idx_bnb != -1, "BNB_ROCM_VERSION not found in worker.py"
        assert idx_sec2 != -1, "'# ── 2. Now import ML libraries' marker not found in worker.py"
        assert idx_bnb < idx_sec2, (
            "BNB_ROCM_VERSION must be set before section 2 ML imports "
            f"(found at {idx_bnb}, section 2 at {idx_sec2})"
        )

    def test_grouped_mm_patch_guarded_by_windows_and_hip_check(self):
        """_grouped_mm patch must only apply on Windows + HIP torch."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        # Must check sys.platform == "win32"
        assert 'sys.platform == "win32"' in source
        # Must gate on HIP version — code uses getattr chain: "version" and "hip"
        assert '"version"' in source and '"hip"' in source

    def test_hip_ver_at_least_helper_defined(self):
        """_hip_ver_at_least helper must be defined inside the Windows ROCm block."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        assert "def _hip_ver_at_least(major: int, minor: int)" in source

    def test_grouped_mm_patch_gated_on_hip_lt_713(self):
        """_grouped_mm patch must be skipped on HIP >= 7.13 (AMD fixed the bug in ROCm 7.13)."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        # The guard must call _hip_ver_at_least with exactly (7, 13)
        assert "_hip_ver_at_least(7, 13)" in source
        # The patch must be inside the `if not` branch (negated guard)
        assert "if not _hip_ver_at_least(7, 13):" in source

    def test_grouped_mm_hip_713_skip_message_present(self):
        """worker.py must log a message when skipping the patch on HIP >= 7.13."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        assert "HIP >= 7.13" in source
        assert "7.13" in source

    def test_grouped_mm_patch_else_branch_present(self):
        """An else branch must follow the _hip_ver_at_least gate (skip path for 7.13+)."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        # There must be an else: after the if not _hip_ver_at_least(7, 13): block
        gate_idx = source.find("if not _hip_ver_at_least(7, 13):")
        assert gate_idx != -1, "Version gate not found in worker.py"
        # The else: branch must appear after the gate
        else_idx = source.find("else:", gate_idx)
        assert else_idx != -1, "else: branch after _hip_ver_at_least gate not found"

    def test_hip_ver_at_least_handles_amd_version_format(self):
        """_hip_ver_at_least must split on '.' and compare only major.minor (handles '7.13.99004')."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        # Must split the version string and take the first two parts
        assert 'split(".")[:2]' in source or ".split('.')[:2]" in source


# TEST: install_python_stack.py -- _ROCM_TORCH_PKG_SPECS mapping


class TestRocmTorchPkgSpecs:
    """Verify per-tag torch version specs are correct."""

    def test_rocm72_has_torch_211(self):
        """rocm7.2 should specify torch 2.11.x."""
        specs = stack_mod._ROCM_TORCH_PKG_SPECS.get("rocm7.2")
        assert specs is not None
        torch_spec = specs[0]
        assert "2.11" in torch_spec

    def test_default_caps_below_211(self):
        """Default spec (rocm7.1 and earlier) should cap below 2.11."""
        specs = stack_mod._ROCM_TORCH_PKG_SPECS.get("_default")
        assert specs is not None
        torch_spec = specs[0]
        assert "<2.11" in torch_spec

    def test_specs_have_torch_vision_audio(self):
        """Each entry should be a 3-tuple: torch, torchvision, torchaudio."""
        for tag, specs in stack_mod._ROCM_TORCH_PKG_SPECS.items():
            assert len(specs) == 3, f"{tag}: expected (torch, torchvision, torchaudio)"
            assert "torch" in specs[0]
            assert "torchvision" in specs[1]
            assert "torchaudio" in specs[2]

    def test_gfx_to_amd_index_covers_rdna4(self):
        """_GFX_TO_AMD_INDEX_ARCH must cover gfx1200 and gfx1201 (RDNA 4)."""
        mapping = stack_mod._GFX_TO_AMD_INDEX_ARCH
        assert mapping.get("gfx1200") == "gfx120X-all"
        assert mapping.get("gfx1201") == "gfx120X-all"

    def test_gfx_to_amd_index_covers_strix_halo(self):
        """_GFX_TO_AMD_INDEX_ARCH must cover gfx1151 and gfx1150 (RDNA 3.5)."""
        mapping = stack_mod._GFX_TO_AMD_INDEX_ARCH
        assert mapping.get("gfx1151") == "gfx1151"
        assert mapping.get("gfx1150") == "gfx1150"

    def test_gfx_to_amd_index_covers_rdna3(self):
        """_GFX_TO_AMD_INDEX_ARCH must cover gfx1100-gfx1103 (RDNA 3)."""
        mapping = stack_mod._GFX_TO_AMD_INDEX_ARCH
        for arch in ("gfx1100", "gfx1101", "gfx1102", "gfx1103"):
            assert mapping.get(arch) == "gfx110X-all", f"{arch} missing from mapping"


# TEST: setup.ps1 / install.ps1 -- Strix Halo gfx arch detection

_SETUP_PS1_PATH = PACKAGE_ROOT / "studio" / "setup.ps1"
_INSTALL_PS1_PATH = PACKAGE_ROOT / "install.ps1"


class TestStrixHaloGfxArchDetection:
    """Verify that setup.ps1 and install.ps1 have robust gfx arch detection
    for Strix Halo / iGPU users who only have the HIP runtime (no hipinfo)."""

    def test_amd_smi_static_asic_attempted_in_setup(self):
        """setup.ps1 must try 'amd-smi static --asic' when list output lacks gfx arch."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "static --asic" in source

    def test_amd_smi_static_asic_attempted_in_install(self):
        """install.ps1 must try 'amd-smi static --asic' when list output lacks gfx arch."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "static --asic" in source

    def test_env_var_override_in_setup(self):
        """setup.ps1 must honour UNSLOTH_ROCM_GFX_ARCH as a manual arch override."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "UNSLOTH_ROCM_GFX_ARCH" in source

    def test_env_var_override_in_install(self):
        """install.ps1 must honour UNSLOTH_ROCM_GFX_ARCH as a manual arch override."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "UNSLOTH_ROCM_GFX_ARCH" in source

    def test_name_arch_table_covers_strix_halo_in_setup(self):
        """setup.ps1 name→arch table must map 890M / Strix Halo to gfx1151."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "gfx1151" in source
        assert "890M" in source or "Strix Halo" in source

    def test_name_arch_table_covers_strix_halo_in_install(self):
        """install.ps1 name→arch table must map 890M / Strix Halo to gfx1151."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "gfx1151" in source
        assert "890M" in source or "Strix Halo" in source

    def test_name_arch_table_covers_strix_point_in_setup(self):
        """setup.ps1 name→arch table must map 880M / Strix Point to gfx1150."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "gfx1150" in source
        assert "880M" in source or "Strix Point" in source

    def test_name_arch_table_covers_strix_point_in_install(self):
        """install.ps1 name→arch table must map 880M / Strix Point to gfx1150."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "gfx1150" in source
        assert "880M" in source or "Strix Point" in source

    def test_name_arch_table_covers_rdna3_phoenix_in_setup(self):
        """setup.ps1 name→arch table must map 780M / Phoenix to gfx1103."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "gfx1103" in source
        assert "780M" in source or "Phoenix" in source

    def test_wmi_does_not_set_hasrocm_in_setup(self):
        """WMI block in setup.ps1 must NOT set $HasROCm = $true (no runtime confirmation)."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        # Find the WMI block and confirm HasROCm is not set inside it
        wmi_idx = source.find("Win32_VideoController")
        assert wmi_idx != -1, "WMI block not found in setup.ps1"
        # The nearest HasROCm = $true must not appear between the WMI block
        # and the closing brace of that if-block.  We check by confirming
        # $HasROCm = $true does NOT appear within 300 chars of the WMI call.
        wmi_context = source[wmi_idx : wmi_idx + 300]
        assert "$HasROCm = $true" not in wmi_context

    def test_gfx_arch_regex_parses_from_amd_smi_output(self):
        """Both files must use the gfx\\d+[a-z]? regex to parse arch from amd-smi output."""
        for path in (_SETUP_PS1_PATH, _INSTALL_PS1_PATH):
            source = path.read_text(encoding = "utf-8")
            # The regex pattern used to match gfx arches
            assert (
                "gfx\\d+" in source or r"gfx\d+" in source
            ), f"gfx arch regex not found in {path.name}"


# TEST: HIP SDK tool path resolution via HIP_PATH / ROCM_PATH env vars


class TestHipSdkEnvPathResolution:
    """Verify that both install scripts resolve hipinfo/hipconfig via HIP_PATH
    and ROCM_PATH when the tools are not on $PATH, and emit explicit warnings."""

    # ── hipinfo resolution ────────────────────────────────────────────────────

    def test_setup_checks_hip_path_for_hipinfo(self):
        """setup.ps1 must reference HIP_PATH when resolving hipinfo."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "HIP_PATH" in source
        assert "hipinfo" in source

    def test_install_checks_hip_path_for_hipinfo(self):
        """install.ps1 must reference HIP_PATH when resolving hipinfo."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "HIP_PATH" in source
        assert "hipinfo" in source

    def test_setup_checks_rocm_path_as_hipinfo_fallback(self):
        """setup.ps1 must also check ROCM_PATH as a secondary hipinfo fallback."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "ROCM_PATH" in source
        # Confirm the fallback pattern: HIP_PATH ?? ROCM_PATH (or equivalent elseif)
        assert "ROCM_PATH" in source and "HIP_PATH" in source

    def test_install_checks_rocm_path_as_hipinfo_fallback(self):
        """install.ps1 must also check ROCM_PATH as a secondary hipinfo fallback."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "ROCM_PATH" in source
        assert "ROCM_PATH" in source and "HIP_PATH" in source

    def test_setup_resolves_hipinfo_via_bin_subdir(self):
        """setup.ps1 must join the env var root with 'bin\\hipinfo.exe'."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert r"bin\hipinfo.exe" in source

    def test_install_resolves_hipinfo_via_bin_subdir(self):
        """install.ps1 must join the env var root with 'bin\\hipinfo.exe'."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert r"bin\hipinfo.exe" in source

    # ── hipinfo not-on-PATH warning ───────────────────────────────────────────

    def test_setup_warns_when_hipinfo_not_on_path(self):
        """setup.ps1 must warn when hipinfo is found via env var but not on PATH."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "hipinfo not on PATH" in source

    def test_install_warns_when_hipinfo_not_on_path(self):
        """install.ps1 must warn when hipinfo is found via env var but not on PATH."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "hipinfo not on PATH" in source

    # ── warn when HIP_PATH set but exe missing ────────────────────────────────

    def test_setup_warns_when_hip_path_set_but_exe_missing(self):
        """setup.ps1 must warn when HIP_PATH is set but hipinfo.exe is not present."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        # The warning must mention that the SDK install may be incomplete
        assert "incomplete" in source or "not found at" in source

    def test_install_warns_when_hip_path_set_but_exe_missing(self):
        """install.ps1 must warn when HIP_PATH is set but hipinfo.exe is not present."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "incomplete" in source or "not found at" in source

    # ── hipinfo runtime error warning ─────────────────────────────────────────

    def test_setup_warns_on_hipinfo_nonzero_exit(self):
        """setup.ps1 must warn when hipinfo runs but returns a non-zero exit code."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "HIP runtime error" in source or "runtime error" in source.lower()

    def test_install_warns_on_hipinfo_nonzero_exit(self):
        """install.ps1 must warn when hipinfo runs but returns a non-zero exit code."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "HIP runtime error" in source or "runtime error" in source.lower()

    # ── hipconfig resolution ──────────────────────────────────────────────────

    def test_setup_resolves_hipconfig_via_bin_subdir(self):
        """setup.ps1 must also fall back to HIP_PATH/bin/hipconfig.exe for version detection."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert r"bin\hipconfig.exe" in source

    def test_install_resolves_hipconfig_via_bin_subdir(self):
        """install.ps1 must also fall back to HIP_PATH/bin/hipconfig.exe for version detection."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert r"bin\hipconfig.exe" in source

    def test_setup_warns_when_hipconfig_not_on_path(self):
        """setup.ps1 must warn when hipconfig is found via env var but not on PATH."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "hipconfig not on PATH" in source

    def test_install_warns_when_hipconfig_not_on_path(self):
        """install.ps1 must warn when hipconfig is found via env var but not on PATH."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "hipconfig not on PATH" in source

    # ── PATH fix hint ─────────────────────────────────────────────────────────

    def test_setup_provides_path_fix_hint(self):
        """setup.ps1 must tell the user how to add the HIP bin dir to PATH."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        # Should mention adding to PATH or SetEnvironmentVariable
        assert "PATH" in source and ("SetEnvironmentVariable" in source or "Add" in source)

    def test_install_provides_path_fix_hint(self):
        """install.ps1 must tell the user how to add the HIP bin dir to PATH."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "PATH" in source and ("SetEnvironmentVariable" in source or "Add" in source)


# TEST: HIP SDK detected substep -- path + hipconfig version shown in terminal


class TestHipSdkDetectedSubstep:
    """Verify that both scripts print HIP SDK path and full hipconfig version
    as substeps under the gpu step when AMD ROCm is successfully detected."""

    def test_setup_prints_hip_sdk_path_substep(self):
        """setup.ps1 must print an 'HIP SDK:' substep showing the resolved path."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "HIP SDK:" in source

    def test_install_prints_hip_sdk_path_substep(self):
        """install.ps1 must print an 'HIP SDK:' substep showing the resolved path."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "HIP SDK:" in source

    def test_setup_shows_hipconfig_full_version(self):
        """setup.ps1 must capture and display the full hipconfig version string."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "ROCmVersionFull" in source or "hipconfig:" in source

    def test_install_shows_hipconfig_full_version(self):
        """install.ps1 must capture and display the full hipconfig version string."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "ROCmVersionFull" in source or "hipconfig:" in source

    def test_setup_captures_full_version_not_just_major_minor(self):
        """setup.ps1 must store the raw hipconfig output line, not just major.minor."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "ROCmVersionFull" in source

    def test_install_captures_full_version_not_just_major_minor(self):
        """install.ps1 must store the raw hipconfig output line, not just major.minor."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "ROCmVersionFull" in source

    def test_setup_uses_hip_path_or_rocm_path_for_sdk_display(self):
        """setup.ps1 HIP SDK path substep must check HIP_PATH then ROCM_PATH."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "HIP_PATH" in source and "ROCM_PATH" in source

    def test_install_uses_hip_path_or_rocm_path_for_sdk_display(self):
        """install.ps1 HIP SDK path substep must check HIP_PATH then ROCM_PATH."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "HIP_PATH" in source and "ROCM_PATH" in source

    def test_setup_rocm_step_uses_full_version(self):
        """setup.ps1 'rocm' step label must prefer the full version string."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "ROCmVersionFull" in source and "rocm" in source


# TEST: install.sh -- Strix Halo rocm7.1 → rocm7.2 override

_INSTALL_SH_PATH = PACKAGE_ROOT / "install.sh"
_SETUP_SH_PATH = PACKAGE_ROOT / "studio" / "setup.sh"


class TestStrixRocm71Override:
    """Verify install.sh skips Radeon repo and routes to AMD arch-specific index
    for gfx1151/gfx1150 when ROCm 7.1 would otherwise be selected (known _grouped_mm segfault).
    AMD's repo.amd.com/rocm/whl/gfx1151/ serves torch 2.11+rocm7.13 which has the real fix."""

    def test_strix_gfx_detection_in_install_sh(self):
        """install.sh must detect gfx1151 and gfx1150 for the override."""
        source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
        assert "gfx1151" in source and "gfx1150" in source

    def test_rocm71_override_to_amd_arch_index_in_install_sh(self):
        """install.sh must override TORCH_INDEX_URL to AMD arch-specific index for Strix."""
        source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
        # The override must route to AMD's arch-specific index (repo.amd.com/rocm/whl)
        assert "repo.amd.com/rocm/whl" in source
        assert "_strix_gfx" in source
        # The URL must incorporate the detected gfx arch so gfx1151 → .../gfx1151/
        strix_idx = source.find("_amd_strix_base")
        assert strix_idx != -1
        ctx = source[strix_idx : strix_idx + 500]
        assert "_strix_gfx" in ctx

    def test_radeon_repo_bypassed_for_strix_in_install_sh(self):
        """install.sh must set _amd_gpu_radeon=false when Strix + ROCm 7.1 detected."""
        source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
        assert "_amd_gpu_radeon=false" in source

    def test_strix_override_warns_with_moe_utils_reference(self):
        """install.sh must emit a [WARN] mentioning the moe_utils segfault."""
        source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
        assert "moe_utils" in source or "_grouped_mm" in source

    def test_strix_override_only_fires_on_rocm71(self):
        """install.sh must scope the Strix override to rocm7.1 only (not rocm7.2+)."""
        source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
        strix_idx = source.find("_strix_gfx")
        assert strix_idx != -1
        # Look back for the rocm7.1 pattern within 600 chars before _strix_gfx
        context_before = source[max(0, strix_idx - 2400) : strix_idx]
        assert "rocm7.1" in context_before

    def test_torch_constraint_updated_for_strix_amd_index(self):
        """install.sh must set TORCH_CONSTRAINT>=2.11 when routing Strix to AMD index."""
        source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
        assert "TORCH_CONSTRAINT" in source and "2.11" in source

    def test_amd_rocm_mirror_env_var_respected(self):
        """install.sh must honour UNSLOTH_AMD_ROCM_MIRROR for air-gapped installs."""
        source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
        assert "UNSLOTH_AMD_ROCM_MIRROR" in source

    def test_tauri_family_recognises_amd_arch_url(self):
        """_tauri_torch_index_family must return a rocm* family for AMD arch-specific URLs."""
        source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
        # The function must have a case branch for repo.amd.com/rocm/whl/gfx* URLs
        assert "rocm/whl/gfx" in source


# TEST: setup.sh -- gcc-install-dir fix for Ubuntu 24.04 + ROCm 7.x clang-20


class TestSetupShGccInstallDir:
    """Verify setup.sh applies the --gcc-install-dir flag when building llama.cpp
    with HIP on Ubuntu 24.04+ to work around ROCm 7.x clang-20 header path bug."""

    def test_gcc_install_dir_search_loop_present(self):
        """setup.sh must iterate gcc versions 14→11 to find one with C++ headers."""
        source = _SETUP_SH_PATH.read_text(encoding = "utf-8")
        assert "_GCC_INSTALL_DIR" in source
        assert "/usr/lib/gcc/x86_64-linux-gnu" in source

    def test_gcc_install_dir_checks_include_dir(self):
        """setup.sh must check that the gcc dir has an 'include' subdirectory."""
        source = _SETUP_SH_PATH.read_text(encoding = "utf-8")
        assert "include" in source and "_GCC_INSTALL_DIR" in source

    def test_gcc_install_dir_appended_to_cmake_hip_flags(self):
        """setup.sh must pass --gcc-install-dir via CMAKE_HIP_FLAGS."""
        source = _SETUP_SH_PATH.read_text(encoding = "utf-8")
        assert "CMAKE_HIP_FLAGS" in source
        assert "gcc-install-dir" in source

    def test_gcc_install_dir_only_applied_in_hip_build_block(self):
        """The --gcc-install-dir fix must only apply in the HIP/ROCm build branch."""
        source = _SETUP_SH_PATH.read_text(encoding = "utf-8")
        # GGML_HIP=ON must appear before gcc-install-dir in the source
        hip_idx = source.find("GGML_HIP=ON")
        gcc_idx = source.find("gcc-install-dir")
        assert hip_idx != -1 and gcc_idx != -1
        assert hip_idx < gcc_idx

    def test_gcc_install_dir_logs_substep(self):
        """setup.sh must print a substep when the gcc install dir is resolved."""
        source = _SETUP_SH_PATH.read_text(encoding = "utf-8")
        assert "gcc install dir" in source or "GCC_INSTALL_DIR" in source


# TEST: main.py -- BNB_ROCM_VERSION server startup + distributed stubs

_MAIN_PY_PATH = PACKAGE_ROOT / "studio" / "backend" / "main.py"
_HARDWARE_PY_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"


class TestServerStartupRocmFixes:
    """Verify main.py sets BNB_ROCM_VERSION before any bitsandbytes import and
    hardware.py injects torch._C._distributed_c10d stubs before torch.distributed."""

    # ── BNB_ROCM_VERSION in server process ────────────────────────────────────

    def test_main_py_sets_bnb_rocm_version(self):
        """main.py must set BNB_ROCM_VERSION in the server process before imports."""
        source = _MAIN_PY_PATH.read_text(encoding = "utf-8")
        assert "BNB_ROCM_VERSION" in source

    def test_main_py_bnb_detection_scoped_to_win32(self):
        """main.py BNB_ROCM_VERSION logic must be inside the win32 platform guard."""
        source = _MAIN_PY_PATH.read_text(encoding = "utf-8")
        win32_idx = source.find('sys.platform == "win32"')
        bnb_idx = source.find("BNB_ROCM_VERSION")
        assert win32_idx != -1 and bnb_idx != -1
        assert win32_idx < bnb_idx

    def test_main_py_bnb_dll_detection_uses_glob(self):
        """main.py must scan for libbitsandbytes_rocm*.dll to find the right version."""
        source = _MAIN_PY_PATH.read_text(encoding = "utf-8")
        assert "libbitsandbytes_rocm" in source

    def test_main_py_bnb_falls_back_to_72(self):
        """main.py must fall back to BNB_ROCM_VERSION='72' when no DLL is found."""
        source = _MAIN_PY_PATH.read_text(encoding = "utf-8")
        assert '"72"' in source or "'72'" in source

    def test_main_py_bnb_only_set_when_not_already_in_env(self):
        """main.py must not override an existing BNB_ROCM_VERSION env var."""
        source = _MAIN_PY_PATH.read_text(encoding = "utf-8")
        assert '"BNB_ROCM_VERSION" not in os.environ' in source

    # ── hipInfo.exe PATH prepend (bitsandbytes arch-probe fix) ────────────────
    # bitsandbytes' get_rocm_gpu_arch() runs `hipinfo.exe` via subprocess PATH
    # at import time. The AMD torch wheel ships hipInfo.exe in the venv
    # Scripts dir, which is on PATH only for activated venvs -- Studio and the
    # installer launch python directly, so without the prepend every bnb
    # import logs "Could not detect ROCm GPU architecture: [WinError 2]".

    def test_main_py_prepends_hipinfo_dir_to_path(self):
        """main.py must make hipInfo.exe resolvable before bnb imports."""
        source = _MAIN_PY_PATH.read_text(encoding = "utf-8")
        assert "hipInfo.exe" in source
        # The prepend must come before the BNB_ROCM_VERSION block (both run
        # pre-import; order documents that bnb sees the fixed PATH).
        assert source.find("hipInfo.exe") < source.find("BNB_ROCM_VERSION")

    def test_main_py_hipinfo_prepend_gated_on_file_presence(self):
        """Only AMD ROCm wheels ship hipInfo.exe; NVIDIA/CPU hosts must be
        untouched, so the prepend must check the file exists first."""
        source = _MAIN_PY_PATH.read_text(encoding = "utf-8")
        assert 'os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe"))' in source

    def test_worker_py_prepends_hipinfo_dir_to_path(self):
        """worker.py must mirror the prepend for standalone-spawned workers."""
        source = _WORKER_PATH.read_text(encoding = "utf-8")
        assert "hipInfo.exe" in source

    def test_install_stack_prepends_hipinfo_dir_to_path(self):
        """install_python_stack.py must prepend so the installer's child
        import checks inherit a PATH where bnb's probe succeeds."""
        source = _STACK_PATH.read_text(encoding = "utf-8")
        assert "hipInfo.exe" in source

    # ── torch._C._distributed_c10d stubs in hardware.py ──────────────────────

    def test_hardware_py_injects_distributed_c10d_stub(self):
        """hardware.py must inject torch._C._distributed_c10d into sys.modules."""
        source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8")
        assert "_distributed_c10d" in source

    def test_hardware_py_stub_injected_before_distributed_import(self):
        """The sys.modules stub must be injected BEFORE import torch.distributed."""
        source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8")
        c10d_idx = source.find("_distributed_c10d")
        dist_idx = source.find("import torch.distributed")
        assert c10d_idx != -1 and dist_idx != -1
        assert c10d_idx < dist_idx

    def test_hardware_py_stub_uses_types_moduletype(self):
        """hardware.py must create the stub with types.ModuleType."""
        source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8")
        assert "ModuleType" in source

    def test_hardware_py_stub_scoped_to_win32(self):
        """hardware.py distributed stub injection must be gated on win32."""
        source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8")
        assert 'platform == "win32"' in source or "win32" in source

    def test_hardware_py_stub_exposes_fake_process_group(self):
        """hardware.py stub must set FakeProcessGroup so torch.distributed doesn't raise AttributeError."""
        source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8")
        assert "FakeProcessGroup" in source

    def test_hardware_py_stub_exposes_process_group(self):
        """hardware.py stub must set ProcessGroup on the c10d stub."""
        source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8")
        assert "ProcessGroup" in source

    def test_hardware_py_stub_uses_setattr_for_symbols(self):
        """hardware.py must use setattr to populate stub symbols dynamically."""
        source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8")
        assert "setattr" in source

    def test_hardware_py_stub_all_c10d_siblings_covered(self):
        """hardware.py must stub all three torch._C._distributed_* submodules."""
        source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8")
        assert "_distributed_c10d" in source
        assert "_distributed_autograd" in source
        assert "_distributed_rpc" in source


# TEST: install.ps1 / setup.ps1 -- HipSdkInstalled flag (SDK found, device inaccessible)


class TestHipSdkInstalledButDeviceInaccessible:
    """Verify that when hipinfo is found but exits non-zero (device not ROCm-accessible),
    both scripts distinguish this from 'HIP SDK not found' and emit the correct message."""

    def test_install_ps1_has_hip_sdk_installed_flag(self):
        """install.ps1 must track HipSdkInstalled separately from HasROCm."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "HipSdkInstalled" in source

    def test_setup_ps1_has_hip_sdk_installed_flag(self):
        """setup.ps1 must track HipSdkInstalled separately from HasROCm."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "HipSdkInstalled" in source

    def test_install_ps1_sets_flag_when_hipinfo_binary_found(self):
        """install.ps1 must set HipSdkInstalled=true inside the 'if ($hipinfoExe)' block."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        # HipSdkInstalled must be assigned inside the hipinfoExe block
        hipinfo_block_idx = source.find("if ($hipinfoExe)")
        sdk_flag_idx = source.find("$HipSdkInstalled = $true", hipinfo_block_idx)
        assert hipinfo_block_idx != -1 and sdk_flag_idx != -1
        assert sdk_flag_idx > hipinfo_block_idx

    def test_setup_ps1_sets_flag_when_hipinfo_binary_found(self):
        """setup.ps1 must set HipSdkInstalled=true inside the 'if ($hipinfoExe)' block."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        hipinfo_block_idx = source.find("if ($hipinfoExe)")
        sdk_flag_idx = source.find("$HipSdkInstalled = $true", hipinfo_block_idx)
        assert hipinfo_block_idx != -1 and sdk_flag_idx != -1
        assert sdk_flag_idx > hipinfo_block_idx

    def test_install_ps1_version_capture_runs_when_sdk_installed(self):
        """install.ps1 must capture hipconfig version when HipSdkInstalled even if HasROCm is false."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "HasROCm -or $HipSdkInstalled" in source or "$HipSdkInstalled" in source

    def test_setup_ps1_version_capture_runs_when_sdk_installed(self):
        """setup.ps1 must capture hipconfig version when HipSdkInstalled even if HasROCm is false."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "HasROCm -or $HipSdkInstalled" in source or "$HipSdkInstalled" in source

    def test_install_ps1_distinct_message_for_sdk_found_but_device_inaccessible(self):
        """install.ps1 must show 'not ROCm-accessible' message (not 'HIP SDK not found') when SDK present."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "not ROCm-accessible" in source

    def test_setup_ps1_distinct_message_for_sdk_found_but_device_inaccessible(self):
        """setup.ps1 must show 'not ROCm-accessible' message (not 'HIP SDK not found') when SDK present."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "not ROCm-accessible" in source

    def test_install_ps1_driver_guidance_in_sdk_found_branch(self):
        """install.ps1 must tell user this is a driver issue, not an SDK issue."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "driver issue" in source

    def test_setup_ps1_driver_guidance_in_sdk_found_branch(self):
        """setup.ps1 must tell user this is a driver issue, not an SDK issue."""
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "driver issue" in source

    def test_install_ps1_cpu_hint_distinguishes_driver_vs_no_sdk(self):
        """install.ps1 CPU-only hint must say 'GPU not ROCm-accessible' not 'require the HIP SDK' when SDK found."""
        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
        assert "GPU not ROCm-accessible" in source


# TEST: --rocm-gfx forwarding -- setup.sh/setup.ps1 hand their resolved gfx arch
# to install_llama_prebuilt.py so the lemonade HIP prebuilt is selected even when
# the installer's own hipinfo/amd-smi probe cannot report it.

_SETUP_SH_PATH = PACKAGE_ROOT / "studio" / "setup.sh"


class TestNormalizeForwardedGfx:
    """A forwarded gfx string is reduced to a single clean gfx token."""

    def test_plain_token(self):
        assert _normalize_forwarded_gfx("gfx1151") == "gfx1151"

    def test_uppercase_normalized(self):
        assert _normalize_forwarded_gfx("GFX1151") == "gfx1151"

    def test_extracts_from_noise(self):
        assert _normalize_forwarded_gfx("gcnArchName: gfx942") == "gfx942"

    def test_malformed_is_ignored(self):
        assert _normalize_forwarded_gfx("not-a-gpu") is None

    def test_empty_and_none(self):
        assert _normalize_forwarded_gfx("") is None
        assert _normalize_forwarded_gfx(None) is None


class TestApplyHostOverrides:
    """Forwarded ROCm detection is folded into the host profile correctly."""

    def test_forwarded_gfx_fills_empty_probe(self):
        # amd-smi-only / name-inferred host: installer probe found no gfx.
        host = rocm_host(rocm_gfx_target = None)
        out = _apply_host_overrides(host, override_rocm_gfx = "gfx1151")
        assert out.has_rocm is True
        assert out.rocm_gfx_target == "gfx1151"

    def test_forwarded_gfx_implies_rocm(self):
        # A CPU-looking host with a forwarded gfx is an AMD host.
        out = _apply_host_overrides(cpu_host(), override_rocm_gfx = "gfx1200")
        assert out.has_rocm is True
        assert out.rocm_gfx_target == "gfx1200"

    def test_forwarded_gfx_is_authoritative(self):
        # setup already applied visible-device selection; its value wins.
        host = rocm_host(rocm_gfx_target = "gfx1100")
        out = _apply_host_overrides(host, override_rocm_gfx = "gfx1151")
        assert out.rocm_gfx_target == "gfx1151"

    def test_has_rocm_only_keeps_probe_gfx(self):
        out = _apply_host_overrides(cpu_host(), override_has_rocm = True)
        assert out.has_rocm is True
        assert out.rocm_gfx_target is None

    def test_malformed_forwarded_gfx_falls_back_to_has_rocm(self):
        out = _apply_host_overrides(cpu_host(), override_has_rocm = True, override_rocm_gfx = "junk")
        assert out.has_rocm is True
        assert out.rocm_gfx_target is None

    def test_no_overrides_leaves_host_unchanged(self):
        host = nvidia_host()
        assert _apply_host_overrides(host) is host


class TestRocmGfxForwarding:
    """setup.sh / setup.ps1 forward their resolved gfx; the installer accepts it."""

    def test_installer_exposes_rocm_gfx_arg(self):
        source = _PREBUILT_PATH.read_text(encoding = "utf-8")
        assert '"--rocm-gfx"' in source
        # Defaults to the env override so a standalone run still works.
        assert 'os.environ.get("UNSLOTH_ROCM_GFX_ARCH")' in source

    def test_setup_sh_forwards_rocm_gfx(self):
        source = _SETUP_SH_PATH.read_text(encoding = "utf-8")
        assert "--rocm-gfx" in source
        assert '"$_setup_gfx"' in source

    def test_setup_sh_forwards_has_rocm(self):
        # When AMD is detected but gfx resolution fails, setup.sh must still
        # forward --has-rocm so the installer knows ROCm is present.
        source = _SETUP_SH_PATH.read_text(encoding = "utf-8")
        assert "--has-rocm" in source
        assert "_setup_amd_detected" in source

    def test_setup_ps1_forwards_rocm_gfx(self):
        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
        assert "--rocm-gfx" in source
        assert "$script:ROCmGfxArch" in source


if __name__ == "__main__":
    pytest.main([__file__, "-v"])