tests/conftest.py

# SPDX-FileCopyrightText: 2026 SZL Holdings
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations

import os
import sys
from pathlib import Path

import pytest
import torch

_REPO_ROOT = Path(__file__).resolve().parents[1]
_TORCH_EXT = _REPO_ROOT / "torch-ext"


def pytest_configure(config: pytest.Config) -> None:
    config.addinivalue_line(
        "markers",
        "kernels_ci: kernel-builder CI tests (keep the marked set under 60s)",
    )


def _load_via_get_kernel():
    kernels = pytest.importorskip("kernels")
    return kernels.get_kernel(
        "SZLHOLDINGS/YARQA-ATTN",
        revision="main",
        trust_remote_code=True,
    )


@pytest.fixture(scope="session")
def kernel_mod():
    """Load the kernel the way Hub users will: ``get_kernel``.

    kernel-builder testshell sets ``LOCAL_KERNELS`` — load failures there
    are real failures, not skips. Source-tree developers may set
    ``SZL_SOURCE_TREE_TESTS=1`` to import ``torch-ext/yarqa_attn``
    directly. That path is labeled and is not a fabricated Hub load.
    """
    if os.environ.get("SZL_SOURCE_TREE_TESTS") == "1":
        sys.path.insert(0, str(_TORCH_EXT))
        import yarqa_attn

        return yarqa_attn
    if os.environ.get("LOCAL_KERNELS"):
        return _load_via_get_kernel()
    try:
        return _load_via_get_kernel()
    except Exception as exc:
        pytest.skip(
            "get_kernel could not load SZLHOLDINGS/YARQA-ATTN "
            f"({type(exc).__name__}: {exc}). Not a pass. Set "
            "SZL_SOURCE_TREE_TESTS=1 to exercise the source tree, or run "
            "under kernel-builder testshell (LOCAL_KERNELS)."
        )


@pytest.fixture(scope="session")
def cpu_device() -> torch.device:
    return torch.device("cpu")