File size: 3,223 Bytes
c28ffff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | 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()
|