bvandeusen
4272a19d40
The single _FETCH_TIMEOUT=3000s meant different things per host: a TOTAL wall-clock for mega (subprocess), but only a per-read socket timeout for HTTP hosts (requests' timeout is the idle gap between bytes, never a total). So a stalled HTTP connection tied up a download-worker slot AND the per-host serialize lock for ~50 min before failing (operator-flagged 2026-06-17). Split into two limits in external_fetch: - read timeout (_READ_TIMEOUT=60s, with _CONNECT_TIMEOUT=30s) → requests gets (connect, read); a stalled socket now fails in ~60s. - total budget (_TOTAL_TIMEOUT=30min) → enforced as a wall-clock deadline across chunks in _stream_to_file (HTTP has no total-download timeout), and passed as the subprocess total for mega. fetch_external() signature: timeout= → read_timeout=/total_timeout=. gdrive (gdown) self-manages; the celery hard limit is the outer backstop. Also lowered the per-host lock TTL 3600→2400 so a worker that dies holding it can't wedge a host's links much past one fetch's budget. Each external link is already one Celery task (sweep enqueues one fetch_external_link.delay per link), so these budgets are per-link. Tests: total-budget-exceeded cleans the .part; HTTP gets (connect, read); mega gets the total. Worker fakes updated to **kwargs. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
182 lines
6.2 KiB
Python
182 lines
6.2 KiB
Python
"""Unit tests for external_fetch — no network/gdown/MEGAcmd (seams stubbed)."""
|
|
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import requests
|
|
|
|
import backend.app.services.external_fetch as ef
|
|
from backend.app.services.external_fetch import fetch_external
|
|
|
|
|
|
def _resp(*, status=200, headers=None, chunks=(b"hello",), text=""):
|
|
class _R:
|
|
def __init__(self):
|
|
self.status_code = status
|
|
self.headers = headers or {}
|
|
self.text = text
|
|
|
|
def iter_content(self, chunk_size=65536):
|
|
yield from chunks
|
|
|
|
return _R()
|
|
|
|
|
|
def test_dropbox_forces_dl1_and_downloads(tmp_path, monkeypatch):
|
|
seen = {}
|
|
|
|
def fake_get(url, *, timeout, headers=None, stream=True):
|
|
seen["url"] = url
|
|
return _resp(headers={"Content-Disposition": 'attachment; filename="film.zip"'})
|
|
|
|
monkeypatch.setattr(ef, "_http_get", fake_get)
|
|
res = fetch_external(
|
|
"dropbox", "https://www.dropbox.com/s/abc/film.zip?dl=0", tmp_path
|
|
)
|
|
assert res.ok
|
|
assert "dl=1" in seen["url"]
|
|
assert res.files[0].name == "film.zip"
|
|
assert res.files[0].read_bytes() == b"hello"
|
|
|
|
|
|
def test_pixeldrain_hits_api_endpoint(tmp_path, monkeypatch):
|
|
seen = {}
|
|
|
|
def fake_get(url, *, timeout, headers=None, stream=True):
|
|
seen["url"] = url
|
|
return _resp()
|
|
|
|
monkeypatch.setattr(ef, "_http_get", fake_get)
|
|
res = fetch_external("pixeldrain", "https://pixeldrain.com/u/ems4UomN", tmp_path)
|
|
assert res.ok
|
|
assert seen["url"] == "https://pixeldrain.com/api/file/ems4UomN"
|
|
|
|
|
|
def test_mediafire_scrapes_then_downloads(tmp_path, monkeypatch):
|
|
calls = []
|
|
|
|
def fake_get(url, *, timeout, headers=None, stream=True):
|
|
calls.append(url)
|
|
if "mediafire.com/file" in url:
|
|
return _resp(
|
|
text='<a id="downloadButton" '
|
|
'href="https://download1234.mediafire.com/abc/film.zip">dl</a>'
|
|
)
|
|
return _resp(headers={"Content-Disposition": 'filename="film.zip"'})
|
|
|
|
monkeypatch.setattr(ef, "_http_get", fake_get)
|
|
res = fetch_external("mediafire", "https://www.mediafire.com/file/x/film.zip", tmp_path)
|
|
assert res.ok
|
|
assert res.files[0].name == "film.zip"
|
|
assert calls[1].startswith("https://download1234.mediafire.com/")
|
|
|
|
|
|
def test_mediafire_missing_link_errors(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(ef, "_http_get", lambda *a, **k: _resp(text="<html>none</html>"))
|
|
res = fetch_external("mediafire", "https://www.mediafire.com/file/x", tmp_path)
|
|
assert not res.ok
|
|
assert "not found" in res.error
|
|
|
|
|
|
def test_gdrive_uses_gdown(tmp_path, monkeypatch):
|
|
def fake_gdown(url, out_dir):
|
|
p = Path(out_dir) / "drive.bin"
|
|
p.write_bytes(b"xyz")
|
|
return str(p)
|
|
|
|
monkeypatch.setattr(ef, "_gdown_download", fake_gdown)
|
|
res = fetch_external("gdrive", "https://drive.google.com/file/d/ID/view", tmp_path)
|
|
assert res.ok
|
|
assert res.files[0].name == "drive.bin"
|
|
assert res.bytes == 3
|
|
|
|
|
|
def test_gdrive_no_file_errors(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(ef, "_gdown_download", lambda url, out_dir: None)
|
|
res = fetch_external("gdrive", "https://drive.google.com/file/d/ID/view", tmp_path)
|
|
assert not res.ok
|
|
assert "gdown" in res.error
|
|
|
|
|
|
def test_mega_collects_new_files(tmp_path, monkeypatch):
|
|
def fake_mega(url, out_dir, *, timeout):
|
|
(Path(out_dir) / "film.mp4").write_bytes(b"vid")
|
|
|
|
monkeypatch.setattr(ef, "_run_mega_get", fake_mega)
|
|
res = fetch_external("mega", "https://mega.nz/file/x#k", tmp_path)
|
|
assert res.ok
|
|
assert res.files[0].name == "film.mp4"
|
|
|
|
|
|
def test_unsupported_host(tmp_path):
|
|
res = fetch_external("nope", "https://x", tmp_path)
|
|
assert not res.ok
|
|
assert "unsupported" in res.error
|
|
|
|
|
|
def test_non_200_errors(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(ef, "_http_get", lambda *a, **k: _resp(status=404))
|
|
res = fetch_external("pixeldrain", "https://pixeldrain.com/u/x", tmp_path)
|
|
assert not res.ok
|
|
assert "404" in res.error
|
|
|
|
|
|
def test_transport_error_captured(tmp_path, monkeypatch):
|
|
def boom(*a, **k):
|
|
raise requests.ConnectionError("down")
|
|
|
|
monkeypatch.setattr(ef, "_http_get", boom)
|
|
res = fetch_external("dropbox", "https://www.dropbox.com/s/x?dl=0", tmp_path)
|
|
assert not res.ok
|
|
assert "transport" in res.error
|
|
|
|
|
|
def test_should_stop_aborts_and_cleans_part(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(ef, "_http_get", lambda *a, **k: _resp(chunks=(b"a", b"b")))
|
|
res = fetch_external(
|
|
"pixeldrain", "https://pixeldrain.com/u/x", tmp_path,
|
|
should_stop=lambda: True,
|
|
)
|
|
assert not res.ok
|
|
assert list(tmp_path.glob("*.part")) == []
|
|
|
|
|
|
def test_total_budget_exceeded_aborts_and_cleans_part(tmp_path, monkeypatch):
|
|
# A negative total_timeout puts the deadline in the past, so the first
|
|
# chunk's budget check trips — simulates a slow trickle blowing the budget.
|
|
monkeypatch.setattr(ef, "_http_get", lambda *a, **k: _resp(chunks=(b"a", b"b")))
|
|
res = fetch_external(
|
|
"pixeldrain", "https://pixeldrain.com/u/x", tmp_path, total_timeout=-1,
|
|
)
|
|
assert not res.ok
|
|
assert "budget" in res.error
|
|
assert list(tmp_path.glob("*.part")) == []
|
|
|
|
|
|
def test_http_get_uses_connect_and_read_timeout_tuple(tmp_path, monkeypatch):
|
|
# A stalled socket must fail on the short read timeout, not the total — so
|
|
# the HTTP layer gets (connect, read), never the total budget.
|
|
seen = {}
|
|
|
|
def fake_get(url, *, timeout, headers=None, stream=True):
|
|
seen["timeout"] = timeout
|
|
return _resp()
|
|
|
|
monkeypatch.setattr(ef, "_http_get", fake_get)
|
|
fetch_external("pixeldrain", "https://pixeldrain.com/u/x", tmp_path, read_timeout=42)
|
|
assert seen["timeout"] == (ef._CONNECT_TIMEOUT, 42)
|
|
|
|
|
|
def test_mega_uses_total_timeout(tmp_path, monkeypatch):
|
|
# megatools is a subprocess: its only timeout is a total wall-clock, so the
|
|
# total budget (not the read timeout) applies to it.
|
|
seen = {}
|
|
|
|
def fake_mega(url, out_dir, *, timeout):
|
|
seen["timeout"] = timeout
|
|
(Path(out_dir) / "film.mp4").write_bytes(b"vid")
|
|
|
|
monkeypatch.setattr(ef, "_run_mega_get", fake_mega)
|
|
fetch_external("mega", "https://mega.nz/file/x#k", tmp_path, total_timeout=123)
|
|
assert seen["timeout"] == 123
|