sientia-dataops-orchestrator_temporal/e2e/helpers.py

"""
Shared helpers for orchestrator E2E tests (Temporal workflows + Mongo + Redis + Postgres).
"""

import asyncio
import json
import re
import uuid
from datetime import UTC, datetime, timedelta
from email.message import EmailMessage, Message
from pathlib import Path
from typing import Any

from pymongo import MongoClient
from redis import Redis
from sqlalchemy import text
from sqlalchemy.engine import Engine
from temporalio.client import Client

SCENARIO_INPUTS_DIR = Path(__file__).parent / 'scenario_inputs'
ORCHESTRATOR_TASK_QUEUE = 'orchestrator-test-queue'

MONGO_COLLECTIONS = (
    'notification_queue',
    'receiver_groups',
    'orchestrated_schedules',
    'pipelines',
    'opc_servers',
    'opc-servers',
)

DATETIME_FORMAT_MS_WITH_TZ = '%Y-%m-%d %H:%M:%S.%f%z'

_TIMESTAMP_MARKER_PATTERN = re.compile(r'^@now(?:([+-])(\d+)([smhd]))?$')
_TIMESTAMP_UNIT_TO_KWARG = {'s': 'seconds', 'm': 'minutes', 'h': 'hours', 'd': 'days'}


def _resolve_timestamp_marker(value: Any) -> Any:
    """
    Convert ``@now`` / ``@now-1h`` markers into timezone-aware datetimes.

    Args:
        value: Any JSON value. Only strings matching the marker pattern are converted.

    Return:
        Any: The resolved datetime or the original value unchanged.
    """
    if not isinstance(value, str):
        return value
    match = _TIMESTAMP_MARKER_PATTERN.match(value)
    if not match:
        return value
    sign, amount, unit = match.groups()
    now = datetime.now(UTC)
    if sign is None:
        return now
    delta = timedelta(**{_TIMESTAMP_UNIT_TO_KWARG[unit]: int(amount)})
    return now + delta if sign == '+' else now - delta


def _resolve_payload(payload: Any) -> Any:
    """Recursively walk a JSON-like structure resolving ``@now`` timestamp markers."""
    if isinstance(payload, dict):
        return {key: _resolve_payload(value) for key, value in payload.items()}
    if isinstance(payload, list):
        return [_resolve_payload(item) for item in payload]
    return _resolve_timestamp_marker(payload)


def load_scenario_input(name: str, **overrides: Any) -> dict[str, Any]:
    """
    Load a scenario JSON file from e2e/scenario_inputs and apply overrides.

    Strings matching ``@now`` or ``@now[+-]<int>[smhd]`` (anywhere in the payload)
    are converted to timezone-aware ``datetime`` instances. This lets scenario
    files declare relative timestamps such as ``"updated_at": "@now-1h"``.

    Args:
        name: File name (with or without .json suffix).
        **overrides: Top-level keys to replace in the loaded dict.

    Return:
        dict[str, Any]: Scenario payload with timestamp markers resolved.
    """
    file_name = name if name.endswith('.json') else f'{name}.json'
    file_path = SCENARIO_INPUTS_DIR / file_name
    with file_path.open('r', encoding='utf-8') as handle:
        payload = json.load(handle)
    payload = _resolve_payload(payload)
    payload.update(overrides)
    return payload


def make_workflow_id(prefix: str) -> str:
    """Build a unique workflow id using a prefix and UUID suffix."""
    return f'{prefix}-{uuid.uuid4().hex[:12]}'


async def start_and_await_workflow(
    client: Client,
    workflow_run,
    input_data: dict[str, Any],
    workflow_id: str,
    *,
    task_queue: str = ORCHESTRATOR_TASK_QUEUE,
    timeout: float = 120.0,
) -> Any:
    """
    Start a workflow and wait for its result.

    Args:
        client: Temporal client (default namespace).
        workflow_run: Workflow run method (e.g. Orchestrator.run).
        input_data: Workflow input payload.
        workflow_id: Unique workflow id.
        task_queue: Task queue for the orchestrator worker.
        timeout: Max seconds to wait for completion.

    Return:
        Workflow result value.
    """
    handle = await client.start_workflow(
        workflow_run,
        input_data,
        id=workflow_id,
        task_queue=task_queue,
    )
    return await asyncio.wait_for(handle.result(), timeout=timeout)


def seed_pipelines(
    mongo_uri: str,
    database: str,
    pipelines: list[dict[str, Any]],
) -> None:
    """Insert pipeline documents into the test Mongo database."""
    client = MongoClient(mongo_uri)
    try:
        collection = client[database]['pipelines']
        if pipelines:
            collection.insert_many(pipelines)
    finally:
        client.close()


