diff --git a/bec_ipython_client/bec_ipython_client/progressbar.py b/bec_ipython_client/bec_ipython_client/progressbar.py index 154854666..7e88cfe83 100644 --- a/bec_ipython_client/bec_ipython_client/progressbar.py +++ b/bec_ipython_client/bec_ipython_client/progressbar.py @@ -316,5 +316,7 @@ def set_finished(self, device): Args: device (str): device name """ + if device not in self.devices: + return device_index = self.devices.index(device) self._progress.advance(self._tasks[device_index], self.NUM_STEPS) diff --git a/bec_lib/bec_lib/bl_state_machine.py b/bec_lib/bec_lib/bl_state_machine.py new file mode 100644 index 000000000..9bd6966ba --- /dev/null +++ b/bec_lib/bec_lib/bl_state_machine.py @@ -0,0 +1,120 @@ +""" +Module for managing beamline states based on configuration files. + +Example of the YAML configuration file: +``` yaml +bl_transition_states: + bl_state_class: AggregatedState + config: + evaluation_method: any + states: + alignment: + devices: + samx: + value: 0 + abs_tol: 0.1 + signals: + velocity: + value: 5 + abs_tol: 0.1 + bpm4i: + value: 100 + abs_tol: 10 + samy: + at: in + abs_tol: 0.1 +samx_within_limits: + bl_state_class: DeviceWithinLimitsState + config: + device: samx + signal: samx + low_limit: -1 + high_limit: 1 + tolerance: 0.1 +``` + +""" + +from __future__ import annotations + +import yaml + +from bec_lib import bl_states +from bec_lib.bl_state_manager import BeamlineStateManager, _state_class_for_state_type + + +class BeamlineStateMachine: + + def __init__(self, manager: BeamlineStateManager) -> None: + self._manager = manager + self._configs: dict[str, bl_states.BeamlineStateConfig] = {} + + def load_from_config( + self, + config_path: str | None = None, + config_dict: dict | None = None, + flush: bool = True, + skip_existing: bool = False, + ) -> None: + """ + Load a state configuration from a YAML file or a dictionary. If None or both are provided, + an error will be raised. Config must be a dictionary/YAML file with top-level beamline + state names and entries containing bl_state_class/config. Any previously loaded states will be + cleared if flush is True. + + Args: + config_path (str | None): The path to the YAML configuration file. + config_dict (dict | None): A dictionary containing the configuration. + flush (bool): If True, existing states in the manager will be cleared before loading new ones. + skip_existing (bool): If True, existing states in the manager will be skipped during loading. + """ + self._check_inputs(config_path=config_path, config_dict=config_dict) + config_dict = self._load_config(config_path=config_path, config_dict=config_dict) + # Check first if the config is valid before clearing the existing states + configs = self._parse_config(config_dict=config_dict) + if flush: + self._manager.clear_all() + for config in configs: + self._manager.add(config, skip_existing=skip_existing) + + def _check_inputs(self, config_path: str | None, config_dict: dict | None) -> None: + if (config_path is None and config_dict is None) or ( + config_path is not None and config_dict is not None + ): + raise ValueError("Either config_path or config_dict must be provided, but not both.") + + def _load_config(self, config_path: str | None, config_dict: dict | None) -> dict: + if config_path: + with open(config_path, "r", encoding="utf-8") as f: + loaded_config = yaml.safe_load(f) + if loaded_config is None: + raise ValueError("Config file is empty.") + return loaded_config + if config_dict is None: + raise ValueError("config_dict must be provided when config_path is not set.") + return config_dict + + def _parse_config(self, config_dict: dict) -> list[bl_states.BeamlineStateConfig]: + parsed_configs: list[bl_states.BeamlineStateConfig] = [] + for state_name, entry in config_dict.items(): + state_config_class = self._resolve_config_class(entry["bl_state_class"]) + config = entry["config"] + if not isinstance(config, dict): + raise ValueError(f"Config for state {state_name!r} must be a dictionary.") + if "name" in config and config["name"] != state_name: + raise ValueError( + f"Config name {config['name']!r} does not match top-level state name {state_name!r}." + ) + config = {**config, "name": state_name} + parsed_configs.append(state_config_class(**config)) + return parsed_configs + + def _resolve_config_class( + self, bl_state_class: str | type[bl_states.BeamlineStateConfig] + ) -> type[bl_states.BeamlineStateConfig]: + if isinstance(bl_state_class, str): + resolved_class = _state_class_for_state_type(bl_state_class) + else: + resolved_class = bl_state_class + config = resolved_class.CONFIG_CLASS + return config diff --git a/bec_lib/bec_lib/bl_state_manager.py b/bec_lib/bec_lib/bl_state_manager.py index 2604ab19a..69abd6092 100644 --- a/bec_lib/bec_lib/bl_state_manager.py +++ b/bec_lib/bec_lib/bl_state_manager.py @@ -2,12 +2,15 @@ import inspect import time +from collections.abc import Mapping, Sequence from inspect import Parameter, Signature -from typing import TYPE_CHECKING, TypedDict +from typing import TYPE_CHECKING, Any, TypedDict from pydantic import BaseModel from rich.console import Console from rich.table import Table +from rich.text import Text +from rich.tree import Tree from bec_lib import bl_states, messages from bec_lib.endpoints import MessageEndpoints @@ -73,6 +76,47 @@ class BeamlineStateGet(TypedDict): label: str +def _add_parameter_to_tree(parent: Tree, name: str, value: Any) -> None: + """Add a configuration value to a Rich tree without interpreting it as markup.""" + if isinstance(value, Mapping): + branch = parent.add(Text(str(name), style="cyan")) + if not value: + branch.add(Text("{}", style="grey70")) + return + for child_name, child_value in value.items(): + _add_parameter_to_tree(branch, str(child_name), child_value) + return + + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + branch = parent.add(Text(str(name), style="cyan")) + if not value: + branch.add(Text("[]", style="grey70")) + return + for index, child_value in enumerate(value): + _add_parameter_to_tree(branch, f"[{index}]", child_value) + return + + line = Text() + line.append(f"{name}: ", style="cyan") + line.append(repr(value), style="grey70") + parent.add(line) + + +def _truncate_text(value: str, max_length: int) -> str: + """Truncate text to a deterministic maximum length.""" + if len(value) <= max_length: + return value + return f"{value[: max_length - 1]}…" + + +def _summarize_label(label: str, max_labels: int = 3, max_length: int = 64) -> str: + """Bound the potentially large set of matching labels shown by ``show_all``.""" + labels = label.split("|") + if len(labels) > max_labels: + label = f"{'|'.join(labels[:max_labels])}|… (+{len(labels) - max_labels})" + return _truncate_text(label.replace("\n", " "), max_length=max_length) + + class BeamlineStateClientBase: """Base class for beamline state clients.""" @@ -118,6 +162,24 @@ def get(self) -> BeamlineStateGet: msg = msg_container["data"] return {"status": msg.status, "label": msg.label} + def describe(self) -> None: + """Pretty print the complete, validated configuration for this state.""" + title = Text(self._state.name, style="bold magenta") + title.append(f" ({self._state.state_type})", style="grey70") + tree = Tree(title) + parameters = self._state.model_dump(exclude={"name"}, exclude_none=True) + if ( + isinstance(self._state, bl_states.AggregatedStateConfig) + and self._state.evaluation_method is None + ): + parameters = {"evaluation_method": None, **parameters} + if not parameters: + tree.add(Text("No parameters", style="grey70")) + else: + for name, value in parameters.items(): + _add_parameter_to_tree(tree, name, value) + Console().print(tree) + def remove(self) -> None: """ Remove the current beamline state. @@ -221,17 +283,27 @@ def _wait_for_initial_state(self, state_name: str, timeout_s: float = 5.0) -> No ##### Public API ######### ########################## - def add(self, state: bl_states.BeamlineStateConfig) -> None: + def add(self, state: bl_states.BeamlineStateConfig, skip_existing: bool = False) -> None: """ Add a new beamline state to the manager. Args: state (BeamlineStateConfig): The beamline state to add. + skip_existing (bool): If True, existing states in the manager will be skipped during loading. """ - + if skip_existing and state.name in self._states: + return self._add_state(state) self._publish_states() self._wait_for_initial_state(state.name) + def clear_all(self) -> None: + """ + Clear all beamline states from the manager. + """ + for state_name in list(self._states.keys()): + self._delete_state(state_name) + self._publish_states() + def delete(self, state_name: str) -> None: """ Delete a beamline state from the manager. @@ -257,12 +329,6 @@ def show_all(self): Pretty print all beamline states using rich. """ - def _format_parameters(state_config: bl_states.BeamlineStateConfig) -> str: - parameter_dict = state_config.model_dump(exclude={"name"}, exclude_none=True) - if not parameter_dict: - return "-" - return "\n".join(f"{key}={value}" for key, value in parameter_dict.items()) - def _status_style(status_value: str) -> str: status_styles = {"valid": "green3", "invalid": "red3", "warning": "yellow3"} return status_styles.get(status_value.lower(), "grey50") @@ -276,7 +342,8 @@ def _status_style(status_value: str) -> str: table.add_column("Label") for state in self._states.values(): - params = _format_parameters(state) + state_class = _state_class_for_state_type(state.state_type) + params = state_class.format_config_summary(state) status = ( getattr(self, state.name).get() if hasattr(self, state.name) @@ -288,8 +355,8 @@ def _status_style(status_value: str) -> str: str(state.name), str(state.state_type), str(params), - f"[{status_style}]{status_value}[/{status_style}]", - f"[{status_style}]{str(status.get('label', ''))}[/{status_style}]", + Text(status_value, style=status_style), + Text(_summarize_label(str(status.get("label", ""))), style=status_style), ) console.print(table) diff --git a/bec_lib/bec_lib/bl_states.py b/bec_lib/bec_lib/bl_states.py index 8e25070f5..242920454 100644 --- a/bec_lib/bec_lib/bl_states.py +++ b/bec_lib/bec_lib/bl_states.py @@ -1,12 +1,16 @@ +"""Module defining beamline states and their evaluation logic.""" + from __future__ import annotations import functools import keyword import traceback from abc import ABC, abstractmethod -from typing import Annotated, Callable, ClassVar, Generic, Type, TypeVar, cast +from dataclasses import dataclass +from typing import Annotated, Any, Callable, ClassVar, Generic, Literal, Type, TypeVar, cast -from pydantic import BaseModel, field_validator, model_validator +import yaml +from pydantic import BaseModel, Field, field_validator, model_validator from bec_lib import messages from bec_lib.alarm_handler import Alarms @@ -194,6 +198,97 @@ def validate_limits(self) -> DeviceWithinLimitsStateConfig: return self +class SignalConfig(BaseModel): + """ + Target value for a signal inside a named machine state. + Value specifies the target value for the signal, alternatively + at may be used to pass a user parameter of the signal as a dynamic input. + Please note, only one of value or at can be specified, not both. + """ + + value: float | int | str | bool | None = None + abs_tol: float = Field(default=0.0, ge=0.0) + at: str | None = None # Optional input for user parameter + + @model_validator(mode="after") + def validate_config(self) -> SignalConfig: + """ + Validate that either value or at is provided, but not both. + """ + if self.value is not None and self.at is not None: + raise ValueError("Cannot specify both 'value' and 'at' for a signal configuration.") + if self.value is None and self.at is None: + raise ValueError("Either 'value' or 'at' must be specified for a signal configuration.") + return self + + +class DeviceConfig(BaseModel): + """ + Configuration for a device inside a named machine state. + Value specifies the target value for the device (positioner), alternatively + at may be used to pass a user parameter of the device as a dynamic input. + Please note, only one of value or at can be specified, not both. + """ + + abs_tol: float = Field(default=0.0, ge=0.0) + value: float | int | str | bool | None = None + at: str | None = None # Optional input for user parameter + low_limit: SignalConfig | None = None + high_limit: SignalConfig | None = None + signals: dict[str, SignalConfig] | None = None + + @model_validator(mode="after") + def validate_config(self) -> DeviceConfig: + """ + Validate that either value, low_limit, high_limit, or signals are provided. + """ + if self.value is not None and self.at is not None: + raise ValueError("Cannot specify both 'value' and 'at' for a device configuration.") + if ( + self.value is None + and self.at is None + and self.low_limit is None + and self.high_limit is None + and self.signals is None + ): + raise ValueError( + "At least one of value, at, low_limit, high_limit, or signals must be provided." + ) + return self + + +class SubDeviceStateConfig(BaseModel): + """ + Configuration for a sub-state with a specific label. + This is a device/signal mapping to either a DeviceConfig or SignalConfig. + """ + + devices: dict[str, DeviceConfig | SignalConfig] + transition_metadata: dict[str, Any] | None = None + + +class AggregatedStateConfig(BeamlineStateConfig): + """ + Configuration for a state machine driven by multiple device signals. + + Keys of the states dictionary are the labels of the different states. + ``evaluation_method`` controls how matching labels determine validity: + ``any`` requires one or more, ``all`` requires every configured label, + ``exclusive`` requires exactly one, and ``None`` disables validation. + """ + + state_type: ClassVar[str] = "AggregatedState" + + evaluation_method: Literal["any", "all", "exclusive"] | None = Field( + default="any", + description=( + "How matching labels determine validity. Use null to disable validation and always " + "report valid after successful initialization." + ), + ) + states: dict[str, SubDeviceStateConfig] + + C = TypeVar("C", bound=BeamlineStateConfig) D = TypeVar("D", bound=DeviceStateConfig) @@ -222,6 +317,21 @@ def update_parameters(self, **kwargs) -> None: """Update the configuration parameters of the state.""" self.config = self.CONFIG_CLASS(**{**self.config.model_dump(), **kwargs}) + @staticmethod + def format_config_summary(config: BeamlineStateConfig) -> str: + """Return a compact, single-line configuration summary for overview displays.""" + parameters = config.model_dump(exclude={"name"}, exclude_none=True) + if not parameters: + return "-" + + summary_parts = [] + for name, value in parameters.items(): + value_text = str(value).replace("\n", " ") + if len(value_text) > 60: + value_text = f"{value_text[:59]}…" + summary_parts.append(f"{name}={value_text}") + return ", ".join(summary_parts) + @abstractmethod def evaluate(self, *args, **kwargs) -> messages.BeamlineStateMessage | None: """Evaluate the state and return its state.""" @@ -408,6 +518,555 @@ def _update_device_state(self, msg_obj: MessageObject) -> messages.BeamlineState return self.evaluate(msg) +SignalSource = Literal["readback", "configuration", "limits"] + + +@dataclass(frozen=True) +class ResolvedStateSignal: + label: str + device_name: str + signal_name: str + expected_value: float | int | str | bool | None + at: str | None + abs_tolerance: float | int + source: SignalSource + + +class AggregatedState(BeamlineState[AggregatedStateConfig]): + """Beamline state that infers the current named state from multiple device signals.""" + + CONFIG_CLASS = AggregatedStateConfig + + @staticmethod + def format_config_summary(config: BeamlineStateConfig) -> str: + """Return a bounded summary of an aggregated-state configuration.""" + if not isinstance(config, AggregatedStateConfig): + return BeamlineState.format_config_summary(config) + + device_names = { + device_name for state in config.states.values() for device_name in state.devices + } + transition_count = sum( + state.transition_metadata is not None for state in config.states.values() + ) + evaluation_method = config.evaluation_method or "null" + + def count_phrase(count: int, noun: str) -> str: + suffix = "" if count == 1 else "s" + return f"{count} {noun}{suffix}" + + summary = ( + f"{evaluation_method} · {count_phrase(len(config.states), 'label')} · " + f"{count_phrase(len(device_names), 'device')} · " + f"{count_phrase(AggregatedState._count_config_requirements(config), 'requirement')}" + ) + if transition_count: + summary += f" · {count_phrase(transition_count, 'transition')}" + return summary + + @staticmethod + def _count_config_requirements(config: AggregatedStateConfig) -> int: + """Count individual signal requirements without resolving devices.""" + requirement_count = 0 + for state in config.states.values(): + for device_config in state.devices.values(): + if isinstance(device_config, SignalConfig): + requirement_count += 1 + continue + requirement_count += int( + device_config.value is not None or device_config.at is not None + ) + requirement_count += int(device_config.low_limit is not None) + requirement_count += int(device_config.high_limit is not None) + requirement_count += len(device_config.signals or {}) + return requirement_count + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + # Mapping from signal updates to affected state labels, used for efficient evaluation when a signal update is received + self._signal_info_to_labels: dict[tuple[str, SignalSource, str], set[str]] = {} + # Mapping from state labels to the list of signal requirements that define that state + self._requirements_for_label: dict[str, list[ResolvedStateSignal]] = {} + # Set of subscriptions to signal updates + self._subscriptions: set[tuple[str, SignalSource]] = set() + # Cache of the latest signal values + self._signal_value_cache: dict[tuple[str, SignalSource, str], Any] = {} + # List of currently active state labels + self._current_labels: list[str] = [] + + @staticmethod + def _endpoint(device: str, source: SignalSource): + """Static method to get the appropriate message endpoint based on the signal source.""" + if source == "readback": + return MessageEndpoints.device_readback(device) + if source == "configuration": + return MessageEndpoints.device_read_configuration(device) + if source == "limits": + return MessageEndpoints.device_limits(device) + raise ValueError( + f"Invalid signal source '{source}', please use 'readback', 'configuration', or 'limits'." + ) + + def _get_device_manager(self): + """Get the device manager.""" + if self.device_manager is None: + # pylint: disable=import-outside-toplevel + from bec_lib.client import BECClient + + bec = BECClient() + return bec.device_manager + return self.device_manager + + @staticmethod + def _get_signal_source(signal_info: dict[str, Any], error_prefix: str) -> SignalSource: + """ + Determine the signal source (readback, configuration, or limits) based on the signal information. + + Args: + signal_info (dict[str, Any]): The signal information dictionary containing at least the "kind_str" key. + error_prefix (str): A prefix to use in error messages for better context. + + Returns: + SignalSource: A string literal indicating the signal source, one of "readback", "configuration", or "limits". + """ + kind_str = str(signal_info.get("kind_str", "")).lower() + if "hinted" in kind_str or "normal" in kind_str: + return "readback" + if "config" in kind_str: + return "configuration" + raise ValueError( + f"{error_prefix} Unsupported kind: '{kind_str}' for signal : \n {yaml.dump(signal_info, indent=4)}" + ) + + @staticmethod + def _resolve_signal( + device_name: str, signal_name: str, device_manager: DeviceManagerBase, error_prefix: str + ) -> tuple[str, SignalSource]: + """ + Resolve the signal information for a given device and signal name. + + Args: + device_name (str): The name of the device. + signal_name (str): The name of the signal. + device_manager (DeviceManagerBase): The device manager instance. + error_prefix (str): A prefix to use in error messages for better context. + + Returns: + tuple[str, SignalSource]: A tuple containing the object name and the signal source. + """ + devices = device_manager.devices + try: + if not isinstance(device_name, str): + raise ValueError( + f"{error_prefix} Device name must be a string, got {type(device_name)}" + ) + device_obj: DeviceBase = devices[device_name] + except KeyError: + raise ValueError(f"{error_prefix} Device '{device_name}' not found.") from None + + # Special handling for limits, as they are not regular signals. + if signal_name in ["low_limit", "low_limit_travel"]: + return "low", "limits" + if signal_name in ["high_limit", "high_limit_travel"]: + return "high", "limits" + + signal_info = None + # This case is relevant if we are looking at a Signal directly + if device_name == signal_name and len(device_obj.root._info["signals"]) == 0: + signal_info = {"obj_name": signal_name, "kind_str": "hinted"} + # Case where we have a signal specified as a dotted name, e.g. + elif "." in signal_name: + try: + signal_obj = devices[signal_name] + except AttributeError: + raise ValueError( + f"{error_prefix} Signal '{signal_name}' not found for device '{device_name}'." + ) from None + if signal_obj.parent != device_obj: + raise ValueError( + f"{error_prefix} Signal '{signal_name}' does not belong to device '{device_name}'." + ) + signal_component = ".".join(signal_name.split(".")[1:]) + signal_info = device_obj.root._info["signals"].get(signal_component) + # Case where the signal is specified as the signal + else: + signal_info = device_obj.root._info["signals"].get(signal_name) + if signal_info is None: + for candidate in device_obj.root._info["signals"].values(): + if candidate.get("obj_name") == signal_name: + signal_info = candidate + break + + if signal_info is None: + raise ValueError( + f"{error_prefix} Signal '{signal_name}' not found for device '{device_name}'." + ) + + obj_name = signal_info.get("obj_name") + signal_source = AggregatedState._get_signal_source(signal_info, error_prefix) + return obj_name, signal_source + + @staticmethod + def _validate_user_parameter( + device_name: str, + config: DeviceConfig | SignalConfig, + device_manager: DeviceManagerBase, + error_prefix: str, + ) -> None: + """ + Validate a user parameter referenced by a signal or device configuration. + + Args: + device_name (str): The name of the device. + config (DeviceConfig | SignalConfig): The configuration containing the reference. + device_manager (DeviceManagerBase): The device manager instance. + error_prefix (str): A prefix to use in error messages for better context. + """ + if config.at is None: + return + + dev_obj = device_manager.devices.get(device_name, None) + if dev_obj is None: + raise ValueError(f"{error_prefix} Device '{device_name}' not found in device manager.") + if dev_obj.user_parameter.get(config.at) is None: + raise ValueError( + f"{error_prefix} User parameter '{config.at}' for device '{device_name}' not found." + ) + + @staticmethod + def _build_requirement_from_config( + device_name: str, + signal_name: str, + config: DeviceConfig | SignalConfig, + label: str, + device_manager: DeviceManagerBase, + error_prefix: str, + ) -> ResolvedStateSignal: + """Validate a target configuration and build its resolved signal requirement.""" + AggregatedState._validate_user_parameter(device_name, config, device_manager, error_prefix) + return AggregatedState._build_requirement_for_signal( + device_name=device_name, + signal_name=signal_name, + value=config.value, + at=config.at, + abs_tol=config.abs_tol, + label=label, + device_manager=device_manager, + error_prefix=error_prefix, + ) + + @staticmethod + def get_state_requirements( + label: str, + state_config: SubDeviceStateConfig, + device_manager: DeviceManagerBase, + error_prefix: str, + ) -> list[ResolvedStateSignal]: + """ + Get the state requirements for a given label and state configuration. + + Args: + label (str): The label for the state. + state_config (SubDeviceStateConfig): The state configuration. + device_manager (DeviceManagerBase): The device manager instance. + error_prefix (str): A prefix to use in error messages for better context. + + Returns: + list[ResolvedStateSignal]: A list of resolved state signals. + """ + state_requirements: list[ResolvedStateSignal] = [] + for device_name, config in state_config.devices.items(): + if isinstance(config, SignalConfig): + state_requirements.append( + AggregatedState._build_requirement_from_config( + device_name, device_name, config, label, device_manager, error_prefix + ) + ) + elif isinstance(config, DeviceConfig): + if config.value is not None or config.at is not None: + state_requirements.append( + AggregatedState._build_requirement_from_config( + device_name, device_name, config, label, device_manager, error_prefix + ) + ) + if config.low_limit is not None: + state_requirements.append( + AggregatedState._build_requirement_from_config( + device_name, + "low_limit", + config.low_limit, + label, + device_manager, + error_prefix, + ) + ) + if config.high_limit is not None: + state_requirements.append( + AggregatedState._build_requirement_from_config( + device_name, + "high_limit", + config.high_limit, + label, + device_manager, + error_prefix, + ) + ) + for signal_name, signal_config in (config.signals or {}).items(): + state_requirements.append( + AggregatedState._build_requirement_from_config( + device_name, + signal_name, + signal_config, + label, + device_manager, + error_prefix, + ) + ) + + return state_requirements + + def _build_rules(self) -> None: + """Build the internal rules and mappings for state evaluation based on the configuration.""" + signal_info_to_labels = {} + requirements_for_label = {} + subscriptions = set() + + for label, device_config in self.config.states.items(): + requirements = self.get_state_requirements( + label, device_config, self._get_device_manager(), self._error_prefix + ) + + for requirement in requirements: + key = (requirement.device_name, requirement.source, requirement.signal_name) + subscriptions.add((requirement.device_name, requirement.source)) + signal_info_to_labels.setdefault(key, set()).add(label) + + requirements_for_label[label] = requirements + + # Commit only after every label has resolved successfully. + self._signal_info_to_labels = signal_info_to_labels + self._requirements_for_label = requirements_for_label + self._subscriptions = subscriptions + + @staticmethod + def _build_requirement_for_signal( + device_name: str, + signal_name: str, + value: Any | None, + at: str | None, + abs_tol: float, + label: str, + device_manager: DeviceManagerBase, + error_prefix: str, + ) -> ResolvedStateSignal: + """ + Build a ResolvedStateSignal for a given device, signal, and expected value. + + Args: + device_name (str): The name of the device. + signal_name (str): The name of the signal. + value (Any | None): The expected value for the signal. + abs_tol (float): The absolute tolerance for comparing the signal value. + label (str): The label of the state that this requirement belongs to. + device_manager (DeviceManagerBase): The device manager instance. + error_prefix (str): A prefix to use in error messages for better context. + at (str | None): The user parameter for the expected value, if applicable. + + Returns: + ResolvedStateSignal: The resolved state signal requirement. + """ + if value is None and at is None: + raise ValueError( + f"{error_prefix} For device '{device_name}' and signal '{signal_name}', " + f"either 'value' or 'at' must be specified for state label '{label}'." + ) + resolved_signal_name, source = AggregatedState._resolve_signal( + device_name, signal_name, device_manager, error_prefix + ) + cfg = { + "device_name": device_name, + "signal_name": resolved_signal_name, + "expected_value": value, + "abs_tolerance": abs_tol, + "label": label, + "source": source, + "at": at, + } + return ResolvedStateSignal(**cfg) + + def start(self) -> None: + if self.started: + return + + if self.connector is None: + raise RuntimeError("Redis connector is not set.") + msg = None + try: + self._build_rules() + self._signal_value_cache.clear() + self._current_labels.clear() + affected_labels = self._fill_cache() + msg = self.evaluate(affected_labels=affected_labels) + except Exception as exc: + self._signal_info_to_labels.clear() + self._requirements_for_label.clear() + self._subscriptions.clear() + self._signal_value_cache.clear() + self._current_labels.clear() + self._handle_state_exception(exc) + self.started = False + return # Do not proceed if there was an exception during rule building or cache filling + + if msg is not None: + self._emit_state(msg) + for device, source in self._subscriptions: + self.connector.register( + self._endpoint(device, source), + cb=self._update_aggregated_state, + device=device, + source=source, + ) + super().start() + + def _fill_cache(self) -> set[str]: + """Fill the signal value cache with the current values and return the set of affected state labels.""" + affected_labels: set[str] = set() + for device, source in self._subscriptions: + endpoint = self._endpoint(device, source) + msg = self.connector.get(endpoint) + if msg is not None: + affected_labels.update(self._cache_message(device, source, msg)) + return affected_labels + + def _cache_message( + self, device: str, source: SignalSource, msg: messages.DeviceMessage + ) -> set[str]: + """Cache the signal values from a device message and return the set of affected state labels.""" + affected_labels: set[str] = set() + for signal_name, signal_data in msg.signals.items(): + key = (device, source, signal_name) + labels = self._signal_info_to_labels.get(key) + if labels is None: # signal not relevant for any state + continue + self._signal_value_cache[key] = signal_data.get("value") + affected_labels.update(labels) + return affected_labels + + def stop(self) -> None: + """Stop the state manager and unregister all subscriptions.""" + if not self.started: + return + if self.connector is not None: + for device, source in self._subscriptions: + self.connector.unregister( + self._endpoint(device, source), cb=self._update_aggregated_state + ) + super().stop() + + def _update_aggregated_state( + self, msg_obj: MessageObject, device: str, source: SignalSource, **_kwargs + ) -> None: + """Update the aggregated state based on a new device message.""" + try: + msg: messages.DeviceMessage = msg_obj.value # type: ignore ; we know it's a DeviceMessage + affected_labels = self._cache_message(device, source, msg) + if affected_labels: + state_msg = self.evaluate(affected_labels=affected_labels) + if state_msg is not None: + self._emit_state(state_msg) + except Exception as exc: + self._handle_state_exception(exc) + + def evaluate( + self, affected_labels: set[str] | None = None + ) -> messages.BeamlineStateMessage | None: + """ + Evaluate affected and currently matching labels using cached signal values. + + All matching labels are retained and reported. The configured + ``evaluation_method`` determines whether the set of matching labels is valid. + + Args: + affected_labels (set[str] | None): The set of state labels that are affected by + the latest signal update. If None, all states will be evaluated. + + Returns: + messages.BeamlineStateMessage | None: The resulting state message after evaluation, or None + if no state could be evaluated. + """ + if affected_labels is None: + return None + # We need to always extend the affected labels with the current labels, + # as the signal that updated might be not relevant for the currently active state, + # but the state should still be checked for validity. + affected_labels.update(self._current_labels) + matching_labels = sorted(label for label in affected_labels if self._label_matches(label)) + self._current_labels = matching_labels + + status = "valid" if self._matching_labels_are_valid(matching_labels) else "invalid" + return messages.BeamlineStateMessage( + name=self.config.name, + status=status, + label="|".join(matching_labels) if matching_labels else "No matching state", + ) + + def _matching_labels_are_valid(self, matching_labels: list[str]) -> bool: + """Evaluate the matching-label set according to the configured method.""" + evaluation_method = self.config.evaluation_method + if evaluation_method is None: + return True + if evaluation_method == "any": + return bool(matching_labels) + if evaluation_method == "all": + return bool(self.config.states) and len(matching_labels) == len(self.config.states) + if evaluation_method == "exclusive": + return len(matching_labels) == 1 + raise ValueError(f"Unsupported evaluation method: {evaluation_method!r}") + + def _label_matches(self, label: str) -> bool: + """Check if the given label matches the current signal values based on the defined requirements.""" + requirements = self._requirements_for_label.get(label, []) + return bool(requirements) and all( + self._requirement_matches(requirement) for requirement in requirements + ) + + @staticmethod + def get_expected_value( + requirement: ResolvedStateSignal, device_manager: DeviceManagerBase + ) -> Any: + """Get the expected value for a requirement, resolving user parameters if necessary.""" + if requirement.expected_value is not None: + return requirement.expected_value + return device_manager.devices[requirement.device_name].user_parameter[requirement.at] + + def _requirement_matches(self, requirement: ResolvedStateSignal) -> bool: + """Check if the given requirement matches the current signal values.""" + key = (requirement.device_name, requirement.source, requirement.signal_name) + cached_value = self._signal_value_cache.get(key, None) + if cached_value is None: + return False + + expected_value = self.get_expected_value(requirement, self._get_device_manager()) + try: + # Cast to float to make sure comparison with abs works as expected. + value = float(cached_value) + comparison_value = float(expected_value) + return abs(value - comparison_value) <= requirement.abs_tolerance + # Catch TypeError and ValueError in case the value is not a number or cannot be cast to float, + # in that case we fall back to exact equality. + except (TypeError, ValueError): + try: + result = cached_value == expected_value + except (TypeError, ValueError): + return False + # In case this comparison runs on comparing two arrays. + # We do not consider this comparsion as valid currently. + try: + return bool(result) + except (TypeError, ValueError): + return False + + class DeviceWithinLimitsState(DeviceBeamlineState[DeviceWithinLimitsStateConfig]): """ A state that checks if a device signal is within limits. @@ -426,6 +1085,19 @@ class DeviceWithinLimitsState(DeviceBeamlineState[DeviceWithinLimitsStateConfig] CONFIG_CLASS = DeviceWithinLimitsStateConfig + @staticmethod + def format_config_summary(config: BeamlineStateConfig) -> str: + """Return a readable summary of the monitored signal and its limits.""" + if not isinstance(config, DeviceWithinLimitsStateConfig): + return BeamlineState.format_config_summary(config) + signal = config.signal or "auto" + low_limit = config.low_limit if config.low_limit is not None else "-∞" + high_limit = config.high_limit if config.high_limit is not None else "∞" + return ( + f"{config.device} · signal={signal} · limits=[{low_limit}, {high_limit}] · " + f"tolerance={config.tolerance}" + ) + def evaluate( self, msg: messages.DeviceMessage, *args, **kwargs ) -> messages.BeamlineStateMessage: diff --git a/bec_lib/bec_lib/client.py b/bec_lib/bec_lib/client.py index eaabd1a68..4236399ea 100644 --- a/bec_lib/bec_lib/client.py +++ b/bec_lib/bec_lib/client.py @@ -20,6 +20,7 @@ from bec_lib.alarm_handler import AlarmHandler, Alarms from bec_lib.bec_service import BECService +from bec_lib.bl_state_machine import BeamlineStateMachine from bec_lib.bl_state_manager import BeamlineStateManager from bec_lib.builtin_actor_hli import BuiltinActorHli from bec_lib.callback_handler import CallbackHandler, EventType @@ -162,6 +163,7 @@ def __init__( self._username = "" self._system_user = "" self.beamline_states = None + self.state_machine = None self.messaging: MessagingContainer = None # type: ignore self.builtin_actors: BuiltinActorHli = None # type: ignore @@ -244,6 +246,7 @@ def _start_services(self): self.dap = DAPPlugins(self) self._update_username() self.beamline_states = BeamlineStateManager(client=self) + self.state_machine = BeamlineStateMachine(manager=self.beamline_states) def alarms(self, severity=Alarms.WARNING): """get the next alarm with at least the specified severity""" diff --git a/bec_lib/bec_lib/tests/utils.py b/bec_lib/bec_lib/tests/utils.py index 29144401a..1a9b5cfbc 100644 --- a/bec_lib/bec_lib/tests/utils.py +++ b/bec_lib/bec_lib/tests/utils.py @@ -22,6 +22,7 @@ if TYPE_CHECKING: # pragma: no cover from bec_lib.alarm_handler import Alarms + from bec_lib.bec_service import ServiceConfig dir_path = os.path.dirname(bec_lib.__file__) @@ -430,8 +431,10 @@ def get_device_info_mock(device_name, device_class) -> messages.DeviceInfoMessag return messages.DeviceInfoMessage( device="rt_controller", info=positioner_info_mock_with_user_access(device_name) ) - elif device_name == "samx": - return messages.DeviceInfoMessage(device="samx", info=positioner_info_mock(device_name)) + elif device_name in ["samx", "samy"]: + return messages.DeviceInfoMessage( + device=device_name, info=positioner_info_mock(device_name) + ) elif device_name == "dyn_signals": return DYN_SIGNALS_MSG elif device_name == "eiger": @@ -440,19 +443,7 @@ def get_device_info_mock(device_name, device_class) -> messages.DeviceInfoMessag if device_base_class == "positioner": signals = positioner_info_mock(device_name)["device_info"]["signals"] elif device_base_class == "signal": - signals = { - device_name: { - "metadata": { - "connected": True, - "read_access": True, - "write_access": False, - "timestamp": 0, - "status": None, - "severity": None, - "precision": None, - } - } - } + signals = {} else: signals = {} dev_info = { diff --git a/bec_lib/tests/test_beamline_states.py b/bec_lib/tests/test_beamline_states.py index 30bdfcc34..d2d5d722b 100644 --- a/bec_lib/tests/test_beamline_states.py +++ b/bec_lib/tests/test_beamline_states.py @@ -5,13 +5,17 @@ import time from unittest import mock +import numpy as np import pytest +import yaml from pydantic import BaseModel from bec_lib import bl_states, messages +from bec_lib.bl_state_machine import BeamlineStateMachine from bec_lib.bl_state_manager import ( BeamlineStateClientBase, BeamlineStateManager, + _summarize_label, build_signature_from_model, ) from bec_lib.endpoints import MessageEndpoints @@ -24,6 +28,7 @@ def state_manager(connected_connector): client = mock.MagicMock() client.connector = connected_connector manager = BeamlineStateManager(client) + yield manager @@ -47,7 +52,9 @@ def test_beamline_state_config_valid_name(self): config = bl_states.BeamlineStateConfig(name="sample_x_limits") assert config.name == "sample_x_limits" - @pytest.mark.parametrize("invalid_name", ["state-name", "class", "add", "remove", "show_all"]) + @pytest.mark.parametrize( + "invalid_name", ["state-name", "class", "add", "remove", "show_all"] + ) def test_beamline_state_config_invalid_name(self, invalid_name): with pytest.raises(ValueError): bl_states.BeamlineStateConfig(name=invalid_name) @@ -59,18 +66,50 @@ def test_device_state_config_keeps_string_device_and_signal(self): def test_device_state_config_accepts_matching_signal_device(self, dm_with_devices): config = bl_states.DeviceStateConfig( - name="state", device=dm_with_devices.devices.bpm4i, signal=dm_with_devices.devices.bpm4i + name="state", + device=dm_with_devices.devices.bpm4i, + signal=dm_with_devices.devices.bpm4i, ) assert config.device == "bpm4i" assert config.signal == "bpm4i" - def test_device_state_config_rejects_mismatched_signal_for_signal_device(self, dm_with_devices): + def test_device_state_config_rejects_mismatched_signal_for_signal_device( + self, dm_with_devices + ): with pytest.raises(ValueError, match="does not match signal device"): bl_states.DeviceStateConfig( name="state", device=dm_with_devices.devices.bpm4i, signal="bpm5i" ) + def test_aggregated_state_config_defaults_to_any_evaluation(self): + config = bl_states.AggregatedStateConfig( + name="evaluation", states={"label": {"devices": {"bpm4i": {"value": 0}}}} + ) + + assert config.evaluation_method == "any" + + @pytest.mark.parametrize("evaluation_method", ["any", "all", "exclusive", None]) + def test_aggregated_state_config_accepts_evaluation_method(self, evaluation_method): + config = bl_states.AggregatedStateConfig( + name="evaluation", + evaluation_method=evaluation_method, + states={"label": {"devices": {"bpm4i": {"value": 0}}}}, + ) + + assert config.evaluation_method == evaluation_method + + @pytest.mark.parametrize("evaluation_method", ["invalid", "null", True]) + def test_aggregated_state_config_rejects_invalid_evaluation_method( + self, evaluation_method + ): + with pytest.raises(ValueError, match="evaluation_method"): + bl_states.AggregatedStateConfig( + name="evaluation", + evaluation_method=evaluation_method, + states={"label": {"devices": {"bpm4i": {"value": 0}}}}, + ) + class TestBeamlineStateBase: def test_beamline_state_initialization_and_update(self): @@ -88,6 +127,11 @@ def evaluate(self, *args, **kwargs): assert state.connector is None assert state._last_state is None + def test_default_config_summary(self): + config = bl_states.DeviceStateConfig(name="state", device="samx") + + assert bl_states.BeamlineState.format_config_summary(config) == "device=samx" + class TestDeviceBeamlineState: def test_start_requires_connector(self, dm_with_devices): @@ -103,7 +147,9 @@ def test_start_requires_connector(self, dm_with_devices): with pytest.raises(RuntimeError, match="Redis connector is not set"): state.start() - def test_start_registers_device_callback(self, connected_connector, dm_with_devices): + def test_start_registers_device_callback( + self, connected_connector, dm_with_devices + ): state = bl_states.DeviceWithinLimitsState( name="sample_x_limits", device="samx", @@ -121,7 +167,9 @@ def test_start_registers_device_callback(self, connected_connector, dm_with_devi MessageEndpoints.device_readback("samx"), cb=state._update_device_state ) - def test_stop_unregisters_device_callback(self, connected_connector, dm_with_devices): + def test_stop_unregisters_device_callback( + self, connected_connector, dm_with_devices + ): state = bl_states.DeviceWithinLimitsState( name="sample_x_limits", device="samx", @@ -154,7 +202,8 @@ def test_update_device_state_publishes_when_state_changes( ) msg = messages.DeviceMessage( - signals={"samx": {"value": 5.0, "timestamp": 1.0}}, metadata={"stream": "primary"} + signals={"samx": {"value": 5.0, "timestamp": 1.0}}, + metadata={"stream": "primary"}, ) state._update_device_state(MessageObject(value=msg, topic="test")) @@ -169,6 +218,43 @@ def test_update_device_state_publishes_when_state_changes( class TestConcreteStates: + + @pytest.fixture(scope="function") + def aggregated_state_config(self, dm_with_devices): + """Fixture for an test aggregated state configuration.""" + dm_with_devices.devices["samx"].user_parameter["test"] = 0 + return bl_states.AggregatedStateConfig( + name="alignment", + states={ + "alignment": { + "devices": { + "samx": { + "value": 0, + "abs_tol": 0.1, + "low_limit": {"value": -20, "abs_tol": 0.1}, + "high_limit": {"value": 20, "abs_tol": 0.1}, + }, + "bpm4i": {"value": 0, "abs_tol": 0.1}, + } + }, + "measurement": { + "devices": { + "samx": { + "value": 19, + "abs_tol": 0.1, + "low_limit": {"value": -20, "abs_tol": 0.1}, + "high_limit": {"value": 20, "abs_tol": 0.1}, + "signals": {"velocity": {"value": 5, "abs_tol": 0.1}}, + }, + "bpm4i": {"value": 2, "abs_tol": 0.1}, + } + }, + "test": {"devices": {"bpm4i": {"value": 0, "abs_tol": 0.1}}}, + "string_state": {"devices": {"bpm3i": {"value": "ok"}}}, + "state_with_user_param": {"devices": {"samx": {"at": "test"}}}, + }, + ) + def test_device_within_limits_state_valid_and_invalid( self, connected_connector, dm_with_devices ): @@ -184,10 +270,12 @@ def test_device_within_limits_state_valid_and_invalid( state.start() valid_msg = messages.DeviceMessage( - signals={"samx": {"value": 5.0, "timestamp": 1.0}}, metadata={"stream": "primary"} + signals={"samx": {"value": 5.0, "timestamp": 1.0}}, + metadata={"stream": "primary"}, ) invalid_msg = messages.DeviceMessage( - signals={"samx": {"value": 11.0, "timestamp": 2.0}}, metadata={"stream": "primary"} + signals={"samx": {"value": 11.0, "timestamp": 2.0}}, + metadata={"stream": "primary"}, ) assert state.evaluate(valid_msg).status == "valid" @@ -206,13 +294,16 @@ def test_device_within_limits_state(self, connected_connector, dm_with_devices): state.start() valid = messages.DeviceMessage( - signals={"samx": {"value": 5.0, "timestamp": 1.0}}, metadata={"stream": "primary"} + signals={"samx": {"value": 5.0, "timestamp": 1.0}}, + metadata={"stream": "primary"}, ) warning = messages.DeviceMessage( - signals={"samx": {"value": 0.05, "timestamp": 2.0}}, metadata={"stream": "primary"} + signals={"samx": {"value": 0.05, "timestamp": 2.0}}, + metadata={"stream": "primary"}, ) invalid = messages.DeviceMessage( - signals={"samx": {"value": 11.0, "timestamp": 3.0}}, metadata={"stream": "primary"} + signals={"samx": {"value": 11.0, "timestamp": 3.0}}, + metadata={"stream": "primary"}, ) missing = messages.DeviceMessage( signals={"samx": {"timestamp": 4.0}}, metadata={"stream": "primary"} @@ -239,12 +330,750 @@ def test_device_within_limits_state_accepts_signal_backed_device( state.start() msg = messages.DeviceMessage( - signals={"bpm4i": {"value": 5.0, "timestamp": 1.0}}, metadata={"stream": "primary"} + signals={"bpm4i": {"value": 5.0, "timestamp": 1.0}}, + metadata={"stream": "primary"}, ) assert state.signal_name == "bpm4i" assert state.evaluate(msg).status == "valid" + def test_aggregated_state_init( + self, connected_connector, dm_with_devices, aggregated_state_config + ): + """ + Test the initialization of the AggregatedState. + + Based on the provided configuration, we expect certain callbacks to be registered with the + Redis connector. This test checks this which essentially checks the proper functionality + of the 'start' method. + """ + + state = bl_states.AggregatedState( + name=aggregated_state_config.name, + config=aggregated_state_config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + state.start() + # We should now have subscriptions on samx limits, readback and read_configuration, and bpm4i & bpm4i + info = [ + MessageEndpoints.device_readback("samx"), + MessageEndpoints.device_read_configuration("samx"), + MessageEndpoints.device_limits("samx"), + MessageEndpoints.device_readback("bpm4i"), + MessageEndpoints.device_readback("bpm3i"), + ] + for endpoint in info: + assert endpoint.endpoint in state.connector._topics_cb + + def test_aggregated_state_evaluation( + self, connected_connector, dm_with_devices, aggregated_state_config + ): + """ + Test the evaluation of the AggregatedState when receiving message updates. This should trigger a state evaluation for + the affected labels and the current state, and if the state changes, a new state should be published. + """ + state = bl_states.AggregatedState( + name=aggregated_state_config.name, + config=aggregated_state_config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + state.start() + + with ( + mock.patch.object(state, "evaluate", return_value=None) as evaluate, + mock.patch.object(state, "_emit_state") as emit_state, + ): + # Test triggering evaluation for multiple labels + # samx affects alignment and measurement, so both should be evaluated. + msg_with_2_states = messages.DeviceMessage( + signals={"samx": {"value": 5.0, "timestamp": 1.0}} + ) + msg_obj = MessageObject( + value=msg_with_2_states, + topic=MessageEndpoints.device_readback("samx").endpoint, + ) + state._update_aggregated_state(msg_obj, device="samx", source="readback") + evaluate.assert_called_once_with( + affected_labels=set( + ["state_with_user_param", "alignment", "measurement"] + ) + ) + emit_state.assert_not_called() # As evaluate is mocked to return None, _emit_state should not be called + + def test_aggregated_state_evaluate( + self, connected_connector, dm_with_devices, aggregated_state_config + ): + """ + Test the evaluate method. + We manually cache the relevant messages and then call evaluate with the affected label. + We then check if the output message has the expected status and label, and if the current labels are updated correctly. + """ + state = bl_states.AggregatedState( + name=aggregated_state_config.name, + config=aggregated_state_config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + state._build_rules() + # Assume that we are currently in test + state._current_labels = ["test"] + state._cache_message( + "samx", + "readback", + messages.DeviceMessage( + signals={"samx": {"value": 0, "timestamp": 1.0}}, + metadata={"stream": "primary"}, + ), + ) + state._cache_message( + "samx", + "configuration", + messages.DeviceMessage( + signals={"samx_velocity": {"value": 5, "timestamp": 1.0}}, + metadata={"stream": "baseline"}, + ), + ) + state._cache_message( + "samx", + "limits", + messages.DeviceMessage( + signals={ + "low": {"value": -20, "timestamp": 1.0}, + "high": {"value": 20, "timestamp": 1.0}, + }, + metadata={"stream": "baseline"}, + ), + ) + state._cache_message( + "bpm4i", + "readback", + messages.DeviceMessage( + signals={"bpm4i": {"value": 0, "timestamp": 1.0}}, + metadata={"stream": "primary"}, + ), + ) + + msg = state.evaluate(affected_labels={"alignment"}) + + assert msg.status == "valid" + assert msg.label == "alignment|test" + assert set(state._current_labels) == set(["alignment", "test"]) + dm_with_devices.devices["samx"].user_parameter["test"] = 0 + msg = state.evaluate(affected_labels={"alignment", "state_with_user_param"}) + assert msg.status == "valid" + assert msg.label == "alignment|state_with_user_param|test" + assert set(state._current_labels) == set( + ["alignment", "state_with_user_param", "test"] + ) + + dm_with_devices.devices["samx"].user_parameter["test"] = 1 + msg = state.evaluate(affected_labels={"state_with_user_param"}) + assert msg.status == "valid" + assert msg.label == "alignment|test" + assert state._current_labels == ["alignment", "test"] + + state._cache_message( + "samx", + "readback", + messages.DeviceMessage( + signals={"samx": {"value": 3, "timestamp": 2.0}}, + metadata={"stream": "primary"}, + ), + ) + + msg = state.evaluate(affected_labels={"alignment"}) + + assert msg.status == "valid" + assert msg.label == "test" + assert state._current_labels == ["test"] + + state._cache_message( + "bpm4i", + "readback", + messages.DeviceMessage( + signals={"bpm4i": {"value": 2, "timestamp": 2.0}}, + metadata={"stream": "primary"}, + ), + ) + + msg = state.evaluate(affected_labels={"alignment", "test", "measurement"}) + + assert msg.status == "invalid" + assert msg.label == "No matching state" + assert state._current_labels == [] + + @pytest.mark.parametrize( + ("evaluation_method", "targets", "expected_status", "expected_label"), + [ + ("any", (0, 1), "valid", "first"), + ("any", (0, 0), "valid", "first|second"), + ("any", (1, 2), "invalid", "No matching state"), + ("all", (0, 0), "valid", "first|second"), + ("all", (0, 1), "invalid", "first"), + ("all", (1, 2), "invalid", "No matching state"), + ("exclusive", (0, 1), "valid", "first"), + ("exclusive", (0, 0), "invalid", "first|second"), + ("exclusive", (1, 2), "invalid", "No matching state"), + (None, (0, 0), "valid", "first|second"), + (None, (1, 2), "valid", "No matching state"), + ], + ) + def test_aggregated_state_evaluation_method( + self, + connected_connector, + dm_with_devices, + evaluation_method, + targets, + expected_status, + expected_label, + ): + config = bl_states.AggregatedStateConfig( + name="evaluation", + evaluation_method=evaluation_method, + states={ + "first": {"devices": {"bpm4i": {"value": targets[0]}}}, + "second": {"devices": {"bpm4i": {"value": targets[1]}}}, + }, + ) + state = bl_states.AggregatedState( + config=config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + state._build_rules() + state._cache_message( + "bpm4i", + "readback", + messages.DeviceMessage(signals={"bpm4i": {"value": 0, "timestamp": 1}}), + ) + + msg = state.evaluate(affected_labels={"first", "second"}) + + assert msg.status == expected_status + assert msg.label == expected_label + assert state._current_labels == ( + [] if expected_label == "No matching state" else expected_label.split("|") + ) + + def test_aggregated_state_preserves_incremental_label_evaluation( + self, connected_connector, dm_with_devices + ): + config = bl_states.AggregatedStateConfig( + name="evaluation", + states={ + "affected": {"devices": {"samx": {"value": 1}}}, + "current": {"devices": {"bpm4i": {"value": 0}}}, + "unrelated": {"devices": {"bpm3i": {"value": "ok"}}}, + }, + ) + state = bl_states.AggregatedState( + config=config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + state._build_rules() + state._signal_value_cache.update( + { + ("samx", "readback", "samx"): 1, + ("bpm4i", "readback", "bpm4i"): 0, + ("bpm3i", "readback", "bpm3i"): "ok", + } + ) + state._current_labels = ["current"] + + with mock.patch.object( + state, "_label_matches", wraps=state._label_matches + ) as label_matches: + msg = state.evaluate(affected_labels={"affected"}) + + assert {call.args[0] for call in label_matches.call_args_list} == { + "affected", + "current", + } + assert msg.status == "valid" + assert msg.label == "affected|current" + + def test_aggregated_state_runtime_evaluation_method_update( + self, connected_connector, dm_with_devices + ): + config = bl_states.AggregatedStateConfig( + name="evaluation", + evaluation_method="any", + states={ + "alignment": {"devices": {"bpm4i": {"value": 0}}}, + "duplicate": {"devices": {"bpm4i": {"value": 0}}}, + }, + ) + state = bl_states.AggregatedState( + config=config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + connected_connector.set_and_publish( + MessageEndpoints.device_readback("bpm4i"), + messages.DeviceMessage(signals={"bpm4i": {"value": 0, "timestamp": 1}}), + ) + state.start() + initial = connected_connector.xread( + MessageEndpoints.beamline_state("evaluation"), from_start=True + )[-1]["data"] + assert initial.status == "valid" + assert initial.label == "alignment|duplicate" + + state.update_parameters(evaluation_method="exclusive") + state.restart() + + updated = connected_connector.xread( + MessageEndpoints.beamline_state("evaluation"), from_start=True + )[-1]["data"] + assert updated.status == "invalid" + assert updated.label == "alignment|duplicate" + + def test_aggregated_state_exception_handling( + self, connected_connector, dm_with_devices, aggregated_state_config + ): + """ + Test that if an exception is raised during the evaluation of the state, this is properly handled and an alarm is raised. + We check that the evaluate method is called and that if it raises an exception, the raise_alarm method of the connector + is called, and a state with status "unknown" and label "broken state" is published. + """ + state = bl_states.AggregatedState( + name=aggregated_state_config.name, + config=aggregated_state_config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + state.start() + msg = messages.DeviceMessage( + signals={"samx": {"value": 0, "timestamp": 1.0}}, + metadata={"stream": "primary"}, + ) + msg_obj = MessageObject( + value=msg, topic=MessageEndpoints.device_readback("samx").endpoint + ) + + with ( + mock.patch.object( + state, "evaluate", side_effect=RuntimeError("broken state") + ) as evaluate, + mock.patch.object(connected_connector, "raise_alarm") as raise_alarm, + ): + state._update_aggregated_state(msg_obj, device="samx", source="readback") + + evaluate.assert_called_once_with( + affected_labels={"state_with_user_param", "alignment", "measurement"} + ) + raise_alarm.assert_called_once() + out = connected_connector.xread( + MessageEndpoints.beamline_state("alignment"), from_start=True + ) + assert out[-1]["data"].status == "unknown" + assert out[-1]["data"].label == "broken state" + assert state.raised_warning is True + + def test_aggregated_state_transitions_between_labels( + self, connected_connector, dm_with_devices, aggregated_state_config + ): + """ + Test the transitions between different labels of the aggregated state. We simulate the messages that would trigger + the transitions and check that the output message has the expected status and label, and that the current labels are updated correctly. + """ + state = bl_states.AggregatedState( + name=aggregated_state_config.name, + config=aggregated_state_config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + state.start() + + def update(device, source, signals): + msg = messages.DeviceMessage( + signals=signals, metadata={"stream": "primary"} + ) + msg_obj = MessageObject( + value=msg, topic=state._endpoint(device, source).endpoint + ) + state._update_aggregated_state(msg_obj, device=device, source=source) + out = connected_connector.xread( + MessageEndpoints.beamline_state("alignment"), from_start=True + ) + return out[-1]["data"] + + msg = update( + "samx", "configuration", {"samx_velocity": {"value": 5, "timestamp": 1.0}} + ) + assert msg.status == "invalid" + + update( + "samx", + "limits", + { + "low": {"value": -20, "timestamp": 1.0}, + "high": {"value": 20, "timestamp": 1.0}, + }, + ) + update("samx", "readback", {"samx": {"value": 0, "timestamp": 1.0}}) + msg = update("bpm4i", "readback", {"bpm4i": {"value": 0, "timestamp": 1.0}}) + assert msg.status == "valid" + assert msg.label == "alignment|state_with_user_param|test" + + msg = update("samx", "readback", {"samx": {"value": 19, "timestamp": 2.0}}) + assert msg.status == "valid" + assert msg.label == "test" + + msg = update("bpm4i", "readback", {"bpm4i": {"value": 2, "timestamp": 2.0}}) + assert msg.status == "valid" + assert msg.label == "measurement" + + @pytest.mark.parametrize( + ("cached_value", "expected_value", "at", "abs_tolerance", "matches"), + [ + (1.05, 1.0, None, 0.1, True), + (1.2, 1.0, None, 0.1, False), + (5, 5, None, 0.0, True), + (np.int64(5), 5, None, 0.0, True), + (np.float64(1.05), 1.0, None, 0.1, True), + ("ok", "ok", None, 0.0, True), + ("not-ok", "ok", None, 0.0, False), + ([1, 2], 1, None, 0.0, False), + (np.array([1.0, 2.0]), 1.0, None, 0.1, False), + (np.array([1.0, 2.0]), np.array([1.0, 2.0]), None, 0.0, False), + ], + ) + def test_aggregated_state_requirement_matches( + self, + connected_connector, + dm_with_devices, + aggregated_state_config, + cached_value, + expected_value, + at, + abs_tolerance, + matches, + ): + """ + Test the evaluation of requirements in the aggregated state. We manually set the signal value + cache and then call the _requirement_matches method with a requirement, and check if the output is as expected. + """ + state = bl_states.AggregatedState( + name=aggregated_state_config.name, + config=aggregated_state_config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + requirement = bl_states.ResolvedStateSignal( + label="alignment", + device_name="bpm4i", + signal_name="bpm4i", + expected_value=expected_value, + at=at, + abs_tolerance=abs_tolerance, + source="readback", + ) + state._signal_value_cache[("bpm4i", "readback", "bpm4i")] = cached_value + + assert state._requirement_matches(requirement) is matches + + def test_device_config_requires_at_least_one_target(self): + with pytest.raises(ValueError, match="At least one of value"): + bl_states.DeviceConfig() + + @pytest.mark.parametrize( + "config_class", [bl_states.SignalConfig, bl_states.DeviceConfig] + ) + def test_target_config_rejects_negative_abs_tolerance(self, config_class): + with pytest.raises(ValueError, match="greater than or equal to 0"): + config_class(value=1, abs_tol=-0.1) + + @pytest.mark.parametrize( + "config_class", [bl_states.SignalConfig, bl_states.DeviceConfig] + ) + def test_target_config_rejects_value_and_at(self, config_class): + with pytest.raises(ValueError, match="Cannot specify both 'value' and 'at'"): + config_class(value=1, at="target") + + def test_signal_config_requires_value_or_at(self): + with pytest.raises(ValueError, match="Either 'value' or 'at'"): + bl_states.SignalConfig() + + @pytest.mark.parametrize("target", [0, False, ""]) + def test_aggregated_state_resolves_falsy_user_parameter( + self, dm_with_devices, target + ): + dm_with_devices.devices["samx"].user_parameter["target"] = target + config = bl_states.SubDeviceStateConfig(devices={"samx": {"at": "target"}}) + + requirements = bl_states.AggregatedState.get_state_requirements( + "test", config, dm_with_devices, "test" + ) + + assert len(requirements) == 1 + resolved = bl_states.AggregatedState.get_expected_value( + requirements[0], dm_with_devices + ) + assert resolved == target + assert type(resolved) is type(target) + + @pytest.mark.parametrize("target", [0, False, ""]) + def test_aggregated_state_keeps_falsy_literal_target(self, dm_with_devices, target): + config = bl_states.SubDeviceStateConfig(devices={"samx": {"value": target}}) + + requirements = bl_states.AggregatedState.get_state_requirements( + "test", config, dm_with_devices, "test" + ) + + assert len(requirements) == 1 + assert requirements[0].expected_value == target + assert type(requirements[0].expected_value) is type(target) + + def test_aggregated_state_resolves_nested_user_parameters(self, dm_with_devices): + dm_with_devices.devices["samx"].user_parameter.update( + {"low": -5, "high": 5, "speed": 2} + ) + config = bl_states.SubDeviceStateConfig( + devices={ + "samx": { + "low_limit": {"at": "low"}, + "high_limit": {"at": "high"}, + "signals": {"velocity": {"at": "speed"}}, + } + } + ) + + requirements = bl_states.AggregatedState.get_state_requirements( + "test", config, dm_with_devices, "test" + ) + + assert [(req.signal_name, req.at) for req in requirements] == [ + ("low", "low"), + ("high", "high"), + ("samx_velocity", "speed"), + ] + assert [ + bl_states.AggregatedState.get_expected_value(req, dm_with_devices) + for req in requirements + ] == [-5, 5, 2] + + @pytest.mark.parametrize( + ("target_config", "user_parameters"), + [ + ({"low_limit": {"at": "missing"}}, {}), + ({"low_limit": {"at": "missing"}}, {"missing": None}), + ({"high_limit": {"at": "missing"}}, {}), + ({"high_limit": {"at": "missing"}}, {"missing": None}), + ({"signals": {"velocity": {"at": "missing"}}}, {}), + ({"signals": {"velocity": {"at": "missing"}}}, {"missing": None}), + ], + ) + def test_aggregated_state_rejects_missing_nested_user_parameter( + self, dm_with_devices, target_config, user_parameters + ): + dm_with_devices.devices["samx"].user_parameter.update(user_parameters) + config = bl_states.SubDeviceStateConfig(devices={"samx": target_config}) + + with pytest.raises(ValueError, match="User parameter 'missing'"): + bl_states.AggregatedState.get_state_requirements( + "test", config, dm_with_devices, "test" + ) + + def test_aggregated_state_endpoint_rejects_unknown_source(self): + with pytest.raises(ValueError, match="Invalid signal source"): + bl_states.AggregatedState._endpoint("samx", "unknown") + + def test_aggregated_state_get_device_manager_falls_back_to_client(self): + state = bl_states.AggregatedState( + name="alignment", states={"label": {"devices": {"samx": {"value": 0}}}} + ) + client = mock.MagicMock() + + with mock.patch("bec_lib.client.BECClient", return_value=client): + assert state._get_device_manager() is client.device_manager + + def test_aggregated_state_get_signal_source_rejects_unsupported_kind(self): + with pytest.raises(ValueError, match="Unsupported kind"): + bl_states.AggregatedState._get_signal_source( + {"kind_str": "omitted", "obj_name": "samx_unused"}, "test" + ) + + def test_aggregated_state_resolve_signal_edge_cases(self, dm_with_devices): + assert bl_states.AggregatedState._resolve_signal( + "samx", "low_limit_travel", dm_with_devices, "test" + ) == ("low", "limits") + assert bl_states.AggregatedState._resolve_signal( + "samx", "high_limit_travel", dm_with_devices, "test" + ) == ("high", "limits") + assert bl_states.AggregatedState._resolve_signal( + "samx", "samx_velocity", dm_with_devices, "test" + ) == ("samx_velocity", "configuration") + + with pytest.raises(ValueError, match="Device 'missing' not found"): + bl_states.AggregatedState._resolve_signal( + "missing", "missing", dm_with_devices, "test" + ) + with pytest.raises(ValueError, match="Device name must be a string"): + bl_states.AggregatedState._resolve_signal( + 1, "samx", dm_with_devices, "test" + ) + with pytest.raises(ValueError, match="Signal 'missing_signal' not found"): + bl_states.AggregatedState._resolve_signal( + "samx", "missing_signal", dm_with_devices, "test" + ) + with pytest.raises(ValueError, match="Unsupported kind"): + bl_states.AggregatedState._resolve_signal( + "samx", "unused", dm_with_devices, "test" + ) + + def test_aggregated_state_resolve_dotted_signal_edge_cases(self, dm_with_devices): + assert bl_states.AggregatedState._resolve_signal( + "samx", "samx.velocity", dm_with_devices, "test" + ) == ("samx_velocity", "configuration") + + with pytest.raises(ValueError, match="does not belong"): + bl_states.AggregatedState._resolve_signal( + "samx", "samy.velocity", dm_with_devices, "test" + ) + + devices = mock.MagicMock() + devices.__getitem__.side_effect = [ + dm_with_devices.devices["samx"], + AttributeError, + ] + manager = mock.MagicMock(devices=devices) + with pytest.raises(ValueError, match="Signal 'samx.missing' not found"): + bl_states.AggregatedState._resolve_signal( + "samx", "samx.missing", manager, "test" + ) + + def test_aggregated_state_start_edge_cases( + self, connected_connector, dm_with_devices, aggregated_state_config + ): + state = bl_states.AggregatedState( + config=aggregated_state_config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + state.started = True + with mock.patch.object(state, "_build_rules") as build_rules: + state.start() + build_rules.assert_not_called() + + state = bl_states.AggregatedState( + config=aggregated_state_config, + redis_connector=None, + device_manager=dm_with_devices, + ) + with pytest.raises(RuntimeError, match="Redis connector is not set"): + state.start() + + def test_aggregated_state_start_handles_rule_build_error( + self, connected_connector, dm_with_devices, aggregated_state_config + ): + state = bl_states.AggregatedState( + config=aggregated_state_config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + + with ( + mock.patch.object( + state, "_build_rules", side_effect=RuntimeError("bad rules") + ), + mock.patch.object(state, "_handle_state_exception") as handle_exception, + mock.patch.object(connected_connector, "register") as register, + ): + state.start() + + handle_exception.assert_called_once() + register.assert_not_called() + assert state.started is False + assert state._subscriptions == set() + assert state._requirements_for_label == {} + assert state._signal_info_to_labels == {} + + def test_aggregated_state_fill_cache_uses_existing_messages( + self, connected_connector, dm_with_devices, aggregated_state_config + ): + state = bl_states.AggregatedState( + config=aggregated_state_config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + state._build_rules() + connected_connector.set_and_publish( + MessageEndpoints.device_readback("samx"), + messages.DeviceMessage(signals={"samx": {"value": 0, "timestamp": 1.0}}), + ) + + affected_labels = state._fill_cache() + + assert affected_labels == {"alignment", "measurement", "state_with_user_param"} + assert state._signal_value_cache[("samx", "readback", "samx")] == 0 + + def test_aggregated_state_cache_ignores_irrelevant_signals( + self, connected_connector, dm_with_devices, aggregated_state_config + ): + state = bl_states.AggregatedState( + config=aggregated_state_config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + state._build_rules() + + affected_labels = state._cache_message( + "samx", + "readback", + messages.DeviceMessage( + signals={"samx_unused": {"value": 1, "timestamp": 1.0}}, + metadata={"stream": "primary"}, + ), + ) + + assert affected_labels == set() + assert ("samx", "readback", "samx_unused") not in state._signal_value_cache + + def test_aggregated_state_stop_unregisters_subscriptions( + self, connected_connector, dm_with_devices, aggregated_state_config + ): + state = bl_states.AggregatedState( + config=aggregated_state_config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + state.start() + + with mock.patch.object(connected_connector, "unregister") as unregister: + state.stop() + + assert unregister.call_count == len(state._subscriptions) + assert state.started is False + + def test_aggregated_state_stop_is_noop_before_start( + self, connected_connector, dm_with_devices, aggregated_state_config + ): + state = bl_states.AggregatedState( + config=aggregated_state_config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + + with mock.patch.object(connected_connector, "unregister") as unregister: + state.stop() + + unregister.assert_not_called() + + def test_aggregated_state_evaluate_without_affected_labels( + self, connected_connector, dm_with_devices, aggregated_state_config + ): + state = bl_states.AggregatedState( + config=aggregated_state_config, + redis_connector=connected_connector, + device_manager=dm_with_devices, + ) + + assert state.evaluate() is None + class TestBeamlineStateManager: def test_manager_registers_for_state_updates(self, connected_connector): @@ -254,7 +1083,9 @@ def test_manager_registers_for_state_updates(self, connected_connector): with mock.patch.object(connected_connector, "register") as register: BeamlineStateManager(client) - register.assert_called_once_with(MessageEndpoints.available_beamline_states(), cb=mock.ANY) + register.assert_called_once_with( + MessageEndpoints.available_beamline_states(), cb=mock.ANY + ) def test_manager_is_ready_when_no_state_update_exists(self, connected_connector): client = mock.MagicMock() @@ -324,15 +1155,25 @@ def test_on_state_update_creates_client_attribute(self, state_manager): assert "sample_y_limits" in state_manager._states assert isinstance( - state_manager._states["sample_y_limits"], bl_states.DeviceWithinLimitsStateConfig + state_manager._states["sample_y_limits"], + bl_states.DeviceWithinLimitsStateConfig, + ) + assert isinstance( + getattr(state_manager, "sample_y_limits"), BeamlineStateClientBase ) - assert isinstance(getattr(state_manager, "sample_y_limits"), BeamlineStateClientBase) - def test_update_parameters_from_client_updates_state_and_publishes(self, state_manager): + def test_update_parameters_from_client_updates_state_and_publishes( + self, state_manager + ): config = messages.BeamlineStateConfig( name="limits", state_type="DeviceWithinLimitsState", - parameters={"name": "limits", "device": "samx", "low_limit": 0.0, "high_limit": 10.0}, + parameters={ + "name": "limits", + "device": "samx", + "low_limit": 0.0, + "high_limit": 10.0, + }, ) update = messages.AvailableBeamlineStatesMessage(states=[config]) state_manager._on_state_update({"data": update}, parent=state_manager) @@ -347,11 +1188,41 @@ def test_update_parameters_from_client_updates_state_and_publishes(self, state_m assert out assert isinstance(out[-1]["data"], messages.AvailableBeamlineStatesMessage) - def test_external_parameter_update_refreshes_existing_client_state(self, state_manager): + def test_runtime_evaluation_method_is_published(self, state_manager): + config = messages.BeamlineStateConfig( + name="evaluation", + state_type="AggregatedState", + parameters={ + "name": "evaluation", + "evaluation_method": "any", + "states": {"alignment": {"devices": {"bpm4i": {"value": 0}}}}, + }, + ) + state_manager._on_state_update( + {"data": messages.AvailableBeamlineStatesMessage(states=[config])}, + parent=state_manager, + ) + + state_manager.evaluation.update_parameters(evaluation_method="exclusive") + + update = state_manager._connector.xread( + MessageEndpoints.available_beamline_states(), from_start=True + )[-1]["data"] + published = next(state for state in update.states if state.name == "evaluation") + assert published.parameters["evaluation_method"] == "exclusive" + + def test_external_parameter_update_refreshes_existing_client_state( + self, state_manager + ): initial = messages.BeamlineStateConfig( name="limits", state_type="DeviceWithinLimitsState", - parameters={"name": "limits", "device": "samx", "low_limit": 0.0, "high_limit": 10.0}, + parameters={ + "name": "limits", + "device": "samx", + "low_limit": 0.0, + "high_limit": 10.0, + }, ) state_manager._on_state_update( {"data": messages.AvailableBeamlineStatesMessage(states=[initial])}, @@ -377,7 +1248,10 @@ def test_external_parameter_update_refreshes_existing_client_state(self, state_m assert state_manager._states["limits"].low_limit == 1.0 assert state_manager._states["limits"].high_limit == 9.0 assert state_manager._states["limits"].tolerance == 0.25 - assert state_manager.limits._state.model_dump(exclude_none=True) == updated.parameters + assert ( + state_manager.limits._state.model_dump(exclude_none=True) + == updated.parameters + ) def test_client_get_returns_unknown_without_status_message(self, state_manager): config = messages.BeamlineStateConfig( @@ -394,7 +1268,10 @@ def test_client_get_returns_unknown_without_status_message(self, state_manager): state_manager._on_state_update({"data": update}, parent=state_manager) result = state_manager.sample_y_limits.get() - assert result == {"status": "unknown", "label": "No state information available."} + assert result == { + "status": "unknown", + "label": "No state information available.", + } def test_client_get_returns_latest_status_message(self, state_manager): config = messages.BeamlineStateConfig( @@ -489,6 +1366,9 @@ def test_show_all_prints_table(self, state_manager, capsys): state = bl_states.DeviceWithinLimitsStateConfig( name="sample_y_limits", device="samy", low_limit=0.0, high_limit=10.0 ) + assert bl_states.DeviceWithinLimitsState.format_config_summary(state) == ( + "samy · signal=auto · limits=[0.0, 10.0] · tolerance=0.1" + ) with mock.patch.object(state_manager, "_wait_for_initial_state"): state_manager.add(state) @@ -496,3 +1376,193 @@ def test_show_all_prints_table(self, state_manager, capsys): captured = capsys.readouterr() assert "sample_y_limits" in (captured.out + captured.err) + + @staticmethod + def _aggregated_display_config(): + return bl_states.AggregatedStateConfig( + name="beamline_mode", + evaluation_method="all", + states={ + "alignment": { + "devices": { + "samx": {"value": 0, "signals": {"velocity": {"value": 5}}} + }, + "transition_metadata": {"description": "Prepare alignment"}, + }, + "measurement": {"devices": {"bpm4i": {"value": 100}}}, + "parked": { + "devices": { + "samx": { + "low_limit": {"value": -10}, + "high_limit": {"value": 10}, + } + } + }, + "test": {"devices": {"samy": {"at": "in"}}}, + }, + ) + + def test_show_all_summarizes_aggregated_state(self, state_manager, capsys): + state = self._aggregated_display_config() + state_manager._add_state(state) + + assert bl_states.AggregatedState.format_config_summary(state) == ( + "all · 4 labels · 3 devices · 6 requirements · 1 transition" + ) + assert _summarize_label("alignment|measurement|parked|test") == ( + "alignment|measurement|parked|… (+1)" + ) + + with mock.patch.object( + state_manager.beamline_mode, + "get", + return_value={ + "status": "valid", + "label": "alignment|measurement|parked|test", + }, + ): + state_manager.show_all() + + output = capsys.readouterr().out + assert "beamline_mode" in output + assert "4 labels" in output + assert "requirements" in output + assert "transition_metadata" not in output + assert "velocity" not in output + + def test_state_client_describe_prints_nested_parameters( + self, state_manager, capsys + ): + state_manager._add_state(self._aggregated_display_config()) + + state_manager.beamline_mode.describe() + + output = capsys.readouterr().out + for expected in ( + "beamline_mode", + "AggregatedState", + "evaluation_method", + "'all'", + "states", + "alignment", + "devices", + "samx", + "signals", + "velocity", + "transition_metadata", + "Prepare alignment", + ): + assert expected in output + + +class TestStateMachine: + + @pytest.fixture() + def state_machine(self, state_manager): + state_machine = BeamlineStateMachine(manager=state_manager) + return state_machine + + @pytest.fixture() + def config_dict(self): + return { + "alignment": { + "bl_state_class": "AggregatedState", + "config": { + "evaluation_method": "any", + "states": { + "alignment": { + "devices": { + "samx": { + "value": 0, + "abs_tol": 0.1, + "signals": { + "velocity": {"value": 5, "abs_tol": 0.1} + }, + } + } + } + }, + }, + } + } + + @pytest.mark.timeout(30) + def test_load_from_config_with_dict( + self, state_machine: BeamlineStateMachine, tmp_path, config_dict + ): + """Test loading configuration from a dictionary or file.""" + + # Load valid configuration from dictionary + with mock.patch.object(state_machine._manager, "add") as manager_add: + state_machine.load_from_config(config_path=None, config_dict=config_dict) + manager_add.assert_called_once_with( + bl_states.AggregatedStateConfig( + name="alignment", + states=config_dict["alignment"]["config"]["states"], + ), + skip_existing=False, + ) + # Loading with both config_path and config_dict should raise an error + with pytest.raises(ValueError): + state_machine.load_from_config( + config_path="path/to/config.yaml", config_dict=config_dict + ) + # Loading with neither config_path nor config_dict should raise an error + with pytest.raises(ValueError): + state_machine.load_from_config(config_path=None, config_dict=None) + + # Loading from file should work. + config_path = tmp_path / "config.yaml" + with open(config_path, "w", encoding="utf-8") as f: + yaml.dump(config_dict, f) + state_machine.load_from_config(config_path=str(config_path)) + manager_add.assert_called_with( + bl_states.AggregatedStateConfig( + name="alignment", + states=config_dict["alignment"]["config"]["states"], + ), + skip_existing=False, + ) + + def test_load_from_config_rejects_invalid_evaluation_method( + self, state_machine, config_dict + ): + config_dict["alignment"]["config"]["evaluation_method"] = "invalid" + + with ( + mock.patch.object(state_machine._manager, "clear_all") as clear_all, + pytest.raises(ValueError, match="evaluation_method"), + ): + state_machine.load_from_config(config_dict=config_dict) + + clear_all.assert_not_called() + + def test_load_from_yaml_accepts_null_evaluation_method( + self, state_machine, config_dict, tmp_path + ): + config_dict["alignment"]["config"]["evaluation_method"] = None + config_path = tmp_path / "state_config.yaml" + with open(config_path, "w", encoding="utf-8") as stream: + yaml.safe_dump(config_dict, stream) + + with mock.patch.object(state_machine._manager, "add") as manager_add: + state_machine.load_from_config(config_path=str(config_path)) + + loaded_config = manager_add.call_args.args[0] + assert loaded_config.evaluation_method is None + + def test_load_from_config_forwards_skip_existing( + self, state_machine: BeamlineStateMachine, config_dict + ): + """Test that skip_existing is forwarded to the state manager.""" + with mock.patch.object(state_machine._manager, "add") as manager_add: + state_machine.load_from_config( + config_dict=config_dict, flush=False, skip_existing=True + ) + + manager_add.assert_called_once_with( + bl_states.AggregatedStateConfig( + name="alignment", states=config_dict["alignment"]["config"]["states"] + ), + skip_existing=True, + ) diff --git a/bec_server/bec_server/device_server/devices/device_serializer.py b/bec_server/bec_server/device_server/devices/device_serializer.py index 63bdc1e4e..1caf254b7 100644 --- a/bec_server/bec_server/device_server/devices/device_serializer.py +++ b/bec_server/bec_server/device_server/devices/device_serializer.py @@ -202,6 +202,8 @@ def get_device_info( "kind_str": signal_obj.kind.name, "doc": doc, "describe": signal_obj.describe().get(signal_obj.name, {}), + "read_access": getattr(signal_obj, "read_access", None), + "write_access": getattr(signal_obj, "write_access", None), # pylint: disable=protected-access "metadata": signal_obj._metadata, "labels": sorted(signal_obj._ophyd_labels_), diff --git a/bec_server/bec_server/scan_server/scans/state_transition_scan.py b/bec_server/bec_server/scan_server/scans/state_transition_scan.py new file mode 100644 index 000000000..344dde4af --- /dev/null +++ b/bec_server/bec_server/scan_server/scans/state_transition_scan.py @@ -0,0 +1,330 @@ +""" +Updated move scan implementation for coordinated motor repositioning commands. + +Scan procedure: + - prepare_scan + - open_scan + - stage + - pre_scan + - scan_core + - at_each_point (optionally called by scan_core) + - post_scan + - unstage + - close_scan + - on_exception (called if any exception is raised during the scan) +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Tuple + +from bec_lib.alarm_handler import AlarmBase, Alarms +from bec_lib.bl_states import AggregatedState, SubDeviceStateConfig +from bec_lib.device import DeviceBase, Positioner, Signal +from bec_lib.endpoints import MessageEndpoints +from bec_lib.logger import bec_logger +from bec_lib.messages import AlarmMessage, ErrorInfo +from bec_server.scan_server.scans.scan_base import ScanBase +from bec_server.scan_server.scans.scan_modifier import scan_hook + +if TYPE_CHECKING: + from bec_lib.bl_states import AggregatedStateConfig, ResolvedStateSignal + from bec_lib.messages import AvailableBeamlineStatesMessage + +logger = bec_logger.logger + + +class StateTransitionScanError(AlarmBase): + """Exception raised when an RPC call fails.""" + + def __init__(self, exc_type: str, message: str, compact_message: str) -> None: + alarm = AlarmMessage( + severity=Alarms.MAJOR, + info=ErrorInfo( + exception_type=exc_type, + error_message=message, + compact_error_message=compact_message, + ), + ) + super().__init__(alarm, Alarms.MAJOR, handled=False) + + +class StateTransitionScan(ScanBase): + + # Scan Type: Hardware triggered or software triggered? + # If the main trigger and readout logic is done within the at_each_point method in scan_core, choose SOFTWARE_TRIGGERED. + # If the main trigger and readout logic is implemented on a device that is simply kicked off in this scan, choose HARDWARE_TRIGGERED. + # This primarily serves as information for devices: The device may need to react differently if a software trigger is expected + # for every point. + scan_type = None + + # Scan name: This is the name of the scan, e.g. "line_scan". This is used for display purposes and to identify the scan type in user interfaces. + # Choose a descriptive name that does not conflict with existing scan names. + scan_name = "_v4_state_transition" + + # We set is_scan to False to separate this class from the other scans in the user interface + is_scan = False + + def __init__(self, *args, state_name: str, target_label: str, **kwargs): + """ + State transition scan that moves a motor in between two states. + The main purpose of this scan is to be used in conjunction with state + management in BEC, and transitioning the beamline in-between different aggregated states. + """ + super().__init__(**kwargs) + self.state_name = state_name + self.target_label = target_label + + # We need to sort the devices and signals in the config, and identify which of them are motor setpoint/readback pairs + # and which of them are just readouts and thereby can not be set within the transition. + self._signals_to_set: list[Tuple[Signal, Any]] = [] + self._limits_to_set: dict[str, Tuple[Positioner, float, float]] = {} + self._devices_to_set: list[Tuple[Positioner, float]] = [] + + # pylint: disable=protected-access + @scan_hook + def prepare_scan(self): + """ + Prepare the scan. This can include any steps that need to be executed + before the scan is opened, such as preparing the positions (if not done already) + or setting up the devices. + """ + # Check if the state and the target label exists, if yes, fetch the configuration for the target state + self.config_for_label = self._fetch_config_for_label(self.state_name, self.target_label) + requirements: list[ResolvedStateSignal] = AggregatedState.get_state_requirements( + self.target_label, self.config_for_label, self.device_manager, "StateTransitionScan" + ) + self._signals_to_set, self._limits_to_set, self._devices_to_set = ( + self._fetch_devices_signals_and_limits_to_set(requirements) + ) + + self.update_scan_info(scan_report_devices=[dev for dev, _ in self._devices_to_set]) + + @scan_hook + def open_scan(self): + """ + Open the scan. + This step must call self.actions.open_scan() to ensure that a new scan is + opened. Make sure to prepare the scan metadata before, either in + prepare_scan() or in open_scan() itself and call self.update_scan_info(...) + to update the scan metadata if needed. + """ + + @scan_hook + def stage(self): + """ + Stage the devices for the upcoming scan. The stage logic is typically + implemented on the device itself (i.e. by the device's stage method). + However, if there are any additional steps that need to be executed before + staging the devices, they can be implemented here. + """ + + @scan_hook + def pre_scan(self): + """ + Pre-scan steps to be executed before the main scan logic. + This is typically the last chance to prepare the devices before the core scan + logic is executed. For example, this is a good place to initialize time-criticial + devices, e.g. devices that have a short timeout. + The pre-scan logic is typically implemented on the device itself. + """ + + @scan_hook + def scan_core(self): + """ + Core scan logic to be executed during the scan. + This is where the main scan logic should be implemented. + """ + # Motors + motors = [element[0] for element in self._devices_to_set] + target_positions = [element[1] for element in self._devices_to_set] + current_positions = self.components.get_start_positions(motors) + # First set signals, then scan report instructions, otherwise the IPython progress will + # otherwise resolve on the RID from the scan together with the "set" of the signal. + logger.warning(f"Request moving motors {motors}.") + self.actions.add_scan_report_instruction_readback( + devices=motors, start=current_positions, stop=target_positions + ) + self.components.move_and_wait(motors, target_positions) + # Limits & Signals + self._set_signals() + self._set_limits() + + @scan_hook + def at_each_point(self): + """ + Logic to be executed at each point during the scan. This is called by the step_scan method at each point. + + Args: + motors (list[str | DeviceBase]): List of motor names or device instances being moved. + positions (np.ndarray): Current positions of the motors, shape (len(motors),). + last_positions (np.ndarray | None): Previous positions of the motors, shape (len(motors),) or None if this is the first point. + """ + + @scan_hook + def post_scan(self): + """ + Post-scan steps to be executed after the main scan logic. + """ + + @scan_hook + def unstage(self): + """Unstage the scan by executing post-scan steps.""" + + @scan_hook + def close_scan(self): + """Close the scan.""" + + @scan_hook + def on_exception(self, exception: Exception): + """ + Handle exceptions that occur during the scan. + This is a good place to implement any cleanup logic that needs to be executed in case of an exception, + such as returning the devices to a safe state or moving the motors back to their starting position. + """ + + ################# + ## Custom Methods + ################# + + def _set_signals(self): + """Method to set signals in the transition.""" + for signal_obj, target_value in self._signals_to_set: + # Check if signal is writable before setting it, if not skip. + if self._check_if_signal_has_write_access(signal_obj): + signal_obj.set(target_value).wait() + + def _set_limits(self): + """Method to set limits for devices in the transition.""" + for dev_name, (dev_obj, low_limit, high_limit) in self._limits_to_set.items(): + dev_obj.limits = [low_limit, high_limit] + + def _check_if_signal_has_write_access(self, signal_obj: Signal) -> bool: + """ + Check if a signal has write access based on its signal information. The issue here is that signals of a + device follow a slightly different pattern. Therefore, we have to first check "_info" for signals + and if that is empty, check '_signal_info' for sub-signals of devices. + + Args: + signal_obj (Signal): Signal object to check. + Returns: + bool: True if the signal has write access, False otherwise. + """ + write_access = signal_obj._info.get("write_access", None) + if write_access is None: + write_access = signal_obj._signal_info.get("write_access", False) + return write_access + + def _fetch_devices_signals_and_limits_to_set( + self, requirements: list[ResolvedStateSignal] + ) -> Tuple[dict, dict, dict]: + """ + This method fetches the device signals, limits and devices to set based on a list of state requirements. + It returns a tuple containing three elements: + - signals_to_set (list[Tuple[Signal, Any]]): List of signals to set with their target values. + - limits_to_set (dict[str, Tuple[Positioner, float, float]]): Dictionary of devices and their limits to set. + - devices_to_set (list[Tuple[Positioner, float]]): List of devices to set with their target positions. + + Args: + requirements (list[ResolvedStateSignal]): List of state requirements to fetch the device signals and limits for. + + Returns: + Tuple containing: + - signals_to_set (list[Tuple[Signal, Any]]): List of signals to set with their target values. + - limits_to_set (dict[str, Tuple[Positioner, float, float]]): Dictionary of devices and their limits to set. + - devices_to_set (list[Tuple[Positioner, float]]): List of devices to set with their target positions. + """ + _signals_to_set: list[Tuple[Signal, Any]] = [] + _limits_to_set: dict[str, Tuple[Positioner, float, float]] = {} + _devices_to_set: list[Tuple[Positioner, float]] = [] + for req in requirements: + dev_obj: DeviceBase = self.device_manager.devices.get(req.device_name) + # Device not found + if dev_obj is None: + raise StateTransitionScanError( + exc_type="DeviceNotFound", + message=f"Device {req.device_name} not found in device manager.", + compact_message=f"Device {req.device_name} not found.", + ) + try: + expected_value = AggregatedState.get_expected_value(req, self.device_manager) + except Exception: + raise StateTransitionScanError( + exc_type="ExpectedValueError", + message=f"Failed to get expected value for device {req.device_name} and signal {req.signal_name} with expected value {req.expected_value} or user_parameter {req.at}.", + compact_message=f"Failed to get expected value for device {req.device_name}.", + ) + # First we handle Signals logic + if isinstance(dev_obj, Signal): + _signals_to_set.append((dev_obj, expected_value)) + continue + # Positioner and Device logic. Devices must implement .set for this to work, otherwise we can not set them and we raise an error + if isinstance(dev_obj, DeviceBase): + # Handle motor-specific logic here + # First we handle logic for motions of the motor. Device_name and signal_name will be equivalent here + if req.signal_name == req.device_name: + _devices_to_set.append((dev_obj, expected_value)) + continue + if req.source == "limits": + if req.device_name not in _limits_to_set: + _limits_to_set[req.device_name] = ( + dev_obj, + dev_obj.low_limit, + dev_obj.high_limit, + ) + if req.signal_name == "low": + _limits_to_set[req.device_name] = ( + dev_obj, + expected_value, + _limits_to_set[req.device_name][2], + ) + elif req.signal_name == "high": + _limits_to_set[req.device_name] = ( + dev_obj, + _limits_to_set[req.device_name][1], + expected_value, + ) + continue + signal_obj = self._get_signal_object(dev_obj, req.signal_name) + if signal_obj is None: + raise StateTransitionScanError( + exc_type="SignalNotFound", + message=f"Signal {req.signal_name} for device {req.device_name} not found in device manager.", + compact_message=f"Signal {req.signal_name} for device {req.device_name} not found.", + ) + _signals_to_set.append((signal_obj, expected_value)) + continue + # Return the collected signals, limits and devices to set + return _signals_to_set, _limits_to_set, _devices_to_set + + def _get_signal_object(self, device_obj: DeviceBase, signal_name: str) -> Signal: + for component_name, info in device_obj._info["signals"].items(): + if info["obj_name"] == signal_name: + return getattr(device_obj, component_name) + + def _fetch_config_for_label(self, state_name: str, target_label: str) -> SubDeviceStateConfig: + available_states_msg: AvailableBeamlineStatesMessage = self.redis_connector.get_last( + MessageEndpoints.available_beamline_states() + ) + if available_states_msg is None: + raise ValueError( + "No available beamline states found in Redis. Cannot fetch configuration for state transition scan." + ) + configs = [ + state for state in available_states_msg["data"].states if state.name == state_name + ] + if len(configs) == 0: + raise ValueError(f"State {state_name} not found in available states.") + elif len(configs) > 1: # Should not be possible, but just in case + raise ValueError(f"Multiple states with name {state_name} found in available states.") + config: AggregatedStateConfig = configs[0] + if config.state_type != "AggregatedState": + raise ValueError( + f"State {state_name} is not an aggregated state. Transitions are only supported for aggregated states." + ) + available_labels = list(config.parameters["states"].keys()) + if target_label not in available_labels: + raise ValueError( + f"Target label {target_label} not found in state {state_name}. Available labels: {available_labels}" + ) + return SubDeviceStateConfig.model_validate(config.parameters["states"][target_label]) diff --git a/bec_server/bec_server/scan_server/tests/scan_fixtures.py b/bec_server/bec_server/scan_server/tests/scan_fixtures.py index c3bc2508b..5365b5bbd 100644 --- a/bec_server/bec_server/scan_server/tests/scan_fixtures.py +++ b/bec_server/bec_server/scan_server/tests/scan_fixtures.py @@ -127,6 +127,18 @@ def dotted_name(self): def limits(self): return self._limits + @limits.setter + def limits(self, value): + self._limits = tuple(value) + + @property + def low_limit(self): + return self._limits[0] + + @property + def high_limit(self): + return self._limits[1] + @property def enabled(self): return self._enabled @@ -234,6 +246,18 @@ def dotted_name(self): def limits(self): return self._limits + @limits.setter + def limits(self, value): + self._limits = tuple(value) + + @property + def low_limit(self): + return self._limits[0] + + @property + def high_limit(self): + return self._limits[1] + @property def enabled(self): return self._enabled @@ -368,6 +392,7 @@ def v4_scan_assembler(readout_priority: ReadoutPriorityContainer, device_manager def _assemble_scan(scan_type, *scan_args, **scan_kwargs): scan_id = scan_kwargs.pop("scan_id", "scan-id-test") + connector = scan_kwargs.pop("connector", None) or ConnectorMock("") try: scan_cls = scan_classes[scan_type] @@ -375,7 +400,6 @@ def _assemble_scan(scan_type, *scan_args, **scan_kwargs): available = ", ".join(sorted(scan_classes)) raise KeyError(f"Unknown scan type '{scan_type}'. Available: {available}") from exc - connector = ConnectorMock("") instruction_handler = InstructionHandler(connector) device_names = sorted( set(_infer_v4_device_names(scan_cls, scan_args, scan_kwargs)) diff --git a/bec_server/tests/tests_scan_server/scans_v4/test_state_transition_scan.py b/bec_server/tests/tests_scan_server/scans_v4/test_state_transition_scan.py new file mode 100644 index 000000000..06a081419 --- /dev/null +++ b/bec_server/tests/tests_scan_server/scans_v4/test_state_transition_scan.py @@ -0,0 +1,305 @@ +from unittest import mock + +import pytest +from ophyd_devices.sim.sim_positioner import SimPositioner + +from bec_lib import messages +from bec_lib.device import Positioner, _PermissiveDeviceModel +from bec_lib.endpoints import MessageEndpoints +from bec_lib.messages import AvailableBeamlineStatesMessage, BeamlineStateConfig +from bec_server.device_server.devices.device_serializer import get_device_info +from bec_server.scan_server.tests.scan_hook_tests import run_scan_tests + + +def assert_scan_open_called(scan): + scan.actions.open_scan = mock.MagicMock() + scan.open_scan() + scan.actions.open_scan.assert_not_called() + + +def assert_scan_stage_called(scan): + scan.actions.stage_all_devices = mock.MagicMock() + scan.stage() + scan.actions.stage_all_devices.assert_not_called() + + +def assert_scan_pre_scan_called(scan): + scan.actions.pre_scan_all_devices = mock.MagicMock() + scan.pre_scan() + scan.actions.pre_scan_all_devices.assert_not_called() + + +def assert_scan_unstage_called(scan): + scan.actions.unstage_all_devices = mock.MagicMock() + scan.unstage() + scan.actions.unstage_all_devices.assert_not_called() + + +def assert_scan_close_scan_called(scan, status=None): + scan.actions.close_scan = mock.MagicMock() + scan.close_scan() + scan.actions.close_scan.assert_not_called() + + +ACQUIRE_DEFAULT_HOOK_TESTS = [ + ("open_scan", [assert_scan_open_called]), + ("stage", [assert_scan_stage_called]), + ("pre_scan", [assert_scan_pre_scan_called]), + ("unstage", [assert_scan_unstage_called]), + ("close_scan", [assert_scan_close_scan_called]), +] + + +@pytest.fixture +def state_transition_connector(connected_connector): + connected_connector.xadd( + MessageEndpoints.available_beamline_states(), + { + "data": messages.AvailableBeamlineStatesMessage( + states=[ + messages.BeamlineStateConfig( + name="test", + title="Test state", + state_type="AggregatedState", + parameters={ + "states": { + "alignment": { + "devices": { + "samx": { + "value": 1.5, + "low_limit": {"value": -2}, + "high_limit": {"value": 2}, + "signals": {"velocity": {"value": 0.5}}, + }, + "samy": { + "value": 0.5, + "low_limit": {"value": -1}, + "high_limit": {"value": 1}, + }, + } + } + } + }, + ) + ] + ) + }, + ) + return connected_connector + + +def publish_state_config(connector, *, label, devices): + """Publish one aggregated-state label for transition tests.""" + connector.xadd( + MessageEndpoints.available_beamline_states(), + { + "data": AvailableBeamlineStatesMessage( + states=[ + BeamlineStateConfig( + name="test", + title="Test state", + state_type="AggregatedState", + parameters={"states": {label: {"devices": devices}}}, + ) + ] + ) + }, + ) + + +@pytest.fixture +def simulated_samx(device_manager): + # dev_obj = SimPositioner(name="samx") + name = "samx" + dev = SimPositioner(name=name) + config = _PermissiveDeviceModel( + enabled=True, + deviceClass="ophyd_devices.sim.sim_positioner.SimPositioner", + readoutPriority="baseline", + ) + info = get_device_info(dev, connect=True) + dev_man_obj = Positioner( + name=name, + info=info, + config=config, + class_name=config.deviceClass, + parent=device_manager, + ) + return dev_man_obj + + +@pytest.mark.timeout(20) +@pytest.mark.parametrize(("hook_name", "hook_tests"), ACQUIRE_DEFAULT_HOOK_TESTS) +def test_state_transition_default_hooks( + v4_scan_assembler, + state_transition_connector, + nth_done_status_mock, + hook_name, + hook_tests, +): + """Test default hooks open_scan, stage, pre_scan, unstage, and close_scan for the StateTransitionScan.""" + scan = v4_scan_assembler( + "_v4_state_transition", + state_name="test", + target_label="alignment", + connector=state_transition_connector, + ) + + run_scan_tests( + scan, [(hook_name, hook_tests)], nth_done_status_mock=nth_done_status_mock + ) + + +@pytest.mark.timeout(20) +def test_state_transition_prepare_scan( + v4_scan_assembler, state_transition_connector, device_manager, simulated_samx +): + """Test prepare scan hook for the StateTransitionScan.""" + device_manager.add_device(simulated_samx, replace=True) + scan = v4_scan_assembler( + "_v4_state_transition", + state_name="test", + target_label="alignment", + connector=state_transition_connector, + ) + + scan.prepare_scan() + + devices_to_set = {(device.name, value) for device, value in scan._devices_to_set} + limits_to_set = { + device_name: (device.name, low_limit, high_limit) + for device_name, (device, low_limit, high_limit) in scan._limits_to_set.items() + } + signals_to_set = { + (signal.full_name, value) for signal, value in scan._signals_to_set + } + + assert devices_to_set == {("samx", 1.5), ("samy", 0.5)} + assert limits_to_set == {"samx": ("samx", -2, 2), "samy": ("samy", -1, 1)} + assert signals_to_set == {("samx_velocity", 0.5)} + + +@pytest.mark.timeout(20) +def test_state_transition_scan_core( + v4_scan_assembler, state_transition_connector, device_manager, simulated_samx +): + device_manager.add_device(simulated_samx, replace=True) + scan = v4_scan_assembler( + "_v4_state_transition", + state_name="test", + target_label="alignment", + connector=state_transition_connector, + ) + scan.prepare_scan() + signal_by_name = {signal.full_name: signal for signal, _ in scan._signals_to_set} + velocity_set_status = mock.MagicMock() + signal_by_name["samx_velocity"].set = mock.MagicMock( + return_value=velocity_set_status + ) + with ( + mock.patch.object( + scan.components, "get_start_positions", return_value=[0, 0] + ) as mock_get_start_positions, + mock.patch.object(scan.components, "move_and_wait") as mock_move_and_wait, + mock.patch.object( + scan.actions, "add_scan_report_instruction_readback" + ) as mock_add_scan_report_instruction_readback, + mock.patch.object(scan, "_set_limits") as mock_set_limits, + ): + scan.scan_core() + mock_get_start_positions.assert_called_once() + mock_add_scan_report_instruction_readback.assert_called_once_with( + devices=[ + scan.device_manager.devices["samx"], + scan.device_manager.devices["samy"], + ], + start=[0, 0], + stop=[1.5, 0.5], + ) + mock_move_and_wait.assert_called_once_with( + [scan.device_manager.devices["samx"], scan.device_manager.devices["samy"]], + [1.5, 0.5], + ) + signal_by_name["samx_velocity"].set.assert_called_once_with(0.5) + velocity_set_status.wait.assert_called_once_with() + mock_set_limits.assert_called_once() + + +@pytest.mark.timeout(20) +def test_state_transition_prepare_scan_resolves_user_parameters( + v4_scan_assembler, connected_connector, device_manager, simulated_samx +): + simulated_samx._config["userParameter"] = { + "position": 0, + "low": -3, + "high": 3, + "velocity": False, + } + device_manager.add_device(simulated_samx, replace=True) + publish_state_config( + connected_connector, + label="dynamic", + devices={ + "samx": { + "at": "position", + "low_limit": {"at": "low"}, + "high_limit": {"at": "high"}, + "signals": {"velocity": {"at": "velocity"}}, + } + }, + ) + scan = v4_scan_assembler( + "_v4_state_transition", + state_name="test", + target_label="dynamic", + connector=connected_connector, + ) + + scan.prepare_scan() + + assert [(device.name, value) for device, value in scan._devices_to_set] == [ + ("samx", 0) + ] + assert type(scan._devices_to_set[0][1]) is int + assert { + name: (device.name, low, high) + for name, (device, low, high) in scan._limits_to_set.items() + } == {"samx": ("samx", -3, 3)} + assert [(signal.full_name, value) for signal, value in scan._signals_to_set] == [ + ("samx_velocity", False) + ] + assert type(scan._signals_to_set[0][1]) is bool + + +@pytest.mark.timeout(20) +@pytest.mark.parametrize( + "target_config", + [ + {"at": "missing"}, + {"low_limit": {"at": "missing"}}, + {"high_limit": {"at": "missing"}}, + {"signals": {"velocity": {"at": "missing"}}}, + ], +) +def test_state_transition_prepare_scan_rejects_missing_user_parameter( + v4_scan_assembler, + connected_connector, + device_manager, + simulated_samx, + target_config, +): + device_manager.add_device(simulated_samx, replace=True) + publish_state_config( + connected_connector, + label="dynamic", + devices={"samx": target_config}, + ) + scan = v4_scan_assembler( + "_v4_state_transition", + state_name="test", + target_label="dynamic", + connector=connected_connector, + ) + + with pytest.raises(ValueError, match="User parameter 'missing'"): + scan.prepare_scan() diff --git a/bec_server/tests/tests_scan_server/test_beamline_state_manager.py b/bec_server/tests/tests_scan_server/test_beamline_state_manager.py index 345136334..d157a3ba4 100644 --- a/bec_server/tests/tests_scan_server/test_beamline_state_manager.py +++ b/bec_server/tests/tests_scan_server/test_beamline_state_manager.py @@ -163,6 +163,35 @@ def test_state_manager_rejects_abstract_state_type(state_manager): state_manager.update_states(msg) +def test_state_manager_restarts_aggregated_state_after_evaluation_method_update(state_manager): + def state_message(evaluation_method): + return messages.AvailableBeamlineStatesMessage( + states=[ + messages.BeamlineStateConfig( + name="evaluation", + state_type="AggregatedState", + parameters={ + "name": "evaluation", + "evaluation_method": evaluation_method, + "states": { + "alignment": {"devices": {"bpm4i": {"value": 0}}} + }, + }, + ) + ] + ) + + with mock.patch.object(bl_states.AggregatedState, "start"): + state_manager.update_states(state_message("any")) + + state = state_manager._states["evaluation"] + with mock.patch.object(state, "restart") as restart: + state_manager.update_states(state_message("exclusive")) + + assert state.config.evaluation_method == "exclusive" + restart.assert_called_once_with() + + @pytest.mark.timeout(5) def test_states_restarted_when_device_config_updated( state_manager, connected_connector, fake_bl_states