tellupdater.py: adding listening to more TELL events

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()
 

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([{