229 lines
15 KiB
Python
229 lines
15 KiB
Python
import io
|
||
import json
|
||
import os
|
||
import struct
|
||
import tempfile
|
||
import unittest
|
||
from pathlib import Path
|
||
from unittest.mock import patch
|
||
|
||
from host.common import safe_filename, select_streams, validate_download, youtube_url
|
||
from host.media import encoding_plan, ffmpeg_command
|
||
from host.protocol import MAX_MESSAGE, read_message, write_message
|
||
from host.worker import publish_file
|
||
from host.main import NativeHost
|
||
from host.session import apply_cookies, validate_cookies, public_error
|
||
|
||
|
||
class HostTests(unittest.TestCase):
|
||
def test_fragment_access_failure_retries_and_preserves_the_original_error(self):
|
||
from unittest.mock import MagicMock
|
||
from yt_dlp.utils import DownloadError
|
||
from host.worker import download
|
||
first, second = MagicMock(), MagicMock()
|
||
first.__enter__.return_value = first; second.__enter__.return_value = second
|
||
handles = []
|
||
def fail(ydl):
|
||
# Reproduce a downloader retaining its .part handle during failure.
|
||
part = Path(ydl.params["outtmpl"].replace("%(ext)s", "mp4") + ".part")
|
||
handles.append(part.open("wb"))
|
||
raise DownloadError("HTTP Error 403: Forbidden")
|
||
def create_ydl(params):
|
||
current = first if not handles else second
|
||
current.params = params
|
||
current.extract_info.side_effect = lambda *args, **kwargs: fail(current)
|
||
return current
|
||
with tempfile.TemporaryDirectory() as folder:
|
||
try:
|
||
with patch("yt_dlp.YoutubeDL", side_effect=create_ydl), patch("host.worker.ydl_options", return_value={}), \
|
||
patch("host.worker.time.sleep"), patch("host.worker.emit") as emit:
|
||
with self.assertRaisesRegex(DownloadError, "HTTP Error 403"):
|
||
download({"url": "https://youtu.be/jNQXAC9IVRw", "folder": folder}, [])
|
||
self.assertEqual(len(handles), 2)
|
||
self.assertTrue(any("повторяю загрузку" in str(call) for call in emit.call_args_list))
|
||
finally:
|
||
for handle in handles: handle.close()
|
||
|
||
def test_transient_youtube_challenge_is_retried_once_without_browser_cookies(self):
|
||
from unittest.mock import MagicMock
|
||
from yt_dlp.utils import DownloadError
|
||
first, second = MagicMock(), MagicMock()
|
||
first.__enter__.return_value = first; second.__enter__.return_value = second
|
||
first.extract_info.side_effect = DownloadError("Sign in to confirm you’re not a bot")
|
||
second.extract_info.return_value = {"id": "jNQXAC9IVRw"}
|
||
with patch("yt_dlp.YoutubeDL", side_effect=[first, second]), patch("host.worker.ydl_options", return_value={}), \
|
||
patch("host.worker.summarize", side_effect=lambda info: info), patch("host.worker.time.sleep"):
|
||
from host.worker import inspect
|
||
self.assertEqual(inspect("https://youtu.be/jNQXAC9IVRw")["id"], "jNQXAC9IVRw")
|
||
first.cookiejar.set_cookie.assert_not_called(); second.cookiejar.set_cookie.assert_not_called()
|
||
|
||
def test_browser_session_is_youtube_only_and_never_saved_to_a_cookie_file(self):
|
||
from urllib.request import Request
|
||
from yt_dlp import YoutubeDL
|
||
from yt_dlp.cookies import YoutubeDLCookieJar
|
||
cookie = {"name": "SAPISID", "value": "synthetic-cookie", "domain": ".youtube.com", "path": "/", "secure": True, "httpOnly": True}
|
||
with patch.object(YoutubeDLCookieJar, "save", side_effect=AssertionError("Cookies must stay in memory")):
|
||
with YoutubeDL({"quiet": True}) as ydl:
|
||
apply_cookies(ydl, [cookie])
|
||
self.assertIsNone(ydl.cookiejar.filename)
|
||
request = Request("https://www.youtube.com/watch?v=jNQXAC9IVRw")
|
||
ydl.cookiejar.add_cookie_header(request)
|
||
self.assertIn("SAPISID=synthetic-cookie", request.get_header("Cookie"))
|
||
unrelated = Request("https://www.google.com/")
|
||
ydl.cookiejar.add_cookie_header(unrelated)
|
||
self.assertIsNone(unrelated.get_header("Cookie"))
|
||
for invalid in [{**cookie, "domain": ".google.com"}, {**cookie, "domain": "youtube.com.evil.test"},
|
||
{**cookie, "value": "secret\r\nInjected: header"}, {**cookie, "expirationDate": float("nan")}, {**cookie, "name": "bad=name"}]:
|
||
with self.assertRaises(ValueError): validate_cookies([invalid])
|
||
|
||
def test_native_inspection_passes_cookies_only_to_worker_and_not_to_reply(self):
|
||
with tempfile.TemporaryDirectory() as folder, patch("host.main.ROOT", Path(folder)):
|
||
output = io.BytesIO(); host = NativeHost(output)
|
||
cookie = {"name": "SID", "value": "synthetic-session", "domain": ".youtube.com", "path": "/"}
|
||
try:
|
||
with patch.object(host, "worker", return_value={"title": "Test"}) as worker:
|
||
host.handle({"id": "inspect-auth", "action": "inspect", "payload": {"url": "https://youtu.be/jNQXAC9IVRw", "browserCookies": [cookie]}})
|
||
self.assertEqual(worker.call_args.args[1]["browserCookies"][0]["value"], cookie["value"])
|
||
self.assertNotIn(b"synthetic-session", output.getvalue())
|
||
self.assertNotIn("browserCookies", json.dumps(host.jobs))
|
||
finally: host.close()
|
||
|
||
def test_youtube_login_error_is_distinct_from_helper_connection_failure(self):
|
||
message = public_error(ValueError("ERROR: Sign in to confirm you’re not a bot"))
|
||
self.assertIn("Использовать вход в YouTube", message)
|
||
self.assertIn("Помощник работает", message)
|
||
self.assertNotIn("cookies-from-browser", message)
|
||
self.assertIn("отклонил вход", public_error(ValueError("Sign in to confirm"), True))
|
||
|
||
def test_youtube_urls_are_canonical_and_do_not_allow_arbitrary_hosts(self):
|
||
expected = "https://www.youtube.com/watch?v=jNQXAC9IVRw"
|
||
for value in [expected + "&list=other", "https://youtu.be/jNQXAC9IVRw?t=2", "https://www.youtube.com/shorts/jNQXAC9IVRw"]:
|
||
self.assertEqual(youtube_url(value), expected)
|
||
for value in ["https://youtube.com.evil.test/watch?v=jNQXAC9IVRw", "https://127.0.0.1/watch?v=jNQXAC9IVRw", "https://youtube.com@evil.test/watch?v=jNQXAC9IVRw", "file:///test", "https://www.youtube.com:444/watch?v=jNQXAC9IVRw", "https://www.youtube.com/@channel"]:
|
||
with self.assertRaises(ValueError): youtube_url(value)
|
||
|
||
def test_separate_audio_is_added_and_video_only_stays_single(self):
|
||
info = {"formats": [{"format_id": "18", "vcodec": "h264", "acodec": "aac", "height": 360},
|
||
{"format_id": "137", "vcodec": "h264", "acodec": "none", "height": 1080},
|
||
{"format_id": "140", "vcodec": "none", "acodec": "aac"}]}
|
||
base = {"mode": "video", "videoId": "137", "audioId": "best"}
|
||
self.assertEqual(select_streams(info, base), "137+140")
|
||
self.assertEqual(select_streams(info, {**base, "mode": "video_only"}), "137")
|
||
self.assertEqual(select_streams(info, {**base, "videoId": "18"}), "18")
|
||
self.assertEqual(select_streams(info, {**base, "mode": "audio"}), "140")
|
||
|
||
def test_download_does_not_accept_format_selector_injection(self):
|
||
with tempfile.TemporaryDirectory() as folder:
|
||
base = {"url": "https://youtu.be/jNQXAC9IVRw", "folder": folder, "videoId": "137+bestaudio"}
|
||
with self.assertRaises(ValueError): validate_download(base)
|
||
self.assertEqual(validate_download({**base, "videoId": "137"})["container"], "mp4")
|
||
with self.assertRaises(ValueError): validate_download({**base, "videoId": "137", "container": "exe"})
|
||
|
||
def test_gpu_is_used_only_when_video_needs_transcoding(self):
|
||
media = {"streams": [{"codec_type": "video", "codec_name": "vp9", "pix_fmt": "yuv420p"}, {"codec_type": "audio", "codec_name": "opus"}]}
|
||
options = {"mode": "video", "container": "mp4", "codec": "auto", "engine": "auto"}
|
||
plan = encoding_plan(options, media, ["h264", "hevc", "av1"])
|
||
self.assertEqual((plan["video"], plan["audio"], plan["engine"]), ("h264", "aac", "nvenc"))
|
||
copied = encoding_plan({**options, "container": "mkv"}, media, ["h264"])
|
||
self.assertEqual((copied["video"], copied["audio"], copied["engine"]), ("copy", "copy", "copy"))
|
||
cpu = encoding_plan(options, media, [])
|
||
self.assertEqual(cpu["engine"], "cpu")
|
||
with self.assertRaises(ValueError): encoding_plan({**options, "engine": "nvenc"}, media, [])
|
||
with self.assertRaises(ValueError): encoding_plan({**options, "codec": "copy"}, media, [])
|
||
|
||
def test_hdr_is_not_silently_converted_to_h264(self):
|
||
media = {"streams": [{"codec_type": "video", "codec_name": "vp9", "pix_fmt": "yuv420p10le", "color_transfer": "smpte2084"}, {"codec_type": "audio", "codec_name": "opus"}]}
|
||
options = {"mode": "video", "container": "mp4", "codec": "auto", "engine": "auto"}
|
||
self.assertEqual(encoding_plan(options, media, ["hevc"])["video"], "hevc")
|
||
with self.assertRaises(ValueError): encoding_plan({**options, "codec": "h264"}, media, ["h264"])
|
||
|
||
def test_amf_requires_measured_codec_and_10bit_support(self):
|
||
media = {"streams": [{"codec_type": "video", "codec_name": "vp9", "pix_fmt": "yuv420p"}, {"codec_type": "audio", "codec_name": "opus"}]}
|
||
options = {"mode": "video", "container": "mp4", "codec": "hevc", "engine": "auto", "preset": "balanced", "bitrate": 256}
|
||
self.assertEqual(encoding_plan(options, media, {"amf": ["hevc"]})["engine"], "amf")
|
||
self.assertEqual(encoding_plan(options, media, {"nvenc": ["hevc"], "amf": ["hevc"]})["engine"], "nvenc")
|
||
with self.assertRaisesRegex(ValueError, "AMD AMF"):
|
||
encoding_plan({**options, "engine": "amf"}, media, {})
|
||
media["streams"][0].update(pix_fmt="yuv420p10le", color_transfer="smpte2084")
|
||
self.assertEqual(encoding_plan(options, media, {"amf": ["hevc"]})["engine"], "cpu")
|
||
plan = encoding_plan({**options, "engine": "amf"}, media, {"amf": ["hevc"], "amf10bit": ["hevc"]})
|
||
self.assertTrue(plan["tenBit"])
|
||
with patch("host.media.tool", return_value="ffmpeg"):
|
||
command = ffmpeg_command("source.mkv", "out.mp4", options, plan)
|
||
self.assertIn("hevc_amf", command); self.assertIn("p010le", command); self.assertIn("main10", command)
|
||
self.assertIn("d3d11va", command); self.assertNotIn("cuda", command)
|
||
|
||
def test_audio_and_video_only_output_maps(self):
|
||
source = {"streams": [{"codec_type": "video", "codec_name": "h264"}, {"codec_type": "audio", "codec_name": "aac"}]}
|
||
options = {"mode": "video_only", "container": "mp4", "codec": "auto", "engine": "auto", "preset": "balanced", "bitrate": 256}
|
||
with patch("host.media.tool", return_value="ffmpeg"):
|
||
command = ffmpeg_command("source.mkv", "out.mp4", options, encoding_plan(options, source, []))
|
||
self.assertIn("-an", command)
|
||
self.assertNotIn("0:a:0", command)
|
||
audio = encoding_plan({**options, "mode": "audio", "container": "mp3"}, source, [])
|
||
self.assertEqual((audio["video"], audio["audio"], audio["engine"]), (None, "libmp3lame", "cpu"))
|
||
|
||
def test_protocol_uses_utf8_byte_lengths_and_handles_partial_reads(self):
|
||
message = {"text": "Русский путь 🎬"}
|
||
stream = io.BytesIO()
|
||
write_message(stream, message)
|
||
self.assertEqual(struct.unpack("<I", stream.getvalue()[:4])[0], len(stream.getvalue()) - 4)
|
||
class Partial(io.BytesIO):
|
||
def read(self, count=-1): return super().read(min(count, 2))
|
||
self.assertEqual(read_message(Partial(stream.getvalue())), message)
|
||
self.assertIsNone(read_message(io.BytesIO()))
|
||
with self.assertRaises(ValueError): read_message(io.BytesIO(struct.pack("<I", MAX_MESSAGE + 1)))
|
||
with self.assertRaises(EOFError): read_message(io.BytesIO(struct.pack("<I", 5) + b"{}"))
|
||
|
||
def test_filename_is_safe_and_existing_files_are_never_overwritten(self):
|
||
self.assertEqual(safe_filename("../CON:bad/name", "video"), "_CON_bad_name")
|
||
self.assertEqual(safe_filename("CON", "video"), "_CON")
|
||
with tempfile.TemporaryDirectory() as folder:
|
||
root = Path(folder); original = root / "video.mp4"; original.write_bytes(b"original")
|
||
source = root / "temp.mp4"; source.write_bytes(b"new")
|
||
saved = publish_file(source, root, "video", "mp4")
|
||
self.assertEqual(saved.name, "video (1).mp4")
|
||
self.assertEqual(original.read_bytes(), b"original")
|
||
self.assertEqual(saved.read_bytes(), b"new")
|
||
|
||
def test_dependency_install_refreshes_capabilities_before_reporting_completion(self):
|
||
with tempfile.TemporaryDirectory() as folder, patch("host.main.ROOT", Path(folder)):
|
||
host = NativeHost(io.BytesIO())
|
||
host._capabilities = {"ready": False}
|
||
host.jobs["install"] = {"id": "install", "status": "running", "kind": "dependency"}
|
||
try:
|
||
with patch.object(host, "worker", return_value={"title": "Node.js"}), patch.object(host, "event") as event:
|
||
event.side_effect = lambda *args, **kwargs: self.assertIsNone(host._capabilities)
|
||
host.run_job("install", {"dependency": "node"}, "dependency_install")
|
||
self.assertEqual(event.call_args.kwargs["status"], "complete")
|
||
with patch.object(host, "worker", side_effect=ValueError("install failed")), patch.object(host, "event") as event:
|
||
host._capabilities = {"ready": False}
|
||
event.side_effect = lambda *args, **kwargs: self.assertIsNone(host._capabilities)
|
||
host.run_job("install", {"dependency": "node"}, "dependency_install")
|
||
self.assertEqual(event.call_args.kwargs["status"], "error")
|
||
finally:
|
||
host.close()
|
||
|
||
def test_completed_file_can_be_revealed_after_native_host_restart(self):
|
||
with tempfile.TemporaryDirectory() as folder, patch("host.main.ROOT", Path(folder)):
|
||
saved = Path(folder) / "finished.mp4"
|
||
saved.write_bytes(b"media")
|
||
first = NativeHost(io.BytesIO())
|
||
first.jobs["completed"] = {"id": "completed", "status": "running", "createdAt": 1}
|
||
first.event("completed", status="complete", result={"path": str(saved)})
|
||
first.close()
|
||
output = io.BytesIO()
|
||
restarted = NativeHost(output)
|
||
try:
|
||
with patch("host.main.subprocess.Popen") as launch:
|
||
restarted.handle({"id": "reveal", "action": "reveal", "payload": {"jobId": "completed"}})
|
||
output.seek(0)
|
||
self.assertTrue(read_message(output)["ok"])
|
||
if os.name == "nt":
|
||
self.assertEqual(launch.call_args.args[0][-1], str(saved))
|
||
finally:
|
||
restarted.close()
|
||
|
||
|
||
if __name__ == "__main__": unittest.main()
|