diff --git a/ophyd_devices/interfaces/base_classes/psi_device_base.py b/ophyd_devices/interfaces/base_classes/psi_device_base.py index 9d0bdf7..3a6fa07 100644 --- a/ophyd_devices/interfaces/base_classes/psi_device_base.py +++ b/ophyd_devices/interfaces/base_classes/psi_device_base.py @@ -41,9 +41,9 @@ class PSIDeviceBase(Device): def __init__( self, + prefix: str = "", *, name: str, - prefix: str = "", scan_info: ScanInfo | None = None, device_manager: DeviceManagerBase | None = None, **kwargs, @@ -58,7 +58,17 @@ class PSIDeviceBase(Device): # Make sure device_manager is not passed to super().__init__ if not specified # This is to avoid issues with ophyd.OphydObject.__init__ when the parent is ophyd.Device # and the device_manager is passed to it. This will cause a TypeError. + parent = kwargs.get("parent") + ancestor = parent + while ancestor is not None and (device_manager is None or scan_info is None): + if device_manager is None: + device_manager = getattr(ancestor, "device_manager", None) + if scan_info is None: + scan_info = getattr(ancestor, "scan_info", None) + ancestor = getattr(ancestor, "parent", None) self.device_manager = device_manager + if scan_info is not None: + self.scan_info = scan_info sig = inspect.signature(super().__init__) if "device_manager" in sig.parameters: super().__init__(device_manager=device_manager, prefix=prefix, name=name, **kwargs) diff --git a/tests/test_psi_device_base.py b/tests/test_psi_device_base.py index ab1c388..3fc7635 100644 --- a/tests/test_psi_device_base.py +++ b/tests/test_psi_device_base.py @@ -5,12 +5,16 @@ import time from unittest import mock import pytest -from ophyd import Device +from bec_server.device_server.devices.devicemanager import DeviceManagerDS +from bec_server.device_server.tests.utils import DMMock +from ophyd import Component as Cpt +from ophyd import Device, Staged from ophyd.status import StatusBase from ophyd_devices.interfaces.base_classes.psi_device_base import DeviceStoppedError, PSIDeviceBase from ophyd_devices.sim.sim_camera import SimCamera from ophyd_devices.sim.sim_positioner import SimPositioner +from ophyd_devices.tests.utils import get_mock_scan_info # pylint: disable=redefined-outer-name # pylint: disable=protected-access @@ -24,6 +28,24 @@ class SimDevice(PSIDeviceBase, Device): """Simulated Device with PSI Device Base""" +class ParentDevice(PSIDeviceBase): + """PSI device with a PSI subdevice.""" + + child = Cpt(PSIDeviceBase, "child:") + + +class NonPSIParentDevice(Device): + """Plain ophyd device with a PSI subdevice.""" + + child = Cpt(PSIDeviceBase, "child:") + + +class NestedParentDevice(PSIDeviceBase): + """PSI device with a plain ophyd child that contains a PSI subdevice.""" + + container = Cpt(NonPSIParentDevice, "container:") + + @pytest.fixture def device_positioner(): """Fixture for Device""" @@ -82,6 +104,113 @@ def test_psi_device_base_init_with_device_manager(): assert device_2.device_manager is dm +def test_psi_device_base_can_be_created_as_component(): + """Test PSIDeviceBase compatibility with ophyd Component.create_component.""" + parent = ParentDevice("root:", name="parent") + + assert parent.child.name == "parent_child" + assert parent.child.prefix == "root:child:" + assert parent.child.parent is parent + + +def test_psi_subdevice_inherits_bec_context(): + """Test PSI subdevices use the parent's BEC context by default.""" + dm = mock.MagicMock() + scan_info = mock.MagicMock() + parent = ParentDevice("root:", name="parent", device_manager=dm, scan_info=scan_info) + + assert parent.child.device_manager is dm + assert parent.child.scan_info is scan_info + + +def test_psi_subdevice_under_plain_ophyd_parent_uses_mock_context(): + """Test PSI subdevices outside BEC context keep the top-level mock fallback.""" + parent = NonPSIParentDevice("root:", name="parent") + + assert parent.child.device_manager is None + assert parent.child.scan_info is not None + assert parent.child.scan_info is not getattr(parent, "scan_info", None) + + +def test_psi_subdevice_under_plain_ophyd_parent_can_inherit_explicit_context(): + """Test plain parents can host PSI subdevices when they expose BEC context.""" + dm = mock.MagicMock() + scan_info = mock.MagicMock() + + class ContextParentDevice(Device): + child = Cpt(PSIDeviceBase, "child:") + + def __init__(self, *args, **kwargs): + self.device_manager = dm + self.scan_info = scan_info + super().__init__(*args, **kwargs) + + parent = ContextParentDevice("root:", name="parent") + + assert parent.child.device_manager is dm + assert parent.child.scan_info is scan_info + + +def test_psi_subdevice_walks_parent_chain_for_bec_context(): + """Test nested PSI subdevices inherit BEC context from higher ancestors.""" + dm = mock.MagicMock() + scan_info = mock.MagicMock() + parent = NestedParentDevice("root:", name="parent", device_manager=dm, scan_info=scan_info) + + assert parent.container.child.device_manager is dm + assert parent.container.child.scan_info is scan_info + + +def test_psi_subdevice_context_with_bec_device_manager_construction(): + """Test BEC device-manager construction passes context to nested PSI subdevices.""" + dm = DMMock() + dm.scan_info = get_mock_scan_info(device=None) + config = { + "name": "parent", + "deviceClass": "tests.test_psi_device_base.NestedParentDevice", + "deviceConfig": {"prefix": "root:"}, + } + + with mock.patch.object(DeviceManagerDS, "_get_device_class", return_value=NestedParentDevice): + parent, leftover_config = DeviceManagerDS.construct_device_obj(config, dm) + + assert leftover_config == {} + assert parent.device_manager is dm + assert parent.scan_info is dm.scan_info + assert parent.container.child.device_manager is dm + assert parent.container.child.scan_info is dm.scan_info + + +def test_psi_subdevice_follows_parent_stage_and_unstage(): + """Test ophyd side-effects when PSIDeviceBase is used as a subdevice.""" + parent = ParentDevice("root:", name="parent") + + assert parent.staged == Staged.no + assert parent.child.staged == Staged.no + + staged = parent.stage() + assert staged == [parent, parent.child] + assert parent.staged == Staged.yes + assert parent.child.staged == Staged.yes + + unstaged = parent.unstage() + assert unstaged == [parent.child, parent] + assert parent.staged == Staged.no + assert parent.child.staged == Staged.no + + +def test_psi_subdevice_stop_is_propagated_when_connected(): + """Test parent stop propagates to connected PSI subdevices.""" + parent = ParentDevice("root:", name="parent") + + with mock.patch.object(PSIDeviceBase, "connected", new_callable=mock.PropertyMock) as connected: + connected.return_value = True + parent.stop() + + assert parent.stopped is True + assert parent.child.stopped is True + + def test_on_stage_hook(device): """Test user method hooks""" with mock.patch.object(device, "on_stage") as mock_on_stage: