""" 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', 'connectors', 'connector_tags', ) 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[+-][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_tag_database( mongo_uri: str, database: str, connectors: list[dict[str, Any]], connector_tags: list[dict[str, Any]], ) -> None: """Insert tag database documents (connectors, connector_tags) into the test Mongo database.""" client = MongoClient(mongo_uri) try: if connectors: client[database]['connectors'].insert_many(connectors) if connector_tags: client[database]['connector_tags'].insert_many(connector_tags) 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 '')