mcbebop/tests/test_rtp.py

"""The RTP packetiser and the capture format, against bytes.

Synthetic NALs throughout: a NAL is a header byte and a payload as far as
RFC 6184 is concerned, so real H.264 would only make the assertions harder to
read. The one test that needs real video is the ffmpeg decode in
`test_sim_video.py`, which is the only one that can tell whether any of this
is actually correct.
"""

import random
import struct
from itertools import pairwise

import pytest

from mcbebop.media import capture, rtp

# header bytes: NRI 3 is "most important", which parameter sets carry.
SPS = bytes([0x67]) + b"\x42\xc0\x1e"
PPS = bytes([0x68]) + b"\xce\x3c\x80"
SEI = bytes([0x06]) + b"\x05\x02\x00\x00\x80"
# A slice whose first bit is set, so first_mb_in_slice is 0 and it opens a
# picture. 0x65 is an IDR at NRI 3, 0x41 a non-IDR at NRI 2.
IDR = bytes([0x65, 0x88]) + b"\xaa" * 40
SLICE = bytes([0x41, 0x9A]) + b"\xbb" * 40
# first bit clear, so first_mb_in_slice is not 0: a continuation slice.
SLICE_2 = bytes([0x41, 0x1A]) + b"\xcc" * 40


def annex_b(*nals: bytes, four_byte: bool = False) -> bytes:
    code = b"\x00\x00\x00\x01" if four_byte else b"\x00\x00\x01"
    return b"".join(code + nal for nal in nals)


# -- Annex-B parsing -----------------------------------------------------
def test_three_and_four_byte_start_codes_both_split():
    assert list(rtp.iter_nals(annex_b(SPS, PPS))) == [SPS, PPS]
    assert list(rtp.iter_nals(annex_b(SPS, PPS, four_byte=True))) == [SPS, PPS]


def test_trailing_zeros_belong_to_the_start_code_not_the_nal():
    # A fourth zero before 00 00 01 is a trailing byte of the previous NAL.
    # Counting it would change the length and so change where the MTU falls.
    data = b"\x00\x00\x01" + SPS + b"\x00\x00\x00\x00\x01" + PPS
    assert list(rtp.iter_nals(data)) == [SPS, PPS]


def test_a_stream_with_no_start_code_yields_nothing():
    assert list(rtp.iter_nals(b"\xde\xad\xbe\xef")) == []
    with pytest.raises(ValueError, match="Annex-B"):
        rtp.AnnexBStream.from_bytes(b"\xde\xad\xbe\xef")


def test_leading_garbage_before_the_first_start_code_is_skipped():
    assert list(rtp.iter_nals(b"junk" + annex_b(SPS))) == [SPS]


def test_nal_type_reads_the_low_five_bits():
    assert rtp.nal_type(SPS) == 7
    assert rtp.nal_type(PPS) == 8
    assert rtp.nal_type(IDR) == 5
    assert rtp.nal_type(SLICE) == 1


# -- access units --------------------------------------------------------
def test_parameter_sets_attach_to_the_picture_that_follows_them():
    units = rtp.access_units([SPS, PPS, IDR, SLICE])
    assert [u.nals for u in units] == [(SPS, PPS, IDR), (SLICE,)]
    assert units[0].is_idr and units[0].carries_parameter_sets
    assert not units[1].is_idr and not units[1].carries_parameter_sets


def test_a_continuation_slice_stays_in_the_same_access_unit():
    # Splitting a multi-slice picture per slice would put the marker bit and
    # the timestamp step in the wrong places.
    units = rtp.access_units([IDR, SLICE_2, SLICE])
    assert [u.nals for u in units] == [(IDR, SLICE_2), (SLICE,)]


def test_the_stream_keeps_the_parameter_sets_for_reuse():
    stream = rtp.AnnexBStream.from_bytes(annex_b(SPS, PPS, IDR, SLICE, SLICE))
    assert stream.parameter_sets == (SPS, PPS)
    assert len(stream.units) == 3


def test_a_stream_without_parameter_sets_warns(caplog):
    rtp.AnnexBStream.from_bytes(annex_b(IDR))
    assert "joining late" in caplog.text


# -- RTP headers ---------------------------------------------------------
def unpack(packet):
    return rtp.parse_packet(packet)


