import hashlib import tempfile import threading 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_reconnect_replay_never_duplicates_any_transfer_step(self): for step in ( job_pb2.JOB_STEP_KIND_SOURCE_STAGE, job_pb2.JOB_STEP_KIND_SYNCTHING_TRANSFER, job_pb2.JOB_STEP_KIND_TARGET_MATERIALIZE, job_pb2.JOB_STEP_KIND_QB_VERIFY, job_pb2.JOB_STEP_KIND_STAGING_CLEANUP, ): with self.subTest(step=step), 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="reconnect 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, ) executor.assign(control_pb2.AssignJobCommand( job=definition, expected_job_revision=1, expected_last_event_sequence=0, )) entered = threading.Event() release = threading.Event() def execute_step(*_args): entered.set() release.wait(1) return None executor._execute_step = Mock(side_effect=execute_step) command = control_pb2.ExecuteStepCommand( job_id=definition.job_id, expected_job_revision=1, expected_last_event_sequence=1, step=step, attempt=1, ) results: list[list[control_pb2.JobEvent]] = [] first = threading.Thread( target=lambda: results.append(executor.execute(command)) ) second = threading.Thread( target=lambda: results.append(executor.execute(command)) ) first.start() self.assertTrue(entered.wait(1)) second.start() release.set() first.join(1) second.join(1) self.assertFalse(first.is_alive()) self.assertFalse(second.is_alive()) self.assertEqual(executor._execute_step.call_count, 1) self.assertEqual(len(results), 2) self.assertEqual(results[0][-1].event_id, results[1][-1].event_id) 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 test_first_partial_progress_is_durable(self): """A sparse transfer must not lose its only initial observation.""" 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, 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=0, )) def partial_progress(_definition, _step, progress): progress(0.001, 10, 10_000, "still receiving") with patch.object(executor, "_execute_step", partial_progress): events = executor.execute(control_pb2.ExecuteStepCommand( job_id=definition.job_id, expected_job_revision=1, expected_last_event_sequence=1, step=job_pb2.JOB_STEP_KIND_SYNCTHING_TRANSFER, attempt=1, )) self.assertEqual(len(events), 3) self.assertEqual(events[1].type, control_pb2.JOB_EVENT_TYPE_PROGRESS) self.assertEqual(events[1].progress.bytes_complete, 10) 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()