Add callback support

Author: tpoliaw

src/blueapi/client/client.py

@@ -1,5 +1,7 @@
+import itertools
 import logging
 import time
+from collections.abc import Iterable
 from concurrent.futures import Future
 from functools import cached_property
 from itertools import chain
@@ -196,6 +198,8 @@ class BlueapiClient:
     _rest: BlueapiRestClient
     _events: EventBusClient | None
     _instrument_session: str | None = None
+    _callbacks: dict[int, OnAnyEvent]
+    _callback_id: itertools.count
 
     def __init__(
         self,
@@ -204,6 +208,8 @@ def __init__(
     ):
         self._rest = rest
         self._events = events
+        self._callbacks = {}
+        self._callback_id = itertools.count()
 
     @cached_property
     @start_as_current_span(TRACER)
@@ -258,6 +264,22 @@ def instrument_session(self, session: str):
         log.debug("Setting instrument_session to %s", session)
         self._instrument_session = session
 
+    def with_instrument_session(self, session: str) -> Self:
+        self.instrument_session = session
+        return self
+
+    def add_callback(self, callback: OnAnyEvent) -> int:
+        cb_id = next(self._callback_id)
+        self._callbacks[cb_id] = callback
+        return cb_id
+
+    def remove_callback(self, id: int):
+        self._callbacks.pop(id)
+
+    @property
+    def callbacks(self) -> Iterable[OnAnyEvent]:
+        return self._callbacks.values()
+
     @property
     @start_as_current_span(TRACER)
     def state(self) -> WorkerState:
@@ -355,6 +377,13 @@ def inner_on_event(event: AnyEvent, ctx: MessageContext) -> None:
             if relates_to_task:
                 if on_event is not None:
                     on_event(event)
+                for cb in self._callbacks.values():
+                    try:
+                        cb(event)
+                    except Exception as e:
+                        log.error(
+                            f"Callback ({cb}) failed for event: {event}", exc_info=e
+                        )
                 if isinstance(event, WorkerEvent) and (
                     (event.is_complete()) and (ctx.correlation_id == task_id)
                 ):

tests/unit_tests/client/test_client.py

@@ -764,3 +764,82 @@ def test_plan_invalid_param_mapping(args, kwargs, msg):
     with pytest.raises(TypeError, match=msg):
         plan(*args, **kwargs)
     client.run_task.assert_not_called()
+
+
+def test_adding_removing_callback(client):
+    def callback(*a, **kw):
+        pass
+
+    cb_id = client.add_callback(callback)
+    assert len(client.callbacks) == 1
+    client.remove_callback(cb_id)
+    assert len(client.callbacks) == 0
+
+
+@pytest.mark.parametrize(
+    "test_event",
+    [
+        WorkerEvent(
+            state=WorkerState.RUNNING,
+            task_status=TaskStatus(
+                task_id="foo",
+                task_complete=False,
+                task_failed=False,
+            ),
+        ),
+        ProgressEvent(task_id="foo"),
+        DataEvent(name="start", doc={}, task_id="0000-1111"),
+    ],
+)
+def test_client_callbacks(
+    client_with_events: BlueapiClient,
+    mock_rest: Mock,
+    mock_events: MagicMock,
+    test_event: AnyEvent,
+):
+    callback = Mock()
+    client_with_events.add_callback(callback)
+    mock_rest.create_task.return_value = TaskResponse(task_id="foo")
+    mock_rest.update_worker_task.return_value = TaskResponse(task_id="foo")
+
+    ctx = Mock()
+    ctx.correlation_id = "foo"
+
+    def subscribe(on_event: Callable[[AnyEvent, MessageContext], None]):
+        on_event(test_event, ctx)
+        on_event(COMPLETE_EVENT, ctx)
+
+    mock_events.subscribe_to_all_events = subscribe  # type: ignore
+
+    client_with_events.run_task(TaskRequest(name="foo", instrument_session="cm12345-1"))
+
+    assert callback.mock_calls == [call(test_event), call(COMPLETE_EVENT)]
+
+
+def test_client_callback_failures(
+    client_with_events: BlueapiClient,
+    mock_rest: Mock,
+    mock_events: MagicMock,
+):
+    failing_callback = Mock(side_effect=ValueError("Broken callback"))
+    callback = Mock()
+    client_with_events.add_callback(failing_callback)
+    client_with_events.add_callback(callback)
+    mock_rest.create_task.return_value = TaskResponse(task_id="foo")
+    mock_rest.update_worker_task.return_value = TaskResponse(task_id="foo")
+
+    ctx = Mock()
+    ctx.correlation_id = "foo"
+
+    evt = DataEvent(name="start", doc={}, task_id="foo")
+
+    def subscribe(on_event: Callable[[AnyEvent, MessageContext], None]):
+        on_event(evt, ctx)
+        on_event(COMPLETE_EVENT, ctx)
+
+    mock_events.subscribe_to_all_events = subscribe  # type: ignore
+
+    client_with_events.run_task(TaskRequest(name="foo", instrument_session="cm12345-1"))
+
+    assert failing_callback.mock_calls == [call(evt), call(COMPLETE_EVENT)]
+    assert callback.mock_calls == [call(evt), call(COMPLETE_EVENT)]