f

Author: wakonig

bec_lib/bec_lib/messaging_services.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import enum
 import mimetypes
 import os
 from abc import ABC
@@ -9,6 +10,7 @@
 from bec_lib.endpoints import MessageEndpoints
 
 if TYPE_CHECKING:
+    from bec_lib.connector import MessageObject
     from bec_lib.redis_connector import RedisConnector
 
 # Type variable for the message object class
@@ -126,6 +128,7 @@ class MessagingService(ABC, Generic[MessageObjectT]):
     def __init__(self, redis_connector: RedisConnector) -> None:
         self._redis_connector = redis_connector
         self._scopes: set[str] = set()
+        self._auto_notifications: dict[str, list[str]] = {}
         self._enabled = False
         self._default_scope: str | list[str] | None = None
         self._service_config: messages.AvailableMessagingServicesMessage | None = None
@@ -134,6 +137,12 @@ def __init__(self, redis_connector: RedisConnector) -> None:
             cb=self._on_new_scope_change_msg,
             from_start=True,
         )
+        self._redis_connector.register(
+            MessageEndpoints.notification_config(), cb=self._on_notification_config_change_msg
+        )
+        config_msg = self._redis_connector.get(MessageEndpoints.notification_config())
+        if config_msg is not None:
+            self._update_auto_notifications(config_msg)
 
     def set_default_scope(self, scope: str | list[str] | None) -> None:
         """
@@ -146,6 +155,133 @@ def set_default_scope(self, scope: str | list[str] | None) -> None:
             raise ValueError(f"Scope '{scope}' is not available for this messaging service.")
         self._default_scope = scope
 
+    def set_auto_notifications(
+        self,
+        event_type: Literal["new_scan", "scan_completed", "alarm", "scan_interlock"] | str,
+        enabled: bool,
+        scopes: list[str] | str | None = None,
+    ) -> None:
+        """
+        Set automatic notifications for a specific event type.
+
+        Args:
+            event_type (Literal["new_scan", "scan_completed", "alarm", "scan_interlock"] | str): The type of event to set notifications for.
+            enabled (bool): Whether to enable or disable notifications for the event.
+            scopes (list[str] | str | None): The scopes to apply the notifications to.
+        """
+        event_name = event_type.value if isinstance(event_type, enum.Enum) else event_type
+        scopes_list: list[str] = []
+        if scopes is not None:
+            if isinstance(scopes, str):
+                scopes_list = [scopes]
+            else:
+                scopes_list = scopes
+
+            for scope in scopes_list:
+                if scope not in self._scopes:
+                    raise ValueError(
+                        f"Scope '{scope}' is not available for this messaging service."
+                    )
+        else:
+            if self._default_scope is not None:
+                scopes_list = (
+                    [self._default_scope]
+                    if isinstance(self._default_scope, str)
+                    else self._default_scope
+                )
+            elif not self._SUPPORTS_EMPTY_SCOPES:
+                raise ValueError(
+                    "Scopes must be provided when there is no default scope and empty scopes are not supported."
+                )
+
+        if enabled:
+            # merge with existing scopes if already enabled for this event type
+            existing_scopes = set(self._auto_notifications.get(event_name, []))
+            existing_scopes.update(scopes_list)
+            self._auto_notifications[event_name] = list(existing_scopes)
+        else:
+            # if disabling, remove the scopes for this event type, or the entire event type if no scopes are provided
+            if scopes_list:
+                existing_scopes = set(self._auto_notifications.get(event_name, []))
+                existing_scopes.difference_update(scopes_list)
+                if existing_scopes:
+                    self._auto_notifications[event_name] = list(existing_scopes)
+                else:
+                    self._auto_notifications.pop(event_name, None)
+            else:
+                self._auto_notifications.pop(event_name, None)
+
+        self._sync_auto_notifications_config(event_name, enabled=enabled, scopes=scopes_list)
+
+    def _sync_auto_notifications_config(
+        self, event_name: str, enabled: bool, scopes: list[str]
+    ) -> None:
+        config_msg = self._redis_connector.get(MessageEndpoints.notification_config())
+        if config_msg is None:
+            config_msg = messages.NotificationConfigMessage()
+
+        routes = {name: list(targets) for name, targets in config_msg.routes.items()}
+        event_routes = list(routes.get(event_name, []))
+
+        if enabled:
+            scopes_to_add = scopes or [None]
+            for scope in scopes_to_add:
+                target = messages.NotificationServiceTarget(
+                    service_name=self._SERVICE_NAME, scope=scope
+                )
+                if not any(existing == target for existing in event_routes):
+                    event_routes.append(target)
+        else:
+            if scopes:
+                event_routes = [
+                    target
+                    for target in event_routes
+                    if not (target.service_name == self._SERVICE_NAME and target.scope in scopes)
+                ]
+            else:
+                event_routes = [
+                    target for target in event_routes if target.service_name != self._SERVICE_NAME
+                ]
+
+        if event_routes:
+            routes[event_name] = event_routes
+        else:
+            routes.pop(event_name, None)
+
+        self._redis_connector.set_and_publish(
+            MessageEndpoints.notification_config(),
+            messages.NotificationConfigMessage(routes=routes, metadata=config_msg.metadata),
+        )
+
+    def _on_notification_config_change_msg(
+        self,
+        message: (
+            MessageObject[messages.NotificationConfigMessage]
+            | dict[str, messages.NotificationConfigMessage]
+        ),
+    ) -> None:
+        config_msg = message.value if hasattr(message, "value") else message["data"]
+        self._update_auto_notifications(config_msg)
+
+    def _update_auto_notifications(self, config_msg: messages.NotificationConfigMessage) -> None:
+        auto_notifications: dict[str, list[str]] = {}
+        for event_name, targets in config_msg.routes.items():
+            scopes: list[str] = []
+            has_matching_service = False
+            for target in targets:
+                if target.service_name != self._SERVICE_NAME:
+                    continue
+                has_matching_service = True
+                if isinstance(target.scope, str):
+                    scopes.append(target.scope)
+                elif isinstance(target.scope, list):
+                    scopes.extend(target.scope)
+
+            if has_matching_service:
+                auto_notifications[event_name] = list(dict.fromkeys(scopes))
+
+        self._auto_notifications = auto_notifications
+
     def _on_new_scope_change_msg(
         self, message: dict[str, messages.AvailableMessagingServicesMessage]
     ) -> None:

bec_lib/tests/test_messaging_service.py

@@ -1,7 +1,10 @@
+import time
+
 import pytest
 
 from bec_lib import messages
 from bec_lib.endpoints import MessageEndpoints
+from bec_lib.messaging_hooks import MessagingEvent
 from bec_lib.messaging_services import (
     MessageServiceObject,
     SciLogMessagingService,
@@ -552,3 +555,154 @@ def test_scilog_add_text_bold_and_color(scilog_message, connected_connector):
         text_part.content
         == '<p><mark class="pen-red"><strong>Beamline checks failed</strong></mark></p>'
     )
+
+
+def test_set_auto_notifications_persists_notification_config(scilog_service, connected_connector):
+    scilog_service.set_auto_notifications(MessagingEvent.SCAN, enabled=True, scopes="default")
+
+    config_msg = connected_connector.get(MessageEndpoints.notification_config())
+    assert config_msg == messages.NotificationConfigMessage(
+        routes={
+            "new_scan": [
+                messages.NotificationServiceTarget(service_name="scilog", scope="default")
+            ]
+        }
+    )
+    assert scilog_service._auto_notifications == {"new_scan": ["default"]}  # pylint: disable=protected-access
+
+
+def test_set_auto_notifications_merges_with_existing_routes(scilog_service, connected_connector):
+    connected_connector.set_and_publish(
+        MessageEndpoints.notification_config(),
+        messages.NotificationConfigMessage(
+            routes={
+                "new_scan": [
+                    messages.NotificationServiceTarget(service_name="signal", scope="beamline-ops")
+                ]
+            }
+        ),
+    )
+
+    scilog_service.set_auto_notifications(MessagingEvent.SCAN, enabled=True, scopes="default")
+
+    config_msg = connected_connector.get(MessageEndpoints.notification_config())
+    assert config_msg == messages.NotificationConfigMessage(
+        routes={
+            "new_scan": [
+                messages.NotificationServiceTarget(service_name="signal", scope="beamline-ops"),
+                messages.NotificationServiceTarget(service_name="scilog", scope="default"),
+            ]
+        }
+    )
+    assert scilog_service._auto_notifications == {"new_scan": ["default"]}  # pylint: disable=protected-access
+
+
+def test_set_auto_notifications_disable_removes_only_matching_service_scope(
+    scilog_service, connected_connector
+):
+    connected_connector.set_and_publish(
+        MessageEndpoints.notification_config(),
+        messages.NotificationConfigMessage(
+            routes={
+                "new_scan": [
+                    messages.NotificationServiceTarget(service_name="signal", scope="beamline-ops"),
+                    messages.NotificationServiceTarget(service_name="scilog", scope="default"),
+                ]
+            }
+        ),
+    )
+
+    scilog_service.set_auto_notifications(MessagingEvent.SCAN, enabled=False, scopes="default")
+
+    config_msg = connected_connector.get(MessageEndpoints.notification_config())
+    assert config_msg == messages.NotificationConfigMessage(
+        routes={
+            "new_scan": [
+                messages.NotificationServiceTarget(service_name="signal", scope="beamline-ops")
+            ]
+        }
+    )
+    assert scilog_service._auto_notifications == {}  # pylint: disable=protected-access
+
+
+def test_set_auto_notifications_uses_default_scope_when_scopes_omitted(
+    scilog_service, connected_connector
+):
+    scilog_service.set_default_scope("default")
+
+    scilog_service.set_auto_notifications(MessagingEvent.SCAN, enabled=True)
+
+    config_msg = connected_connector.get(MessageEndpoints.notification_config())
+    assert config_msg == messages.NotificationConfigMessage(
+        routes={
+            "new_scan": [
+                messages.NotificationServiceTarget(service_name="scilog", scope="default")
+            ]
+        }
+    )
+    assert scilog_service._auto_notifications == {"new_scan": ["default"]}  # pylint: disable=protected-access
+
+
+def test_messaging_service_tracks_external_notification_config_updates(
+    scilog_service, connected_connector
+):
+    connected_connector.set_and_publish(
+        MessageEndpoints.notification_config(),
+        messages.NotificationConfigMessage(
+            routes={
+                "new_scan": [
+                    messages.NotificationServiceTarget(service_name="signal", scope="beamline-ops"),
+                    messages.NotificationServiceTarget(service_name="scilog", scope="default"),
+                ],
+                "alarm": [
+                    messages.NotificationServiceTarget(
+                        service_name="scilog", scope=["default", "secondary"]
+                    )
+                ],
+            }
+        ),
+    )
+
+    deadline = time.time() + 1
+    while (
+        time.time() < deadline
+        and scilog_service._auto_notifications  # pylint: disable=protected-access
+        != {"new_scan": ["default"], "alarm": ["default", "secondary"]}
+    ):
+        time.sleep(0.01)
+
+    assert scilog_service._auto_notifications == {  # pylint: disable=protected-access
+        "new_scan": ["default"],
+        "alarm": ["default", "secondary"],
+    }
+
+
+def test_messaging_service_loads_notification_config_on_init(connected_connector):
+    connected_connector.set_and_publish(
+        MessageEndpoints.notification_config(),
+        messages.NotificationConfigMessage(
+            routes={
+                "new_scan": [
+                    messages.NotificationServiceTarget(service_name="scilog", scope="default")
+                ]
+            }
+        ),
+    )
+
+    service = SciLogMessagingService(connected_connector)
+    available_services = messages.AvailableMessagingServicesMessage(
+        config=messages.MessagingConfig(
+            signal=messages.MessagingServiceScopeConfig(enabled=False),
+            teams=messages.MessagingServiceScopeConfig(enabled=False),
+            scilog=messages.MessagingServiceScopeConfig(enabled=True),
+        ),
+        deployment_services=[
+            messages.SciLogServiceInfo(
+                id="test_scilog", scope="default", enabled=True, logbook_id="test_logbook"
+            )
+        ],
+        session_services=[],
+    )
+    service._on_new_scope_change_msg(message={"data": available_services})
+
+    assert service._auto_notifications == {"new_scan": ["default"]}  # pylint: disable=protected-access