TheRock/rocm-systems/shared/kpack/tests/test_wheel_splitter.py

"""
Tests for the wheel splitter.

Unit tests verify metadata parsing and generation without requiring a toolchain.
Integration tests use real fat binaries from test_assets to verify the full
split pipeline.
"""

import hashlib
import os
import shutil
import stat
import zipfile
from pathlib import Path

import pytest

from rocm_kpack.wheel_splitter import (
    FatBinaryInfo,
    InvalidWheelError,
    NoFatBinariesError,
    SplitResult,
    WheelIdentity,
    WheelSplitter,
    WheelSplitError,
    _arch_to_bundle_key,
    _bundle_key_to_dist_name,
    compute_dist_info_name,
    compute_wheel_filename,
    generate_device_metadata,
    generate_record,
    generate_wheel_file,
    parse_wheel_identity,
    zip_wheel,
)

# =============================================================================
# Fixtures
# =============================================================================


@pytest.fixture
def mock_wheel(tmp_path: Path):
    """Create a minimal mock wheel directory structure (no fat binaries)."""
    wheel_dir = tmp_path / "wheel"
    wheel_dir.mkdir()

    # Create torch/ overlay root with some files
    torch_dir = wheel_dir / "torch"
    torch_dir.mkdir()
    (torch_dir / "__init__.py").write_text("# torch\n")
    lib_dir = torch_dir / "lib"
    lib_dir.mkdir()
    (lib_dir / "libtorch_cpu.so").write_bytes(b"\x00" * 100)

    # Create dist-info
    dist_info = wheel_dir / "torch-2.10.0+rocm7.1.dist-info"
    dist_info.mkdir()
    (dist_info / "METADATA").write_text(
        "Metadata-Version: 2.4\n"
        "Name: torch\n"
        "Version: 2.10.0+rocm7.1\n"
        "Summary: Tensors and Dynamic neural networks\n"
    )
    (dist_info / "WHEEL").write_text(
        "Wheel-Version: 1.0\n"
        "Generator: setuptools 80.9.0\n"
        "Root-Is-Purelib: false\n"
        "Tag: cp313-cp313-manylinux_2_28_x86_64\n"
    )
    (dist_info / "RECORD").write_text("")
    (dist_info / "top_level.txt").write_text("torch\n")

    return wheel_dir


@pytest.fixture
def fat_binary_wheel(tmp_path: Path, test_assets_dir: Path):
    """Create a mock wheel directory with a real fat binary from test assets."""
    wheel_dir = tmp_path / "fat_wheel"
    wheel_dir.mkdir()

    # Create torch/ overlay root
    torch_dir = wheel_dir / "torch"
    torch_dir.mkdir()
    (torch_dir / "__init__.py").write_text("# torch\n")
    lib_dir = torch_dir / "lib"
    lib_dir.mkdir()

    # Copy a real fat binary as a simulated torch library
    source_binary = (
        test_assets_dir / "bundled_binaries/linux/cov5/libtest_kernel_single.so"
    )
    if not source_binary.exists():
        pytest.skip("Test assets not available")
    shutil.copy2(source_binary, lib_dir / "libfat_test.so")

    # Also add a non-fat binary
    (lib_dir / "libplain.so").write_bytes(b"\x00" * 100)

    # Create dist-info
    dist_info = wheel_dir / "torch-2.10.0+rocm7.1.dist-info"
    dist_info.mkdir()
    (dist_info / "METADATA").write_text(
        "Metadata-Version: 2.4\n"
        "Name: torch\n"
        "Version: 2.10.0+rocm7.1\n"
        "Summary: Tensors and Dynamic neural networks\n"
    )
    (dist_info / "WHEEL").write_text(
        "Wheel-Version: 1.0\n"
        "Generator: setuptools 80.9.0\n"
        "Root-Is-Purelib: false\n"
        "Tag: cp313-cp313-manylinux_2_28_x86_64\n"
    )
    (dist_info / "RECORD").write_text("")
    (dist_info / "top_level.txt").write_text("torch\n")

    return wheel_dir


# =============================================================================
# Unit tests — no toolchain needed
# =============================================================================


