path: root/packages/meshbay-node/tests/test_stream_audio_track_selection.py

"""
The viewer picks which audio track is streamed.

A dubbed film carries several audio tracks and the streaming path used to map
`0:a:0` unconditionally, so it played in whichever language happened to be
muxed first. Across a real library that is overwhelmingly one language, and the
others could not be reached at all.

The tracks here are told apart by **amplitude**, not by their tags: each is the
same tone at a different volume, so an assertion about which track was served
is measured from the decoded audio of the reassembled stream and cannot be
satisfied by mapping the wrong one. Tags would only prove that the node copied
a string it was given.

Like the other streaming tests, these spawn real ffmpeg/ffprobe against small
synthetic files rather than asserting against the source text.
"""

import re
import shutil
import subprocess
from pathlib import Path

import pytest
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from meshbay_common.crypto import generate_gek
from meshbay_common.webcrypto import chunk_key_aes, decrypt_chunk_aes
from meshbay_node.indexer.group_index import GroupIndex
from meshbay_node.media_probe import probe_video as _probe_video
from meshbay_node.transport.webrtc_server import WebRTCPeerSession

from conftest import needs_subprocess, one_root

_HAVE_FFMPEG = shutil.which("ffmpeg") and shutil.which("ffprobe")
pytestmark = [
    pytest.mark.asyncio,
    pytest.mark.skipif(not _HAVE_FFMPEG, reason="ffmpeg/ffprobe not installed"),
    needs_subprocess,
]

# Each track's tone is attenuated by a different amount, far enough apart that
# an AAC round trip cannot blur one into another. Index in this list is the
# audio ordinal the node is asked for.
_TRACK_GAIN = [1.0, 0.1, 0.01]          # 0 dB, -20 dB, -40 dB
_TRACK_LANG = ["fre", "eng", "spa"]
_TRACK_TITLE = ["Surround", "Original", None]


def _make_multitrack_clip(path: Path) -> None:
    """~1s of H264 video plus three audio tracks at descending volumes.

    The video is muxed first, so the audio streams sit at container indices 1,
    2 and 3 while their audio *ordinals* are 0, 1 and 2 — the gap that
    `-map 0:a:<n>` is indexed by and that a probe reading `s["index"]` would
    get wrong.
    """
    args = ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
            "-f", "lavfi", "-i", "testsrc=size=320x240:rate=25:duration=1"]
    for gain in _TRACK_GAIN:
        args += ["-f", "lavfi",
                 "-i", f"sine=frequency=440:duration=1:sample_rate=48000,"
                       f"volume={gain}"]
    args += ["-map", "0:v:0"]
    for i in range(len(_TRACK_GAIN)):
        args += ["-map", f"{i + 1}:a:0"]
    args += ["-c:v", "libx264", "-preset", "ultrafast", "-c:a", "aac"]
    for i, lang in enumerate(_TRACK_LANG):
        args += [f"-metadata:s:a:{i}", f"language={lang}"]
        if _TRACK_TITLE[i]:
            args += [f"-metadata:s:a:{i}", f"title={_TRACK_TITLE[i]}"]
    args.append(str(path))
    subprocess.run(args, check=True, capture_output=True)


def _session(tmp_path: Path, video_path: Path, gek: bytes):
    import blake3
    file_bytes = video_path.read_bytes()
    file_id = blake3.blake3(file_bytes).hexdigest()

    sk_node = Ed25519PrivateKey.generate()
    index = GroupIndex(group_id="g" * 32, sk_node=sk_node, gek=gek)
    from meshbay_common.protocol import IndexEntry
    index.add_entry(IndexEntry(
        id=file_id, name=video_path.name, path=video_path.parent.name,
        size=len(file_bytes), type="video", added_at=0))

    session = WebRTCPeerSession.__new__(WebRTCPeerSession)
    session._ctx = {
        "roots": one_root(video_path.parent),
        "index": index,
        "gek": gek,
        "sk_node": sk_node,
        "max_concurrent_streams": 4,
    }
    session._group_id = None
    session._user_id = "tester"
    session._stream_stopped = False
    session._stream_keepalives = 0
    session.sent = []
    session._send = session.sent.append
    session._audit = lambda *a, **k: None
    return session, file_id


