357 lines
11 KiB
Python
357 lines
11 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 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'
|
|
DATETIME_FORMAT_WITH_TZ = '%Y-%m-%d %H:%M:%S%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 ``sientia_do.notifications.models.Notification``
|
|
writes into ``notification_queue.timestamp``: ``DATETIME_FORMAT_WITH_TZ``
|
|
(e.g. ``"2026-05-22 16:47:02+0000"``) — no microseconds and no colon in
|
|
the timezone offset.
|
|
|
|
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_WITH_TZ``
|
|
(e.g. ``"2024-06-01 10:30:00+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: str | 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 string in ``DATETIME_FORMAT_WITH_TZ``
|
|
because that is exactly what ``sientia_do.notifications.models.Notification``
|
|
writes into ``notification_queue`` in production (``now().strftime(
|
|
DATETIME_FORMAT_WITH_TZ)``). Tests intentionally use this same format so we
|
|
surface, rather than hide, real production behavior in downstream
|
|
activities.
|
|
|
|
Args:
|
|
notification_id: Unique identifier for the notification.
|
|
level: Notification level (e.g. ``"ERROR"``, ``"WARNING"``).
|
|
timestamp: Optional production-format timestamp string. ``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
|
|
).strftime(DATETIME_FORMAT_WITH_TZ)
|
|
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: 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': ignore_models or [],
|
|
}
|