diff --git a/src/aare/daq/tellupdater.py b/src/aare/daq/tellupdater.py index d853de86..de43c319 100644 --- a/src/aare/daq/tellupdater.py +++ b/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() diff --git a/tests/unit/daq/test_tellupdater.py b/tests/unit/daq/test_tellupdater.py index 30bca29d..15a84c48 100644 --- a/tests/unit/daq/test_tellupdater.py +++ b/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([{