def _reassemble(sent: list[dict], gek: bytes, file_id: str) -> bytes:
    file_hash = bytes.fromhex(file_id)
    segments = sorted(
        (m for m in sent if m.get("type") == "stream_data"),
        key=lambda m: m["segment_index"])
    out = b""
    for m in segments:
        key = chunk_key_aes(gek, file_hash, m["segment_index"])
        out += decrypt_chunk_aes(key, m["nonce"], m["ct"])
    return out


def _mean_volume_db(path: Path) -> float:
    """What ffmpeg's volumedetect measures in the decoded audio."""
    proc = subprocess.run(
        ["ffmpeg", "-hide_banner", "-i", str(path), "-af", "volumedetect",
         "-f", "null", "-"],
        capture_output=True, text=True)
    m = re.search(r"mean_volume:\s*(-?\d+(?:\.\d+)?) dB", proc.stderr)
    assert m, f"volumedetect said nothing usable: {proc.stderr[-400:]}"
    return float(m.group(1))


async def _stream(tmp_path: Path, clip: Path, msg_extra: dict):
    gek = generate_gek()
    session, file_id = _session(tmp_path, clip, gek)
    await session._stream_video_inner(
        {"file_id": file_id, "start": 0, "credits": 0, **msg_extra})
    errors = [m for m in session.sent if m.get("type") == "error"]
    assert not errors, f"streaming must not fail: {errors}"
    init = next(m for m in session.sent if m.get("type") == "stream_init")
    return session, file_id, gek, init


async def _streamed_volume(tmp_path: Path, clip: Path, msg_extra: dict, tag: str):
    session, file_id, gek, init = await _stream(tmp_path, clip, msg_extra)
    out = tmp_path / f"out-{tag}.mp4"
    out.write_bytes(_reassemble(session.sent, gek, file_id))
    return _mean_volume_db(out), init


async def test_probe_enumerates_every_track_by_ordinal_not_container_index(tmp_path):
    """The trap this feature is one wrong line away from.

    `-map 0:a:1` counts audio streams; `s["index"]` counts every stream in the
    container. With a video stream muxed first the two never agree, and a probe
    reporting container indices would make the client ask for track 1 and be
    served track 2 — silently, since both are real audio.
    """
    clip = tmp_path / "clip.mkv"
    _make_multitrack_clip(clip)

    probe = await _probe_video(str(clip))

    assert [tr.ordinal for tr in probe.audio_tracks] == [0, 1, 2]
    assert [tr.language for tr in probe.audio_tracks] == _TRACK_LANG
    assert probe.has_audio is True
    # The container really does disagree, or the assertion above proves nothing.
    raw = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "a",
         "-show_entries", "stream=index", "-of", "csv=p=0", str(clip)],
        check=True, capture_output=True, text=True)
    assert [int(x) for x in raw.stdout.split()] == [1, 2, 3], (
        "the fixture must put audio at container indices that differ from the "
        "ordinals, or it cannot tell the two apart")


async def test_probe_reports_the_title_tag_when_the_muxer_wrote_one(tmp_path):
    """Two tracks in one language are one menu entry repeated without it."""
    clip = tmp_path / "clip.mkv"
    _make_multitrack_clip(clip)

    probe = await _probe_video(str(clip))

    assert [tr.title for tr in probe.audio_tracks] == _TRACK_TITLE
    assert all(tr.channels == 1 for tr in probe.audio_tracks)


async def test_the_requested_audio_track_is_the_one_streamed(tmp_path):
    """Measured from the decoded audio, not from a tag the node echoed back."""
    clip = tmp_path / "clip.mkv"
    _make_multitrack_clip(clip)

    first, init_first = await _streamed_volume(tmp_path, clip, {"audio_track": 0}, "0")
    third, init_third = await _streamed_volume(tmp_path, clip, {"audio_track": 2}, "2")

    assert init_first["audio_track"] == 0
    assert init_third["audio_track"] == 2
    # -40 dB of attenuation between them; anything under 15 dB of measured
    # separation means the same track was served twice.
    assert first - third > 15, (
        f"track 0 ({first} dB) and track 2 ({third} dB) decoded to the same "
        "loudness, so the requested track was not the one mapped")


