import os import stat import subprocess import tempfile import threading import time 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, canonical_playlist_path, player_format_selector, ) from controller.models import SlotProtectedError 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 def run_next(self): future, function, args, kwargs = self.calls.pop(0) try: future.set_result(function(*args, **kwargs)) except BaseException as exc: future.set_exception(exc) return future class FakeYtDlp: def __init__( self, media_base: Path, *, fail_download=False, no_media=False, rewrite_cookie=False, ): self.media_base = media_base self.fail_download = fail_download self.no_media = no_media self.rewrite_cookie = rewrite_cookie self.calls = [] self.cookie_observations = [] def __call__(self, arguments, **kwargs): self.calls.append((list(arguments), dict(kwargs))) if "--cookies" in arguments: cookie_path = Path(arguments[arguments.index("--cookies") + 1]) self.cookie_observations.append( ( cookie_path, cookie_path.read_text(encoding="utf-8"), stat.S_IMODE(cookie_path.stat().st_mode), ) ) if self.rewrite_cookie: cookie_path.write_text("updated by yt-dlp", encoding="utf-8") if "--flat-playlist" in arguments: return subprocess.CompletedProcess( arguments, 0, stdout="Test Playlist|||First Video|||playlist-id\n", stderr="", ) if self.fail_download: raise subprocess.CalledProcessError(7, arguments) if not self.no_media: output_index = arguments.index("-o") + 1 output_template = Path(arguments[output_index]) directory = output_template.parent directory.mkdir(parents=True, exist_ok=True) (directory / "001 - First [one].mp4").write_bytes(b"video-one") (directory / "002 - Second [two].webm").write_bytes(b"video-two") return subprocess.CompletedProcess(arguments, 0, stdout="", stderr="") class DownloadManagerTests(unittest.TestCase): def setUp(self): self.tempdir = tempfile.TemporaryDirectory() base = Path(self.tempdir.name) self.media = base / "media" self.data = base / "data" self.cookie = base / "cookies.txt" self.cookie.write_text("cookie data", encoding="utf-8") self.config = ControllerConfig( storage=StorageConfig( media_base=self.media, data_base=self.data, retention_seconds=100, ), downloads=DownloadConfig( yt_dlp_binary="yt-dlp-test", cookie_path=self.cookie, cookie_required=True, max_workers=1, process_timeout_seconds=30, max_height=720, ), ) self.store = ControllerStore(self.config.storage).open() def tearDown(self): self.store.close() self.tempdir.cleanup() def manager(self, runner, events=None): event_list = events if events is not None else [] return DownloadManager( self.config, self.store, runner=runner, executor=ImmediateExecutor(), on_event=event_list.append, ) def test_success_uses_configured_cookie_and_atomically_builds_playlist(self): events = [] self.cookie.chmod(0o444) runner = FakeYtDlp(self.media, rewrite_cookie=True) manager = self.manager(runner, events) result = manager.enqueue("https://example.test/playlist", "Alice", "Service") self.assertEqual(result.status, "reserved") self.assertEqual([event.state for event in events], ["queued", "downloading", "ready"]) slot = self.store.get_slot(1) self.assertEqual(slot.media_state, "ready") self.assertEqual(slot.title, "Test Playlist") self.assertEqual(slot.item_count, 2) playlist = canonical_playlist_path(self.media, 1) self.assertEqual( playlist.read_text(encoding="utf-8"), "001 - First [one].mp4\n002 - Second [two].webm\n", ) self.assertEqual(list(playlist.parent.glob(".playlist.*.tmp")), []) for arguments, _ in runner.calls: self.assertIn("--ignore-config", arguments) self.assertIn("--no-js-runtimes", arguments) self.assertEqual(arguments[arguments.index("--js-runtimes") + 1], "node") self.assertEqual( arguments[arguments.index("--cache-dir") + 1], "/tmp/yt-dlp-cache", ) self.assertIn("--cookies", arguments) self.assertNotIn(str(self.cookie), arguments) self.assertEqual(len(runner.cookie_observations), 2) working_paths = [item[0] for item in runner.cookie_observations] self.assertEqual(len(set(working_paths)), 2) self.assertTrue( all(item[1] == "cookie data" for item in runner.cookie_observations) ) self.assertTrue( all(item[2] == 0o600 for item in runner.cookie_observations) ) self.assertTrue(all(not path.exists() for path in working_paths)) self.assertEqual(self.cookie.read_text(encoding="utf-8"), "cookie data") download_arguments = runner.calls[-1][0] self.assertIn("bv*[height<=720]+ba/b[height<=720]", download_arguments) @staticmethod def _select_player_format(max_height, formats): with YoutubeDL({"quiet": True}) as ydl: selector = ydl.build_format_selector( player_format_selector(max_height) ) selected = list(ydl._select_formats(formats, selector)) return selected def test_player_selector_chooses_separate_true_1080_video_and_audio(self): formats = [ { "format_id": "22", "url": "https://example.test/720.mp4", "ext": "mp4", "height": 720, "width": 1280, "vcodec": "avc1.64001f", "acodec": "mp4a.40.2", "tbr": 2500, "filesize": 1, }, { "format_id": "399", "url": "https://example.test/1080.mp4", "ext": "mp4", "height": 1080, "width": 1920, "vcodec": "av01.0.08M.08", "acodec": "none", "tbr": 4500, "filesize": 1, }, { "format_id": "251", "url": "https://example.test/audio.webm", "ext": "webm", "vcodec": "none", "acodec": "opus", "tbr": 160, "filesize": 1, }, ] selected = self._select_player_format(1080, formats) self.assertEqual(len(selected), 1) self.assertEqual(selected[0]["height"], 1080) self.assertEqual( [item["format_id"] for item in selected[0]["requested_formats"]], ["399", "251"], ) self.assertEqual(selected[0]["vcodec"], "av01.0.08M.08") self.assertEqual(selected[0]["acodec"], "opus") def test_player_selector_preserves_lower_source_resolution(self): formats = [ { "format_id": "18", "url": "https://example.test/360.mp4", "ext": "mp4", "height": 360, "width": 640, "vcodec": "avc1.42001e", "acodec": "mp4a.40.2", "tbr": 700, "filesize": 1, }, { "format_id": "22", "url": "https://example.test/720.mp4", "ext": "mp4", "height": 720, "width": 1280, "vcodec": "avc1.64001f", "acodec": "mp4a.40.2", "tbr": 2500, "filesize": 1, }, ] selected = self._select_player_format(1080, formats) self.assertEqual(len(selected), 1) self.assertEqual(selected[0]["format_id"], "22") self.assertEqual(selected[0]["height"], 720) def test_player_selector_never_uses_unfiltered_over_height_fallback(self): formats = [ { "format_id": "401", "url": "https://example.test/2160.mp4", "ext": "mp4", "height": 2160, "width": 3840, "vcodec": "av01.0.12M.08", "acodec": "none", "tbr": 12000, "filesize": 1, }, { "format_id": "251", "url": "https://example.test/audio.webm", "ext": "webm", "vcodec": "none", "acodec": "opus", "tbr": 160, "filesize": 1, }, ] self.assertEqual(self._select_player_format(1080, formats), []) def test_worker_failure_is_contained_and_persisted(self): events = [] runner = FakeYtDlp(self.media, fail_download=True) manager = self.manager(runner, events) result = manager.enqueue("https://example.test/failure", None, None) self.assertEqual(result.status, "reserved") slot = self.store.get_slot(1) self.assertEqual(slot.media_state, "error") self.assertIn("exit code 7", slot.last_error) self.assertEqual(events[-1].state, "error") self.assertTrue( all(not observation[0].exists() for observation in runner.cookie_observations) ) def test_missing_required_cookie_fails_job_without_escaping(self): self.cookie.unlink() events = [] runner = FakeYtDlp(self.media) manager = self.manager(runner, events) manager.enqueue("https://example.test/no-cookie", None, None) self.assertEqual(self.store.get_slot(1).media_state, "error") self.assertIn("required yt-dlp cookie", self.store.get_slot(1).last_error) self.assertEqual(runner.calls, []) def test_delete_moves_only_canonical_slot_and_keeps_other_identity(self): first = self.store.reserve_download("https://example.test/one", None, None, job_id="one") second = self.store.reserve_download("https://example.test/two", None, None, job_id="two") self.store.fail_download("one", 1, first.generation, "failed") self.store.fail_download("two", 2, second.generation, "failed") slot_one = self.media / "slot1" slot_two = self.media / "slot2" slot_one.mkdir(parents=True) slot_two.mkdir(parents=True) (slot_one / "one.mp4").write_bytes(b"one") (slot_two / "two.mp4").write_bytes(b"two") manager = self.manager(FakeYtDlp(self.media)) self.assertTrue(manager.delete_slot(1, reason="manual")) self.assertFalse(slot_one.exists()) self.assertTrue((slot_two / "two.mp4").exists()) self.assertEqual(self.store.get_slot(1).media_state, "empty") self.assertEqual(self.store.get_slot(2).url, "https://example.test/two") def test_delete_refuses_live_download_and_playback_guard(self): reservation = self.store.reserve_download( "https://example.test/live", None, None, job_id="live" ) manager = self.manager(FakeYtDlp(self.media)) with self.assertRaises(SlotProtectedError): manager.delete_slot(1, reason="manual") self.store.fail_download("live", 1, reservation.generation, "cancelled") self.store.protect_slot(1, "player_active", guard_id="active") with self.assertRaises(SlotProtectedError): manager.delete_slot(1, reason="manual") def test_symlink_slot_is_unlinked_without_touching_target(self): reservation = self.store.reserve_download( "https://example.test/link", None, None, job_id="link" ) self.store.fail_download("link", 1, reservation.generation, "failed") target = Path(self.tempdir.name) / "outside" target.mkdir() protected = target / "keep.txt" protected.write_text("keep", encoding="utf-8") self.media.mkdir(parents=True) os.symlink(target, self.media / "slot1") manager = self.manager(FakeYtDlp(self.media)) self.assertTrue(manager.delete_slot(1, reason="manual")) self.assertEqual(protected.read_text(encoding="utf-8"), "keep") def test_interrupted_deletion_is_recoverable(self): reservation = self.store.reserve_download( "https://example.test/recover", None, None, job_id="recover" ) self.store.fail_download("recover", 1, reservation.generation, "failed") token = self.store.begin_delete(1, reason="manual") slot = self.media / "slot1" slot.mkdir(parents=True) (slot / "partial.mp4").write_bytes(b"partial") self.assertEqual(self.store.pending_deletions(), [ type(token)(slot=1, generation=token.generation, reason="recovered_interrupted_deletion") ]) manager = self.manager(FakeYtDlp(self.media)) self.assertEqual(manager.recover_pending_deletions(), [1]) self.assertFalse(slot.exists()) self.assertEqual(self.store.get_slot(1).media_state, "empty") def test_interrupted_deletion_recovery_defers_a_new_playback_guard(self): reservation = self.store.reserve_download( "https://example.test/protected-recovery", None, None, job_id="protected-recovery", ) self.store.fail_download( "protected-recovery", 1, reservation.generation, "failed", ) self.store.begin_delete(1, reason="manual") slot = self.media / "slot1" slot.mkdir(parents=True) (slot / "still-playing.mp4").write_bytes(b"playing") self.store.protect_slot(1, "player_active", guard_id="fresh-player") manager = self.manager(FakeYtDlp(self.media)) self.assertEqual(manager.recover_pending_deletions(), []) self.assertTrue(slot.exists()) self.assertEqual(self.store.get_slot(1).media_state, "deleting") self.store.clear_playback_guards(reason="fresh terminal idle evidence") self.assertEqual(manager.recover_pending_deletions(), [1]) self.assertFalse(slot.exists()) def test_async_deletion_failure_remains_durable_and_is_retryable(self): reservation = self.store.reserve_download( "https://example.test/retry-delete", None, None, job_id="retry-delete", ) self.store.fail_download( "retry-delete", 1, reservation.generation, "failed", ) slot = self.media / "slot1" slot.mkdir(parents=True) (slot / "retry.mp4").write_bytes(b"retry") events = [] manager = self.manager(FakeYtDlp(self.media), events) original_recover = manager._recover_deletion attempts = 0 def fail_once(token): nonlocal attempts attempts += 1 if attempts == 1: raise OSError("transient media failure") return original_recover(token) manager._recover_deletion = fail_once token = manager.schedule_delete_slot(1, reason="manual") self.assertEqual(token.slot, 1) self.assertEqual(self.store.get_slot(1).media_state, "deleting") self.assertEqual(events[-1].state, "delete_error") self.assertEqual(len(self.store.pending_deletions()), 1) self.assertTrue(manager.schedule_pending_deletion_recovery()) self.assertEqual(attempts, 2) self.assertEqual(events[-1].state, "deleted") self.assertEqual(self.store.get_slot(1).media_state, "empty") self.assertFalse(slot.exists()) def test_async_pending_recovery_retries_after_guard_is_cleared(self): reservation = self.store.reserve_download( "https://example.test/deferred-async-delete", None, None, job_id="deferred-async-delete", ) self.store.fail_download( "deferred-async-delete", 1, reservation.generation, "failed", ) self.store.begin_delete(1, reason="manual") slot = self.media / "slot1" slot.mkdir(parents=True) (slot / "active.mp4").write_bytes(b"active") self.store.protect_slot(1, "player_active", guard_id="active-player") events = [] manager = self.manager(FakeYtDlp(self.media), events) self.assertTrue(manager.schedule_pending_deletion_recovery()) self.assertEqual(events[-1].state, "delete_deferred") self.assertEqual(self.store.get_slot(1).media_state, "deleting") self.assertTrue(slot.exists()) self.store.clear_playback_guards(reason="fresh terminal idle evidence") self.assertTrue(manager.schedule_pending_deletion_recovery()) self.assertEqual(events[-1].state, "deleted") self.assertEqual(self.store.get_slot(1).media_state, "empty") self.assertFalse(slot.exists()) def test_guarded_recovery_restores_crash_staged_trash_to_canonical_slot(self): reservation = self.store.reserve_download( "https://example.test/restore-trash", None, None, job_id="restore-trash", ) self.store.fail_download( "restore-trash", 1, reservation.generation, "failed", ) token = self.store.begin_delete(1, reason="expired") trash = ( self.media / ".controller-trash" / f"slot1.generation{token.generation}.crash" ) trash.mkdir(parents=True) (trash / "still-active.mp4").write_bytes(b"active") self.store.protect_slot(1, "player_active", guard_id="fresh-active") events = [] manager = self.manager(FakeYtDlp(self.media), events) self.assertTrue(manager.schedule_pending_deletion_recovery()) canonical = self.media / "slot1" self.assertEqual(events[-1].state, "delete_deferred") self.assertTrue((canonical / "still-active.mp4").exists()) self.assertFalse(trash.exists()) self.assertEqual(self.store.get_slot(1).media_state, "deleting") def test_queued_deletion_rechecks_guard_before_filesystem_mutation(self): reservation = self.store.reserve_download( "https://example.test/queued-delete-guard", None, None, job_id="queued-delete-guard", ) self.store.fail_download( "queued-delete-guard", 1, reservation.generation, "failed", ) self.store.begin_delete(1, reason="manual") slot = self.media / "slot1" slot.mkdir(parents=True) (slot / "must-survive.mp4").write_bytes(b"active") events = [] executor = QueuedExecutor() manager = DownloadManager( self.config, self.store, runner=FakeYtDlp(self.media), executor=executor, on_event=events.append, ) self.assertTrue(manager.schedule_pending_deletion_recovery()) self.store.protect_slot(1, "player_active", guard_id="late-player") executor.run_next() self.assertEqual(events[-1].state, "delete_deferred") self.assertEqual(self.store.get_slot(1).media_state, "deleting") self.assertTrue((slot / "must-survive.mp4").exists()) self.store.clear_playback_guards(reason="fresh terminal idle evidence") self.assertTrue(manager.schedule_pending_deletion_recovery()) executor.run_next() self.assertEqual(events[-1].state, "deleted") self.assertEqual(self.store.get_slot(1).media_state, "empty") self.assertFalse(slot.exists()) def test_shutdown_terminates_tracked_process_without_waiting_for_job_timeout(self): manager = DownloadManager(self.config, self.store) errors = [] def run_long_process(): try: manager._run_command( ["/bin/sh", "-c", "sleep 60"], check=True, text=True, timeout=60, capture_output=True, ) except BaseException as exc: errors.append(exc) worker = threading.Thread(target=run_long_process) worker.start() deadline = time.monotonic() + 2 while time.monotonic() < deadline: with manager._lock: if manager._processes: break time.sleep(0.01) else: self.fail("test process did not start") started = time.monotonic() manager.shutdown(wait=False) worker.join(timeout=3) self.assertFalse(worker.is_alive()) self.assertLess(time.monotonic() - started, 3) self.assertTrue(errors) with manager._lock: self.assertEqual(manager._processes, set()) if __name__ == "__main__": unittest.main()