class TestParseWheelIdentity:
    def test_basic(self, mock_wheel: Path):
        identity = parse_wheel_identity(mock_wheel)
        assert identity.name == "torch"
        assert identity.version == "2.10.0+rocm7.1"
        assert identity.python_tag == "cp313"
        assert identity.abi_tag == "cp313"
        assert identity.platform_tag == "manylinux_2_28_x86_64"
        assert identity.dist_info_name == "torch-2.10.0+rocm7.1.dist-info"

    def test_no_dist_info(self, tmp_path: Path):
        empty_dir = tmp_path / "empty"
        empty_dir.mkdir()
        with pytest.raises(InvalidWheelError, match="Expected exactly one .dist-info"):
            parse_wheel_identity(empty_dir)

    def test_multiple_dist_info(self, tmp_path: Path):
        wheel_dir = tmp_path / "multi"
        wheel_dir.mkdir()
        (wheel_dir / "a-1.0.dist-info").mkdir()
        (wheel_dir / "b-2.0.dist-info").mkdir()
        with pytest.raises(InvalidWheelError, match="Expected exactly one .dist-info"):
            parse_wheel_identity(wheel_dir)

    def test_missing_metadata(self, tmp_path: Path):
        wheel_dir = tmp_path / "no_meta"
        wheel_dir.mkdir()
        dist_info = wheel_dir / "pkg-1.0.dist-info"
        dist_info.mkdir()
        (dist_info / "WHEEL").write_text("Tag: cp313-cp313-linux_x86_64\n")
        with pytest.raises(InvalidWheelError, match="No METADATA"):
            parse_wheel_identity(wheel_dir)

    def test_missing_wheel_file(self, tmp_path: Path):
        wheel_dir = tmp_path / "no_wheel"
        wheel_dir.mkdir()
        dist_info = wheel_dir / "pkg-1.0.dist-info"
        dist_info.mkdir()
        (dist_info / "METADATA").write_text("Name: pkg\nVersion: 1.0\n")
        with pytest.raises(InvalidWheelError, match="No WHEEL"):
            parse_wheel_identity(wheel_dir)


class TestComputeWheelFilename:
    def test_basic(self):
        identity = WheelIdentity(
            name="torch",
            version="2.10.0+rocm7.1",
            python_tag="cp313",
            abi_tag="cp313",
            platform_tag="manylinux_2_28_x86_64",
            dist_info_name="torch-2.10.0+rocm7.1.dist-info",
        )
        filename = compute_wheel_filename(
            "amd-torch-device-gfx942", "2.10.0+rocm7.1", identity
        )
        assert filename == (
            "amd_torch_device_gfx942-2.10.0+rocm7.1"
            "-cp313-cp313-manylinux_2_28_x86_64.whl"
        )

    def test_underscore_normalization(self):
        identity = WheelIdentity(
            name="pkg",
            version="1.0",
            python_tag="py3",
            abi_tag="none",
            platform_tag="any",
            dist_info_name="pkg-1.0.dist-info",
        )
        filename = compute_wheel_filename("my-package", "1.0", identity)
        assert filename == "my_package-1.0-py3-none-any.whl"


class TestComputeDistInfoName:
    def test_basic(self):
        assert (
            compute_dist_info_name("amd-torch-device-gfx942", "2.10.0+rocm7.1")
            == "amd_torch_device_gfx942-2.10.0+rocm7.1.dist-info"
        )


