Implement dodal connect ALL to test connection to all beamlines

src/dodal/cli.py

@@ -8,6 +8,7 @@
 from dodal.utils import make_all_devices
 
 from . import __version__
+from .common.beamlines.beamline_utils import clear_devices
 
 
 @click.group(invoke_without_command=True)
@@ -21,7 +22,7 @@ def main(ctx: click.Context) -> None:
 @main.command(name="connect")
 @click.argument(
     "beamline",
-    type=click.Choice(list(all_beamline_names())),
+    type=click.Choice(list(all_beamline_names()) + ["ALL"]),
     required=True,
 )
 @click.option(
@@ -31,6 +32,15 @@ def main(ctx: click.Context) -> None:
     help="Attempt to connect to devices marked as skipped",
     default=False,
 )
+@click.option(
+    "--include-training",
+    is_flag=True,
+    help="Include beamlines from the training module",
+    default=False,
+)
+@click.option(
+    "--include-sim", is_flag=True, help="Include sim beamlines sxx", default=False
+)
 @click.option(
     "-s",
     "--sim-backend",
@@ -39,33 +49,63 @@ def main(ctx: click.Context) -> None:
     "attempt any I/O. Useful as a a dry-run.",
     default=False,
 )
-def connect(beamline: str, all: bool, sim_backend: bool) -> None:
+def connect(
+    beamline: str,
+    all: bool,
+    sim_backend: bool,
+    include_training: bool,
+    include_sim: bool,
+) -> None:
     """Initialises a beamline module, connects to all devices, reports
     any connection issues."""
 
-    os.environ["BEAMLINE"] = beamline
+    if beamline == "ALL":
+        beamlines = sorted(all_beamline_names())
+        module_names = {bl: module_name_for_beamline(bl) for bl in beamlines}
+
+        def included(bl: str):
+            return not (
+                (bl.startswith("s") and not include_sim)
+                or ("training_rig" == module_names[bl] and not include_training)
+            )
+
+        beamlines = [bl for bl in beamlines if included(bl)]
+    else:
+        beamlines = [beamline]
+        module_names = {beamline: module_name_for_beamline(beamline)}
+
+    success = True
+
+    for beamline in beamlines:
+        clear_devices()
+        os.environ["BEAMLINE"] = beamline
+
+        module_name = module_names[beamline]
+        full_module_path = f"dodal.beamlines.{module_name}"
 
-    module_name = module_name_for_beamline(beamline)
-    full_module_path = f"dodal.beamlines.{module_name}"
+        # We need to make a RunEngine to allow ophyd-async devices to connect.
+        # See https://blueskyproject.io/ophyd-async/main/explanations/event-loop-choice.html
+        RunEngine()
 
-    # We need to make a RunEngine to allow ophyd-async devices to connect.
-    # See https://blueskyproject.io/ophyd-async/main/explanations/event-loop-choice.html
-    RunEngine()
+        print("***************************************************************")
+        print(f"Attempting connection to {beamline} (using {full_module_path})")
+        print("***************************************************************")
+        devices, exceptions = make_all_devices(
+            full_module_path,
+            include_skipped=all,
+            fake_with_ophyd_sim=sim_backend,
+        )
+        sim_statement = " (sim mode)" if sim_backend else ""
 
-    print(f"Attempting connection to {beamline} (using {full_module_path})")
-    devices, exceptions = make_all_devices(
-        full_module_path,
-        include_skipped=all,
-        fake_with_ophyd_sim=sim_backend,
-    )
-    sim_statement = " (sim mode)" if sim_backend else ""
+        print(f"{len(devices)} devices connected{sim_statement}:")
+        connected_devices = "\n".join(
+            sorted([f"\t{device_name}" for device_name in devices.keys()])
+        )
+        print(connected_devices)
 
-    print(f"{len(devices)} devices connected{sim_statement}:")
-    connected_devices = "\n".join(
-        sorted([f"\t{device_name}" for device_name in devices.keys()])
-    )
-    print(connected_devices)
+        # If exceptions have occurred, this will print details of the relevant PVs
+        if len(exceptions) > 0:
+            success = False
+            print(str(NotConnected(exceptions)))
 
-    # If exceptions have occurred, this will print details of the relevant PVs
-    if len(exceptions) > 0:
-        raise NotConnected(exceptions)
+    exit(0 if success else 1)

tests/unit_tests/test_cli.py

@@ -1,9 +1,10 @@
 import os
-from unittest.mock import patch
+from collections.abc import Sequence
+from unittest.mock import MagicMock, patch
 
 import pytest
 from click.testing import CliRunner, Result
-from ophyd_async.core import Device, NotConnected
+from ophyd_async.core import Device
 
 from dodal import __version__
 from dodal.cli import main
@@ -88,12 +89,42 @@ def test_cli_connect_when_devices_error(
     runner: CliRunner,
     devices: tuple[dict[str, AnyDevice], dict[str, Exception]],
 ):
-    with pytest.raises(NotConnected):
-        _mock_connect(
-            EXAMPLE_BEAMLINE,
-            runner=runner,
-            devices=SOME_SUCCESSFUL_DEVICES,
-        )
+    result = _mock_connect(
+        EXAMPLE_BEAMLINE,
+        runner=runner,
+        devices=SOME_SUCCESSFUL_DEVICES,
+    )
+    assert result.exit_code != 0
+
+
+@pytest.mark.parametrize(
+    "args, expected_beamlines",
+    [[["ALL"], ["i03"]], [["--include-training", "ALL"], ["i03", "training_rig"]]],
+)
+@patch(
+    "dodal.cli.all_beamline_names", new=MagicMock(return_value=["i03", "training_rig"])
+)
+@patch.dict(os.environ, clear=True)
+def test_cli_all_include_training(
+    runner: CliRunner, args: Sequence[str], expected_beamlines: Sequence[str]
+):
+    result = _mock_connect(*args, runner=runner, devices=ALL_SUCCESSFUL_DEVICES)
+    for bl in expected_beamlines:
+        assert f"Attempting connection to {bl}" in result.stdout
+
+
+@pytest.mark.parametrize(
+    "args, expected_beamlines",
+    [[["ALL"], ["i03"]], [["--include-sim", "ALL"], ["i03", "s03"]]],
+)
+@patch("dodal.cli.all_beamline_names", new=MagicMock(return_value=["i03", "s03"]))
+@patch.dict(os.environ, clear=True)
+def test_cli_all_include_sim(
+    runner: CliRunner, args: Sequence[str], expected_beamlines: Sequence[str]
+):
+    result = _mock_connect(*args, runner=runner, devices=ALL_SUCCESSFUL_DEVICES)
+    for bl in expected_beamlines:
+        assert f"Attempting connection to {bl}" in result.stdout
 
 
 def _mock_connect(