tests/test_artifacts.py

from pathlib import Path
import tempfile
import unittest
from unittest import mock


class ArtifactTests(unittest.TestCase):
    def test_checkpoint_path_rejects_absolute_parent_and_empty_paths(self):
        import artifacts

        for value in ("/tmp/model", "../model", "a/../../model", ""):
            with self.subTest(value=value), self.assertRaises(ValueError):
                artifacts.normalize_checkpoint_path(value)

    def test_download_checkpoint_validates_required_layout(self):
        import artifacts

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            checkpoint = root / "checkpoints/30000"
            (checkpoint / "params").mkdir(parents=True)
            (checkpoint / "assets/ur_demo").mkdir(parents=True)
            (checkpoint / "assets/ur_demo/norm_stats.json").write_text("{}")
            with mock.patch.object(artifacts, "snapshot_download", return_value=str(root)):
                paths = artifacts.download_checkpoint("owner/model", "checkpoints/30000")
            self.assertEqual(paths.checkpoint, checkpoint)
            self.assertEqual(paths.norm_stats, checkpoint / "assets/ur_demo/norm_stats.json")

    def test_download_checkpoint_accepts_named_asset_directory(self):
        import artifacts

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            checkpoint = root / "3000"
            (checkpoint / "params/ocdbt.process_0/d").mkdir(parents=True)
            stats = checkpoint / "assets/F-Fer/ur-1/norm_stats.json"
            stats.parent.mkdir(parents=True)
            stats.write_text("{}")
            with mock.patch.object(artifacts, "snapshot_download", return_value=str(root)):
                paths = artifacts.download_checkpoint("owner/model", "3000")
            self.assertEqual(paths.norm_stats, stats)

    def test_download_checkpoint_reports_missing_weights(self):
        import artifacts

        with tempfile.TemporaryDirectory() as directory:
            with mock.patch.object(artifacts, "snapshot_download", return_value=directory):
                with self.assertRaisesRegex(FileNotFoundError, "params/|model.safetensors"):
                    artifacts.download_checkpoint("owner/model", "checkpoint")


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