468 lines
17 KiB
Python
468 lines
17 KiB
Python
import hashlib
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
from uuid import uuid4
|
|
|
|
from archive_clients.bencode import encode
|
|
from archive_clients.jobs import (
|
|
ClientJobExecutor,
|
|
JobExecutionError,
|
|
_resource_fingerprint,
|
|
)
|
|
from archive_clients.syncthing import RouteSetupError
|
|
from archive_clients.resources import normalize_resource
|
|
from archive_clients.state import ClientStore
|
|
from archive_control.v1 import control_pb2, job_pb2
|
|
|
|
|
|
class CompleteSyncthing:
|
|
def __init__(self):
|
|
self.posts = []
|
|
|
|
def get_json(self, path):
|
|
if path.startswith("/rest/db/completion?"):
|
|
return {"completion": 100}
|
|
if path.startswith("/rest/db/need?"):
|
|
return {"progress": [], "queued": [], "rest": []}
|
|
raise AssertionError(path)
|
|
|
|
def post(self, path):
|
|
self.posts.append(path)
|
|
|
|
|
|
class SlowRescanSyncthing(CompleteSyncthing):
|
|
def post(self, path):
|
|
raise RouteSetupError("Syncthing API is unavailable")
|
|
|
|
|
|
class ClientJobHappyPathTests(unittest.TestCase):
|
|
def test_capacity_guard_fails_before_data_movement(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
store = ClientStore(root / "client.db")
|
|
store.initialize()
|
|
executor = ClientJobExecutor(
|
|
client_id="archive-1",
|
|
qbittorrent=Mock(),
|
|
store=store,
|
|
qb_root=root,
|
|
qb_api_root=Path("/downloads"),
|
|
route_path=lambda _: root,
|
|
syncthing_transport=Mock(),
|
|
sparse_supported=True,
|
|
free_space_reserve_bytes=100,
|
|
)
|
|
with patch(
|
|
"archive_clients.jobs.shutil.disk_usage",
|
|
return_value=Mock(free=109),
|
|
), self.assertRaisesRegex(
|
|
JobExecutionError,
|
|
"109 bytes available, 110 bytes required including reserve",
|
|
):
|
|
executor._require_space(root, 10)
|
|
self.assertLessEqual(
|
|
{path.name for path in root.iterdir()},
|
|
{"client.db", "client.db-wal", "client.db-shm"},
|
|
)
|
|
|
|
def test_archive_and_unarchive_five_step_execution(self):
|
|
for operation in (
|
|
job_pb2.JOB_OPERATION_ARCHIVE,
|
|
job_pb2.JOB_OPERATION_UNARCHIVE,
|
|
):
|
|
with self.subTest(operation=operation), tempfile.TemporaryDirectory() as directory:
|
|
self._run_transfer(Path(directory), operation)
|
|
|
|
def test_same_filesystem_source_stage_does_not_require_payload_space(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
self._run_transfer(
|
|
Path(directory),
|
|
job_pb2.JOB_OPERATION_ARCHIVE,
|
|
source_stage_free_bytes=32 * 1024 * 1024 + 1024,
|
|
content=b"x" * 4096,
|
|
)
|
|
|
|
def test_mount_boundary_requires_copy_space_even_with_same_device(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
source_root = root / "source"
|
|
destination_root = root / "destination"
|
|
source_root.mkdir()
|
|
destination_root.mkdir()
|
|
(source_root / "fixture.bin").write_bytes(b"fixture")
|
|
with patch(
|
|
"archive_clients.jobs._mount_id",
|
|
side_effect=("source-mount", "destination-mount"),
|
|
):
|
|
required = ClientJobExecutor._copy_required_bytes(
|
|
source_root,
|
|
destination_root,
|
|
(("fixture.bin", 7),),
|
|
)
|
|
self.assertEqual(required, 7)
|
|
|
|
def test_source_stage_survives_a_timed_out_syncthing_rescan(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
self._run_transfer(
|
|
Path(directory),
|
|
job_pb2.JOB_OPERATION_ARCHIVE,
|
|
syncthing=SlowRescanSyncthing(),
|
|
)
|
|
|
|
def test_target_step_advances_past_replayed_source_completion(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
store = ClientStore(root / "client.db")
|
|
store.initialize()
|
|
definition = job_pb2.JobDefinition(
|
|
job_id=str(uuid4()),
|
|
idempotency_key=str(uuid4()),
|
|
operation=job_pb2.JOB_OPERATION_ARCHIVE,
|
|
resource_display_name="fixture",
|
|
transfer={
|
|
"source_client_id": "cache-1",
|
|
"target_client_id": "archive-1",
|
|
"route_id": "route-1",
|
|
},
|
|
)
|
|
definition.resource_id.info_hash_v1_hex = "a" * 40
|
|
definition.created_at.GetCurrentTime()
|
|
executor = ClientJobExecutor(
|
|
client_id="archive-1",
|
|
qbittorrent=Mock(),
|
|
store=store,
|
|
qb_root=root,
|
|
qb_api_root=Path("/downloads"),
|
|
route_path=lambda _: root,
|
|
syncthing_transport=Mock(),
|
|
sparse_supported=True,
|
|
poll_interval=0,
|
|
)
|
|
executor.assign(control_pb2.AssignJobCommand(
|
|
job=definition,
|
|
expected_job_revision=1,
|
|
expected_last_event_sequence=1,
|
|
))
|
|
source_complete = executor._event(
|
|
definition,
|
|
sequence=5,
|
|
revision=4,
|
|
event_type=control_pb2.JOB_EVENT_TYPE_STEP_SUCCEEDED,
|
|
state=job_pb2.JOB_STATE_RUNNING,
|
|
committed=False,
|
|
step=job_pb2.JOB_STEP_KIND_SOURCE_STAGE,
|
|
step_state=job_pb2.STEP_STATE_SUCCEEDED,
|
|
)
|
|
executor._record(definition, source_complete)
|
|
|
|
with patch.object(executor, "_wait_for_syncthing"):
|
|
events = executor.execute(control_pb2.ExecuteStepCommand(
|
|
job_id=definition.job_id,
|
|
expected_job_revision=4,
|
|
expected_last_event_sequence=5,
|
|
step=job_pb2.JOB_STEP_KIND_SYNCTHING_TRANSFER,
|
|
attempt=1,
|
|
))
|
|
|
|
self.assertEqual(events[0].sequence, 6)
|
|
self.assertEqual(
|
|
events[0].progress.step,
|
|
job_pb2.JOB_STEP_KIND_SYNCTHING_TRANSFER,
|
|
)
|
|
self.assertEqual(events[-1].sequence, 7)
|
|
self.assertEqual(
|
|
events[-1].type,
|
|
control_pb2.JOB_EVENT_TYPE_STEP_SUCCEEDED,
|
|
)
|
|
|
|
def _run_transfer(
|
|
self,
|
|
root: Path,
|
|
operation: int,
|
|
*,
|
|
source_stage_free_bytes: int | None = None,
|
|
content: bytes = b"archive-control-happy-path",
|
|
syncthing: CompleteSyncthing | None = None,
|
|
):
|
|
source_root = root / "source"
|
|
target_root = root / "target"
|
|
route_root = root / "route"
|
|
source_root.mkdir()
|
|
target_root.mkdir()
|
|
route_root.mkdir()
|
|
(source_root / "fixture.bin").write_bytes(content)
|
|
info = {
|
|
b"length": len(content),
|
|
b"name": b"fixture.bin",
|
|
b"piece length": 16384,
|
|
b"pieces": hashlib.sha1(content).digest(),
|
|
}
|
|
metainfo = encode({b"info": info})
|
|
torrent_hash = hashlib.sha1(encode(info)).hexdigest()
|
|
resource = normalize_resource(
|
|
{
|
|
"hash": torrent_hash,
|
|
"name": "fixture.bin",
|
|
"state": "uploading",
|
|
},
|
|
[{
|
|
"index": 0,
|
|
"name": "fixture.bin",
|
|
"size": len(content),
|
|
"completed": len(content),
|
|
"priority": 1,
|
|
}],
|
|
metainfo,
|
|
)
|
|
source_id, target_id = (
|
|
("cache-1", "archive-1")
|
|
if operation == job_pb2.JOB_OPERATION_ARCHIVE
|
|
else ("archive-1", "cache-1")
|
|
)
|
|
definition = job_pb2.JobDefinition(
|
|
job_id=str(uuid4()),
|
|
idempotency_key=str(uuid4()),
|
|
operation=operation,
|
|
resource_display_name="fixture.bin",
|
|
transfer={
|
|
"source_client_id": source_id,
|
|
"target_client_id": target_id,
|
|
"route_id": "route-1",
|
|
"requested_logical_bytes": len(content),
|
|
"transfer_delta_logical_bytes": len(content),
|
|
},
|
|
)
|
|
definition.resource_id.info_hash_v1_hex = torrent_hash
|
|
definition.created_at.GetCurrentTime()
|
|
definition.transfer.requested_files.ranges.add(first=0, last=0)
|
|
definition.transfer.transfer_delta_files.ranges.add(first=0, last=0)
|
|
definition.transfer.source_fingerprint.CopyFrom(
|
|
_resource_fingerprint(resource, source_id)
|
|
)
|
|
|
|
source_store = ClientStore(root / "source.db")
|
|
target_store = ClientStore(root / "target.db")
|
|
source_store.initialize()
|
|
target_store.initialize()
|
|
source_qb = Mock()
|
|
source_qb.get_resource.return_value = resource
|
|
target_qb = Mock()
|
|
target_qb.get_resource.side_effect = [
|
|
None, None, resource, resource,
|
|
]
|
|
syncthing = syncthing or CompleteSyncthing()
|
|
source = ClientJobExecutor(
|
|
client_id=source_id,
|
|
qbittorrent=source_qb,
|
|
store=source_store,
|
|
qb_root=source_root,
|
|
qb_api_root=Path("/downloads"),
|
|
route_path=lambda _: route_root,
|
|
syncthing_transport=syncthing,
|
|
sparse_supported=True,
|
|
poll_interval=0,
|
|
)
|
|
target = ClientJobExecutor(
|
|
client_id=target_id,
|
|
qbittorrent=target_qb,
|
|
store=target_store,
|
|
qb_root=target_root,
|
|
qb_api_root=Path("/downloads"),
|
|
route_path=lambda _: route_root,
|
|
syncthing_transport=syncthing,
|
|
sparse_supported=True,
|
|
poll_interval=0,
|
|
)
|
|
|
|
source_assigned = source.assign(control_pb2.AssignJobCommand(
|
|
job=definition,
|
|
expected_job_revision=1,
|
|
expected_last_event_sequence=0,
|
|
))
|
|
target_assigned = target.assign(control_pb2.AssignJobCommand(
|
|
job=definition,
|
|
expected_job_revision=1,
|
|
expected_last_event_sequence=1,
|
|
))
|
|
self.assertEqual(source_assigned[0].sequence, 1)
|
|
self.assertEqual(target_assigned[0].sequence, 2)
|
|
|
|
cursor_revision = 1
|
|
cursor_sequence = 2
|
|
pipeline = (
|
|
(source, job_pb2.JOB_STEP_KIND_SOURCE_STAGE),
|
|
(target, job_pb2.JOB_STEP_KIND_SYNCTHING_TRANSFER),
|
|
(target, job_pb2.JOB_STEP_KIND_TARGET_MATERIALIZE),
|
|
(target, job_pb2.JOB_STEP_KIND_QB_VERIFY),
|
|
(source, job_pb2.JOB_STEP_KIND_STAGING_CLEANUP),
|
|
)
|
|
final = None
|
|
for executor, step in pipeline:
|
|
command = control_pb2.ExecuteStepCommand(
|
|
job_id=definition.job_id,
|
|
expected_job_revision=cursor_revision,
|
|
expected_last_event_sequence=cursor_sequence,
|
|
step=step,
|
|
attempt=1,
|
|
)
|
|
if executor is source and step == job_pb2.JOB_STEP_KIND_SOURCE_STAGE and source_stage_free_bytes is not None:
|
|
with patch(
|
|
"archive_clients.jobs.shutil.disk_usage",
|
|
return_value=Mock(free=source_stage_free_bytes),
|
|
):
|
|
events = executor.execute(command)
|
|
else:
|
|
events = executor.execute(command)
|
|
self.assertEqual(
|
|
[event.sequence for event in events],
|
|
list(range(
|
|
cursor_sequence + 1,
|
|
cursor_sequence + len(events) + 1,
|
|
)),
|
|
)
|
|
cursor_sequence = events[-1].sequence
|
|
cursor_revision = events[-1].job_revision
|
|
final = events[-1]
|
|
|
|
self.assertEqual((target_root / "fixture.bin").read_bytes(), content)
|
|
target_qb.add_stopped_with_retry.assert_called_once_with(
|
|
metainfo, "/downloads", torrent_hash
|
|
)
|
|
target_qb.set_selection.assert_called_once_with(torrent_hash, [0], 1)
|
|
target_qb.recheck_and_wait.assert_called_once()
|
|
target_qb.start.assert_called_once_with(torrent_hash)
|
|
self.assertIsNotNone(final)
|
|
self.assertTrue(final.committed)
|
|
self.assertEqual(final.type, control_pb2.JOB_EVENT_TYPE_SUCCEEDED)
|
|
self.assertFalse(
|
|
(route_root / ".archive-control/jobs" / definition.job_id).exists()
|
|
)
|
|
|
|
replay = source.execute(control_pb2.ExecuteStepCommand(
|
|
job_id=definition.job_id,
|
|
expected_job_revision=cursor_revision - 2,
|
|
expected_last_event_sequence=cursor_sequence - 2,
|
|
step=job_pb2.JOB_STEP_KIND_STAGING_CLEANUP,
|
|
attempt=1,
|
|
))
|
|
self.assertEqual(replay[-1].event_id, final.event_id)
|
|
|
|
def test_step_failure_is_durable_and_reports_clear_reason(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
store = ClientStore(root / "client.db")
|
|
store.initialize()
|
|
definition = job_pb2.JobDefinition(
|
|
job_id=str(uuid4()),
|
|
idempotency_key=str(uuid4()),
|
|
operation=job_pb2.JOB_OPERATION_ARCHIVE,
|
|
resource_display_name="fixture",
|
|
transfer={
|
|
"source_client_id": "cache-1",
|
|
"target_client_id": "archive-1",
|
|
"route_id": "route-1",
|
|
},
|
|
)
|
|
definition.resource_id.info_hash_v1_hex = "a" * 40
|
|
definition.created_at.GetCurrentTime()
|
|
executor = ClientJobExecutor(
|
|
client_id="cache-1",
|
|
qbittorrent=Mock(),
|
|
store=store,
|
|
qb_root=root,
|
|
qb_api_root=Path("/downloads"),
|
|
route_path=lambda _: root,
|
|
syncthing_transport=Mock(),
|
|
sparse_supported=True,
|
|
poll_interval=0,
|
|
)
|
|
executor.assign(control_pb2.AssignJobCommand(
|
|
job=definition,
|
|
expected_job_revision=1,
|
|
expected_last_event_sequence=0,
|
|
))
|
|
command = control_pb2.ExecuteStepCommand(
|
|
job_id=definition.job_id,
|
|
expected_job_revision=1,
|
|
expected_last_event_sequence=1,
|
|
step=job_pb2.JOB_STEP_KIND_SOURCE_STAGE,
|
|
attempt=1,
|
|
)
|
|
with patch.object(
|
|
executor,
|
|
"_execute_step",
|
|
side_effect=JobExecutionError(
|
|
"partfile cannot be handled safely"
|
|
),
|
|
):
|
|
events = executor.execute(command)
|
|
|
|
self.assertEqual(
|
|
[event.type for event in events],
|
|
[
|
|
control_pb2.JOB_EVENT_TYPE_STEP_STARTED,
|
|
control_pb2.JOB_EVENT_TYPE_FAILED,
|
|
],
|
|
)
|
|
self.assertEqual(events[-1].state, job_pb2.JOB_STATE_FAILED)
|
|
self.assertIn("partfile", events[-1].error.message)
|
|
self.assertEqual(
|
|
executor.execute(command)[-1].event_id,
|
|
events[-1].event_id,
|
|
)
|
|
self.assertEqual(store.list_active_job_cursors(), [])
|
|
|
|
def test_active_source_cancellation_reports_durable_rollback(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
store = ClientStore(root / "client.db")
|
|
store.initialize()
|
|
definition = job_pb2.JobDefinition(
|
|
job_id=str(uuid4()),
|
|
idempotency_key=str(uuid4()),
|
|
operation=job_pb2.JOB_OPERATION_ARCHIVE,
|
|
resource_display_name="fixture",
|
|
transfer={
|
|
"source_client_id": "cache-1",
|
|
"target_client_id": "archive-1",
|
|
"route_id": "route-1",
|
|
},
|
|
)
|
|
definition.resource_id.info_hash_v1_hex = "a" * 40
|
|
definition.created_at.GetCurrentTime()
|
|
executor = ClientJobExecutor(
|
|
client_id="cache-1",
|
|
qbittorrent=Mock(),
|
|
store=store,
|
|
qb_root=root,
|
|
qb_api_root=Path("/downloads"),
|
|
route_path=lambda _: root / "route",
|
|
syncthing_transport=Mock(),
|
|
sparse_supported=True,
|
|
poll_interval=0,
|
|
)
|
|
executor.assign(control_pb2.AssignJobCommand(
|
|
job=definition,
|
|
expected_job_revision=1,
|
|
expected_last_event_sequence=0,
|
|
))
|
|
events = executor.cancel(control_pb2.CancelJobCommand(
|
|
job_id=definition.job_id,
|
|
expected_job_revision=2,
|
|
expected_last_event_sequence=1,
|
|
reason="test",
|
|
))
|
|
self.assertEqual(
|
|
events[-1].type,
|
|
control_pb2.JOB_EVENT_TYPE_ROLLBACK_SUCCEEDED,
|
|
)
|
|
self.assertEqual(
|
|
events[-1].state, job_pb2.JOB_STATE_ROLLING_BACK
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|