TheRock/rocm-systems/shared/kpack/tests/common/test_kpack_transform.py

"""
Common tests for kpack transformation using the generic API.

These tests use the format-agnostic rocm_kpack API and are parameterized
to run against both ELF (Linux) and PE/COFF (Windows) binaries.
"""

from pathlib import Path

import pytest

from rocm_kpack import (
    kpack_offload_binary,
    read_kpack_ref_marker,
    is_fat_binary,
    detect_binary_format,
    NotFatBinaryError,
)


class TestFormatDetection:
    """Tests for binary format detection."""

    def test_detect_format(self, single_arch_exe: Path, asset_config: tuple[str, str]):
        """Test that format is correctly detected."""
        platform, _ = asset_config
        expected_format = "coff" if platform == "windows" else "elf"
        assert detect_binary_format(single_arch_exe) == expected_format

    def test_is_fat_binary_true(self, single_arch_exe: Path):
        """Test that fat binaries are detected as such."""
        assert is_fat_binary(single_arch_exe) is True

    def test_is_fat_binary_false(self, host_only_exe: Path):
        """Test that host-only binaries are not fat binaries."""
        assert is_fat_binary(host_only_exe) is False


class TestKpackOffloadBinary:
    """Tests for generic kpack_offload_binary."""

    def test_transform_single_arch_exe(
        self, single_arch_exe: Path, tmp_path: Path, asset_config: tuple[str, str]
    ):
        """Test transforming a single-arch executable."""
        platform, _ = asset_config
        output = tmp_path / "output.exe"

        result = kpack_offload_binary(
            input_path=single_arch_exe,
            output_path=output,
            kpack_search_paths=["kernels.kpack"],
            kernel_name="test_kernel",
        )

        # Verify output exists
        assert output.exists()

        # Verify format is preserved
        expected_format = "coff" if platform == "windows" else "elf"
        assert result["format"] == expected_format

        # Verify result has expected keys
        assert "removed" in result
        assert "original_size" in result
        assert "new_size" in result

    def test_transform_multi_arch_exe(
        self, multi_arch_exe: Path, tmp_path: Path, asset_config: tuple[str, str]
    ):
        """Test transforming a multi-arch executable."""
        platform, _ = asset_config
        output = tmp_path / "output.exe"

        original_size = multi_arch_exe.stat().st_size

        result = kpack_offload_binary(
            input_path=multi_arch_exe,
            output_path=output,
            kpack_search_paths=["kernels.kpack"],
            kernel_name="test_kernel",
        )

        assert output.exists()
        assert result["original_size"] == original_size
        assert result["new_size"] == output.stat().st_size

    def test_transform_shared_lib(
        self, single_arch_lib: Path, tmp_path: Path, asset_config: tuple[str, str]
    ):
        """Test transforming a shared library."""
        platform, _ = asset_config
        ext = ".dll" if platform == "windows" else ".so"
        output = tmp_path / f"output{ext}"

        result = kpack_offload_binary(
            input_path=single_arch_lib,
            output_path=output,
            kpack_search_paths=["libs/kernels.kpack"],
            kernel_name="lib_kernel",
        )

        assert output.exists()
        assert result["new_size"] > 0


class TestKpackRefMarker:
    """Tests for marker section read/write."""

    def test_marker_roundtrip(self, single_arch_exe: Path, tmp_path: Path):
        """Test that marker can be written and read back."""
        output = tmp_path / "output.exe"

        kpack_paths = ["path/to/kernels.kpack", "fallback.kpack"]
        kernel_name = "my_test_kernel"

        kpack_offload_binary(
            input_path=single_arch_exe,
            output_path=output,
            kpack_search_paths=kpack_paths,
            kernel_name=kernel_name,
        )

        # Read marker back
        marker = read_kpack_ref_marker(output)

        assert marker is not None
        assert marker["kpack_search_paths"] == kpack_paths
        assert marker["kernel_name"] == kernel_name

    def test_marker_missing_in_original(self, single_arch_exe: Path):
        """Test that untransformed binaries have no marker."""
        marker = read_kpack_ref_marker(single_arch_exe)
        assert marker is None


class TestNotFatBinary:
    """Tests for NotFatBinaryError."""

    def test_host_only_exe_raises(self, host_only_exe: Path, tmp_path: Path):
        """Test that host-only executable raises NotFatBinaryError."""
        output = tmp_path / "output.exe"

        with pytest.raises(NotFatBinaryError):
            kpack_offload_binary(
                input_path=host_only_exe,
                output_path=output,
                kpack_search_paths=["test.kpack"],
                kernel_name="test",
            )

    def test_host_only_lib_raises(self, host_only_lib: Path, tmp_path: Path):
        """Test that host-only library raises NotFatBinaryError."""
        output = tmp_path / "output"

        with pytest.raises(NotFatBinaryError):
            kpack_offload_binary(
                input_path=host_only_lib,
                output_path=output,
                kpack_search_paths=["test.kpack"],
                kernel_name="test",
            )