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