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()