import hashlib import os import tempfile import unittest from pathlib import Path from unittest import mock from uuid import uuid4 from archive_clients.state import ClientStore from archive_clients.transfer import ( FileMaterializer, TransferIntegrityError, canonical_message_json, cleanup_partial_transfer, load_published_transfer, materialize_transfer, stage_transfer, ) from archive_control.v1 import transfer_pb2 class TransferHappyPathTests(unittest.TestCase): def setUp(self): self.temporary = tempfile.TemporaryDirectory() self.root = Path(self.temporary.name) self.source = self.root / "source" self.sync = self.root / "sync" self.target = self.root / "target" self.source.mkdir() self.sync.mkdir() self.target.mkdir() (self.source / "album").mkdir() (self.source / "album" / "one.bin").write_bytes(b"one" * 4096) (self.source / "album" / "two.bin").write_bytes(b"two" * 2048) self.metainfo = self.root / "source.torrent" self.metainfo.write_bytes(b"d4:infod4:name5:albume") self.store = ClientStore(self.root / "state.db") self.store.initialize() def tearDown(self): self.temporary.cleanup() def test_stage_publish_and_receiver_materialization_hardlink(self): manifest = self._manifest() published = stage_transfer( manifest, source_root=self.source, sync_root=self.sync, store=self.store, artifact_sources={"metainfo/source.torrent": self.metainfo}, ) staged_one = published.job_directory / "payload/album/one.bin" self.assertEqual( staged_one.stat().st_ino, (self.source / "album/one.bin").stat().st_ino, ) self.assertTrue(published.ready_path.is_file()) self.assertEqual( hashlib.sha256(published.manifest_path.read_bytes()).hexdigest(), published.manifest_sha256_hex, ) loaded = load_published_transfer(published.job_directory) self.assertEqual(loaded.manifest.job_id, manifest.job_id) self.assertEqual( loaded.manifest.files[0].source_staging_method, transfer_pb2.MATERIALIZATION_METHOD_HARD_LINK, ) self.assertEqual( loaded.manifest.artifacts[0].sha256_hex, hashlib.sha256(self.metainfo.read_bytes()).hexdigest(), ) target_manifest = materialize_transfer( loaded, target_root=self.target, store=self.store, ) self.assertEqual( (self.target / "album/one.bin").read_bytes(), (self.source / "album/one.bin").read_bytes(), ) self.assertEqual( target_manifest.files[0].target_method, transfer_pb2.MATERIALIZATION_METHOD_HARD_LINK, ) self.assertFalse(target_manifest.files[0].target_preexisted) def test_stage_and_materialize_are_idempotent_after_store_reopen(self): manifest = self._manifest() first = stage_transfer( manifest, source_root=self.source, sync_root=self.sync, store=self.store, artifact_sources={"metainfo/source.torrent": self.metainfo}, ) first_target = materialize_transfer( first, target_root=self.target, store=self.store ) reopened = ClientStore(self.root / "state.db") second = stage_transfer( manifest, source_root=self.source, sync_root=self.sync, store=reopened, artifact_sources={"metainfo/source.torrent": self.metainfo}, ) second_target = materialize_transfer( second, target_root=self.target, store=reopened ) self.assertEqual( first.manifest_sha256_hex, second.manifest_sha256_hex ) self.assertEqual( first_target.files[0].target_method, second_target.files[0].target_method, ) self.assertEqual( len(reopened.file_operation_rows(manifest.job_id)), 5 ) def test_same_size_preexisting_target_is_reused(self): manifest = self._manifest() published = stage_transfer( manifest, source_root=self.source, sync_root=self.sync, store=self.store, artifact_sources={"metainfo/source.torrent": self.metainfo}, ) (self.target / "album").mkdir() preexisting = self.target / "album/one.bin" preexisting.write_bytes((self.source / "album/one.bin").read_bytes()) inode = preexisting.stat().st_ino result = materialize_transfer( published, target_root=self.target, store=self.store ) self.assertEqual(preexisting.stat().st_ino, inode) self.assertEqual( result.files[0].target_method, transfer_pb2.MATERIALIZATION_METHOD_PREEXISTING_REUSED, ) self.assertTrue(result.files[0].target_preexisted) def test_reflink_fallback_records_reflink_method(self): destination_root = self.root / "reflink-target" destination_root.mkdir() def fake_reflink(source, destination, operation_id): del operation_id destination.write_bytes(source.read_bytes()) with mock.patch( "archive_clients.transfer._reflink", side_effect=fake_reflink ): result = FileMaterializer( self.store, allow_hardlink=False, allow_reflink=True ).materialize( job_id=str(uuid4()), operation_id="reflink-operation", source_root=self.source, source_relative_path="album/one.bin", destination_root=destination_root, destination_relative_path="album/one.bin", expected_bytes=(self.source / "album/one.bin").stat().st_size, allow_preexisting_reuse=False, sparse_supported=True, ) self.assertEqual( result.method, transfer_pb2.MATERIALIZATION_METHOD_REFLINK ) def test_copy_fallback_preserves_sparse_file(self): sparse = self.source / "album/sparse.bin" with sparse.open("wb") as output: output.seek(4 * 1024 * 1024) output.write(b"tail") destination_root = self.root / "copy-target" destination_root.mkdir() result = FileMaterializer( self.store, allow_hardlink=False, allow_reflink=False ).materialize( job_id=str(uuid4()), operation_id="copy-operation", source_root=self.source, source_relative_path="album/sparse.bin", destination_root=destination_root, destination_relative_path="album/sparse.bin", expected_bytes=sparse.stat().st_size, allow_preexisting_reuse=False, sparse_supported=True, ) copied = destination_root / "album/sparse.bin" self.assertEqual(copied.read_bytes()[-4:], b"tail") self.assertEqual(copied.stat().st_size, sparse.stat().st_size) self.assertEqual( result.method, transfer_pb2.MATERIALIZATION_METHOD_COPY ) self.assertTrue(result.sparse) self.assertLess( os.stat(copied).st_blocks * 512, os.stat(copied).st_size ) def test_partial_cleanup_removes_only_journalled_temporary(self): job_id = str(uuid4()) job_root = self.sync / ".archive-control/jobs" / job_id destination = job_root / "payload/album/one.bin" destination.parent.mkdir(parents=True) operation_id = "interrupted-copy" self.store.begin_file_operation( operation_id, job_id, '{"destination":"payload/album/one.bin"}', ) temporary = destination.with_name( ".one.bin.archive-control-" + hashlib.sha256(operation_id.encode()).hexdigest()[:16] + ".tmp" ) temporary.write_bytes(b"partial") unrelated = destination.parent / "keep-me" unrelated.write_bytes(b"unrelated") cleanup_partial_transfer( job_root, job_id=job_id, store=self.store ) self.assertFalse(temporary.exists()) self.assertTrue(unrelated.is_file()) def test_canonical_manifest_json_is_stable(self): manifest = self._manifest() first = canonical_message_json(manifest) second = canonical_message_json(self._manifest(job_id=manifest.job_id)) self.assertEqual(first, second) self.assertEqual(first, canonical_message_json(manifest)) def test_manifest_rejects_case_colliding_target_paths(self): manifest = self._manifest() manifest.files[1].target_canonical_path = "album/ONE.bin" with self.assertRaisesRegex( TransferIntegrityError, "case-colliding target" ): stage_transfer( manifest, source_root=self.source, sync_root=self.sync, store=self.store, artifact_sources={ "metainfo/source.torrent": self.metainfo }, ) def _manifest(self, job_id=None): manifest = transfer_pb2.TransferManifest( manifest_version=1, job_id=job_id or str(uuid4()), source_client_id="cache-1", target_client_id="archive-1", route_id="route-1", ) manifest.resource_id.info_hash_v1_hex = "a" * 40 manifest.created_at.seconds = 1_700_000_000 for index, relative in enumerate( ("album/one.bin", "album/two.bin") ): size = (self.source / relative).stat().st_size entry = manifest.files.add( file_index=index, payload_relative_path=f"payload/{relative}", target_canonical_path=relative, logical_bytes=size, ) del entry artifact = manifest.artifacts.add( kind=transfer_pb2.ARTIFACT_KIND_TORRENT_FILE, payload_relative_path="metainfo/source.torrent", logical_bytes=self.metainfo.stat().st_size, format="application/x-bittorrent", ) del artifact return manifest if __name__ == "__main__": unittest.main()