tests/test_data_and_device.py

import os
from pathlib import Path
import tempfile
import unittest
from unittest.mock import patch

import numpy as np

from dropout_decay.data import load_cached_splits
from dropout_decay.experiments.device import assert_mps_only


class DataAndDeviceTests(unittest.TestCase):
    def test_load_cached_splits_uses_expected_train_val_partition(self):
        with tempfile.TemporaryDirectory() as tmp:
            cache_dir = Path(tmp)
            (cache_dir / "tokenizer-v16.json").write_text("{}", encoding="utf-8")
            np.save(cache_dir / "tokens-v16-uint16.npy", np.arange(100, dtype=np.uint16))

            tokenizer, splits = load_cached_splits(
                cache_dir=cache_dir,
                vocab_size=16,
                max_required_train_tokens=80,
                val_tokens=10,
                allow_short_corpus=False,
            )

            self.assertEqual(tokenizer.vocab_size, 16)
            self.assertEqual(len(splits.train), 90)
            self.assertEqual(len(splits.val), 10)
            self.assertEqual(splits.tokenizer_path, cache_dir / "tokenizer-v16.json")

    def test_load_cached_splits_rejects_missing_cache(self):
        with tempfile.TemporaryDirectory() as tmp:
            with self.assertRaises(FileNotFoundError):
                load_cached_splits(
                    cache_dir=Path(tmp),
                    vocab_size=16,
                    max_required_train_tokens=10,
                    val_tokens=10,
                    allow_short_corpus=False,
                )

    def test_mps_guard_rejects_fallback_before_backend_checks(self):
        env = dict(os.environ)
        env["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
        with patch.dict(os.environ, env, clear=True):
            with self.assertRaises(SystemExit) as caught:
                assert_mps_only()

        self.assertIn("MPS_FALLBACK", str(caught.exception))


if __name__ == "__main__":
    unittest.main()