def test_a_small_nal_becomes_exactly_one_single_nal_packet():
    p = rtp.Packetiser(mtu=1400, ssrc=0x11223344, seq=7, timestamp=9000)
    packets = p.packetise([IDR])
    assert len(packets) == 1
    pt, seq, ts, ssrc, marker, payload = unpack(packets[0])
    assert (pt, seq, ts, ssrc, marker) == (96, 7, 9000, 0x11223344, True)
    assert payload == IDR  # the NAL travels whole, header byte included


def test_the_header_is_twelve_bytes_of_version_two():
    packets = rtp.Packetiser().packetise([IDR])
    assert packets[0][0] == 0x80  # V=2, no padding, no extension, no CSRC
    assert len(packets[0]) == 12 + len(IDR)


def test_parse_packet_rejects_anything_that_is_not_version_two():
    with pytest.raises(ValueError, match="version 2"):
        rtp.parse_packet(b"\x40" + b"\x00" * 15)
    with pytest.raises(ValueError, match="at least"):
        rtp.parse_packet(b"\x80\x60")


# -- FU-A ----------------------------------------------------------------
def big_nal(size: int, header: int = 0x65) -> bytes:
    return bytes([header]) + bytes((i % 251) + 1 for i in range(size - 1))


def test_an_oversized_nal_fragments_into_well_formed_fu_a():
    nal = big_nal(4000)
    p = rtp.Packetiser(mtu=1400)
    packets = p.packetise([nal])
    payloads = [unpack(pkt)[5] for pkt in packets]
    assert len(payloads) == 3  # 1398 payload bytes per fragment, 3999 to carry

    for payload in payloads:
        assert payload[0] & 0x1F == rtp.NAL_FU_A
        assert payload[0] & 0x60 == nal[0] & 0x60  # NRI is preserved
        assert payload[1] & 0x1F == rtp.nal_type(nal)  # the carried type
        assert payload[1] & 0x20 == 0  # the reserved bit must be zero

    starts = [bool(p[1] & 0x80) for p in payloads]
    ends = [bool(p[1] & 0x40) for p in payloads]
    assert starts == [True, False, False]
    assert ends == [False, False, True]


def test_fu_a_fragments_reassemble_into_the_original_nal():
    nal = big_nal(4000)
    payloads = [unpack(p)[5] for p in rtp.Packetiser(mtu=1400).packetise([nal])]
    rebuilt = bytes([(payloads[0][0] & 0xE0) | (payloads[0][1] & 0x1F)])
    rebuilt += b"".join(p[2:] for p in payloads)
    assert rebuilt == nal


def test_no_fragment_exceeds_the_mtu():
    for mtu in (3, 50, 1400):
        packets = rtp.Packetiser(mtu=mtu).packetise([big_nal(4000)])
        assert max(len(unpack(p)[5]) for p in packets) <= mtu


def test_a_nal_exactly_at_the_mtu_is_not_fragmented():
    nal = big_nal(100)
    assert len(rtp.Packetiser(mtu=100, aggregate=False).packetise([nal])) == 1
    assert len(rtp.Packetiser(mtu=99, aggregate=False).packetise([nal])) == 2


def test_an_mtu_too_small_to_carry_a_payload_byte_is_refused():
    with pytest.raises(ValueError, match="at least 3"):
        rtp.Packetiser(mtu=2)


# -- STAP-A --------------------------------------------------------------
def test_parameter_sets_aggregate_into_one_stap_a():
    packets = rtp.Packetiser(mtu=1400).packetise([SPS, PPS, SEI, IDR])
    assert len(packets) == 2  # one STAP-A for the headers, one for the picture
    payload = unpack(packets[0])[5]
    assert payload[0] & 0x1F == rtp.NAL_STAP_A
    assert payload[0] & 0x60 == 0x60  # the highest NRI of what it carries

    carried, pos = [], 1
    while pos < len(payload):
        (length,) = struct.unpack_from(">H", payload, pos)
        carried.append(payload[pos + 2 : pos + 2 + length])
        pos += 2 + length
    assert carried == [SPS, PPS, SEI]


def test_a_picture_is_never_aggregated():
    # A STAP-A carrying a slice would be legal, but it is not what this does,
    # and the docstring says so. Two small slices stay two packets.
    packets = rtp.Packetiser(mtu=1400).packetise([IDR, SLICE_2])
    assert len(packets) == 2
    assert all(unpack(p)[5][0] & 0x1F != rtp.NAL_STAP_A for p in packets)


