File size: 1,814 Bytes
4fb3364 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | """共享测试配置和 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)
|