sientia-dataops-laborious_temporal/e2e/opc_test_server.py

"""
In-process OPC UA server for E2E tests (asyncua).

Provides writable prediction/confidence nodes and optional write faults
(Tier-1 BadSessionIdInvalid via PreWrite callback).
"""

from __future__ import annotations

import socket
from dataclasses import dataclass
from typing import TYPE_CHECKING

from asyncua import Server, ua
from asyncua.common.callback import CallbackType
from asyncua.common.utils import ServiceError

if TYPE_CHECKING:
    from asyncua.common.node import Node


UNKNOWN_NODE_ID = 'ns=99;i=9999'


@dataclass(frozen=True)
class OpcE2ENodeIds:
    """NodeId strings used in opc_output_config for E2E workflows."""

    prediction: str
    confidence: str
    unknown: str = UNKNOWN_NODE_ID


class OpcE2ETestServer:
    """
    Ephemeral asyncua server with Laborious E2E variables and controllable faults.

    Args:
        host: Bind address (default 127.0.0.1).
    """

    def __init__(self, host: str = '127.0.0.1') -> None:
        self._host = host
        self._server: Server | None = None
        self._prediction_node: Node | None = None
        self._confidence_node: Node | None = None
        self._session_bad_on_write = False
        self._url: str | None = None
        self._node_ids: OpcE2ENodeIds | None = None

    @property
    def url(self) -> str:
        if self._url is None:
            raise RuntimeError('OPC E2E server is not started')
        return self._url

    @property
    def node_ids(self) -> OpcE2ENodeIds:
        if self._node_ids is None:
            raise RuntimeError('OPC E2E server is not started')
        return self._node_ids

    def set_session_bad_on_write(self, enabled: bool) -> None:
        """
        When enabled, every client Write is rejected with BadSessionIdInvalid.

        Args:
            enabled (bool): Turn Tier-1 session fault injection on or off.
        """
        self._session_bad_on_write = enabled

    async def start(self) -> OpcE2ENodeIds:
        """
        Start the OPC UA server on a free TCP port.

        Return:
            OpcE2ENodeIds: NodeId strings for prediction and confidence tags.
        """
        port = _free_port(self._host)
        self._url = f'opc.tcp://{self._host}:{port}/freeopcua/server/'

        server = Server()
        server.set_endpoint(self._url)
        await server.init()
        server.iserver.callback_service.addListener(
            CallbackType.PreWrite,
            self._pre_write_callback,
        )

        idx = await server.register_namespace('http://sientia.test/laborious-e2e')
        e2e_object = await server.nodes.objects.add_object(idx, 'LaboriousE2E')
        prediction = await e2e_object.add_variable(
            idx,
            'Prediction',
            ua.Variant(0.0, ua.VariantType.Float),
        )
        confidence = await e2e_object.add_variable(
            idx,
            'Confidence',
            ua.Variant(0.0, ua.VariantType.Float),
        )
        await prediction.set_writable()
        await confidence.set_writable()

        await server.start()
        self._server = server
        self._prediction_node = prediction
        self._confidence_node = confidence
        self._node_ids = OpcE2ENodeIds(
            prediction=prediction.nodeid.to_string(),
            confidence=confidence.nodeid.to_string(),
        )
        return self._node_ids

    async def stop(self) -> None:
        """Stop the OPC UA server and release the listening port."""
        if self._server is not None:
            await self._server.stop()
            self._server = None
        self._prediction_node = None
        self._confidence_node = None
        self._url = None
        self._node_ids = None
        self._session_bad_on_write = False

    async def read_prediction(self) -> float:
        """
        Read the current prediction variable value from the address space.

        Return:
            float: Stored prediction value.
        """
        if self._prediction_node is None:
            raise RuntimeError('OPC E2E server is not started')
        value = await self._prediction_node.read_value()
        return float(value)

    async def read_confidence(self) -> float:
        """
        Read the current confidence variable value from the address space.

        Return:
            float: Stored confidence value.
        """
        if self._confidence_node is None:
            raise RuntimeError('OPC E2E server is not started')
        value = await self._confidence_node.read_value()
        return float(value)

    async def _pre_write_callback(self, _event, _service) -> None:
        if self._session_bad_on_write:
            raise ServiceError(ua.StatusCodes.BadSessionIdInvalid)


def _free_port(host: str) -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.bind((host, 0))
        return int(sock.getsockname()[1])


def build_opc_output_config(
    node_ids: OpcE2ENodeIds,
    *,
    prediction_tag: str | None = None,
    confidence_tag: str | None = None,
    prediction_only: bool = False,
    server_key: str = '1',
) -> dict[str, dict]:
    """
    Build opc_output_config for PredictionsBatch using real server NodeIds.

    Args:
        node_ids (OpcE2ENodeIds): Node ids from OpcE2ETestServer.
        prediction_tag (str | None): Override prediction NodeId (default: node_ids.prediction).
        confidence_tag (str | None): Override confidence NodeId (default: node_ids.confidence).
        prediction_only (bool): When True, omit confidence_tags (single write per activity).
        server_key (str): OPC server id key in opc_output_config.

    Return:
        dict: opc_output_config payload for workflow input.
    """
    pred = prediction_tag if prediction_tag is not None else node_ids.prediction
    conf = confidence_tag if confidence_tag is not None else node_ids.confidence
    server_config: dict = {
        'prediction_tags': {pred: {'data_type': 'float'}},
    }
    if not prediction_only:
        server_config['confidence_tags'] = {conf: {'data_type': 'float'}}
    return {server_key: server_config}