From 060c812408c56513c9286af645773b1fd6bbb2c7 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Thu, 27 Aug 2026 17:30:32 +0100 Subject: [PATCH 1/7] Merge run_task handling --- src/blueapi/cli/cli.py | 24 +++++++++++++++--------- src/blueapi/client/client.py | 27 +++++++++++++++++++++++---- src/blueapi/client/rest.py | 2 +- 3 files changed, 39 insertions(+), 14 deletions(-) diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py index c3f9bde13..22d71a01a 100644 --- a/src/blueapi/cli/cli.py +++ b/src/blueapi/cli/cli.py @@ -322,14 +322,17 @@ def on_event( @click.argument("name", type=str) @click.argument("parameters", type=ParametersType(), default={}, required=False) @click.option( - "--ws", + "--ws/--stomp", type=bool, is_flag=True, - default=False, + default=None, # the classic three-state boolean help=textwrap.dedent(""" - Run the plan in the foreground using the (experimental) websocket connection - to monitor progress. Allows plans to be run without a message bus and associated - configuration. + Method used to monitor the progress of plans when run in the foreground. + --stomp requires stomp configuration, --ws uses the (currently experimental) + websocket connection. + + If neither is specified, stomp is used if configuration is present and + websockets are used if not. Has no effect if --bg is also passed as the plan will not be monitored. """), @@ -361,7 +364,7 @@ def run_plan( name: str, timeout: float | None, foreground: bool, - ws: bool, + ws: bool | None, instrument_session: str, parameters: TaskParameters, ) -> None: @@ -390,10 +393,13 @@ def on_event(event: AnyEvent) -> None: client.add_callback(on_event) - if ws: - resp = client.run_blocking(task) - else: + if ws is None: + # no preference was given so use whichever we have config for resp = client.run_task(task) + elif ws: + resp = client.run_ws(task) + else: + resp = client.run_stomp(task) match resp.result: case TaskResult(result=None, type="NoneType"): diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index e393da659..9b1ca5988 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -11,11 +11,13 @@ from bluesky_stomp.messaging import MessageContext, StompClient from bluesky_stomp.models import Broker from observability_utils.tracing import get_tracer, start_as_current_span +from pydantic import HttpUrl from blueapi.config import ( ApplicationConfig, ConfigLoader, MissingStompConfigurationError, + RestConfig, ) from blueapi.core.bluesky_types import DataEvent from blueapi.service.authentication import SessionCacheManager, SessionManager @@ -107,6 +109,11 @@ def from_config(cls, config: ApplicationConfig) -> Self: else: return cls(rest) + @classmethod + def for_host(cls, host: str) -> Self: + conf = ApplicationConfig(api=RestConfig(url=HttpUrl(host))) + return cls.from_config(conf) + @cached_property @start_as_current_span(TRACER) def plans(self) -> PlanCache: @@ -311,10 +318,22 @@ def get_active_task(self) -> WorkerTask: return self.active_task - @start_as_current_span(TRACER, "request") - def run_blocking( + @start_as_current_span(TRACER, "task", "timeout") + def run_task( + self, + task: TaskRequest, + on_event: OnAnyEvent | None = None, + timeout: float | None = None, + ) -> TaskStatus: + if self._events is None: + return self.run_ws(task, on_event) + else: + return self.run_stomp(task, on_event) + + def run_ws( self, request: TaskRequest, on_event: OnAnyEvent | None = None ) -> TaskStatus: + log.info("Running plan via websocket") for event in self._rest.run_blocking(request): if on_event is not None: on_event(event) @@ -332,13 +351,13 @@ def run_blocking( return event.task_status raise BlueskyRemoteControlError("Connection closed before plan completed.") - @start_as_current_span(TRACER, "task", "timeout") - def run_task( + def run_stomp( self, task: TaskRequest, on_event: OnAnyEvent | None = None, timeout: float | None = None, ) -> TaskStatus: + log.info("Running plan via stomp") """ Synchronously run a task, requires a message bus connection diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py index da6b4b6dd..166620ac2 100644 --- a/src/blueapi/client/rest.py +++ b/src/blueapi/client/rest.py @@ -370,7 +370,7 @@ def run_blocking( additional_headers=headers, user_agent_header=USER_AGENT, ) as ws: - ws.send(Submit(task=req).model_dump_json()) + ws.send(Submit(task=req).model_dump_json(fallback=_task_model_fallback)) for message in ws: event = ControlResponse.validate_json(message) match event: From 5cdf9f787e1dfe4f8fdde492a0a50227f7628c8b Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Fri, 28 Aug 2026 10:20:33 +0100 Subject: [PATCH 2/7] Lint et al --- src/blueapi/cli/cli.py | 2 +- src/blueapi/client/client.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py index 22d71a01a..81dc83637 100644 --- a/src/blueapi/cli/cli.py +++ b/src/blueapi/cli/cli.py @@ -397,7 +397,7 @@ def on_event(event: AnyEvent) -> None: # no preference was given so use whichever we have config for resp = client.run_task(task) elif ws: - resp = client.run_ws(task) + resp = client.run_blocking(task) else: resp = client.run_stomp(task) diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index 9b1ca5988..8239289a8 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -326,11 +326,11 @@ def run_task( timeout: float | None = None, ) -> TaskStatus: if self._events is None: - return self.run_ws(task, on_event) + return self.run_blocking(task, on_event) else: return self.run_stomp(task, on_event) - def run_ws( + def run_blocking( self, request: TaskRequest, on_event: OnAnyEvent | None = None ) -> TaskStatus: log.info("Running plan via websocket") From 54e2401cee39cfdf2eea6c5b081c16038194f063 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Fri, 28 Aug 2026 10:49:36 +0100 Subject: [PATCH 3/7] Patch up tests --- src/blueapi/client/client.py | 11 ++++++----- tests/unit_tests/cli/test_cli.py | 2 ++ tests/unit_tests/client/test_client.py | 4 ++-- 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index 8239289a8..053cd9406 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -318,7 +318,6 @@ def get_active_task(self) -> WorkerTask: return self.active_task - @start_as_current_span(TRACER, "task", "timeout") def run_task( self, task: TaskRequest, @@ -330,11 +329,12 @@ def run_task( else: return self.run_stomp(task, on_event) + @start_as_current_span(TRACER, "task") def run_blocking( - self, request: TaskRequest, on_event: OnAnyEvent | None = None + self, task: TaskRequest, on_event: OnAnyEvent | None = None ) -> TaskStatus: - log.info("Running plan via websocket") - for event in self._rest.run_blocking(request): + log.debug("Running plan via websocket") + for event in self._rest.run_blocking(task): if on_event is not None: on_event(event) for cb in self._callbacks.values(): @@ -351,13 +351,14 @@ def run_blocking( return event.task_status raise BlueskyRemoteControlError("Connection closed before plan completed.") + @start_as_current_span(TRACER, "task", "timeout") def run_stomp( self, task: TaskRequest, on_event: OnAnyEvent | None = None, timeout: float | None = None, ) -> TaskStatus: - log.info("Running plan via stomp") + log.debug("Running plan via stomp") """ Synchronously run a task, requires a message bus connection diff --git a/tests/unit_tests/cli/test_cli.py b/tests/unit_tests/cli/test_cli.py index 28cb4fbc9..1840216ff 100644 --- a/tests/unit_tests/cli/test_cli.py +++ b/tests/unit_tests/cli/test_cli.py @@ -247,6 +247,7 @@ def test_submit_plan_without_stomp(runner: CliRunner): config_path, "controller", "run", + "--stomp", "-i", "cm12345-1", "sleep", @@ -448,6 +449,7 @@ def test_cannot_run_plans_without_stomp_config(runner: CliRunner): [ "controller", "run", + "--stomp", "-i", "cm12345-1", "sleep", diff --git a/tests/unit_tests/client/test_client.py b/tests/unit_tests/client/test_client.py index 3293c3610..e7ff70b51 100644 --- a/tests/unit_tests/client/test_client.py +++ b/tests/unit_tests/client/test_client.py @@ -353,7 +353,7 @@ def test_cannot_run_task_without_message_bus(client: BlueapiClient): MissingStompConfigurationError, match="Stomp configuration required to run plans is missing or disabled", ): - client.run_task(TaskRequest(name="foo", instrument_session="cm12345-1")) + client.run_stomp(TaskRequest(name="foo", instrument_session="cm12345-1")) def test_run_task_sets_up_control( @@ -631,7 +631,7 @@ def test_cannot_run_task_span_ok( match="Stomp configuration required to run plans is missing or disabled", ): with asserting_span_exporter(exporter, "grun_task"): - client.run_task(TaskRequest(name="foo", instrument_session="cm12345-1")) + client.run_stomp(TaskRequest(name="foo", instrument_session="cm12345-1")) def test_instrument_session_required(client): From 717687d9b2d4eefd7723985b89355aecbdaa230a Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Fri, 28 Aug 2026 14:18:03 +0100 Subject: [PATCH 4/7] Reverse event config check --- src/blueapi/client/client.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index 053cd9406..d9f4c3e42 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -324,10 +324,10 @@ def run_task( on_event: OnAnyEvent | None = None, timeout: float | None = None, ) -> TaskStatus: - if self._events is None: - return self.run_blocking(task, on_event) - else: + if self._events: return self.run_stomp(task, on_event) + else: + return self.run_blocking(task, on_event) @start_as_current_span(TRACER, "task") def run_blocking( From 81363ba133ccb16957ee767ff38dec387bc7bd2e Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Fri, 28 Aug 2026 14:18:47 +0100 Subject: [PATCH 5/7] Expand for_host method and add test --- src/blueapi/client/client.py | 7 +++++-- tests/unit_tests/client/test_client.py | 5 +++++ 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index d9f4c3e42..a8626736f 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -110,8 +110,11 @@ def from_config(cls, config: ApplicationConfig) -> Self: return cls(rest) @classmethod - def for_host(cls, host: str) -> Self: - conf = ApplicationConfig(api=RestConfig(url=HttpUrl(host))) + def for_host(cls, host: str | None = None) -> Self: + if host: + conf = ApplicationConfig(api=RestConfig(url=HttpUrl(host))) + else: + conf = ApplicationConfig() return cls.from_config(conf) @cached_property diff --git a/tests/unit_tests/client/test_client.py b/tests/unit_tests/client/test_client.py index e7ff70b51..26e98dcd9 100644 --- a/tests/unit_tests/client/test_client.py +++ b/tests/unit_tests/client/test_client.py @@ -65,6 +65,11 @@ def test_client_from_config(): assert bc._rest._config.url == HttpUrl("http://example.com:8082") +def test_client_for_host(): + bc = BlueapiClient.for_host("http://custom.example.com:1234") + assert bc._rest._config.url == HttpUrl("http://custom.example.com:1234") + + def test_get_plans(client: BlueapiClient): assert PlanResponse(plans=[p.model for p in client.plans]) == PLANS From fb63cc256cc3feea4600581be24d1c2d381135e6 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Fri, 28 Aug 2026 15:22:05 +0100 Subject: [PATCH 6/7] Shuffle tests around --- tests/unit_tests/client/test_client.py | 52 +++++++++++++++++++------- 1 file changed, 38 insertions(+), 14 deletions(-) diff --git a/tests/unit_tests/client/test_client.py b/tests/unit_tests/client/test_client.py index 26e98dcd9..e7e31f930 100644 --- a/tests/unit_tests/client/test_client.py +++ b/tests/unit_tests/client/test_client.py @@ -27,7 +27,7 @@ BlueskyRequestError, ServiceUnavailableError, ) -from blueapi.config import MissingStompConfigurationError +from blueapi.config import MissingStompConfigurationError, RestConfig from blueapi.core import DataEvent from blueapi.service.model import ( DeviceModel, @@ -70,6 +70,11 @@ def test_client_for_host(): assert bc._rest._config.url == HttpUrl("http://custom.example.com:1234") +def test_client_default_for_host(): + bc = BlueapiClient.for_host() + assert bc._rest._config.url == RestConfig().url + + def test_get_plans(client: BlueapiClient): assert PlanResponse(plans=[p.model for p in client.plans]) == PLANS @@ -353,7 +358,7 @@ def test_resume( ) -def test_cannot_run_task_without_message_bus(client: BlueapiClient): +def test_cannot_run_stomp_without_message_bus(client: BlueapiClient): with pytest.raises( MissingStompConfigurationError, match="Stomp configuration required to run plans is missing or disabled", @@ -361,7 +366,7 @@ def test_cannot_run_task_without_message_bus(client: BlueapiClient): client.run_stomp(TaskRequest(name="foo", instrument_session="cm12345-1")) -def test_run_task_sets_up_control( +def test_run_stomp_sets_up_control( client_with_events: BlueapiClient, mock_rest: Mock, mock_events: MagicMock, @@ -372,14 +377,16 @@ def test_run_task_sets_up_control( ctx.correlation_id = "foo" mock_events.subscribe_to_all_events = lambda on_event: on_event(COMPLETE_EVENT, ctx) - client_with_events.run_task(TaskRequest(name="foo", instrument_session="cm12345-1")) + client_with_events.run_stomp( + TaskRequest(name="foo", instrument_session="cm12345-1") + ) mock_rest.create_task.assert_called_once_with( TaskRequest(name="foo", instrument_session="cm12345-1") ) mock_rest.update_worker_task.assert_called_once_with(WorkerTask(task_id="foo")) -def test_run_task_fails_on_failing_event( +def test_run_stomp_fails_on_failing_event( client_with_events: BlueapiClient, mock_rest: Mock, mock_events: MagicMock, @@ -392,7 +399,7 @@ def test_run_task_fails_on_failing_event( mock_events.subscribe_to_all_events = lambda on_event: on_event(FAILED_EVENT, ctx) on_event = Mock() - outcome = client_with_events.run_task( + outcome = client_with_events.run_stomp( TaskRequest(name="foo", instrument_session="cm12345-1"), on_event=on_event, ) @@ -421,7 +428,7 @@ def test_run_task_fails_on_failing_event( DataEvent(name="start", doc={}, task_id="0000-1111"), ], ) -def test_run_task_calls_event_callback( +def test_run_stomp_calls_event_callback( client_with_events: BlueapiClient, mock_rest: Mock, mock_events: MagicMock, @@ -440,7 +447,7 @@ def callback(on_event: Callable[[AnyEvent, MessageContext], None]): mock_events.subscribe_to_all_events = callback # type: ignore mock_on_event = Mock() - client_with_events.run_task( + client_with_events.run_stomp( TaskRequest(name="foo", instrument_session="cm12345-1"), on_event=mock_on_event ) @@ -463,7 +470,7 @@ def callback(on_event: Callable[[AnyEvent, MessageContext], None]): object(), ], ) -def test_run_task_ignores_non_matching_events( +def test_run_stomp_ignores_non_matching_events( client_with_events: BlueapiClient, mock_rest: Mock, mock_events: MagicMock, @@ -482,7 +489,7 @@ def callback(on_event: Callable[[AnyEvent, MessageContext], None]): mock_events.subscribe_to_all_events = callback mock_on_event = Mock() - client_with_events.run_task( + client_with_events.run_stomp( TaskRequest(name="foo", instrument_session="cm12345-1"), on_event=mock_on_event ) @@ -522,6 +529,19 @@ def test_scripting_interface_raises_exceptions(): demo_plan() +@pytest.mark.parametrize( + "events,method", [(None, "run_blocking"), (EventBusClient(Mock()), "run_stomp")] +) +def test_run_test_implementation_switching(events: EventBusClient | None, method: str): + client = Mock() + client._events = events + + task = Mock() + BlueapiClient.run_task(client, task) + + getattr(client, method).assert_called_once_with(task, None) + + def test_oidc_config_property(client, mock_rest): assert client.oidc_config == mock_rest.get_oidc_config() @@ -635,7 +655,7 @@ def test_cannot_run_task_span_ok( MissingStompConfigurationError, match="Stomp configuration required to run plans is missing or disabled", ): - with asserting_span_exporter(exporter, "grun_task"): + with asserting_span_exporter(exporter, "run_stomp"): client.run_stomp(TaskRequest(name="foo", instrument_session="cm12345-1")) @@ -644,7 +664,7 @@ def test_instrument_session_required(client): _ = client.instrument_session -def test_setting_instrument_session(client): +def test_setting_instrument_session(client: BlueapiClient): # This looks like a completely pointless test but instrument_session is a # property with some logic so it's not purely to get coverage up client.instrument_session = "cm12345-4" @@ -703,7 +723,9 @@ def subscribe(on_event: Callable[[AnyEvent, MessageContext], None]): mock_events.subscribe_to_all_events = subscribe # type: ignore - client_with_events.run_task(TaskRequest(name="foo", instrument_session="cm12345-1")) + client_with_events.run_stomp( + TaskRequest(name="foo", instrument_session="cm12345-1") + ) assert callback.mock_calls == [call(test_event), call(COMPLETE_EVENT)] @@ -731,7 +753,9 @@ def subscribe(on_event: Callable[[AnyEvent, MessageContext], None]): mock_events.subscribe_to_all_events = subscribe # type: ignore - client_with_events.run_task(TaskRequest(name="foo", instrument_session="cm12345-1")) + client_with_events.run_stomp( + 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)] From 63663d1d49e0ef714da952ae463069a7f49329f4 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Wed, 2 Sep 2026 16:17:46 +0100 Subject: [PATCH 7/7] Re-add tracing span on run_task It's probably useful to have a single search to cover both --- src/blueapi/client/client.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index a8626736f..b0258141f 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -321,6 +321,7 @@ def get_active_task(self) -> WorkerTask: return self.active_task + @start_as_current_span(TRACER, "task") def run_task( self, task: TaskRequest,