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.

tests/conftest.py

@@ -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 (
-from tests.helpers import TINY_CONFIG, RandomTokenDataset, make_tiny_config
+    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())
+    return build_test_tokenizer(vocab_size)
-    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
 
 
 @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(temp_dir, "config.json")
-    config_path = os.path.join(test_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():

tests/data/conftest.py

@@ -1,21 +1,15 @@
 import json
 import os
-import tempfile
 
 import pytest
-from tokenizers import Tokenizer, models, pre_tokenizers, trainers
 
-from astrai.config.preprocess_config import (
-    InputConfig,
-    PipelineConfig,
-    ProcessingConfig,
-)
 from astrai.preprocessing.builder import (
     MultiOutputMaskBuilder,
     SectionedMaskBuilder,
     SingleOutputMaskBuilder,
 )
-from astrai.tokenize import AutoTokenizer
+from tests.data.factories import make_grpo_config
+from tests.helpers import build_test_tokenizer
 
 _SPECIAL_TOKENS_CONFIG = {
     "bos_token": "<|begin_of_sentence|>",
@@ -41,55 +35,42 @@ _CHAT_TEMPLATE = (
     "{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
 )
 
-_CHAT_SECTIONS = [{"field": "messages", "action": "$role", "template": True}]
 
-_INSTRUCTION_SECTIONS = [
+_CHAT_TOKENIZER_DATA = [
-    {"field": "prompt", "action": "mask", "add_special_tokens": True},
+    "hello world",
-    {"field": "response", "action": "train"},
+    "Hi there!",
+    "You are helpful.",
+    "What is 2+2?",
+    "Tell me a story about dragons and knights.",
+    "Sure, here is a tale.",
+    "Translate to French: Hello",
+    "Bonjour",
+    "Artificial Intelligence is a field of computer science.",
+    "system",
+    "user",
+    "assistant",
+    "<|im_start|>",
+    "<|im_end|>",
+    *[chr(i) for i in range(32, 127)],
 ]
 
-_TEXT_SECTIONS = [{"field": "text", "action": "train"}]
+_CHAT_TOKENIZER_MAP = {
-
+    "bos_token": "<|begin_of_sentence|>",
-_GRPO_RESPONSE_SECTIONS = [{"field": "responses", "action": "train"}]
+    "eos_token": "<|end_of_sentence|>",
+    "pad_token": "<|_pad_|>",
+    "unk_token": "<|_unk_|>",
+}
 
 
 def _build_chat_tokenizer():
-    tok = Tokenizer(models.BPE())
+    return build_test_tokenizer(
-    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
-    tr = trainers.BpeTrainer(
         vocab_size=512,
-        min_frequency=1,
         special_tokens=_SPECIAL_TOKENS,
+        special_token_map=_CHAT_TOKENIZER_MAP,
+        add_prefix_space=False,
+        train_data=_CHAT_TOKENIZER_DATA,
+        chat_template=_CHAT_TEMPLATE,
     )
-    train_data = [
-        "hello world",
-        "Hi there!",
-        "You are helpful.",
-        "What is 2+2?",
-        "Tell me a story about dragons and knights.",
-        "Sure, here is a tale.",
-        "Translate to French: Hello",
-        "Bonjour",
-        "Artificial Intelligence is a field of computer science.",
-        "system",
-        "user",
-        "assistant",
-        "<|im_start|>",
-        "<|im_end|>",
-        *[chr(i) for i in range(32, 127)],
-    ]
-    tok.train_from_iterator(train_data, tr)
-
-    auto_tok = AutoTokenizer()
-    auto_tok._tokenizer = tok
-    auto_tok._special_token_map = {
-        "bos_token": "<|begin_of_sentence|>",
-        "eos_token": "<|end_of_sentence|>",
-        "pad_token": "<|_pad_|>",
-        "unk_token": "<|_unk_|>",
-    }
-    auto_tok.set_chat_template(_CHAT_TEMPLATE)
-    return auto_tok
 
 
 @pytest.fixture(scope="session")
@@ -97,116 +78,11 @@ def chat_tokenizer():
     return _build_chat_tokenizer()
 
 
-@pytest.fixture
+def _write_tokenizer_dir(dir_path, tokenizer, tokenizer_config):
-def temp_dir():
+    """Persist a tokenizer plus ``tokenizer_config.json`` into *dir_path*."""
-    d = tempfile.mkdtemp()
+    tokenizer._tokenizer.save(os.path.join(dir_path, "tokenizer.json"))
-    yield d
+    with open(os.path.join(dir_path, "tokenizer_config.json"), "w") as f:
-    import shutil
+        json.dump(tokenizer_config, f)
-
-    shutil.rmtree(d, ignore_errors=True)
-
-
-def make_chat_config():
-    return PipelineConfig(
-        input=InputConfig(sections=_CHAT_SECTIONS),
-        mask={"system": "mask", "user": "mask", "assistant": "train"},
-        mask_default="mask",
-        preprocessing=ProcessingConfig(max_seq_len=2048),
-    )
-
-
-def make_instruction_config():
-    return PipelineConfig(
-        input=InputConfig(sections=_INSTRUCTION_SECTIONS),
-        mask={"prompt": "mask", "response": "train"},
-        mask_default="mask",
-        preprocessing=ProcessingConfig(max_seq_len=2048),
-    )
-
-
-def make_text_config():
-    return PipelineConfig(
-        input=InputConfig(sections=_TEXT_SECTIONS),
-        preprocessing=ProcessingConfig(
-            max_seq_len=2048, min_chars=1, max_chars=2_000_000
-        ),
-    )
-
-
-def make_dpo_chat_config():
-    return PipelineConfig(
-        input=InputConfig(
-            sources={
-                "chosen": {
-                    "sections": [
-                        {"field": "chosen", "action": "$role", "template": True}
-                    ]
-                },
-                "rejected": {
-                    "sections": [
-                        {"field": "rejected", "action": "$role", "template": True}
-                    ]
-                },
-            }
-        ),
-        mask={"user": "mask", "assistant": "train"},
-        mask_default="mask",
-        preprocessing=ProcessingConfig(max_seq_len=2048),
-    )
-
-
-def make_grpo_config():
-    return PipelineConfig(
-        input=InputConfig(
-            sources={
-                "prompts": {
-                    "sections": [
-                        {"field": "prompt", "action": "mask", "template": True}
-                    ]
-                },
-                "responses": {
-                    "sections": _GRPO_RESPONSE_SECTIONS,
-                    "list_field": True,
-                    "mask_key": "masks",
-                },
-                "rewards": {
-                    "sections": [{"field": "rewards", "action": "value"}],
-                },
-            }
-        ),
-        mask={"user": "mask", "assistant": "train"},
-        mask_default="mask",
-        preprocessing=ProcessingConfig(max_seq_len=2048),
-    )
-
-
-def make_grpo_no_template_config():
-    return PipelineConfig(
-        input=InputConfig(
-            sources={
-                "prompts": {
-                    "sections": [
-                        {
-                            "field": "prompt",
-                            "action": "mask",
-                            "add_special_tokens": True,
-                        }
-                    ]
-                },
-                "responses": {
-                    "sections": _GRPO_RESPONSE_SECTIONS,
-                    "list_field": True,
-                    "mask_key": "masks",
-                },
-                "rewards": {
-                    "sections": [{"field": "rewards", "action": "value"}],
-                },
-            }
-        ),
-        mask={"user": "mask", "assistant": "train"},
-        mask_default="mask",
-        preprocessing=ProcessingConfig(max_seq_len=2048),
-    )
 
 
 @pytest.fixture
@@ -228,11 +104,11 @@ def multi_builder():
 def tokenizer_dir(temp_dir, test_tokenizer):
     d = os.path.join(temp_dir, "tok")
     os.makedirs(d, exist_ok=True)
-    test_tokenizer._tokenizer.save(os.path.join(d, "tokenizer.json"))
+    _write_tokenizer_dir(
-    with open(os.path.join(d, "tokenizer_config.json"), "w") as f:
+        d,
-        json.dump(
+        test_tokenizer,
-            {"special_tokens": {"pad_token": "<|_pad_|>", "unk_token": "<|_unk_|>"}}, f
+        {"special_tokens": {"pad_token": "<|_pad_|>", "unk_token": "<|_unk_|>"}},
-        )
+    )
     return d
 
 
@@ -240,10 +116,9 @@ def tokenizer_dir(temp_dir, test_tokenizer):
 def chat_tokenizer_dir(temp_dir, chat_tokenizer):
     d = os.path.join(temp_dir, "tok")
     os.makedirs(d, exist_ok=True)
-    chat_tokenizer._tokenizer.save(os.path.join(d, "tokenizer.json"))
+    _write_tokenizer_dir(
-    with open(os.path.join(d, "tokenizer_config.json"), "w") as f:
+        d,
-        json.dump(
+        chat_tokenizer,
-            {"special_tokens": _SPECIAL_TOKENS_CONFIG, "chat_template": _CHAT_TEMPLATE},
+        {"special_tokens": _SPECIAL_TOKENS_CONFIG, "chat_template": _CHAT_TEMPLATE},
-            f,
+    )
-        )
     return d

tests/data/factories.py

@@ -0,0 +1,86 @@
+"""Test data builders for preprocessing and dataset scenarios."""
+
+from astrai.config.preprocess_config import (
+    InputConfig,
+    PipelineConfig,
+    ProcessingConfig,
+)
+
+CHAT_SECTIONS = [{"field": "messages", "action": "$role", "template": True}]
+INSTRUCTION_SECTIONS = [
+    {"field": "prompt", "action": "mask", "add_special_tokens": True},
+    {"field": "response", "action": "train"},
+]
+TEXT_SECTIONS = [{"field": "text", "action": "train"}]
+GRPO_RESPONSE_SECTIONS = [{"field": "responses", "action": "train"}]
+
+
+def make_pipeline_config(sections, *, mask=None, preprocessing=None, sources=None):
+    """Build a pipeline config with the common test defaults."""
+    return PipelineConfig(
+        input=InputConfig(sections=sections, sources=sources),
+        mask={} if mask is None else mask,
+        mask_default="mask",
+        preprocessing=preprocessing or ProcessingConfig(max_seq_len=2048),
+    )
+
+
+def make_chat_config():
+    return make_pipeline_config(
+        CHAT_SECTIONS,
+        mask={"system": "mask", "user": "mask", "assistant": "train"},
+    )
+
+
+def make_instruction_config():
+    return make_pipeline_config(
+        INSTRUCTION_SECTIONS,
+        mask={"prompt": "mask", "response": "train"},
+    )
+
+
+def make_text_config():
+    return make_pipeline_config(
+        TEXT_SECTIONS,
+        preprocessing=ProcessingConfig(
+            max_seq_len=2048, min_chars=1, max_chars=2_000_000
+        ),
+    )
+
+
+def make_dpo_chat_config():
+    sources = {
+        name: {"sections": [{"field": name, "action": "$role", "template": True}]}
+        for name in ("chosen", "rejected")
+    }
+    return make_pipeline_config(
+        None,
+        mask={"user": "mask", "assistant": "train"},
+        sources=sources,
+    )
+
+
+def make_grpo_config(*, template=True):
+    prompt_section = {"field": "prompt", "action": "mask"}
+    if template:
+        prompt_section["template"] = True
+    else:
+        prompt_section["add_special_tokens"] = True
+    sources = {
+        "prompts": {"sections": [prompt_section]},
+        "responses": {
+            "sections": GRPO_RESPONSE_SECTIONS,
+            "list_field": True,
+            "mask_key": "masks",
+        },
+        "rewards": {"sections": [{"field": "rewards", "action": "value"}]},
+    }
+    return make_pipeline_config(
+        None,
+        mask={"user": "mask", "assistant": "train"},
+        sources=sources,
+    )
+
+
+def make_grpo_no_template_config():
+    return make_grpo_config(template=False)

tests/data/test_dataset.py

@@ -24,7 +24,7 @@ from astrai.serialization import (
     load_bin,
     save_bin,
 )
-from tests.data.conftest import make_grpo_no_template_config
+from tests.data.factories import make_grpo_config
 
 
 def _rand_seq(length, vocab=1000):
@@ -797,7 +797,7 @@ def test_grpo_builder_preserves_response_boundaries(base_test_env):
     _save_test_tokenizer(base_test_env["test_dir"], tokenizer)
 
     builder = SectionedMaskBuilder()
-    config = make_grpo_no_template_config()
+    config = make_grpo_config(template=False)
     config.preprocessing.max_seq_len = 128
 
     item = {

tests/data/test_preprocess_builder.py

@@ -12,10 +12,10 @@ from astrai.preprocessing.builder import (
     SectionedMaskBuilder,
     SingleOutputMaskBuilder,
 )
-from tests.data.conftest import (
+from tests.data.factories import (
-    _CHAT_SECTIONS,
+    CHAT_SECTIONS,
-    _INSTRUCTION_SECTIONS,
+    INSTRUCTION_SECTIONS,
-    _TEXT_SECTIONS,
+    TEXT_SECTIONS,
     make_chat_config,
     make_dpo_chat_config,
     make_grpo_config,
@@ -101,7 +101,7 @@ def test_chat_uniform_masking(
     mask_rules, mask_default, expect_nonzero, chat_tokenizer, builder
 ):
     config = PipelineConfig(
-        input=InputConfig(sections=_CHAT_SECTIONS),
+        input=InputConfig(sections=CHAT_SECTIONS),
         mask=mask_rules,
         mask_default=mask_default,
         preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -128,7 +128,7 @@ def test_chat_empty_messages(chat_tokenizer, builder):
 
 def test_chat_domain_extraction(chat_tokenizer, builder):
     config = PipelineConfig(
-        input=InputConfig(sections=_CHAT_SECTIONS),
+        input=InputConfig(sections=CHAT_SECTIONS),
         mask={"assistant": "train"},
         mask_default="mask",
         preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -147,7 +147,7 @@ def test_chat_domain_extraction(chat_tokenizer, builder):
 
 def test_chat_truncation(chat_tokenizer, builder):
     config = PipelineConfig(
-        input=InputConfig(sections=_CHAT_SECTIONS),
+        input=InputConfig(sections=CHAT_SECTIONS),
         mask={"assistant": "train"},
         mask_default="mask",
         preprocessing=ProcessingConfig(max_seq_len=10),
@@ -237,7 +237,7 @@ def test_text_empty(test_tokenizer, builder):
 
 def test_text_too_short(test_tokenizer, builder):
     config = PipelineConfig(
-        input=InputConfig(sections=_TEXT_SECTIONS),
+        input=InputConfig(sections=TEXT_SECTIONS),
         preprocessing=ProcessingConfig(min_chars=100),
     )
     assert builder.build({"text": "short"}, config, test_tokenizer) is None
@@ -245,7 +245,7 @@ def test_text_too_short(test_tokenizer, builder):
 
 def test_text_truncation(test_tokenizer, builder):
     config = PipelineConfig(
-        input=InputConfig(sections=_TEXT_SECTIONS),
+        input=InputConfig(sections=TEXT_SECTIONS),
         preprocessing=ProcessingConfig(max_seq_len=3, min_chars=1),
     )
     item = {"text": "This is a very long text that should be truncated"}
@@ -255,7 +255,7 @@ def test_text_truncation(test_tokenizer, builder):
 
 def test_sectioned_chat(chat_tokenizer, builder):
     config = PipelineConfig(
-        input=InputConfig(sections=_CHAT_SECTIONS),
+        input=InputConfig(sections=CHAT_SECTIONS),
         mask={"system": "mask", "user": "mask", "assistant": "train"},
         mask_default="mask",
         preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -275,7 +275,7 @@ def test_sectioned_chat(chat_tokenizer, builder):
 
 def test_sectioned_instruction(test_tokenizer, builder):
     config = PipelineConfig(
-        input=InputConfig(sections=_INSTRUCTION_SECTIONS),
+        input=InputConfig(sections=INSTRUCTION_SECTIONS),
         preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=0),
     )
     item = {"prompt": "Q: Why?", "response": "A: Because."}
@@ -288,7 +288,7 @@ def test_sectioned_instruction(test_tokenizer, builder):
 
 def test_sectioned_text(test_tokenizer, builder):
     config = PipelineConfig(
-        input=InputConfig(sections=_TEXT_SECTIONS),
+        input=InputConfig(sections=TEXT_SECTIONS),
         preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=1),
     )
     item = {"text": "Hello world, this is a test."}
@@ -299,7 +299,7 @@ def test_sectioned_text(test_tokenizer, builder):
 
 def test_sectioned_text_too_short(test_tokenizer, builder):
     config = PipelineConfig(
-        input=InputConfig(sections=_TEXT_SECTIONS),
+        input=InputConfig(sections=TEXT_SECTIONS),
         preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=100),
     )
     assert builder.build({"text": "short"}, config, test_tokenizer) is None

tests/data/test_preprocess_config.py

@@ -4,9 +4,9 @@ from astrai.config.preprocess_config import (
     InputConfig,
     PipelineConfig,
 )
-from tests.data.conftest import (
+from tests.data.factories import (
-    _INSTRUCTION_SECTIONS,
+    INSTRUCTION_SECTIONS,
-    _TEXT_SECTIONS,
+    TEXT_SECTIONS,
     make_dpo_chat_config,
 )
 
@@ -43,26 +43,26 @@ def test_from_dict_flat():
 
 def test_to_dict_roundtrip():
     config = PipelineConfig(
-        input=InputConfig(sections=_INSTRUCTION_SECTIONS),
+        input=InputConfig(sections=INSTRUCTION_SECTIONS),
         mask={"prompt": "mask", "response": "train"},
         mask_default="mask",
     )
     d = config.to_dict()
     config2 = PipelineConfig.from_dict(d)
-    assert config2.input.sections == _INSTRUCTION_SECTIONS
+    assert config2.input.sections == INSTRUCTION_SECTIONS
     assert config2.mask == {"prompt": "mask", "response": "train"}
 
 
 def test_to_file_from_file(temp_dir):
     config = PipelineConfig(
-        input=InputConfig(sections=_TEXT_SECTIONS),
+        input=InputConfig(sections=TEXT_SECTIONS),
         mask={"text": "train"},
         mask_default="mask",
     )
     path = os.path.join(temp_dir, "config.json")
     config.to_file(path)
     loaded = PipelineConfig.from_file(path)
-    assert loaded.input.sections == _TEXT_SECTIONS
+    assert loaded.input.sections == TEXT_SECTIONS
     assert loaded.mask == {"text": "train"}
 
 

tests/data/test_preprocess_pipeline.py

@@ -9,10 +9,10 @@ from astrai.config.preprocess_config import (
 )
 from astrai.preprocessing.packing import PackingStrategyFactory
 from astrai.preprocessing.pipeline import Pipeline, filter_by_length
-from tests.data.conftest import (
+from tests.data.factories import (
-    _CHAT_SECTIONS,
+    CHAT_SECTIONS,
-    _INSTRUCTION_SECTIONS,
+    INSTRUCTION_SECTIONS,
-    _TEXT_SECTIONS,
+    TEXT_SECTIONS,
     make_dpo_chat_config,
     make_grpo_no_template_config,
 )
@@ -54,7 +54,7 @@ def test_full_chat_pipeline(temp_dir, chat_tokenizer_dir):
         )
 
     config = PipelineConfig(
-        input=InputConfig(sections=_CHAT_SECTIONS),
+        input=InputConfig(sections=CHAT_SECTIONS),
         mask={"system": "mask", "user": "mask", "assistant": "train"},
         mask_default="mask",
         preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -97,7 +97,7 @@ def test_full_text_pipeline(temp_dir, tokenizer_dir):
         )
 
     config = PipelineConfig(
-        input=InputConfig(sections=_TEXT_SECTIONS),
+        input=InputConfig(sections=TEXT_SECTIONS),
         preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=10),
         output=OutputConfig(storage_format="bin"),
     )
@@ -138,7 +138,7 @@ def test_full_instruction_pipeline(temp_dir, tokenizer_dir):
         )
 
     config = PipelineConfig(
-        input=InputConfig(sections=_INSTRUCTION_SECTIONS),
+        input=InputConfig(sections=INSTRUCTION_SECTIONS),
         mask={"prompt": "mask", "response": "train"},
         mask_default="mask",
         preprocessing=ProcessingConfig(max_seq_len=2048),
@@ -164,7 +164,7 @@ def test_dtype_override(temp_dir, tokenizer_dir):
         f.write(json.dumps({"prompt": "Q", "response": "A"}) + "\n")
 
     config = PipelineConfig(
-        input=InputConfig(sections=_INSTRUCTION_SECTIONS),
+        input=InputConfig(sections=INSTRUCTION_SECTIONS),
         mask={"prompt": "mask", "response": "train"},
         mask_default="mask",
         preprocessing=ProcessingConfig(max_seq_len=2048),

tests/helpers.py

@@ -4,10 +4,12 @@ import json
 import os
 
 import torch
+from tokenizers import Tokenizer, models, pre_tokenizers, trainers
 from torch.utils.data import Dataset
 
 from astrai.config.model_config import AutoRegressiveLMConfig
 from astrai.model.transformer import AutoRegressiveLM
+from astrai.tokenize import AutoTokenizer
 
 TINY_CONFIG = dict(
     vocab_size=1000,
@@ -57,6 +59,44 @@ def make_model(device, **cfg_overrides):
     return model, cfg
 
 
+def build_test_tokenizer(
+    vocab_size: int = 1000,
+    *,
+    special_tokens=("<unk>", "<pad>"),
+    special_token_map=None,
+    add_prefix_space: bool = True,
+    train_data=None,
+    chat_template: str | None = None,
+) -> AutoTokenizer:
+    """Build a lightweight BPE ``AutoTokenizer`` for tests.
+
+    ``special_token_map`` defaults to ``{"unk_token", "pad_token"}``
+    pointing at the first two special tokens.
+    """
+    tokenizer = Tokenizer(models.BPE())
+    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(
+        add_prefix_space=add_prefix_space
+    )
+    trainer = trainers.BpeTrainer(
+        vocab_size=vocab_size,
+        min_frequency=1,
+        special_tokens=list(special_tokens),
+    )
+    tokenizer.train_from_iterator(
+        train_data if train_data is not None else [chr(i) for i in range(256)],
+        trainer,
+    )
+    auto_tokenizer = AutoTokenizer()
+    auto_tokenizer._tokenizer = tokenizer
+    auto_tokenizer._special_token_map = special_token_map or {
+        "unk_token": special_tokens[0],
+        "pad_token": special_tokens[1],
+    }
+    if chat_template is not None:
+        auto_tokenizer.set_chat_template(chat_template)
+    return auto_tokenizer
+
+
 def make_frozen(model, device):
     """Create a frozen, eval-mode copy of *model* with identical weights."""
     cfg = make_rollout_config()

tests/inference/test_engine.py

@@ -8,6 +8,17 @@ from astrai.inference import STOP
 from astrai.inference.engine import GenerateResult, InferenceEngine
 
 
+def _make_engine_mocks(decode=None):
+    """Build the standard mock model/tokenizer pair used by engine tests."""
+    mock_model = MagicMock()
+    mock_tokenizer = MagicMock()
+    mock_tokenizer.encode.return_value = [1, 2, 3]
+    mock_tokenizer.stop_ids = [0]
+    if decode is not None:
+        mock_tokenizer.decode.return_value = decode
+    return mock_model, mock_tokenizer
+
+
 def test_result_append_single():
     r = GenerateResult(count=1)
     r.append("hello", 0)
@@ -102,11 +113,7 @@ def test_result_get_results():
 
 
 def test_engine_generate_non_streaming_single():
-    mock_model = MagicMock()
+    mock_model, mock_tokenizer = _make_engine_mocks(decode="response")
-    mock_tokenizer = MagicMock()
-    mock_tokenizer.encode.return_value = [1, 2, 3]
-    mock_tokenizer.decode.return_value = "response"
-    mock_tokenizer.stop_ids = [0]
 
     with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
         instance = MockSched.return_value
@@ -125,11 +132,7 @@ def test_engine_generate_non_streaming_single():
 
 
 def test_engine_generate_streaming_yields_tokens():
-    mock_model = MagicMock()
+    mock_model, mock_tokenizer = _make_engine_mocks(decode="tok")
-    mock_tokenizer = MagicMock()
-    mock_tokenizer.encode.return_value = [1, 2, 3]
-    mock_tokenizer.decode.return_value = "tok"
-    mock_tokenizer.stop_ids = [0]
 
     callbacks_saved = []
 
@@ -154,11 +157,7 @@ def test_engine_generate_streaming_yields_tokens():
 
 
 def test_engine_generate_non_streaming_batch():
-    mock_model = MagicMock()
+    mock_model, mock_tokenizer = _make_engine_mocks(decode="r")
-    mock_tokenizer = MagicMock()
-    mock_tokenizer.encode.return_value = [1, 2, 3]
-    mock_tokenizer.decode.return_value = "r"
-    mock_tokenizer.stop_ids = [0]
 
     with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
         instance = MockSched.return_value
@@ -177,10 +176,7 @@ def test_engine_generate_non_streaming_batch():
 
 
 def test_engine_generate_zero_max_tokens_returns_empty():
-    mock_model = MagicMock()
+    mock_model, mock_tokenizer = _make_engine_mocks()
-    mock_tokenizer = MagicMock()
-    mock_tokenizer.encode.return_value = [1, 2, 3]
-    mock_tokenizer.stop_ids = [0]
 
     with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
         instance = MockSched.return_value
@@ -192,8 +188,7 @@ def test_engine_generate_zero_max_tokens_returns_empty():
 
 
 def test_engine_generate_zero_max_tokens_stream_is_empty():
-    mock_model = MagicMock()
+    mock_model, mock_tokenizer = _make_engine_mocks()
-    mock_tokenizer = MagicMock()
 
     with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
         instance = MockSched.return_value
@@ -203,8 +198,7 @@ def test_engine_generate_zero_max_tokens_stream_is_empty():
 
 
 def test_engine_passes_backend_to_scheduler():
-    mock_model = MagicMock()
+    mock_model, mock_tokenizer = _make_engine_mocks()
-    mock_tokenizer = MagicMock()
 
     with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
         InferenceEngine(
@@ -218,8 +212,7 @@ def test_engine_passes_backend_to_scheduler():
 
 
 def test_generate_captures_calling_backend_context():
-    mock_model = MagicMock()
+    mock_model, mock_tokenizer = _make_engine_mocks()
-    mock_tokenizer = MagicMock()
     captured = []
 
     with patch("astrai.inference.engine.InferenceScheduler") as MockSched:

tests/module/test_lora.py

@@ -1,9 +1,6 @@
-import tempfile
-
 import pytest
 import torch
 
-from astrai.config.model_config import AutoRegressiveLMConfig
 from astrai.model import AutoRegressiveLM
 from astrai.model.components.linear import Linear
 from astrai.model.components.lora import (
@@ -16,22 +13,20 @@ from astrai.model.components.lora import (
     merge_lora,
     save_lora,
 )
+from tests.helpers import make_tiny_config
 
-MODEL_KWARGS = dict(
+LORA_MODEL_KWARGS = dict(
     vocab_size=1000,
     hidden_size=64,
     num_attention_heads=4,
     num_key_value_heads=2,
     intermediate_size=128,
-    num_hidden_layers=2,
     max_position_embeddings=32,
-    rms_norm_eps=1e-5,
 )
 
 
 def _make_model(**kwargs):
-    kw = {**MODEL_KWARGS, **kwargs}
+    config = make_tiny_config(**{**LORA_MODEL_KWARGS, **kwargs})
-    config = AutoRegressiveLMConfig(**kw)
     model = AutoRegressiveLM(config)
     model.eval()
     return model
@@ -227,7 +222,7 @@ def test_state_dict_after_inject_consistent_with_original():
     assert len(lora_keys) > 0
 
 
-def test_save_load_roundtrip():
+def test_save_load_roundtrip(temp_dir):
     model = _make_model()
     cfg = inject_lora(model, r=4, alpha=8, target_modules={"q_proj"})
 
@@ -240,20 +235,19 @@ def test_save_load_roundtrip():
     with torch.no_grad():
         out_src = model(x)["logits"].clone()
 
-    with tempfile.TemporaryDirectory() as tmpdir:
+    save_lora(model, temp_dir, cfg)
-        save_lora(model, tmpdir, cfg)
 
-        model2 = _make_model()
+    model2 = _make_model()
-        model2.load_state_dict(model.state_dict(), strict=False)
+    model2.load_state_dict(model.state_dict(), strict=False)
-        load_lora(model2, tmpdir)
+    load_lora(model2, temp_dir)
 
-        with torch.no_grad():
+    with torch.no_grad():
-            out_dst = model2(x)["logits"]
+        out_dst = model2(x)["logits"]
 
-        torch.testing.assert_close(out_src, out_dst)
+    torch.testing.assert_close(out_src, out_dst)
 
 
-def test_save_after_merge_raises():
+def test_save_after_merge_raises(temp_dir):
     model = _make_model()
     cfg = inject_lora(model, r=4, alpha=8, target_modules={"q_proj"})
 
@@ -262,16 +256,14 @@ def test_save_after_merge_raises():
             if isinstance(m, LoRALinear):
                 m.lora_B.fill_(0.5)
 
-    with tempfile.TemporaryDirectory() as tmpdir:
+    save_lora(model, temp_dir, cfg)
-        save_lora(model, tmpdir, cfg)
+    merge_lora(model)
-        merge_lora(model)
 
-        with tempfile.TemporaryDirectory() as tmpdir2:
+    with pytest.raises(RuntimeError, match="No LoRA parameters"):
-            with pytest.raises(RuntimeError, match="No LoRA parameters"):
+        save_lora(model, temp_dir, cfg)
-                save_lora(model, tmpdir2, cfg)
 
 
-def test_load_lora_on_already_injected():
+def test_load_lora_on_already_injected(temp_dir):
     model = _make_model()
     inject_lora(model, r=4, alpha=8, target_modules={"q_proj"})
 
@@ -280,18 +272,17 @@ def test_load_lora_on_already_injected():
             if isinstance(m, LoRALinear):
                 m.lora_B.fill_(0.5)
 
-    with tempfile.TemporaryDirectory() as tmpdir:
+    save_lora(model, temp_dir, LoRAConfig(r=4, alpha=8, target_modules=("q_proj",)))
-        save_lora(model, tmpdir, LoRAConfig(r=4, alpha=8, target_modules=("q_proj",)))
 
-        model2 = _make_model()
+    model2 = _make_model()
-        model2.load_state_dict(model.state_dict(), strict=False)
+    model2.load_state_dict(model.state_dict(), strict=False)
-        inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"})
+    inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"})
 
-        load_lora(model2, tmpdir)
+    load_lora(model2, temp_dir)
-        assert _get_lora_count(model2) > 0
+    assert _get_lora_count(model2) > 0
 
 
-def test_load_lora_mismatched_r_raises():
+def test_load_lora_mismatched_r_raises(temp_dir):
     model = _make_model()
     cfg = inject_lora(model, r=8, alpha=16, target_modules={"q_proj"})
 
@@ -300,15 +291,14 @@ def test_load_lora_mismatched_r_raises():
             if isinstance(m, LoRALinear):
                 m.lora_B.fill_(0.5)
 
-    with tempfile.TemporaryDirectory() as tmpdir:
+    save_lora(model, temp_dir, cfg)
-        save_lora(model, tmpdir, cfg)
 
-        model2 = _make_model()
+    model2 = _make_model()
-        model2.load_state_dict(model.state_dict(), strict=False)
+    model2.load_state_dict(model.state_dict(), strict=False)
-        inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"})
+    inject_lora(model2, r=4, alpha=8, target_modules={"q_proj"})
 
-        with pytest.raises(RuntimeError, match="size mismatch"):
+    with pytest.raises(RuntimeError, match="size mismatch"):
-            load_lora(model2, tmpdir)
+        load_lora(model2, temp_dir)
 
 
 def test_merge_preserves_output():

tests/module/test_model_forward.py

@@ -39,6 +39,21 @@ def _make_model(config=None) -> AutoRegressiveLM:
     return AutoRegressiveLM(config)
 
 
+def _make_batch(config, batch_size=2, seq_len=8, with_extra=False):
+    """Build a random token batch, optionally with position ids and loss mask."""
+    vocab = config.vocab_size
+    batch = {
+        "input_ids": torch.randint(0, vocab, (batch_size, seq_len)),
+        "target_ids": torch.randint(0, vocab, (batch_size, seq_len)),
+    }
+    if with_extra:
+        batch["position_ids"] = (
+            torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1)
+        )
+        batch["loss_mask"] = torch.ones(batch_size, seq_len, dtype=torch.bool)
+    return batch
+
+
 def test_model_forward_contract_uses_dense_training_and_packed_inference():
     from astrai.inference.cache import PagePool, TaskCacheManager
     from astrai.inference.workspace import InferenceWorkspace
@@ -180,13 +195,6 @@ class TestSEQStrategyMoE:
         self.model = _make_model(self.config).to(device)
         self.model.train()
 
-    def _make_batch(self, batch_size=2, seq_len=8):
-        vocab = self.config.vocab_size
-        input_ids = torch.randint(0, vocab, (batch_size, seq_len))
-        # target = input shifted right
-        target_ids = torch.randint(0, vocab, (batch_size, seq_len))
-        return {"input_ids": input_ids, "target_ids": target_ids}
-
     def test_compute_loss_returns_scalar(self):
         """compute_loss should return a scalar tensor."""
         strategy = SEQStrategy(
@@ -194,7 +202,7 @@ class TestSEQStrategyMoE:
             self.device,
             moe_aux_loss_coef=0.01,
         )
-        loss = strategy.compute_loss(self._make_batch())
+        loss = strategy.compute_loss(_make_batch(self.config))
         assert loss.ndim == 0
         assert loss.requires_grad
 
@@ -205,7 +213,7 @@ class TestSEQStrategyMoE:
             self.device,
             moe_aux_loss_coef=0.01,
         )
-        output = strategy.compute_loss_output(self._make_batch())
+        output = strategy.compute_loss_output(_make_batch(self.config))
 
         assert "loss" in output
         assert "metrics" in output
@@ -225,7 +233,7 @@ class TestSEQStrategyMoE:
             self.device,
             moe_aux_loss_coef=0.01,
         )
-        strategy.compute_loss_output(self._make_batch())
+        strategy.compute_loss_output(_make_batch(self.config))
 
         moe_metrics = strategy._moe_metrics
         assert moe_metrics, "_moe_metrics should not be empty for MoE model"
@@ -246,7 +254,7 @@ class TestSEQStrategyMoE:
             self.device,
             moe_aux_loss_coef=0.0,
         )
-        output = strategy.compute_loss_output(self._make_batch())
+        output = strategy.compute_loss_output(_make_batch(self.config))
         metrics = output["metrics"]
 
         # task_loss and loss should be equal (aux weighted by zero)
@@ -268,7 +276,7 @@ class TestSEQStrategyMoE:
             self.device,
             moe_aux_loss_coef=0.01,
         )
-        output = strategy.compute_loss_output(self._make_batch())
+        output = strategy.compute_loss_output(_make_batch(self.config))
         assert output["metrics"]["loss"] > output["metrics"]["task_loss"] + 1e-12
 
     def test_factory_creates_strategy_with_coef(self):
@@ -294,7 +302,7 @@ class TestSEQStrategyMoE:
             self.device,
             moe_aux_loss_coef=0.01,
         )
-        output = strategy.compute_loss_output(self._make_batch())
+        output = strategy.compute_loss_output(_make_batch(self.config))
         metrics = output["metrics"]
 
         assert "moe_aux_loss" not in metrics
@@ -313,19 +321,6 @@ class TestSFTStrategyMoE:
         self.model = _make_model(self.config).to(device)
         self.model.train()
 
-    def _make_batch(self, batch_size=2, seq_len=8):
-        vocab = self.config.vocab_size
-        input_ids = torch.randint(0, vocab, (batch_size, seq_len))
-        target_ids = torch.randint(0, vocab, (batch_size, seq_len))
-        position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1)
-        loss_mask = torch.ones(batch_size, seq_len, dtype=torch.bool)
-        return {
-            "input_ids": input_ids,
-            "target_ids": target_ids,
-            "position_ids": position_ids,
-            "loss_mask": loss_mask,
-        }
-
     def test_compute_loss_output_with_aux_loss(self):
         """SFTStrategy produces MoE metrics when coef > 0."""
         strategy = SFTStrategy(
@@ -333,7 +328,7 @@ class TestSFTStrategyMoE:
             self.device,
             moe_aux_loss_coef=0.01,
         )
-        output = strategy.compute_loss_output(self._make_batch())
+        output = strategy.compute_loss_output(_make_batch(self.config, with_extra=True))
 
         metrics = output["metrics"]
         assert "moe_aux_loss" in metrics
@@ -351,7 +346,7 @@ class TestSFTStrategyMoE:
             self.device,
             moe_aux_loss_coef=0.0,
         )
-        output = strategy.compute_loss_output(self._make_batch())
+        output = strategy.compute_loss_output(_make_batch(self.config, with_extra=True))
         metrics = output["metrics"]
 
         assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)

tests/trainer/test_grpo_e2e.py

@@ -1,123 +0,0 @@
-"""End-to-end integration test for online GRPO rollout."""
-
-import os
-from functools import partial
-
-import pytest
-import torch
-from torch.utils.data import Dataset
-
-from astrai.config import TrainConfig
-from astrai.model.transformer import AutoRegressiveLM
-from astrai.trainer.rollout import BaseRewardModel
-from astrai.trainer.schedule import SchedulerFactory
-from astrai.trainer.trainer import Trainer
-from tests.helpers import CHAT_TEMPLATE
-
-
-class InstructionDataset(Dataset):
-    """Toy instruction/input dataset for online RL rollout.
-
-    Each sample has an ``instruction`` and an optional ``input``; the
-    RolloutGenerator renders both through the tokenizer's chat template
-    so the prompt matches the SFT-trained format.
-    """
-
-    _SAMPLES = [
-        {"instruction": "Hello", "input": ""},
-        {"instruction": "Tell me a story", "input": "about dragons"},
-        {"instruction": "Summarize", "input": "the article"},
-        {"instruction": "Translate", "input": "to French: hi"},
-    ]
-
-    def __len__(self):
-        return len(self._SAMPLES)
-
-    def __getitem__(self, idx):
-        return dict(self._SAMPLES[idx])
-
-
-class LengthRewardModel(BaseRewardModel):
-    """Rewards each response by its (non-pad) token count.
-
-    Gives the group-normalized advantage a non-degenerate signal.
-    """
-
-    def score(self, prompts, responses):
-        B = len(prompts)
-        G = len(responses[0]) if B else 0
-        rewards = torch.zeros(B, G)
-        for i in range(B):
-            for g in range(G):
-                rewards[i, g] = float(len(responses[i][g]))
-        return rewards
-
-
-def instruction_collate_fn(batch):
-    """Stack a list of instruction/input dicts into a batch dict of lists."""
-    return {
-        "instruction": [b["instruction"] for b in batch],
-        "input": [b.get("input", "") for b in batch],
-    }
-
-
-def _model_fn(model_config):
-    return AutoRegressiveLM(model_config).to(dtype=torch.float32)
-
-
-def _optimizer_fn(m):
-    return torch.optim.AdamW(m.parameters(), lr=1e-4)
-
-
-def _scheduler_fn(optim):
-    return SchedulerFactory.create(
-        "cosine", optim, warmup_steps=1, lr_decay_steps=4, min_rate=0.05
-    )
-
-
-@pytest.mark.integration
-def test_online_grpo_end_to_end(base_test_env):
-    """Run one epoch of online GRPO with KV-cache-backed rollout."""
-    test_dir = base_test_env["test_dir"]
-    device = base_test_env["device"]
-    tokenizer = base_test_env["tokenizer"]
-    model_config = base_test_env["transformer_config"]
-
-    tokenizer.set_chat_template(CHAT_TEMPLATE)
-    tokenizer.save_pretrained(test_dir)
-
-    model_fn = partial(_model_fn, model_config)
-    optimizer_fn = _optimizer_fn
-    scheduler_fn = _scheduler_fn
-
-    dataset = InstructionDataset()
-
-    train_config = TrainConfig(
-        strategy="online_grpo",
-        model_fn=model_fn,
-        dataset=dataset,
-        optimizer_fn=optimizer_fn,
-        scheduler_fn=scheduler_fn,
-        ckpt_dir=os.path.join(test_dir, "ckpt"),
-        n_epoch=1,
-        batch_per_device=2,
-        ckpt_interval=100,
-        grad_accum_steps=1,
-        random_seed=42,
-        device_type=device,
-        nprocs=1,
-        parallel_mode="none",
-        strategy_kwargs={"clip_eps": 0.2, "kl_coef": 0.01, "group_size": 2},
-        rollout_interval=1,
-        rollout_temperature=1.0,
-        rollout_top_k=0,
-        rollout_top_p=1.0,
-        rollout_max_tokens=4,
-        reward_model_fn=LengthRewardModel,
-        collate_fn=instruction_collate_fn,
-    )
-
-    trainer = Trainer(train_config)
-    trainer.train(param_path=test_dir)
-
-    assert os.path.isdir(os.path.join(test_dir, "ckpt"))

tests/trainer/test_online_e2e.py

@@ -1,4 +1,4 @@
-"""End-to-end integration test for online DPO rollout."""
+"""End-to-end integration tests for online GRPO/DPO rollout."""
 
 import os
 from functools import partial
@@ -40,7 +40,7 @@ class InstructionDataset(Dataset):
 class LengthRewardModel(BaseRewardModel):
     """Rewards each response by its (non-pad) token count.
 
-    Enough for DPO to distinguish chosen/rejected from the rollout group.
+    Gives the group-normalized advantage a non-degenerate signal.
     """
 
     def score(self, prompts, responses):
@@ -75,31 +75,34 @@ def _scheduler_fn(optim):
     )
 
 
+_ONLINE_STRATEGIES = [
+    pytest.param(
+        "online_grpo",
+        {"clip_eps": 0.2, "kl_coef": 0.01, "group_size": 2},
+        id="grpo",
+    ),
+    pytest.param("online_dpo", {"beta": 0.1, "group_size": 2}, id="dpo"),
+]
+
+
 @pytest.mark.integration
-def test_online_dpo_end_to_end(base_test_env):
+@pytest.mark.parametrize(("strategy", "strategy_kwargs"), _ONLINE_STRATEGIES)
-    """Run one epoch of online DPO with KV-cache-backed rollout."""
+def test_online_rollout_end_to_end(base_test_env, strategy, strategy_kwargs):
+    """Run one epoch of online RL rollout with KV-cache-backed generation."""
     test_dir = base_test_env["test_dir"]
     device = base_test_env["device"]
     tokenizer = base_test_env["tokenizer"]
     model_config = base_test_env["transformer_config"]
 
-    # Equip tokenizer with a chat template so RolloutGenerator can
-    # render instruction/input via apply_chat_template.
     tokenizer.set_chat_template(CHAT_TEMPLATE)
     tokenizer.save_pretrained(test_dir)
 
-    model_fn = partial(_model_fn, model_config)
-    optimizer_fn = _optimizer_fn
-    scheduler_fn = _scheduler_fn
-
-    dataset = InstructionDataset()
-
     train_config = TrainConfig(
-        strategy="online_dpo",
+        strategy=strategy,
-        model_fn=model_fn,
+        model_fn=partial(_model_fn, model_config),
-        dataset=dataset,
+        dataset=InstructionDataset(),
-        optimizer_fn=optimizer_fn,
+        optimizer_fn=_optimizer_fn,
-        scheduler_fn=scheduler_fn,
+        scheduler_fn=_scheduler_fn,
         ckpt_dir=os.path.join(test_dir, "ckpt"),
         n_epoch=1,
         batch_per_device=2,
@@ -109,7 +112,7 @@ def test_online_dpo_end_to_end(base_test_env):
         device_type=device,
         nprocs=1,
         parallel_mode="none",
-        strategy_kwargs={"beta": 0.1, "group_size": 2},
+        strategy_kwargs=strategy_kwargs,
         rollout_interval=1,
         rollout_temperature=1.0,
         rollout_top_k=0,