"""Multi-file Archive.org scheduling, fair budgets and safe restarts.""" import asyncio import json import re import threading import time from concurrent.futures import ThreadPoolExecutor import httpx import pytest from fastapi.testclient import TestClient from app import archive_download, main, store class BytesStream(httpx.SyncByteStream): def __init__(self, data): self.data = data def __iter__(self): yield self.data @pytest.fixture def setup(tmp_path, monkeypatch): monkeypatch.setattr(store, 'DATA', tmp_path) monkeypatch.setattr(store, 'DB_PATH', tmp_path / 'queue.sqlite3') monkeypatch.setenv('ARCHIVE_CONNECTIONS', '50') monkeypatch.setenv('WEB_USER', 'admin') monkeypatch.setenv('WEB_PASSWORD', 'unit-test-password') monkeypatch.setattr(archive_download, 'MIN_SEGMENT_BYTES', 16) store.init_db() return tmp_path def multi_job(names=('A.iso', 'B.iso', 'C.iso')): urls = [f'https://archive.org/download/test/{name}' for name in names] return main.new_job('archive_many', json.dumps(urls), 'Multi', 'raw') def mock_archive(monkeypatch, files, first_request_barrier=None, delay=0): lock = threading.Lock() active = 0 max_active = 0 seen = {} requests = [] def handler(req): nonlocal active, max_active name = req.url.path.rsplit('/', 1)[-1] data = files[name] with lock: active += 1 max_active = max(max_active, active) first = name not in seen seen[name] = seen.get(name, 0) + 1 requests.append((name, req.headers.get('range'))) try: if first and first_request_barrier: first_request_barrier.wait(timeout=5) if delay: time.sleep(delay) span = req.headers.get('range') if span: match = re.fullmatch(r'bytes=(\d+)-(\d*)', span) assert match, span begin = int(match.group(1)) end = int(match.group(2)) if match.group(2) else len(data)-1 if begin >= len(data): return httpx.Response(416, headers={'Content-Range': f'bytes */{len(data)}'}) assert 0 <= begin <= end < len(data) return httpx.Response(206, stream=BytesStream(data[begin:end+1]), headers={ 'Content-Range': f'bytes {begin}-{end}/{len(data)}', 'Content-Length': str(end - begin + 1), 'ETag': '"r1"'}) return httpx.Response(200, stream=BytesStream(data), headers={'Content-Length': str(len(data))}) finally: with lock: active -= 1 monkeypatch.setattr(archive_download, '_client', lambda: httpx.Client( transport=httpx.MockTransport(handler), follow_redirects=False)) return requests, lambda: max_active def test_three_files_start_together_and_share_fifty_connections(setup, monkeypatch): names = ('A.iso', 'B.iso', 'C.iso') contents = {name: bytes([i+10]) * 1024 for i, name in enumerate(names)} barrier = threading.Barrier(3) requests, max_connections = mock_archive(monkeypatch, contents, barrier, delay=0.005) job = multi_job(names) asyncio.run(archive_download.run_one(job)) for name in names: assert (setup / 'raw' / job['id'] / name).read_bytes() == contents[name] data = store.get(job['id']) assert data['state'] == 'preparing' assert data['progress'] == 100 assert data['part_index'] == 3 assert data['total'] == 3 * 1024 assert data['completed'] == 3 * 1024 assert max_connections() <= 50 assert max_connections() > 1 assert sum(1 for _, item in requests if item == 'bytes=0-0') == 3 log = (setup / 'logs' / f'{job["id"]}.download.log').read_text() assert 'start 3 plików równolegle' in log assert '17, 17, 16' in log def test_existing_fifty_segment_manifests_do_not_exceed_global_limit(setup, monkeypatch): files = {'A.iso': b'a' * 1024, 'B.iso': b'b' * 1024, 'C.iso': b'c' * 1024} requests, max_connections = mock_archive(monkeypatch, files, delay=0.002) job = multi_job() root = setup / 'raw' / job['id'] root.mkdir(parents=True) for url in json.loads(job['source']): target = root / (url.rsplit('/', 1)[-1] + '.part') _, manifest = archive_download._prepare_segments(target, url, 1024, '"r1"', 50) assert len(manifest['ranges']) == 50 archive_download._transfer_many(job) assert sum(bool(item and item != 'bytes=0-0') for _, item in requests) == 150 assert max_connections() <= 50 assert store.get(job['id'])['state'] == 'preparing' for name, data in files.items(): assert (root / name).read_bytes() == data def test_skip_completed_files_and_resume_incomplete_file(setup, monkeypatch): files = {'A.iso': b'a' * 128, 'B.iso': b'b' * 128, 'C.iso': b'c' * 128} requests, _ = mock_archive(monkeypatch, files) job = multi_job() raw = setup / 'raw' / job['id'] raw.mkdir(parents=True) (raw / 'A.iso').write_bytes(files['A.iso']) (raw / 'B.iso.part').write_bytes(files['B.iso'][:32]) archive_download._transfer_many(job) assert store.get(job['id'])['state'] == 'preparing' assert not any(name == 'A.iso' for name, _ in requests) assert not any(name == 'B.iso' and value == 'bytes=0-31' for name, value in requests) for name, data in files.items(): assert (raw / name).read_bytes() == data def test_pause_stops_all_files_and_resume_keeps_segments(setup, monkeypatch): files = {'A.iso': b'a' * 1024, 'B.iso': b'b' * 1024} gate = threading.Event() release = threading.Event() original = httpx.MockTransport def handler(req): name = req.url.path.rsplit('/', 1)[-1] data = files[name] header = req.headers.get('range', '') match = re.fullmatch(r'bytes=(\d+)-(\d*)', header) assert match start = int(match.group(1)) end = int(match.group(2)) if match.group(2) else len(data)-1 if (start, end) != (0, 0) and not gate.is_set(): gate.set() assert release.wait(5), 'download thread was not released' return httpx.Response(206, stream=BytesStream(data[start:end+1]), headers={ 'ETag': '"r1"', 'Content-Range': f'bytes {start}-{end}/{len(data)}'}) monkeypatch.setattr(archive_download, '_client', lambda: httpx.Client( transport=original(handler), follow_redirects=False)) job = multi_job(('A.iso', 'B.iso')) with ThreadPoolExecutor(max_workers=1) as executor: fut = executor.submit(archive_download._transfer_many, job) assert gate.wait(5) store.update(job['id'], state='paused') release.set() with pytest.raises(archive_download.ArchiveStopped): fut.result(timeout=10) assert store.get(job['id'])['state'] == 'paused' store.update(job['id'], state='queued') archive_download._transfer_many(store.get(job['id'])) assert store.get(job['id'])['state'] == 'preparing' assert (setup / 'raw' / job['id'] / 'A.iso').read_bytes() == files['A.iso'] assert (setup / 'raw' / job['id'] / 'B.iso').read_bytes() == files['B.iso'] def test_duplicate_filenames_fail_without_writing(setup): j = multi_job(('A/Game.iso', 'B/Game.iso')) with pytest.raises(archive_download.ArchiveDownloadError, match='powtarzające'): archive_download._transfer_many(j) assert not (setup / 'raw' / j['id']).exists() def test_retry_partial_archive_not_mistaken_for_ready_conversion(setup, monkeypatch): j = multi_job() raw = setup / 'raw' / j['id'] raw.mkdir(parents=True) (raw / 'A.iso').write_bytes(b'a') store.update(j['id'], state='error', progress=99.99) async def idle_worker(): await asyncio.sleep(1000) monkeypatch.setattr(main, 'worker', idle_worker) with TestClient(main.app) as client: response = client.post(f'/api/jobs/{j["id"]}/retry', auth=('admin', 'unit-test-password')) assert response.status_code == 200 assert response.json()['state'] == 'queued' def test_legacy_aria2_file_is_not_accepted_as_finished(setup, monkeypatch): job = multi_job(('A.iso', 'B.iso')) raw = setup / 'raw' / job['id'] raw.mkdir(parents=True) (raw / 'A.iso').write_bytes(b'old incomplete file') (raw / 'A.iso.aria2').write_bytes(b'aria2 control file') files = {'A.iso': b'A' * 128, 'B.iso': b'B' * 128} requests, _ = mock_archive(monkeypatch, files) archive_download._transfer_many(job) assert ('A.iso', 'bytes=0-0') in requests assert (raw / 'A.iso').read_bytes() == files['A.iso'] assert (raw / 'B.iso').read_bytes() == files['B.iso'] def test_aggregate_progress_is_weighted_by_bytes_not_file_count(setup): j = multi_job() tracker = archive_download._BatchProgress(j['id'], ['A.iso', 'B.iso', 'C.iso'], 50) tracker.report(0, 100, 100) tracker.report(1, 10, 1000) tracker.report(2, 20, 2000) tracker.publish() data = store.get(j['id']) assert data['total'] == 3100 assert data['completed'] == 130 assert data['progress'] == 4.19 assert 'limit 50 połączeń' in data['note']