AareDAQ/tests/unit/gui/test_prediction_subscriber.py

"""Checks the binary-frame sniff in _try_parse_json: JPEG frames (0xff lead
byte) must be rejected without attempting utf-8 decode (previously raised
UnicodeDecodeError on every frame)."""

from typing import cast

from aare.gui.threads.prediction_subscriber import PredictionSubscriber


def _parse(part: bytes):
    # Called unbound: _try_parse_json touches no instance state, so no
    # QThread/zmq construction is needed. cast keeps basedpyright happy
    # about the stand-in self.
    return PredictionSubscriber._try_parse_json(cast(PredictionSubscriber, object()), part)


def test_binary_jpeg_frame_is_not_json():
    assert _parse(b"\xff\xd8\xff\xe0somejpegbytes") is None


def test_json_dict_is_parsed():
    assert _parse(b'{"encoding": "jpeg"}') == {"encoding": "jpeg"}


def test_json_non_dict_is_rejected():
    assert _parse(b"[1, 2]") is None
    assert _parse(b"") is None


# --- Frame pacing and the multipart-safe conflate -------------------------
# Real socket on an inproc endpoint: the drain and the in-flight cap are the
# fix, so they run against libzmq rather than a mock.

import time

import cv2
import numpy as np
import pytest
import zmq
from PySide6.QtGui import QImage


@pytest.fixture
def sub(qapp):
    s = PredictionSubscriber("inproc://samcam-test")
    yield s
    # run() nulls both on its own cleanup; only tear down what is left.
    if s._sock is not None:
        s._sock.close(0)
    if s._ctx is not None:
        s._ctx.term()


def test_emit_image_caps_frames_in_flight(sub):
    image = QImage(4, 4, QImage.Format.Format_RGB888)
    for _ in range(sub.MAX_FRAMES_IN_FLIGHT + 1):
        sub._emit_image(image)
    assert sub._frames_in_flight == sub.MAX_FRAMES_IN_FLIGHT
    assert sub._dropped == 1

    for _ in range(sub.MAX_FRAMES_IN_FLIGHT + 1):
        sub.notify_frame_displayed()
    assert sub._frames_in_flight == 0  # floors at zero, never negative


def test_recv_latest_keeps_only_the_newest_message(sub):
    pub = sub._ctx.socket(zmq.PUB)
    pub.bind("inproc://samcam-test")
    try:
        # Slow joiner: publish probes until the subscription has propagated.
        deadline = time.monotonic() + 5
        while True:
            pub.send_multipart([b"probe"])
            try:
                sub._sock.recv_multipart(zmq.NOBLOCK)
                break
            except zmq.Again:
                assert time.monotonic() < deadline, "subscription never propagated"
                time.sleep(0.01)
        while True:  # flush any extra probes
            try:
                sub._sock.recv_multipart(zmq.NOBLOCK)
            except zmq.Again:
                break
        sub._dropped = 0

        for payload in (b"old", b"stale", b"newest"):
            pub.send_multipart([b"hdr", payload])
        assert sub._recv_latest() == [b"hdr", b"newest"]
        assert sub._dropped == 2
    finally:
        pub.close(0)


def test_run_decodes_one_frame_and_reports_the_drops(sub, qtbot):
    ok, jpeg = cv2.imencode(".jpg", np.zeros((4, 6, 3), dtype=np.uint8))
    assert ok
    frame = [b'{"encoding": "jpeg"}', jpeg.tobytes()]
    calls = []

    def one_frame_then_stop():
        calls.append(1)
        if len(calls) == 1:
            return frame
        sub.running = False
        raise zmq.Again

    sub._recv_latest = one_frame_then_stop  # type: ignore[method-assign]
    sub._fps_window_start = time.perf_counter() - 10  # force the fps report
    sub._dropped = 3
    images = []
    sub.image.connect(images.append)

    sub.run()

    assert len(images) == 1 and images[0].width() == 6 and images[0].height() == 4
    assert sub._frames_in_flight == 1
    assert sub._dropped == 0  # reported alongside the fps, then reset
    assert sub._camera_available
    assert sub._sock is None and sub._ctx is None  # run() cleaned up