diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py index c3f9bde13..81dc83637 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: + 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_blocking(task) else: - resp = client.run_task(task) + 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..b0258141f 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,14 @@ def from_config(cls, config: ApplicationConfig) -> Self: else: return cls(rest) + @classmethod + 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 @start_as_current_span(TRACER) def plans(self) -> PlanCache: @@ -311,11 +321,24 @@ def get_active_task(self) -> WorkerTask: return self.active_task - @start_as_current_span(TRACER, "request") + @start_as_current_span(TRACER, "task") + def run_task( + self, + task: TaskRequest, + on_event: OnAnyEvent | None = None, + timeout: float | None = None, + ) -> TaskStatus: + 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( - self, request: TaskRequest, on_event: OnAnyEvent | None = None + self, task: TaskRequest, on_event: OnAnyEvent | None = None ) -> TaskStatus: - 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(): @@ -333,12 +356,13 @@ def run_blocking( 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.debug("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: 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..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, @@ -65,6 +65,16 @@ 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_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 @@ -348,15 +358,15 @@ 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", ): - 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( +def test_run_stomp_sets_up_control( client_with_events: BlueapiClient, mock_rest: Mock, mock_events: MagicMock, @@ -367,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, @@ -387,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, ) @@ -416,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, @@ -435,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 ) @@ -458,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, @@ -477,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 ) @@ -517,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() @@ -630,8 +655,8 @@ 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"): - client.run_task(TaskRequest(name="foo", instrument_session="cm12345-1")) + with asserting_span_exporter(exporter, "run_stomp"): + client.run_stomp(TaskRequest(name="foo", instrument_session="cm12345-1")) def test_instrument_session_required(client): @@ -639,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" @@ -698,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)] @@ -726,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)]