airfRANS-model-exploration/tests/test_training_config.py

41 lines
1.6 KiB
Python

from __future__ import annotations
import tempfile
import unittest
from pathlib import Path
from airfrans_frontier.training.config import load_training_config
class TrainingConfigTests(unittest.TestCase):
def test_config_loader_accepts_mlp_tiny(self) -> None:
config = load_training_config("configs/mlp_tiny.toml")
self.assertEqual(config.run.name, "mlp_tiny")
self.assertEqual(config.model.type, "mlp")
self.assertEqual(config.loss.type, "normalized_mse")
self.assertEqual(config.device.type, "cuda")
self.assertTrue(config.data.root.is_absolute())
self.assertEqual(config.data.source, "local")
self.assertIsNone(config.data.hf_repo_id)
self.assertIsNone(config.data.cache_dir)
def test_config_loader_accepts_huggingface_data_source(self) -> None:
config = load_training_config("configs/aggressive_smoke.toml")
self.assertEqual(config.data.source, "huggingface")
self.assertEqual(config.data.hf_repo_id, "zacheryasc/airfrans-processed")
self.assertEqual(config.data.hf_path_prefix, "processed/full")
self.assertTrue(config.data.cache_dir is not None)
def test_config_loader_rejects_missing_section(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
config_path = Path(tmp) / "bad.toml"
config_path.write_text("[run]\nname = 'bad'\n")
with self.assertRaisesRegex(ValueError, r"missing \[data\] section"):
load_training_config(config_path)
if __name__ == "__main__":
unittest.main()