File size: 3,128 Bytes
768b4cb | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """
pytest 全局配置和共享 fixtures
提供所有测试模块共用的 fixture，包括临时目录、mock 对象等。
"""
import os
import sys
import json
import tempfile
import pytest
from pathlib import Path
from unittest.mock import MagicMock, patch
# 确保 hos_optimizer 包可被导入
sys.path.insert(0, str(Path(__file__).parent.parent))
@pytest.fixture
def tmp_dir():
"""创建临时目录，测试结束后自动清理"""
with tempfile.TemporaryDirectory() as tmpdir:
yield tmpdir
@pytest.fixture
def sample_yaml_config(tmp_dir):
"""创建一个示例 YAML 配置文件"""
config_path = os.path.join(tmp_dir, "test_config.yaml")
content = (
"backend: vllm\n"
"model:\n"
" path: /tmp/test_model\n"
" format: awq\n"
" dtype: float16\n"
"inference:\n"
" max_model_len: 512\n"
" gpu_memory_utilization: 0.9\n"
"sampling:\n"
" temperature: 0.7\n"
" top_p: 0.9\n"
)
with open(config_path, "w", encoding="utf-8") as f:
f.write(content)
return config_path
@pytest.fixture
def sample_alpaca_dataset(tmp_dir):
"""创建一个示例 Alpaca 格式数据集文件"""
dataset_path = os.path.join(tmp_dir, "dataset.json")
data = [
{
"instruction": "什么是网络安全？",
"input": "",
"output": "网络安全是指保护计算机网络免受未经授权的访问。"
},
{
"instruction": "解释SQL注入",
"input": "请举例说明",
"output": "SQL注入是通过在输入中插入恶意SQL代码来攻击数据库。"
},
]
with open(dataset_path, "w", encoding="utf-8") as f:
json.dump(data, f, ensure_ascii=False)
return dataset_path
@pytest.fixture
def sample_sharegpt_dataset(tmp_dir):
"""创建一个示例 ShareGPT 格式数据集文件"""
dataset_path = os.path.join(tmp_dir, "sharegpt_dataset.json")
data = [
{
"conversations": [
{"from": "human", "value": "你好"},
{"from": "gpt", "value": "你好！有什么可以帮助你的吗？"},
]
},
]
with open(dataset_path, "w", encoding="utf-8") as f:
json.dump(data, f, ensure_ascii=False)
return dataset_path
@pytest.fixture
def mock_torch():
"""Mock torch 模块，避免实际加载模型"""
with patch.dict("sys.modules", {
"torch": MagicMock(),
"torch.cuda": MagicMock(),
}):
yield sys.modules["torch"]
@pytest.fixture
def mock_transformers():
"""Mock transformers 库"""
with patch.dict("sys.modules", {
"transformers": MagicMock(),
"transformers.AutoModelForCausalLM": MagicMock(),
"transformers.AutoTokenizer": MagicMock(),
"transformers.BitsAndBytesConfig": MagicMock(),
"transformers.TrainingArguments": MagicMock(),
"transformers.Trainer": MagicMock(),
"transformers.DataCollatorForSeq2Seq": MagicMock(),
}):
yield sys.modules["transformers"]
|