m3u8dl/N_m3u8DL_py/tests/test_api.py

from __future__ import annotations

import shutil
from pathlib import Path

import pytest

from n_m3u8dl_py import (
    DownloadCancelledError,
    DownloadEventKind,
    DownloadRequest,
    N_m3u8DL,
    SelectionError,
    TrackSelection,
)


def test_public_api_inspects_dash_with_stable_track_ids() -> None:
    root = Path(__file__).parents[2]
    manifest = root / "src/N_m3u8DL-RE.Tests/Resources/Dash/Manifest_DuplicateSegments.mpd"
    info = N_m3u8DL().inspect(manifest)

    assert info.protocol == "DASH"
    assert info.tracks[0].track_id == "dash:0"
    assert info.tracks[0].segments_count == 3
    assert info.tracks[0].duration == 6


def test_api_downloads_hls_with_default_auto_selection_and_events(tmp_path: Path) -> None:
    (tmp_path / "one.ts").write_bytes(b"one")
    (tmp_path / "two.ts").write_bytes(b"two")
    manifest = tmp_path / "movie.m3u8"
    manifest.write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXTINF:1,\ntwo.ts\n#EXT-X-ENDLIST\n")
    events = []
    client = N_m3u8DL()

    info = client.inspect(manifest, on_event=events.append)
    result = client.download(
        info,
        DownloadRequest(
            output_dir=tmp_path / "output",
            temporary_dir=tmp_path / "temp",
            file_name="movie",
            binary_merge=True,
        ),
        on_event=events.append,
    )

    assert result.selected_track_ids == ("hls:0",)
    assert result.files[0].path.read_bytes() == b"onetwo"
    assert result.files[0].size == 6
    assert [event.kind for event in events] == [
        DownloadEventKind.INSPECT_STARTED,
        DownloadEventKind.INSPECT_COMPLETED,
        DownloadEventKind.DOWNLOAD_STARTED,
        DownloadEventKind.PLAYLIST_LOADED,
        DownloadEventKind.TRACK_STARTED,
        DownloadEventKind.SEGMENT_COMPLETED,
        DownloadEventKind.SEGMENT_COMPLETED,
        DownloadEventKind.TRACK_COMPLETED,
        DownloadEventKind.DOWNLOAD_COMPLETED,
    ]
    segment_events = [event for event in events if event.kind is DownloadEventKind.SEGMENT_COMPLETED]
    assert [(event.completed_segments, event.total_segments) for event in segment_events] == [(1, 2), (2, 2)]


def test_api_download_url_is_one_call_convenience(tmp_path: Path) -> None:
    (tmp_path / "one.ts").write_bytes(b"one")
    manifest = tmp_path / "movie.m3u8"
    manifest.write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXT-X-ENDLIST\n")

    result = N_m3u8DL().download_url(
        manifest,
        DownloadRequest(output_dir=tmp_path / "output", temporary_dir=tmp_path / "temp", file_name="one-call", binary_merge=True),
    )

    assert result.files[0].path.read_bytes() == b"one"


def test_api_downloads_segments_without_merging_or_ffmpeg(tmp_path: Path) -> None:
    (tmp_path / "one.ts").write_bytes(b"one")
    (tmp_path / "two.ts").write_bytes(b"two")
    manifest = tmp_path / "movie.m3u8"
    manifest.write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXTINF:1,\ntwo.ts\n#EXT-X-ENDLIST\n")

    result = N_m3u8DL().download_segments_url(
        manifest,
        DownloadRequest(
            output_dir=tmp_path / "output",
            temporary_dir=tmp_path / "temporary",
            file_name="movie",
            ffmpeg_path="/not-used-by-download-segments",
        ),
    )

    track = result.tracks[0]
    assert [path.read_bytes() for path in track.segments] == [b"one", b"two"]
    assert track.directory.is_dir()
    assert not list((tmp_path / "output").glob("movie.*"))
    shutil.rmtree(result.temporary_dir.parent)
    assert not result.temporary_dir.exists()


def test_api_downloads_only_explicit_track_id(tmp_path: Path) -> None:
    (tmp_path / "video.ts").write_bytes(b"video")
    (tmp_path / "audio.ts").write_bytes(b"audio")
    (tmp_path / "video.m3u8").write_text("#EXTM3U\n#EXTINF:1,\nvideo.ts\n#EXT-X-ENDLIST\n")
    (tmp_path / "audio.m3u8").write_text("#EXTM3U\n#EXTINF:1,\naudio.ts\n#EXT-X-ENDLIST\n")
    manifest = tmp_path / "master.m3u8"
    manifest.write_text(
        "#EXTM3U\n"
        '#EXT-X-MEDIA:TYPE=AUDIO,GROUP-ID="aud",NAME="English",LANGUAGE="en",URI="audio.m3u8"\n'
        '#EXT-X-STREAM-INF:BANDWIDTH=1000,RESOLUTION=640x360,AUDIO="aud"\n'
        "video.m3u8\n"
    )
    client = N_m3u8DL()
    info = client.inspect(manifest)
    video = next(track for track in info.tracks if track.resolution == "640x360")

    result = client.download(
        info,
        DownloadRequest(
            output_dir=tmp_path / "output",
            temporary_dir=tmp_path / "temp",
            file_name="video-only",
            binary_merge=True,
            selection=TrackSelection(track_ids=(video.track_id,)),
        ),
    )

    assert result.selected_track_ids == (video.track_id,)
    assert len(result.files) == 1
    assert result.files[0].path.read_bytes() == b"video"


def test_api_rejects_unknown_track_id(tmp_path: Path) -> None:
    (tmp_path / "one.ts").write_bytes(b"one")
    manifest = tmp_path / "movie.m3u8"
    manifest.write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXT-X-ENDLIST\n")
    client = N_m3u8DL()
    info = client.inspect(manifest)

    with pytest.raises(SelectionError, match="Unknown track ID"):
        client.download(info, DownloadRequest(selection=TrackSelection(track_ids=("hls:404",))))


def test_api_event_callback_can_cancel_download(tmp_path: Path) -> None:
    (tmp_path / "one.ts").write_bytes(b"one")
    manifest = tmp_path / "movie.m3u8"
    manifest.write_text("#EXTM3U\n#EXTINF:1,\none.ts\n#EXT-X-ENDLIST\n")
    events = []

    def on_event(event):
        events.append(event)
        return False if event.kind is DownloadEventKind.TRACK_STARTED else None

    client = N_m3u8DL()
    info = client.inspect(manifest, on_event=on_event)
    with pytest.raises(DownloadCancelledError):
        client.download(
            info,
            DownloadRequest(output_dir=tmp_path / "output", temporary_dir=tmp_path / "temp", file_name="cancelled"),
            on_event=on_event,
        )

    assert events[-1].kind is DownloadEventKind.CANCELLED
    assert not (tmp_path / "output" / "cancelled.ts").exists()