feat: reserve downloads and protect slot cleanup
This commit is contained in:
Media
@@ -0,0 +1,199 @@
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
import unittest
|
||||
from concurrent.futures import Future
|
||||
from pathlib import Path
|
||||
|
||||
from controller.config import ControllerConfig, DownloadConfig, StorageConfig
|
||||
from controller.downloads import DownloadManager, canonical_playlist_path
|
||||
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 FakeYtDlp:
|
||||
def __init__(self, media_base: Path, *, fail_download=False, no_media=False):
|
||||
self.media_base = media_base
|
||||
self.fail_download = fail_download
|
||||
self.no_media = no_media
|
||||
self.calls = []
|
||||
|
||||
def __call__(self, arguments, **kwargs):
|
||||
self.calls.append((list(arguments), dict(kwargs)))
|
||||
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 = []
|
||||
runner = FakeYtDlp(self.media)
|
||||
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("--cookies", arguments)
|
||||
self.assertIn(str(self.cookie), arguments)
|
||||
download_arguments = runner.calls[-1][0]
|
||||
self.assertIn("bv*[height<=720]+ba/b[height<=720]/b", download_arguments)
|
||||
|
||||
def test_worker_failure_is_contained_and_persisted(self):
|
||||
events = []
|
||||
manager = self.manager(FakeYtDlp(self.media, fail_download=True), 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")
|
||||
|
||||
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")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user