sientia-dataops-opc-ingestor/e2e/helpers.py

"""Shared helpers for the gateway E2E tests (OPC UA + Redis + MongoDB + Postgres + MinIO + NATS)."""

import asyncio
import copy
import inspect
import json
import os
import time
from collections.abc import Awaitable, Callable
from pathlib import Path
from typing import Any

from asyncua import ua
from asyncua.ua.ua_binary import Primitives, nodeid_to_binary
from pymongo.database import Database
from sqlalchemy import text
from sqlalchemy.engine import Engine

from ingestor.config import GatewayConfig
from ingestor.gateway_runtime import GatewayRuntime
from ingestor.ingestor import Ingestor

SCENARIO_INPUTS_DIR = Path(__file__).parent / 'scenario_inputs'

# Collections the gateway reads (connectors, opc_certificates) or writes
# (readings, notification_queue). Dropped before every test with the whole database.
NOTIFICATIONS = 'notification_queue'
CERT_BUCKET = 'opc-certificates'
PREDICTIONS_TABLE = 'predictions'


def _resolve_markers(value: Any, servers: dict) -> Any:
    """Replace ``@server:<name>`` with the URL of that in-process OPC server."""
    if isinstance(value, dict):
        return {k: _resolve_markers(v, servers) for k, v in value.items()}
    if isinstance(value, list):
        return [_resolve_markers(v, servers) for v in value]
    if isinstance(value, str) and value.startswith('@server:'):
        return servers[value.removeprefix('@server:')].url
    return value


def load_scenario_input(name: str, servers: dict | None = None, **overrides: Any) -> Any:
    """
    Load a JSON scenario input from ``e2e/scenario_inputs``.

    Args:
        name: File name (e.g. ``slot_single_server.json``).
        servers: ``{name: OpcTestServer}`` used to resolve ``@server:<name>`` markers.
        **overrides: Top-level keys replaced in the loaded dict.

    Return:
        Any: The resolved JSON content.
    """
    data = json.loads((SCENARIO_INPUTS_DIR / name).read_text(encoding='utf-8'))
    data = _resolve_markers(data, servers or {})
    if overrides:
        data = {**data, **copy.deepcopy(overrides)}
    return data


async def wait_until(
    predicate: Callable[[], Any | Awaitable[Any]], timeout: float = 15.0, interval: float = 0.1
) -> Any:
    """
    Poll ``predicate`` (sync or async) until it returns a truthy value.

    Sleeps on the event loop between attempts, so OPC subscription callbacks and NATS
    handlers keep running while the test waits.

    Raises:
        AssertionError: When the predicate is still falsy after ``timeout`` seconds.
    """
    deadline = time.monotonic() + timeout
    while True:
        result = predicate()
        if inspect.isawaitable(result):
            result = await result
        if result:
            return result
        if time.monotonic() > deadline:
            raise AssertionError(f'condition not met within {timeout}s')
        await asyncio.sleep(interval)


async def stays_false(
    predicate: Callable[[], Any | Awaitable[Any]], duration: float = 1.5, interval: float = 0.1
) -> bool:
    """True when ``predicate`` never becomes truthy during ``duration`` seconds."""
    deadline = time.monotonic() + duration
    while time.monotonic() < deadline:
        result = predicate()
        if inspect.isawaitable(result):
            result = await result
        if result:
            return False
        await asyncio.sleep(interval)
    return True