class TestGenerateDeviceMetadata:
    def test_target_level(self):
        identity = WheelIdentity(
            name="torch",
            version="2.10.0+rocm7.1",
            python_tag="cp313",
            abi_tag="cp313",
            platform_tag="manylinux_2_28_x86_64",
            dist_info_name="torch-2.10.0+rocm7.1.dist-info",
        )
        metadata = generate_device_metadata(identity, "gfx942", "amd-torch-device", [])
        assert "Name: amd-torch-device-gfx942" in metadata
        assert "Version: 2.10.0+rocm7.1" in metadata
        assert "Requires-Dist: torch == 2.10.0+rocm7.1" in metadata

    def test_family_level(self):
        identity = WheelIdentity(
            name="torch",
            version="2.10.0+rocm7.1",
            python_tag="cp313",
            abi_tag="cp313",
            platform_tag="manylinux_2_28_x86_64",
            dist_info_name="torch-2.10.0+rocm7.1.dist-info",
        )
        metadata = generate_device_metadata(identity, "gfx11", "amd-torch-device", [])
        assert "Name: amd-torch-device-gfx11" in metadata

    def test_sub_family_level(self):
        identity = WheelIdentity(
            name="torch",
            version="2.10.0+rocm7.1",
            python_tag="cp313",
            abi_tag="cp313",
            platform_tag="manylinux_2_28_x86_64",
            dist_info_name="torch-2.10.0+rocm7.1.dist-info",
        )
        metadata = generate_device_metadata(identity, "gfx12_0", "amd-torch-device", [])
        # gfx12_0 becomes gfx12-0 in dist name (underscore to hyphen)
        assert "Name: amd-torch-device-gfx12-0" in metadata

    def test_gfx11_sub_family_level(self):
        identity = WheelIdentity(
            name="torch",
            version="2.10.0+rocm7.1",
            python_tag="cp313",
            abi_tag="cp313",
            platform_tag="manylinux_2_28_x86_64",
            dist_info_name="torch-2.10.0+rocm7.1.dist-info",
        )
        metadata = generate_device_metadata(identity, "gfx110x", "amd-torch-device", [])
        assert "Name: amd-torch-device-gfx110x" in metadata
        assert "Summary: AMD device kernels (RDNA3)" in metadata

    def test_with_device_requires_dist(self):
        identity = WheelIdentity(
            name="torch",
            version="2.10.0+rocm7.1",
            python_tag="cp313",
            abi_tag="cp313",
            platform_tag="manylinux_2_28_x86_64",
            dist_info_name="torch-2.10.0+rocm7.1.dist-info",
        )
        metadata = generate_device_metadata(
            identity,
            "gfx942",
            "amd-torch-device",
            ["rocm-sdk-device-@GFXARCH@ == 7.1"],
        )
        assert "Requires-Dist: rocm-sdk-device-gfx942 == 7.1" in metadata

    def test_gfxarch_placeholder_expansion(self):
        identity = WheelIdentity(
            name="torch",
            version="1.0",
            python_tag="py3",
            abi_tag="none",
            platform_tag="any",
            dist_info_name="torch-1.0.dist-info",
        )
        metadata = generate_device_metadata(
            identity,
            "gfx1100",
            "amd-torch-device",
            ["some-dep-@GFXARCH@ >= 2.0"],
        )
        assert "Requires-Dist: some-dep-gfx1100 >= 2.0" in metadata

    def test_gfxarch_dep_dropped_for_family(self):
        # @GFXARCH@-templated deps name per-target packages (e.g.
        # rocm-sdk-device-<target>) which only exist for TARGET-level
        # bundles. Family/sub-family device wheels are co-installed with a
        # target wheel that already carries those deps, so the line must be
        # dropped here.
        identity = WheelIdentity(
            name="torch",
            version="2.10.0+rocm7.13",
            python_tag="cp313",
            abi_tag="cp313",
            platform_tag="manylinux_2_28_x86_64",
            dist_info_name="torch-2.10.0+rocm7.13.dist-info",
        )
        family_metadata = generate_device_metadata(
            identity,
            "gfx11",
            "amd-torch-device",
            ["rocm-sdk-device-@GFXARCH@ == 7.13", "static-dep == 1.0"],
        )
        assert "rocm-sdk-device" not in family_metadata
        assert "Requires-Dist: static-dep == 1.0" in family_metadata

        sub_family_metadata = generate_device_metadata(
            identity,
            "gfx12_0",
            "amd-torch-device",
            ["rocm-sdk-device-@GFXARCH@ == 7.13"],
        )
        assert "rocm-sdk-device" not in sub_family_metadata


class TestArchToBundleKey:
    def test_bare_arch(self):
        assert _arch_to_bundle_key("gfx942") == "gfx942"

    def test_xnack_plus(self):
        assert _arch_to_bundle_key("gfx942:xnack+") == "gfx942"

    def test_xnack_minus(self):
        assert _arch_to_bundle_key("gfx90a:xnack-") == "gfx90a"

    def test_no_xnack(self):
        assert _arch_to_bundle_key("gfx1100") == "gfx1100"

    def test_family_key(self):
        assert _arch_to_bundle_key("gfx11") == "gfx11"

    def test_sub_family_key(self):
        assert _arch_to_bundle_key("gfx12_0") == "gfx12_0"


class TestBundleKeyToDistName:
    def test_target_level(self):
        assert (
            _bundle_key_to_dist_name("amd-torch-device", "gfx942")
            == "amd-torch-device-gfx942"
        )

    def test_family_level(self):
        assert (
            _bundle_key_to_dist_name("amd-torch-device", "gfx11")
            == "amd-torch-device-gfx11"
        )

    def test_sub_family_level(self):
        assert (
            _bundle_key_to_dist_name("amd-torch-device", "gfx110x")
            == "amd-torch-device-gfx110x"
        )

    def test_structural_sub_family_level(self):
        # Underscore in bundle key becomes hyphen in dist name.
        assert (
            _bundle_key_to_dist_name("amd-torch-device", "gfx12_0")
            == "amd-torch-device-gfx12-0"
        )

    def test_various_targets(self):
        cases = [
            ("gfx900", "amd-torch-device-gfx900"),
            ("gfx90a", "amd-torch-device-gfx90a"),
            ("gfx1100", "amd-torch-device-gfx1100"),
            ("gfx1151", "amd-torch-device-gfx1151"),
            ("gfx9", "amd-torch-device-gfx9"),
            ("gfx9_4", "amd-torch-device-gfx9-4"),
        ]
        for bundle_key, expected in cases:
            assert (
                _bundle_key_to_dist_name("amd-torch-device", bundle_key) == expected
            ), f"Failed for {bundle_key}"


