From 501a810c11f4a24377d3c778e00bd26cc42fd40d Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Mon, 4 May 2026 08:53:47 +0200 Subject: [PATCH] tellupdater.py: adding listening to more TELL events --- src/aare/daq/tellupdater.py | 110 ++++++++++++++++++++++++++++- tests/unit/daq/test_tellupdater.py | 40 +++++++++++ 2 files changed, 147 insertions(+), 3 deletions(-) diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index 7e329ddf..d853de86 100644 --- a/src/aare/daq/tellupdater.py +++ b/src/aare/daq/tellupdater.py @@ -6,6 +6,7 @@ import websocket import sseclient import time from aareDB.models import PuckWithTellPosition +from pydantic import BaseModel from aare.devices.tell_client import TellClient @@ -23,6 +24,107 @@ aare_db = None current_pucks = [] last_pucks_sent = [] +DEWAR_CONTENT_UPDATE_EVENT = "DewarContentUpdate" + +TRACKED_TELL_EVENT_VALUES = { + "Gripper detection": { + "Pin still in Gripper", + "No Pin in Gripper", + "Pin is lost", + }, +} + +TRACKED_MOTION_EVENT_VALUES = { + "Motion Task": { + "dry" + }, +} +TRACKED_MOTION_SYNC_EVENTS = { + "Motion Sync": { + "Robot Clear after mount", + "Gripper cooling down", + "Sample put on Gonio", + "Sample get from Gonio", + "Sample put on Puck", + "Sample get on Puck", + }, +} + +latest_tell_events = {} +TELL_JOURNAL_PREFIX = "[TELL][JOURNAL]" + + +class TellSsePayload(BaseModel): + value: str | None = None + eventValue: str | None = None + data: str | None = None + message: str | None = None + state: str | None = None + + @property + def normalized_value(self) -> str | None: + for candidate in ( + self.value, + self.eventValue, + self.data, + self.message, + self.state, + ): + if candidate is None: + continue + stripped = candidate.strip() + if stripped: + return stripped + return None + + +def _normalize_event_data(data): + if data is None: + return None + + if not isinstance(data, str): + data = str(data) + + stripped = data.strip() + if not stripped: + return None + + try: + payload = TellSsePayload.model_validate_json(stripped) + except Exception: + return stripped + + return payload.normalized_value + + +def extract_tracked_tell_event(event_name, event_data): + normalized_data = _normalize_event_data(event_data) + + if event_name in TRACKED_TELL_EVENT_VALUES: + allowed_values = TRACKED_TELL_EVENT_VALUES[event_name] + if normalized_data in allowed_values: + return event_name, normalized_data + return None + + if event_name in TRACKED_MOTION_EVENT_VALUES: + allowed_values = TRACKED_MOTION_EVENT_VALUES[event_name] + if normalized_data in allowed_values: + return event_name, normalized_data + return None + + if event_name in TRACKED_MOTION_SYNC_EVENTS: + allowed_values = TRACKED_MOTION_SYNC_EVENTS[event_name] + if normalized_data in allowed_values: + return event_name, normalized_data + return None + + return None + + +def record_tell_event(event_name, event_value): + latest_tell_events[event_name] = event_value + logger.info(f"{TELL_JOURNAL_PREFIX} event={event_name} value={event_value}") + def listen_to_sse(): print("[SSE][listen_to_sse] Entered listen_to_sse()") if not tell_client.url: @@ -56,8 +158,7 @@ def listen_to_sse(): events_iter = client.events() if hasattr(client, "events") else iter(client) for event in events_iter: # print(f"event = {event.event} with data: {event.data}") - if event.event == "DewarContentUpdate": - on_sse_event(event) + on_sse_event(event) except Exception as exc: print(f"[SSE][listen_to_sse][ERROR] Connection lost or failed: {exc}") @@ -98,7 +199,10 @@ def handle_tell_change_event(): def on_sse_event(event): print(f"[EVENT] SSE event: {event.event}") print(f"[EVENT] Data: {event.data}") - if event.event == "DewarContentUpdate": + tracked_event = extract_tracked_tell_event(event.event, event.data) + if tracked_event is not None: + record_tell_event(*tracked_event) + if event.event == DEWAR_CONTENT_UPDATE_EVENT: print("Dewar content changed, updating...") handle_tell_change_event() diff --git a/tests/unit/daq/test_tellupdater.py b/tests/unit/daq/test_tellupdater.py index c9a000b8..30bca29d 100644 --- a/tests/unit/daq/test_tellupdater.py +++ b/tests/unit/daq/test_tellupdater.py @@ -55,6 +55,46 @@ def test_on_sse_event(): tellupdater.on_sse_event(event) mock_handle.assert_called_once() +def test_extract_tracked_tell_event_with_json_payload(): + tracked = tellupdater.extract_tracked_tell_event( + "Gripper detection", + json.dumps({"value": "Pin still in Gripper"}), + ) + assert tracked == ("Gripper detection", "Pin still in Gripper") + +def test_extract_tracked_tell_event_for_standalone_message(): + tracked = tellupdater.extract_tracked_tell_event( + "Motion Sync", + json.dumps({"message": "Sample put on Gonio"}), + ) + assert tracked == ("Motion Sync", "Sample put on Gonio") + +def test_on_sse_event_tracks_interesting_event(): + class MockEvent: + def __init__(self, event, data): + self.event = event + self.data = data + + tellupdater.latest_tell_events.clear() + event = MockEvent("Motion Task", json.dumps({"value": "dry"})) + + with patch("aare.daq.tellupdater.handle_tell_change_event") as mock_handle: + tellupdater.on_sse_event(event) + mock_handle.assert_not_called() + + assert tellupdater.latest_tell_events["Motion Task"] == "dry" + +def test_record_tell_event_logs_to_journal(): + tellupdater.latest_tell_events.clear() + + with patch("aare.daq.tellupdater.logger.info") as mock_info: + tellupdater.record_tell_event("Motion Sync", "Sample put on Gonio") + + assert tellupdater.latest_tell_events["Motion Sync"] == "Sample put on Gonio" + mock_info.assert_called_once_with( + "[TELL][JOURNAL] event=Motion Sync value=Sample put on Gonio" + ) + @patch("aare.daq.tellupdater.websocket.WebSocketApp") def test_on_message(mock_ws_app): message = json.dumps([{