refactor: eliminate test duplication via shared helpers

- Add tests/helpers.py with shared config, dataset, tokenizer, executor, and assertion helpers
- Replace 15 copies of device one-liner with session-scoped fixture
- Collapse 5 near-identical Dataset subclasses into RandomTokenDataset
- Remove duplicate _make_config/_make_model/_make_frozen and FakeTokenizer/FakeExecutor definitions
- Make test_callbacks and test_early_stopping use existing train_config_factory
- Replace 6 duplicate meta.json read blocks with load_shard_meta
- Fix mkdtemp leaks in test_lora.py with TemporaryDirectory

tests/module/test_tie_weight.py

@@ -1,6 +1,5 @@
 import json
 import os
-import tempfile
 
 import pytest
 import safetensors.torch as st
@@ -8,43 +7,13 @@ import torch
 
 from astrai.config.model_config import AutoRegressiveLMConfig
 from astrai.model.transformer import AutoRegressiveLM
+from tests.helpers import TINY_CONFIG
 
 
-@pytest.fixture
-def transformer_test_env():
-    test_dir = tempfile.mkdtemp(prefix="transformer_test_")
-    config_path = os.path.join(test_dir, "config.json")
+def test_tie_weight_init(base_test_env):
+    config_path = base_test_env["config_path"]
 
-    config = {
-        "vocab_size": 1000,
-        "hidden_size": 8,
-        "num_attention_heads": 2,
-        "num_key_value_heads": 1,
-        "intermediate_size": 16,
-        "max_position_embeddings": 64,
-        "num_hidden_layers": 2,
-        "rms_norm_eps": 1e-5,
-    }
-
-    with open(config_path, "w") as f:
-        json.dump(config, f)
-
-    yield {"test_dir": test_dir, "config_path": config_path, "config": config}
-
-    if os.path.exists(test_dir):
-        try:
-            for file in os.listdir(test_dir):
-                os.remove(os.path.join(test_dir, file))
-            os.rmdir(test_dir)
-        except Exception:
-            pass
-
-
-def test_tie_weight_init(transformer_test_env):
-    config_path = transformer_test_env["config_path"]
-    config_data = transformer_test_env["config"].copy()
-
-    # case 1: tie weight
+    config_data = TINY_CONFIG.copy()
     config_data["tie_word_embeddings"] = True
 
     with open(config_path, "w") as f:
@@ -62,7 +31,6 @@ def test_tie_weight_init(transformer_test_env):
     assert torch.equal(model.lm_head.weight, model.embed_tokens.weight)
     assert not torch.equal(model.lm_head.weight, original_weight)
 
-    # case 2: not tie weight
     config_data["tie_word_embeddings"] = False
 
     with open(config_path, "w") as f:
@@ -81,13 +49,11 @@ def test_tie_weight_init(transformer_test_env):
     assert not torch.equal(model.lm_head.weight, original_weight)
 
 
-def test_model_save_load_with_tie_weight(transformer_test_env):
-    test_dir = transformer_test_env["test_dir"]
+def test_model_save_load_with_tie_weight(base_test_env):
+    test_dir = base_test_env["test_dir"]
     model_path = os.path.join(test_dir, "model.safetensors")
 
-    config_data = transformer_test_env["config"].copy()
-
-    # case 1: tie weight
+    config_data = TINY_CONFIG.copy()
     config_data["tie_word_embeddings"] = True
     config_path = os.path.join(test_dir, "config.json")
 
@@ -107,7 +73,6 @@ def test_model_save_load_with_tie_weight(transformer_test_env):
     assert model.lm_head.weight.data_ptr() == model.embed_tokens.weight.data_ptr()
     assert "lm_head.weight" not in model.state_dict()
 
-    # case 2: not tie weight (form tie-weight state dict load)
     config_data["tie_word_embeddings"] = False
     with open(config_path, "w") as f:
         json.dump(config_data, f)