Merge pull request #23 from DiamondLightSource/testRigFixes

Tidy up util module(s) and tests

Author: DominicOram

pyproject.toml

@@ -20,7 +20,6 @@ dependencies = [
     "dataclasses-json",
     "pillow",
     "requests",
-    "pydantic",
 ]
 dynamic = ["version"]
 license.file = "LICENSE"

src/dodal/adsim.py

@@ -1,22 +1,17 @@
-import socket
-
-from pydantic import BaseSettings
+import os
 
 from dodal.devices.adsim import SimStage
 from dodal.devices.areadetector import AdSimDetector
 
+from .utils import get_hostname
 
-# Settings can be customized via environment variables
-class Settings(BaseSettings):
-    pv_prefix: str = socket.gethostname().split(".")[0]
-
-
-_settings = Settings()
+# Default prefix to hostname unless overriden with export PREFIX=<prefix>
+PREFIX: str = os.environ.get("PREFIX", get_hostname())
 
 
 def stage(name: str = "sim_motors") -> SimStage:
-    return SimStage(name=name, prefix=f"{_settings.pv_prefix}-MO-SIM-01:")
+    return SimStage(name=name, prefix=f"{PREFIX}-MO-SIM-01:")
 
 
 def det(name: str = "adsim") -> AdSimDetector:
-    return AdSimDetector(name=name, prefix=f"{_settings.pv_prefix}-AD-SIM-01:")
+    return AdSimDetector(name=name, prefix=f"{PREFIX}-AD-SIM-01:")

src/dodal/p45.py

@@ -1,30 +1,22 @@
-from pydantic import BaseSettings
-
 from dodal.devices.areadetector import AdAravisDetector
 from dodal.devices.p45 import Choppers, TomoStageWithStretchAndSkew
 
+from .utils import BeamlinePrefix
 
-# Settings can be customized via environment variables
-class Settings(BaseSettings):
-    pv_prefix: str = "BL45P"
-
-
-_settings = Settings()
+PREFIX: str = BeamlinePrefix("p45").beamline_prefix
 
 
-def sample_sample(name: str = "sample_stage") -> TomoStageWithStretchAndSkew:
-    return TomoStageWithStretchAndSkew(
-        name=name, prefix=f"{_settings.pv_prefix}-MO-STAGE-01:"
-    )
+def sample(name: str = "sample_stage") -> TomoStageWithStretchAndSkew:
+    return TomoStageWithStretchAndSkew(name=name, prefix=f"{PREFIX}-MO-STAGE-01:")
 
 
 def choppers(name: str = "chopper") -> Choppers:
-    return Choppers(name=name, prefix=f"{_settings.pv_prefix}-MO-CHOP-01:")
+    return Choppers(name=name, prefix=f"{PREFIX}-MO-CHOP-01:")
 
 
 def det(name: str = "det") -> AdAravisDetector:
-    return AdAravisDetector(name=name, prefix=f"{_settings.pv_prefix}-EA-MAP-01:")
+    return AdAravisDetector(name=name, prefix=f"{PREFIX}-EA-MAP-01:")
 
 
 def diff(name: str = "diff") -> AdAravisDetector:
-    return AdAravisDetector(name=name, prefix=f"{_settings.pv_prefix}-EA-DIFF-01:")
+    return AdAravisDetector(name=name, prefix=f"{PREFIX}-EA-DIFF-01:")

src/dodal/util.py

@@ -1,7 +1,48 @@
+import inspect
+import socket
 from collections import namedtuple
 from dataclasses import dataclass
+from importlib import import_module
+from inspect import signature
 from os import environ
-from typing import Optional
+from types import ModuleType
+from typing import Any, Callable, Dict, Iterable, Optional, Type, Union
+
+from bluesky.protocols import (
+    Checkable,
+    Configurable,
+    Flyable,
+    HasHints,
+    HasName,
+    HasParent,
+    Movable,
+    Pausable,
+    Readable,
+    Stageable,
+    Stoppable,
+    Subscribable,
+    Triggerable,
+    WritesExternalAssets,
+)
+
+#: Protocols defining interface to hardware
+BLUESKY_PROTOCOLS = [
+    Checkable,
+    Flyable,
+    HasHints,
+    HasName,
+    HasParent,
+    Movable,
+    Pausable,
+    Readable,
+    Stageable,
+    Stoppable,
+    Subscribable,
+    WritesExternalAssets,
+    Configurable,
+    Triggerable,
+]
+
 
 Point2D = namedtuple("Point2D", ["x", "y"])
 Point3D = namedtuple("Point3D", ["x", "y", "z"])
