overlay/tests/test_kernels.py

"""Tests for kernel stubs.

Verifies that:
  1. Every kernel stub file exists on disk.
  2. Python stub files contain a module-level docstring.
  3. Python stub files do NOT define a callable with that name
     (they are stubs — Phase 2 will implement them).

Run:
    uv run pytest tests/test_kernels.py -v
"""
import os
import pytest

import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
KERNEL_DIR = os.path.join(_REPO, "kernels")

# ---------------------------------------------------------------------------
# Existence checks — one per stub file
# ---------------------------------------------------------------------------

_ALL_STUBS = [
    ("triton", "ssd_exp_trap.py"),
    ("triton", "sinkhorn_fused.py"),
    ("triton", "bcnorm_fused.py"),
    ("triton", "oja_update.py"),
    ("tilelang", "ssd_mimo_prefill.py"),
    ("tilelang", "mhc_kernels.py"),
    ("cuda", "hash_kernel.cu"),
    ("cuda", "decode_kernels.cu"),
]

_PYTHON_STUBS = [
    ("triton", "ssd_exp_trap.py"),
    ("triton", "sinkhorn_fused.py"),
    ("triton", "bcnorm_fused.py"),
    ("triton", "oja_update.py"),
    ("tilelang", "ssd_mimo_prefill.py"),
    ("tilelang", "mhc_kernels.py"),
]

_CUDA_STUBS = [
    ("cuda", "hash_kernel.cu"),
    ("cuda", "decode_kernels.cu"),
]


@pytest.mark.parametrize("subdir,filename", _ALL_STUBS)
def test_kernel_stub_exists(subdir: str, filename: str) -> None:
    """Each kernel stub file must exist on disk."""
    path = os.path.join(KERNEL_DIR, subdir, filename)
    assert os.path.exists(path), (
        f"Missing kernel stub: kernels/{subdir}/{filename}\n"
        f"(Full path: {path})"
    )


@pytest.mark.parametrize("subdir,filename", _PYTHON_STUBS)
def test_python_stub_has_docstring(subdir: str, filename: str) -> None:
    """Python kernel stubs must have a module-level docstring."""
    path = os.path.join(KERNEL_DIR, subdir, filename)
    with open(path) as fh:
        content = fh.read()
    assert '"""' in content or "'''" in content, (
        f"No docstring found in kernels/{subdir}/{filename}"
    )


@pytest.mark.parametrize("subdir,filename", _PYTHON_STUBS)
def test_python_stub_is_non_empty(subdir: str, filename: str) -> None:
    """Python stub files must contain at least some text (not empty)."""
    path = os.path.join(KERNEL_DIR, subdir, filename)
    assert os.path.getsize(path) > 0, (
        f"kernels/{subdir}/{filename} is empty"
    )


@pytest.mark.parametrize("subdir,filename", _CUDA_STUBS)
def test_cuda_stub_has_comment(subdir: str, filename: str) -> None:
    """CUDA stub files must contain a comment describing their purpose."""
    path = os.path.join(KERNEL_DIR, subdir, filename)
    with open(path) as fh:
        content = fh.read()
    assert "/*" in content or "//" in content, (
        f"No comment found in kernels/{subdir}/{filename}"
    )


def test_kernel_dir_structure() -> None:
    """kernels/ directory contains triton/, tilelang/, and cuda/ subdirectories."""
    for subdir in ("triton", "tilelang", "cuda"):
        path = os.path.join(KERNEL_DIR, subdir)
        assert os.path.isdir(path), f"Missing kernels/{subdir}/ directory"


def test_triton_stub_count() -> None:
    """kernels/triton/ contains exactly the expected number of stubs."""
    triton_dir = os.path.join(KERNEL_DIR, "triton")
    py_files = [f for f in os.listdir(triton_dir) if f.endswith(".py")]
    expected = {name for _, name in _PYTHON_STUBS if _ == "triton"}
    assert expected.issubset(set(py_files)), (
        f"Missing triton stubs: {expected - set(py_files)}"
    )


def test_tilelang_stub_count() -> None:
    """kernels/tilelang/ contains exactly the expected number of stubs."""
    tilelang_dir = os.path.join(KERNEL_DIR, "tilelang")
    py_files = [f for f in os.listdir(tilelang_dir) if f.endswith(".py")]
    expected = {name for _, name in _PYTHON_STUBS if _ == "tilelang"}
    assert expected.issubset(set(py_files)), (
        f"Missing tilelang stubs: {expected - set(py_files)}"
    )


def test_cuda_stub_count() -> None:
    """kernels/cuda/ contains exactly the expected number of stubs."""
    cuda_dir = os.path.join(KERNEL_DIR, "cuda")
    cu_files = [f for f in os.listdir(cuda_dir) if f.endswith(".cu")]
    expected = {name for _, name in _CUDA_STUBS}
    assert expected.issubset(set(cu_files)), (
        f"Missing CUDA stubs: {expected - set(cu_files)}"
    )


# ---------------------------------------------------------------------------
# Content-quality checks for Python stubs
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("subdir,filename", _PYTHON_STUBS)
def test_stub_mentions_phase(subdir: str, filename: str) -> None:
    """Python stubs should document which Phase will implement them."""
    path = os.path.join(KERNEL_DIR, subdir, filename)
    with open(path) as fh:
        content = fh.read()
    assert "Phase" in content, (
        f"kernels/{subdir}/{filename} should mention 'Phase 1' or 'Phase 2' in its docs"
    )