From e396b7e6a0ac68e057276a3edae5dcf51e8485c4 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Mon, 4 May 2026 11:01:44 +0200 Subject: [PATCH] tellupdater.py: writing tell events in redis and updating test cases --- src/aare/daq/tellupdater.py | 50 ++++++++++++++++++++++++++++-- tests/unit/daq/test_tellupdater.py | 40 ++++++++++++++++++++++++ 2 files changed, 87 insertions(+), 3 deletions(-) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index de43c319..e1ab97c7 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -1,6 +1,9 @@ import os import json import threading +from collections import deque +from datetime import datetime, timezone +from typing import Any, cast import requests import websocket @@ -17,6 +20,7 @@ from aare.devices.tell_client import TellClient from aare.common.exception_handler import TellCommunicationError, TellConnectionException from aare.common.logger_config import setup_logger from aare.daq.aaredb import AareWrapper +from aare.daq.config import BeamlineConfig from aare.common.beamline import MXBeamline, mx_beamline logger = setup_logger("aareDAQ") @@ -24,6 +28,7 @@ logger = setup_logger("aareDAQ") # Initialize TELL client and DB wrapper to None so they can be patched in tests tell_client = None aare_db = None +config = None # Track current state current_pucks = [] @@ -56,7 +61,9 @@ TRACKED_MOTION_SYNC_EVENTS = { } latest_tell_events = {} +tell_event_history = deque(maxlen=25) TELL_JOURNAL_PREFIX = "[TELL][JOURNAL]" +TELL_EVENTS_REDIS_KEY_SUFFIX = "tell_events" SSE_RECONNECT_DELAY_S = 5 SSE_RECOVERABLE_EXCEPTIONS = ( RequestException, @@ -94,6 +101,19 @@ class TellSsePayload(BaseModel): return None +class TellEventRecord(BaseModel): + timestamp: str + class_: str + event: str + + def redis_dict(self) -> dict[str, str]: + return { + "timestamp": self.timestamp, + "class": self.class_, + "event": self.event, + } + + def _normalize_event_data(data): if data is None: return None @@ -137,8 +157,32 @@ def extract_tracked_tell_event(event_name, event_data): return None +def set_tell_events_in_redis(events: list[dict[str, str]]) -> None: + if config is None: + logger.debug("[REDIS] BeamlineConfig unavailable; skipping TELL event write") + return + + redis_client = getattr(config, "_BeamlineConfig__client", None) + beamline_key = getattr(config, "_BeamlineConfig__bl", None) + if redis_client is None or beamline_key is None: + logger.error("[REDIS] BeamlineConfig internals unavailable; skipping TELL event write") + return + + redis_key = f"{beamline_key}:{TELL_EVENTS_REDIS_KEY_SUFFIX}" + cast(Any, redis_client).set(redis_key, json.dumps(events)) + logger.info(f"[REDIS] Written TELL events to: {redis_key}") + + def record_tell_event(event_name, event_value): + event_record = TellEventRecord( + timestamp=datetime.now(timezone.utc).isoformat(), + class_=event_name, + event=event_value, + ) + latest_tell_events[event_name] = event_value + tell_event_history.append(event_record) + set_tell_events_in_redis([item.redis_dict() for item in tell_event_history]) logger.info(f"{TELL_JOURNAL_PREFIX} event={event_name} value={event_value}") @@ -327,10 +371,10 @@ def main(): if __name__ == "__main__": # Configuration - tell_client = TellClient(bl=mx_beamline()) - aare_db = AareWrapper(bl=mx_beamline()) - beamline = mx_beamline() + tell_client = TellClient(bl=beamline) + aare_db = AareWrapper(bl=beamline) + config = BeamlineConfig(bl=beamline) SLOT_IDENTIFIER = beamline.value.upper() # WS_URL = f"wss://mx-db-01.psi.ch/dispatcher/protected_router/wstell/ws/slot/{SLOT_IDENTIFIER}" WS_URL = f"wss://mx-aaredb-dmz-01.psi.ch/dispatcher/protected_router/wstell/ws/slot/{SLOT_IDENTIFIER}" diff --git a/tests/unit/daq/test_tellupdater.py b/tests/unit/daq/test_tellupdater.py index 15a84c48..ffcd51ff 100644 --- a/tests/unit/daq/test_tellupdater.py +++ b/tests/unit/daq/test_tellupdater.py @@ -86,6 +86,7 @@ def test_on_sse_event_tracks_interesting_event(): def test_record_tell_event_logs_to_journal(): tellupdater.latest_tell_events.clear() + tellupdater.tell_event_history.clear() with patch("aare.daq.tellupdater.logger.info") as mock_info: tellupdater.record_tell_event("Motion Sync", "Sample put on Gonio") @@ -95,6 +96,45 @@ def test_record_tell_event_logs_to_journal(): "[TELL][JOURNAL] event=Motion Sync value=Sample put on Gonio" ) +def test_record_tell_event_writes_history_to_redis(): + tellupdater.latest_tell_events.clear() + tellupdater.tell_event_history.clear() + + mock_config = MagicMock() + mock_config._BeamlineConfig__bl = "x10sa" + + with patch("aare.daq.tellupdater.config", mock_config): + tellupdater.record_tell_event("Motion Task", "dry") + + mock_config._BeamlineConfig__client.set.assert_called_once() + redis_key, redis_value = mock_config._BeamlineConfig__client.set.call_args.args + assert redis_key == "x10sa:tell_events" + + payload = json.loads(redis_value) + assert len(payload) == 1 + assert payload[0]["class"] == "Motion Task" + assert payload[0]["event"] == "dry" + assert "timestamp" in payload[0] + +def test_record_tell_event_keeps_last_25_events(): + tellupdater.latest_tell_events.clear() + tellupdater.tell_event_history.clear() + + mock_config = MagicMock() + mock_config._BeamlineConfig__bl = "x10sa" + + with patch("aare.daq.tellupdater.config", mock_config): + for idx in range(30): + tellupdater.record_tell_event("Motion Sync", f"event-{idx}") + + redis_key, redis_value = mock_config._BeamlineConfig__client.set.call_args.args + assert redis_key == "x10sa:tell_events" + + payload = json.loads(redis_value) + assert len(payload) == 25 + assert payload[0]["event"] == "event-5" + assert payload[-1]["event"] == "event-29" + def test_on_message_logs_json_decode_error(): with patch("aare.daq.tellupdater.logger.exception") as mock_exception: tellupdater.on_message(None, "{invalid")