first commit

tests/test_archive_multi_file.py

@@ -0,0 +1,228 @@
+"""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']