"""Tests for ``dl.write_export`` — the ``--export`` JSON sidecar. Regression: DRM-free tracks never pass through ``prepare_drm``, so ``write_export`` must accept ``drm=None`` (and DRM systems without ``to_dict``/``content_keys`` such as ClearKey) and still record the track/manifest/chapter/attachment info that ``unshackle import`` rebuilds a download from. """ from __future__ import annotations import json from pathlib import Path from types import SimpleNamespace from uuid import UUID from unshackle.commands.dl import dl from unshackle.core.drm.clearkey import ClearKey from unshackle.core.import_service import ImportService from unshackle.core.titles import Movie from unshackle.core.tracks import Audio, Chapter, Subtitle, Video KID = UUID(hex="00000000000000000000000000000001") class StubService: """Stands in for the service class slot on Movie; never instantiated.""" class StubDRM: """Minimal licensed-DRM shape: ``to_dict`` plus filled ``content_keys``.""" def __init__(self) -> None: self.content_keys = {KID: "aa" * 16} def to_dict(self) -> dict: return {"system": "Widevine", "pssh_b64": "AAAA"} def make_dl() -> dl: # __new__ skips the CLI-driven __init__; write_export only needs `service`. instance = dl.__new__(dl) instance.service = "EXAMPLE" return instance def make_title() -> Movie: title = Movie(id_="movie-1", service=StubService, name="Example Movie", year=2024, language="en") title.tracks.add( Video( id_="v1", url="https://example.test/v1.mp4", language="en", codec=Video.Codec.AVC, range_=Video.Range.SDR, width=1920, height=1080, bitrate=5_000_000, ) ) title.tracks.add( Audio( id_="a1", url="https://example.test/a1.mp4", language="en", codec=Audio.Codec.AAC, bitrate=128_000, ) ) title.tracks.add( Subtitle( id_="s1", url="https://example.test/s1.vtt", language="en", codec=Subtitle.Codec.WebVTT, ) ) return title def read_export(path: Path) -> dict: return json.loads(path.read_text(encoding="utf8")) def test_drm_free_track_exports(tmp_path: Path) -> None: """The reported bug: DRM-free downloads produced no usable export.""" export = tmp_path / "export.json" title = make_title() video = title.tracks.videos[0] make_dl().write_export(export, title, video) doc = read_export(export) assert doc["version"] == 2 assert doc["service"] == "EXAMPLE" tinfo = doc["titles"]["movie-1"] assert tinfo["meta"]["name"] == "Example Movie" assert set(tinfo["tracks"]) == {"v1", "a1", "s1"} assert "keys" not in tinfo["tracks"]["v1"] assert "drm" not in tinfo["tracks"]["v1"] def test_clearkey_drm_exports_track_without_keys(tmp_path: Path) -> None: """ClearKey has no to_dict/content_keys; the track info must still export.""" export = tmp_path / "export.json" title = make_title() video = title.tracks.videos[0] make_dl().write_export(export, title, video, ClearKey(key="bb" * 16)) doc = read_export(export) track = doc["titles"]["movie-1"]["tracks"]["v1"] assert "keys" not in track assert "drm" not in track def test_post_download_write_keeps_licensed_keys(tmp_path: Path) -> None: """The drm=None write after download must not clobber prepare_drm's DRM/keys.""" export = tmp_path / "export.json" title = make_title() video = title.tracks.videos[0] runner = make_dl() runner.write_export(export, title, video, StubDRM()) # prepare_drm runner.write_export(export, title, video) # post-download hook track = read_export(export)["titles"]["movie-1"]["tracks"]["v1"] assert track["drm"] == [{"system": "Widevine", "pssh_b64": "AAAA"}] assert track["keys"] == {KID.hex: "aa" * 16} def test_drm_free_export_roundtrips_through_import_service(tmp_path: Path) -> None: """A DRM-free export must rebuild via ImportService: titles, tracks (no DRM), empty key pool (resolve_server_keys no-op) and chapters.""" export = tmp_path / "export.json" title = make_title() title.tracks.chapters.add(Chapter("00:00:10.000", "Intro")) runner = make_dl() for track in [*title.tracks.videos, *title.tracks.audio, *title.tracks.subtitles]: runner.write_export(export, title, track) # ImportService only touches ctx.parent.params (proxy flags) when building its session. ctx = SimpleNamespace(parent=None, params={}) svc = ImportService(ctx, "EXAMPLE", "movie-1", str(export)) titles = list(svc.get_titles()) assert len(titles) == 1 movie = titles[0] assert isinstance(movie, Movie) assert movie.name == "Example Movie" assert movie.year == 2024 tracks = svc.get_tracks(movie) assert {t.id for t in tracks} == {"v1", "a1", "s1"} assert all(not t.drm for t in tracks) assert svc.key_pool() == {} movie.tracks = tracks svc.resolve_server_keys(movie) assert all(not t.drm for t in movie.tracks) # Chapters.add auto-inserts a nameless 00:00:00 baseline chapter; it round-trips too. assert [c.name for c in svc.get_chapters(movie)] == [None, "Intro"] def test_keyless_content_keys_writes_no_keys_entry(tmp_path: Path) -> None: """A DRM object with empty content_keys must not create an empty keys map.""" export = tmp_path / "export.json" title = make_title() video = title.tracks.videos[0] drm = StubDRM() drm.content_keys = {} make_dl().write_export(export, title, video, drm) track = read_export(export)["titles"]["movie-1"]["tracks"]["v1"] assert track["drm"] == [{"system": "Widevine", "pssh_b64": "AAAA"}] assert "keys" not in track