class TestRewriteHostMetadata:
    """Test _rewrite_host_metadata via WheelSplitter."""

    def _make_host_staging(self, tmp_path: Path) -> tuple[Path, WheelIdentity]:
        """Create a minimal host staging directory with METADATA."""
        staging = tmp_path / "host_staging"
        staging.mkdir()
        dist_info = staging / "torch-2.10.0+rocm7.1.dist-info"
        dist_info.mkdir()
        (dist_info / "METADATA").write_text(
            "Metadata-Version: 2.4\n"
            "Name: torch\n"
            "Version: 2.10.0+rocm7.1\n"
            "Summary: Tensors and Dynamic neural networks\n"
            "\n"
            "A description body that comes after the header section.\n"
        )
        identity = WheelIdentity(
            name="torch",
            version="2.10.0+rocm7.1",
            python_tag="cp313",
            abi_tag="cp313",
            platform_tag="manylinux_2_28_x86_64",
            dist_info_name="torch-2.10.0+rocm7.1.dist-info",
        )
        return staging, identity

    def _make_splitter(self) -> WheelSplitter:
        from rocm_kpack.binutils import Toolchain

        return WheelSplitter(
            device_package_prefix="amd-torch-device",
            overlay_root="torch/",
            toolchain=Toolchain(),
        )

    def test_rocm_bootstrap_dep(self, tmp_path: Path):
        staging, identity = self._make_host_staging(tmp_path)
        splitter = self._make_splitter()
        splitter._rewrite_host_metadata(staging, identity, {"gfx942"})

        metadata = (staging / identity.dist_info_name / "METADATA").read_text()
        assert "Requires-Dist: rocm-bootstrap" in metadata

    def test_extras_for_single_target(self, tmp_path: Path):
        staging, identity = self._make_host_staging(tmp_path)
        splitter = self._make_splitter()
        # gfx942 packaging chain: gfx942 -> gfx9_4 -> gfx9
        # Only gfx942 has a device wheel (is in bundle_keys)
        splitter._rewrite_host_metadata(staging, identity, {"gfx942"})

        metadata = (staging / identity.dist_info_name / "METADATA").read_text()
        assert "Provides-Extra: device-gfx942" in metadata
        assert (
            'Requires-Dist: amd-torch-device-gfx942 == 2.10.0+rocm7.1; extra == "device-gfx942"'
            in metadata
        )
        assert "Provides-Extra: device-all" in metadata

    def test_classic_has_no_variant_markers(self, tmp_path: Path):
        staging, identity = self._make_host_staging(tmp_path)
        splitter = self._make_splitter()
        splitter._rewrite_host_metadata(staging, identity, {"gfx942"})

        metadata = (staging / identity.dist_info_name / "METADATA").read_text()
        assert "variant_properties" not in metadata

    def test_variant_markers_when_enabled(self, tmp_path: Path):
        staging, identity = self._make_host_staging(tmp_path)
        splitter = self._make_splitter()
        splitter._rewrite_host_metadata(
            staging, identity, {"gfx942"}, include_variant_markers=True
        )

        metadata = (staging / identity.dist_info_name / "METADATA").read_text()
        assert (
            "Requires-Dist: amd-torch-device-gfx942 == 2.10.0+rocm7.1; "
            '"amd :: gfx_arch :: gfx942" in variant_properties'
        ) in metadata

    def test_chain_with_family_wheel(self, tmp_path: Path):
        staging, identity = self._make_host_staging(tmp_path)
        splitter = self._make_splitter()
        # gfx1100 chain: gfx1100 -> gfx110x -> gfx11
        # If gfx11 and gfx1100 both have device wheels:
        splitter._rewrite_host_metadata(staging, identity, {"gfx1100", "gfx11"})

        metadata = (staging / identity.dist_info_name / "METADATA").read_text()
        # Both should appear as deps under the gfx1100 extra
        assert (
            'Requires-Dist: amd-torch-device-gfx1100 == 2.10.0+rocm7.1; extra == "device-gfx1100"'
            in metadata
        )
        assert (
            'Requires-Dist: amd-torch-device-gfx11 == 2.10.0+rocm7.1; extra == "device-gfx1100"'
            in metadata
        )
        # "device-all" should include both
        assert (
            'Requires-Dist: amd-torch-device-gfx11 == 2.10.0+rocm7.1; extra == "device-all"'
            in metadata
        )
        assert (
            'Requires-Dist: amd-torch-device-gfx1100 == 2.10.0+rocm7.1; extra == "device-all"'
            in metadata
        )

    def test_chain_with_gfx110x_wheel(self, tmp_path: Path):
        staging, identity = self._make_host_staging(tmp_path)
        splitter = self._make_splitter()
        splitter._rewrite_host_metadata(staging, identity, {"gfx1100", "gfx110x"})

        metadata = (staging / identity.dist_info_name / "METADATA").read_text()
        assert (
            'Requires-Dist: amd-torch-device-gfx1100 == 2.10.0+rocm7.1; extra == "device-gfx1100"'
            in metadata
        )
        assert (
            'Requires-Dist: amd-torch-device-gfx110x == 2.10.0+rocm7.1; extra == "device-gfx1100"'
            in metadata
        )
        assert (
            'Requires-Dist: amd-torch-device-gfx110x == 2.10.0+rocm7.1; extra == "device-all"'
            in metadata
        )

    def test_family_only_not_target_extra(self, tmp_path: Path):
        """Family-level bundle keys don't get their own Provides-Extra."""
        staging, identity = self._make_host_staging(tmp_path)
        splitter = self._make_splitter()
        # gfx11 is a family, not a target - no Provides-Extra for it directly
        splitter._rewrite_host_metadata(staging, identity, {"gfx11"})

        metadata = (staging / identity.dist_info_name / "METADATA").read_text()
        # No target-level extras since gfx11 is a family
        assert "Provides-Extra: device-gfx11" not in metadata
        # But "device-all" should still include it
        assert "Provides-Extra: device-all" in metadata
        assert (
            'Requires-Dist: amd-torch-device-gfx11 == 2.10.0+rocm7.1; extra == "device-all"'
            in metadata
        )

    def test_preserves_existing_metadata(self, tmp_path: Path):
        staging, identity = self._make_host_staging(tmp_path)
        splitter = self._make_splitter()
        splitter._rewrite_host_metadata(staging, identity, {"gfx942"})

        metadata = (staging / identity.dist_info_name / "METADATA").read_text()
        # Original fields should still be there
        assert "Name: torch" in metadata
        assert "Version: 2.10.0+rocm7.1" in metadata
        assert "Summary: Tensors and Dynamic neural networks" in metadata

    def test_headers_before_body(self, tmp_path: Path):
        staging, identity = self._make_host_staging(tmp_path)
        splitter = self._make_splitter()
        splitter._rewrite_host_metadata(staging, identity, {"gfx942"})

        metadata = (staging / identity.dist_info_name / "METADATA").read_text()
        # Split into header and body at the first blank line
        header_body = metadata.split("\n\n", 1)
        assert len(header_body) == 2, "METADATA should have header and body sections"
        header_section, body_section = header_body
        # All injected headers must be in the header section, not the body
        assert "Requires-Dist: rocm-bootstrap" in header_section
        assert "Provides-Extra: device-gfx942" in header_section
        assert "Provides-Extra: device-all" in header_section
        # Body should still contain the description
        assert "description body" in body_section


