72 lines
4.3 KiB
Python
72 lines
4.3 KiB
Python
import hashlib
|
|
import io
|
|
import json
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from paper_atlas.io.snapshots import download
|
|
|
|
|
|
class Response(io.BytesIO):
|
|
def __init__(self,body,status=200,headers=None):
|
|
super().__init__(body);self.status=status;self.headers=headers or {'Content-Length':str(len(body))}
|
|
|
|
|
|
class DownloadTests(unittest.TestCase):
|
|
def test_verified_download_is_reused_without_replacing_it(self):
|
|
body=b'public metadata';sha=hashlib.sha256(body).hexdigest()
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
path=Path(tmp)/'snapshot.gz'
|
|
with patch('paper_atlas.io.snapshots.urllib.request.urlopen',return_value=Response(body)):
|
|
result=download('https://example.org/snapshot.gz',path,sha256=sha,expected_bytes=len(body),min_free=0)
|
|
self.assertEqual(result['status'],'verified')
|
|
with patch('paper_atlas.io.snapshots.urllib.request.urlopen',side_effect=AssertionError('No network expected')):
|
|
reused=download('https://example.org/snapshot.gz',path,sha256=sha,min_free=0)
|
|
self.assertEqual(reused['status'],'already_verified')
|
|
self.assertEqual(json.loads(Path(str(path)+'.manifest.json').read_text())['sha256'],sha)
|
|
|
|
def test_bad_checksum_never_publishes_a_snapshot(self):
|
|
with tempfile.TemporaryDirectory() as tmp,patch('paper_atlas.io.snapshots.urllib.request.urlopen',return_value=Response(b'wrong bytes')):
|
|
path=Path(tmp)/'snapshot.gz'
|
|
with self.assertRaisesRegex(ValueError,'SHA256 mismatch'):
|
|
download('https://example.org/snapshot.gz',path,sha256='0'*64,min_free=0)
|
|
self.assertFalse(path.exists());self.assertTrue(Path(str(path)+'.part').exists())
|
|
|
|
def test_resume_requires_the_same_identity_and_a_real_partial_response(self):
|
|
body=b'abcdef';sha=hashlib.sha256(body).hexdigest();url='https://example.org/snapshot.gz'
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
path=Path(tmp)/'snapshot.gz';part=Path(str(path)+'.part');part.write_bytes(body[:3])
|
|
Path(str(part)+'.json').write_text(json.dumps({'url':url,'expected_sha256':sha,'expected_md5':None,'expected_bytes':6}))
|
|
with self.assertRaisesRegex(ValueError,'same snapshot'):
|
|
download(url,path,sha256=sha,expected_bytes=6,min_free=0)
|
|
with patch('paper_atlas.io.snapshots.urllib.request.urlopen',return_value=Response(body)):
|
|
with self.assertRaisesRegex(ValueError,'resume range'):
|
|
download(url,path,sha256=sha,expected_bytes=6,resume=True,min_free=0)
|
|
self.assertEqual(part.read_bytes(),b'abc')
|
|
with patch('paper_atlas.io.snapshots.urllib.request.urlopen',return_value=Response(b'def',206,{'Content-Range':'bytes 3-5/6'})):
|
|
download(url,path,sha256=sha,expected_bytes=6,resume=True,min_free=0)
|
|
self.assertEqual(path.read_bytes(),body)
|
|
|
|
def test_wrong_http_length_is_rejected_before_writing(self):
|
|
with tempfile.TemporaryDirectory() as tmp,patch('paper_atlas.io.snapshots.urllib.request.urlopen',return_value=Response(b'abc')):
|
|
path=Path(tmp)/'snapshot.gz'
|
|
with self.assertRaisesRegex(ValueError,'HTTP size'):
|
|
download('https://example.org/snapshot.gz',path,sha256='0'*64,expected_bytes=4,min_free=0)
|
|
self.assertFalse(path.exists());self.assertFalse(Path(str(path)+'.part').exists())
|
|
|
|
def test_empty_partial_and_missing_final_manifest_can_recover(self):
|
|
body=b'abc';sha=hashlib.sha256(body).hexdigest();url='https://example.org/snapshot.gz'
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
path=Path(tmp)/'snapshot.gz';part=Path(str(path)+'.part');part.write_bytes(b'')
|
|
Path(str(part)+'.json').write_text(json.dumps({'url':url,'expected_sha256':sha,'expected_md5':None,'expected_bytes':3}))
|
|
with patch('paper_atlas.io.snapshots.urllib.request.urlopen',return_value=Response(body)):
|
|
download(url,path,sha256=sha,expected_bytes=3,resume=True,min_free=0)
|
|
metadata=Path(str(path)+'.manifest.json');metadata.unlink()
|
|
download(url,path,sha256=sha,expected_bytes=3,min_free=0)
|
|
self.assertEqual(json.loads(metadata.read_text())['sha256'],sha)
|
|
|
|
|
|
if __name__=='__main__':unittest.main()
|