diff --git a/src/aare/gui/main_window.py b/src/aare/gui/main_window.py index 1094658f..40b98ec4 100644 --- a/src/aare/gui/main_window.py +++ b/src/aare/gui/main_window.py @@ -32,8 +32,10 @@ from PySide6.QtGui import ( QColor, QCursor, QGuiApplication, + QImage, QKeySequence, QPalette, + QPixmap, ) from PySide6.QtWidgets import ( QAbstractButton, @@ -1049,9 +1051,7 @@ class MainWindow(QMainWindow): if sample_feed_addr is not None: logger.debug(f"Starting prediction subscriber thread {sample_feed_addr}") self.prediction_thread = PredictionSubscriber(pred_zmq_url=sample_feed_addr, topic=b"") - self.prediction_thread.image.connect(self.sample_camera.update_pixmap) - self.prediction_thread.image.connect(self.compact_sample_camera.update_pixmap) - self.prediction_thread.image.connect(self.portrait_sample_camera.update_pixmap) + self.prediction_thread.image.connect(self._on_sample_camera_frame) self.prediction_thread.prediction.connect(self.sample_camera.update_detections) self.prediction_thread.prediction.connect(self.compact_sample_camera.update_detections) self.prediction_thread.prediction.connect(self.portrait_sample_camera.update_detections) @@ -1609,6 +1609,33 @@ class MainWindow(QMainWindow): settings.setValue("samcam/compact_overlay_legend", overlay["compact_overlay_legend"]) settings.setValue("samcam/target_color", overlay["target_color"]) + @Slot(QImage) + def _on_sample_camera_frame(self, image: QImage) -> None: + """Hand one frame to whichever sample camera view is on screen. + + The subscriber emits from its own thread, so this runs queued on the + GUI thread. Painting only the visible view costs one repaint per frame + instead of three, and the acknowledgement at the end is what paces the + subscriber to the rate the GUI can actually keep up with. + + The subscriber sends a QImage: QPixmap is a GUI-thread-only class, so + it is built here rather than in the worker. + """ + try: + pixmap = QPixmap.fromImage(image) + for view in ( + self.sample_camera, + self.compact_sample_camera, + self.portrait_sample_camera, + ): + if view.isVisible(): + view.update_pixmap(pixmap) + finally: + # The slot is only connected once the subscriber exists, but the + # attribute is Optional so the guard is what the type checker needs. + if self.prediction_thread is not None: + self.prediction_thread.notify_frame_displayed() + @Slot(bool) def _on_sample_camera_availability_changed(self, available: bool) -> None: if available: diff --git a/src/aare/gui/styles.py b/src/aare/gui/styles.py index ed30b0cf..ec619d1d 100644 --- a/src/aare/gui/styles.py +++ b/src/aare/gui/styles.py @@ -471,9 +471,6 @@ LEGEND_BG = "#eff1f5" # base LEGEND_TEXT = "#4c4f69" # text TOOLTIP_TEXT = "#4c4f69" # camera coords tooltip pen — NOT the QToolTip popup SCALE_BAR_GREY = "#8c8fa1" # hover HUD scale bar (Latte overlay1 grey) -MARK_TOOLTIP_GOLD = "#df8e1d" # yellow -MARK_TOOLTIP_ORANGE = "#fe640b" # peach -MARK_TOOLTIP_RED = "#d20f39" # red MARK_BADGE_BG = "#fe640b" # peach # Prediction class overlay colors. The old pure-green vs CSS-green split @@ -503,22 +500,11 @@ BOOKMARK_COLORS = { "lime": "#179299", # teal — Latte has one green; teal keeps the pair distinct } -# -- Busy overlay (per-source color coding) --------------------------------- +# -- Busy overlay: PSI red for every busy state, blue for robot cooling ------ # Catppuccin Latte accents; BORDER/DOT are 25%/20% mixes toward Latte base. BUSY_YELLOW = "#df8e1d" # yellow BUSY_YELLOW_BORDER = "#ebd8bf" BUSY_YELLOW_DOT = "#ecddca" -BUSY_YELLOW_TEXT_DARK = "#4c4f69" # text -BUSY_PURPLE = "#8839ef" # mauve -BUSY_PURPLE_BORDER = "#d5c3f4" -BUSY_PURPLE_DOT = "#daccf4" -BUSY_RED_BADGE = "#d20f39" # red -BUSY_RED_FILL = "#d20f39" # red -BUSY_RED_BORDER = "#e8b8c6" -BUSY_RED_DOT = "#e9c4cf" -BUSY_ORANGE = "#fe640b" # peach -BUSY_ORANGE_BORDER = "#f3ceba" -BUSY_ORANGE_DOT = "#f2d5c6" BUSY_BLUE = "#1e66f5" # blue BUSY_BLUE_BORDER = "#bbcef5" BUSY_BLUE_DOT = "#c5d5f5" diff --git a/src/aare/gui/threads/prediction_subscriber.py b/src/aare/gui/threads/prediction_subscriber.py index 85b5a24f..e5c40980 100644 --- a/src/aare/gui/threads/prediction_subscriber.py +++ b/src/aare/gui/threads/prediction_subscriber.py @@ -8,7 +8,7 @@ from aarecommon.config.logger import setup_logger from aarecommon.math.autofocus import focus_measure_edges from aarecommon.models.models import DAQStatusModel from PySide6.QtCore import QThread, Signal, Slot -from PySide6.QtGui import QImage, QPixmap +from PySide6.QtGui import QImage from aare.gui.constants import LOGGER_NAME @@ -18,18 +18,29 @@ logger = setup_logger(LOGGER_NAME) class PredictionSubscriber(QThread): prediction = Signal(dict) target_point = Signal(dict) - image = Signal(QPixmap) + image = Signal(QImage) focus_measure = Signal(float) fps_measure = Signal(float) camera_availability_changed = Signal(bool) camera_error = Signal(str) + #: How many frames may sit unpainted in the GUI event queue (see _emit_image). + MAX_FRAMES_IN_FLIGHT = 2 + def __init__(self, pred_zmq_url: str, topic: bytes | str = b"", parent=None): super().__init__(parent) self._ctx = zmq.Context() self._sock = self._ctx.socket(zmq.SUB) self._sock.setsockopt(zmq.RCVTIMEO, 500) self._sock.setsockopt(zmq.LINGER, 0) + # Cap the backlog the socket may hold. Both must be set before connect() + # to take effect on the pipe. Defaults let the SUB queue grow to 1000 + # messages (20 s of video at 50 Hz) with megabytes more hidden in the + # kernel buffer, which is the backlog _recv_latest would then have to + # throw away every frame. RCVBUF stays well above the ~2.5 MB/s the + # stream actually needs, so it costs no throughput. + self._sock.setsockopt(zmq.RCVHWM, 4) + self._sock.setsockopt(zmq.RCVBUF, 512 * 1024) self._emit_images = True self.running = True @@ -50,6 +61,9 @@ class PredictionSubscriber(QThread): self._camera_available = False self._last_camera_error: str | None = None + self._dropped = 0 # stale frames skipped since the last fps report + self._frames_in_flight = 0 + if isinstance(topic, str): self._sock.setsockopt_string(zmq.SUBSCRIBE, topic) elif isinstance(topic, bytes): @@ -132,9 +146,34 @@ class PredictionSubscriber(QThread): return None - def _rgb_to_pixmap(self, rgb: np.ndarray) -> QPixmap: - qimage = QImage(rgb.data, rgb.shape[1], rgb.shape[0], QImage.Format.Format_RGB888).copy() - return QPixmap.fromImage(qimage) + def _rgb_to_qimage(self, rgb: np.ndarray) -> QImage: + # .copy() because QImage does not own the numpy buffer it wraps. + return QImage(rgb.data, rgb.shape[1], rgb.shape[0], QImage.Format.Format_RGB888).copy() + + def _emit_image(self, image: QImage) -> None: + """Emit a frame only while the GUI is keeping up, dropping it otherwise. + + `image` is a queued cross-thread signal and Qt's event queue has no + high-water mark, so if the GUI paints slower than frames arrive the + queue grows without bound and the picture on screen falls further + behind for as long as the stream runs. Capping the frames in flight + paces the feed to whatever the GUI can actually sustain. + + The cap is 2 rather than 1 so that one frame can be queued while the + GUI paints the other: at 1 the GUI idles until the worker finishes + decoding the next frame. The second slot costs one frame of latency + (20 ms at 50 Hz) and buys back that idle time. + """ + if self._frames_in_flight >= self.MAX_FRAMES_IN_FLIGHT: + self._dropped += 1 + return + + self._frames_in_flight += 1 + self.image.emit(image) + + def notify_frame_displayed(self) -> None: + """Called from the GUI thread once a frame has been handed to the views.""" + self._frames_in_flight = max(0, self._frames_in_flight - 1) def _emit_focus_measure_if_enabled(self, rgb: np.ndarray) -> None: if not self._measure_focus: @@ -151,13 +190,34 @@ class PredictionSubscriber(QThread): sharpness = focus_measure_edges(gray, self._focus_mask) self.focus_measure.emit(sharpness) + def _recv_latest(self) -> list[bytes]: + """Return the newest queued message, discarding the backlog behind it. + + zmq.CONFLATE is unusable here: it keeps only the last *part*, which + shreds multipart messages. recv_multipart is atomic though — a + non-blocking recv either yields a whole message or raises Again — so + draining in a loop is a multipart-safe conflate. The drain runs after + the (slow) decode of the previous frame, so it eats exactly the + backlog that piled up while we were busy, and the stream is always + consumed at line rate no matter how long a frame takes to decode. + """ + # Bind locally: run() nulls the attribute on cleanup, so it is Optional + # to the type checker even though it cannot be None while run() loops. + sock = self._sock + assert sock is not None, "_recv_latest called after the socket was closed" + parts = sock.recv_multipart() # honours RCVTIMEO, may raise zmq.Again + while True: + try: + parts = sock.recv_multipart(zmq.NOBLOCK) + except zmq.Again: + return parts + self._dropped += 1 + def run(self): - self._debug_last_log_ts = time.perf_counter() - self._debug_msg_count = 0 try: while self.running: try: - parts = self._sock.recv_multipart() + parts = self._recv_latest() except zmq.Again: now = time.perf_counter() elapsed = now - self._fps_window_start @@ -185,6 +245,13 @@ class PredictionSubscriber(QThread): if elapsed >= self._fps_emit_period_s: fps = self._fps_frame_count / elapsed if elapsed > 0 else 0.0 self.fps_measure.emit(float(fps)) + if self._dropped: + logger.debug( + "Sample camera: skipped %d stale frame(s), showing %.1f fps", + self._dropped, + fps, + ) + self._dropped = 0 self._fps_window_start = now self._fps_frame_count = 0 @@ -221,7 +288,7 @@ class PredictionSubscriber(QThread): self._set_camera_available(True) self._emit_focus_measure_if_enabled(rgb) if self.running: - self.image.emit(self._rgb_to_pixmap(rgb)) + self._emit_image(self._rgb_to_qimage(rgb)) elif self._emit_images: self._set_camera_available( False, "Sample camera feed unavailable: no frame header in zmq stream" diff --git a/src/aare/gui/widgets/busy_overlay.py b/src/aare/gui/widgets/busy_overlay.py index cbf356b7..19b6fd6d 100644 --- a/src/aare/gui/widgets/busy_overlay.py +++ b/src/aare/gui/widgets/busy_overlay.py @@ -11,23 +11,12 @@ from aare.gui.styles import ( BUSY_BLUE, BUSY_BLUE_BORDER, BUSY_BLUE_DOT, - BUSY_ORANGE, - BUSY_ORANGE_BORDER, - BUSY_ORANGE_DOT, BUSY_PSI_RED, BUSY_PSI_RED_BORDER, BUSY_PSI_RED_DOT, - BUSY_PURPLE, - BUSY_PURPLE_BORDER, - BUSY_PURPLE_DOT, - BUSY_RED_BADGE, - BUSY_RED_BORDER, - BUSY_RED_DOT, - BUSY_RED_FILL, BUSY_YELLOW, BUSY_YELLOW_BORDER, BUSY_YELLOW_DOT, - BUSY_YELLOW_TEXT_DARK, SHADOW, WHITE, qcolor, @@ -172,6 +161,21 @@ def draw_busy_status_text( painter.drawText(QPoint(x, baseline), style.text) +# The one "not yours" badge. Module-level so the sample camera can paint the +# same box while the session state is still unknown (no status yet), instead +# of a differently colored hand-rolled one. +VIEWING_MODE_STYLE = BusyOverlayStyle( + text="Viewing mode, Click here to grab the baton", + badge_bg=BUSY_YELLOW, + badge_fg=WHITE, + overlay_fill=qcolor(BUSY_YELLOW, 195), + overlay_border=qcolor(BUSY_YELLOW_BORDER, 235), + overlay_text=qcolor(WHITE), + accent_dot=BUSY_YELLOW_DOT, + subtext="Grab the baton if you need to interact with GUI", +) + + def build_busy_overlay_style( *, is_busy: bool, @@ -184,16 +188,7 @@ def build_busy_overlay_style( SessionsStateEnum.PendingYouToElse, SessionsStateEnum.Vacant, }: - return BusyOverlayStyle( - text="Viewing mode, Click here to grab the baton", - badge_bg=BUSY_YELLOW, - badge_fg=WHITE, - overlay_fill=qcolor(BUSY_YELLOW, 195), - overlay_border=qcolor(BUSY_YELLOW_BORDER, 235), - overlay_text=qcolor(WHITE), - accent_dot=BUSY_YELLOW_DOT, - subtext="Grab the baton if you need to interact with GUI", - ) + return VIEWING_MODE_STYLE # Auto loop centering is busy + SampleAlignment in /status — the exact # combo the callers hide (manual omega moves look identical), so it @@ -201,77 +196,36 @@ def build_busy_overlay_style( # automation progress says the Center step is running. Checked before # is_busy because the callers pass is_busy=False during alignment. if auto_centering: - return BusyOverlayStyle( - text="AUTO CENTERING", - badge_bg=BUSY_PURPLE, - badge_fg=WHITE, - overlay_fill=qcolor(BUSY_PURPLE, 190), - overlay_border=qcolor(BUSY_PURPLE_BORDER, 235), - overlay_text=qcolor(WHITE), - accent_dot=BUSY_PURPLE_DOT, - animate=True, - ) + return _animated("AUTO CENTERING", BUSY_PSI_RED, BUSY_PSI_RED_BORDER, BUSY_PSI_RED_DOT) if not is_busy: return None activity_value = str(getattr(getattr(tell_state, "activity", None), "value", "") or "").lower() - if activity_value == "mounting": - return BusyOverlayStyle( - text="ROBOT MOUNTING", - badge_bg=BUSY_RED_BADGE, - badge_fg=WHITE, - overlay_fill=qcolor(BUSY_RED_FILL, 185), - overlay_border=qcolor(BUSY_RED_BORDER, 230), - overlay_text=qcolor(WHITE), - accent_dot=BUSY_RED_DOT, - animate=True, - ) - - if activity_value == "unmounting": - return BusyOverlayStyle( - text="ROBOT UNMOUNTING", - badge_bg=BUSY_ORANGE, - badge_fg=WHITE, - overlay_fill=qcolor(BUSY_ORANGE, 190), - overlay_border=qcolor(BUSY_ORANGE_BORDER, 230), - overlay_text=qcolor(WHITE), - accent_dot=BUSY_ORANGE_DOT, - animate=True, - ) - - if activity_value == "drying": - return BusyOverlayStyle( - text="ROBOT DRYING", - badge_bg=BUSY_YELLOW, - badge_fg=BUSY_YELLOW_TEXT_DARK, - overlay_fill=qcolor(BUSY_YELLOW, 195), - overlay_border=qcolor(BUSY_YELLOW_BORDER, 235), - overlay_text=qcolor(BUSY_YELLOW_TEXT_DARK), - accent_dot=BUSY_YELLOW_DOT, - animate=True, - ) - + # Cooling is the one robot phase that keeps its own (blue) color: it is + # long, harmless and the operator should read it as "wait" not "danger". + # Every other busy state shares the PSI red so the badge reads as one + # "hands off" signal instead of a per-activity rainbow. if activity_value == "cooling": - return BusyOverlayStyle( - text="ROBOT COOLING", - badge_bg=BUSY_BLUE, - badge_fg=WHITE, - overlay_fill=qcolor(BUSY_BLUE, 190), - overlay_border=qcolor(BUSY_BLUE_BORDER, 235), - overlay_text=qcolor(WHITE), - accent_dot=BUSY_BLUE_DOT, - animate=True, - ) + return _animated("ROBOT COOLING", BUSY_BLUE, BUSY_BLUE_BORDER, BUSY_BLUE_DOT) + text = { + "mounting": "ROBOT MOUNTING", + "unmounting": "ROBOT UNMOUNTING", + "drying": "ROBOT DRYING", + }.get(activity_value, "BEAMLINE BUSY") + return _animated(text, BUSY_PSI_RED, BUSY_PSI_RED_BORDER, BUSY_PSI_RED_DOT) + + +def _animated(text: str, base: str, border: str, dot: str) -> BusyOverlayStyle: return BusyOverlayStyle( - text="BEAMLINE BUSY", - badge_bg=BUSY_PSI_RED, + text=text, + badge_bg=base, badge_fg=WHITE, - overlay_fill=qcolor(BUSY_PSI_RED, 195), - overlay_border=qcolor(BUSY_PSI_RED_BORDER, 235), + overlay_fill=qcolor(base, 190), + overlay_border=qcolor(border, 235), overlay_text=qcolor(WHITE), - accent_dot=BUSY_PSI_RED_DOT, + accent_dot=dot, animate=True, ) diff --git a/src/aare/gui/widgets/camera_image.py b/src/aare/gui/widgets/camera_image.py index d4093c74..03454282 100644 --- a/src/aare/gui/widgets/camera_image.py +++ b/src/aare/gui/widgets/camera_image.py @@ -50,9 +50,6 @@ from aare.gui.styles import ( LEGEND_BG, LEGEND_TEXT, MARK_BADGE_BG, - MARK_TOOLTIP_GOLD, - MARK_TOOLTIP_ORANGE, - MARK_TOOLTIP_RED, MARKER_GREEN, PATH_END, PATH_START, @@ -65,6 +62,7 @@ from aare.gui.styles import ( qcolor, ) from aare.gui.widgets.busy_overlay import ( + VIEWING_MODE_STYLE, BusyOverlayStyle, build_busy_overlay_style, draw_busy_badge, @@ -374,27 +372,31 @@ class SampleCameraImageLabel(QGraphicsView): painter.restore() return - # SESSION VACANT / GUEST MODE badges double as the click target for the - # grab/request menu, same as the _draw_session_overlay badge they hide. - session_badge = self._session_state in ( - SessionsStateEnum.Vacant, - SessionsStateEnum.OwnedByElse, - SessionsStateEnum.PendingYouToElse, + # The session badge doubles as the click target for the grab/request + # menu, same as the _draw_session_overlay badge it hides. + self._session_badge_rect = draw_busy_badge( + painter, + self.viewport().width(), + self.viewport().height(), + style, + fill=self._badge_fill(style), ) - fill = QColor(style.overlay_fill) - if session_badge and self._session_badge_hovered: - # Hover: darker in the light themes, brighter in Sunset. - fill = fill.lighter(125) if self._dark_theme else fill.darker(115) - - bg_rect = draw_busy_badge( - painter, self.viewport().width(), self.viewport().height(), style, fill=fill - ) - self._session_badge_rect = bg_rect if session_badge else None - painter.restore() + def _badge_fill(self, style: BusyOverlayStyle) -> QColor: + fill = QColor(style.overlay_fill) + if self._session_badge_hovered: + # Hover: darker in the light themes, brighter in Sunset. + fill = fill.lighter(125) if self._dark_theme else fill.darker(115) + return fill + def _draw_session_overlay(self, painter: QPainter): + """Session badge for when the busy overlay has nothing to draw, i.e. + the session state is still unknown (no status yet, or one without a + session). Same yellow "viewing mode" badge as the known not-owned + states, so the click target looks identical whichever path paints it + (it used to be a gold/orange/red "Guest Mode" box of its own).""" if self._busy_overlay_style is not None: # Busy overlay drew (and owns) the session badge rect — don't clobber. return @@ -408,49 +410,13 @@ class SampleCameraImageLabel(QGraphicsView): painter.save() painter.resetTransform() - - font = QFont() - font.setPointSize(24) - font.setBold(True) - painter.setFont(font) - - if self._session_state == SessionsStateEnum.Vacant: - bg_color = qcolor(MARK_TOOLTIP_GOLD, 180) - text = "Session Vacant" - elif self._session_state == SessionsStateEnum.PendingYouToElse: - bg_color = qcolor(MARK_TOOLTIP_ORANGE, 150) - text = "Baton Requested..." - else: - bg_color = qcolor(MARK_TOOLTIP_RED, 150) - text = "Guest Mode" - - fm = QFontMetrics(font) - text_rect = fm.boundingRect(text) - padding = 16 - - vw = self.viewport().width() - vh = self.viewport().height() - - bg_w = text_rect.width() + 2 * padding - bg_h = text_rect.height() + 2 * padding - position_x = int((vw - bg_w) / 2) - position_y = int((vh - bg_h) / 2) - - bg_rect = QRect(position_x, position_y, bg_w, bg_h) - # Clicking the badge opens the session (grab/request) menu. - self._session_badge_rect = bg_rect - - if self._session_badge_hovered: - # Same hover polarity as the busy-overlay badge. - bg_color = bg_color.lighter(125) if self._dark_theme else bg_color.darker(115) - - painter.setPen(QPen(qcolor(WHITE, 220))) - painter.setBrush(bg_color) - painter.drawRoundedRect(bg_rect, 10, 10) - - painter.setPen(QPen(qcolor(WHITE))) - painter.drawText(QPoint(position_x + padding, position_y + padding + fm.ascent()), text) - + self._session_badge_rect = draw_busy_badge( + painter, + self.viewport().width(), + self.viewport().height(), + VIEWING_MODE_STYLE, + fill=self._badge_fill(VIEWING_MODE_STYLE), + ) painter.restore() def _draw_camera_unavailable_overlay(self, painter: QPainter): diff --git a/tests/unit/gui/test_busy_overlay.py b/tests/unit/gui/test_busy_overlay.py index 6a71154b..cf188c2d 100644 --- a/tests/unit/gui/test_busy_overlay.py +++ b/tests/unit/gui/test_busy_overlay.py @@ -50,3 +50,28 @@ def test_auto_centering_flag_overrides_alignment_gate(): auto_centering=True, ) assert viewing is not None and viewing.text.startswith("Viewing mode") + + +def test_robot_cooling_is_the_only_blue_busy_state(): + from types import SimpleNamespace + from typing import cast + + from aarecommon.models.tell import TellStateModel + + from aare.gui.styles import BUSY_BLUE, BUSY_PSI_RED + + def _style(activity: str): + # Only .activity.value is read; cast keeps basedpyright off the stand-in. + tell = cast( + TellStateModel, cast(object, SimpleNamespace(activity=SimpleNamespace(value=activity))) + ) + style = build_busy_overlay_style(is_busy=True, tell_state=tell) + assert style is not None + return style + + assert _style("cooling").text == "ROBOT COOLING" + assert _style("cooling").badge_bg == BUSY_BLUE + # Every other busy state shares the one PSI red. + for activity in ("mounting", "unmounting", "drying", "unknown"): + assert _style(activity).badge_bg == BUSY_PSI_RED + assert _style("unknown").text == "BEAMLINE BUSY" diff --git a/tests/unit/gui/test_camera_image.py b/tests/unit/gui/test_camera_image.py index ac9c866e..be91b41a 100644 --- a/tests/unit/gui/test_camera_image.py +++ b/tests/unit/gui/test_camera_image.py @@ -289,3 +289,11 @@ def test_autoscale_fits_from_the_first_frame(camera): camera._autoscale = False camera._scaling() assert camera.transform().isIdentity() + + +def test_unknown_session_paints_the_viewing_mode_badge(camera): + # No status yet: the busy overlay has nothing, the session overlay must + # still hand out the same click target as the known not-owned states. + assert camera._busy_overlay_style is None + camera.grab() + assert camera._session_badge_rect is not None diff --git a/tests/unit/gui/test_main_window.py b/tests/unit/gui/test_main_window.py index 4020f28a..e709e3ea 100644 --- a/tests/unit/gui/test_main_window.py +++ b/tests/unit/gui/test_main_window.py @@ -606,3 +606,39 @@ def test_nonstaff_beamline_gate_popups(qtbot, mock_ui_state): with patch("aare.gui.main_window.QToolTip") as tip: qtbot.mousePress(aux_bar, Qt.MouseButton.LeftButton, pos=aux_bar.tabRect(1).center()) assert tip.showText.called, "aux-puck tab click must explain the lock" + + +def test_sample_camera_frame_paints_visible_views_and_acks(qtbot, mock_ui_state, monkeypatch): + from PySide6.QtGui import QImage + + with ( + patch("requests.get"), + patch("aare.gui.main_window.DAQWorker"), + patch("aare.gui.main_window.PredictionSubscriber"), + patch("aare.gui.main_window.VideoThread"), + patch("aare.gui.main_window.JFJochDBusClient"), + patch("aare.gui.main_window.jwt.decode") as mock_jwt, + ): + mock_jwt.return_value = { + "sub": "testuser", + "staff": True, + "pgroups": ["p123"], + "session": 15, + } + win = _make_window(qtbot) + + views = (win.sample_camera, win.compact_sample_camera, win.portrait_sample_camera) + # Nothing is shown in tests, so pretend every view is on screen to reach + # the paint branch. + monkeypatch.setattr(type(win.sample_camera), "isVisible", lambda self: True) + win.prediction_thread = MagicMock() + + win._on_sample_camera_frame(QImage(4, 6, QImage.Format.Format_RGB888)) + + assert all(v.pixmap_item.pixmap().width() == 4 for v in views) + # The ack is what paces the subscriber; it must fire after every frame. + win.prediction_thread.notify_frame_displayed.assert_called_once() + + # No subscriber (GUI started without a sample feed): the slot must not blow up. + win.prediction_thread = None + win._on_sample_camera_frame(QImage(4, 6, QImage.Format.Format_RGB888)) diff --git a/tests/unit/gui/test_prediction_subscriber.py b/tests/unit/gui/test_prediction_subscriber.py index c57552fb..092dfed4 100644 --- a/tests/unit/gui/test_prediction_subscriber.py +++ b/tests/unit/gui/test_prediction_subscriber.py @@ -25,3 +25,96 @@ def test_json_dict_is_parsed(): 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