class TestVariantWheel:
    """Tests for PEP 817 variant wheel generation."""

    def _make_host_staging(self, tmp_path: Path) -> tuple[Path, WheelIdentity]:
        staging = tmp_path / "host_staging"
        staging.mkdir()
        dist_info = staging / "torch-2.10.0+rocm7.1.dist-info"
        dist_info.mkdir()
        # Write a classic-rewritten METADATA (extras only, no variant markers)
        (dist_info / "METADATA").write_text(
            "Metadata-Version: 2.4\n"
            "Name: torch\n"
            "Version: 2.10.0+rocm7.1\n"
            "Requires-Dist: rocm-bootstrap\n"
            "Provides-Extra: device-gfx942\n"
            'Requires-Dist: amd-torch-device-gfx942 == 2.10.0+rocm7.1; extra == "device-gfx942"\n'
            "Provides-Extra: device-all\n"
            'Requires-Dist: amd-torch-device-gfx942 == 2.10.0+rocm7.1; extra == "device-all"\n'
            "\n"
            "Description body.\n"
        )
        (dist_info / "WHEEL").write_text(
            "Wheel-Version: 1.0\nTag: cp313-cp313-manylinux_2_28_x86_64\n"
        )
        (dist_info / "RECORD").write_text("")
        # Add a dummy file so the wheel isn't empty
        (staging / "torch").mkdir()
        (staging / "torch" / "__init__.py").write_text("")
        identity = WheelIdentity(
            name="torch",
            version="2.10.0+rocm7.1",
            python_tag="cp313",
            abi_tag="cp313",
            platform_tag="manylinux_2_28_x86_64",
            dist_info_name="torch-2.10.0+rocm7.1.dist-info",
        )
        return staging, identity

    def _make_splitter(self) -> WheelSplitter:
        from rocm_kpack.binutils import Toolchain

        return WheelSplitter(
            device_package_prefix="amd-torch-device",
            overlay_root="torch/",
            toolchain=Toolchain(),
            generate_variant_wheel=True,
        )

    def test_add_variant_markers(self, tmp_path: Path):
        staging, identity = self._make_host_staging(tmp_path)
        splitter = self._make_splitter()
        splitter._add_variant_markers_to_metadata(staging, identity, {"gfx942"})

        metadata = (staging / identity.dist_info_name / "METADATA").read_text()
        # Should have the extras line AND the variant marker line
        assert 'extra == "device-gfx942"' in metadata
        assert ('"amd :: gfx_arch :: gfx942" in variant_properties') in metadata
        # "device-all" extra should NOT get variant markers
        assert '"amd :: gfx_arch :: all"' not in metadata

    def test_variant_json(self, tmp_path: Path):
        import json

        staging, identity = self._make_host_staging(tmp_path)
        splitter = self._make_splitter()
        splitter._write_variant_json(staging, identity, {"gfx942", "gfx11"})

        variant_path = staging / identity.dist_info_name / "variant.json"
        assert variant_path.exists()
        data = json.loads(variant_path.read_text())
        assert "providers" in data
        assert "amd" in data["providers"]
        assert "variants" in data
        # gfx942 is a target, gfx11 is a family — only targets get variants
        assert "gfx942" in data["variants"]
        assert "gfx11" not in data["variants"]

    def test_create_variant_wheel_directory(self, tmp_path: Path):
        import json

        staging, identity = self._make_host_staging(tmp_path)
        splitter = self._make_splitter()
        output_dir = tmp_path / "output"
        output_dir.mkdir()

        variant_path = splitter._create_variant_wheel(
            staging, identity, {"gfx942"}, output_dir, "directory"
        )

        assert variant_path.exists()
        assert "variant" in variant_path.name

        # Check METADATA has variant markers
        metadata = (variant_path / identity.dist_info_name / "METADATA").read_text()
        assert "variant_properties" in metadata
        assert 'extra == "device-gfx942"' in metadata

        # Check variant.json exists
        variant_json = variant_path / identity.dist_info_name / "variant.json"
        assert variant_json.exists()
        data = json.loads(variant_json.read_text())
        assert "gfx942" in data["variants"]