class Pod:
    """
    One gateway pod, wired exactly as ``ingestor/app.py`` does (minus the Prometheus
    HTTP server and the sleep between iterations): ``Ingestor.prepare_ingestor``, then
    ``GatewayRuntime`` on its data manager and connection pool. ``tick`` runs one main
    loop iteration (``Ingestor.loop`` + ``GatewayRuntime.idle_sweep``).

    Every component reads its settings from the environment at construction time, so
    the caller sets the env (``HOSTNAME`` included) before ``start``.
    """

    def __init__(self, name: str, runtime: bool = True):
        self.name = name
        self.with_runtime = runtime
        self.ingestor: Ingestor | None = None
        self.runtime: GatewayRuntime | None = None

    async def start(self) -> 'Pod':
        os.environ['HOSTNAME'] = self.name
        self.ingestor = Ingestor()
        await self.ingestor.prepare_ingestor()
        if self.with_runtime:
            manager = self.ingestor.ingestor_manager
            assert manager is not None
            self.runtime = GatewayRuntime(
                config=GatewayConfig.from_env(),
                data_manager=manager.data_manager,
                pool=manager.pool,
                logger=self.ingestor.logger,
                notification_handler=self.ingestor.notification_handler,
                metrics_controller=self.ingestor.metrics_controller,
                metadata=self.ingestor.metadata,
            )
            await self.runtime.start()
            consumer = self.runtime.nats_consumer
            if consumer.nc is not None:
                # PING/PONG: the server has processed the SUB, so a command published
                # right after start is not lost for lack of a subscriber.
                await consumer.nc.flush()
        return self

    @property
    def manager(self):
        assert self.ingestor is not None and self.ingestor.ingestor_manager is not None
        return self.ingestor.ingestor_manager

    @property
    def slots(self) -> set[str]:
        """Slots this pod currently manages."""
        return set(self.manager.managed_tags)

    async def tick(self, times: int = 1) -> None:
        assert self.ingestor is not None
        for _ in range(times):
            await self.ingestor.loop()
            if self.runtime is not None:
                await self.runtime.idle_sweep()

    async def stop(self) -> None:
        if self.runtime is not None:
            await self.runtime.stop()
            self.runtime = None
        if self.ingestor is not None:
            await self.ingestor.shutdown()
            self.ingestor.notification_handler.shutdown()
            self.ingestor = None


def redis_json(redis_client, key: str) -> Any:
    """Value stored by the gateway's RedisRepository (JSON-encoded), or None."""
    raw = redis_client.get(key)
    return None if raw is None else json.loads(raw)


def readings(db: Database, collection: str, **query: Any) -> list[dict]:
    """Documents the read path persisted into ``collection`` (oldest first)."""
    return list(db[collection].find(query, {'_id': 0}).sort('inserted_at', 1))


def notifications(db: Database, prefix: str = '') -> list[dict]:
    """Notifications sent by the gateway whose ``notification_id`` starts with ``prefix``."""
    return [
        n
        for n in db[NOTIFICATIONS].find({}, {'_id': 0})
        if str(n.get('notification_id', '')).startswith(prefix)
    ]


def insert_prediction(engine: Engine, model_id: int, timestamp: str, write_status: str | None):
    """One ``sientia_data.predictions`` row, as the Laborious inserts it before publishing C2."""
    with engine.begin() as conn:
        conn.execute(
            text(
                'INSERT INTO sientia_data.predictions (model_id, prediction, '
                'prediction_confidence, response_time, prediction_status, "timestamp", '
                'write_status) VALUES (:model_id, 1.5, 0.9, 0.1, :status, :ts, :write_status)'
            ),
            {
                'model_id': model_id,
                'status': 'success',
                'ts': timestamp,
                'write_status': write_status,
            },
        )


def write_status(engine: Engine, model_id: int, timestamp: str) -> str | None:
    """``write_status`` of a prediction row (``'<missing>'`` when the row does not exist)."""
    with engine.begin() as conn:
        row = conn.execute(
            text(
                'SELECT write_status FROM sientia_data.predictions '
                'WHERE model_id = :model_id AND "timestamp" = :ts'
            ),
            {'model_id': model_id, 'ts': timestamp},
        ).first()
    return '<missing>' if row is None else row[0]


def prediction_count(engine: Engine) -> int:
    with engine.begin() as conn:
        return conn.execute(text('SELECT count(*) FROM sientia_data.predictions')).scalar_one()


# DataValue encoding mask bits (OPC UA Part 6, 5.2.2.17).
STATUS_CODE_BIT = 0x02
SOURCE_TIMESTAMP_BIT = 0x04


def datavalue_mask(request: bytes, addr: str) -> int:
    """
    Encoding mask of the DataValue written to ``addr`` in a serialized WriteRequest.

    An omitted StatusCode decodes as Good on the server (OPC UA default), so whether
    the gateway sent the quality is only visible on the wire. WriteValue = NodeId,
    AttributeId, IndexRange, DataValue: the byte right after the first three is the mask.
    """
    prefix = (
        nodeid_to_binary(ua.NodeId.from_string(addr))
        + Primitives.UInt32.pack(ua.AttributeIds.Value)
        + Primitives.String.pack(None)
    )
    return request[request.index(prefix) + len(prefix)]