@@ -15,6 +56,10 @@ def get_beamline_name(ixx: str) -> str:
         return bl
 
 
+def get_hostname() -> str:
+    return socket.gethostname().split(".")[0]
+
+
 @dataclass
 class BeamlinePrefix:
     ixx: str
@@ -24,3 +69,29 @@ def __post_init__(self):
         self.suffix = self.ixx[0].upper() if not self.suffix else self.suffix
         self.beamline_prefix = f"BL{self.ixx[1:3]}{self.suffix}"
         self.insertion_prefix = f"SR{self.ixx[1:3]}{self.suffix}"
+
+
+def make_all_devices(module: Union[str, ModuleType, None] = None) -> Dict[str, Any]:
+    if isinstance(module, str) or module is None:
+        module = import_module(module or __name__)
+    factories = collect_factories(module)
+    return {device.name: device for device in map(lambda factory: factory(), factories)}
+
+
+def collect_factories(module: ModuleType) -> Iterable[Callable[..., Any]]:
+    for var in module.__dict__.values():
+        if callable(var) and _is_device_factory(var):
+            yield var
+
+
+def _is_device_factory(func: Callable[..., Any]) -> bool:
+    return_type = signature(func).return_annotation
+    return _is_device_type(return_type)
+
+
+def _is_device_type(obj: Type[Any]) -> bool:
+    is_class = inspect.isclass(obj)
+    follows_protocols = any(
+        map(lambda protocol: isinstance(obj, protocol), BLUESKY_PROTOCOLS)
+    )
+    return is_class and follows_protocols

src/dodal/utils.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Any, Dict, List, Tuple, cast
 
 import bluesky.plan_stubs as bps

tests/__init__.py

@@ -0,0 +1,42 @@
+from unittest.mock import MagicMock, patch
+
+from bluesky.protocols import Readable
+from ophyd import EpicsMotor
+
+from dodal.utils import collect_factories, get_hostname, make_all_devices
+
+
+def test_finds_device_factories() -> None:
+    import tests.fake_beamline as fake_beamline
+
+    factories = set(collect_factories(fake_beamline))
+
+    from tests.fake_beamline import device_a, device_b, device_c
+
+    assert {device_a, device_b, device_c} == factories
+
+
+def test_makes_devices() -> None:
+    import tests.fake_beamline as fake_beamline
+
+    devices = make_all_devices(fake_beamline)
+    assert {"readable", "motor", "cryo"} == devices.keys()
+
+
+def test_makes_devices_with_module_name() -> None:
+    devices = make_all_devices("tests.fake_beamline")
+    assert {"readable", "motor", "cryo"} == devices.keys()
+
+
+def test_get_hostname() -> None:
+    with patch("dodal.utils.socket.gethostname") as mock:
+        mock.return_value = "a.b.c"
+        assert get_hostname() == "a"
+
+
+def device_a() -> Readable:
+    return MagicMock()
+
+
+def device_b() -> EpicsMotor:
+    return MagicMock()

tests/devices/system_tests/test_aperturescatterguard_system.py

@@ -0,0 +1,28 @@
+from unittest.mock import MagicMock
+
+from bluesky.protocols import Readable
+from ophyd import EpicsMotor
+
+from dodal.devices.cryostream import Cryo
+
+
+def device_a() -> Readable:
+    return _mock_with_name("readable")
+
+
+def device_b() -> EpicsMotor:
+    return _mock_with_name("motor")
+
+
+def device_c() -> Cryo:
+    return _mock_with_name("cryo")
+
+
+def not_device() -> int:
+    return 5
+
+
+def _mock_with_name(name: str) -> MagicMock:
+    mock = MagicMock()
+    mock.name = name
+    return mock