From e8d1a86cd850d6a70ce9a5ca821124e900893036 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Thu, 27 Aug 2026 17:30:32 +0100 Subject: [PATCH 1/6] 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 e275cdf74..74225f895 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 66a79d8dd..592b1a944 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -14,11 +14,13 @@ 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 @@ -283,6 +285,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: @@ -487,10 +494,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) @@ -508,13 +527,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 6e03040c4..a96c39ba1 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 418b7cb5fc11bf977a83d7cc09128b234c1772a8 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Fri, 28 Aug 2026 10:20:33 +0100 Subject: [PATCH 2/6] 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 74225f895..17261c6f4 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 592b1a944..9f2f2d934 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -502,11 +502,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 91e2ba21f62174fde5f89ef6ca9fff8fb828e441 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Fri, 28 Aug 2026 10:49:36 +0100 Subject: [PATCH 3/6] 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 9f2f2d934..f2d1771c3 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -494,7 +494,6 @@ def get_active_task(self) -> WorkerTask: return self.active_task - @start_as_current_span(TRACER, "task", "timeout") def run_task( self, task: TaskRequest, @@ -506,11 +505,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(): @@ -527,13 +527,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 aa540d391..580187fb3 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 d5b0493ad..93b225cc0 100644 --- a/tests/unit_tests/client/test_client.py +++ b/tests/unit_tests/client/test_client.py @@ -439,7 +439,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( @@ -717,7 +717,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 9ab72e92653c5150082bea9944680f4758a1e5d1 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Fri, 28 Aug 2026 14:18:03 +0100 Subject: [PATCH 4/6] 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 f2d1771c3..9ad575a58 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -500,10 +500,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 95a918f310066f9f990eced65875c13501768aa8 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Fri, 28 Aug 2026 14:18:47 +0100 Subject: [PATCH 5/6] 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 9ad575a58..247bfb7ad 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -286,8 +286,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 93b225cc0..7ee62fe4b 100644 --- a/tests/unit_tests/client/test_client.py +++ b/tests/unit_tests/client/test_client.py @@ -151,6 +151,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 5b426af53cf598970eae0d809d550a279e9235e8 Mon Sep 17 00:00:00 2001 From: Peter Holloway Date: Fri, 28 Aug 2026 15:22:05 +0100 Subject: [PATCH 6/6] 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 7ee62fe4b..6593fab9c 100644 --- a/tests/unit_tests/client/test_client.py +++ b/tests/unit_tests/client/test_client.py @@ -28,7 +28,7 @@ NotFoundError, ServiceUnavailableError, ) -from blueapi.config import MissingStompConfigurationError +from blueapi.config import MissingStompConfigurationError, RestConfig from blueapi.core import DataEvent from blueapi.service.model import ( DeviceModel, @@ -156,6 +156,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 @@ -439,7 +444,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", @@ -447,7 +452,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, @@ -458,14 +463,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, @@ -478,7 +485,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, ) @@ -507,7 +514,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, @@ -526,7 +533,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 ) @@ -549,7 +556,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, @@ -568,7 +575,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 ) @@ -608,6 +615,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() @@ -721,7 +741,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")) @@ -730,7 +750,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" @@ -1002,7 +1022,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)] @@ -1030,7 +1052,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)]