refactor: deduplicate and restructure test suite
- extract preprocessing config factories into tests/data/factories.py - keep conftest.py fixtures-only; stop importing builders from it - promote temp_dir fixture to root conftest for cross-directory reuse - unify duplicate BPE tokenizer builders into build_test_tokenizer - merge grpo/dpo online e2e tests into one parametrized integration test - extract engine mock factory and shared model batch builders - drop local tempfile usage in favor of shared fixtures No behavior change: 519 tests pass.
This commit is contained in:
+20
-21
@@ -5,12 +5,15 @@ import tempfile
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from tokenizers import Tokenizer, models, pre_tokenizers, trainers
|
||||
|
||||
from astrai.extension import KERNEL_NAMES, is_available
|
||||
from astrai.model.transformer import AutoRegressiveLM
|
||||
from astrai.tokenize import AutoTokenizer
|
||||
from tests.helpers import TINY_CONFIG, RandomTokenDataset, make_tiny_config
|
||||
from tests.helpers import (
|
||||
TINY_CONFIG,
|
||||
RandomTokenDataset,
|
||||
build_test_tokenizer,
|
||||
make_tiny_config,
|
||||
)
|
||||
|
||||
CUDA_AVAIL = torch.cuda.is_available()
|
||||
KERNEL_AVAIL = CUDA_AVAIL and all(is_available(k) for k in KERNEL_NAMES)
|
||||
@@ -30,18 +33,9 @@ def device():
|
||||
return "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
|
||||
def create_test_tokenizer(vocab_size: int = 1000) -> AutoTokenizer:
|
||||
def create_test_tokenizer(vocab_size: int = 1000):
|
||||
"""Create a simple tokenizer for testing purposes."""
|
||||
tokenizer = Tokenizer(models.BPE())
|
||||
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel()
|
||||
trainer = trainers.BpeTrainer(
|
||||
vocab_size=vocab_size, min_frequency=1, special_tokens=["<unk>", "<pad>"]
|
||||
)
|
||||
tokenizer.train_from_iterator([chr(i) for i in range(256)], trainer)
|
||||
auto_tokenizer = AutoTokenizer()
|
||||
auto_tokenizer._tokenizer = tokenizer
|
||||
auto_tokenizer._special_token_map = {"unk_token": "<unk>", "pad_token": "<pad>"}
|
||||
return auto_tokenizer
|
||||
return build_test_tokenizer(vocab_size)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
@@ -59,24 +53,29 @@ def test_model(device):
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def base_test_env(test_model, test_tokenizer):
|
||||
def temp_dir():
|
||||
"""Function-scoped temporary directory, cleaned up after each test."""
|
||||
d = tempfile.mkdtemp()
|
||||
yield d
|
||||
shutil.rmtree(d, ignore_errors=True)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def base_test_env(test_model, test_tokenizer, temp_dir):
|
||||
"""Function-scoped test environment with isolated temp directory."""
|
||||
test_dir = tempfile.mkdtemp()
|
||||
config_path = os.path.join(test_dir, "config.json")
|
||||
config_path = os.path.join(temp_dir, "config.json")
|
||||
with open(config_path, "w") as f:
|
||||
json.dump(TINY_CONFIG, f)
|
||||
|
||||
yield {
|
||||
return {
|
||||
"device": test_model["device"],
|
||||
"test_dir": str(test_dir),
|
||||
"test_dir": temp_dir,
|
||||
"config_path": config_path,
|
||||
"transformer_config": test_model["config"],
|
||||
"model": test_model["model"],
|
||||
"tokenizer": test_tokenizer,
|
||||
}
|
||||
|
||||
shutil.rmtree(test_dir)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def random_dataset():
|
||||
|
||||
Reference in New Issue
Block a user