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) 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()