| import hashlib |
| import json |
| from pathlib import Path |
| import sys |
| import tempfile |
| import unittest |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| from attach_model_bundle import install, plan_install |
|
|
|
|
| class InstallerTests(unittest.TestCase): |
| def setUp(self): |
| self.temporary = tempfile.TemporaryDirectory() |
| self.addCleanup(self.temporary.cleanup) |
| root = Path(self.temporary.name) |
| self.bundle, self.code = root / 'bundle', root / 'code' |
| for path in (self.bundle/'models/base', self.bundle/'checkpoints/hoi', self.bundle/'code_support', self.code): |
| path.mkdir(parents=True, exist_ok=True) |
| (self.code/'project_paths.py').write_text('existing code root\n') |
| (self.code/'loader.py').write_text('old known version\n') |
| (self.bundle/'code_support/loader.py').write_text('new shard-aware version\n') |
| old = hashlib.sha256((self.code/'loader.py').read_bytes()).hexdigest() |
| new = hashlib.sha256((self.bundle/'code_support/loader.py').read_bytes()).hexdigest() |
| (self.bundle/'code_support/loader_patch.json').write_text(json.dumps( |
| {'files': [{'path': 'loader.py', 'previous_sha256': old, 'new_sha256': new}]})) |
| (self.bundle/'bundle_manifest.json').write_text('{"status":"complete"}') |
|
|
| def test_dry_run_makes_no_changes(self): |
| operations = install(self.bundle, self.code) |
| self.assertEqual(len(operations), 3) |
| self.assertEqual((self.code/'loader.py').read_text(), 'old known version\n') |
| self.assertFalse((self.code/'models').exists()) |
|
|
| def test_apply_and_idempotent(self): |
| install(self.bundle, self.code, apply=True) |
| self.assertEqual((self.code/'models/base').resolve(), self.bundle/'models/base') |
| self.assertEqual((self.code/'checkpoints/hoi').resolve(), self.bundle/'checkpoints/hoi') |
| self.assertEqual((self.code/'loader.py').read_text(), 'new shard-aware version\n') |
| self.assertEqual(plan_install(self.bundle, self.code), []) |
|
|
| def test_local_code_changes_rejected_before_any_write(self): |
| (self.code/'loader.py').write_text('important local edits\n') |
| with self.assertRaisesRegex(ValueError, 'local code changes'): |
| install(self.bundle, self.code, apply=True) |
| self.assertFalse((self.code/'models').exists()) |
| self.assertEqual((self.code/'loader.py').read_text(), 'important local edits\n') |
|
|
| def test_existing_weight_directory_preserved(self): |
| (self.code/'models/base').mkdir(parents=True) |
| (self.code/'models/base/my_weight').write_text('keep') |
| with self.assertRaisesRegex(ValueError, 'not be overwritten'): |
| install(self.bundle, self.code, apply=True) |
| self.assertEqual((self.code/'models/base/my_weight').read_text(), 'keep') |
|
|
| def test_unsafe_patch_path_rejected(self): |
| path = self.bundle/'code_support/loader_patch.json' |
| value = json.loads(path.read_text()) |
| value['files'][0]['path'] = '../outside.py' |
| path.write_text(json.dumps(value)) |
| with self.assertRaisesRegex(ValueError, 'unsafe'): |
| plan_install(self.bundle, self.code) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