class TestGenerateWheelFile:
    def test_basic(self):
        identity = WheelIdentity(
            name="torch",
            version="2.10.0+rocm7.1",
            python_tag="cp313",
            abi_tag="cp313",
            platform_tag="manylinux_2_28_x86_64",
            dist_info_name="torch-2.10.0+rocm7.1.dist-info",
        )
        content = generate_wheel_file(identity)
        assert "Wheel-Version: 1.0" in content
        assert "Tag: cp313-cp313-manylinux_2_28_x86_64" in content
        assert "Root-Is-Purelib: false" in content
        assert "Generator: rocm-kpack-split-python-wheels" in content


class TestGenerateRecord:
    def test_basic(self, tmp_path: Path):
        wheel_dir = tmp_path / "wheel"
        wheel_dir.mkdir()
        dist_info = wheel_dir / "pkg-1.0.dist-info"
        dist_info.mkdir()

        # Create a file with known content
        test_content = b"hello world"
        (wheel_dir / "test.py").write_bytes(test_content)
        expected_hash = hashlib.sha256(test_content).hexdigest()

        # Create empty RECORD
        (dist_info / "RECORD").write_text("")

        record = generate_record(wheel_dir, "pkg-1.0.dist-info")
        lines = record.strip().split("\n")

        # Should have 2 entries: test.py and RECORD itself
        assert len(lines) == 2

        # RECORD entry has no hash
        record_line = [l for l in lines if "RECORD" in l][0]
        assert record_line == "pkg-1.0.dist-info/RECORD,,"

        # test.py has hash
        test_line = [l for l in lines if "test.py" in l][0]
        assert f"sha256={expected_hash}" in test_line
        assert f",{len(test_content)}" in test_line


class TestWheelArchivePermissions:
    def test_extract_and_repack_preserves_exact_executable_mode(
        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ):
        from rocm_kpack.binutils import Toolchain

        input_wheel = tmp_path / "input.whl"
        executable_path = "torch/bin/torch_shm_manager"
        plain_path = "torch/__init__.py"
        extracted_root = input_wheel.parent / (input_wheel.stem + ".tmp_extract")
        chmod_calls: dict[str, int] = {}
        path_type = type(tmp_path)
        original_chmod = path_type.chmod

        def recording_chmod(self: Path, mode: int) -> None:
            chmod_calls[self.relative_to(extracted_root).as_posix()] = mode
            if os.name != "nt":
                original_chmod(self, mode)

        with zipfile.ZipFile(input_wheel, "w") as zf:
            executable_info = zipfile.ZipInfo(executable_path)
            executable_info.external_attr = (stat.S_IFREG | 0o751) << 16
            zf.writestr(executable_info, b"#!/bin/sh\n")

            plain_info = zipfile.ZipInfo(plain_path)
            plain_info.external_attr = (stat.S_IFREG | 0o640) << 16
            zf.writestr(plain_info, b"# torch\n")

        monkeypatch.setattr(path_type, "chmod", recording_chmod)

        splitter = WheelSplitter(
            device_package_prefix="amd-torch-device",
            overlay_root="torch/",
            toolchain=Toolchain(),
        )
        extracted_dir, is_temporary = splitter._resolve_input(input_wheel)
        assert is_temporary

        try:
            extracted_executable = extracted_dir / executable_path
            if os.name != "nt":
                assert stat.S_IMODE(extracted_executable.stat().st_mode) == 0o751
                assert (
                    stat.S_IMODE((extracted_dir / plain_path).stat().st_mode) == 0o640
                )

            output_wheel = tmp_path / "output.whl"
            zip_wheel(extracted_dir, output_wheel)
        finally:
            shutil.rmtree(extracted_dir)

        assert chmod_calls[executable_path] == 0o751
        assert chmod_calls[plain_path] == 0o640

        if os.name == "nt":
            return

        with zipfile.ZipFile(output_wheel, "r") as zf:
            output_executable_mode = stat.S_IMODE(
                zf.getinfo(executable_path).external_attr >> 16
            )
            output_plain_mode = stat.S_IMODE(zf.getinfo(plain_path).external_attr >> 16)

        assert output_executable_mode == 0o751
        assert output_plain_mode == 0o640


