293 lines
10 KiB
Python
293 lines
10 KiB
Python
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()
|