229 lines
9.0 KiB
Python
229 lines
9.0 KiB
Python
"""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']
|