async def test_no_audio_track_asked_for_still_means_the_first(tmp_path):
    """A client that says nothing gets exactly what it got before."""
    clip = tmp_path / "clip.mkv"
    _make_multitrack_clip(clip)

    silent, init_silent = await _streamed_volume(tmp_path, clip, {}, "default")
    explicit, _ = await _streamed_volume(tmp_path, clip, {"audio_track": 0}, "explicit")

    assert init_silent["audio_track"] == 0
    assert abs(silent - explicit) < 2


async def test_an_out_of_range_track_falls_back_to_the_first_and_says_so(tmp_path):
    """The client's list can predate the file being replaced on disk.

    A viewer who asked for the second track of a file that now has one wants
    the film, not an error — and `stream_init` has to report the track actually
    used, the same way it reports the `start` actually used.
    """
    clip = tmp_path / "clip.mkv"
    _make_multitrack_clip(clip)

    got, init = await _streamed_volume(tmp_path, clip, {"audio_track": 9}, "oob")
    expected, _ = await _streamed_volume(tmp_path, clip, {"audio_track": 0}, "base")

    assert init["audio_track"] == 0, "stream_init must not echo the impossible request"
    assert abs(got - expected) < 2


async def test_a_malformed_track_number_is_not_an_error(tmp_path):
    clip = tmp_path / "clip.mkv"
    _make_multitrack_clip(clip)

    for bad in ("two", None, -1, 1.5):
        _, _, _, init = await _stream(tmp_path, clip, {"audio_track": bad})
        assert init["audio_track"] in (0, 1), f"{bad!r} produced {init['audio_track']!r}"


async def test_stream_init_lists_the_tracks_for_the_client_to_choose_from(tmp_path):
    """The list is the whole capability negotiation.

    The player draws its selector from this and from nothing else — there is no
    version check in it — so a node that sends no list gets no selector and is
    never sent an `audio_track` it would ignore and answer in the wrong
    language.
    """
    clip = tmp_path / "clip.mkv"
    _make_multitrack_clip(clip)

    _, _, _, init = await _stream(tmp_path, clip, {})

    assert [tr["i"] for tr in init["audio_tracks"]] == [0, 1, 2]
    assert [tr["lang"] for tr in init["audio_tracks"]] == _TRACK_LANG
    assert init["audio_tracks"][0]["title"] == "Surround"
    assert init["audio_tracks"][2]["title"] is None


async def test_a_single_track_file_reports_a_list_of_one(tmp_path):
    """Not an empty list: the player decides on length, and a one-entry list is
    how it knows there is nothing to choose rather than nothing to report."""
    clip = tmp_path / "mono.mkv"
    subprocess.run(
        ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
         "-f", "lavfi", "-i", "testsrc=size=320x240:rate=25:duration=1",
         "-f", "lavfi", "-i", "sine=frequency=440:duration=1:sample_rate=48000",
         "-c:v", "libx264", "-preset", "ultrafast", "-c:a", "aac", str(clip)],
        check=True, capture_output=True)

    _, _, _, init = await _stream(tmp_path, clip, {})

    assert len(init["audio_tracks"]) == 1
    assert init["audio_track"] == 0


async def test_a_file_with_no_audio_reports_no_track_at_all(tmp_path):
    clip = tmp_path / "silent.mkv"
    subprocess.run(
        ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
         "-f", "lavfi", "-i", "testsrc=size=320x240:rate=25:duration=1",
         "-c:v", "libx264", "-preset", "ultrafast", "-an", str(clip)],
        check=True, capture_output=True)

    _, _, _, init = await _stream(tmp_path, clip, {"audio_track": 1})

    assert init["audio_tracks"] == []
    assert init["audio_track"] is None