Source code for csaxs_bec.devices.sim.simulated_beamline_devices

"""This module contains simulated versions of beamline devices relevant for running simulated experiments."""

from __future__ import annotations

from bec_lib.logger import bec_logger
from ophyd import Device
from ophyd_devices import StatusBase
from ophyd_devices.tests.utils import patched_device

from csaxs_bec.devices.epics.delay_generator_csaxs.ddg_1 import DDG1
from csaxs_bec.devices.epics.fast_shutter import cSAXSFastEpicsShutter

logger = bec_logger.logger




class cSAXSSimulatedFastShutter(Device):
    """
    Simulated version of the cSAXS Fast Shutter.

    To set up initial signal to specific values, please use the example below:
    fsh.shutter._read_pv.mock_data = 1 #
    -> with 'shutter' being a signal of the shutter device

    """

    def __new__(cls, *args, **kwargs):
        with patched_device(cSAXSFastEpicsShutter, *args, **kwargs) as fsh:
            # fsh.shutter._read_pv.mock_data = 1 # Set mock data if needed
            return fsh





class DDG1WithPatchedStage(DDG1):
    """
    Simulated version of the DDG1 with patched stage method.
    We inherit here to be able to overwrite the on_stage method if needed, without compromising the rest of the device's functionality.
    """



    def on_stage(self):
        logger.warning(
            f"Staging of {self.name}. This is a simulated device of {self.__class__.__name__}."
        )
        super().on_stage()




    def on_trigger(self):
        """
        We overwrite and patch the on_trigger method. This resolves immediately now and
        sets the status to finished, instead of waiting for the trigger logic to complete.
        """
        # Change logic if needed
        status = StatusBase(obj=self)
        status.set_finished()
        return status






class SimulatedDDG1(Device):

    def __new__(cls, *args, **kwargs):
        with patched_device(DDG1WithPatchedStage, *args, **kwargs) as ddg1:
            return ddg1