Code import - branch 0.6.0
This commit is contained in:
356
e2e/helpers.py
Normal file
356
e2e/helpers.py
Normal file
@@ -0,0 +1,356 @@
|
||||
"""
|
||||
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 [],
|
||||
}
|
||||
Reference in New Issue
Block a user