Hyperion 1474: Save the Panda (#702)

* (DiamondLightSource/hyperion#1474) Add cli tool to save pandas

* (DiamondLightSource/hyperion#1474) fixes to save_panda script

* Add unit test for saving the panda

* Make pyright happy

* Changes to make command more specific to pandas

* Add additional unit test coverage of argument parsing, force option

* Fix CI test failure due to polluting the environment

* Construct only the necessary devices when saving the panda

* Additional code coverage

pyproject.toml

@@ -69,6 +69,7 @@ dev = [
 
 [project.scripts]
 dodal = "dodal.__main__:main"
+save-panda = "dodal.devices.util.save_panda:main"
 
 [project.urls]
 GitHub = "https://github.com/DiamondLightSource/dodal"

src/dodal/devices/util/save_panda.py

@@ -0,0 +1,87 @@
+import argparse
+import os
+import sys
+from argparse import ArgumentParser
+from pathlib import Path
+from typing import cast
+
+from bluesky.run_engine import RunEngine
+from ophyd_async.core import Device, save_device
+from ophyd_async.panda import phase_sorter
+
+from dodal.beamlines import module_name_for_beamline
+from dodal.utils import make_device
+
+
+def main(argv: list[str]):
+    """CLI Utility to save the panda configuration."""
+    parser = ArgumentParser(description="Save an ophyd_async panda to yaml")
+    parser.add_argument(
+        "--beamline", help="beamline to save from e.g. i03. Defaults to BEAMLINE"
+    )
+    parser.add_argument(
+        "--device-name",
+        help='name of the device. The default is "panda"',
+        default="panda",
+    )
+    parser.add_argument(
+        "-f",
+        "--force",
+        action=argparse.BooleanOptionalAction,
+        help="Force overwriting an existing file",
+    )
+    parser.add_argument("output_file", help="output filename")
+
+    # this exit()s with message/help unless args parsed successfully
+    args = parser.parse_args(argv[1:])
+
+    beamline = args.beamline
+    device_name = args.device_name
+    output_file = args.output_file
+    force = args.force
+
+    if beamline:
+        os.environ["BEAMLINE"] = beamline
+    else:
+        beamline = os.environ.get("BEAMLINE", None)
+
+    if not beamline:
+        sys.stderr.write("BEAMLINE not set and --beamline not specified\n")
+        return 1
+
+    if Path(output_file).exists() and not force:
+        sys.stderr.write(
+            f"Output file {output_file} already exists and --force not specified."
+        )
+        return 1
+
+    _save_panda(beamline, device_name, output_file)
+
+    print("Done.")
+    return 0
+
+
+def _save_panda(beamline, device_name, output_file):
+    RE = RunEngine()
+    print("Creating devices...")
+    module_name = module_name_for_beamline(beamline)
+    try:
+        devices = make_device(f"dodal.beamlines.{module_name}", device_name)
+    except Exception as error:
+        sys.stderr.write(f"Couldn't create device {device_name}: {error}\n")
+        sys.exit(1)
+
+    panda = devices[device_name]
+    print(f"Saving to {output_file} from {device_name} on {beamline}...")
+    _save_panda_to_file(RE, cast(Device, panda), output_file)
+
+
+def _save_panda_to_file(RE: RunEngine, panda: Device, path: str):
+    def save_to_file():
+        yield from save_device(panda, path, sorter=phase_sorter)
+
+    RE(save_to_file())
+
+
+if __name__ == "__main__":  # pragma: no cover
+    sys.exit(main(sys.argv))

src/dodal/utils.py

@@ -102,6 +102,32 @@ def wrapper(*args, **kwds) -> T:
     return decorator
 
 
+def make_device(
+    module: str | ModuleType,
+    device_name: str,
+    **kwargs,
+) -> dict[str, AnyDevice]:
+    """Make a single named device and its dependencies from the given beamline module.
+
+    Args:
+        module (str | ModuleType): The module to make devices from.
+        device_name: Name of the device to construct
+        **kwargs: Arguments passed on to every device factory
+
+    Returns:
+        dict[str, AnyDevice]: A dict mapping device names to the constructed devices
+    """
+    if isinstance(module, str):
+        module = import_module(module)
+
+    device_collector = {}
+    factories = collect_factories(module)
+    device_collector[device_name] = _make_one_device(
+        module, device_name, device_collector, factories, **kwargs
+    )
+    return device_collector
+
+
 def make_all_devices(
     module: str | ModuleType | None = None, include_skipped: bool = False, **kwargs
 ) -> tuple[dict[str, AnyDevice], dict[str, Exception]]:
@@ -326,3 +352,30 @@ def get_run_number(directory: str, prefix: str = "") -> int:
         return 1
     else:
         return _find_next_run_number_from_files(nexus_file_names)
+
+
+def _make_one_device(
+    module: ModuleType,
+    device_name: str,
+    devices: dict[str, AnyDevice],
+    factories: dict[str, AnyDeviceFactory],
+    **kwargs,
+) -> AnyDevice:
+    factory = factories.get(device_name)
+    if not factory:
+        raise ValueError(f"Unable to find factory for {device_name}")
+
+    dependencies = list(extract_dependencies(factories, device_name))
+    for dependency_name in dependencies:
+        if dependency_name not in devices:
+            try:
+                devices[dependency_name] = _make_one_device(
+                    module, dependency_name, devices, factories, **kwargs
+                )
+            except Exception as e:
+                raise RuntimeError(
+                    f"Unable to construct device {dependency_name}"
+                ) from e
+
+    params = {name: devices[name] for name in dependencies}
+    return factory(**params, **kwargs)

tests/devices/unit_tests/util/test_save_panda.py

@@ -0,0 +1,144 @@
+import os
+from unittest.mock import MagicMock, patch
+
+import pytest
+from bluesky.simulators import RunEngineSimulator
+from ophyd_async.panda import phase_sorter
+
+from dodal.devices.util.save_panda import _save_panda, main
+
+
+def test_save_panda():
+    sim_run_engine = RunEngineSimulator()
+    panda = MagicMock()
+    with (
+        patch(
+            "dodal.devices.util.save_panda.make_device", return_value={"panda": panda}
+        ) as mock_make_device,
+        patch(
+            "dodal.devices.util.save_panda.RunEngine",
+            return_value=MagicMock(side_effect=sim_run_engine.simulate_plan),
+        ),
+        patch("dodal.devices.util.save_panda.save_device") as mock_save_device,
+    ):
+        _save_panda("i03", "panda", "test/file.yml")
+
+        mock_make_device.assert_called_with("dodal.beamlines.i03", "panda")
+        mock_save_device.assert_called_with(panda, "test/file.yml", sorter=phase_sorter)
+
+
+@patch(
+    "dodal.devices.util.save_panda.sys.exit",
+    side_effect=AssertionError("This exception expected"),
+)
+def test_save_panda_failure_to_create_device_exits_with_failure_code(mock_exit):
+    with patch(
+        "dodal.devices.util.save_panda.make_device",
+        side_effect=ValueError("device does not exist"),
+    ):
+        with pytest.raises(AssertionError):
+            _save_panda("i03", "panda", "test/file.yml")
+
+    assert mock_exit.called_once_with(1)
+
+
+@patch("dodal.devices.util.save_panda._save_panda")
+@pytest.mark.parametrize(
+    "beamline, args, expected_beamline, expected_device_name, expected_output_file, "
+    "expected_return_value",
+    [
+        ("i03", ["my_file_name.yml"], "i03", "panda", "my_file_name.yml", 0),
+        (
+            "i02",
+            ["--beamline=i04", "my_file_name.yml"],
+            "i04",
+            "panda",
+            "my_file_name.yml",
+            0,
+        ),
+        (
+            None,
+            ["--beamline=i04", "my_file_name.yml"],
+            "i04",
+            "panda",
+            "my_file_name.yml",
+            0,
+        ),
+        (
+            "i03",
+            ["--device-name=my_panda", "my_file_name.yml"],
+            "i03",
+            "my_panda",
+            "my_file_name.yml",
+            0,
+        ),
+        (
+            None,
+            ["--device-name=my_panda", "my_file_name.yml"],
+            "i03",
+            "my_panda",
+            "my_file_name.yml",
+            1,
+        ),
+    ],
+)
+def test_main(
+    mock_save_panda: MagicMock,
+    beamline: str,
+    args: list[str],
+    expected_beamline,
+    expected_device_name,
+    expected_output_file,
+    expected_return_value,
+):
+    args.insert(0, "save_panda")
+    env_patch = {}
+    if beamline:
+        env_patch["BEAMLINE"] = beamline
+
+    with patch.dict(os.environ, env_patch):
+        return_value = main(args)
+
+    assert return_value == expected_return_value
+    if not expected_return_value:
+        mock_save_panda.assert_called_with(
+            expected_beamline, expected_device_name, expected_output_file
+        )
+
+
+@pytest.mark.parametrize(
+    "file_exists, force, save_panda_called, expected_return_value",
+    [
+        (True, True, True, 0),
+        (False, False, True, 0),
+        (True, False, False, 1),
+        (True, True, True, 0),
+    ],
+)
+@patch("dodal.devices.util.save_panda._save_panda")
+@patch("dodal.devices.util.save_panda.Path", autospec=True)
+def test_file_exists_check(
+    mock_path: MagicMock,
+    mock_save_panda: MagicMock,
+    file_exists: bool,
+    force: bool,
+    save_panda_called: bool,
+    expected_return_value: int,
+):
+    exists = mock_path.return_value.exists
+    exists.return_value = file_exists
+    argv = ["save_panda", "--beamline=i03", "test_output_file.yml"]
+    if force:
+        argv.insert(1, "--force")
+
+    with patch.dict("os.environ"):
+        return_value = main(argv)
+
+    mock_path.assert_called_with("test_output_file.yml")
+    exists.assert_called_once()
+    if save_panda_called:
+        mock_save_panda.assert_called_with("i03", "panda", "test_output_file.yml")
+    else:
+        mock_save_panda.assert_not_called()
+
+    assert return_value == expected_return_value

tests/fake_beamline_broken_dependency.py

@@ -0,0 +1,24 @@
+from unittest.mock import MagicMock
+
+from bluesky.protocols import Readable
+from ophyd import EpicsMotor
+
+from dodal.devices.cryostream import Cryo
+
+
+def device_x() -> Readable:
+    return _mock_with_name("readable")
+
+
+def device_y() -> EpicsMotor:
+    raise AssertionError("Test failure")
+
+
+def device_z(device_x: Readable, device_y: EpicsMotor) -> Cryo:
+    return _mock_with_name("cryo")
+
+
+def _mock_with_name(name: str) -> MagicMock:
+    mock = MagicMock()
+    mock.name = name
+    return mock

tests/test_utils.py

@@ -13,6 +13,7 @@
     get_hostname,
     get_run_number,
     make_all_devices,
+    make_device,
 )
 
 
@@ -90,6 +91,46 @@ def test_some_devices_when_some_factories_raise_exceptions() -> None:
     )
 
 
+def test_make_device_with_dependency():
+    import tests.fake_beamline_dependencies as fake_beamline
+
+    devices = make_device(fake_beamline, "device_z")
+    assert devices.keys() == {"device_x", "device_y", "device_z"}
+
+
+def test_make_device_no_dependency():
+    import tests.fake_beamline_dependencies as fake_beamline
+
+    devices = make_device(fake_beamline, "device_x")
+    assert devices.keys() == {"device_x"}
+
+
+def test_make_device_with_exception():
+    import tests.fake_beamline_all_devices_raise_exception as fake_beamline
+
+    with pytest.raises(ValueError):
+        make_device(fake_beamline, "device_c")
+
+
+def test_make_device_with_module_name():
+    devices = make_device("tests.fake_beamline", "device_a")
+    assert {"device_a"} == devices.keys()
+
+
+def test_make_device_no_factory():
+    import tests.fake_beamline_dependencies as fake_beamline
+
+    with pytest.raises(ValueError):
+        make_device(fake_beamline, "this_device_does_not_exist")
+
+
+def test_make_device_dependency_throws():
+    import tests.fake_beamline_broken_dependency as fake_beamline
+
+    with pytest.raises(RuntimeError):
+        make_device(fake_beamline, "device_z")
+
+
 def device_a() -> Readable:
     return MagicMock()