vitor-aignosi
11e350a6ce
Add e2e tag-database scenarios and cover worker/email/formatter gaps - Add e2e scenarios 1.9.1-1.9.4 for PI Web API through the tag database: scouter and output tag resolution, mixed-connector schedule hold, and OPC slot unchanged by PI connectors. Seed connectors and connector_tags from a new scenario input. - Tighten the happy-path e2e to assert both schedules and that no PI output config is produced when the tag database is unused. - Drop dead `if self.server` / `if notification_handler` guards in the email reconnect and worker shutdown paths, since both are always set. - Log the traceback inside the message on worker failure, as the sientia_do Logger.error does not accept exc_info. - Add unit tests for the worker failure exit, Prometheus startup failure, email reconnect without a password, and formatter, mongo, temporal manager and orchestrator edge cases.
418 lines
14 KiB
Python
418 lines
14 KiB
Python
"""
|
|
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[+-]<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_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 '')
|