class TestKpackSearchPattern:
    """Test the search pattern computation logic."""

    def _compute(self, overlay_relative_path: str, group_name: str) -> str:
        """Helper to call the private method via a throwaway splitter."""
        from rocm_kpack.binutils import Toolchain

        splitter = WheelSplitter(
            device_package_prefix="test",
            overlay_root="torch/",
            toolchain=Toolchain(),
        )
        return splitter._compute_kpack_search_pattern(overlay_relative_path, group_name)

    def test_binary_in_subdir(self):
        # lib/libtorch_hip.so -> depth=1 -> ../.kpack/torch_@GFXARCH@.kpack
        pattern = self._compute("lib/libtorch_hip.so", "torch")
        assert pattern == "../.kpack/torch_@GFXARCH@.kpack"

    def test_binary_at_root(self):
        # libtorch_hip.so -> depth=0 -> .kpack/torch_@GFXARCH@.kpack
        pattern = self._compute("libtorch_hip.so", "torch")
        assert pattern == ".kpack/torch_@GFXARCH@.kpack"

    def test_binary_deeply_nested(self):
        # a/b/c/binary.so -> depth=3 -> ../../../.kpack/torch_@GFXARCH@.kpack
        pattern = self._compute("a/b/c/binary.so", "torch")
        assert pattern == "../../../.kpack/torch_@GFXARCH@.kpack"


# =============================================================================
# Integration tests — require toolchain and test assets
# =============================================================================


