tellupdater.py: better handling of exceptions and logging for TELL events.

src/aare/daq/tellupdater.py

@@ -2,16 +2,21 @@ import os
 import json
 import threading
 
+import requests
 import websocket
 import sseclient
 import time
+import ssl
 from aareDB.models import PuckWithTellPosition
-from pydantic import BaseModel
+from aareDB.exceptions import ApiException
+from pydantic import BaseModel, ValidationError
+from requests import RequestException
 
 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  # Make sure the import path fits your project
+from aare.daq.aaredb import AareWrapper
 from aare.common.beamline import MXBeamline, mx_beamline
 
 logger = setup_logger("aareDAQ")
@@ -52,6 +57,17 @@ TRACKED_MOTION_SYNC_EVENTS = {
 
 latest_tell_events = {}
 TELL_JOURNAL_PREFIX = "[TELL][JOURNAL]"
+SSE_RECONNECT_DELAY_S = 5
+SSE_RECOVERABLE_EXCEPTIONS = (
+    RequestException,
+    OSError,
+    ssl.SSLError,
+)
+WS_RECOVERABLE_EXCEPTIONS = (
+    OSError,
+    ssl.SSLError,
+    websocket.WebSocketException,
+)
 
 
 class TellSsePayload(BaseModel):
@@ -91,7 +107,7 @@ def _normalize_event_data(data):
 
     try:
         payload = TellSsePayload.model_validate_json(stripped)
-    except Exception:
+    except (ValidationError, json.JSONDecodeError, ValueError, TypeError):
         return stripped
 
     return payload.normalized_value
@@ -125,22 +141,28 @@ 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 get_dmz_client_cert_paths(bl: MXBeamline) -> tuple[str, str]:
+    beamline = bl.value.lower()
+    cert_dir = "/etc/ssl/certs/secrets"
+    return (
+        f"{cert_dir}/mx-{beamline}-queue-01_from_dmz-01.crt",
+        f"{cert_dir}/mx-{beamline}-queue-01_from_dmz-01.key",
+    )
+
 def listen_to_sse():
-    print("[SSE][listen_to_sse] Entered listen_to_sse()")
+    logger.info("[SSE] listener started")
     if not tell_client.url:
-        print(f"[SSE][WARN] No TELL URL configured – SSE listener not started. (tell_client.url={tell_client.url})")
+        logger.warning("[SSE] No TELL URL configured; listener not started")
         return
     sse_url = tell_client.url + "/events"
+    cert_pair = get_dmz_client_cert_paths(mx_beamline())
 
     while True:
         try:
-            print(f"[SSE][INFO] Attempting to connect to {sse_url}...")
+            logger.info(f"[SSE] Connecting to {sse_url}")
             if sse_url.startswith("https://mx-aaredb-dmz-01"): #"https://mx-db-01"
                 # mTLS path
-                import requests
-                #cert_pair = ("/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.crt", "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_db-01.key")
-                cert_pair = ("/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.crt",
-                             "/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.key")
                 #ca_root = "/etc/ssl/certs/secrets/mx-db-01_Full_Chain_CA.pem"
                 ca_root = "/etc/ssl/certs/secrets/mx-aaredb-dmz-01_Full_Chain_CA.pem"
                 response = requests.get(sse_url, stream=True, cert=cert_pair, verify=ca_root)
@@ -151,7 +173,7 @@ def listen_to_sse():
                 # SSEClient will handle the simple HTTP GET itself
                 client = sseclient.SSEClient(sse_url)
 
-            print("[SSE][listen_to_sse] Initial detected pucks fetch on connect")
+            logger.info("[SSE] Connected; refreshing detected pucks")
             handle_tell_change_event()
 
             # Compatibility: some SSEClient versions are iterable, others expose .events().
@@ -159,11 +181,11 @@ def listen_to_sse():
             for event in events_iter:
                 # print(f"event = {event.event} with data: {event.data}")
                 on_sse_event(event)
-        except Exception as exc:
-            print(f"[SSE][listen_to_sse][ERROR] Connection lost or failed: {exc}")
-        
-        print("[SSE][INFO] Reconnecting to SSE in 5 seconds...")
-        time.sleep(5)
+        except SSE_RECOVERABLE_EXCEPTIONS:
+            logger.exception("[SSE] Recoverable failure while consuming TELL event stream")
+
+        logger.info(f"[SSE] Reconnecting in {SSE_RECONNECT_DELAY_S} seconds")
+        time.sleep(SSE_RECONNECT_DELAY_S)
 
 def compare_and_report_change(old, new, key_func):
     """
@@ -192,27 +214,28 @@ def handle_tell_change_event():
     try:
         detected_pucks = tell_client.get_detected_pucks()
         aare_db.set_pucks_beamline(detected_pucks)
-        print("[SSE][INFO] Updated pucks in database after TELL change.")
-    except Exception as exc:
-        print(f"[SSE][ERROR] {exc}")
+        logger.info("[SSE] Updated pucks in database after TELL change")
+    except (TellCommunicationError, TellConnectionException, RequestException):
+        logger.exception("[SSE] Failed to fetch detected pucks from TELL")
+    except (ApiException, json.JSONDecodeError, KeyError, TypeError, ValueError):
+        logger.exception("[SSE] Failed to update puck state after TELL change")
 
 def on_sse_event(event):
-    print(f"[EVENT] SSE event: {event.event}")
-    print(f"[EVENT] Data: {event.data}")
+    logger.debug(f"[EVENT] SSE event={event.event} data={event.data!r}")
     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...")
+        logger.info("[SSE] Dewar content changed; updating")
         handle_tell_change_event()
 
 def on_message(ws, message):
     #print("[WS][MESSAGE]", message)
     try:
         data = json.loads(message)
-        print("[WS][MESSAGE]", data)
+        logger.debug(f"[WS] Received message: {data!r}")
         if not isinstance(data, list):
-            print("[WS][ERROR] Unknown message format")
+            logger.error(f"[WS] Unexpected message type: {type(data).__name__}")
             return
 
         new_pucks = [PuckWithTellPosition(**item) for item in data]
@@ -226,9 +249,8 @@ def on_message(ws, message):
             handle_tell_change_event()
             # Small wait gives SSE/DB handler a moment to settle (optional)
             time.sleep(0.2)
-        except Exception as exc:
-
-            print(f"[WS][WARN] handle_tell_change_event() failed: {exc}")
+        except (TellCommunicationError, TellConnectionException, RequestException, ApiException):
+            logger.exception("[WS] Failed to refresh TELL state before websocket update")
 
         # Compare lists and ignore reference puck at X1 if needed
         joined, left = compare_and_report_change(current_pucks, new_pucks, lambda p: getattr(p, "puck_name", None))
@@ -237,36 +259,38 @@ def on_message(ws, message):
         current_pucks = new_pucks
 
         if joined or left:
-            print(f"[WS][INFO] Changes detected joined={joined} left={left}")
+            logger.info(f"[WS] Changes detected joined={joined} left={left}")
             # Send only when there is a real change
             ws_update_samples_info(new_pucks)
-            print("[WS][INFO] Sent current pucks state to TELL.")
+            logger.info("[WS] Sent current pucks state to TELL")
         else:
-            print("[WS][INFO] No change, skipping send to TELL.")
-    except Exception as exc:
-        print("[WS][ERROR] Failed to parse websocket message:", exc)
+            logger.info("[WS] No change; skipping send to TELL")
+    except json.JSONDecodeError:
+        logger.exception("[WS] Failed to decode websocket message as JSON")
+    except (ValidationError, TypeError, ValueError, KeyError):
+        logger.exception("[WS] Failed to process websocket message")
 
 def on_error(ws, error):
-    print("[WS][ERROR]", error)
+    logger.error(f"[WS] Error: {error}")
 
 def on_close(ws, close_status_code, close_msg):
-    print("[WS][CLOSE]", close_status_code, close_msg)
+    logger.info(f"[WS] Closed status={close_status_code} message={close_msg}")
 
 def on_open(ws):
-    print("[WS][OPEN] WebSocket opened.")
+    logger.info("[WS] WebSocket opened")
 
 
 def main():
     # Start SSE listener in a separate background thread
     sse_thread = threading.Thread(target=listen_to_sse, daemon=True)
     sse_thread.start()
+    cert_file, key_file = get_dmz_client_cert_paths(mx_beamline())
 
     """
         Main function to initiate WebSocket connection with mTLS.
         """
     while True:
         try:
-            import ssl
             # 1. Create a modern SSL Context for a TLS Client
             context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
 
@@ -274,10 +298,9 @@ def main():
             context.load_verify_locations(cafile="/etc/ssl/certs/secrets/mx-aaredb-dmz-01_Full_Chain_CA.pem")
 
             # 3. Load the Client Certificate and Key (mTLS)
-            # Using the 'dmz-01' paths that worked in your curl
             context.load_cert_chain(
-                certfile="/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.crt",
-                keyfile="/etc/ssl/certs/secrets/mx-x10sa-queue-01_from_dmz-01.key"
+                certfile=cert_file,
+                keyfile=key_file
             )
 
             # Optional: Ensure hostname matching is active (recommended)
@@ -293,14 +316,14 @@ def main():
             )
 
             # 4. Pass the context directly via sslopt
-            print(f"[WS][INFO] Connecting to {WS_URL} using mTLS...")
+            logger.info(f"[WS] Connecting to {WS_URL} using mTLS")
             ws.run_forever(sslopt={"context": context})
 
-        except Exception as e:
-            print(f"[MAIN][ERROR] WebSocket connection failed: {e}")
+        except WS_RECOVERABLE_EXCEPTIONS:
+            logger.exception("[MAIN] Recoverable WebSocket connection failure")
 
-        print("[WS][INFO] Reconnecting in 5 seconds...")
-        time.sleep(5)
+        logger.info(f"[WS] Reconnecting in {SSE_RECONNECT_DELAY_S} seconds")
+        time.sleep(SSE_RECONNECT_DELAY_S)
 
 if __name__ == "__main__":
     # Configuration
@@ -313,6 +336,6 @@ if __name__ == "__main__":
     WS_URL = f"wss://mx-aaredb-dmz-01.psi.ch/dispatcher/protected_router/wstell/ws/slot/{SLOT_IDENTIFIER}"
     # WS_URL = f"wss://localhost:8001/protected_router/wstell/ws/slot/{SLOT_IDENTIFIER}"
     WS_HEADERS = [f"X-Shared-Password: {os.getenv('AAREDB_SHARED_PASSWORD')}"]
-    print(WS_HEADERS)
+    logger.debug(f"[WS] Headers configured: {WS_HEADERS}")
 
     main()

tests/unit/daq/test_tellupdater.py

@@ -95,6 +95,26 @@ def test_record_tell_event_logs_to_journal():
         "[TELL][JOURNAL] event=Motion Sync value=Sample put on Gonio"
     )
 
+def test_on_message_logs_json_decode_error():
+    with patch("aare.daq.tellupdater.logger.exception") as mock_exception:
+        tellupdater.on_message(None, "{invalid")
+
+    mock_exception.assert_called_once_with(
+        "[WS] Failed to decode websocket message as JSON"
+    )
+
+@patch("aare.daq.tellupdater.tell_client")
+@patch("aare.daq.tellupdater.aare_db")
+def test_handle_tell_change_event_logs_failure(mock_db, mock_tell):
+    mock_tell.get_detected_pucks.side_effect = ValueError("boom")
+
+    with patch("aare.daq.tellupdater.logger.exception") as mock_exception:
+        tellupdater.handle_tell_change_event()
+
+    mock_exception.assert_called_once_with(
+        "[SSE] Failed to update puck state after TELL change"
+    )
+
 @patch("aare.daq.tellupdater.websocket.WebSocketApp")
 def test_on_message(mock_ws_app):
     message = json.dumps([{