def seed_opc_servers(
    mongo_uri: str,
    database: str,
    servers: list[dict[str, Any]],
    *,
    collection: str = 'opc_servers',
) -> None:
    """Insert OPC server documents into the test Mongo database."""
    client = MongoClient(mongo_uri)
    try:
        coll = client[database][collection]
        if servers:
            coll.insert_many(servers)
    finally:
        client.close()


def seed_receiver_groups(
    mongo_uri: str,
    database: str,
    groups: list[dict[str, Any]],
) -> None:
    """Insert receiver group documents into the test Mongo database."""
    client = MongoClient(mongo_uri)
    try:
        collection = client[database]['receiver_groups']
        if groups:
            collection.insert_many(groups)
    finally:
        client.close()


def seed_notifications(
    mongo_uri: str,
    database: str,
    notifications: list[dict[str, Any]],
) -> None:
    """Insert notification_queue documents into the test Mongo database."""
    client = MongoClient(mongo_uri)
    try:
        collection = client[database]['notification_queue']
        if notifications:
            collection.insert_many(notifications)
    finally:
        client.close()


def seed_orchestrated_schedules(
    mongo_uri: str,
    database: str,
    schedules: list[dict[str, Any]],
) -> None:
    """Insert orchestrated_schedules tracking documents."""
    client = MongoClient(mongo_uri)
    try:
        collection = client[database]['orchestrated_schedules']
        if schedules:
            collection.insert_many(schedules)
    finally:
        client.close()


def seed_opc_slots(redis_client: Redis, slots: dict[str, str]) -> None:
    """Write OPC slot keys (slot:opc_tags:*) in Redis."""
    for key, value in slots.items():
        redis_client.set(key, value)


def seed_active_ingestors(redis_client: Redis, ingestor_keys: list[str]) -> None:
    """Seed heartbeat:ingestor:* keys so load_active_ingestors returns ingestors."""
    for key in ingestor_keys:
        redis_client.set(key, '1')


def seed_last_timestamp(redis_client: Redis, mail_type: str, value: str) -> None:
    """
    Set notification_last_timestamp for a mail type, JSON-encoded.

    Values must be JSON-encoded so ``redis_repository.get`` (which calls
    ``json.loads`` on the raw payload) can deserialize them. The value must
    follow the exact format ``put_last_data_timestamp``
    (``orchestrator/activities/slot_manager.py``) writes into Redis today:
    ``DATETIME_FORMAT_MS_WITH_TZ`` (e.g. ``"2026-05-22 16:47:02.000000+0000"``)
    — microseconds included, no colon in the timezone offset. ``load_latest_data``
    parses this same format back into a ``datetime`` for the ``$gt`` comparison,
    so the two sides must agree exactly.

    Args:
        redis_client: Redis client connected to the test instance.
        mail_type: Mail type identifier (e.g. ``"Alerts"``, ``"Reports"``).
        value: Timestamp string in ``DATETIME_FORMAT_MS_WITH_TZ``
            (e.g. ``"2024-06-01 10:30:00.000000+0000"``).
    """
    redis_client.set(f'notification_last_timestamp:{mail_type}', json.dumps(value))


def seed_notification_cache(
    redis_client: Redis,
    trigger: str,
    notification_id: str,
    *,
    sent_at: str | None = None,
    ttl: int | None = None,
) -> None:
    """
    Pre-seed alerts sent cache entry, JSON-encoded.

    The value must be JSON-encoded because ``filter_notification_alerts``
    reads via ``redis_repository.get`` (which applies ``json.loads``) and
    parses the resulting string with ``DATETIME_FORMAT_MS_WITH_TZ``.

    Args:
        redis_client: Redis client connected to the test instance.
        trigger: Schedule/trigger name used to compose the cache key.
        notification_id: Notification id used to compose the cache key.
        sent_at: Optional timestamp string in ``DATETIME_FORMAT_MS_WITH_TZ``.
        ttl: Optional TTL in seconds for the cache entry.
    """
    key = f'{trigger}:{notification_id}'
    value = sent_at or datetime.now(UTC).strftime(DATETIME_FORMAT_MS_WITH_TZ)
    encoded = json.dumps(value)
    if ttl is not None:
        redis_client.set(key, encoded, ex=ttl)
    else:
        redis_client.set(key, encoded)


