Merge run_task handling

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"):

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
 

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:

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",

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,16 +741,16 @@ 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):
     with pytest.raises(MissingInstrumentSessionError):
         _ = 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)]