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 assert service.bootstrap()["settings"]["max_active_ffmpeg"] == 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, "max_active_ffmpeg": 3, }, method="PUT", ) assert settings[0] == 200 assert settings[1]["max_active_downloads"] == 2 assert settings[1]["max_active_ffmpeg"] == 3 assert service.store.max_active_downloads() == 2 assert service.store.max_active_ffmpeg() == 3 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] paused = _request(base, f"/api/tasks/{task[1]['id']}/pause", {}) assert paused[0] == 200 assert paused[1]["status"] == "paused" resumed = _request(base, f"/api/tasks/{task[1]['id']}/resume", {}) assert resumed[0] == 200 assert resumed[1]["status"] == "running" 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_pause_waits_for_running_item_then_resumes_at_download_queue_tail(tmp_path: Path) -> None: store = TaskState(tmp_path) 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)), ) claimed = store.claim_next_item() assert claimed is not None _, first = claimed assert store.pause_task(task_id) pausing = store.get_task(task_id) assert pausing is not None assert pausing["status"] == "pausing" assert pausing["items"][1]["status"] == "paused" store.mark_item_paused(int(first["id"]), "download") paused = store.get_task(task_id) assert paused is not None assert paused["status"] == "paused" resumed = store.resume_task(task_id) assert resumed == ([int(first["id"]), int(paused["items"][1]["id"])], []) next_claimed = store.claim_next_item() assert next_claimed is not None _, next_item = next_claimed assert next_item["id"] == first["id"] def test_cancelling_a_paused_task_reaches_cancelled(tmp_path: Path) -> None: store = TaskState(tmp_path) task_id = store.create_task( title="batch", base_dir=tmp_path, output_dir=tmp_path / "batch", proxy="", items=decode_download_items(json.dumps([_item("http://127.0.0.1")])), ) assert store.pause_task(task_id) assert store.get_task(task_id)["status"] == "paused" # type: ignore[index] assert store.cancel_task(task_id) cancelled = store.get_task(task_id) assert cancelled is not None assert cancelled["status"] == "cancelled" 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", "--max-active-ffmpeg", "4"]) config = config_from_args(args) assert config.max_active_downloads == 3 assert config.max_active_ffmpeg == 4 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()