308 lines
12 KiB
Python
308 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import functools
|
|
import http.server
|
|
import json
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
from urllib.request import Request, urlopen
|
|
|
|
from m3u8downloaderd.models import DATE_TEMPLATE, DEFAULT_PROXY, decode_download_items
|
|
from m3u8downloaderd.tasks import TaskState
|
|
from m3u8downloaderd.web import DownloadHTTPServer, DownloadService, ServiceConfig, build_argument_parser, config_from_args
|
|
from m3u8downloaderd.worker import TaskRunner
|
|
|
|
|
|
def _item(base_url: str) -> dict[str, str]:
|
|
return {
|
|
"code": "episode-1",
|
|
"title": "episode",
|
|
"href": "/episode-1",
|
|
"image_src": f"{base_url}/cover.jpg",
|
|
"m3u8_url": f"{base_url}/video.m3u8",
|
|
"m3u8_referer": "https://example.test/watch/episode-1",
|
|
}
|
|
|
|
|
|
def test_http_api_validates_and_snapshots_settings(tmp_path: Path) -> None:
|
|
config = ServiceConfig("127.0.0.1", 0, tmp_path / "downloads", 1)
|
|
config.download_root.mkdir()
|
|
service = DownloadService(config)
|
|
assert service.bootstrap()["settings"]["proxy"] == DEFAULT_PROXY
|
|
assert service.bootstrap()["settings"]["max_active_downloads"] == 1
|
|
server = DownloadHTTPServer(config, service)
|
|
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
thread.start()
|
|
try:
|
|
base = f"http://127.0.0.1:{server.server_port}"
|
|
invalid = _request(base, "/api/validate", {"payload": "[]"})
|
|
assert invalid[0] == 400
|
|
assert "non-empty" in invalid[1]["error"]
|
|
|
|
validation = _request(base, "/api/validate", {"payload": json.dumps([_item("http://127.0.0.1")])})
|
|
assert validation == (200, {"valid": True, "count": 1})
|
|
settings = _request(
|
|
base,
|
|
"/api/settings",
|
|
{
|
|
"default_title_template": "YYYY-MM-DD",
|
|
"default_directory": str(config.download_root),
|
|
"proxy": "http://127.0.0.1:7890",
|
|
"max_active_downloads": 2,
|
|
},
|
|
method="PUT",
|
|
)
|
|
assert settings[0] == 200
|
|
assert settings[1]["max_active_downloads"] == 2
|
|
assert service.store.max_active_downloads() == 2
|
|
task = _request(base, "/api/tasks", {"title": "batch", "directory": str(config.download_root), "payload": json.dumps([_item("http://127.0.0.1")])})
|
|
assert task[0] == 201
|
|
assert task[1]["proxy"] == "http://127.0.0.1:7890"
|
|
assert task[1]["output_dir"] == str(config.download_root / "batch")
|
|
assert "m3u8_referer" not in task[1]["items"][0]
|
|
service.store.log("error", "Expected log entry")
|
|
logs = _get(base, "/api/logs?level=error")
|
|
assert logs[0] == 200
|
|
assert logs[1]["logs"][0]["message"] == "Expected log entry"
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|
|
thread.join()
|
|
service.stop()
|
|
|
|
|
|
def test_runner_creates_mp4_with_attached_cover(tmp_path: Path) -> None:
|
|
source = tmp_path / "source"
|
|
source.mkdir()
|
|
_make_hls_fixture(source)
|
|
root = tmp_path / "downloads"
|
|
root.mkdir()
|
|
config = ServiceConfig("127.0.0.1", 0, root, 1)
|
|
service = DownloadService(config)
|
|
service.update_settings({"proxy": ""})
|
|
service.start()
|
|
try:
|
|
with _http_server(source) as base_url:
|
|
task = service.create_task({"title": "collection", "directory": str(root), "payload": json.dumps([_item(base_url)])})
|
|
final = _wait_for_terminal(service, task["id"])
|
|
assert final["status"] == "completed", final
|
|
item = final["items"][0]
|
|
assert item["status"] == "completed", item
|
|
output = Path(item["output_path"])
|
|
assert output == root / "collection" / "episode.mp4"
|
|
assert output.is_file()
|
|
streams = subprocess.run(
|
|
["ffprobe", "-v", "error", "-show_streams", "-of", "json", str(output)],
|
|
capture_output=True,
|
|
text=True,
|
|
check=True,
|
|
).stdout
|
|
assert '"attached_pic": 1' in streams
|
|
finally:
|
|
service.stop()
|
|
|
|
|
|
def test_tasks_and_settings_are_not_retained_after_restart(tmp_path: Path) -> None:
|
|
root = tmp_path / "downloads"
|
|
root.mkdir()
|
|
config = ServiceConfig("127.0.0.1", 0, root, 1)
|
|
first = DownloadService(config)
|
|
first.update_settings({"default_title_template": "batch", "proxy": ""})
|
|
first.create_task({"title": "batch", "directory": str(root), "payload": json.dumps([_item("http://127.0.0.1")])})
|
|
|
|
restarted = DownloadService(config)
|
|
|
|
assert restarted.list_tasks() == []
|
|
assert restarted.bootstrap()["settings"]["default_title_template"] == DATE_TEMPLATE
|
|
assert restarted.bootstrap()["settings"]["proxy"] == DEFAULT_PROXY
|
|
assert restarted.bootstrap()["settings"]["max_active_downloads"] == 1
|
|
|
|
|
|
def test_runner_runs_multiple_videos_from_one_task_concurrently(tmp_path: Path) -> None:
|
|
store = TaskState(tmp_path, max_active_downloads=2)
|
|
runner = TaskRunner(store, worker_capacity=2, ffmpeg="ffmpeg")
|
|
lock = threading.Lock()
|
|
both_started = threading.Event()
|
|
release = threading.Event()
|
|
active_downloads = 0
|
|
peak_downloads = 0
|
|
|
|
def fake_download(task: dict[str, object], item: dict[str, object]) -> None:
|
|
nonlocal active_downloads, peak_downloads
|
|
with lock:
|
|
active_downloads += 1
|
|
peak_downloads = max(peak_downloads, active_downloads)
|
|
if active_downloads == 2:
|
|
both_started.set()
|
|
release.wait(timeout=5)
|
|
store.update_item(int(item["id"]), status="completed", stage="Completed", completed_at="now")
|
|
with lock:
|
|
active_downloads -= 1
|
|
|
|
runner._download_and_finalize = fake_download # type: ignore[method-assign]
|
|
payload = [_item("http://127.0.0.1"), {**_item("http://127.0.0.1"), "code": "episode-2", "title": "episode-2"}]
|
|
task_id = store.create_task(
|
|
title="batch",
|
|
base_dir=tmp_path,
|
|
output_dir=tmp_path / "batch",
|
|
proxy="",
|
|
items=decode_download_items(json.dumps(payload)),
|
|
)
|
|
runner.start()
|
|
try:
|
|
assert both_started.wait(timeout=3)
|
|
release.set()
|
|
deadline = time.monotonic() + 3
|
|
while time.monotonic() < deadline:
|
|
task = store.get_task(task_id)
|
|
assert task is not None
|
|
if task["status"] == "completed":
|
|
break
|
|
time.sleep(0.05)
|
|
else:
|
|
raise AssertionError("Task did not complete")
|
|
finally:
|
|
release.set()
|
|
runner.stop()
|
|
|
|
assert peak_downloads == 2
|
|
|
|
|
|
def test_cover_download_prefers_img2_s1080_url(tmp_path: Path) -> None:
|
|
store = TaskState(tmp_path)
|
|
runner = TaskRunner(store, ffmpeg="ffmpeg")
|
|
source_url = "https://cdn.example.test/img2/s720/cover.jpg?token=abc"
|
|
expected_url = "https://cdn.example.test/img2/s1080/cover.jpg?token=abc"
|
|
downloaded = tmp_path / "cover.jpg"
|
|
downloaded.write_bytes(b"cover")
|
|
requested_urls: list[str] = []
|
|
|
|
def fake_download(url: str, output_dir: Path, item_id: int, proxy: str) -> Path:
|
|
requested_urls.append(url)
|
|
return downloaded
|
|
|
|
runner._download_cover_url = fake_download # type: ignore[method-assign]
|
|
|
|
assert runner._download_cover(source_url, tmp_path, 2, "", task_id=1) == downloaded
|
|
assert requested_urls == [expected_url]
|
|
assert store.logs("warning") == []
|
|
|
|
|
|
def test_cover_download_falls_back_to_original_img2_url_and_warns(tmp_path: Path) -> None:
|
|
store = TaskState(tmp_path)
|
|
runner = TaskRunner(store, ffmpeg="ffmpeg")
|
|
source_url = "https://cdn.example.test/img2/s720/cover.jpg"
|
|
preferred_url = "https://cdn.example.test/img2/s1080/cover.jpg"
|
|
downloaded = tmp_path / "cover.jpg"
|
|
downloaded.write_bytes(b"cover")
|
|
requested_urls: list[str] = []
|
|
|
|
def fake_download(url: str, output_dir: Path, item_id: int, proxy: str) -> Path:
|
|
requested_urls.append(url)
|
|
if url == preferred_url:
|
|
raise OSError("1080p image unavailable")
|
|
return downloaded
|
|
|
|
runner._download_cover_url = fake_download # type: ignore[method-assign]
|
|
|
|
assert runner._download_cover(source_url, tmp_path, 2, "", task_id=1) == downloaded
|
|
assert requested_urls == [preferred_url, source_url]
|
|
warning = store.logs("warning")
|
|
assert len(warning) == 1
|
|
assert warning[0]["task_id"] == 1
|
|
assert warning[0]["item_id"] == 2
|
|
assert "falling back to the original URL" in warning[0]["message"]
|
|
|
|
|
|
def test_cover_download_keeps_original_url_without_img2_size_path(tmp_path: Path) -> None:
|
|
store = TaskState(tmp_path)
|
|
runner = TaskRunner(store, ffmpeg="ffmpeg")
|
|
source_url = "https://cdn.example.test/cover.jpg?path=/img2/s720/"
|
|
downloaded = tmp_path / "cover.jpg"
|
|
downloaded.write_bytes(b"cover")
|
|
requested_urls: list[str] = []
|
|
|
|
def fake_download(url: str, output_dir: Path, item_id: int, proxy: str) -> Path:
|
|
requested_urls.append(url)
|
|
return downloaded
|
|
|
|
runner._download_cover_url = fake_download # type: ignore[method-assign]
|
|
|
|
assert runner._download_cover(source_url, tmp_path, 2, "", task_id=1) == downloaded
|
|
assert requested_urls == [source_url]
|
|
|
|
|
|
def test_cli_uses_max_active_downloads() -> None:
|
|
args = build_argument_parser().parse_args(["--max-active-downloads", "3"])
|
|
|
|
assert config_from_args(args).max_active_downloads == 3
|
|
|
|
|
|
def _get(base: str, path: str) -> tuple[int, dict[str, object]]:
|
|
with urlopen(f"{base}{path}") as response:
|
|
return response.status, json.loads(response.read())
|
|
|
|
|
|
def _request(base: str, path: str, payload: dict[str, object], method: str = "POST") -> tuple[int, dict[str, object]]:
|
|
request = Request(
|
|
f"{base}{path}",
|
|
data=json.dumps(payload).encode(),
|
|
method=method,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
try:
|
|
with urlopen(request) as response:
|
|
return response.status, json.loads(response.read())
|
|
except Exception as error:
|
|
response = error
|
|
return response.code, json.loads(response.read()) # type: ignore[attr-defined]
|
|
|
|
|
|
def _wait_for_terminal(service: DownloadService, task_id: int, timeout: float = 30) -> dict[str, object]:
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
task = service.get_task(task_id)
|
|
assert task is not None
|
|
if task["status"] in {"completed", "partial", "failed", "cancelled"}:
|
|
return task
|
|
time.sleep(0.1)
|
|
raise AssertionError("Timed out waiting for download task")
|
|
|
|
|
|
def _make_hls_fixture(directory: Path) -> None:
|
|
subprocess.run(
|
|
[
|
|
"ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
|
|
"-f", "lavfi", "-i", "testsrc2=size=160x90:rate=25:duration=1",
|
|
"-f", "lavfi", "-i", "sine=frequency=880:duration=1",
|
|
"-shortest", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-c:a", "aac",
|
|
"-f", "mpegts", str(directory / "segment.ts"),
|
|
],
|
|
check=True,
|
|
)
|
|
subprocess.run(
|
|
[
|
|
"ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
|
|
"-f", "lavfi", "-i", "color=c=orange:s=64x64:d=1", "-frames:v", "1", str(directory / "cover.jpg"),
|
|
],
|
|
check=True,
|
|
)
|
|
(directory / "video.m3u8").write_text("#EXTM3U\n#EXTINF:1,\nsegment.ts\n#EXT-X-ENDLIST\n")
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _http_server(directory: Path):
|
|
handler = functools.partial(http.server.SimpleHTTPRequestHandler, directory=str(directory))
|
|
server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), handler)
|
|
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
thread.start()
|
|
try:
|
|
yield f"http://127.0.0.1:{server.server_port}"
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|
|
thread.join()
|