Create a CLI for connecting to a beamline (#519)

* Create a CLI for connecting to a beamline

* Move ALL_BEAMLINES environment variable out into common module

* Remove post-init connection from CLI

* Add CLI tests

* Fix formatting

* Auto generate list of beamlines

* Allow for multiple beamlines per dodal module

* Note the CLI in the README

* Delete old system tests

* Add connection tests to PR template checklist

* Sync beamline module and BEAMLINE environment variable in CLI

Co-authored-by: Dominic Oram <dominic.oram@diamond.ac.uk>

* Explain the need for a RunEngine in the CLI

* Print starting connection message

* Fix CLI tests

---------

Co-authored-by: Dominic Oram <dominic.oram@diamond.ac.uk>

README.rst

@@ -13,6 +13,26 @@ Documentation  https://DiamondLightSource.github.io/dodal
 Releases       https://github.com/DiamondLightSource/dodal/releases
 ============== ==============================================================
 
+Testing Connectivity
+--------------------
+
+You can test your connection to a beamline if it's PVs are visible to your machine with:
+
+.. code:: shell
+
+    # On any workstation:
+    dodal connect <BEAMLINE>
+
+    # On a beamline workstation, this should suffice:
+    dodal connect ${BEAMLINE}
+
+
+For more options, including a list of valid beamlines, type
+
+.. code:: shell
+
+    dodal connect --help
+
 
 .. |code_ci| image:: https://github.com/DiamondLightSource/dodal/actions/workflows/code.yml/badge.svg?branch=main
     :target: https://github.com/DiamondLightSource/dodal/actions/workflows/code.yml

pull_request_template.md

@@ -7,4 +7,5 @@ Fixes #ISSUE
 ### Checks for reviewer
 - [ ] Would the PR title make sense to a scientist on a set of release notes
 - [ ] If a new device has been added does it follow the [standards](https://github.com/DiamondLightSource/dodal/wiki/Device-Standards)
-- [ ] If changing the API for a pre-existing device, ensure that any beamlines using this device have updated their Bluesky plans accordingly
+- [ ] If changing the API for a pre-existing device, ensure that any beamlines using this device have updated their Bluesky plans accordingly
+- [ ] Have the connection tests for the relevant beamline(s) been run via `dodal connect ${BEAMLINE}`

pyproject.toml

@@ -12,6 +12,7 @@ classifiers = [
 ]
 description = "Ophyd devices and other utils that could be used across DLS beamlines"
 dependencies = [
+    "click",
     "ophyd",
     "ophyd-async>=0.3a5",
     "bluesky",
@@ -58,6 +59,9 @@ dev = [
     "types-aiofiles",
 ]
 
+[project.scripts]
+dodal = "dodal.__main__:main"
+
 [project.urls]
 GitHub = "https://github.com/DiamondLightSource/dodal"
 

src/dodal/__main__.py

@@ -1,16 +1,8 @@
-from argparse import ArgumentParser
-
-from . import __version__
+from .cli import main
 
 __all__ = ["main"]
 
 
-def main(args=None):
-    parser = ArgumentParser()
-    parser.add_argument("-v", "--version", action="version", version=__version__)
-    args = parser.parse_args(args)
-
-
 # test with: python -m dodal
 if __name__ == "__main__":
     main()

src/dodal/beamlines/__init__.py

@@ -0,0 +1,85 @@
+import importlib.util
+from functools import lru_cache
+from pathlib import Path
+from typing import Iterable, Mapping
+
+# Where beamline names (per the ${BEAMLINE} environment variable don't always
+# match up, we have to map between them bidirectionally). The most common use case is
+# beamlines with a "-"" in the name such as "i04-1", which is not valid in a Python
+# module name. Add any new beamlines whose name differs from their module name to this
+# dictionary, which maps ${BEAMLINE} to dodal.beamlines.<MODULE NAME>
+_BEAMLINE_NAME_OVERRIDES = {
+    "i04-1": "i04_1",
+    "i20-1": "i20_1",
+    "s03": "i03",
+}
+
+
+def all_beamline_modules() -> Iterable[str]:
+    """
+    Get the names of all importable modules in beamlines
+
+    Returns:
+        Iterable[str]: Generator of beamline module names
+    """
+
+    # This is done by inspecting file names rather than modules to avoid
+    # premature importing
+    spec = importlib.util.find_spec(__name__)
+    if spec is not None:
+        search_paths = [Path(path) for path in spec.submodule_search_locations]
+        for path in search_paths:
+            for subpath in path.glob("**/*"):
+                if (
+                    subpath.name.endswith(".py")
+                    and subpath.name != "__init__.py"
+                    and ("__pycache__" not in str(subpath))
+                ):
+                    yield subpath.with_suffix("").name
+    else:
+        raise KeyError(f"Unable to find {__name__} module")
+
+
+def all_beamline_names() -> Iterable[str]:
+    """
+    Get the names of all beamlines as per the ${BEAMLINE} environment variable
+
+    Returns:
+        Iterable[str]: Generator of beamline names that dodal supports
+    """
+    inverse_mapping = _module_name_overrides()
+    for module_name in all_beamline_modules():
+        yield from inverse_mapping.get(module_name, set()).union({module_name})
+
+
+@lru_cache
+def _module_name_overrides() -> Mapping[str, set[str]]:
+    """
+    Get the inverse of _BEAMLINE_NAME_OVERRIDES so that modules can be mapped back to
+    beamlines. _BEAMLINE_NAME_OVERRIDES is expected to be a constant so the return
+    value is cached.
+
+    Returns:
+        Mapping[str, set[str]]: A dictionary mapping the name of a dodal module to the
+        set of beamlines it supports.
+    """
+
+    inverse_mapping: dict[str, set[str]] = {}
+    for beamline, module in _BEAMLINE_NAME_OVERRIDES.items():
+        inverse_mapping[module] = inverse_mapping.get(module, set()).union({beamline})
+    return inverse_mapping
+
+
+def module_name_for_beamline(beamline: str) -> str:
+    """
+    Get the module name for a particular beamline, it may differ from the beamline
+    name e.g. i04-1 -> i04_1
+
+    Args:
+        beamline: The beamline name as per the ${BEAMLINE} environment variable
+
+    Returns:
+        str: The importable module name
+    """
+
+    return _BEAMLINE_NAME_OVERRIDES.get(beamline, beamline)

src/dodal/cli.py

@@ -0,0 +1,62 @@
+import os
+
+import click
+from bluesky.run_engine import RunEngine
+
+from dodal.beamlines import all_beamline_names, module_name_for_beamline
+from dodal.utils import make_all_devices
+
+from . import __version__
+
+
+@click.group(invoke_without_command=True)
+@click.version_option(version=__version__, message="%(version)s")
+@click.pass_context
+def main(ctx: click.Context) -> None:
+    if ctx.invoked_subcommand is None:
+        print("Please invoke subcommand!")
+
+
+@main.command(name="connect")
+@click.argument(
+    "beamline",
+    type=click.Choice(list(all_beamline_names())),
+    required=True,
+)
+@click.option(
+    "-a",
+    "--all",
+    is_flag=True,
+    help="Attempt to connect to devices marked as skipped",
+    default=False,
+)
+@click.option(
+    "-s",
+    "--sim-backend",
+    is_flag=True,
+    help="Connect to a sim backend, this initializes all device objects but does not "
+    "attempt any I/O. Useful as a a dry-run.",
+    default=False,
+)
+def connect(beamline: str, all: bool, sim_backend: bool) -> None:
+    """Initialises a beamline module, connects to all devices, reports
+    any connection issues."""
+
+    os.environ["BEAMLINE"] = beamline
+
+    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()
+
+    print(f"Attempting connection to {beamline} (using {full_module_path})")
+    devices = 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}): ")
+    print("\n".join([f"\t{key}" for key in devices.keys()]))

src/dodal/common/beamlines/beamline_utils.py

@@ -44,7 +44,7 @@ def active_device_is_same_type(
     return inspect.isclass(device) and isinstance(active_device, device)
 
 
-def _wait_for_connection(
+def wait_for_connection(
     device: AnyDevice,
     timeout: float = DEFAULT_CONNECTION_TIMEOUT,
     mock: bool = False,
@@ -109,7 +109,7 @@ def device_instantiation(
         )
         ACTIVE_DEVICES[name] = device_instance
         if wait:
-            _wait_for_connection(device_instance, mock=fake)
+            wait_for_connection(device_instance, mock=fake)
 
     else:
         if not active_device_is_same_type(already_existing_device, device_factory):

tests/beamlines/unit_tests/test_mapping.py

@@ -0,0 +1,35 @@
+import pytest
+
+from dodal.beamlines import (
+    all_beamline_names,
+    module_name_for_beamline,
+)
+
+
+@pytest.mark.parametrize(
+    "beamline,expected_module",
+    {
+        "i03": "i03",
+        "s03": "i03",
+        "i04-1": "i04_1",
+        "i22": "i22",
+    }.items(),
+)
+def test_beamline_name_mapping(beamline: str, expected_module: str):
+    assert module_name_for_beamline(beamline) == expected_module
+
+
+def test_all_beamline_names_includes_non_overridden_modules():
+    beamlines = set(all_beamline_names())
+    assert "i22" in beamlines
+
+
+def test_all_beamline_names_includes_overriden_modules():
+    beamlines = set(all_beamline_names())
+    assert "i04-1" in beamlines
+
+
+def test_all_beamline_names_includes_overriden_default_modules():
+    beamlines = set(all_beamline_names())
+    assert "i03" in beamlines
+    assert "s03" in beamlines

tests/common/beamlines/test_beamline_utils.py

@@ -97,7 +97,7 @@ def test_wait_for_v1_device_connection_passes_through_timeout(kwargs, expected_t
     device = OphydV1Device(name="")
     device.wait_for_connection = MagicMock()
 
-    beamline_utils._wait_for_connection(device, **kwargs)
+    beamline_utils.wait_for_connection(device, **kwargs)
 
     device.wait_for_connection.assert_called_once_with(timeout=expected_timeout)
 
@@ -113,7 +113,7 @@ def test_wait_for_v2_device_connection_passes_through_timeout(
     device = OphydV2Device()
     device.connect = MagicMock()
 
-    beamline_utils._wait_for_connection(device, **kwargs)
+    beamline_utils.wait_for_connection(device, **kwargs)
 
     device.connect.assert_called_once_with(
         mock=ANY,

tests/common/beamlines/test_device_instantiation.py

@@ -2,18 +2,19 @@
 
 import pytest
 
+from dodal.beamlines import all_beamline_modules
 from dodal.common.beamlines import beamline_utils
 from dodal.utils import BLUESKY_PROTOCOLS, make_all_devices
 
-ALL_BEAMLINES = {"i03", "i04", "i04_1", "i20_1", "i22", "i23", "i24", "p38", "p45"}
-
 
 def follows_bluesky_protocols(obj: Any) -> bool:
     return any((isinstance(obj, protocol) for protocol in BLUESKY_PROTOCOLS))
 
 
 @pytest.mark.parametrize(
-    "module_and_devices_for_beamline", ALL_BEAMLINES, indirect=True
+    "module_and_devices_for_beamline",
+    set(all_beamline_modules()),
+    indirect=True,
 )
 def test_device_creation(RE, module_and_devices_for_beamline):
     """
@@ -31,7 +32,9 @@ def test_device_creation(RE, module_and_devices_for_beamline):
 
 
 @pytest.mark.parametrize(
-    "module_and_devices_for_beamline", ALL_BEAMLINES, indirect=True
+    "module_and_devices_for_beamline",
+    set(all_beamline_modules()),
+    indirect=True,
 )
 def test_devices_are_identical(RE, module_and_devices_for_beamline):
     """