import json import subprocess import tempfile import unittest from concurrent.futures import Future from pathlib import Path from yt_dlp import YoutubeDL from controller.config import ControllerConfig, DownloadConfig, StorageConfig from controller.downloads import ( DownloadManager, mp4_export_format_selector, safe_export_filename, ) from controller.store import ControllerStore class ImmediateExecutor: def submit(self, function, *args, **kwargs): future = Future() try: future.set_result(function(*args, **kwargs)) except BaseException as exc: future.set_exception(exc) return future class QueuedExecutor: def __init__(self): self.calls = [] def submit(self, function, *args, **kwargs): future = Future() self.calls.append((future, function, args, kwargs)) return future class FakeExportRunner: def __init__( self, *, video_codec="h264", audio_codec="aac", pixel_format="yuv420p", container="mp4", height=1080, playlist_items=1, fail_download=False, ): self.video_codec = video_codec self.audio_codec = audio_codec self.pixel_format = pixel_format self.container = container self.height = height self.playlist_items = playlist_items self.fail_download = fail_download self.calls = [] def __call__(self, arguments, **kwargs): arguments = list(arguments) self.calls.append((arguments, dict(kwargs))) program = Path(arguments[0]).name if program == "yt-dlp-test": if self.fail_download: raise subprocess.CalledProcessError(9, arguments) template = Path(arguments[arguments.index("-o") + 1]) template.parent.mkdir(parents=True, exist_ok=True) if "--audio-format" in arguments: extension = "." + arguments[arguments.index("--audio-format") + 1] else: extension = { "mp4": ".mp4", "matroska": ".mkv", "webm": ".webm", }[self.container] for index in range(1, self.playlist_items + 1): (template.parent / f"{index:03d} - Tést Item {index} [id{index}]{extension}").write_bytes( f"source-{index}".encode() ) return subprocess.CompletedProcess(arguments, 0, stdout="", stderr="") if program == "ffmpeg": Path(arguments[-1]).write_bytes(b"converted") return subprocess.CompletedProcess(arguments, 0, stdout="", stderr="") if program == "ffprobe": path = Path(arguments[-1]) converted = path.is_file() and path.read_bytes() == b"converted" if path.suffix == ".m4a": payload = { "streams": [{"codec_type": "audio", "codec_name": "aac"}], "format": {"format_name": "mov,mp4,m4a,3gp,3g2,mj2"}, } elif path.suffix == ".mp3": payload = { "streams": [{"codec_type": "audio", "codec_name": "mp3"}], "format": {"format_name": "mp3"}, } else: payload = { "streams": [ { "codec_type": "video", "codec_name": "h264" if converted else self.video_codec, "pix_fmt": "yuv420p" if converted else self.pixel_format, "width": 1920 if self.height == 1080 else 1280, "height": self.height, }, { "codec_type": "audio", "codec_name": "aac" if converted else self.audio_codec, }, ], "format": { "format_name": "mov,mp4,m4a,3gp,3g2,mj2" if converted or self.container == "mp4" else self.container }, } return subprocess.CompletedProcess( arguments, 0, stdout=json.dumps(payload), stderr="", ) raise AssertionError(f"unexpected command: {arguments!r}") @property def yt_dlp_calls(self): return [arguments for arguments, _ in self.calls if Path(arguments[0]).name == "yt-dlp-test"] @property def ffmpeg_calls(self): return [arguments for arguments, _ in self.calls if Path(arguments[0]).name == "ffmpeg"] class UtilityExportTests(unittest.TestCase): def setUp(self): self.tempdir = tempfile.TemporaryDirectory() base = Path(self.tempdir.name) self.media = base / "media" self.data = base / "data" self.config = ControllerConfig( storage=StorageConfig(media_base=self.media, data_base=self.data), downloads=DownloadConfig( yt_dlp_binary="yt-dlp-test", cookie_path=None, max_workers=1, process_timeout_seconds=30, max_height=1080, export_directory=Path("exports"), export_max_height=1080, max_queued_exports=5, ), ) self.store = ControllerStore(self.config.storage).open() def tearDown(self): self.store.close() self.tempdir.cleanup() def manager(self, runner, events=None): return DownloadManager( self.config, self.store, runner=runner, executor=ImmediateExecutor(), export_executor=ImmediateExecutor(), on_event=(events if events is not None else []).append, ) def assert_player_slots_empty(self): self.assertTrue( all(slot.media_state == "empty" for slot in self.store.list_slots()) ) def test_compatible_mp4_exports_without_slot_or_transcode(self): events = [] runner = FakeExportRunner() result = self.manager(runner, events).enqueue_export( "https://www.youtube.com/watch?v=one", "mp4", "Alice", "Slides" ) self.assertEqual(result.status, "queued") job = self.store.get_export_job(result.job_id) self.assertEqual(job.status, "completed") self.assertEqual(len(job.output_paths), 1) self.assertTrue(job.output_paths[0].startswith(f"exports/video/{result.job_id}/")) self.assertTrue(job.output_paths[0].endswith(".mp4")) self.assertTrue((self.media / job.output_paths[0]).is_file()) self.assertEqual(events, []) self.assertEqual(runner.ffmpeg_calls, []) self.assert_player_slots_empty() arguments = runner.yt_dlp_calls[0] self.assertIn("--yes-playlist", arguments) self.assertIn("--windows-filenames", arguments) self.assertIn("--no-overwrites", arguments) self.assertEqual( arguments[arguments.index("-f") + 1], "bv*[height<=1080]+ba/b[height<=1080]", ) self.assertEqual( arguments[arguments.index("-S") + 1], "res:1080,vcodec:h264,acodec:aac", ) def test_incompatible_mp4_transcodes_only_incompatible_codecs(self): runner = FakeExportRunner( video_codec="vp9", audio_codec="opus", pixel_format="yuv420p", container="webm", ) result = self.manager(runner).enqueue_export( "https://www.youtube.com/watch?v=vp9", "mp4", None, None ) self.assertEqual(self.store.get_export_job(result.job_id).status, "completed") self.assertEqual(len(runner.ffmpeg_calls), 1) arguments = runner.ffmpeg_calls[0] self.assertIn("libx264", arguments) self.assertIn("-crf", arguments) self.assertEqual(arguments[arguments.index("-crf") + 1], "18") self.assertIn("yuv420p", arguments) self.assertIn("aac", arguments) self.assertIn("192k", arguments) self.assertNotIn("-vf", arguments) self.assertNotIn("scale", " ".join(arguments)) self.assert_player_slots_empty() def test_compatible_h264_aac_in_other_container_is_remuxed_only(self): runner = FakeExportRunner(container="matroska") result = self.manager(runner).enqueue_export( "https://www.youtube.com/watch?v=mkv", "mp4", None, None ) self.assertEqual(self.store.get_export_job(result.job_id).status, "completed") arguments = runner.ffmpeg_calls[0] self.assertEqual(arguments[arguments.index("-c:v") + 1], "copy") self.assertEqual(arguments[arguments.index("-c:a") + 1], "copy") self.assertNotIn("libx264", arguments) def test_lower_resolution_mp4_is_preserved_without_scaling(self): runner = FakeExportRunner(height=720) result = self.manager(runner).enqueue_export( "https://www.youtube.com/watch?v=720", "mp4", None, None ) self.assertEqual(self.store.get_export_job(result.job_id).status, "completed") self.assertEqual(runner.ffmpeg_calls, []) def test_over_height_mp4_is_rejected_even_if_downloader_misbehaves(self): runner = FakeExportRunner(height=2160) result = self.manager(runner).enqueue_export( "https://www.youtube.com/watch?v=2160", "mp4", None, None ) job = self.store.get_export_job(result.job_id) self.assertEqual(job.status, "failed") self.assertIn("maximum height", job.error) self.assertFalse((self.media / "exports" / "video" / result.job_id).exists()) self.assert_player_slots_empty() def test_audio_exports_use_intentional_audio_only_strategies(self): for mode, expected_selector in ( ("m4a", "ba/b"), ("mp3", "ba/b"), ): with self.subTest(mode=mode): events = [] runner = FakeExportRunner() result = self.manager(runner, events).enqueue_export( f"https://www.youtube.com/watch?v={mode}", mode, None, None ) job = self.store.get_export_job(result.job_id) self.assertEqual(job.status, "completed") self.assertEqual(len(job.output_paths), 1) self.assertTrue(job.output_paths[0].endswith(f".{mode}")) arguments = runner.yt_dlp_calls[0] self.assertEqual(arguments[arguments.index("-f") + 1], expected_selector) self.assertIn("-x", arguments) self.assertEqual( arguments[arguments.index("--audio-format") + 1], mode ) self.assertEqual( arguments[arguments.index("--audio-quality") + 1], "0" ) if mode == "m4a": self.assertEqual( arguments[arguments.index("-S") + 1], "abr,acodec:aac", ) self.assertEqual(runner.ffmpeg_calls, []) self.assertEqual(events, []) self.assert_player_slots_empty() def test_playlist_items_publish_as_individual_collision_safe_files(self): runner = FakeExportRunner(playlist_items=3) result = self.manager(runner).enqueue_export( "https://www.youtube.com/playlist?list=three", "m4a", None, None ) job = self.store.get_export_job(result.job_id) self.assertEqual(job.status, "completed") self.assertEqual(len(job.output_paths), 3) self.assertEqual(len(set(job.output_paths)), 3) self.assertTrue(all((self.media / path).is_file() for path in job.output_paths)) def test_export_failure_and_expiry_never_mutate_player_or_export_files(self): slot_one = self.media / "slot1" slot_one.mkdir(parents=True) protected_file = slot_one / "active.mp4" protected_file.write_bytes(b"active") failed = self.manager(FakeExportRunner(fail_download=True)).enqueue_export( "https://www.youtube.com/watch?v=fail", "mp3", None, None ) self.assertEqual(self.store.get_export_job(failed.job_id).status, "failed") self.assertEqual(protected_file.read_bytes(), b"active") self.assert_player_slots_empty() completed = self.manager(FakeExportRunner()).enqueue_export( "https://www.youtube.com/watch?v=keep", "m4a", None, None ) export_path = self.media / self.store.get_export_job(completed.job_id).output_paths[0] self.manager(FakeExportRunner()).expire_due_slots() self.assertTrue(export_path.is_file()) self.assertEqual(protected_file.read_bytes(), b"active") def test_symlinked_export_root_cannot_redirect_into_player_slot(self): slot_one = self.media / "slot1" slot_one.mkdir(parents=True) protected_file = slot_one / "active.mp4" protected_file.write_bytes(b"active") (self.media / "exports").symlink_to(slot_one, target_is_directory=True) result = self.manager(FakeExportRunner()).enqueue_export( "https://www.youtube.com/watch?v=escape", "m4a", None, None ) self.assertEqual(self.store.get_export_job(result.job_id).status, "failed") self.assertIn("unsafe export directory", self.store.get_export_job(result.job_id).error) self.assertEqual(protected_file.read_bytes(), b"active") def test_export_work_uses_dedicated_executor_capacity(self): player_executor = QueuedExecutor() export_executor = QueuedExecutor() manager = DownloadManager( self.config, self.store, runner=FakeExportRunner(), executor=player_executor, export_executor=export_executor, ) manager.enqueue_export( "https://www.youtube.com/watch?v=export", "mp4", None, None ) manager.enqueue( "https://www.youtube.com/watch?v=player", None, None ) self.assertEqual(len(export_executor.calls), 1) self.assertEqual(len(player_executor.calls), 1) def test_restart_reconciles_published_export_and_cleans_manifest(self): runner = FakeExportRunner() manager = self.manager(runner) job_id = "a" * 32 self.store.reserve_export_job( "https://www.youtube.com/watch?v=recover", "mp4", None, None, max_active=5, job_id=job_id, ) self.store.mark_export_running(job_id) self.store.fail_export_job(job_id, "database completion failed") work = self.media / "exports" / ".work" / job_id published = self.media / "exports" / "video" / job_id work.mkdir(parents=True) published.mkdir(parents=True) filename = "Recovered.mp4" (published / filename).write_bytes(b"compatible") manager._write_export_manifest( work, job_id=job_id, mode="mp4", filenames=[filename], ) manager.reconcile_export_artifacts() job = self.store.get_export_job(job_id) self.assertEqual(job.status, "completed") self.assertEqual( job.output_paths, (f"exports/video/{job_id}/{filename}",), ) self.assertFalse(work.exists()) self.assertTrue((published / filename).is_file()) def test_restart_atomically_publishes_manifested_ready_export(self): runner = FakeExportRunner() manager = self.manager(runner) job_id = "b" * 32 self.store.reserve_export_job( "https://www.youtube.com/watch?v=staged", "m4a", None, None, max_active=5, job_id=job_id, ) self.store.mark_export_running(job_id) self.store.fail_export_job(job_id, "interrupted after conversion") work = self.media / "exports" / ".work" / job_id ready = work / "ready" ready.mkdir(parents=True) filename = "Recovered.m4a" (ready / filename).write_bytes(b"aac") manager._write_export_manifest( work, job_id=job_id, mode="m4a", filenames=[filename], ) manager.reconcile_export_artifacts() job = self.store.get_export_job(job_id) output = self.media / job.output_paths[0] self.assertEqual(job.status, "completed") self.assertTrue(output.is_file()) self.assertFalse(work.exists()) def test_restart_discards_only_known_incomplete_export_staging(self): manager = self.manager(FakeExportRunner()) job_id = "c" * 32 self.store.reserve_export_job( "https://www.youtube.com/watch?v=incomplete", "mp3", None, None, max_active=5, job_id=job_id, ) self.store.fail_export_job(job_id, "interrupted download") work = self.media / "exports" / ".work" / job_id work.mkdir(parents=True) (work / "partial.part").write_bytes(b"partial") unknown = self.media / "exports" / ".work" / ("d" * 32) unknown.mkdir() (unknown / "operator-file").write_bytes(b"keep") manager.reconcile_export_artifacts() self.assertFalse(work.exists()) self.assertTrue((unknown / "operator-file").is_file()) def test_safe_filename_blocks_traversal_preserves_unicode_and_deduplicates(self): used = set() first = safe_export_filename( "../../slot1/CON:Tést\x00?", ".mp4", used_casefolded=used ) second = safe_export_filename( "../../slot1/CON:Tést\x00?", ".mp4", used_casefolded=used ) long_name = safe_export_filename("🎬" * 500, ".mp3") self.assertNotIn("/", first) self.assertNotIn("\\", first) self.assertNotIn("..", Path(first).parts) self.assertIn("Tést", first) self.assertNotEqual(first.casefold(), second.casefold()) self.assertTrue(second.endswith(" (2).mp4")) self.assertLessEqual(len(Path(long_name).stem.encode("utf-8")), 180) def test_mp4_sorting_keeps_resolution_ahead_of_codec_preference(self): formats = [ { "format_id": "av1-1080", "url": "https://example.test/av1", "ext": "webm", "height": 1080, "width": 1920, "vcodec": "av01.0.08M.08", "acodec": "none", "tbr": 4000, "filesize": 1, }, { "format_id": "h264-720", "url": "https://example.test/h264", "ext": "mp4", "height": 720, "width": 1280, "vcodec": "avc1.64001f", "acodec": "none", "tbr": 3000, "filesize": 1, }, { "format_id": "aac", "url": "https://example.test/aac", "ext": "m4a", "vcodec": "none", "acodec": "mp4a.40.2", "tbr": 128, "filesize": 1, }, ] with YoutubeDL( { "quiet": True, "format_sort": ["res:1080", "vcodec:h264", "acodec:aac"], } ) as ydl: info = {"formats": formats} ydl.sort_formats(info) selector = ydl.build_format_selector(mp4_export_format_selector(1080)) selected = list(ydl._select_formats(info["formats"], selector)) self.assertEqual(selected[0]["height"], 1080) self.assertEqual(selected[0]["requested_formats"][0]["format_id"], "av1-1080") def test_mp4_sorting_prefers_native_h264_aac_at_equal_quality(self): formats = [ { "format_id": "vp9-1080", "url": "https://example.test/vp9", "ext": "webm", "height": 1080, "width": 1920, "vcodec": "vp9", "acodec": "none", "tbr": 4200, "filesize": 1, }, { "format_id": "h264-1080", "url": "https://example.test/h264", "ext": "mp4", "height": 1080, "width": 1920, "vcodec": "avc1.640028", "acodec": "none", "tbr": 4000, "filesize": 1, }, { "format_id": "opus", "url": "https://example.test/opus", "ext": "webm", "vcodec": "none", "acodec": "opus", "tbr": 160, "filesize": 1, }, { "format_id": "aac", "url": "https://example.test/aac", "ext": "m4a", "vcodec": "none", "acodec": "mp4a.40.2", "tbr": 128, "filesize": 1, }, ] with YoutubeDL( { "quiet": True, "format_sort": ["res:1080", "vcodec:h264", "acodec:aac"], } ) as ydl: info = {"formats": formats} ydl.sort_formats(info) selector = ydl.build_format_selector(mp4_export_format_selector(1080)) selected = list(ydl._select_formats(info["formats"], selector)) self.assertEqual( [item["format_id"] for item in selected[0]["requested_formats"]], ["h264-1080", "aac"], ) def test_m4a_sorting_keeps_audio_quality_ahead_of_native_aac_preference(self): formats = [ { "format_id": "aac-low", "url": "https://example.test/aac-low", "ext": "m4a", "vcodec": "none", "acodec": "mp4a.40.2", "abr": 48, "tbr": 48, "filesize": 1, }, { "format_id": "opus-high", "url": "https://example.test/opus-high", "ext": "webm", "vcodec": "none", "acodec": "opus", "abr": 160, "tbr": 160, "filesize": 1, }, ] with YoutubeDL( {"quiet": True, "format_sort": ["abr", "acodec:aac"]} ) as ydl: info = {"formats": formats} ydl.sort_formats(info) selected = list( ydl._select_formats( info["formats"], ydl.build_format_selector("ba/b"), ) ) self.assertEqual(selected[0]["format_id"], "opus-high") if __name__ == "__main__": unittest.main()