init
This commit is contained in:
root
@@ -0,0 +1,307 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user