From 753f4c7fab565234e455f19e3ca123e44171e042 Mon Sep 17 00:00:00 2001 From: Willem Melching Date: Fri, 29 Jul 2022 17:46:19 +0200 Subject: [PATCH] casync unittests (#25306) * simple extract test * add test with loopback * test for chunk reuse * test seeding * add instructions * cleanup * install casync * make more tests work without a loopback --- system/hardware/tici/casync.py | 33 +++-- system/hardware/tici/tests/__init__.py | 0 system/hardware/tici/tests/test_casync.py | 153 ++++++++++++++++++++++ tools/ubuntu_setup.sh | 1 + 4 files changed, 175 insertions(+), 12 deletions(-) create mode 100644 system/hardware/tici/tests/__init__.py create mode 100755 system/hardware/tici/tests/test_casync.py diff --git a/system/hardware/tici/casync.py b/system/hardware/tici/casync.py index 9dff64239e..8ae42fa714 100755 --- a/system/hardware/tici/casync.py +++ b/system/hardware/tici/casync.py @@ -40,10 +40,12 @@ class ChunkReader(ABC): class FileChunkReader(ChunkReader): """Reads chunks from a local file""" def __init__(self, fn: str) -> None: - super().__init__() self.f = open(fn, 'rb') + def __del__(self): + self.f.close() + def read(self, chunk: Chunk) -> bytes: self.f.seek(chunk.offset) return self.f.read(chunk.length) @@ -61,19 +63,24 @@ class RemoteChunkReader(ChunkReader): sha_hex = chunk.sha.hex() url = os.path.join(self.url, sha_hex[:4], sha_hex + ".cacnk") - for i in range(CHUNK_DOWNLOAD_RETRIES): - try: - resp = self.session.get(url, timeout=CHUNK_DOWNLOAD_TIMEOUT) - break - except Exception: - if i == CHUNK_DOWNLOAD_RETRIES - 1: - raise - time.sleep(CHUNK_DOWNLOAD_TIMEOUT) + if os.path.isfile(url): + with open(url, 'rb') as f: + contents = f.read() + else: + for i in range(CHUNK_DOWNLOAD_RETRIES): + try: + resp = self.session.get(url, timeout=CHUNK_DOWNLOAD_TIMEOUT) + break + except Exception: + if i == CHUNK_DOWNLOAD_RETRIES - 1: + raise + time.sleep(CHUNK_DOWNLOAD_TIMEOUT) - resp.raise_for_status() + resp.raise_for_status() + contents = resp.content decompressor = lzma.LZMADecompressor(format=lzma.FORMAT_AUTO) - return decompressor.decompress(resp.content) + return decompressor.decompress(contents) def parse_caibx(caibx_path: str) -> List[Chunk]: @@ -120,6 +127,7 @@ def parse_caibx(caibx_path: str) -> List[Chunk]: chunks.append(Chunk(sha, offset, length)) offset = new_offset + caibx.close() return chunks @@ -139,7 +147,8 @@ def extract(target: List[Chunk], progress: Optional[Callable[[int], None]] = None): stats: Dict[str, int] = defaultdict(int) - with open(out_path, 'wb') as out: + mode = 'rb+' if os.path.exists(out_path) else 'wb' + with open(out_path, mode) as out: for cur_chunk in target: # Find source for desired chunk diff --git a/system/hardware/tici/tests/__init__.py b/system/hardware/tici/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/system/hardware/tici/tests/test_casync.py b/system/hardware/tici/tests/test_casync.py new file mode 100755 index 0000000000..8724575ad6 --- /dev/null +++ b/system/hardware/tici/tests/test_casync.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +import os +import unittest +import tempfile +import subprocess + +import system.hardware.tici.casync as casync + +# dd if=/dev/zero of=/tmp/img.raw bs=1M count=2 +# sudo losetup -f /tmp/img.raw +# losetup -a | grep img.raw +LOOPBACK = os.environ.get('LOOPBACK', None) + + +class TestCasync(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.tmpdir = tempfile.TemporaryDirectory() + + # Build example contents + chunk_a = [i % 256 for i in range(1024)] * 512 + chunk_b = [(256 - i) % 256 for i in range(1024)] * 512 + zeroes = [0] * (1024 * 128) + contents = chunk_a + chunk_b + zeroes + chunk_a + + cls.contents = bytes(contents) + + # Write to file + cls.orig_fn = os.path.join(cls.tmpdir.name, 'orig.bin') + with open(cls.orig_fn, 'wb') as f: + f.write(cls.contents) + + # Create casync files + cls.manifest_fn = os.path.join(cls.tmpdir.name, 'orig.caibx') + cls.store_fn = os.path.join(cls.tmpdir.name, 'store') + subprocess.check_output(["casync", "make", "--compression=xz", "--store", cls.store_fn, cls.manifest_fn, cls.orig_fn]) + + target = casync.parse_caibx(cls.manifest_fn) + hashes = [c.sha.hex() for c in target] + + # Ensure we have chunk reuse + assert len(hashes) > len(set(hashes)) + + def setUp(self): + # Clear target_lo + if LOOPBACK is not None: + self.target_lo = LOOPBACK + with open(self.target_lo, 'wb') as f: + f.write(b"0" * len(self.contents)) + + self.target_fn = os.path.join(self.tmpdir.name, next(tempfile._get_candidate_names())) + self.seed_fn = os.path.join(self.tmpdir.name, next(tempfile._get_candidate_names())) + + def tearDown(self): + for fn in [self.target_fn, self.seed_fn]: + try: + os.unlink(fn) + except FileNotFoundError: + pass + + def test_simple_extract(self): + target = casync.parse_caibx(self.manifest_fn) + + sources = [('remote', casync.RemoteChunkReader(self.store_fn), casync.build_chunk_dict(target))] + stats = casync.extract(target, sources, self.target_fn) + + with open(self.target_fn, 'rb') as target_f: + self.assertEqual(target_f.read(), self.contents) + + self.assertEqual(stats['remote'], len(self.contents)) + + def test_seed(self): + target = casync.parse_caibx(self.manifest_fn) + + # Populate seed with half of the target contents + with open(self.seed_fn, 'wb') as seed_f: + seed_f.write(self.contents[:len(self.contents) // 2]) + + sources = [('seed', casync.FileChunkReader(self.seed_fn), casync.build_chunk_dict(target))] + sources += [('remote', casync.RemoteChunkReader(self.store_fn), casync.build_chunk_dict(target))] + stats = casync.extract(target, sources, self.target_fn) + + with open(self.target_fn, 'rb') as target_f: + self.assertEqual(target_f.read(), self.contents) + + self.assertGreater(stats['seed'], 0) + self.assertLess(stats['remote'], len(self.contents)) + + def test_already_done(self): + """Test that an already flashed target doesn't download any chunks""" + target = casync.parse_caibx(self.manifest_fn) + + with open(self.target_fn, 'wb') as f: + f.write(self.contents) + + sources = [('target', casync.FileChunkReader(self.target_fn), casync.build_chunk_dict(target))] + sources += [('remote', casync.RemoteChunkReader(self.store_fn), casync.build_chunk_dict(target))] + + stats = casync.extract(target, sources, self.target_fn) + + with open(self.target_fn, 'rb') as f: + self.assertEqual(f.read(), self.contents) + + self.assertEqual(stats['target'], len(self.contents)) + + def test_chunk_reuse(self): + """Test that chunks that are reused are only downloaded once""" + target = casync.parse_caibx(self.manifest_fn) + + # Ensure target exists + with open(self.target_fn, 'wb'): + pass + + sources = [('target', casync.FileChunkReader(self.target_fn), casync.build_chunk_dict(target))] + sources += [('remote', casync.RemoteChunkReader(self.store_fn), casync.build_chunk_dict(target))] + + stats = casync.extract(target, sources, self.target_fn) + + with open(self.target_fn, 'rb') as f: + self.assertEqual(f.read(), self.contents) + + self.assertLess(stats['remote'], len(self.contents)) + + @unittest.skipUnless(LOOPBACK, "requires loopback device") + def test_lo_simple_extract(self): + target = casync.parse_caibx(self.manifest_fn) + sources = [('remote', casync.RemoteChunkReader(self.store_fn), casync.build_chunk_dict(target))] + + stats = casync.extract(target, sources, self.target_lo) + + with open(self.target_lo, 'rb') as target_f: + self.assertEqual(target_f.read(len(self.contents)), self.contents) + + self.assertEqual(stats['remote'], len(self.contents)) + + @unittest.skipUnless(LOOPBACK, "requires loopback device") + def test_lo_chunk_reuse(self): + """Test that chunks that are reused are only downloaded once""" + target = casync.parse_caibx(self.manifest_fn) + + sources = [('target', casync.FileChunkReader(self.target_lo), casync.build_chunk_dict(target))] + sources += [('remote', casync.RemoteChunkReader(self.store_fn), casync.build_chunk_dict(target))] + + stats = casync.extract(target, sources, self.target_lo) + + with open(self.target_lo, 'rb') as f: + self.assertEqual(f.read(len(self.contents)), self.contents) + + self.assertLess(stats['remote'], len(self.contents)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/ubuntu_setup.sh b/tools/ubuntu_setup.sh index bd12b878f1..379df37adb 100755 --- a/tools/ubuntu_setup.sh +++ b/tools/ubuntu_setup.sh @@ -14,6 +14,7 @@ function install_ubuntu_common_requirements() { autoconf \ build-essential \ ca-certificates \ + casync \ clang \ cmake \ make \