474c2ef42c
Enhance orchestration configuration and documentation. Added `RUNTIME` variable to `.env.example`, updated `.gitignore` to exclude `openspec/` and `.cursor/`, and modified `README.md` to clarify queue naming conventions and runtime handling. Refactored activities to use synchronous database and email handling, improving performance and consistency. Updated test cases to reflect these changes and ensure compatibility with new activity definitions.
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 [],
|
|
}
|