tests/conftest.py

"""共享测试配置和 fixtures."""

import os
import pytest


def _is_embedder_cache_available():
    """检查 embedding 模型是否已缓存在本地。"""
    cache_dirs = [
        os.path.join(os.path.expanduser("~"), ".cache", "torch", "sentence_transformers", "all-MiniLM-L6-v2"),
        os.path.join(os.path.expanduser("~"), ".cache", "huggingface", "hub", "models--sentence-transformers--all-MiniLM-L6-v2"),
    ]
    for d in cache_dirs:
        if os.path.isdir(d) and os.listdir(d):
            return True
    return False


# Custom marker: tests that need the embedding model
# Usage: @pytest.mark.requires_embedder
# Custom marker: tests that need GPU + training deps
# Usage: @pytest.mark.requires_train
def pytest_configure(config):
    config.addinivalue_line(
        "markers", "requires_embedder: mark test as requiring the sentence-transformers model"
    )
    config.addinivalue_line(
        "markers", "requires_train: mark test as requiring GPU + QLoRA training deps"
    )


def _gpu_available():
    try:
        import torch
        return torch.cuda.is_available()
    except ImportError:
        return False


def pytest_collection_modifyitems(config, items):
    """Skip tests that require embedder if model not cached."""
    if not _is_embedder_cache_available():
        skip_embedder = pytest.mark.skip(reason="Embedding model not cached locally (network required for first download)")
        for item in items:
            if "requires_embedder" in item.keywords:
                item.add_marker(skip_embedder)

    if not _gpu_available():
        skip_train = pytest.mark.skip(reason="CUDA GPU not available (required for QLoRA training test)")
        for item in items:
            if "requires_train" in item.keywords:
                item.add_marker(skip_train)