from __future__ import annotations
|
|
import contextlib
|
import io
|
import json
|
import os
|
import pathlib
|
import shutil
|
import subprocess
|
import sys
|
import tempfile
|
import unittest
|
import importlib.util
|
from unittest import mock
|
|
PROJECT_DEV = pathlib.Path(__file__).resolve().parents[2]
|
if str(PROJECT_DEV) not in sys.path:
|
sys.path.insert(0, str(PROJECT_DEV))
|
|
from bili_authenticated_extension.constants import CANONICAL_URL, EXTENSION_BUILD, TARGET_BVID # noqa: E402
|
from bili_authenticated_extension.worker import ( # noqa: E402
|
CancelRequested,
|
HostConfig,
|
SubprocessPolicy,
|
WorkerError,
|
bootstrap_ytdlp,
|
build_cookie_stream,
|
close_cookie_stream,
|
merge_local_streams,
|
prepare_download_info,
|
probe_mkv,
|
remux_single_to_mkv,
|
run_authenticated_task,
|
run_frozen_bridge,
|
sha256_file,
|
validate_download_info,
|
verify_frozen_ytdlp,
|
ytdlp_options,
|
)
|
|
|
def valid_start() -> dict:
|
now_ms = int(__import__("time").time() * 1000)
|
return {
|
"schema": 2,
|
"type": "start",
|
"extension_build": EXTENSION_BUILD,
|
"target": TARGET_BVID,
|
"canonical_url": CANONICAL_URL,
|
"cookie_store_id": "0",
|
"prepare_id": "d" * 32,
|
"page_proof": {
|
"target": TARGET_BVID,
|
"canonical_url": CANONICAL_URL,
|
"task_nonce": "b" * 32,
|
"observed_at_unix_ms": now_ms,
|
"observed_duration_ms": 3_133_950,
|
"video_width": 1920,
|
"video_height": 1080,
|
"ready_state": 4,
|
"eme_present": False,
|
},
|
"cookies": [{
|
"name": "SYNTHETIC_COOKIE_NAME",
|
"value": "SYNTHETIC_COOKIE_VALUE",
|
"domain": ".bilibili.com",
|
"host_only": False,
|
"path": "/",
|
"secure": True,
|
"http_only": True,
|
"same_site": "unspecified",
|
"session": True,
|
"expiration_unix": None,
|
"store_id": "0",
|
"partition_key": None,
|
}],
|
}
|
|
|
class FakeYDL:
|
def __init__(self) -> None:
|
self.params = {}
|
self.extract_count = 0
|
self.processed_object = None
|
video = {
|
"format_id": "v",
|
"url": "https://example.invalid/video",
|
"protocol": "https",
|
"vcodec": "h264",
|
"acodec": "none",
|
"has_drm": False,
|
}
|
audio = {
|
"format_id": "a",
|
"url": "https://example.invalid/audio",
|
"protocol": "http_dash_segments",
|
"vcodec": "none",
|
"acodec": "aac",
|
"has_drm": False,
|
}
|
self.info = {
|
"id": TARGET_BVID,
|
"_type": "video",
|
"duration": 3133.95,
|
"is_live": False,
|
"has_drm": False,
|
"formats": [video, audio],
|
}
|
self.selected = {
|
"id": TARGET_BVID,
|
"vcodec": "h264",
|
"acodec": "aac",
|
"requested_formats": [video, audio],
|
}
|
|
def extract_info(self, url, download=False, process=True):
|
self.extract_count += 1
|
assert url == CANONICAL_URL and download is False and process is True
|
return self.info
|
|
def build_format_selector(self, spec):
|
assert spec == "bestvideo+bestaudio/best"
|
return object()
|
|
def _get_formats(self, info):
|
assert info is self.info
|
return info["formats"]
|
|
def _select_formats(self, formats, selector):
|
del formats, selector
|
yield self.selected
|
|
@staticmethod
|
def _copy_infodict(info):
|
return dict(info)
|
|
def process_info(self, info):
|
self.processed_object = info
|
|
|
class WorkerContractTests(unittest.TestCase):
|
@classmethod
|
def setUpClass(cls) -> None:
|
cls.ffmpeg = pathlib.Path(shutil.which("ffmpeg") or "")
|
cls.ffprobe = pathlib.Path(shutil.which("ffprobe") or "")
|
|
@unittest.skipUnless(importlib.util.find_spec("yt_dlp"), "frozen yt-dlp environment is required")
|
def test_frozen_real_package_and_plugins(self) -> None:
|
root = verify_frozen_ytdlp()
|
self.assertTrue(root.is_dir())
|
yt_dlp, globals_module = bootstrap_ytdlp()
|
self.assertEqual("2026.07.04", yt_dlp.version.__version__)
|
self.assertEqual([], globals_module.plugin_dirs.value)
|
self.assertEqual({}, globals_module.plugin_ies.value)
|
self.assertEqual({}, globals_module.plugin_pps.value)
|
|
@unittest.skipUnless(importlib.util.find_spec("yt_dlp"), "frozen yt-dlp environment is required")
|
def test_default_plugin_directory_sentinel_cannot_load(self) -> None:
|
with tempfile.TemporaryDirectory() as temporary:
|
root = pathlib.Path(temporary)
|
plugin_dir = root / "yt_dlp_plugins" / "extractor"
|
plugin_dir.mkdir(parents=True)
|
marker = root / "loaded.marker"
|
plugin = plugin_dir / "sentinel.py"
|
plugin.write_text(
|
"from pathlib import Path\n"
|
f"Path({str(marker)!r}).write_text('loaded', encoding='ascii')\n",
|
encoding="utf-8",
|
)
|
environment = os.environ.copy()
|
environment["PYTHONPATH"] = os.pathsep.join((str(root), str(PROJECT_DEV)))
|
environment.pop("YTDLP_NO_PLUGINS", None)
|
result = subprocess.run(
|
[sys.executable, "-c", "from bili_authenticated_extension.worker import bootstrap_ytdlp; bootstrap_ytdlp()"],
|
check=False,
|
stdout=subprocess.PIPE,
|
stderr=subprocess.PIPE,
|
env=environment,
|
)
|
self.assertEqual(0, result.returncode, result.stderr)
|
self.assertFalse(marker.exists())
|
self.assertEqual(b"", result.stdout)
|
self.assertEqual(b"", result.stderr)
|
|
@unittest.skipUnless(importlib.util.find_spec("yt_dlp"), "frozen yt-dlp environment is required")
|
def test_real_cookiejar_uses_memory_stream_and_sentinel_is_not_output(self) -> None:
|
yt_dlp, _ = bootstrap_ytdlp()
|
start = valid_start()
|
stream = build_cookie_stream(start)
|
self.assertIsInstance(stream, io.StringIO)
|
capture_out, capture_error = io.StringIO(), io.StringIO()
|
with contextlib.redirect_stdout(capture_out), contextlib.redirect_stderr(capture_error):
|
with yt_dlp.YoutubeDL({"cookiefile": stream, "logger": type("L", (), {"debug": lambda *_: None, "info": lambda *_: None, "warning": lambda *_: None, "error": lambda *_: None})()}) as ydl:
|
names = {cookie.name for cookie in ydl.cookiejar}
|
self.assertIn("SYNTHETIC_COOKIE_NAME", names)
|
close_cookie_stream(stream)
|
self.assertTrue(stream.closed)
|
self.assertNotIn("SYNTHETIC_COOKIE_VALUE", capture_out.getvalue() + capture_error.getvalue())
|
|
def test_one_extract_same_object_and_second_extract_guard(self) -> None:
|
ydl = FakeYDL()
|
def resolver(leaf, _params):
|
return type("HttpFD" if leaf["protocol"] == "https" else "DashSegmentsFD", (), {})
|
|
info, single, signed_urls = prepare_download_info(ydl, downloader_resolver=resolver)
|
self.assertFalse(single)
|
self.assertEqual(2, len(signed_urls))
|
ydl.process_info(info)
|
self.assertIs(info, ydl.processed_object)
|
self.assertEqual(1, ydl.extract_count)
|
with self.assertRaises(WorkerError) as caught:
|
ydl.extract_info(CANONICAL_URL)
|
self.assertEqual("E_SECOND_EXTRACT", caught.exception.code)
|
|
@unittest.skipUnless(importlib.util.find_spec("yt_dlp"), "frozen yt-dlp environment is required")
|
def test_real_downloader_dispatch_and_rejections(self) -> None:
|
bootstrap_ytdlp()
|
valid = FakeYDL().selected | {"id": TARGET_BVID}
|
leaves, single = validate_download_info(valid, {})
|
self.assertEqual(2, len(leaves))
|
self.assertFalse(single)
|
for protocol in ("m3u8", "m3u8_native", "rtmp", "unknown"):
|
changed = json.loads(json.dumps(valid))
|
changed["requested_formats"][0]["protocol"] = protocol
|
with self.subTest(protocol=protocol), self.assertRaises(WorkerError):
|
validate_download_info(changed, {})
|
changed = json.loads(json.dumps(valid))
|
changed["requested_formats"][0]["url"] = "http://example.invalid/video"
|
with self.assertRaises(WorkerError):
|
validate_download_info(changed, {})
|
|
def test_subprocess_policy_rejects_url_header_secret_and_nonlocal(self) -> None:
|
with tempfile.TemporaryDirectory() as temporary:
|
root = pathlib.Path(temporary).resolve()
|
local = root / "input.mp4"
|
local.write_bytes(b"fixture")
|
destination = root / "output.mkv"
|
policy = SubprocessPolicy(root, {self.ffmpeg}, {"SYNTHETIC_SECRET"})
|
policy.check("subprocess.Popen", (str(self.ffmpeg), [str(self.ffmpeg), "-i", str(local), str(destination)], None, {}))
|
for argv in (
|
[str(self.ffmpeg), "-i", "https://example.invalid/signed", str(destination)],
|
[str(self.ffmpeg), "-headers", "SYNTHETIC_SECRET", "-i", str(local), str(destination)],
|
[str(self.ffmpeg), "-i", str(root.parent / "outside.mp4"), str(destination)],
|
):
|
with self.subTest(argv=argv), self.assertRaises(WorkerError):
|
policy.check("subprocess.Popen", (str(self.ffmpeg), argv, None, {}))
|
|
def test_real_audit_hook_blocks_before_createprocess(self) -> None:
|
with tempfile.TemporaryDirectory() as temporary:
|
root = pathlib.Path(temporary)
|
marker = root / "child.marker"
|
script = root / "audit_test.py"
|
child_code = f"import pathlib; pathlib.Path({str(marker)!r}).write_text('ran')"
|
script.write_text(
|
"import pathlib, subprocess, sys\n"
|
"from bili_authenticated_extension.worker import SubprocessPolicy, WorkerError\n"
|
f"root=pathlib.Path({str(root)!r})\n"
|
"SubprocessPolicy(root,{pathlib.Path(sys.executable)}).install()\n"
|
"try:\n"
|
f" subprocess.run([sys.executable,'-c',{child_code!r},'https://example.invalid/signed'])\n"
|
"except WorkerError:\n"
|
" pass\n"
|
"else:\n"
|
" raise SystemExit(9)\n",
|
encoding="utf-8",
|
)
|
environment = os.environ.copy()
|
environment["PYTHONPATH"] = str(PROJECT_DEV)
|
result = subprocess.run(
|
[sys.executable, str(script)],
|
check=False,
|
stdout=subprocess.PIPE,
|
stderr=subprocess.PIPE,
|
env=environment,
|
)
|
self.assertEqual(0, result.returncode, result.stderr)
|
self.assertFalse(marker.exists())
|
self.assertNotIn(b"signed", result.stdout + result.stderr)
|
|
def test_frozen_retry_and_output_options(self) -> None:
|
config = HostConfig(
|
ffmpeg=self.ffmpeg,
|
ffprobe=self.ffprobe,
|
bridge_python=pathlib.Path(sys.executable),
|
bridge_script=PROJECT_DEV / "bili_video_download_bridge.py",
|
batch_json=PROJECT_DEV / "test" / "test_bili_video_download_bridge.py",
|
yt_dlp_executable=pathlib.Path(sys.executable),
|
destination=PROJECT_DEV,
|
)
|
stream = io.StringIO()
|
try:
|
options = ytdlp_options(config, PROJECT_DEV, stream, lambda _: None)
|
self.assertEqual("bestvideo+bestaudio/best", options["format"])
|
self.assertEqual("mkv", options["merge_output_format"])
|
self.assertEqual(20, options["socket_timeout"])
|
self.assertEqual(1, options["extractor_retries"])
|
self.assertEqual(2, options["retries"])
|
self.assertEqual(2, options["fragment_retries"])
|
self.assertEqual(1, options["file_access_retries"])
|
self.assertEqual({}, options["external_downloader"])
|
for sleeper in options["retry_sleep_functions"].values():
|
self.assertEqual(1, sleeper(99))
|
finally:
|
stream.close()
|
|
def test_bridge_result_must_match_persisted_mapping_and_complete_mkv(self) -> None:
|
with tempfile.TemporaryDirectory() as temporary:
|
root = pathlib.Path(temporary)
|
destination = root / "destination"
|
destination.mkdir()
|
formal = destination / f"{TARGET_BVID}.mkv"
|
formal.write_bytes(b"synthetic-mkv")
|
formal_sha = sha256_file(formal)
|
mapping_path = destination / f"{TARGET_BVID}.download.json"
|
persisted = {
|
"bvid": TARGET_BVID,
|
"source": CANONICAL_URL,
|
"local_file": formal.name,
|
"bytes": formal.stat().st_size,
|
"sha256": formal_sha,
|
"duration_seconds": 3133.95,
|
"acquisition_mode": "authorized_browser_file_handoff",
|
}
|
mapping_path.write_text(json.dumps(persisted), encoding="utf-8")
|
item = {"status": "COMPLETE", **persisted}
|
output = json.dumps({"result": "PASS", "items": [item]}).encode("utf-8")
|
config = HostConfig(
|
ffmpeg=self.ffmpeg,
|
ffprobe=self.ffprobe,
|
bridge_python=pathlib.Path(sys.executable),
|
bridge_script=PROJECT_DEV / "bili_video_download_bridge.py",
|
batch_json=PROJECT_DEV / "test" / "test_bili_video_download_bridge.py",
|
yt_dlp_executable=pathlib.Path(sys.executable),
|
destination=destination,
|
)
|
completed = subprocess.CompletedProcess([], 0, stdout=output, stderr=b"")
|
with mock.patch("bili_authenticated_extension.worker._run_local", return_value=completed):
|
self.assertEqual((formal.name, mapping_path.name), run_frozen_bridge(config, root / "candidate.mkv"))
|
persisted["sha256"] = "0" * 64
|
mapping_path.write_text(json.dumps(persisted), encoding="utf-8")
|
with self.assertRaises(WorkerError):
|
run_frozen_bridge(config, root / "candidate.mkv")
|
|
@unittest.skipUnless(shutil.which("ffmpeg") and shutil.which("ffprobe"), "FFmpeg is required")
|
def test_real_synthetic_single_and_dual_local_mkv_paths(self) -> None:
|
with tempfile.TemporaryDirectory() as temporary:
|
root = pathlib.Path(temporary).resolve()
|
combined = root / "combined.mp4"
|
video = root / "video.mp4"
|
audio = root / "audio.m4a"
|
commands = [
|
[str(self.ffmpeg), "-v", "error", "-f", "lavfi", "-i", "color=c=blue:s=320x180:r=10", "-f", "lavfi", "-i", "sine=frequency=440:sample_rate=48000", "-t", "1", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-c:a", "aac", str(combined)],
|
[str(self.ffmpeg), "-v", "error", "-f", "lavfi", "-i", "color=c=red:s=320x180:r=10", "-t", "1", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-an", str(video)],
|
[str(self.ffmpeg), "-v", "error", "-f", "lavfi", "-i", "sine=frequency=880:sample_rate=48000", "-t", "1", "-c:a", "aac", str(audio)],
|
]
|
for command in commands:
|
subprocess.run(command, check=True, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
single = root / "single.mkv"
|
dual = root / "dual.mkv"
|
remux_single_to_mkv(self.ffmpeg, combined, single)
|
merge_local_streams(self.ffmpeg, video, audio, dual)
|
probe_mkv(self.ffprobe, single)
|
probe_mkv(self.ffprobe, dual)
|
self.assertGreater(single.stat().st_size, 0)
|
self.assertGreater(dual.stat().st_size, 0)
|
|
def test_cancel_checkpoints_and_commit_arbitration(self) -> None:
|
class TaskYDL(FakeYDL):
|
def __init__(self, run_root: pathlib.Path) -> None:
|
super().__init__()
|
self.run_root = run_root
|
|
def __enter__(self):
|
return self
|
|
def __exit__(self, *_args):
|
return False
|
|
def process_info(self, info):
|
super().process_info(info)
|
(self.run_root / "download.mp4").write_bytes(b"synthetic")
|
|
class YtDlpModule:
|
def __init__(self, run_root: pathlib.Path) -> None:
|
self.run_root = run_root
|
|
def YoutubeDL(self, _options):
|
return TaskYDL(self.run_root)
|
|
def config_for(destination: pathlib.Path) -> HostConfig:
|
return HostConfig(
|
ffmpeg=self.ffmpeg,
|
ffprobe=self.ffprobe,
|
bridge_python=pathlib.Path(sys.executable),
|
bridge_script=PROJECT_DEV / "bili_video_download_bridge.py",
|
batch_json=PROJECT_DEV / "test" / "test_bili_video_download_bridge.py",
|
yt_dlp_executable=pathlib.Path(sys.executable),
|
destination=destination,
|
)
|
|
for point in ("after-probe", "at-commit", "inside-bridge"):
|
with self.subTest(point=point), tempfile.TemporaryDirectory() as temporary:
|
root = pathlib.Path(temporary)
|
stage = root / "stage"
|
destination = root / "destination"
|
destination.mkdir()
|
canceled = False
|
bridge_calls = 0
|
|
def remux(_ffmpeg, _source, output):
|
output.write_bytes(b"synthetic-mkv")
|
|
def probe(_ffprobe, _candidate):
|
nonlocal canceled
|
if point == "after-probe":
|
canceled = True
|
|
def commit_begin():
|
nonlocal canceled
|
if point == "at-commit":
|
canceled = True
|
|
def bridge(_config, _candidate):
|
nonlocal bridge_calls, canceled
|
bridge_calls += 1
|
if point == "inside-bridge":
|
canceled = True
|
(destination / f"{TARGET_BVID}.mkv").write_bytes(b"formal")
|
(destination / f"{TARGET_BVID}.download.json").write_text("{}", encoding="utf-8")
|
return (f"{TARGET_BVID}.mkv", f"{TARGET_BVID}.download.json")
|
|
with (
|
mock.patch(
|
"bili_authenticated_extension.worker.bootstrap_ytdlp",
|
return_value=(YtDlpModule(stage / "run-placeholder"), object()),
|
) as bootstrap,
|
mock.patch("bili_authenticated_extension.worker.remux_single_to_mkv", side_effect=remux),
|
mock.patch("bili_authenticated_extension.worker.probe_mkv", side_effect=probe),
|
mock.patch("bili_authenticated_extension.worker.run_frozen_bridge", side_effect=bridge),
|
mock.patch("bili_authenticated_extension.worker.SubprocessPolicy.install"),
|
mock.patch(
|
"bili_authenticated_extension.worker.prepare_download_info",
|
return_value=({"id": TARGET_BVID}, True, []),
|
),
|
):
|
# The fake downloader needs the actual run directory, which is created by the product helper.
|
def module_factory():
|
run_directory = next(stage.glob("run-*"))
|
return YtDlpModule(run_directory), object()
|
|
bootstrap.side_effect = module_factory
|
if point in {"after-probe", "at-commit"}:
|
with self.assertRaises(CancelRequested):
|
run_authenticated_task(
|
valid_start(),
|
config_for(destination),
|
cancel_check=lambda: canceled,
|
report=lambda *_: None,
|
stage_root=stage,
|
commit_begin=commit_begin,
|
)
|
self.assertEqual(0, bridge_calls)
|
self.assertEqual([], list(destination.iterdir()))
|
else:
|
result = run_authenticated_task(
|
valid_start(),
|
config_for(destination),
|
cancel_check=lambda: canceled,
|
report=lambda *_: None,
|
stage_root=stage,
|
commit_begin=commit_begin,
|
)
|
self.assertEqual(1, bridge_calls)
|
self.assertEqual(f"{TARGET_BVID}.mkv", result[0])
|
self.assertEqual(2, len(list(destination.iterdir())))
|
|
|
if __name__ == "__main__":
|
unittest.main()
|