def count_log_report_rows(engine: Engine, mail_type: str | None = None) -> int:
    """Count rows in sientia_data.log_report, optionally filtered by mail_type."""
    sql = 'SELECT COUNT(*) FROM sientia_data.log_report'
    params: dict[str, Any] = {}
    if mail_type is not None:
        sql += ' WHERE mail_type = :mail_type'
        params['mail_type'] = mail_type
    with engine.connect() as conn:
        return int(conn.execute(text(sql), params).scalar() or 0)


def fetch_log_report(engine: Engine, mail_type: str | None = None) -> list[dict[str, Any]]:
    """Fetch log_report rows as dicts."""
    sql = 'SELECT * FROM sientia_data.log_report'
    params: dict[str, Any] = {}
    if mail_type is not None:
        sql += ' WHERE mail_type = :mail_type'
        params['mail_type'] = mail_type
    with engine.connect() as conn:
        rows = conn.execute(text(sql), params).mappings().all()
    return [dict(row) for row in rows]


def default_notification(
    *,
    notification_id: str,
    level: str = 'ERROR',
    timestamp: datetime | None = None,
    model_name: str = 'model-a',
    model_id: str = '1',
) -> dict[str, Any]:
    """
    Build a minimal notification_queue document mirroring production layout.

    The ``timestamp`` field is stored as a native, timezone-aware ``datetime``
    (BSON Date) because that is exactly what
    ``sientia_do.notifications.models.Notification`` writes into
    ``notification_queue`` in production since SIENTIAPDE-2007
    (``self.timestamp = now()`` — a real ``datetime``, no more string
    formatting). Tests intentionally use this same representation so we
    surface, rather than hide, real production behavior in downstream
    activities (e.g. ``put_last_data_timestamp``'s ``DataFrame(...)['timestamp']
    .max().strftime(...)``, which requires a real ``datetime``, not a string).

    Args:
        notification_id: Unique identifier for the notification.
        level: Notification level (e.g. ``"ERROR"``, ``"WARNING"``).
        timestamp: Optional timezone-aware datetime. ``None`` falls back to a
            fixed sample value.
        model_name: Model name attached to the notification.
        model_id: Model id attached to the notification.

    Return:
        dict[str, Any]: A notification document ready for insertion.
    """
    ts = timestamp if timestamp is not None else datetime(2024, 6, 1, 12, 0, 0, tzinfo=UTC)
    return {
        'notification_id': notification_id,
        'level': level,
        'timestamp': ts,
        'message': f'{level} on {model_name}',
        'trigger': 'test-schedule',
        'block': 'test-block',
        'pipeline': 'test-pipeline',
        'project': 'orchestrator-e2e',
        'model_name': model_name,
        'model_id': model_id,
    }


def default_receiver_group(
    *,
    group_name: str = 'admins',
    members: list[str] | None = None,
    levels: list[str] | None = None,
    contents: list[str] | None = None,
    ignore_models_id: list[str] | None = None,
    model_ids: list[str] | None = None,
) -> dict[str, Any]:
    """Minimal active receiver_groups document."""
    return {
        'group_name': group_name,
        'active': True,
        'members': members or ['admin@example.com'],
        'levels': levels or ['ERROR', 'WARNING', 'INFO'],
        'contents': contents or ['core_alerts', 'persistent_alerts', 'reports'],
        'ignore_models_id': ignore_models_id or [],
        'model_ids': model_ids or [],
    }


def extract_html_body(message: EmailMessage | Message) -> str:
    """
    Return the text/html portion of an email message as a decoded string.

    Walks every part looking for the first text/html payload, decoding it
    according to the part's transfer encoding and charset. Falls back to
    the message's own ``get_content``/raw payload when no HTML part is
    present so callers can still inspect plain-text reports. Shared between
    the Alerts and Reports e2e suites, both of which send mail through the
    same ``build_email_html``/``send_email`` activities.

    Args:
        message: Captured email message returned by the test SMTP server.

    Return:
        str: HTML body content, or an empty string when nothing decodable
            is found.
    """
    if message.is_multipart():
        for part in message.walk():
            if part.get_content_type() != 'text/html':
                continue
            payload = part.get_payload(decode=True)
            if payload is None:
                continue
            charset = part.get_content_charset() or 'utf-8'
            return payload.decode(charset, errors='replace')
    payload = message.get_payload(decode=True)
    if payload is not None:
        charset = message.get_content_charset() or 'utf-8'
        return payload.decode(charset, errors='replace')
    try:
        return message.get_content()
    except (AttributeError, KeyError):
        return str(message.get_payload() or '')