def test_a_lone_parameter_set_goes_out_as_a_single_nal_not_a_stap_a():
    (packet,) = rtp.Packetiser(mtu=1400).packetise([SPS])
    assert unpack(packet)[5] == SPS


def test_aggregation_can_be_switched_off():
    packets = rtp.Packetiser(mtu=1400, aggregate=False).packetise([SPS, PPS, SEI, IDR])
    assert len(packets) == 4
    assert [unpack(p)[5] for p in packets] == [SPS, PPS, SEI, IDR]


def test_aggregation_stops_at_the_mtu_rather_than_overflowing():
    sets = [big_nal(300, header=0x68) for _ in range(5)]
    packets = rtp.Packetiser(mtu=700).packetise(sets)
    assert len(packets) == 3  # 2 + 2 + 1 at 302 bytes of cost each
    assert all(len(unpack(p)[5]) <= 700 for p in packets)


# -- sequence, marker, clock ---------------------------------------------
def test_sequence_numbers_advance_by_one_per_packet():
    p = rtp.Packetiser(mtu=1400, seq=100)
    packets = p.packetise([SPS, PPS, IDR]) + p.packetise([SLICE])
    assert [unpack(pkt)[1] for pkt in packets] == [100, 101, 102]
    assert p.seq == 103


def test_sequence_numbers_wrap_at_sixteen_bits():
    p = rtp.Packetiser(mtu=1400, seq=0xFFFE)
    packets = p.packetise([IDR]) + p.packetise([IDR]) + p.packetise([IDR])
    assert [unpack(pkt)[1] for pkt in packets] == [0xFFFE, 0xFFFF, 0]


def test_the_marker_bit_lands_only_on_the_last_packet_of_an_access_unit():
    p = rtp.Packetiser(mtu=1400)
    packets = p.packetise([SPS, PPS, big_nal(4000)])
    markers = [unpack(pkt)[4] for pkt in packets]
    assert markers == [False, False, False, True]


def test_the_timestamp_is_one_per_access_unit_and_advances_between_them():
    p = rtp.Packetiser(mtu=1400, timestamp=1000)
    first = p.packetise([SPS, PPS, big_nal(4000)], advance=3000)
    second = p.packetise([SLICE], advance=3000)
    assert {unpack(pkt)[2] for pkt in first} == {1000}
    assert {unpack(pkt)[2] for pkt in second} == {4000}


def test_the_timestamp_wraps_at_thirty_two_bits():
    p = rtp.Packetiser(timestamp=(1 << 32) - 1000)
    p.packetise([IDR], advance=3000)
    assert p.timestamp == 2000


def test_the_ssrc_is_the_same_on_every_packet():
    p = rtp.Packetiser(mtu=1400)
    packets = p.packetise([SPS, PPS, big_nal(4000)]) + p.packetise([SLICE])
    assert len({unpack(pkt)[3] for pkt in packets}) == 1


# -- start offsets -------------------------------------------------------
def test_no_offset_starts_at_the_beginning():
    assert rtp.resolve_start_offset(None, 10.0, random.Random(0)) == 0.0


def test_a_random_offset_lands_inside_the_stream_and_repeats_with_a_seed():
    first = rtp.resolve_start_offset("random", 10.0, random.Random(5))
    second = rtp.resolve_start_offset("random", 10.0, random.Random(5))
    assert first == second
    assert 0.0 <= first < 10.0


def test_an_offset_past_the_end_wraps_rather_than_falling_off():
    assert rtp.resolve_start_offset(12.5, 10.0, random.Random(0)) == 2.5


def test_a_nonsense_offset_is_an_error():
    with pytest.raises(ValueError, match="'random'"):
        rtp.resolve_start_offset("middle", 10.0, random.Random(0))
    with pytest.raises(ValueError, match="negative"):
        rtp.resolve_start_offset(-1.0, 10.0, random.Random(0))


# -- the packetised source ----------------------------------------------
@pytest.fixture
def stream():
    # Ten pictures, parameter sets only at the front, which is the case that
    # makes repetition matter.
    nals = [SPS, PPS, IDR] + [SLICE] * 9
    return rtp.AnnexBStream.from_bytes(annex_b(*nals))


