diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py index e275cdf74..17261c6f4 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 66a79d8dd..247bfb7ad 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,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: @@ -487,11 +497,23 @@ def get_active_task(self) -> WorkerTask: return self.active_task - @start_as_current_span(TRACER, "request") + 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(): @@ -509,12 +531,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 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: 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..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, @@ -151,6 +151,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 @@ -434,15 +444,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, @@ -453,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, @@ -473,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, ) @@ -502,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, @@ -521,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 ) @@ -544,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, @@ -563,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 ) @@ -603,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() @@ -716,8 +741,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): @@ -725,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" @@ -997,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)] @@ -1025,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)]