class TestWheelSplitterIntegration:
    """Integration tests that use real fat binaries from test_assets."""

    def test_split_directory_to_directory(
        self, fat_binary_wheel: Path, toolchain, tmp_path: Path
    ):
        """Test splitting an exploded wheel to exploded output directories."""
        output_dir = tmp_path / "output"

        splitter = WheelSplitter(
            device_package_prefix="amd-torch-device",
            overlay_root="torch/",
            toolchain=toolchain,
            verbose=True,
        )

        result = splitter.split(fat_binary_wheel, output_dir, output_format="directory")

        # Should have found architectures
        assert len(result.architectures_found) > 0
        assert result.fat_binaries_processed == 1

        # Host wheel directory should exist
        assert result.host_wheel_path.exists()
        assert result.host_wheel_path.is_dir()

        # Host should have the original files (minus device code)
        host_init = result.host_wheel_path / "torch" / "__init__.py"
        assert host_init.exists()

        # Host fat binary should still exist but be transformed
        host_fat = result.host_wheel_path / "torch" / "lib" / "libfat_test.so"
        assert host_fat.exists()

        # Host RECORD should exist and be non-empty
        host_record = (
            result.host_wheel_path / "torch-2.10.0+rocm7.1.dist-info" / "RECORD"
        )
        assert host_record.exists()
        assert host_record.stat().st_size > 0

        # Device wheels should exist for each bundle key found
        for bundle_key, device_path in result.device_wheel_paths.items():
            assert device_path.exists()
            assert device_path.is_dir()

            # Device wheel should have at least one kpack file
            kpack_dir = device_path / "torch" / ".kpack"
            kpack_files = list(kpack_dir.glob("*.kpack"))
            assert len(kpack_files) > 0
            for kf in kpack_files:
                assert kf.stat().st_size > 0

            # Device wheel should have dist-info with correct naming
            device_name = _bundle_key_to_dist_name("amd-torch-device", bundle_key)
            device_dist_info_name = (
                f"{device_name.replace('-', '_')}-2.10.0+rocm7.1.dist-info"
            )
            device_dist_info = device_path / device_dist_info_name
            assert device_dist_info.exists()

            # Check METADATA
            metadata = (device_dist_info / "METADATA").read_text()
            assert f"Name: {device_name}" in metadata
            assert "Requires-Dist: torch == 2.10.0+rocm7.1" in metadata

            # Check WHEEL
            wheel_content = (device_dist_info / "WHEEL").read_text()
            assert "Tag: cp313-cp313-manylinux_2_28_x86_64" in wheel_content

            # Check RECORD
            record = (device_dist_info / "RECORD").read_text()
            assert ".kpack/" in record

    def test_split_no_fat_binaries(self, mock_wheel: Path, toolchain, tmp_path: Path):
        """Test that splitting a wheel with no fat binaries raises an error."""
        output_dir = tmp_path / "output"

        splitter = WheelSplitter(
            device_package_prefix="amd-torch-device",
            overlay_root="torch/",
            toolchain=toolchain,
        )

        with pytest.raises(NoFatBinariesError):
            splitter.split(mock_wheel, output_dir)

    def test_split_invalid_overlay_root(
        self, mock_wheel: Path, toolchain, tmp_path: Path
    ):
        """Test that an invalid overlay root raises an error."""
        output_dir = tmp_path / "output"

        splitter = WheelSplitter(
            device_package_prefix="amd-torch-device",
            overlay_root="nonexistent/",
            toolchain=toolchain,
        )

        with pytest.raises(InvalidWheelError, match="Overlay root"):
            splitter.split(mock_wheel, output_dir)

    def test_host_record_integrity(
        self, fat_binary_wheel: Path, toolchain, tmp_path: Path
    ):
        """Verify that the host wheel's RECORD has valid hashes."""
        output_dir = tmp_path / "output"

        splitter = WheelSplitter(
            device_package_prefix="amd-torch-device",
            overlay_root="torch/",
            toolchain=toolchain,
        )

        result = splitter.split(fat_binary_wheel, output_dir, output_format="directory")

        host_dir = result.host_wheel_path
        dist_info_name = "torch-2.10.0+rocm7.1.dist-info"
        record_path = host_dir / dist_info_name / "RECORD"
        record_text = record_path.read_text()

        for line in record_text.strip().split("\n"):
            if not line:
                continue
            parts = line.split(",")
            rel_path = parts[0]

            if parts[1] == "":
                # RECORD itself — no hash
                assert rel_path.endswith("RECORD")
                continue

            # Verify hash
            hash_part = parts[1]
            size_part = parts[2]
            assert hash_part.startswith("sha256=")
            expected_hash = hash_part.removeprefix("sha256=")

            file_path = host_dir / rel_path
            assert file_path.exists(), f"RECORD references missing file: {rel_path}"
            actual_hash = hashlib.sha256(file_path.read_bytes()).hexdigest()
            assert actual_hash == expected_hash, f"Hash mismatch for {rel_path}"
            assert int(size_part) == file_path.stat().st_size

    def test_device_requires_dist(
        self, fat_binary_wheel: Path, toolchain, tmp_path: Path
    ):
        """Test that --device-requires-dist with @GFXARCH@ expands correctly."""
        output_dir = tmp_path / "output"

        splitter = WheelSplitter(
            device_package_prefix="amd-torch-device",
            overlay_root="torch/",
            toolchain=toolchain,
            device_requires_dist=["rocm-sdk-device-@GFXARCH@ == 7.1"],
        )

        result = splitter.split(fat_binary_wheel, output_dir, output_format="directory")

        for bundle_key, device_path in result.device_wheel_paths.items():
            device_name = _bundle_key_to_dist_name("amd-torch-device", bundle_key)
            dist_info_name = f"{device_name.replace('-', '_')}-2.10.0+rocm7.1.dist-info"
            metadata = (device_path / dist_info_name / "METADATA").read_text()
            assert f"Requires-Dist: rocm-sdk-device-{bundle_key} == 7.1" in metadata

    def test_split_parallel(self, fat_binary_wheel: Path, toolchain, tmp_path: Path):
        """Test that parallel splitting (jobs=2) produces correct results."""
        output_dir = tmp_path / "output"

        splitter = WheelSplitter(
            device_package_prefix="amd-torch-device",
            overlay_root="torch/",
            toolchain=toolchain,
            verbose=True,
            jobs=2,
        )

        result = splitter.split(fat_binary_wheel, output_dir, output_format="directory")

        # Same assertions as test_split_directory_to_directory
        assert len(result.architectures_found) > 0
        assert result.fat_binaries_processed == 1
        assert result.host_wheel_path.exists()
        assert result.host_wheel_path.is_dir()

        host_fat = result.host_wheel_path / "torch" / "lib" / "libfat_test.so"
        assert host_fat.exists()

        host_record = (
            result.host_wheel_path / "torch-2.10.0+rocm7.1.dist-info" / "RECORD"
        )
        assert host_record.exists()
        assert host_record.stat().st_size > 0

        for bundle_key, device_path in result.device_wheel_paths.items():
            assert device_path.exists()
            kpack_dir = device_path / "torch" / ".kpack"
            kpack_files = list(kpack_dir.glob("*.kpack"))
            assert len(kpack_files) > 0
            for kf in kpack_files:
                assert kf.stat().st_size > 0