def take(source, n):
    out = []
    for item in source.packets():
        out.append(item)
        if len(out) >= n:
            break
    return out


def test_the_source_paces_one_period_per_frame_and_nothing_between_packets(stream):
    source = rtp.PacketisedSource(stream, fps=30.0)
    delays = [delay for delay, _ in take(source, 6)]
    assert delays[0] == pytest.approx(1 / 30)
    assert delays[1] == 0.0  # the second packet of the same frame
    assert sum(1 for d in delays if d > 0) == len([d for d in delays if d > 0])


def test_the_source_loops_without_the_clock_going_backwards(stream):
    source = rtp.PacketisedSource(stream, fps=30.0, parameter_set_period=0)
    packets = [data for _delay, data in take(source, 40)]
    stamps = [unpack(p)[2] for p in packets]
    assert stamps == sorted(stamps), "a timestamp that goes back makes a decoder give up"
    seqs = [unpack(p)[1] for p in packets]
    assert seqs == list(range(seqs[0], seqs[0] + len(seqs)))
    assert len(packets) > len(stream.units), "it must have passed the loop point"


def test_parameter_sets_recur_so_a_viewer_can_join_late(stream):
    source = rtp.PacketisedSource(stream, fps=30.0, parameter_set_period=4)
    payloads = [unpack(data)[5] for _delay, data in take(source, 30)]
    stap_a = [p for p in payloads if p[0] & 0x1F == rtp.NAL_STAP_A]
    # One at the file's own head plus one every fourth frame after it.
    assert len(stap_a) >= 5
    assert all(SPS in p and PPS in p for p in stap_a)


def test_turning_repetition_off_sends_the_parameter_sets_once_per_pass(stream):
    source = rtp.PacketisedSource(stream, fps=30.0, parameter_set_period=0)
    payloads = [unpack(data)[5] for _delay, data in take(source, 25)]
    assert sum(1 for p in payloads if p[0] & 0x1F == rtp.NAL_STAP_A) == 3


def test_a_random_start_offset_begins_mid_stream_and_is_reproducible(stream):
    a = rtp.PacketisedSource(stream, fps=30.0, start_offset="random", seed=11)
    b = rtp.PacketisedSource(stream, fps=30.0, start_offset="random", seed=11)
    assert a.start_index == b.start_index
    assert 0 <= a.start_index < len(stream.units)
    assert rtp.PacketisedSource(stream, fps=30.0, start_offset=0.1).start_index == 3


def test_a_source_needs_a_positive_frame_rate(stream):
    with pytest.raises(ValueError, match="fps"):
        rtp.PacketisedSource(stream, fps=0)


# -- the capture format --------------------------------------------------
def rtp_packet(seq, timestamp, ssrc=0xAABBCCDD, payload=b"\x41\x9a\xff"):
    return struct.pack(">BBHII", 0x80, 96, seq, timestamp, ssrc) + payload


def test_a_capture_round_trips(tmp_path):
    path = tmp_path / "c.rtpcap"
    packets = [(0.0, rtp_packet(1, 9000)), (0.033, rtp_packet(2, 12000))]
    assert capture.write_capture(path, packets) == 2
    assert list(capture.read_capture(path)) == packets


def test_a_capture_creates_its_parent_directory(tmp_path):
    path = tmp_path / "deep" / "c.rtpcap"
    capture.write_capture(path, [(0.0, rtp_packet(1, 0))])
    assert path.is_file()


def test_something_that_is_not_a_capture_is_rejected(tmp_path):
    path = tmp_path / "nope.rtpcap"
    path.write_bytes(b"not a capture at all")
    with pytest.raises(capture.CaptureFormatError, match="does not start"):
        list(capture.read_capture(path))
    short = tmp_path / "short.rtpcap"
    short.write_bytes(b"BEBO")
    with pytest.raises(capture.CaptureFormatError, match="too short"):
        list(capture.read_capture(short))


def test_a_future_capture_version_is_refused_rather_than_misread(tmp_path):
    path = tmp_path / "future.rtpcap"
    path.write_bytes(struct.pack("<8sHH", capture.MAGIC, 99, 0))
    with pytest.raises(capture.CaptureFormatError, match="version 99"):
        list(capture.read_capture(path))


