File size: 2,883 Bytes
b7d1c87
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6455b57
 
 
 
 
 
 
b7d1c87
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
from __future__ import annotations
import hashlib,json,tempfile,unittest
from pathlib import Path
from unittest.mock import patch
from cache import RevisionCache

class CacheTests(unittest.TestCase):
 def test_cold_start_and_corrupt_cache_rebuild(self):
  with tempfile.TemporaryDirectory() as temporary:
   root=Path(temporary);revision="a"*40; bodies={"data/a.parquet":b"parquet-proof","indexes/search.sqlite":b"sqlite-proof"}
   manifest={"dataset_revision":revision,"artifacts":[{"storage":"dataset","path":path,"content_type":"application/vnd.apache.parquet" if path.endswith("parquet") else "application/vnd.sqlite3","size_bytes":len(body),"checksum":{"value":hashlib.sha256(body).hexdigest()}} for path,body in bodies.items()]}
   cache=RevisionCache(root/"cache","manifest","mzx/test")
   def bytes_for(url): return json.dumps(manifest).encode()
   def download(url,target):
    key=next(path for path in bodies if path in url);target.write_bytes(bodies[key]);return len(bodies[key]),hashlib.sha256(bodies[key]).hexdigest()
   with patch.object(cache,"_bytes",side_effect=bytes_for),patch.object(cache,"_download",side_effect=download):
    first=cache.bootstrap();self.assertTrue(first["ready"]);self.assertEqual(first["artifact_count"],2)
    (Path(first["cache_dir"])/"data/a.parquet").write_bytes(b"corrupt")
    second=cache.bootstrap();self.assertTrue(second["ready"]);self.assertEqual((Path(second["cache_dir"])/"data/a.parquet").read_bytes(),b"parquet-proof")
 def test_incomplete_download_never_becomes_ready(self):
  with tempfile.TemporaryDirectory() as temporary:
   revision="b"*40;body=b"expected";manifest={"dataset_revision":revision,"artifacts":[{"storage":"dataset","path":"data/a.parquet","content_type":"application/vnd.apache.parquet","size_bytes":len(body),"checksum":{"value":hashlib.sha256(body).hexdigest()}}]};cache=RevisionCache(Path(temporary),"manifest","mzx/test")
   with patch.object(cache,"_bytes",return_value=json.dumps(manifest).encode()),patch.object(cache,"_download",return_value=(3,"0"*64)):
    with self.assertRaises(ValueError):cache.bootstrap()
   self.assertFalse(cache.status()["ready"]);self.assertFalse(any(Path(temporary).glob("*/ready.json")))
 def test_failed_refresh_keeps_last_verified_revision_ready(self):
  with tempfile.TemporaryDirectory() as temporary:
   cache=RevisionCache(Path(temporary),"manifest","mzx/test");revision="c"*40;active=Path(temporary)/revision;active.mkdir()
   cache._status={"status":"ready","ready":True,"dataset_revision":revision,"artifact_count":2,"cached_bytes":12,"cache_dir":str(active),"error":None}
   with patch.object(cache,"_bytes",side_effect=OSError("temporary network failure")):
    with self.assertRaises(OSError):cache.bootstrap()
   self.assertEqual(cache.status()["dataset_revision"],revision);self.assertTrue(cache.status()["ready"])
if __name__=="__main__":unittest.main()