cesnimda
36 lines
1.8 KiB
Python
36 lines
1.8 KiB
Python
import hashlib
|
|
import io
|
|
import json
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
from prepare import prepare
|
|
|
|
|
|
class PreparationBoundaryTests(unittest.TestCase):
|
|
def run_fixture(self, payload, expected=b'approved'):
|
|
revision='a4f8f3ea906ed274767e9906dbaede7531d660ff'
|
|
manifest={'weights-only':{'sourceRevision':revision,'license':'Apache-2.0 pinned card',
|
|
'files':[{'file':'config.json','url':'https://huggingface.co/sshleifer/distilbart-cnn-12-6/resolve/'+revision+'/config.json',
|
|
'bytes':len(expected),'sha256':hashlib.sha256(expected).hexdigest()}]}}
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
target=Path(directory)/'new-cache'
|
|
with patch.object(Path,'read_text',return_value=json.dumps(manifest)), patch('prepare.shutil.disk_usage',return_value=type('Usage',(),{'free':10*1024**3})()), patch('prepare.urllib.request.urlopen',return_value=io.BytesIO(payload)):
|
|
if payload==expected:
|
|
prepare(target)
|
|
self.assertEqual((target/'config.json').read_bytes(),expected)
|
|
with self.assertRaises(AssertionError):prepare(target)
|
|
self.assertEqual((target/'config.json').read_bytes(),expected)
|
|
else:
|
|
with self.assertRaises(AssertionError):prepare(target)
|
|
self.assertFalse((target/'config.json').exists())
|
|
self.assertTrue((target/'config.json.partial').exists())
|
|
|
|
def test_validated_bytes_and_existing_cache_preserved(self):self.run_fixture(b'approved')
|
|
def test_corrupt_bytes_never_promoted(self):self.run_fixture(b'rejected')
|
|
def test_oversized_bytes_never_promoted(self):self.run_fixture(b'oversized payload')
|
|
|
|
|
|
if __name__=='__main__':unittest.main()
|