105 lines
4.1 KiB
Python
105 lines
4.1 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from yovision_brain.config import parse_input_config
|
|
from yovision_brain.decode import DecoderError, DecoderPipeline, decode_packets
|
|
from yovision_brain.input import CancellationToken, InputPacket, LocalFileInput, SyntheticInput
|
|
|
|
|
|
def synthetic_packets():
|
|
config = parse_input_config(
|
|
{
|
|
"schema": "brain.internal.input/v1",
|
|
"logical_device_id": "synthetic-01",
|
|
"profile": {"id": "main", "width": 2, "height": 1, "fps": 5},
|
|
"source": {"kind": "synthetic", "seed": 3, "frame_count": 2},
|
|
}
|
|
)
|
|
return SyntheticInput(config).packets()
|
|
|
|
|
|
def local_packets(path: Path, *, width: int = 2, height: int = 1, chunk_size: int = 5):
|
|
config = parse_input_config(
|
|
{
|
|
"schema": "brain.internal.input/v1",
|
|
"logical_device_id": "local-01",
|
|
"profile": {"id": "archive", "width": width, "height": height, "fps": 25},
|
|
"source": {"kind": "local_file", "path": str(path), "chunk_size": chunk_size},
|
|
}
|
|
)
|
|
return LocalFileInput(config).packets()
|
|
|
|
|
|
def test_rgb24_pipeline_preserves_order_timestamps_and_metadata() -> None:
|
|
frames = list(decode_packets(synthetic_packets()))
|
|
assert [frame.sequence for frame in frames] == [0, 1]
|
|
assert [frame.timestamp_ns for frame in frames] == [0, 200_000_000]
|
|
assert all(frame.logical_device_id == "synthetic-01" for frame in frames)
|
|
assert all(frame.profile_id == "main" for frame in frames)
|
|
assert all((frame.width, frame.height, frame.pixel_format) == (2, 1, "rgb24") for frame in frames)
|
|
|
|
|
|
def test_rgb24_dimension_change_is_explicit() -> None:
|
|
packets = [
|
|
InputPacket(0, 0, "camera", "main", 1, 1, "rgb24", b"abc"),
|
|
InputPacket(1, 1, "camera", "main", 2, 1, "rgb24", b"abcdef"),
|
|
]
|
|
frames = list(decode_packets(packets))
|
|
assert [frame.dimensions_changed for frame in frames] == [False, True]
|
|
|
|
|
|
def test_invalid_rgb_payload_and_unsupported_format_are_clear() -> None:
|
|
bad = [InputPacket(0, 0, "camera", "main", 2, 2, "rgb24", b"short")]
|
|
with pytest.raises(DecoderError, match="expected 12"):
|
|
list(decode_packets(bad))
|
|
unknown = [InputPacket(0, 0, "camera", "main", 1, 1, "opaque", b"data")]
|
|
with pytest.raises(DecoderError, match="no decoder registered"):
|
|
list(DecoderPipeline().decode(unknown))
|
|
|
|
|
|
def test_y4m_local_video_decodes_across_input_chunks(tmp_path: Path) -> None:
|
|
video = tmp_path / "anonymous.y4m"
|
|
first, second = b"abcdef", b"ghijkl"
|
|
video.write_bytes(b"YUV4MPEG2 W2 H1 F25:1 C444\nFRAME\n" + first + b"FRAME\n" + second)
|
|
frames = list(decode_packets(local_packets(video)))
|
|
assert [frame.payload for frame in frames] == [first, second]
|
|
assert [frame.timestamp_ns for frame in frames] == [0, 40_000_000]
|
|
assert all(frame.pixel_format == "yuv444p" for frame in frames)
|
|
assert all((frame.width, frame.height) == (2, 1) for frame in frames)
|
|
assert all(frame.profile_id == "archive" for frame in frames)
|
|
|
|
|
|
def test_y4m_clean_eof_and_cancellation_are_normal(tmp_path: Path) -> None:
|
|
video = tmp_path / "empty.y4m"
|
|
video.write_bytes(b"YUV4MPEG2 W2 H1 F25:1 C444\n")
|
|
assert list(decode_packets(local_packets(video))) == []
|
|
|
|
token = CancellationToken()
|
|
token.cancel()
|
|
assert list(decode_packets(local_packets(video), token)) == []
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("payload", "message"),
|
|
[
|
|
(b"not-video\n", "expected YUV4MPEG2"),
|
|
(b"YUV4MPEG2 W2 H1 F25:1 C420\n", "only C444"),
|
|
(b"YUV4MPEG2 W2 H1 F25:1 C444\nFRAME\nabc", "truncated Y4M frame"),
|
|
],
|
|
)
|
|
def test_y4m_damage_and_unsupported_content_are_clear(tmp_path: Path, payload: bytes, message: str) -> None:
|
|
video = tmp_path / "broken.y4m"
|
|
video.write_bytes(payload)
|
|
with pytest.raises(DecoderError, match=message):
|
|
list(decode_packets(local_packets(video)))
|
|
|
|
|
|
def test_y4m_profile_dimension_mismatch_is_rejected(tmp_path: Path) -> None:
|
|
video = tmp_path / "mismatch.y4m"
|
|
video.write_bytes(b"YUV4MPEG2 W2 H1 F25:1 C444\n")
|
|
with pytest.raises(DecoderError, match="do not match"):
|
|
list(decode_packets(local_packets(video, width=3)))
|