446 lines
15 KiB
Python
446 lines
15 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import threading
|
||
|
|
from concurrent.futures import ThreadPoolExecutor
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
import cmvr_edge_ai.detection.video as video_module
|
||
|
|
from cmvr_edge_ai.contracts import CameraIntrinsics, ImageFrame
|
||
|
|
from cmvr_edge_ai.core import ComponentContext, Emission, Envelope
|
||
|
|
from cmvr_edge_ai.detection.video import (
|
||
|
|
DecodedVideoFrame,
|
||
|
|
VideoDecodeError,
|
||
|
|
VideoDecoderOperator,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class _BackendLog:
|
||
|
|
created_codecs: list[str] = field(default_factory=list)
|
||
|
|
decoded_packets: list[bytes] = field(default_factory=list)
|
||
|
|
worker_threads: list[int] = field(default_factory=list)
|
||
|
|
close_count: int = 0
|
||
|
|
|
||
|
|
|
||
|
|
class _FakeBackend:
|
||
|
|
def __init__(self, codec: str, log: _BackendLog) -> None:
|
||
|
|
self._log = log
|
||
|
|
self._closed = False
|
||
|
|
log.created_codecs.append(codec)
|
||
|
|
log.worker_threads.append(threading.get_ident())
|
||
|
|
|
||
|
|
def decode(self, data: bytes) -> tuple[DecodedVideoFrame, ...]:
|
||
|
|
if self._closed:
|
||
|
|
raise RuntimeError("fake decoder is closed")
|
||
|
|
self._log.decoded_packets.append(data)
|
||
|
|
self._log.worker_threads.append(threading.get_ident())
|
||
|
|
if data == b"decode-error":
|
||
|
|
raise VideoDecodeError("synthetic packet discontinuity")
|
||
|
|
if data == b"buffered":
|
||
|
|
return ()
|
||
|
|
if data == b"two-frames":
|
||
|
|
return (
|
||
|
|
DecodedVideoFrame(b"first", 2, 2),
|
||
|
|
DecodedVideoFrame(b"second", 3, 2),
|
||
|
|
)
|
||
|
|
return (DecodedVideoFrame(b"decoded:" + data, 2, 2),)
|
||
|
|
|
||
|
|
def close(self) -> None:
|
||
|
|
if not self._closed:
|
||
|
|
self._closed = True
|
||
|
|
self._log.close_count += 1
|
||
|
|
self._log.worker_threads.append(threading.get_ident())
|
||
|
|
|
||
|
|
|
||
|
|
def _factory(log: _BackendLog): # type: ignore[no-untyped-def]
|
||
|
|
return lambda codec: _FakeBackend(codec, log)
|
||
|
|
|
||
|
|
|
||
|
|
def _image(
|
||
|
|
*,
|
||
|
|
codec: str = "H264",
|
||
|
|
key_frame: bool = False,
|
||
|
|
width: int = 640,
|
||
|
|
height: int = 480,
|
||
|
|
data: bytes = b"packet",
|
||
|
|
intrinsics: CameraIntrinsics | None = None,
|
||
|
|
) -> ImageFrame:
|
||
|
|
return ImageFrame(
|
||
|
|
data=data,
|
||
|
|
width=width,
|
||
|
|
height=height,
|
||
|
|
pixel_format="encoded",
|
||
|
|
codec=codec,
|
||
|
|
is_key_frame=key_frame,
|
||
|
|
intrinsics=intrinsics,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _envelope(
|
||
|
|
image: Any,
|
||
|
|
*,
|
||
|
|
source_id: str = "camera-1",
|
||
|
|
session_id: str | None = "session-1",
|
||
|
|
sequence: int = 0,
|
||
|
|
) -> Envelope[Any]:
|
||
|
|
return Envelope(
|
||
|
|
payload=image,
|
||
|
|
schema_name="ImageFrame",
|
||
|
|
schema_version=1,
|
||
|
|
source_id=source_id,
|
||
|
|
session_id=session_id,
|
||
|
|
sequence=sequence,
|
||
|
|
captured_at_ns=123,
|
||
|
|
received_at_ns=456,
|
||
|
|
deadline_ns=999,
|
||
|
|
trace_id="trace-1",
|
||
|
|
attributes={"transport": "test"},
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
async def _setup(operator: VideoDecoderOperator, executor: Any = None) -> None:
|
||
|
|
metadata = {} if executor is None else {"thread_executor": executor}
|
||
|
|
await operator.setup(
|
||
|
|
ComponentContext(
|
||
|
|
pipeline_id="test-pipeline",
|
||
|
|
node_id="decoder",
|
||
|
|
shutdown_event=asyncio.Event(),
|
||
|
|
metadata=metadata,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
async def _keep_restricted_event_loop_responsive() -> None:
|
||
|
|
"""Provide wakeups where sandboxed self-pipe notifications are unavailable."""
|
||
|
|
|
||
|
|
while True:
|
||
|
|
await asyncio.sleep(0.01)
|
||
|
|
|
||
|
|
|
||
|
|
def _run_with_ticker(coroutine: Any) -> Any:
|
||
|
|
async def runner() -> Any:
|
||
|
|
loop = asyncio.get_running_loop()
|
||
|
|
host_executor = ThreadPoolExecutor(max_workers=2)
|
||
|
|
loop.set_default_executor(host_executor)
|
||
|
|
ticker = asyncio.create_task(_keep_restricted_event_loop_responsive())
|
||
|
|
try:
|
||
|
|
return await coroutine
|
||
|
|
finally:
|
||
|
|
ticker.cancel()
|
||
|
|
await asyncio.gather(ticker, return_exceptions=True)
|
||
|
|
host_executor.shutdown(wait=True, cancel_futures=True)
|
||
|
|
# Avoid a second shutdown by asyncio.run() after ownership above.
|
||
|
|
loop._default_executor = None # type: ignore[attr-defined]
|
||
|
|
|
||
|
|
return asyncio.run(runner())
|
||
|
|
|
||
|
|
|
||
|
|
def test_raw_frame_is_passed_through_without_loading_a_backend() -> None:
|
||
|
|
factory_calls: list[str] = []
|
||
|
|
|
||
|
|
def forbidden_factory(codec: str): # type: ignore[no-untyped-def]
|
||
|
|
factory_calls.append(codec)
|
||
|
|
raise AssertionError("raw frames must not construct a decoder")
|
||
|
|
|
||
|
|
operator = VideoDecoderOperator("decoder", {}, backend_factory=forbidden_factory)
|
||
|
|
envelope = _envelope(
|
||
|
|
ImageFrame(
|
||
|
|
data=b"raw-bgr",
|
||
|
|
width=2,
|
||
|
|
height=2,
|
||
|
|
pixel_format="BGR8",
|
||
|
|
codec="none",
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
async def exercise() -> Emission:
|
||
|
|
await _setup(operator)
|
||
|
|
result = await operator.process(envelope)
|
||
|
|
await operator.stop()
|
||
|
|
assert isinstance(result, Emission)
|
||
|
|
return result
|
||
|
|
|
||
|
|
emission = _run_with_ticker(exercise())
|
||
|
|
|
||
|
|
assert emission.port == "frames"
|
||
|
|
assert emission.envelope is envelope
|
||
|
|
assert factory_calls == []
|
||
|
|
|
||
|
|
|
||
|
|
def test_waits_for_keyframe_decodes_off_loop_and_preserves_envelope_metadata() -> None:
|
||
|
|
log = _BackendLog()
|
||
|
|
operator = VideoDecoderOperator("decoder", {}, backend_factory=_factory(log))
|
||
|
|
intrinsics = CameraIntrinsics(100.0, 101.0, 50.0, 51.0)
|
||
|
|
event_loop_thread = threading.get_ident()
|
||
|
|
|
||
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
||
|
|
with ThreadPoolExecutor(max_workers=1) as executor:
|
||
|
|
await _setup(operator, executor)
|
||
|
|
dropped = await operator.process(
|
||
|
|
_envelope(_image(key_frame=False), sequence=10)
|
||
|
|
)
|
||
|
|
key_envelope = _envelope(
|
||
|
|
_image(
|
||
|
|
codec=".H264",
|
||
|
|
key_frame=True,
|
||
|
|
data=b"key",
|
||
|
|
intrinsics=intrinsics,
|
||
|
|
),
|
||
|
|
sequence=11,
|
||
|
|
)
|
||
|
|
decoded = await operator.process(key_envelope)
|
||
|
|
health = await operator.health()
|
||
|
|
await operator.stop()
|
||
|
|
return dropped, key_envelope, decoded, health
|
||
|
|
|
||
|
|
dropped, original, result, health = _run_with_ticker(exercise())
|
||
|
|
|
||
|
|
assert dropped is None
|
||
|
|
assert isinstance(result, tuple) and len(result) == 1
|
||
|
|
emission = result[0]
|
||
|
|
assert emission.port == "frames"
|
||
|
|
assert isinstance(emission.envelope.payload, ImageFrame)
|
||
|
|
assert emission.envelope.payload == ImageFrame(
|
||
|
|
data=b"decoded:key",
|
||
|
|
width=2,
|
||
|
|
height=2,
|
||
|
|
pixel_format="BGR8",
|
||
|
|
codec="none",
|
||
|
|
is_key_frame=False,
|
||
|
|
intrinsics=intrinsics,
|
||
|
|
)
|
||
|
|
assert emission.envelope.source_id == original.source_id
|
||
|
|
assert emission.envelope.session_id == original.session_id
|
||
|
|
assert emission.envelope.sequence == original.sequence
|
||
|
|
assert emission.envelope.trace_id == original.trace_id
|
||
|
|
assert emission.envelope.captured_at_ns == original.captured_at_ns
|
||
|
|
assert emission.envelope.received_at_ns == original.received_at_ns
|
||
|
|
assert emission.envelope.attributes == original.attributes
|
||
|
|
assert log.created_codecs == ["h264"]
|
||
|
|
assert log.decoded_packets == [b"key"]
|
||
|
|
assert log.close_count == 1
|
||
|
|
assert log.worker_threads
|
||
|
|
assert all(thread_id != event_loop_thread for thread_id in log.worker_threads)
|
||
|
|
assert "decoded_frames=1" in health.detail
|
||
|
|
assert "dropped_until_keyframe=1" in health.detail
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
("change", "expected_codec"),
|
||
|
|
[
|
||
|
|
({"source_id": "camera-2"}, "h264"),
|
||
|
|
({"session_id": "session-2"}, "h264"),
|
||
|
|
({"codec": "H265"}, "hevc"),
|
||
|
|
({"width": 800}, "h264"),
|
||
|
|
({"height": 600}, "h264"),
|
||
|
|
({"sequence": 4}, "h264"),
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_stream_identity_change_or_sequence_gap_resets_and_waits_for_keyframe(
|
||
|
|
change: dict[str, Any],
|
||
|
|
expected_codec: str,
|
||
|
|
) -> None:
|
||
|
|
log = _BackendLog()
|
||
|
|
operator = VideoDecoderOperator("decoder", {}, backend_factory=_factory(log))
|
||
|
|
|
||
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
||
|
|
await _setup(operator)
|
||
|
|
first = await operator.process(
|
||
|
|
_envelope(_image(key_frame=True, data=b"initial"), sequence=1)
|
||
|
|
)
|
||
|
|
|
||
|
|
source_id = str(change.get("source_id", "camera-1"))
|
||
|
|
session_id = change.get("session_id", "session-1")
|
||
|
|
sequence = int(change.get("sequence", 2))
|
||
|
|
codec = str(change.get("codec", "H264"))
|
||
|
|
width = int(change.get("width", 640))
|
||
|
|
height = int(change.get("height", 480))
|
||
|
|
changed_image = _image(
|
||
|
|
codec=codec,
|
||
|
|
key_frame=False,
|
||
|
|
width=width,
|
||
|
|
height=height,
|
||
|
|
data=b"must-drop",
|
||
|
|
)
|
||
|
|
dropped = await operator.process(
|
||
|
|
_envelope(
|
||
|
|
changed_image,
|
||
|
|
source_id=source_id,
|
||
|
|
session_id=session_id,
|
||
|
|
sequence=sequence,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
recovered = await operator.process(
|
||
|
|
_envelope(
|
||
|
|
_image(
|
||
|
|
codec=codec,
|
||
|
|
key_frame=True,
|
||
|
|
width=width,
|
||
|
|
height=height,
|
||
|
|
data=b"recovered",
|
||
|
|
),
|
||
|
|
source_id=source_id,
|
||
|
|
session_id=session_id,
|
||
|
|
sequence=sequence + 1,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
await operator.stop()
|
||
|
|
return first, dropped, recovered
|
||
|
|
|
||
|
|
first, dropped, recovered = _run_with_ticker(exercise())
|
||
|
|
|
||
|
|
assert isinstance(first, tuple)
|
||
|
|
assert dropped is None
|
||
|
|
assert isinstance(recovered, tuple)
|
||
|
|
assert log.created_codecs == ["h264", expected_codec]
|
||
|
|
assert log.decoded_packets == [b"initial", b"recovered"]
|
||
|
|
assert log.close_count == 2
|
||
|
|
|
||
|
|
|
||
|
|
def test_recoverable_decode_error_resets_and_suppresses_until_next_keyframe() -> None:
|
||
|
|
log = _BackendLog()
|
||
|
|
operator = VideoDecoderOperator("decoder", {}, backend_factory=_factory(log))
|
||
|
|
|
||
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
||
|
|
await _setup(operator)
|
||
|
|
failed = await operator.process(
|
||
|
|
_envelope(
|
||
|
|
_image(key_frame=True, data=b"decode-error"),
|
||
|
|
sequence=1,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
suppressed = await operator.process(
|
||
|
|
_envelope(_image(key_frame=False, data=b"p-frame"), sequence=2)
|
||
|
|
)
|
||
|
|
recovered = await operator.process(
|
||
|
|
_envelope(_image(key_frame=True, data=b"new-key"), sequence=3)
|
||
|
|
)
|
||
|
|
health = await operator.health()
|
||
|
|
await operator.stop()
|
||
|
|
await operator.stop()
|
||
|
|
return failed, suppressed, recovered, health
|
||
|
|
|
||
|
|
failed, suppressed, recovered, health = _run_with_ticker(exercise())
|
||
|
|
|
||
|
|
assert failed is None
|
||
|
|
assert suppressed is None
|
||
|
|
assert isinstance(recovered, tuple)
|
||
|
|
assert log.created_codecs == ["h264", "h264"]
|
||
|
|
assert log.decoded_packets == [b"decode-error", b"new-key"]
|
||
|
|
assert log.close_count == 2
|
||
|
|
assert "decode_errors=1" in health.detail
|
||
|
|
assert "resets=1" in health.detail
|
||
|
|
|
||
|
|
|
||
|
|
def test_backend_can_emit_zero_or_multiple_decoded_frames() -> None:
|
||
|
|
log = _BackendLog()
|
||
|
|
operator = VideoDecoderOperator("decoder", {}, backend_factory=_factory(log))
|
||
|
|
|
||
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
||
|
|
await _setup(operator)
|
||
|
|
buffered = await operator.process(
|
||
|
|
_envelope(_image(key_frame=True, data=b"buffered"), sequence=1)
|
||
|
|
)
|
||
|
|
multiple = await operator.process(
|
||
|
|
_envelope(_image(key_frame=False, data=b"two-frames"), sequence=2)
|
||
|
|
)
|
||
|
|
await operator.stop()
|
||
|
|
return buffered, multiple
|
||
|
|
|
||
|
|
buffered, multiple = _run_with_ticker(exercise())
|
||
|
|
|
||
|
|
assert buffered is None
|
||
|
|
assert isinstance(multiple, tuple) and len(multiple) == 2
|
||
|
|
assert [item.envelope.payload.data for item in multiple] == [b"first", b"second"]
|
||
|
|
|
||
|
|
|
||
|
|
def test_raw_transition_closes_an_existing_encoded_decoder() -> None:
|
||
|
|
log = _BackendLog()
|
||
|
|
operator = VideoDecoderOperator("decoder", {}, backend_factory=_factory(log))
|
||
|
|
raw_envelope = _envelope(
|
||
|
|
ImageFrame(b"raw", 1, 1, "BGR8", codec="raw"),
|
||
|
|
sequence=2,
|
||
|
|
)
|
||
|
|
|
||
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
||
|
|
await _setup(operator)
|
||
|
|
await operator.process(
|
||
|
|
_envelope(_image(key_frame=True, data=b"key"), sequence=1)
|
||
|
|
)
|
||
|
|
raw = await operator.process(raw_envelope)
|
||
|
|
await operator.stop()
|
||
|
|
return raw
|
||
|
|
|
||
|
|
result = _run_with_ticker(exercise())
|
||
|
|
|
||
|
|
assert isinstance(result, Emission)
|
||
|
|
assert result.envelope is raw_envelope
|
||
|
|
assert log.close_count == 1
|
||
|
|
|
||
|
|
|
||
|
|
def test_rejects_wrong_payload_unsupported_codec_and_invalid_dimensions() -> None:
|
||
|
|
log = _BackendLog()
|
||
|
|
operator = VideoDecoderOperator("decoder", {}, backend_factory=_factory(log))
|
||
|
|
|
||
|
|
async def exercise() -> None:
|
||
|
|
with pytest.raises(TypeError, match="expected ImageFrame"):
|
||
|
|
await operator.process(_envelope("not-an-image"))
|
||
|
|
with pytest.raises(ValueError, match="unsupported encoded video codec"):
|
||
|
|
await operator.process(
|
||
|
|
_envelope(_image(codec="vp9", key_frame=True), sequence=1)
|
||
|
|
)
|
||
|
|
with pytest.raises(ValueError, match="dimensions must be positive"):
|
||
|
|
await operator.process(
|
||
|
|
_envelope(_image(width=0, key_frame=True), sequence=1)
|
||
|
|
)
|
||
|
|
|
||
|
|
_run_with_ticker(exercise())
|
||
|
|
assert log.created_codecs == []
|
||
|
|
|
||
|
|
|
||
|
|
def test_default_backend_imports_pyav_only_for_an_encoded_keyframe(
|
||
|
|
monkeypatch: pytest.MonkeyPatch,
|
||
|
|
) -> None:
|
||
|
|
imported: list[str] = []
|
||
|
|
|
||
|
|
def missing_import(name: str): # type: ignore[no-untyped-def]
|
||
|
|
imported.append(name)
|
||
|
|
raise ImportError("synthetic missing dependency")
|
||
|
|
|
||
|
|
monkeypatch.setattr(video_module, "import_module", missing_import)
|
||
|
|
operator = VideoDecoderOperator("decoder", {})
|
||
|
|
raw = _envelope(ImageFrame(b"raw", 1, 1, "BGR8", codec="none"))
|
||
|
|
|
||
|
|
async def exercise() -> None:
|
||
|
|
assert isinstance(await operator.process(raw), Emission)
|
||
|
|
assert imported == []
|
||
|
|
assert (
|
||
|
|
await operator.process(_envelope(_image(key_frame=False), sequence=1))
|
||
|
|
is None
|
||
|
|
)
|
||
|
|
assert imported == []
|
||
|
|
with pytest.raises(RuntimeError, match="optional video dependencies"):
|
||
|
|
await operator.process(_envelope(_image(key_frame=True), sequence=2))
|
||
|
|
|
||
|
|
_run_with_ticker(exercise())
|
||
|
|
assert imported == ["av"]
|
||
|
|
|
||
|
|
|
||
|
|
def test_setup_rejects_a_non_executor_thread_executor() -> None:
|
||
|
|
operator = VideoDecoderOperator("decoder", {})
|
||
|
|
|
||
|
|
async def exercise() -> None:
|
||
|
|
with pytest.raises(TypeError, match="must be an Executor"):
|
||
|
|
await _setup(operator, object())
|
||
|
|
|
||
|
|
_run_with_ticker(exercise())
|
||
|
|
|
||
|
|
|
||
|
|
def test_constructor_rejects_unknown_parameters() -> None:
|
||
|
|
with pytest.raises(ValueError, match="does not support parameters: typo"):
|
||
|
|
VideoDecoderOperator("decoder", {"typo": True})
|