def test_a_capture_cut_short_still_reads_up_to_the_break(tmp_path, caplog):
    path = tmp_path / "cut.rtpcap"
    capture.write_capture(path, [(0.0, rtp_packet(1, 0)), (0.033, rtp_packet(2, 3000))])
    path.write_bytes(path.read_bytes()[:-5])  # a Ctrl-C mid-write
    got = list(capture.read_capture(path))
    assert len(got) == 1
    assert "mid-packet" in caplog.text


# -- replay --------------------------------------------------------------
@pytest.fixture
def capture_file(tmp_path):
    """Three frames of two packets each, with a gap only between frames."""
    packets = []
    for frame in range(3):
        at = frame * 0.04
        packets.append((at, rtp_packet(frame * 2, 9000 + frame * 3000)))
        packets.append((at + 0.001, rtp_packet(frame * 2 + 1, 9000 + frame * 3000)))
    path = tmp_path / "flight.rtpcap"
    capture.write_capture(path, packets)
    return path


def test_a_replay_reproduces_the_recorded_gaps(capture_file):
    source = capture.ReplaySource.from_path(capture_file)
    delays = [delay for delay, _ in take(source, 6)]
    assert delays[0] == 0.0
    assert delays[1] == pytest.approx(0.001)
    assert delays[2] == pytest.approx(0.039)
    assert delays[3] == pytest.approx(0.001)


def test_a_replay_restamps_the_ssrc_and_renumbers_continuously(capture_file):
    source = capture.ReplaySource.from_path(capture_file, seed=3)
    packets = [data for _delay, data in take(source, 10)]
    assert {unpack(p)[3] for p in packets} == {source.ssrc}
    assert source.ssrc != 0xAABBCCDD
    seqs = [unpack(p)[1] for p in packets]
    assert seqs == list(range(seqs[0], seqs[0] + len(seqs)))


def test_a_replay_keeps_the_recorded_timestamp_deltas(capture_file):
    source = capture.ReplaySource.from_path(capture_file)
    stamps = [unpack(data)[2] for _delay, data in take(source, 6)]
    # Two packets per frame share a timestamp; frames are 3000 ticks apart.
    assert [b - a for a, b in pairwise(stamps)] == [0, 3000, 0, 3000, 0]


def test_a_replay_carries_the_clock_past_the_loop_point(capture_file):
    source = capture.ReplaySource.from_path(capture_file)
    items = take(source, 14)  # six recorded packets, then round again
    stamps = [unpack(data)[2] for _delay, data in items]
    assert stamps == sorted(stamps), "looping must not rewind the clock"
    # The gap the file cannot know: filled with its own median, one frame.
    assert items[6][0] == pytest.approx(0.039)
    assert stamps[6] - stamps[5] == 3000


def test_a_replay_can_start_part_way_in(capture_file):
    source = capture.ReplaySource.from_path(capture_file, start_offset=0.04)
    assert source.start_index == 2
    assert capture.ReplaySource.from_path(capture_file, start_offset="random", seed=2).start_index >= 0


def test_an_empty_capture_is_an_error_not_a_silent_stall(tmp_path):
    path = tmp_path / "empty.rtpcap"
    capture.write_capture(path, [])
    with pytest.raises(capture.CaptureFormatError, match="no packets"):
        capture.ReplaySource.from_path(path)


def test_the_recorder_writes_what_arrives_on_the_port_it_binds(tmp_path):
    import socket
    import threading

    path = tmp_path / "live.rtpcap"
    port = 0
    probe = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    probe.bind(("127.0.0.1", 0))
    port = probe.getsockname()[1]
    probe.close()

    done = threading.Event()
    count = []

    def run():
        count.append(capture.record(path, port=port, host="127.0.0.1", max_packets=2))
        done.set()

    thread = threading.Thread(target=run, daemon=True)
    thread.start()
    sender = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    for i in range(2):
        for _attempt in range(20):
            sender.sendto(rtp_packet(i, i * 3000), ("127.0.0.1", port))
            if done.wait(0.05) or len(count) or i == 0:
                break
    done.wait(5)
    sender.close()
    assert count == [2]
    got = list(capture.read_capture(path))
    assert len(got) == 2
    assert got[0][0] == 0.0  # times are relative to the first packet