gitea_admin/sientia-dataops-orchestrator_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1646

Browse Source 
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.
This commit is contained in:
vitor-aignosi
2026-05-22 15:02:16 -03:00
parent 4d5231ca16
commit 474c2ef42c
57 changed files with 2744 additions and 385 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
.env.example
View File

@@ -29,4 +29,7 @@ PROJECT_NAME="sientia-orchestrator"
TEMPORAL_HOST="temporal-frontend.temporal.svc.cluster.local:7233"
TEMPORAL_NAMESPACE="default"
TEMPORAL_SCOUTER_NAMESPACE="scouter"
TEMPORAL_LABORIOUS_NAMESPACE="laborious"
TEMPORAL_LABORIOUS_NAMESPACE="laborious"
RUNTIME=local
# ACTIVITY_EXECUTOR_MAX_WORKERS=200

4
.github/workflows/quality-gate.yml vendored
View File

@@ -8,9 +8,9 @@ on:
jobs:
quality-gate:
uses: Aignosi/github_workflow_templates/.github/workflows/dataops-module-quality-gate.yml@main
uses: Aignosi/github_workflow_templates/.github/workflows/python-module-quality-gate.yml@main
permissions: write-all
with:
project_name: 'orchestrator'
repositories: 'sientia-dataops-library'
secrets: inherit
secrets: inherit

5
.gitignore vendored
View File

@@ -45,4 +45,7 @@ git_key*
git_log
.env
.env
openspec/
.cursor/

48
README.md
View File

@@ -390,16 +390,17 @@ The orchestrator includes an advanced notification filtering system that prevent
- **Purpose**: Main application orchestrator managing Temporal workers and task queues
- **Responsibilities**:
- Temporal client initialization and connection management with SDK metrics
- Worker lifecycle management and graceful shutdown
- Task queue configuration (orchestrator, alerts, reports) with dedicated workers
- Worker lifecycle management and graceful shutdown via `sientia_do.temporal.worker.prepare_worker`
- Task queue configuration (orchestrator, alerts, reports) with dedicated workers on `<workflow>-<runtime>-queue`
- Prometheus metrics server initialization on HTTP_METRICS_PORT
- Temporal SDK metrics server initialization on HTTP_SDK_METRICS_PORT
- Notification handler setup and configuration
- **Key Features**:
- Multi-queue worker management with three dedicated workers (orchestrator, alerts, reports)
- `RUNTIME` env var (default `legacy`) passed to every `prepare_worker` call
- Sync blocking activities run on `prepare_worker`'s `activity_executor` thread pool; workflows stay `async def`
- Application health metrics (app_up gauge) for Kubernetes liveness/readiness probes
- Graceful shutdown with cleanup procedures for all connections
- Comprehensive error handling and metrics collection
- Graceful shutdown with cleanup procedures for all connections and exit code propagation to Kubernetes
- Parallel worker execution using asyncio.gather
#### **Activities (`orchestrator/activities/`)**
@@ -524,6 +525,11 @@ python -m orchestrator.worker.worker
| `TEMPORAL_NAMESPACE` | Default Temporal namespace | `default` | No |
| `TEMPORAL_SCOUTER_NAMESPACE` | Scouter workflow namespace | `scouter` | No |
| `TEMPORAL_LABORIOUS_NAMESPACE` | Laborious workflow namespace | `laborious` | No |
| `RUNTIME` | Runtime slice for orchestrator worker task queues (`orchestrator-<runtime>-queue`, etc.) | `legacy` | No |
| `ACTIVITY_EXECUTOR_MAX_WORKERS` | Thread pool size for sync activities (all workers) | `200` | No |
| `ORCHESTRATOR_ACTIVITY_EXECUTOR_MAX_WORKERS` | Per-worker override for the Orchestrator worker | falls back to `ACTIVITY_EXECUTOR_MAX_WORKERS` | No |
| `ALERTS_ACTIVITY_EXECUTOR_MAX_WORKERS` | Per-worker override for the Alerts worker | falls back to `ACTIVITY_EXECUTOR_MAX_WORKERS` | No |
| `REPORTS_ACTIVITY_EXECUTOR_MAX_WORKERS` | Per-worker override for the Reports worker | falls back to `ACTIVITY_EXECUTOR_MAX_WORKERS` | No |
| `REDIS_HOST` | Redis server hostname | `localhost` | Yes |
| `REDIS_PORT` | Redis server port | `6379` | Yes |
| `REDIS_USERNAME` | Redis username | `default` | Yes |
@@ -673,16 +679,28 @@ The project maintains comprehensive test coverage including:
### Test Execution
```bash
# Install test dependencies
pip install pytest pytest-cov pytest-asyncio
pip install -r requirements-dev.txt
# Run tests with coverage
# Run unit tests with coverage (default testpaths=tests; E2E excluded)
pytest --cov=orchestrator --cov-report=html
# Run specific test modules
pytest tests/activities/test_mongo_db.py
pytest tests/workflows/test_orchestrator.py
pytest tests/orchestrator/activities/test_mongo_db.py
pytest tests/orchestrator/workflows/test_orchestrator.py
```
### End-to-end tests
E2E tests live in `e2e/` and require **Docker** (testcontainers). They are **not** collected by default `pytest` at the repo root.
```bash
source ./venv/bin/activate
pip install -r requirements-dev.txt
pytest e2e/ --override-ini testpaths=e2e -m e2e -v
```
Coverage is kept separate from unit tests: set `COVERAGE_FILE=.coverage.e2e` when measuring E2E coverage (see `e2e/README.md`).
## 🛡️ Code Quality & Validation
### Overview
@@ -814,6 +832,13 @@ orchestrator/
- Check MongoDB collection configurations
- Verify input data format and required fields
5. **E2E / testcontainers leftovers**
- If a run is interrupted, containers may keep running. List and remove them:
```bash
docker ps -a --filter label=org.testcontainers=true
docker rm -f $(docker ps -aq --filter label=org.testcontainers=true)
```
### Debug Mode
Enable debug logging by setting the log level:
@@ -867,4 +892,11 @@ For support and questions:
---
### Runtime queues and migration
- **Queue naming**: Downstream pipeline schedules target `<workflow_type>-<runtime>-queue` (for example `scouter-legacy-queue`). The orchestrator's own workers register on `orchestrator-<runtime>-queue`, `alerts-<runtime>-queue`, and `reports-<runtime>-queue`.
- **Pipeline `runtime`**: Each pipeline JSON may include `runtime` (default `legacy` via `common_config`). Schedules without `runtime` use `legacy`.
- **Schedule updates**: Changing a pipeline's `runtime` does not update an existing schedule's `task_queue` in place; the next orchestrator tick deletes and recreates the schedule on the new queue.
- **Rollout order**: Deploy scouter/laborious worker fleets bound to the new `<workflow_type>-<runtime>-queue` family before redeploying the orchestrator. On the first tick after upgrade, orphan schedules on old queue names are normalized away and schedules are recreated on the new queues (no in-place rename).
**Note**: The SIENTIA DataOps Orchestrator is designed for production use in enterprise data environments. Ensure proper security configuration and network isolation for production deployments.

35
e2e/README.md Normal file
View File

@@ -0,0 +1,35 @@
# Orchestrator end-to-end tests
End-to-end tests run every external dependency for real (MongoDB, Redis, PostgreSQL via testcontainers; SMTP via in-process `aiosmtpd`; Temporal via `WorkflowEnvironment.start_local()`).
## Requirements
- Docker (for testcontainers)
- Python dev dependencies: `pip install -r requirements-dev.txt`
## Run locally
```bash
source ./venv/bin/activate
pytest e2e/ --override-ini testpaths=e2e -m e2e -v
```
Stop on first failure:
```bash
pytest e2e/ --override-ini testpaths=e2e -m e2e -x
```
## Coverage (separate from unit tests)
```bash
COVERAGE_FILE=.coverage.e2e pytest e2e/ --override-ini testpaths=e2e -m e2e --cov=orchestrator --cov-branch
coverage combine .coverage .coverage.e2e
coverage report
```
Unit tests keep the default `.coverage` file; the E2E run must set `COVERAGE_FILE=.coverage.e2e` so reports do not overwrite each other.
## Scenario catalog
See [scenarios.md](scenarios.md) for numbered scenarios and which test module implements each case.

0
e2e/__init__.py Normal file
View File

446
e2e/conftest.py Normal file
View File

@@ -0,0 +1,446 @@
"""Pytest configuration and fixtures for orchestrator E2E tests."""
import asyncio
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from unittest.mock import MagicMock
import pytest
import pytest_asyncio
from google.protobuf.duration_pb2 import Duration
from pymongo import MongoClient
from redis import Redis
from sientia_do.notifications.handlers import CoreNotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.temporal.worker.prepare_worker import build_queue_name
from sqlalchemy import create_engine
from temporalio.api.enums.v1 import IndexedValueType
from temporalio.api.operatorservice.v1 import AddSearchAttributesRequest
from temporalio.api.workflowservice.v1 import (
DescribeNamespaceRequest,
RegisterNamespaceRequest,
)
from temporalio.client import Client
from temporalio.common import SearchAttributeKey
from temporalio.service import RPCError, RPCStatusCode
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from testcontainers.core.container import DockerContainer
from testcontainers.postgres import PostgresContainer
from e2e.helpers import MONGO_COLLECTIONS, ORCHESTRATOR_TASK_QUEUE
from e2e.smtp_test_server import SmtpTestServer
from e2e.stub_workflows import STUB_WORKFLOW_CLASSES
from orchestrator.activities.activities import Activities
from orchestrator.activities.formatters import schedule_types
from orchestrator.workflows.alerts import Alerts
from orchestrator.workflows.orchestrator import Orchestrator
from orchestrator.workflows.reports import Reports
from orchestrator.workflows.subworkflows.load_notification_package import (
LoadNotificationPackage,
)
from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications
DB_SCHEMA_SQL_PATH = Path(__file__).parent / 'db_schema.sql'
E2E_DATABASE = 'orchestrator_test'
E2E_RUNTIMES = ('legacy', 'gpu')
MANAGED_NAMESPACES = ('scouter', 'laborious')
E2E_SEARCH_ATTRIBUTES = [
SearchAttributeKey.for_keyword('model_id'),
SearchAttributeKey.for_keyword('model_name'),
SearchAttributeKey.for_keyword('orchestrated'),
]
E2E_NAMESPACE_SEARCH_ATTRIBUTES = {
'model_id': IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
'model_name': IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
'orchestrated': IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
}
async def register_namespace_if_missing(env: WorkflowEnvironment, namespace: str) -> None:
"""
Register a Temporal namespace on the local dev server and wait until it is ready.
Args:
env: Session WorkflowEnvironment from start_local().
namespace: Namespace name to register.
"""
service = env.client.service_client
try:
await service.workflow_service.register_namespace(
RegisterNamespaceRequest(
namespace=namespace,
workflow_execution_retention_period=Duration(seconds=86400),
)
)
except RPCError as err:
if err.status != RPCStatusCode.ALREADY_EXISTS:
raise
deadline = time.monotonic() + 5.0
while time.monotonic() < deadline:
try:
await service.workflow_service.describe_namespace(
DescribeNamespaceRequest(namespace=namespace)
)
return
except RPCError:
await asyncio.sleep(0.1)
raise TimeoutError(f'Namespace {namespace} not ready within 5s')
async def ensure_namespace_search_attributes(
env: WorkflowEnvironment, namespace: str
) -> None:
"""
Register the orchestrator search attributes on a namespace, if missing.
The local Temporal dev server only registers search attributes on the default
namespace at start time. Schedules created in additional namespaces fail with
"no mapping defined for search attribute ..." unless we explicitly add the
same attribute mappings to those namespaces via the operator service.
Args:
env: Session WorkflowEnvironment from start_local().
namespace: Namespace where attributes must be available.
"""
service = env.client.service_client
try:
await service.operator_service.add_search_attributes(
AddSearchAttributesRequest(
namespace=namespace,
search_attributes=dict(E2E_NAMESPACE_SEARCH_ATTRIBUTES),
)
)
except RPCError as err:
if err.status != RPCStatusCode.ALREADY_EXISTS:
raise
def temporal_host_from_env(env: WorkflowEnvironment) -> str:
"""Return target host:port for the in-process Temporal dev server."""
return env.client.service_client.config.target_host
def mongo_uri_from_container(mongo_container) -> str:
"""Build a Mongo connection string for the testcontainer."""
port = mongo_container.get_exposed_port(27017)
return f'mongodb://localhost:{port}'
@pytest_asyncio.fixture(scope='session')
def postgres_container():
"""PostgreSQL testcontainer used by all E2E tests."""
postgres = PostgresContainer('postgres:15')
postgres.start()
yield postgres
postgres.stop()
@pytest_asyncio.fixture(scope='session')
def mongo_container():
"""MongoDB testcontainer used by real CoreNotificationHandler."""
mongo = DockerContainer('mongo:7').with_exposed_ports(27017)
mongo.start()
yield mongo
mongo.stop()
@pytest_asyncio.fixture(scope='session')
def redis_container():
"""Redis testcontainer for slot and notification timestamp paths."""
redis = DockerContainer('redis:7').with_exposed_ports(6379)
redis.start()
yield redis
redis.stop()
@pytest_asyncio.fixture
def postgres_engine(postgres_container):
"""SQLAlchemy engine bound to the PostgreSQL testcontainer."""
engine = create_engine(postgres_container.get_connection_url())
yield engine
engine.dispose()
def _create_schema_and_tables(engine):
sql_text = DB_SCHEMA_SQL_PATH.read_text(encoding='utf-8')
with engine.begin() as conn:
conn.exec_driver_sql(sql_text)
@pytest_asyncio.fixture(autouse=True)
def setup_postgres_schema_and_tables(postgres_engine):
"""Recreate Postgres schema from e2e/db_schema.sql before each test."""
_create_schema_and_tables(postgres_engine)
yield
@pytest.fixture
def mongo_uri(mongo_container):
return mongo_uri_from_container(mongo_container)
@pytest_asyncio.fixture(autouse=True)
def reset_mongo_collections(mongo_uri):
"""Drop orchestrator-managed Mongo collections between tests."""
client = MongoClient(mongo_uri)
try:
db = client[E2E_DATABASE]
for name in MONGO_COLLECTIONS:
db[name].drop()
finally:
client.close()
yield
@pytest.fixture
def redis_client(redis_container):
"""Redis client bound to the testcontainer."""
port = int(redis_container.get_exposed_port(6379))
client = Redis(host='localhost', port=port, decode_responses=True)
yield client
client.close()
@pytest_asyncio.fixture(autouse=True)
def reset_redis(redis_client):
"""Flush Redis between tests."""
redis_client.flushdb()
yield
@pytest.fixture(scope='session')
def smtp_server():
"""Session-scoped in-process SMTP server."""
server = SmtpTestServer()
server.start()
yield server
server.stop()
@pytest.fixture(autouse=True)
def smtp_messages_cleanup(smtp_server):
"""Clear captured SMTP messages between tests."""
smtp_server.clear()
yield
@pytest_asyncio.fixture(scope='session')
async def temporal_env():
"""Real Temporal dev server (schedule APIs supported)."""
env = await WorkflowEnvironment.start_local(search_attributes=E2E_SEARCH_ATTRIBUTES)
for namespace in MANAGED_NAMESPACES:
await register_namespace_if_missing(env, namespace)
await ensure_namespace_search_attributes(env, namespace)
yield env
await env.shutdown()
@pytest.fixture(scope='session')
def temporal_host(temporal_env):
return temporal_host_from_env(temporal_env)
@pytest_asyncio.fixture(autouse=True)
async def cleanup_temporal_schedules(temporal_env):
"""Delete orphan schedules in scouter/laborious before each test."""
host = temporal_host_from_env(temporal_env)
for namespace in MANAGED_NAMESPACES:
client = await Client.connect(host, namespace=namespace)
async for schedule in await client.list_schedules():
handle = client.get_schedule_handle(schedule.id)
await handle.delete()
yield
@pytest_asyncio.fixture
def mock_logger():
"""Logger double with readable console output for E2E runs."""
logger = MagicMock(spec=Logger)
logger.info = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
logger.debug = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
logger.error = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
logger.warning = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
logger.custom_info = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
logger.custom_debug = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
logger.custom_error = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
logger.custom_warning = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
return logger
@pytest_asyncio.fixture
async def notification_handler(mock_logger, mongo_container):
"""Real notification handler using MongoDB testcontainer."""
handler = CoreNotificationHandler(
connection_string=mongo_uri_from_container(mongo_container),
database=E2E_DATABASE,
logger=mock_logger,
project_name='orchestrator-e2e',
)
try:
yield handler
finally:
handler.shutdown()
@pytest.fixture
def notification_inserts(notification_handler):
"""Spy on real Mongo insert calls issued by notification handler."""
collection = notification_handler.mongo_collection
original_insert_one = collection.insert_one
spy = MagicMock(wraps=original_insert_one)
collection.insert_one = spy
try:
yield spy
finally:
collection.insert_one = original_insert_one
@pytest_asyncio.fixture
async def test_activities(
postgres_container,
mongo_container,
redis_container,
smtp_server,
mock_logger,
notification_handler,
temporal_host,
):
"""Real Activities wired to testcontainers and in-process SMTP."""
mongo_port = mongo_container.get_exposed_port(27017)
redis_port = int(redis_container.get_exposed_port(6379))
activities = Activities(
temporal_config={
'temporal_host': temporal_host,
'temporal_scouter_namespace': 'scouter',
'temporal_laborious_namespace': 'laborious',
},
redis_config={
'host': 'localhost',
'port': redis_port,
'username': '',
'password': '',
},
mongodb_config={
'connection_string': f'mongodb://localhost:{mongo_port}',
'database_name': E2E_DATABASE,
'ttl_index_seconds': 3600,
},
email_config={
'sender_email': 'e2e@example.com',
'sender_password': '',
'smtp_server': smtp_server.host,
'smtp_port': smtp_server.port,
},
postgres_config={
'host': 'localhost',
'port': int(postgres_container.get_exposed_port(5432)),
'user': postgres_container.username,
'password': postgres_container.password,
'dbname': postgres_container.dbname,
'min_connections': 1,
'max_connections': 5,
},
logger=mock_logger,
notification_handler=notification_handler,
)
await activities.connect_to_temporal()
try:
yield activities
finally:
activities.shutdown()
def _orchestrator_activity_list(activities: Activities) -> list:
return [
activities.load_active_ingestors,
activities.load_opc_slots,
activities.update_slots,
activities.delete_slots,
activities.aggregate_documents_in_mongodb,
activities.find_documents_in_mongodb,
activities.update_pipelines_timestamps,
activities.create_pipelines_timestamps,
activities.delete_pipelines_timestamps,
activities.create_collection_with_ttl_index,
activities.create_schedules,
activities.update_schedules,
activities.delete_schedules,
activities.normalize_schedules,
activities.process_schedules,
activities.process_slots,
activities.create_schedule_config,
activities.create_slot_config,
activities.report_schedule_orchestration,
activities.report_slot_orchestration,
activities.format_schedule_config,
activities.get_last_data_timestamp,
activities.load_latest_data,
activities.put_last_data_timestamp,
activities.filter_notification_alerts,
activities.filter_notification_reports,
activities.build_email_html,
activities.send_email,
activities.format_log_report,
activities.export_data_to_postgres,
activities.store_notification_cache,
]
@pytest_asyncio.fixture
async def orchestrator_worker(temporal_env, test_activities):
"""Worker for orchestrator workflows and all activities on the default namespace."""
with ThreadPoolExecutor(max_workers=32) as activity_executor:
async with Worker(
temporal_env.client,
task_queue=ORCHESTRATOR_TASK_QUEUE,
workflows=[
Orchestrator,
Alerts,
Reports,
LoadNotificationPackage,
ProcessNotifications,
],
activities=_orchestrator_activity_list(test_activities),
activity_executor=activity_executor,
) as worker:
yield worker
@pytest_asyncio.fixture
async def stub_workers(temporal_env):
"""No-op workers on scouter/laborious namespaces for every managed workflow type."""
host = temporal_host_from_env(temporal_env)
worker_contexts: list[Worker] = []
clients: list[Client] = []
stub_types = list(schedule_types.keys()) + [
'xgboost_predictions_batch',
'xgboost_minimal_retrain',
]
try:
for namespace in MANAGED_NAMESPACES:
client = await Client.connect(host, namespace=namespace)
clients.append(client)
queues = {
build_queue_name(workflow_type, runtime)
for workflow_type in stub_types
for runtime in E2E_RUNTIMES
}
for queue in queues:
worker = Worker(
client,
task_queue=queue,
workflows=STUB_WORKFLOW_CLASSES,
)
await worker.__aenter__()
worker_contexts.append(worker)
yield worker_contexts
finally:
for worker in reversed(worker_contexts):
await worker.__aexit__(None, None, None)

33
e2e/db_schema.sql Normal file
View File

@@ -0,0 +1,33 @@
-- =============================================================================
-- E2E test database schema for the ``sientia_data`` namespace.
--
-- SINGLE SOURCE OF TRUTH: mirrors production DDL for tables the orchestrator
-- writes to. Any production DDL change must be pasted into this file (same
-- pattern as sientia-dataops-laborious_temporal/e2e/db_schema.sql).
-- =============================================================================
CREATE SCHEMA IF NOT EXISTS sientia_data;
-- -----------------------------------------------------------------------------
-- sientia_data.log_report
-- Written by ProcessNotifications via export_data_to_postgres.
-- The table is dropped between tests so each scenario starts with a clean
-- slate; the per-test autouse fixture re-runs this script.
-- -----------------------------------------------------------------------------
DROP TABLE IF EXISTS sientia_data.log_report;
CREATE TABLE sientia_data.log_report (
status text,
"timestamp" timestamptz,
groups text,
message text,
level text,
notification_id text,
block text,
schedule text,
pipeline text,
project text,
model_name text,
model_id text,
mail_type text,
created_at timestamptz DEFAULT CURRENT_TIMESTAMP
);

356
e2e/helpers.py Normal file
View 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 [],
}

5
e2e/scenario_inputs/alerts_duplicate.json Normal file
View File

@@ -0,0 +1,5 @@
{
"schedule_name": "alerts-e2e-dup",
"notification_ttl": 300,
"sent_ttl": 600
}

5
e2e/scenario_inputs/alerts_empty.json Normal file
View File

@@ -0,0 +1,5 @@
{
"schedule_name": "alerts-e2e-empty",
"notification_ttl": 300,
"sent_ttl": 600
}

5
e2e/scenario_inputs/alerts_happy_path.json Normal file
View File

@@ -0,0 +1,5 @@
{
"schedule_name": "alerts-e2e",
"notification_ttl": 300,
"sent_ttl": 600
}

5
e2e/scenario_inputs/alerts_persistent.json Normal file
View File

@@ -0,0 +1,5 @@
{
"schedule_name": "alerts-e2e-persistent",
"notification_ttl": 1,
"sent_ttl": 600
}

34
e2e/scenario_inputs/orchestrator_conflict.json Normal file
View File

@@ -0,0 +1,34 @@
{
"workflow_input": {
"schedule_name": "orchestrator-e2e-conflict",
"pipelines_query": {
"collection": "pipelines",
"aggregation": [{"$match": {"active": true}}]
},
"opc_servers_query": {
"collection": "opc_servers",
"filters": {"active": true}
}
},
"pipelines": [
{
"schedule_name": "conflict-pred",
"workflow_type": "predictions_batch",
"runtime": "legacy",
"model_id": "model-1",
"model": {"name": "Model model-1"},
"active": true,
"updated_at": "@now",
"frequency": "1m",
"offset": "0m",
"query": "SELECT 1",
"write_tags": [
{"server_id": "srv-1", "type": "prediction", "addr": "ns=2;s=Tag1"}
]
}
],
"opc_servers": [
{"id": "srv-1", "active": true}
],
"active_ingestors": ["heartbeat:ingestor:1"]
}

34
e2e/scenario_inputs/orchestrator_create_only.json Normal file
View File

@@ -0,0 +1,34 @@
{
"workflow_input": {
"schedule_name": "orchestrator-e2e-create",
"pipelines_query": {
"collection": "pipelines",
"aggregation": [{"$match": {"active": true}}]
},
"opc_servers_query": {
"collection": "opc_servers",
"filters": {"active": true}
}
},
"pipelines": [
{
"schedule_name": "create-only-pred",
"workflow_type": "predictions_batch",
"runtime": "legacy",
"model_id": "model-1",
"model": {"name": "Model model-1"},
"active": true,
"updated_at": "@now",
"frequency": "1m",
"offset": "0m",
"query": "SELECT 1",
"write_tags": [
{"server_id": "srv-1", "type": "prediction", "addr": "ns=2;s=Tag1"}
]
}
],
"opc_servers": [
{"id": "srv-1", "active": true}
],
"active_ingestors": ["heartbeat:ingestor:1"]
}

31
e2e/scenario_inputs/orchestrator_delete_only.json Normal file
View File

@@ -0,0 +1,31 @@
{
"workflow_input": {
"schedule_name": "orchestrator-e2e-delete",
"pipelines_query": {
"collection": "pipelines",
"aggregation": [{"$match": {"active": true}}]
},
"opc_servers_query": {
"collection": "opc_servers",
"filters": {"active": true}
}
},
"pipelines": [
{
"schedule_name": "delete-me",
"workflow_type": "drift",
"runtime": "legacy",
"model_id": "model-1",
"model": {"name": "Model model-1"},
"active": true,
"updated_at": "@now-1h",
"frequency": "1m",
"offset": "0m",
"interval_minutes": 60
}
],
"opc_servers": [
{"id": "srv-1", "active": true}
],
"active_ingestors": ["heartbeat:ingestor:1"]
}

16
e2e/scenario_inputs/orchestrator_empty.json Normal file
View File

@@ -0,0 +1,16 @@
{
"workflow_input": {
"schedule_name": "orchestrator-e2e-empty",
"pipelines_query": {
"collection": "pipelines",
"aggregation": [{"$match": {"active": true}}]
},
"opc_servers_query": {
"collection": "opc_servers",
"filters": {"active": true}
}
},
"pipelines": [],
"opc_servers": [],
"active_ingestors": []
}

46
e2e/scenario_inputs/orchestrator_happy_path.json Normal file
View File

@@ -0,0 +1,46 @@
{
"workflow_input": {
"schedule_name": "orchestrator-e2e",
"pipelines_query": {
"collection": "pipelines",
"aggregation": [{"$match": {"active": true}}]
},
"opc_servers_query": {
"collection": "opc_servers",
"filters": {"active": true}
}
},
"pipelines": [
{
"schedule_name": "pred-legacy",
"workflow_type": "predictions_batch",
"runtime": "legacy",
"model_id": "model-1",
"model": {"name": "Model model-1"},
"active": true,
"updated_at": "@now",
"frequency": "1m",
"offset": "0m",
"query": "SELECT 1",
"write_tags": [
{"server_id": "srv-1", "type": "prediction", "addr": "ns=2;s=Tag1"}
]
},
{
"schedule_name": "drift-gpu",
"workflow_type": "drift",
"runtime": "gpu",
"model_id": "model-2",
"model": {"name": "Model model-2"},
"active": true,
"updated_at": "@now",
"frequency": "1m",
"offset": "0m",
"interval_minutes": 60
}
],
"opc_servers": [
{"id": "srv-1", "active": true, "name": "opc-1"}
],
"active_ingestors": ["heartbeat:ingestor:1"]
}

41
e2e/scenario_inputs/orchestrator_noop.json Normal file
View File

@@ -0,0 +1,41 @@
{
"workflow_input": {
"schedule_name": "orchestrator-e2e-noop",
"pipelines_query": {
"collection": "pipelines",
"aggregation": [{"$match": {"active": true}}]
},
"opc_servers_query": {
"collection": "opc_servers",
"filters": {"active": true}
}
},
"pipelines": [
{
"schedule_name": "noop-pred",
"workflow_type": "predictions_batch",
"runtime": "legacy",
"model_id": "model-1",
"model": {"name": "Model model-1"},
"active": true,
"updated_at": "@now-1h",
"frequency": "1m",
"offset": "0m",
"query": "SELECT 1",
"write_tags": [
{"server_id": "srv-1", "type": "prediction", "addr": "ns=2;s=Tag1"}
]
}
],
"opc_servers": [
{"id": "srv-1", "active": true}
],
"active_ingestors": ["heartbeat:ingestor:1"],
"orchestrated_schedules": [
{
"schedule_name": "noop-pred",
"namespace": "laborious",
"updated_at": "@now-1h"
}
]
}

45
e2e/scenario_inputs/orchestrator_ttl_index.json Normal file
View File

@@ -0,0 +1,45 @@
{
"workflow_input": {
"schedule_name": "orchestrator-e2e-ttl",
"pipelines_query": {
"collection": "pipelines",
"aggregation": [{"$match": {"active": true}}]
},
"opc_servers_query": {
"collection": "opc_servers",
"filters": {"active": true}
}
},
"pipelines": [
{
"schedule_name": "scouter-ttl",
"workflow_type": "scouter",
"runtime": "legacy",
"model_id": "model-1",
"model": {"name": "Model model-1"},
"active": true,
"updated_at": "@now",
"frequency": "1m",
"offset": "0m",
"read_tags": [
{
"server_id": "srv-1",
"tag_name": "Read1",
"tag_address": "ns=2;s=Read1",
"aggr_func": "lts",
"frequency": 1000
}
]
}
],
"opc_servers": [
{
"id": "srv-1",
"active": true,
"server_name": "opc-1",
"url": "opc.tcp://localhost:4840",
"uri": "urn:opcfoundation:UA:DemoServer"
}
],
"active_ingestors": ["heartbeat:ingestor:1"]
}

41
e2e/scenario_inputs/orchestrator_update_only.json Normal file
View File

@@ -0,0 +1,41 @@
{
"workflow_input": {
"schedule_name": "orchestrator-e2e-update",
"pipelines_query": {
"collection": "pipelines",
"aggregation": [{"$match": {"active": true}}]
},
"opc_servers_query": {
"collection": "opc_servers",
"filters": {"active": true}
}
},
"pipelines": [
{
"schedule_name": "update-pred",
"workflow_type": "predictions_batch",
"runtime": "legacy",
"model_id": "model-1",
"model": {"name": "Model model-1"},
"active": true,
"updated_at": "@now",
"frequency": "5m",
"offset": "0m",
"query": "SELECT 1",
"write_tags": [
{"server_id": "srv-1", "type": "prediction", "addr": "ns=2;s=Tag1"}
]
}
],
"opc_servers": [
{"id": "srv-1", "active": true}
],
"active_ingestors": ["heartbeat:ingestor:1"],
"orchestrated_schedules": [
{
"schedule_name": "update-pred",
"namespace": "laborious",
"updated_at": "@now-1h"
}
]
}

3
e2e/scenario_inputs/reports_empty.json Normal file
View File

@@ -0,0 +1,3 @@
{
"schedule_name": "reports-e2e-empty"
}

3
e2e/scenario_inputs/reports_happy_path.json Normal file
View File

@@ -0,0 +1,3 @@
{
"schedule_name": "reports-e2e"
}

3
e2e/scenario_inputs/reports_multi_level.json Normal file
View File

@@ -0,0 +1,3 @@
{
"schedule_name": "reports-e2e-levels"
}

5
e2e/scenario_inputs/subworkflow_load_notification_package.json Normal file
View File

@@ -0,0 +1,5 @@
{
"schedule_name": "load-pkg-e2e",
"mail_type": "Alerts",
"base_data_filter": {"level": "ERROR"}
}

6
e2e/scenario_inputs/subworkflow_process_notifications.json Normal file
View File

@@ -0,0 +1,6 @@
{
"schedule_name": "process-notif-e2e",
"mail_type": "Alerts",
"schema": "sientia_data",
"table_name": "log_report"
}

111
e2e/scenarios.md Normal file
View File

@@ -0,0 +1,111 @@
# E2E Scenario Documentation — Orchestrator
Functional reference for orchestrator E2E scenarios. Tests live under `e2e/`, use `@pytest.mark.e2e`, and run with:
```bash
pytest e2e/ --override-ini testpaths=e2e -m e2e
```
## Execution context
- MongoDB, Redis, PostgreSQL: testcontainers (session-scoped).
- SMTP: in-process `aiosmtpd` (`e2e/smtp_test_server.py`).
- Temporal: `WorkflowEnvironment.start_local()` with stub workers on `scouter` / `laborious`.
- Production code under `orchestrator/**` is not mocked; only `Logger` may be a `MagicMock`.
---
## 1. Orchestrator workflow
Source: `e2e/test_orchestrator_main_workflow.py`
### 1.1.1 Happy path
Pipelines in Mongo → schedules created in correct namespace/task queue, Redis slots written, `orchestrated_schedules` updated.
### 1.2.1 No-op tick
Mongo, Redis, and Temporal already match desired state → no new schedules or slot writes.
### 1.3.1 Create-only
New pipeline only → schedules created, timestamps inserted.
### 1.3.2 Update-only
Existing pipeline with newer `updated_at` → schedule updated in Temporal.
### 1.3.3 Delete-only
Pipeline removed from Mongo → schedule deleted from Temporal.
### 1.4.1 Conflict ordering
Pipeline update and slot delete on same OPC server → slot insert before delete (production ordering).
### 1.5.1 Empty pipelines
No active pipelines → orphan schedules removed, no new orchestration writes.
### 1.6.1 TTL index bootstrap
First run creates TTL index on notification collection used by scouter pipelines.
---
## 2. Alerts workflow
Source: `e2e/test_alerts_main_workflow.py`
### A.1.1 Happy path
ERROR notification → one SMTP message, one `log_report` row, Redis cache key.
### A.1.2 TTL duplicate suppression
Second run with same data and cache seeded → no extra email or log row.
### A.1.3 Persistent escalation
Alert past `notification_ttl` with cache cleared → new email sent.
### A.2.1 Group filtering
Receiver group `levels` / `ignore_models` honored.
### A.3.1 Empty queue
No notifications → no SMTP, no Postgres row.
---
## 3. Reports workflow
Source: `e2e/test_reports_main_workflow.py`
### R.1.1 Happy path
Mixed ERROR/WARNING/INFO → one HTML email with all section headings.
### R.1.2 Per-level rendering
Single-level notifications → only matching section in HTML body.
### R.2.1 Empty queue
No notifications → no SMTP, no Postgres row.
---
## 4. Subworkflows
### LoadNotificationPackage — `e2e/test_subworkflow_load_notification_package.py`
- No prior Redis timestamp → all matching notifications returned, max timestamp stored.
- Prior timestamp → only newer notifications returned.
- Empty Mongo → no Redis timestamp write.
### ProcessNotifications — `e2e/test_subworkflow_process_notifications.py`
- Full round trip: HTML → SMTP → `log_report` in Postgres.
- Empty receiver groups → `{}`.

110
e2e/smtp_test_server.py Normal file
View File

@@ -0,0 +1,110 @@
"""In-process SMTP server for E2E email assertions."""
import asyncio
import socket
import threading
import time
from email import message_from_bytes, policy
from email.message import EmailMessage, Message
from aiosmtpd.controller import Controller
class _CaptureHandler:
"""
aiosmtpd handler that parses every incoming message into an ``EmailMessage``.
The default policy yields a ``Message`` instance, which loses structure
when wrapped into a new ``EmailMessage``. Parsing with ``policy.default``
keeps multipart payloads intact so tests can introspect the HTML body
via ``walk()``/``get_payload(decode=True)``.
"""
def __init__(self):
self.messages: list[EmailMessage | Message] = []
async def handle_DATA(self, server, session, envelope):
message = message_from_bytes(envelope.content, policy=policy.default)
self.messages.append(message)
return '250 OK'
def _reserve_port(host: str = '127.0.0.1') -> int:
"""Reserve a free TCP port on the given host."""
probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
probe.bind((host, 0))
port = probe.getsockname()[1]
probe.close()
return port
class SmtpTestServer:
"""
Wraps aiosmtpd Controller with a dedicated asyncio loop thread for pytest compatibility.
Attributes:
host: Bind host (127.0.0.1).
port: Listening port after start().
messages: Captured outbound messages.
"""
def __init__(self):
self.host = '127.0.0.1'
self.port: int | None = None
self._handler = _CaptureHandler()
self.messages: list[EmailMessage | Message] = self._handler.messages
self._controller: Controller | None = None
self._loop: asyncio.AbstractEventLoop | None = None
self._thread: threading.Thread | None = None
def start(self) -> None:
"""Start the SMTP controller on a reserved port in a background event loop."""
self.port = _reserve_port(self.host)
self._loop = asyncio.new_event_loop()
self._controller = Controller(
self._handler,
hostname=self.host,
port=self.port,
loop=self._loop,
ready_timeout=30,
)
def _run():
asyncio.set_event_loop(self._loop)
if self._controller is None:
raise RuntimeError('SMTP controller not initialized')
self._controller.start()
self._thread = threading.Thread(target=_run, name='e2e-smtp', daemon=True)
self._thread.start()
self._wait_until_ready()
def _wait_until_ready(self, timeout: float = 10.0) -> None:
"""Poll until the SMTP listener accepts TCP connections."""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
if self.port is None:
time.sleep(0.05)
continue
probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
if probe.connect_ex((self.host, self.port)) == 0:
return
finally:
probe.close()
time.sleep(0.05)
raise TimeoutError('SMTP test server did not become ready')
def stop(self) -> None:
"""Stop the controller and background event loop."""
if self._controller is not None:
self._controller.stop()
if self._thread is not None:
self._thread.join(timeout=5)
self._controller = None
self._loop = None
self._thread = None
def clear(self) -> None:
"""Remove all captured messages."""
self.messages.clear()

73
e2e/stub_workflows.py Normal file
View File

@@ -0,0 +1,73 @@
"""No-op Temporal workflows for managed scouter/laborious namespaces in E2E."""
from typing import Any
from temporalio import workflow
@workflow.defn(name='scouter')
class ScouterStub:
@workflow.run
async def run(self, _input_data: dict[str, Any]) -> None:
return None
@workflow.defn(name='pi_web_api_scouter')
class PiWebApiScouterStub:
@workflow.run
async def run(self, _input_data: dict[str, Any]) -> None:
return None
@workflow.defn(name='predictions_batch')
class PredictionsBatchStub:
@workflow.run
async def run(self, _input_data: dict[str, Any]) -> None:
return None
@workflow.defn(name='xgboost_predictions_batch')
class XgboostPredictionsBatchStub:
@workflow.run
async def run(self, _input_data: dict[str, Any]) -> None:
return None
@workflow.defn(name='drift')
class DriftStub:
@workflow.run
async def run(self, _input_data: dict[str, Any]) -> None:
return None
@workflow.defn(name='simple_metrics')
class SimpleMetricsStub:
@workflow.run
async def run(self, _input_data: dict[str, Any]) -> None:
return None
@workflow.defn(name='minimal_retrain')
class MinimalRetrainStub:
@workflow.run
async def run(self, _input_data: dict[str, Any]) -> None:
return None
@workflow.defn(name='xgboost_minimal_retrain')
class XgboostMinimalRetrainStub:
@workflow.run
async def run(self, _input_data: dict[str, Any]) -> None:
return None
STUB_WORKFLOW_CLASSES = [
ScouterStub,
PiWebApiScouterStub,
PredictionsBatchStub,
XgboostPredictionsBatchStub,
DriftStub,
SimpleMetricsStub,
MinimalRetrainStub,
XgboostMinimalRetrainStub,
]

116
e2e/test_alerts_main_workflow.py Normal file
View File

@@ -0,0 +1,116 @@
"""E2E tests for the Alerts main workflow."""
from datetime import UTC, datetime
import pytest
from redis import Redis
from temporalio.testing import WorkflowEnvironment
from e2e.conftest import E2E_DATABASE
from e2e.helpers import (
DATETIME_FORMAT_MS_WITH_TZ,
count_log_report_rows,
default_notification,
default_receiver_group,
fetch_log_report,
load_scenario_input,
make_workflow_id,
seed_notification_cache,
seed_notifications,
seed_receiver_groups,
start_and_await_workflow,
)
from orchestrator.workflows.alerts import Alerts
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_scenario_a_1_1_happy_path(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
mongo_uri: str,
smtp_server,
postgres_engine,
redis_client: Redis,
):
"""A.1.1: ERROR alert sends email, writes log_report, caches notification."""
seed_receiver_groups(mongo_uri, E2E_DATABASE, [default_receiver_group()])
seed_notifications(
mongo_uri,
E2E_DATABASE,
[default_notification(notification_id='alert-1', level='ERROR')],
)
input_data = load_scenario_input('alerts_happy_path.json')
await start_and_await_workflow(
temporal_env.client,
Alerts.run,
input_data,
make_workflow_id('alerts-happy'),
)
assert len(smtp_server.messages) == 1
assert count_log_report_rows(postgres_engine, 'Alerts') == 1
rows = fetch_log_report(postgres_engine, 'Alerts')
assert rows[0]['mail_type'] == 'Alerts'
assert redis_client.get('test-schedule:alert-1') is not None or redis_client.keys('*alert-1*')
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_scenario_a_1_2_duplicate_suppressed(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
mongo_uri: str,
smtp_server,
postgres_engine,
redis_client: Redis,
):
"""A.1.2: Cached notification is not emailed twice within sent_ttl."""
seed_receiver_groups(mongo_uri, E2E_DATABASE, [default_receiver_group()])
notif = default_notification(notification_id='dup-1', level='ERROR')
seed_notifications(mongo_uri, E2E_DATABASE, [notif])
seed_notification_cache(
redis_client,
notif['trigger'],
notif['notification_id'],
sent_at=datetime.now(UTC).strftime(DATETIME_FORMAT_MS_WITH_TZ),
ttl=600,
)
input_data = load_scenario_input('alerts_duplicate.json')
await start_and_await_workflow(
temporal_env.client,
Alerts.run,
input_data,
make_workflow_id('alerts-dup'),
)
assert len(smtp_server.messages) == 0
assert count_log_report_rows(postgres_engine, 'Alerts') == 0
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_scenario_a_3_1_empty_queue(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
mongo_uri: str,
smtp_server,
postgres_engine,
):
"""A.3.1: Empty notification queue short-circuits."""
seed_receiver_groups(mongo_uri, E2E_DATABASE, [default_receiver_group()])
await start_and_await_workflow(
temporal_env.client,
Alerts.run,
load_scenario_input('alerts_empty.json'),
make_workflow_id('alerts-empty'),
)
assert len(smtp_server.messages) == 0
assert count_log_report_rows(postgres_engine, 'Alerts') == 0

226
e2e/test_orchestrator_main_workflow.py Normal file
View File

@@ -0,0 +1,226 @@
"""E2E tests for the Orchestrator main workflow."""
import pytest
from pymongo import MongoClient
from redis import Redis
from sientia_do.temporal.worker.prepare_worker import build_queue_name
from temporalio.client import Client, Schedule, ScheduleActionStartWorkflow, ScheduleSpec
from temporalio.testing import WorkflowEnvironment
from e2e.conftest import E2E_DATABASE
from e2e.helpers import (
load_scenario_input,
make_workflow_id,
seed_active_ingestors,
seed_opc_servers,
seed_orchestrated_schedules,
seed_pipelines,
start_and_await_workflow,
)
from orchestrator.workflows.orchestrator import Orchestrator
async def _schedule_ids(host: str, namespace: str) -> list[str]:
"""Return the list of Temporal schedule ids in a given namespace."""
client = await Client.connect(host, namespace=namespace)
return [schedule.id async for schedule in await client.list_schedules()]
def _apply_orchestrator_seeds(
scenario: dict,
mongo_uri: str,
redis_client: Redis,
) -> None:
"""
Seed Mongo and Redis with the pipelines/opc_servers/ingestors declared in a scenario.
Args:
scenario: Scenario payload returned by ``load_scenario_input``.
mongo_uri: Mongo connection string for the test database.
redis_client: Redis client connected to the test instance.
"""
seed_pipelines(mongo_uri, E2E_DATABASE, scenario.get('pipelines', []))
seed_opc_servers(mongo_uri, E2E_DATABASE, scenario.get('opc_servers', []))
seed_active_ingestors(redis_client, scenario.get('active_ingestors', []))
if scenario.get('orchestrated_schedules'):
seed_orchestrated_schedules(
mongo_uri,
E2E_DATABASE,
scenario['orchestrated_schedules'],
)
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_scenario_1_1_1_happy_path(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
mongo_uri: str,
redis_client: Redis,
temporal_host: str,
):
"""1.1.1: Creates schedules, slots, and orchestrated_schedules entries."""
scenario = load_scenario_input('orchestrator_happy_path')
_apply_orchestrator_seeds(scenario, mongo_uri, redis_client)
await start_and_await_workflow(
temporal_env.client,
Orchestrator.run,
scenario['workflow_input'],
make_workflow_id('orchestrator-happy'),
)
mongo = MongoClient(mongo_uri)
try:
tracked = list(mongo[E2E_DATABASE]['orchestrated_schedules'].find())
names = {doc['schedule_name'] for doc in tracked}
assert names, f'Expected orchestrated_schedules rows, got {tracked}'
finally:
mongo.close()
laborious_schedules = await _schedule_ids(temporal_host, 'laborious')
assert 'pred-legacy' in laborious_schedules or 'drift-gpu' in laborious_schedules, (
f'Expected Temporal schedules in laborious, got {laborious_schedules}'
)
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_scenario_1_3_1_create_only(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
mongo_uri: str,
temporal_host: str,
redis_client: Redis,
):
"""1.3.1: New pipeline creates a Temporal schedule."""
scenario = load_scenario_input('orchestrator_create_only')
_apply_orchestrator_seeds(scenario, mongo_uri, redis_client)
await start_and_await_workflow(
temporal_env.client,
Orchestrator.run,
scenario['workflow_input'],
make_workflow_id('orchestrator-create'),
)
assert 'create-only-pred' in await _schedule_ids(temporal_host, 'laborious')
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_scenario_1_3_3_delete_only(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
mongo_uri: str,
temporal_host: str,
redis_client: Redis,
):
"""1.3.3: Removing pipeline deletes Temporal schedule."""
scenario = load_scenario_input('orchestrator_delete_only')
_apply_orchestrator_seeds(scenario, mongo_uri, redis_client)
await start_and_await_workflow(
temporal_env.client,
Orchestrator.run,
scenario['workflow_input'],
make_workflow_id('orchestrator-seed-delete'),
)
client = await Client.connect(temporal_host, namespace='laborious')
try:
handle = client.get_schedule_handle('delete-me')
await handle.describe()
schedule_exists = True
except Exception:
schedule_exists = False
assert schedule_exists
mongo = MongoClient(mongo_uri)
try:
mongo[E2E_DATABASE]['pipelines'].delete_many({})
finally:
mongo.close()
await start_and_await_workflow(
temporal_env.client,
Orchestrator.run,
scenario['workflow_input'],
make_workflow_id('orchestrator-delete'),
)
assert 'delete-me' not in await _schedule_ids(temporal_host, 'laborious')
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_scenario_1_5_1_empty_pipelines(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
mongo_uri: str,
temporal_host: str,
redis_client: Redis,
):
"""1.5.1: No pipelines → no orchestrated_schedules documents."""
scenario = load_scenario_input('orchestrator_empty')
_apply_orchestrator_seeds(scenario, mongo_uri, redis_client)
client = await Client.connect(temporal_host, namespace='laborious')
await client.create_schedule(
'orphan-schedule',
Schedule(
action=ScheduleActionStartWorkflow(
'drift',
{},
id='orphan-schedule-run',
task_queue=build_queue_name('drift', 'legacy'),
),
spec=ScheduleSpec(),
),
)
await start_and_await_workflow(
temporal_env.client,
Orchestrator.run,
scenario['workflow_input'],
make_workflow_id('orchestrator-empty'),
)
mongo = MongoClient(mongo_uri)
try:
assert mongo[E2E_DATABASE]['orchestrated_schedules'].count_documents({}) == 0
finally:
mongo.close()
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_scenario_1_6_1_ttl_index_bootstrap(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
mongo_uri: str,
redis_client: Redis,
):
"""1.6.1: Scouter pipeline triggers TTL index on notification_queue."""
scenario = load_scenario_input('orchestrator_ttl_index')
_apply_orchestrator_seeds(scenario, mongo_uri, redis_client)
await start_and_await_workflow(
temporal_env.client,
Orchestrator.run,
scenario['workflow_input'],
make_workflow_id('orchestrator-ttl'),
)
mongo = MongoClient(mongo_uri)
try:
indexes = mongo[E2E_DATABASE]['raw_scouter-ttl'].index_information()
assert any('expireAfterSeconds' in info for info in indexes.values())
finally:
mongo.close()

167
e2e/test_reports_main_workflow.py Normal file
View File

@@ -0,0 +1,167 @@
"""E2E tests for the Reports main workflow."""
from email.message import EmailMessage, Message
import pytest
from temporalio.testing import WorkflowEnvironment
from e2e.conftest import E2E_DATABASE
from e2e.helpers import (
count_log_report_rows,
default_notification,
default_receiver_group,
load_scenario_input,
make_workflow_id,
seed_notifications,
seed_receiver_groups,
start_and_await_workflow,
)
from orchestrator.workflows.reports import Reports
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.
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 '')
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_scenario_r_1_1_happy_path(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
mongo_uri: str,
smtp_server,
postgres_engine,
):
"""R.1.1: Mixed-level notifications produce one email with all sections."""
seed_receiver_groups(
mongo_uri,
E2E_DATABASE,
[default_receiver_group(contents=['reports'], levels=['ERROR', 'WARNING', 'INFO'])],
)
seed_notifications(
mongo_uri,
E2E_DATABASE,
[
default_notification(notification_id='r-err', level='ERROR', model_name='m1'),
default_notification(
notification_id='r-warn',
level='WARNING',
model_name='m2',
timestamp='2024-06-01 12:01:00+0000',
),
default_notification(
notification_id='r-info',
level='INFO',
model_name='m3',
timestamp='2024-06-01 12:02:00+0000',
),
],
)
await start_and_await_workflow(
temporal_env.client,
Reports.run,
load_scenario_input('reports_happy_path.json'),
make_workflow_id('reports-happy'),
)
assert len(smtp_server.messages) == 1
body = _extract_html_body(smtp_server.messages[0])
assert 'Errors detected:' in body
assert 'Warnings detected:' in body
assert 'Infos detected:' in body
# format_log_report writes one row per (notification_id, trigger) pair, so the
# three seeded notifications produce three rows even though a single email
# was sent to the receiver group.
assert count_log_report_rows(postgres_engine, 'Reports') == 3
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_scenario_r_1_2_error_section_only(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
mongo_uri: str,
smtp_server,
):
"""R.1.2: Only ERROR notifications → only Errors section in HTML."""
seed_receiver_groups(
mongo_uri,
E2E_DATABASE,
[default_receiver_group(contents=['reports'], levels=['ERROR'])],
)
seed_notifications(
mongo_uri,
E2E_DATABASE,
[default_notification(notification_id='only-err', level='ERROR')],
)
await start_and_await_workflow(
temporal_env.client,
Reports.run,
load_scenario_input('reports_multi_level.json'),
make_workflow_id('reports-error-only'),
)
assert smtp_server.messages, 'Expected at least one report email'
body = _extract_html_body(smtp_server.messages[0])
assert 'Errors detected:' in body
assert 'Warnings detected:' not in body
assert 'Infos detected:' not in body
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_scenario_r_2_1_empty_queue(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
mongo_uri: str,
smtp_server,
postgres_engine,
):
"""R.2.1: Empty queue → no email and no log_report row."""
seed_receiver_groups(mongo_uri, E2E_DATABASE, [default_receiver_group()])
await start_and_await_workflow(
temporal_env.client,
Reports.run,
load_scenario_input('reports_empty.json'),
make_workflow_id('reports-empty'),
)
assert len(smtp_server.messages) == 0
assert count_log_report_rows(postgres_engine, 'Reports') == 0

154
e2e/test_subworkflow_load_notification_package.py Normal file
View File

@@ -0,0 +1,154 @@
"""E2E tests for LoadNotificationPackage subworkflow."""
import pytest
from redis import Redis
from temporalio.testing import WorkflowEnvironment
from e2e.conftest import E2E_DATABASE
from e2e.helpers import (
default_notification,
default_receiver_group,
load_scenario_input,
make_workflow_id,
seed_last_timestamp,
seed_notifications,
seed_receiver_groups,
start_and_await_workflow,
)
from orchestrator.workflows.subworkflows.load_notification_package import LoadNotificationPackage
def _build_metadata(input_data: dict) -> None:
"""Attach the metadata block that the subworkflow expects."""
input_data['metadata'] = {
'metadata': {
'schedule_name': input_data['schedule_name'],
'workflow_name': 'load_notification_package',
'model_name': '-',
'model_id': '-',
}
}
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_load_package_without_prior_timestamp(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
mongo_uri: str,
redis_client: Redis,
):
"""
No Redis timestamp → returns notifications and stores max timestamp.
Notifications are seeded with the exact production format produced by
``sientia_do.notifications.models.Notification`` (string in
``DATETIME_FORMAT_WITH_TZ``, e.g. ``"2024-06-01 11:00:00+0000"``). The
cached "last timestamp" must mirror that representation.
"""
seed_receiver_groups(mongo_uri, E2E_DATABASE, [default_receiver_group()])
seed_notifications(
mongo_uri,
E2E_DATABASE,
[
default_notification(
notification_id='n1',
timestamp='2024-06-01 10:00:00+0000',
),
default_notification(
notification_id='n2',
timestamp='2024-06-01 11:00:00+0000',
),
],
)
input_data = load_scenario_input('subworkflow_load_notification_package.json')
_build_metadata(input_data)
result = await start_and_await_workflow(
temporal_env.client,
LoadNotificationPackage.run,
input_data,
make_workflow_id('load-pkg-none'),
)
assert len(result['notification_package']) == 2
stored = redis_client.get('notification_last_timestamp:Alerts') or ''
assert stored.strip('"') == '2024-06-01 11:00:00+0000'
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_load_package_with_prior_timestamp(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
mongo_uri: str,
redis_client: Redis,
):
"""
Existing timestamp → only newer notifications returned.
Both the seeded "last timestamp" (Redis) and the notification timestamps
(Mongo) follow the production format used by
``sientia_do.notifications.models.Notification`` (string in
``DATETIME_FORMAT_WITH_TZ``). ``load_latest_data`` will translate the
Redis value into a Python ``datetime`` and apply ``{$gt: <Date>}``
against the Mongo string timestamps; this exercise surfaces the real
BSON comparison semantics rather than a synthetic ideal.
"""
seed_receiver_groups(mongo_uri, E2E_DATABASE, [default_receiver_group()])
seed_last_timestamp(redis_client, 'Alerts', '2024-06-01 10:30:00+0000')
seed_notifications(
mongo_uri,
E2E_DATABASE,
[
default_notification(
notification_id='old',
timestamp='2024-06-01 10:00:00+0000',
),
default_notification(
notification_id='new',
timestamp='2024-06-01 11:00:00+0000',
),
],
)
input_data = load_scenario_input('subworkflow_load_notification_package.json')
_build_metadata(input_data)
result = await start_and_await_workflow(
temporal_env.client,
LoadNotificationPackage.run,
input_data,
make_workflow_id('load-pkg-ts'),
)
ids = {n['notification_id'] for n in result['notification_package']}
assert ids == {'new'}
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_load_package_empty_mongo_no_redis_write(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
mongo_uri: str,
redis_client: Redis,
):
"""Empty Mongo → no Redis timestamp write."""
seed_receiver_groups(mongo_uri, E2E_DATABASE, [default_receiver_group()])
input_data = load_scenario_input('subworkflow_load_notification_package.json')
_build_metadata(input_data)
await start_and_await_workflow(
temporal_env.client,
LoadNotificationPackage.run,
input_data,
make_workflow_id('load-pkg-empty'),
)
assert redis_client.get('notification_last_timestamp:Alerts') is None

94
e2e/test_subworkflow_process_notifications.py Normal file
View File

@@ -0,0 +1,94 @@
"""E2E tests for ProcessNotifications subworkflow."""
import pytest
from temporalio.testing import WorkflowEnvironment
from e2e.helpers import (
count_log_report_rows,
default_notification,
fetch_log_report,
load_scenario_input,
make_workflow_id,
start_and_await_workflow,
)
from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications
def _receiver_package():
notif = default_notification(notification_id='proc-1', level='ERROR')
return {
'admins': {
'group_name': 'admins',
'members': ['admin@example.com'],
'status': 'pending',
'notifications': [notif],
}
}
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_process_notifications_round_trip(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
smtp_server,
postgres_engine,
):
"""Email → SMTP → log_report persisted in Postgres."""
input_data = load_scenario_input('subworkflow_process_notifications.json')
input_data['metadata'] = {
'metadata': {
'schedule_name': input_data['schedule_name'],
'workflow_name': 'process_notifications',
'model_name': '-',
'model_id': '-',
}
}
input_data['notification_package'] = _receiver_package()
result = await start_and_await_workflow(
temporal_env.client,
ProcessNotifications.run,
input_data,
make_workflow_id('process-notif'),
)
assert result
assert len(smtp_server.messages) == 1
assert count_log_report_rows(postgres_engine, 'Alerts') == 1
rows = fetch_log_report(postgres_engine, 'Alerts')
assert rows[0]['notification_id'] == 'proc-1'
@pytest.mark.e2e
@pytest.mark.asyncio
async def test_process_notifications_empty_package(
temporal_env: WorkflowEnvironment,
orchestrator_worker,
stub_workers,
smtp_server,
postgres_engine,
):
"""Empty receiver groups returns empty dict."""
input_data = load_scenario_input('subworkflow_process_notifications.json')
input_data['metadata'] = {
'metadata': {
'schedule_name': input_data['schedule_name'],
'workflow_name': 'process_notifications',
'model_name': '-',
'model_id': '-',
}
}
input_data['notification_package'] = {}
result = await start_and_await_workflow(
temporal_env.client,
ProcessNotifications.run,
input_data,
make_workflow_id('process-empty'),
)
assert result == {}
assert len(smtp_server.messages) == 0
assert count_log_report_rows(postgres_engine) == 0

2
orchestrator/activities/activities.py
View File

@@ -6,7 +6,7 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.temporal.activities.postgres import Postgres
from sientia_do.temporal.activities.postgres_sync import Postgres
from orchestrator.activities.email import Email
from orchestrator.activities.formatters import Formatters

4
orchestrator/activities/email.py
View File

@@ -84,7 +84,7 @@ class Email(SientiaMonitoring):
self.close()
@activity.defn(name='build_email_html')
async def build_email_html(self, input_data: dict[str, Any]) -> dict[str, Any]:
def build_email_html(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Build HTML email content for configured receiver groups.
@@ -193,7 +193,7 @@ class Email(SientiaMonitoring):
self.server.sendmail(self.sender_email, receivers, msg.as_string())
@activity.defn(name='send_email')
async def send_email(self, input_data: dict[str, Any]) -> dict[str, Any]:
def send_email(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Send email notifications to configured receiver groups.

56
orchestrator/activities/formatters.py
View File

@@ -2,10 +2,10 @@ from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import json
from collections.abc import Callable, Hashable
from collections.abc import Callable
from logging import Logger
from math import ceil
from typing import Any, TypedDict
from typing import Any, TypedDict, cast
from pandas import DataFrame
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
@@ -152,7 +152,7 @@ class Formatters(SientiaMonitoring):
self.close()
@activity.defn(name='process_schedules')
async def process_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]:
def process_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Process pipeline configurations into Temporal-compatible schedule configurations.
@@ -209,7 +209,7 @@ class Formatters(SientiaMonitoring):
return schedule_config
@activity.defn(name='process_slots')
async def process_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
def process_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Extracts all read tags from input pipelines, divides them into slots and
returns a slot config dictionary. If no ingestor is available, only one slot
@@ -254,7 +254,7 @@ class Formatters(SientiaMonitoring):
slot_config[f'{i}'], notifications = build_tag_config(slot_tags, opc_servers)
if notifications:
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id='ORCHESTRATOR_SERVER_NOT_FOUND_DURING_SLOT_CONFIGURATION',
message=f'Servers {", ".join(notifications)} not found in opc_servers',
@@ -267,7 +267,7 @@ class Formatters(SientiaMonitoring):
slot_tags = tags[last_index:]
slot_config[f'{number_of_slots}'], notifications = build_tag_config(slot_tags, opc_servers)
if notifications:
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id='ORCHESTRATOR_SERVER_NOT_FOUND_DURING_SLOT_CONFIGURATION',
message=f'Servers {", ".join(notifications)} not found in opc_servers',
@@ -281,7 +281,7 @@ class Formatters(SientiaMonitoring):
return slot_config
@activity.defn(name='format_schedule_config')
async def format_schedule_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
def format_schedule_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Format the schedule config to a dictionary with the schedule name as the key.
@@ -358,7 +358,7 @@ class Formatters(SientiaMonitoring):
to_create[namespace][schedule_name] = schedule
@activity.defn(name='create_schedule_config')
async def create_schedule_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
def create_schedule_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Create a schedule config dictionary based on the input data.
@@ -418,7 +418,7 @@ class Formatters(SientiaMonitoring):
return output
@activity.defn(name='create_slot_config')
async def create_slot_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
def create_slot_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Create a slot config dictionary based on the input data.
@@ -459,7 +459,7 @@ class Formatters(SientiaMonitoring):
return output
async def send_success_report(
def send_success_report(
self,
metadata: dict[str, Any],
message: str,
@@ -478,7 +478,7 @@ class Formatters(SientiaMonitoring):
notification_id (str): The ID of the notification to send
attachment (Any | None, optional): Optional attachment content to include
"""
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id=notification_id,
message=message,
@@ -487,7 +487,7 @@ class Formatters(SientiaMonitoring):
attachment_content=json.dumps(attachment, indent=4, sort_keys=True),
)
async def send_error_report(
def send_error_report(
self, metadata: dict[str, Any], message: str, notification_id: str, attachment: str
) -> None:
"""
@@ -502,7 +502,7 @@ class Formatters(SientiaMonitoring):
notification_id (str): The ID of the notification
attachment (str): The attachment content for the notification
"""
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id=notification_id,
message=message,
@@ -569,7 +569,7 @@ class Formatters(SientiaMonitoring):
return success_keys, error_keys
async def manage_and_send_report(
def manage_and_send_report(
self,
metadata: dict[str, Any],
success_keys: list[str],
@@ -591,7 +591,7 @@ class Formatters(SientiaMonitoring):
schedule_data (dict[str, Any]): The schedule data containing items and notification ID
"""
if len(success_keys) > 0:
await self.send_success_report(
self.send_success_report(
metadata=metadata,
message=f'Successfully {schedule_type}: \n {", ".join(success_keys)}',
notification_id=schedule_data['id'],
@@ -606,7 +606,7 @@ class Formatters(SientiaMonitoring):
else:
attachment.append(f'{key}:\n{value["message"]}')
await self.send_error_report(
self.send_error_report(
metadata=metadata,
message=f'Fails on {schedule_type}: \n {", ".join(error_keys)}',
notification_id=f'{schedule_data["id"]}_ERROR',
@@ -614,7 +614,7 @@ class Formatters(SientiaMonitoring):
)
@activity.defn(name='report_schedule_orchestration')
async def report_schedule_orchestration(self, input_data: dict[str, Any]) -> None:
def report_schedule_orchestration(self, input_data: dict[str, Any]) -> None:
"""
Report the orchestration result to the notification handler.
@@ -655,7 +655,7 @@ class Formatters(SientiaMonitoring):
if len(schedule_data['items']) > 0:
success_keys, error_keys = self.parse_report_schedule(schedule_data['items'])
await self.manage_and_send_report(
self.manage_and_send_report(
metadata=metadata,
success_keys=success_keys,
error_keys=error_keys,
@@ -664,7 +664,7 @@ class Formatters(SientiaMonitoring):
)
@activity.defn(name='report_slot_orchestration')
async def report_slot_orchestration(self, input_data: dict[str, Any]) -> None:
def report_slot_orchestration(self, input_data: dict[str, Any]) -> None:
"""
Report the slot orchestration result to the notification handler.
@@ -689,14 +689,14 @@ class Formatters(SientiaMonitoring):
success_keys, error_keys = self.parse_report(inserted_slots)
if len(success_keys) > 0:
await self.send_success_report(
self.send_success_report(
metadata=metadata,
message=f'Inserted slots: \n {", ".join(success_keys)}',
notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS',
)
if len(error_keys) > 0:
await self.send_error_report(
self.send_error_report(
metadata=metadata,
message=f'Failed to insert slots: \n {", ".join(error_keys)}',
notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS',
@@ -707,14 +707,14 @@ class Formatters(SientiaMonitoring):
success_keys, error_keys = self.parse_report(deleted_slots)
if len(success_keys) > 0:
await self.send_success_report(
self.send_success_report(
metadata=metadata,
message=f'Deleted slots: \n {", ".join(success_keys)}',
notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS',
)
if len(error_keys) > 0:
await self.send_error_report(
self.send_error_report(
metadata=metadata,
message=f'Failed to delete slots: \n {", ".join(error_keys)}',
notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS',
@@ -722,7 +722,7 @@ class Formatters(SientiaMonitoring):
)
@activity.defn(name='format_log_report')
async def format_log_report(self, input_data: dict[str, Any]) -> dict[Hashable, Any]:
def format_log_report(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Format the receiver_groups status to a dataframe to be stored in the database.
@@ -736,7 +736,7 @@ class Formatters(SientiaMonitoring):
- metadata (dict): Metadata for logging purposes
Returns:
dict[Hashable, Any]: The formatted log report as a dictionary representation of a DataFrame
dict[str, Any]: The formatted log report as a dictionary representation of a DataFrame
"""
metadata = input_data['metadata']
mail_type = input_data['mail_type']
@@ -776,10 +776,12 @@ class Formatters(SientiaMonitoring):
data_values: DataFrame = DataFrame(list(data.values()))
return data_values.to_dict()
# ``DataFrame.to_dict()`` is typed as ``dict[Hashable, Any]`` in pandas
# stubs, but default orientation uses column names (str keys).
return cast(dict[str, Any], data_values.to_dict())
@activity.defn(name='filter_notification_reports')
async def filter_notification_reports(self, input_data: dict[str, Any]) -> dict[str, Any]:
def filter_notification_reports(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Filter notifications for comprehensive scheduled reports.

57
orchestrator/activities/mongo_db.py
View File

@@ -9,10 +9,9 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.sientia_monitoring import MetricsController, SientiaMonitoring
from sientia_do.repository.mongodb_repository import MongoDBRepository
from sientia_do.repository.mongodb_repository_sync import MongoDBRepository
from sientia_do.temporal.constants import (
DATETIME_FORMAT_MS_WITH_TZ,
DATETIME_FORMAT_WITH_TZ,
now,
)
@@ -84,7 +83,7 @@ class MongoDB(SientiaMonitoring):
@activity.defn(
name='find_documents_in_mongodb',
)
async def find_documents_in_mongodb(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
def find_documents_in_mongodb(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
"""
Find documents in a MongoDB collection based on the provided query parameters.
@@ -113,7 +112,7 @@ class MongoDB(SientiaMonitoring):
)
try:
documents = await self.mongo_db_repository.find(collection_name, filters, metadata)
documents = self.mongo_db_repository.find(collection_name, filters, metadata)
self.info(
f"Loaded {len(documents)} documents from collection '{collection_name}'",
@@ -135,7 +134,7 @@ class MongoDB(SientiaMonitoring):
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id='MONGODB_QUERY_ERROR',
message=f'Failed to execute MongoDB query: {e}',
@@ -148,7 +147,7 @@ class MongoDB(SientiaMonitoring):
raise e
@activity.defn(name='aggregate_documents_in_mongodb')
async def aggregate_documents_in_mongodb(
def aggregate_documents_in_mongodb(
self, input_data: dict[str, Any]
) -> list[dict[str, Any]]:
"""
@@ -181,7 +180,7 @@ class MongoDB(SientiaMonitoring):
)
try:
aggregated_documents = await self.mongo_db_repository.aggregate(
aggregated_documents = self.mongo_db_repository.aggregate(
collection_name, aggregation, metadata
)
@@ -205,7 +204,7 @@ class MongoDB(SientiaMonitoring):
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id='MONGODB_AGGREGATION_ERROR',
message=f'Failed to execute MongoDB aggregation: {e}',
@@ -218,7 +217,7 @@ class MongoDB(SientiaMonitoring):
raise e
@activity.defn(name='update_pipelines_timestamps')
async def update_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
def update_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
"""
Update `updated_at` timestamps for successfully updated pipelines.
@@ -244,12 +243,12 @@ class MongoDB(SientiaMonitoring):
data_filter = {'$or': argument} if argument else {}
try:
await self.mongo_db_repository.update_many(
self.mongo_db_repository.update_many(
'orchestrated_schedules', data_filter, {'$set': {'updated_at': date_now}}, metadata
)
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id='MONGODB_UPDATE_PIPELINES_ERROR',
message=f'Failed to update pipelines timestamps: {e}',
@@ -266,7 +265,7 @@ class MongoDB(SientiaMonitoring):
)
@activity.defn(name='create_pipelines_timestamps')
async def create_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
def create_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
"""
Insert `updated_at` timestamps for newly created pipelines.
@@ -298,12 +297,12 @@ class MongoDB(SientiaMonitoring):
try:
if data_filter:
await self.mongo_db_repository.insert_many(
self.mongo_db_repository.insert_many(
'orchestrated_schedules', data_filter, metadata
)
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id='MONGODB_CREATE_PIPELINES_ERROR',
message=f'Failed to create pipelines timestamps: {e}',
@@ -320,7 +319,7 @@ class MongoDB(SientiaMonitoring):
)
@activity.defn(name='delete_pipelines_timestamps')
async def delete_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
def delete_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
"""
Delete timestamp rows for successfully deleted pipelines.
@@ -345,12 +344,12 @@ class MongoDB(SientiaMonitoring):
data_filter = {'$or': argument} if argument else {}
try:
await self.mongo_db_repository.delete_many(
self.mongo_db_repository.delete_many(
'orchestrated_schedules', data_filter, metadata
)
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id='MONGODB_DELETE_PIPELINES_ERROR',
message=f'Failed to delete pipelines timestamps: {e}',
@@ -367,7 +366,7 @@ class MongoDB(SientiaMonitoring):
)
@activity.defn(name='create_collection_with_ttl_index')
async def create_collection_with_ttl_index(self, input_data: dict[str, Any]) -> None:
def create_collection_with_ttl_index(self, input_data: dict[str, Any]) -> None:
"""
Create collections with TTL indexes for pipeline topics.
@@ -424,7 +423,7 @@ class MongoDB(SientiaMonitoring):
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id='MONGODB_CREATE_COLLECTION_ERROR',
message=f'Failed to create collection {collection} with TTL index: {e}',
@@ -444,7 +443,7 @@ class MongoDB(SientiaMonitoring):
self.debug(f'Created indexes: {created_indexes}', metadata=metadata)
@activity.defn(name='load_latest_data')
async def load_latest_data(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
def load_latest_data(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
"""
Load the latest data from MongoDB collection since a specified timestamp.
@@ -477,16 +476,24 @@ class MongoDB(SientiaMonitoring):
if last_data_timestamp is None:
data_filter = base_data_filter
else:
# ``notification_queue.timestamp`` is stored as a string in
# ``DATETIME_FORMAT_WITH_TZ`` (``Notification`` writes it as
# ``now().strftime(DATETIME_FORMAT_WITH_TZ)``). Coercing
# ``last_data_timestamp`` to ``datetime`` here would force a
# BSON ``String`` vs ``Date`` comparison, which always yields
# ``False`` (``String < Date`` in BSON sort order) and breaks
# incremental loading entirely. Comparing strings preserves the
# intended chronological filter because the format is
# lexicographically ordered when the timezone is fixed
# (``Notification.timestamp`` always uses UTC).
data_filter = {
**base_data_filter,
'timestamp': {
'$gt': datetime.strptime(last_data_timestamp, DATETIME_FORMAT_WITH_TZ)
},
'timestamp': {'$gt': last_data_timestamp},
}
self.debug(f'Data filter: {data_filter}', metadata=metadata)
data = await self.mongo_db_repository.find(collection_name, data_filter, metadata)
data = self.mongo_db_repository.find(collection_name, data_filter, metadata)
self.debug(f'Collected: {data}', metadata=metadata)
@@ -497,7 +504,7 @@ class MongoDB(SientiaMonitoring):
return data
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id='MONGO_LOAD_ERROR',
message=f'Error loading data from MongoDB: {e}',

44
orchestrator/activities/slot_manager.py
View File

@@ -11,7 +11,7 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.sientia_monitoring import MetricsController, SientiaMonitoring
from sientia_do.repository.redis_repository import RedisRepository
from sientia_do.repository.redis_repository_sync import RedisRepository
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
@@ -84,7 +84,7 @@ class SlotManager(SientiaMonitoring):
self.close()
@activity.defn(name='load_opc_slots')
async def load_opc_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
def load_opc_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Load all OPC slots from Redis for current system state assessment.
@@ -112,7 +112,7 @@ class SlotManager(SientiaMonitoring):
opc_slots = {}
try:
slot_keys = await self.redis_repository.keys('slot:opc_tags:*')
slot_keys = self.redis_repository.keys('slot:opc_tags:*')
self.debug(f'Slot keys: {slot_keys}', metadata=metadata)
@@ -123,10 +123,10 @@ class SlotManager(SientiaMonitoring):
decoded_keys = slot_keys
for key in decoded_keys:
opc_slots[key] = await self.redis_repository.get(key)
opc_slots[key] = self.redis_repository.get(key)
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id='REDIS_GET_ERROR',
message=f'Failed to load OPC slots: {e}',
@@ -142,7 +142,7 @@ class SlotManager(SientiaMonitoring):
return opc_slots
@activity.defn(name='load_active_ingestors')
async def load_active_ingestors(self, input_data: dict[str, Any]) -> list[str]:
def load_active_ingestors(self, input_data: dict[str, Any]) -> list[str]:
"""
Load all active ingestors from Redis.
@@ -161,7 +161,7 @@ class SlotManager(SientiaMonitoring):
self.info('Loading active ingestors...', metadata=metadata)
try:
active_ingestors = await self.redis_repository.keys('heartbeat:ingestor:*')
active_ingestors = self.redis_repository.keys('heartbeat:ingestor:*')
self.info(f'Loaded {len(active_ingestors)} active ingestors', metadata=metadata)
@@ -178,7 +178,7 @@ class SlotManager(SientiaMonitoring):
return ingestors
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id='REDIS_GET_ERROR',
message=f'Failed to load active ingestors: {e}',
@@ -190,7 +190,7 @@ class SlotManager(SientiaMonitoring):
raise e
@activity.defn(name='update_slots')
async def update_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
def update_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Update OPC slots in Redis
@@ -213,7 +213,7 @@ class SlotManager(SientiaMonitoring):
for slot in to_insert:
try:
await self.redis_repository.set(f'slot:opc_tags:{slot}', to_insert[slot], ttl=None)
self.redis_repository.set(f'slot:opc_tags:{slot}', to_insert[slot], ttl=None)
report[slot] = {'success': True, 'message': 'Slot updated successfully'}
success_count += 1
except Exception as e:
@@ -227,7 +227,7 @@ class SlotManager(SientiaMonitoring):
return report
@activity.defn(name='delete_slots')
async def delete_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
def delete_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Delete OPC slots from Redis
@@ -251,7 +251,7 @@ class SlotManager(SientiaMonitoring):
for slot in to_delete:
try:
await self.redis_repository.delete(f'slot:opc_tags:{slot}')
self.redis_repository.delete(f'slot:opc_tags:{slot}')
report[slot] = {'success': True, 'message': 'Slot deleted successfully'}
success_count += 1
except Exception as e:
@@ -265,7 +265,7 @@ class SlotManager(SientiaMonitoring):
return report
@activity.defn(name='get_last_data_timestamp')
async def get_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None:
def get_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None:
"""
Get the last data timestamp from Redis.
@@ -285,9 +285,9 @@ class SlotManager(SientiaMonitoring):
key = f'notification_last_timestamp:{input_data["mail_type"]}'
try:
data_hold = await self.redis_repository.get(key)
data_hold = self.redis_repository.get(key)
except Exception as e:
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id='REDIS_GET_ERROR',
message=f'Error getting last data timestamp: {e}',
@@ -305,7 +305,7 @@ class SlotManager(SientiaMonitoring):
return data_hold
@activity.defn(name='put_last_data_timestamp')
async def put_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None:
def put_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None:
"""
Store the last data timestamp in Redis.
@@ -336,9 +336,9 @@ class SlotManager(SientiaMonitoring):
self.debug(f'Last collected timestamp to insert: {last_data_timestamp}', metadata=metadata)
try:
await self.redis_repository.set(key, last_data_timestamp, ttl=60 * 60 * 5)
self.redis_repository.set(key, last_data_timestamp, ttl=60 * 60 * 5)
except Exception as e:
await self.send_notification_async(
self.send_notification(
metadata=metadata,
notification_id='REDIS_SET_ERROR',
message=f'Error setting last data timestamp: {e}',
@@ -351,7 +351,7 @@ class SlotManager(SientiaMonitoring):
return last_data_timestamp
@activity.defn(name='filter_notification_alerts')
async def filter_notification_alerts(self, input_data: dict[str, Any]) -> dict[str, Any]:
def filter_notification_alerts(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Filter notification alerts with intelligent TTL-based duplicate prevention.
@@ -407,7 +407,7 @@ class SlotManager(SientiaMonitoring):
# Check if notification was recently sent
key = f'{notification["trigger"]}:{notification_id}'
last_sent = await self.redis_repository.get(key)
last_sent = self.redis_repository.get(key)
if last_sent is None:
alert_type = 'core_alerts'
@@ -438,7 +438,7 @@ class SlotManager(SientiaMonitoring):
return receiver_groups
@activity.defn(name='store_notification_cache')
async def store_notification_cache(self, input_data: dict[str, Any]) -> None:
def store_notification_cache(self, input_data: dict[str, Any]) -> None:
"""
Store notification cache in Redis to track recently sent notifications.
@@ -463,6 +463,6 @@ class SlotManager(SientiaMonitoring):
status = row['status']
if status == 'sent':
key = f'{row["schedule"]}:{row["notification_id"]}'
await self.redis_repository.set(key, date_now, ttl=sent_ttl)
self.redis_repository.set(key, date_now, ttl=sent_ttl)
self.info('Notification cache stored...', metadata=metadata)

9
orchestrator/activities/temporal_manager.py
View File

@@ -20,6 +20,7 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.sientia_monitoring import MetricsController, SientiaMonitoring
from sientia_do.temporal.worker.prepare_worker import build_queue_name
from orchestrator.utils.converters import parse_frequency
@@ -220,7 +221,9 @@ class TemporalManager(SientiaMonitoring):
workflow_type,
schedule,
id=schedule_name,
task_queue=f'{workflow_type}-queue',
task_queue=build_queue_name(
workflow_type, schedule.get('runtime', 'legacy')
),
execution_timeout=timedelta(seconds=execution_timeout_seconds),
run_timeout=timedelta(seconds=execution_timeout_seconds),
task_timeout=timedelta(seconds=task_timeout_seconds),
@@ -278,6 +281,10 @@ class TemporalManager(SientiaMonitoring):
"""
Update schedules in Temporal.
Does not modify ``task_queue``. The ``ScheduleUpdate`` callback only patches
workflow ``args`` and schedule ``intervals``. A ``runtime`` change requires
delete-then-create on the next orchestrator tick.
Args:
- input_data (dict[str, Any]): The input data containing
the schedules to update.

1
orchestrator/utils/orchestrator_functions.py
View File

@@ -39,6 +39,7 @@ def common_config(config: dict[str, Any]):
'execution_timeout_seconds': config.get('execution_timeout_seconds', 300),
'task_timeout_seconds': config.get('task_timeout_seconds', 300),
'on_conflict': config.get('on_conflict', 'error'),
'runtime': config.get('runtime', 'legacy'),
}

72
orchestrator/worker/prepare_worker.py
View File

@@ -1,72 +0,0 @@
import os
import re
from collections.abc import Sequence
from typing import Any
from sientia_do.observability.logger import Logger
from temporalio.client import Client
from temporalio.worker import PollerBehaviorAutoscaling, Worker
parameters = [
('MAX_CONCURRENT_WORKFLOW_TASKS', '200'),
('MAX_CONCURRENT_ACTIVITIES', '200'),
('MAX_CONCURRENT_LOCAL_ACTIVITIES', '200'),
('MAX_CACHED_WORKFLOWS', '200'),
('WORKFLOW_POLLER_BEHAVIOUR_MINIMUM', '10'),
('WORKFLOW_POLLER_BEHAVIOUR_INITIAL', '100'),
('WORKFLOW_POLLER_BEHAVIOUR_MAXIMUM', '200'),
('ACTIVITY_POLLER_BEHAVIOUR_MINIMUM', '10'),
('ACTIVITY_POLLER_BEHAVIOUR_INITIAL', '100'),
('ACTIVITY_POLLER_BEHAVIOUR_MAXIMUM', '200'),
]
def camel_to_snake(text: str) -> str:
"""Convert camelCase or PascalCase to snake_case."""
text = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', text)
text = re.sub('([a-z0-9])([A-Z])', r'\1_\2', text)
return text.lower()
def prepare_worker(
main_workflow: type,
other_workflows: Sequence[type],
activities: Sequence[Any],
temporal_client: Client,
logger: Logger,
) -> Worker:
main_workflow_name = main_workflow.__name__.upper()
queue_name = f'{camel_to_snake(main_workflow.__name__)}-queue'
local_workflow_parameters = {}
for parameter in parameters:
local_workflow_parameters[parameter[0]] = int(
os.getenv(main_workflow_name + '_' + parameter[0], parameter[1])
)
logger.info(f'Preparing worker for {main_workflow_name} with queue {queue_name}')
return Worker(
temporal_client,
task_queue=queue_name,
workflows=[main_workflow, *other_workflows],
activities=[*activities],
max_concurrent_workflow_tasks=local_workflow_parameters['MAX_CONCURRENT_WORKFLOW_TASKS'],
max_concurrent_activities=local_workflow_parameters['MAX_CONCURRENT_ACTIVITIES'],
max_concurrent_local_activities=local_workflow_parameters[
'MAX_CONCURRENT_LOCAL_ACTIVITIES'
],
max_cached_workflows=local_workflow_parameters['MAX_CACHED_WORKFLOWS'],
workflow_task_poller_behavior=PollerBehaviorAutoscaling(
minimum=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIOUR_MINIMUM'],
initial=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIOUR_INITIAL'],
maximum=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIOUR_MAXIMUM'],
),
activity_task_poller_behavior=PollerBehaviorAutoscaling(
minimum=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIOUR_MINIMUM'],
initial=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIOUR_INITIAL'],
maximum=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIOUR_MAXIMUM'],
),
)

17
orchestrator/worker/worker.py
View File

@@ -9,6 +9,7 @@ with workflow.unsafe.imports_passed_through():
from prometheus_client import start_http_server
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import get_logger
from sientia_do.temporal.worker.prepare_worker import prepare_worker
from orchestrator import metrics
from orchestrator.activities.activities import Activities
@@ -19,7 +20,6 @@ with workflow.unsafe.imports_passed_through():
build_redis_config,
build_temporal_config,
)
from orchestrator.worker.prepare_worker import prepare_worker
from orchestrator.workflows.alerts import Alerts
from orchestrator.workflows.orchestrator import Orchestrator
from orchestrator.workflows.reports import Reports
@@ -136,7 +136,7 @@ async def main():
activities.report_slot_orchestration,
activities.format_schedule_config,
],
logger=logger,
logger=logger
),
prepare_worker(
temporal_client=temporal_client,
@@ -158,7 +158,7 @@ async def main():
# Store notification cache
activities.store_notification_cache,
],
logger=logger,
logger=logger
),
prepare_worker(
temporal_client=temporal_client,
@@ -178,7 +178,7 @@ async def main():
activities.format_log_report,
activities.export_data_to_postgres,
],
logger=logger,
logger=logger
),
]
@@ -188,20 +188,19 @@ async def main():
logger.custom_info('Workers started successfully', metadata=metadata)
exit_code = 0
try:
# This will run the workers and wait for them to complete.
# If an exception occurs in any of the worker handlers, it will be propagated here.
await asyncio.gather(*handlers)
except BaseException as e:
logger.error(f'An unhandled exception occurred: {e}', exc_info=True)
exit_code = 1
finally:
if notification_handler:
notification_handler.shutdown()
if activities:
activities.shutdown()
# Exit with a non-zero status code to indicate failure to Kubernetes
metrics.APP_UP.labels(pod_id=POD_ID).set(0) # Mark app as DOWN
sys.exit(1)
metrics.APP_UP.labels(pod_id=POD_ID).set(0)
sys.exit(exit_code)
def start_prometheus_server():

4
orchestrator/workflows/subworkflows/process_notifications.py
View File

@@ -4,7 +4,7 @@ with workflow.unsafe.imports_passed_through():
from datetime import timedelta
from typing import Any
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.policies import retry_policy
from orchestrator.activities.activities import Activities
@@ -88,7 +88,7 @@ class ProcessNotifications:
'data': log_report,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_MS_WITH_TZ,
'format': DATETIME_FORMAT_WITH_TZ,
},
},
schedule_to_close_timeout=timedelta(seconds=60),

6
pyproject.toml
View File

@@ -59,6 +59,11 @@ ignore = [
"S105", # hardcoded passwords ok in tests
"S106", # hardcoded passwords ok in tests
]
"e2e/**/*.py" = [
"S101", # assert allowed in tests
"S105", # hardcoded passwords ok in tests
"S106", # hardcoded passwords ok in tests
]
[tool.ruff.lint.mccabe]
max-complexity = 15
@@ -120,6 +125,7 @@ markers = [
"asyncio: marks tests as async",
"integration: marks tests as integration tests",
"unit: marks tests as unit tests",
"e2e: end-to-end tests requiring Docker (testcontainers + Temporal local server)",
]
[tool.coverage.run]

5
requirements-dev.txt
View File

@@ -13,6 +13,11 @@ types-requests>=2.31.0 # Type stubs for requests
pytest>=7.4.0 # Testing framework
pytest-cov>=4.1.0 # Coverage plugin for pytest
pytest-asyncio>=0.21.0 # Async test support (already in main requirements)
testcontainers[postgres,mongodb]>=4.0.0
aiosmtpd>=1.4.0
sqlalchemy>=2.0.0
pymongo>=4.6.0
redis>=5.0.0
# Development Tools
ipython>=8.12.0 # Enhanced Python shell

8
requirements-local.txt Normal file
View File

@@ -0,0 +1,8 @@
temporalio
psycopg2-binary
sqlalchemy
redis
pymongo
jinja2
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.12.1
prometheus-client

2
requirements.txt
View File

@@ -4,5 +4,5 @@ sqlalchemy
redis
pymongo
jinja2
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.8.2
sientia_do>=1.12.1
prometheus-client

2
tests/orchestrator/activities/test_activities.py
View File

@@ -12,7 +12,7 @@ from orchestrator.activities.temporal_manager import TemporalManager
@patch('orchestrator.activities.slot_manager.SlotManager.__init__')
@patch('orchestrator.activities.formatters.Formatters.__init__')
@patch('orchestrator.activities.email.Email.__init__')
@patch('sientia_do.temporal.activities.postgres.Postgres.__init__')
@patch('sientia_do.temporal.activities.postgres_sync.Postgres.__init__')
@patch('orchestrator.activities.activities.MetricsController')
def test___init__(
mock_metrics_controller,

20
tests/orchestrator/activities/test_email.py
View File

@@ -1,7 +1,7 @@
from smtplib import SMTPServerDisconnected
from unittest.mock import AsyncMock, MagicMock, call, patch
from pytest import fixture, mark
from pytest import fixture
from orchestrator.activities.email import Email
@@ -19,9 +19,6 @@ def email(smtplib, email_builder):
notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
)
email_builder.send_notification_async = AsyncMock()
email_builder.send_notification = MagicMock()
email.send_notification_async = AsyncMock()
email.send_notification = MagicMock()
email.emit_metric = AsyncMock()
@@ -91,8 +88,7 @@ def test_close(sientia_monitoring_mock, email):
sientia_monitoring_mock.shutdown.assert_called_once()
@mark.asyncio
async def test_build_email_html(email):
def test_build_email_html(email):
email.email_builder.build_email = MagicMock(return_value='test')
input_data = {
**metadata,
@@ -107,7 +103,7 @@ async def test_build_email_html(email):
'mail_type': 'test',
}
response = await email.build_email_html(input_data)
response = email.build_email_html(input_data)
assert response == {
'group_1': {
@@ -253,19 +249,17 @@ def test_try_send_email_reconnect_quit_failure(smtp, email):
raise AssertionError('Expected exception')
@mark.asyncio
async def test_send_email_without_smtp_server(email):
def test_send_email_without_smtp_server(email):
email.smtp_server = None
input_data = {**metadata, 'receiver_groups': {}, 'mail_type': 'test_TYPE'}
response = await email.send_email(input_data)
response = email.send_email(input_data)
assert response == {}
@mark.asyncio
@patch('orchestrator.activities.email.MIMEText')
@patch('orchestrator.activities.email.MIMEMultipart')
async def test_send_email(mimemultipart, mimetext, email):
def test_send_email(mimemultipart, mimetext, email):
side_effect_1 = MagicMock()
side_effect_2 = MagicMock()
mimemultipart.side_effect = [side_effect_1, side_effect_2]
@@ -297,7 +291,7 @@ async def test_send_email(mimemultipart, mimetext, email):
'mail_type': 'test_TYPE',
}
response = await email.send_email(input_data)
response = email.send_email(input_data)
assert response['group_1']['status'] == 'sent'
assert response['group_2']['status'] == 'failed'

69
tests/orchestrator/activities/test_formatters.py
View File

@@ -2,7 +2,7 @@ import json
from unittest.mock import AsyncMock, MagicMock, call, patch
from pandas import DataFrame
from pytest import fixture, mark
from pytest import fixture
from sientia_do.notifications.models import NotificationLevel
from orchestrator.activities.formatters import Formatters
@@ -19,7 +19,6 @@ def formatters():
)
formatters.send_notification = MagicMock()
formatters.send_notification_async = AsyncMock()
formatters.emit_metric = AsyncMock()
formatters.error = MagicMock()
formatters.info = MagicMock()
@@ -37,8 +36,7 @@ metadata = {
}
@mark.asyncio
async def test_process_schedules(formatters):
def test_process_schedules(formatters):
mock_scouter = MagicMock(return_value={'test_scouter': 'test_scouter'})
mock_predictions_batch = MagicMock(
return_value={'test_predictions_batch': 'test_predictions_batch'}
@@ -115,7 +113,7 @@ async def test_process_schedules(formatters):
}
with patch('orchestrator.activities.formatters.schedule_types', mock_schedule_types):
result = await formatters.process_schedules(input_data)
result = formatters.process_schedules(input_data)
assert result == {
'scouter': {
@@ -148,8 +146,7 @@ async def test_process_schedules(formatters):
mock_simple_metrics.assert_called_once_with(input_data['pipelines'][4])
@mark.asyncio
async def test_process_schedules_with_invalid_workflow_type(formatters):
def test_process_schedules_with_invalid_workflow_type(formatters):
"""Test that process_schedules handles invalid workflow types correctly"""
mock_scouter = MagicMock(return_value={'test_scouter': 'test_scouter'})
@@ -193,7 +190,7 @@ async def test_process_schedules_with_invalid_workflow_type(formatters):
}
with patch('orchestrator.activities.formatters.schedule_types', mock_schedule_types):
result = await formatters.process_schedules(input_data)
result = formatters.process_schedules(input_data)
# Assert that error was called for invalid workflow type
formatters.error.assert_called_once_with(
@@ -221,7 +218,6 @@ async def test_process_schedules_with_invalid_workflow_type(formatters):
mock_scouter.assert_any_call(pipelines[2])
@mark.asyncio
@patch(
'orchestrator.activities.formatters.gather_read_tags',
return_value={
@@ -249,7 +245,7 @@ async def test_process_schedules_with_invalid_workflow_type(formatters):
},
)
@patch('orchestrator.activities.formatters.build_tag_config')
async def test_process_slots(mock_build_tag_config, mock_gather_read_tags, formatters):
def test_process_slots(mock_build_tag_config, mock_gather_read_tags, formatters):
input_data = {
'opc_servers': [
{'id': '1', 'server_name': 'test_server_name', 'url': 'test_url', 'uri': 'test_uri'},
@@ -289,7 +285,7 @@ async def test_process_slots(mock_build_tag_config, mock_gather_read_tags, forma
mock_build_tag_config.return_value = (slot_mock, ['2'])
result = await formatters.process_slots(input_data)
result = formatters.process_slots(input_data)
tags = list(mock_gather_read_tags.return_value.values())
@@ -302,7 +298,7 @@ async def test_process_slots(mock_build_tag_config, mock_gather_read_tags, forma
]
)
formatters.send_notification_async.assert_has_calls(
formatters.send_notification.assert_has_calls(
[
call(
metadata=metadata['metadata'],
@@ -327,8 +323,7 @@ async def test_process_slots(mock_build_tag_config, mock_gather_read_tags, forma
}
@mark.asyncio
async def test_format_schedule_config(formatters):
def test_format_schedule_config(formatters):
input_data = {
'schedule_config': [
{'namespace': 'test_namespace1', 'schedule_name': 'test1', 'updated_at': '2021-01-01'},
@@ -337,7 +332,7 @@ async def test_format_schedule_config(formatters):
**metadata,
}
result = await formatters.format_schedule_config(input_data)
result = formatters.format_schedule_config(input_data)
assert result == {
'test_namespace1': {'test1': '2021-01-01'},
@@ -345,8 +340,7 @@ async def test_format_schedule_config(formatters):
}
@mark.asyncio
async def test_create_schedule_config(formatters):
def test_create_schedule_config(formatters):
input_data = {
'current_schedule_config': {
'scouter': {
@@ -372,7 +366,7 @@ async def test_create_schedule_config(formatters):
},
}
result = await formatters.create_schedule_config(input_data)
result = formatters.create_schedule_config(input_data)
assert result == {
'to_create': {
@@ -399,8 +393,7 @@ async def test_create_schedule_config(formatters):
}
@mark.asyncio
async def test_create_slot_config(formatters):
def test_create_slot_config(formatters):
input_data = {
'current_slot_config': {
'1': {'frequency': 60, 'data': {'test': 'test'}},
@@ -409,7 +402,7 @@ async def test_create_slot_config(formatters):
'slot_config': {'1': {'frequency': 60, 'data': {'test': 'test2'}}},
}
result = await formatters.create_slot_config(input_data)
result = formatters.create_slot_config(input_data)
assert result == {
'to_delete': ['2'],
@@ -417,15 +410,14 @@ async def test_create_slot_config(formatters):
}
@mark.asyncio
async def test_send_success_report(formatters):
await formatters.send_success_report(
def test_send_success_report(formatters):
formatters.send_success_report(
metadata=metadata,
message='test_message',
notification_id='test_notification_id',
attachment={'test': 'test'},
)
formatters.send_notification_async.assert_called_once_with(
formatters.send_notification.assert_called_once_with(
metadata=metadata,
notification_id='test_notification_id',
message='test_message',
@@ -435,15 +427,14 @@ async def test_send_success_report(formatters):
)
@mark.asyncio
async def test_send_error_report(formatters):
await formatters.send_error_report(
def test_send_error_report(formatters):
formatters.send_error_report(
metadata=metadata,
message='test_message',
notification_id='test_notification_id',
attachment='test_attachment',
)
formatters.send_notification_async.assert_called_once_with(
formatters.send_notification.assert_called_once_with(
metadata=metadata,
notification_id='test_notification_id',
message='test_message',
@@ -492,8 +483,7 @@ def test_parse_report_schedule(formatters):
)
@mark.asyncio
async def test_report_schedule_orchestration(formatters):
def test_report_schedule_orchestration(formatters):
formatters.parse_report_schedule = MagicMock(side_effect=formatters.parse_report_schedule)
formatters.send_success_report = AsyncMock()
formatters.send_error_report = AsyncMock()
@@ -569,7 +559,7 @@ async def test_report_schedule_orchestration(formatters):
],
}
await formatters.report_schedule_orchestration(input_data)
formatters.report_schedule_orchestration(input_data)
formatters.parse_report_schedule.assert_has_calls(
[
@@ -624,8 +614,7 @@ async def test_report_schedule_orchestration(formatters):
)
@mark.asyncio
async def test_report_slot_orchestration(formatters):
def test_report_slot_orchestration(formatters):
formatters.parse_report = MagicMock(side_effect=formatters.parse_report)
formatters.send_success_report = AsyncMock()
formatters.send_error_report = AsyncMock()
@@ -642,7 +631,7 @@ async def test_report_slot_orchestration(formatters):
},
}
await formatters.report_slot_orchestration(input_data)
formatters.report_slot_orchestration(input_data)
formatters.parse_report.assert_has_calls(
[call(input_data['inserted_slots']), call(input_data['deleted_slots'])]
@@ -679,8 +668,7 @@ async def test_report_slot_orchestration(formatters):
)
@mark.asyncio
async def test_format_log_report(formatters):
def test_format_log_report(formatters):
input_data = {
**metadata,
'receiver_groups': {
@@ -722,7 +710,7 @@ async def test_format_log_report(formatters):
'mail_type': 'test_mail_type',
}
result = await formatters.format_log_report(input_data)
result = formatters.format_log_report(input_data)
expected_result = DataFrame(
[
@@ -747,8 +735,7 @@ async def test_format_log_report(formatters):
assert DataFrame(result).equals(expected_result)
@mark.asyncio
async def test_filter_notification_reports(formatters):
def test_filter_notification_reports(formatters):
input_data = {
**metadata,
'notification_package': [
@@ -766,7 +753,7 @@ async def test_filter_notification_reports(formatters):
],
}
response = await formatters.filter_notification_reports(input_data)
response = formatters.filter_notification_reports(input_data)
assert response == {
'test_group_1': {

131
tests/orchestrator/activities/test_mongo_db.py
View File

@@ -1,9 +1,9 @@
from datetime import datetime
from unittest.mock import ANY, AsyncMock, MagicMock, patch
from pytest import fixture, mark
from pytest import fixture
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ
from orchestrator.activities.mongo_db import MongoDB
@@ -20,7 +20,6 @@ def mongo_db(mongo_mock):
metrics_controller=AsyncMock(),
)
mongo.send_notification = MagicMock()
mongo.send_notification_async = AsyncMock()
mongo.emit_metric = AsyncMock()
return mongo
@@ -63,11 +62,10 @@ def test___del__(mongo_db):
mongo_db.close.assert_called_once()
@mark.asyncio
async def test_find_documents_in_mongodb_success(mongo_db):
def test_find_documents_in_mongodb_success(mongo_db):
input_data = {'collection': 'test_collection', 'filters': {'name': {'$exists': True}}}
mongo_db.mongo_db_repository.find = AsyncMock(
mongo_db.mongo_db_repository.find = MagicMock(
return_value=[
{
'name': 'test1',
@@ -84,7 +82,7 @@ async def test_find_documents_in_mongodb_success(mongo_db):
]
)
result = await mongo_db.find_documents_in_mongodb(
result = mongo_db.find_documents_in_mongodb(
{'query': input_data, 'timestamp_fields': ['timestamp']}
)
@@ -106,16 +104,15 @@ metadata = {
}
@mark.asyncio
async def test_find_documents_in_mongodb_failure(mongo_db):
def test_find_documents_in_mongodb_failure(mongo_db):
input_data = {'collection': 'test_collection', 'filters': {'name': {'$exists': True}}}
mongo_db.mongo_db_repository.find = AsyncMock(side_effect=Exception('Error'))
mongo_db.mongo_db_repository.find = MagicMock(side_effect=Exception('Error'))
try:
await mongo_db.find_documents_in_mongodb({'query': input_data, **metadata})
mongo_db.find_documents_in_mongodb({'query': input_data, **metadata})
except Exception as e:
assert str(e) == 'Error'
mongo_db.send_notification_async.assert_called_once_with(
mongo_db.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='MONGODB_QUERY_ERROR',
message='Failed to execute MongoDB query: Error',
@@ -128,12 +125,11 @@ async def test_find_documents_in_mongodb_failure(mongo_db):
raise AssertionError('Expected an exception to be raised')
@mark.asyncio
async def test_find_documents_in_mongodb_missing_collection(mongo_db):
def test_find_documents_in_mongodb_missing_collection(mongo_db):
input_data = {'query': {'filters': {}}}
try:
await mongo_db.find_documents_in_mongodb(input_data)
mongo_db.find_documents_in_mongodb(input_data)
except ValueError as e:
assert str(e) == 'Collection name must be provided in the query.'
@@ -142,13 +138,12 @@ async def test_find_documents_in_mongodb_missing_collection(mongo_db):
raise AssertionError('Expected a ValueError to be raised')
@mark.asyncio
async def test_aggregate_documents_in_mongodb_success(mongo_db):
def test_aggregate_documents_in_mongodb_success(mongo_db):
input_data = {
'collection': 'test_collection',
'aggregation': [{'$match': {'name': {'$exists': True}}}, {'$project': {'name': 1}}],
}
mongo_db.mongo_db_repository.aggregate = AsyncMock(
mongo_db.mongo_db_repository.aggregate = MagicMock(
return_value=[
{
'name': 'test1',
@@ -165,7 +160,7 @@ async def test_aggregate_documents_in_mongodb_success(mongo_db):
]
)
result = await mongo_db.aggregate_documents_in_mongodb(
result = mongo_db.aggregate_documents_in_mongodb(
{'query': input_data, 'timestamp_fields': ['timestamp']}
)
@@ -180,19 +175,18 @@ async def test_aggregate_documents_in_mongodb_success(mongo_db):
)
@mark.asyncio
async def test_aggregate_documents_in_mongodb_failure(mongo_db):
def test_aggregate_documents_in_mongodb_failure(mongo_db):
input_data = {
'collection': 'test_collection',
'aggregation': [{'$match': {'name': {'$exists': True}}}, {'$project': {'name': 1}}],
}
mongo_db.mongo_db_repository.aggregate = AsyncMock(side_effect=Exception('Error'))
mongo_db.mongo_db_repository.aggregate = MagicMock(side_effect=Exception('Error'))
try:
await mongo_db.aggregate_documents_in_mongodb({'query': input_data, **metadata})
mongo_db.aggregate_documents_in_mongodb({'query': input_data, **metadata})
except Exception as e:
assert str(e) == 'Error'
mongo_db.send_notification_async.assert_called_once_with(
mongo_db.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='MONGODB_AGGREGATION_ERROR',
message='Failed to execute MongoDB aggregation: Error',
@@ -205,12 +199,11 @@ async def test_aggregate_documents_in_mongodb_failure(mongo_db):
raise AssertionError('Expected an exception to be raised')
@mark.asyncio
async def test_aggregate_documents_in_mongodb_missing_collection(mongo_db):
def test_aggregate_documents_in_mongodb_missing_collection(mongo_db):
input_data = {'query': {'aggregation': []}}
try:
await mongo_db.aggregate_documents_in_mongodb(input_data)
mongo_db.aggregate_documents_in_mongodb(input_data)
except ValueError as e:
assert str(e) == 'Collection name must be provided in the query.'
@@ -219,12 +212,11 @@ async def test_aggregate_documents_in_mongodb_missing_collection(mongo_db):
raise AssertionError('Expected a ValueError to be raised')
@mark.asyncio
async def test_aggregate_documents_in_mongodb_missing_aggregation(mongo_db):
def test_aggregate_documents_in_mongodb_missing_aggregation(mongo_db):
input_data = {'query': {'collection': 'test_collection'}}
try:
await mongo_db.aggregate_documents_in_mongodb(input_data)
mongo_db.aggregate_documents_in_mongodb(input_data)
except ValueError as e:
assert str(e) == 'Aggregation must be provided.'
@@ -233,17 +225,16 @@ async def test_aggregate_documents_in_mongodb_missing_aggregation(mongo_db):
raise AssertionError('Expected a ValueError to be raised')
@mark.asyncio
@patch('orchestrator.activities.mongo_db.now')
async def test_update_pipelines_timestamps_success(now_mock, mongo_db):
def test_update_pipelines_timestamps_success(now_mock, mongo_db):
input_data = {
'updated_pipelines': [
{'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
{'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
]
}
mongo_db.mongo_db_repository.update_many = AsyncMock(return_value=MagicMock())
await mongo_db.update_pipelines_timestamps(input_data)
mongo_db.mongo_db_repository.update_many = MagicMock(return_value=MagicMock())
mongo_db.update_pipelines_timestamps(input_data)
mongo_db.mongo_db_repository.update_many.assert_called_once_with(
'orchestrated_schedules',
{
@@ -257,9 +248,8 @@ async def test_update_pipelines_timestamps_success(now_mock, mongo_db):
)
@mark.asyncio
@patch('orchestrator.activities.mongo_db.now')
async def test_update_pipelines_timestamps_failure(now_mock, mongo_db):
def test_update_pipelines_timestamps_failure(now_mock, mongo_db):
input_data = {
'updated_pipelines': [
{'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
@@ -270,10 +260,10 @@ async def test_update_pipelines_timestamps_failure(now_mock, mongo_db):
mongo_db.mongo_db_repository.update_many.side_effect = Exception('Error')
try:
await mongo_db.update_pipelines_timestamps(input_data)
mongo_db.update_pipelines_timestamps(input_data)
except Exception as e:
assert str(e) == 'Error'
mongo_db.send_notification_async.assert_called_once_with(
mongo_db.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='MONGODB_UPDATE_PIPELINES_ERROR',
message='Failed to update pipelines timestamps: Error',
@@ -286,17 +276,16 @@ async def test_update_pipelines_timestamps_failure(now_mock, mongo_db):
raise AssertionError('Expected an exception to be raised')
@mark.asyncio
@patch('orchestrator.activities.mongo_db.now')
async def test_create_pipelines_timestamps_success(now_mock, mongo_db):
def test_create_pipelines_timestamps_success(now_mock, mongo_db):
input_data = {
'created_pipelines': [
{'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
{'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
]
}
mongo_db.mongo_db_repository.insert_many = AsyncMock(return_value=MagicMock())
await mongo_db.create_pipelines_timestamps(input_data)
mongo_db.mongo_db_repository.insert_many = MagicMock(return_value=MagicMock())
mongo_db.create_pipelines_timestamps(input_data)
mongo_db.mongo_db_repository.insert_many.assert_called_once_with(
'orchestrated_schedules',
[
@@ -307,9 +296,8 @@ async def test_create_pipelines_timestamps_success(now_mock, mongo_db):
)
@mark.asyncio
@patch('orchestrator.activities.mongo_db.now')
async def test_create_pipelines_timestamps_failure(now_mock, mongo_db):
def test_create_pipelines_timestamps_failure(now_mock, mongo_db):
input_data = {
'created_pipelines': [
{'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
@@ -319,10 +307,10 @@ async def test_create_pipelines_timestamps_failure(now_mock, mongo_db):
}
mongo_db.mongo_db_repository.insert_many.side_effect = Exception('Error')
try:
await mongo_db.create_pipelines_timestamps(input_data)
mongo_db.create_pipelines_timestamps(input_data)
except Exception as e:
assert str(e) == 'Error'
mongo_db.send_notification_async.assert_called_once_with(
mongo_db.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='MONGODB_CREATE_PIPELINES_ERROR',
message='Failed to create pipelines timestamps: Error',
@@ -335,17 +323,16 @@ async def test_create_pipelines_timestamps_failure(now_mock, mongo_db):
raise AssertionError('Expected an exception to be raised')
@mark.asyncio
@patch('orchestrator.activities.mongo_db.now')
async def test_delete_pipelines_timestamps_success(now_mock, mongo_db):
def test_delete_pipelines_timestamps_success(now_mock, mongo_db):
input_data = {
'deleted_pipelines': [
{'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
{'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
]
}
mongo_db.mongo_db_repository.delete_many = AsyncMock(return_value=MagicMock())
await mongo_db.delete_pipelines_timestamps(input_data)
mongo_db.mongo_db_repository.delete_many = MagicMock(return_value=MagicMock())
mongo_db.delete_pipelines_timestamps(input_data)
mongo_db.mongo_db_repository.delete_many.assert_called_once_with(
'orchestrated_schedules',
{
@@ -358,9 +345,8 @@ async def test_delete_pipelines_timestamps_success(now_mock, mongo_db):
)
@mark.asyncio
@patch('orchestrator.activities.mongo_db.now')
async def test_delete_pipelines_timestamps_failure(now_mock, mongo_db):
def test_delete_pipelines_timestamps_failure(now_mock, mongo_db):
input_data = {
'deleted_pipelines': [
{'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
@@ -370,10 +356,10 @@ async def test_delete_pipelines_timestamps_failure(now_mock, mongo_db):
}
mongo_db.mongo_db_repository.delete_many.side_effect = Exception('Error')
try:
await mongo_db.delete_pipelines_timestamps(input_data)
mongo_db.delete_pipelines_timestamps(input_data)
except Exception as e:
assert str(e) == 'Error'
mongo_db.send_notification_async.assert_called_once_with(
mongo_db.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='MONGODB_DELETE_PIPELINES_ERROR',
message='Failed to delete pipelines timestamps: Error',
@@ -386,8 +372,7 @@ async def test_delete_pipelines_timestamps_failure(now_mock, mongo_db):
raise AssertionError('Expected an exception to be raised')
@mark.asyncio
async def test_create_collection_with_ttl_index_success(mongo_db):
def test_create_collection_with_ttl_index_success(mongo_db):
input_data = {
**metadata,
'pipelines': {
@@ -425,7 +410,7 @@ async def test_create_collection_with_ttl_index_success(mongo_db):
side_effect=[collection_1, collection_2, collection_3]
)
await mongo_db.create_collection_with_ttl_index(input_data)
mongo_db.create_collection_with_ttl_index(input_data)
mongo_db.mongo_db_repository.database.list_collection_names.assert_called_once_with()
@@ -447,8 +432,7 @@ async def test_create_collection_with_ttl_index_success(mongo_db):
collection_3.create_index.assert_not_called()
@mark.asyncio
async def test_create_collection_with_ttl_index_failure(mongo_db):
def test_create_collection_with_ttl_index_failure(mongo_db):
input_data = {**metadata, 'pipelines': {'scouter-pipeline': {'topic': 'raw_scouter_pipeline'}}}
mongo_db.mongo_db_repository.database.list_collection_names.return_value = []
@@ -456,10 +440,10 @@ async def test_create_collection_with_ttl_index_failure(mongo_db):
mongo_db.mongo_db_repository.database.create_collection.side_effect = Exception('Error')
try:
await mongo_db.create_collection_with_ttl_index(input_data)
mongo_db.create_collection_with_ttl_index(input_data)
except Exception as e:
assert str(e) == 'Error'
mongo_db.send_notification_async.assert_called_once_with(
mongo_db.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='MONGODB_CREATE_COLLECTION_ERROR',
message='Failed to create collection raw_scouter_pipeline with TTL index: Error',
@@ -471,10 +455,9 @@ async def test_create_collection_with_ttl_index_failure(mongo_db):
raise AssertionError('Expected an exception to be raised')
@mark.asyncio
async def test_load_latest_data_none_last_data_timestamp(mongo_db):
def test_load_latest_data_none_last_data_timestamp(mongo_db):
"""Test load_latest_data"""
mongo_db.mongo_db_repository.find = AsyncMock(
mongo_db.mongo_db_repository.find = MagicMock(
return_value=[
{
'name': 'test1',
@@ -484,7 +467,7 @@ async def test_load_latest_data_none_last_data_timestamp(mongo_db):
]
)
result = await mongo_db.load_latest_data(
result = mongo_db.load_latest_data(
{
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'collection_name': 'test_collection',
@@ -502,11 +485,10 @@ async def test_load_latest_data_none_last_data_timestamp(mongo_db):
assert result == [{'name': 'test1', 'value': 1, 'timestamp': '2023-01-01 12:00:00+0000'}]
@mark.asyncio
async def test_load_latest_data_not_none_last_data_timestamp(mongo_db):
def test_load_latest_data_not_none_last_data_timestamp(mongo_db):
"""Test load_latest_data"""
mongo_db.mongo_db_repository.find = AsyncMock(
mongo_db.mongo_db_repository.find = MagicMock(
return_value=[
{
'name': 'test1',
@@ -516,7 +498,7 @@ async def test_load_latest_data_not_none_last_data_timestamp(mongo_db):
]
)
result = await mongo_db.load_latest_data(
result = mongo_db.load_latest_data(
{
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'collection_name': 'test_collection',
@@ -529,9 +511,7 @@ async def test_load_latest_data_not_none_last_data_timestamp(mongo_db):
'test_collection',
{
'level': 'ERROR',
'timestamp': {
'$gt': datetime.strptime('2023-01-01 12:00:00+0000', DATETIME_FORMAT_WITH_TZ)
},
'timestamp': {'$gt': '2023-01-01 12:00:00+0000'},
},
{'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
)
@@ -539,13 +519,12 @@ async def test_load_latest_data_not_none_last_data_timestamp(mongo_db):
assert result == [{'name': 'test1', 'value': 1, 'timestamp': '2023-01-01 12:00:00+0000'}]
@mark.asyncio
async def test_load_latest_data_error(mongo_db):
def test_load_latest_data_error(mongo_db):
"""Test load_latest_data"""
mongo_db.mongo_db_repository.find.side_effect = Exception('test')
try:
await mongo_db.load_latest_data(
mongo_db.load_latest_data(
{
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'collection_name': 'test_collection',
@@ -556,7 +535,7 @@ async def test_load_latest_data_error(mongo_db):
except Exception as e:
assert str(e) == 'test'
mongo_db.send_notification_async.assert_called_once_with(
mongo_db.send_notification.assert_called_once_with(
metadata={'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
notification_id='MONGO_LOAD_ERROR',
message='Error loading data from MongoDB: test',

131
tests/orchestrator/activities/test_slot_manager.py
View File

@@ -2,7 +2,7 @@ from datetime import timedelta
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
from pandas import DataFrame
from pytest import fixture, mark
from pytest import fixture
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
@@ -35,21 +35,18 @@ def slot_manager(_redis_mock):
slot_manager.logger = MagicMock()
slot_manager.notification_handler = MagicMock()
slot_manager.send_notification = MagicMock()
slot_manager.send_notification_async = AsyncMock()
slot_manager.emit_metric = AsyncMock()
return slot_manager
@mark.asyncio
async def test_load_opc_slots_no_slot_keys(slot_manager):
slot_manager.redis_repository.keys = AsyncMock(return_value=[])
assert await slot_manager.load_opc_slots(metadata) == {}
def test_load_opc_slots_no_slot_keys(slot_manager):
slot_manager.redis_repository.keys = MagicMock(return_value=[])
assert slot_manager.load_opc_slots(metadata) == {}
@mark.asyncio
async def test_load_opc_slots(slot_manager):
slot_manager.redis_repository.keys = AsyncMock(
def test_load_opc_slots(slot_manager):
slot_manager.redis_repository.keys = MagicMock(
return_value=[
b'slot:opc_tags:1',
b'slot:opc_tags:2',
@@ -57,9 +54,9 @@ async def test_load_opc_slots(slot_manager):
]
)
slot_manager.redis_repository.get = AsyncMock(side_effect=['value1', 'value2', None])
slot_manager.redis_repository.get = MagicMock(side_effect=['value1', 'value2', None])
response = await slot_manager.load_opc_slots(metadata)
response = slot_manager.load_opc_slots(metadata)
assert response == {
'slot:opc_tags:1': 'value1',
@@ -68,9 +65,8 @@ async def test_load_opc_slots(slot_manager):
}
@mark.asyncio
async def test_load_opc_slots_no_decode(slot_manager):
slot_manager.redis_repository.keys = AsyncMock(
def test_load_opc_slots_no_decode(slot_manager):
slot_manager.redis_repository.keys = MagicMock(
return_value=[
'slot:opc_tags:1',
'slot:opc_tags:2',
@@ -78,9 +74,9 @@ async def test_load_opc_slots_no_decode(slot_manager):
]
)
slot_manager.redis_repository.get = AsyncMock(side_effect=['value1', 'value2', None])
slot_manager.redis_repository.get = MagicMock(side_effect=['value1', 'value2', None])
response = await slot_manager.load_opc_slots(metadata)
response = slot_manager.load_opc_slots(metadata)
assert response == {
'slot:opc_tags:1': 'value1',
@@ -89,9 +85,8 @@ async def test_load_opc_slots_no_decode(slot_manager):
}
@mark.asyncio
async def test_load_opc_slots_error(slot_manager):
slot_manager.redis_repository.keys = AsyncMock(
def test_load_opc_slots_error(slot_manager):
slot_manager.redis_repository.keys = MagicMock(
return_value=[
'slot:opc_tags:1',
'slot:opc_tags:2',
@@ -99,13 +94,13 @@ async def test_load_opc_slots_error(slot_manager):
]
)
slot_manager.redis_repository.get = AsyncMock(side_effect=Exception('Test exception'))
slot_manager.redis_repository.get = MagicMock(side_effect=Exception('Test exception'))
try:
await slot_manager.load_opc_slots(metadata)
slot_manager.load_opc_slots(metadata)
except Exception as e:
assert str(e) == 'Test exception'
slot_manager.send_notification_async.assert_called_once_with(
slot_manager.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='REDIS_GET_ERROR',
message='Failed to load OPC slots: Test exception',
@@ -118,9 +113,8 @@ async def test_load_opc_slots_error(slot_manager):
raise AssertionError('Expected an exception to be raised')
@mark.asyncio
async def test_load_active_ingestors(slot_manager):
slot_manager.redis_repository.keys = AsyncMock(
def test_load_active_ingestors(slot_manager):
slot_manager.redis_repository.keys = MagicMock(
return_value=[
b'heartbeat:ingestor:1',
b'heartbeat:ingestor:2',
@@ -128,14 +122,13 @@ async def test_load_active_ingestors(slot_manager):
]
)
response = await slot_manager.load_active_ingestors(metadata)
response = slot_manager.load_active_ingestors(metadata)
assert response == ['heartbeat:ingestor:1', 'heartbeat:ingestor:2', 'heartbeat:ingestor:3']
@mark.asyncio
async def test_load_active_ingestors_error(slot_manager):
slot_manager.redis_repository.keys = AsyncMock(
def test_load_active_ingestors_error(slot_manager):
slot_manager.redis_repository.keys = MagicMock(
return_value=[
'heartbeat:ingestor:1',
'heartbeat:ingestor:2',
@@ -143,13 +136,13 @@ async def test_load_active_ingestors_error(slot_manager):
]
)
slot_manager.redis_repository.keys = AsyncMock(side_effect=Exception('Test exception'))
slot_manager.redis_repository.keys = MagicMock(side_effect=Exception('Test exception'))
try:
await slot_manager.load_active_ingestors(metadata)
slot_manager.load_active_ingestors(metadata)
except Exception as e:
assert str(e) == 'Test exception'
slot_manager.send_notification_async.assert_called_once_with(
slot_manager.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='REDIS_GET_ERROR',
message='Failed to load active ingestors: Test exception',
@@ -162,11 +155,10 @@ async def test_load_active_ingestors_error(slot_manager):
raise AssertionError('Expected an exception to be raised')
@mark.asyncio
async def test_update_slots(slot_manager):
slot_manager.redis_repository.set = AsyncMock(side_effect=[None, Exception('Test exception')])
def test_update_slots(slot_manager):
slot_manager.redis_repository.set = MagicMock(side_effect=[None, Exception('Test exception')])
response = await slot_manager.update_slots({'to_insert': {'1': 'value1', '2': 'value2'}})
response = slot_manager.update_slots({'to_insert': {'1': 'value1', '2': 'value2'}})
slot_manager.redis_repository.set.assert_has_calls(
[call('slot:opc_tags:1', 'value1', ttl=None), call('slot:opc_tags:2', 'value2', ttl=None)]
@@ -178,13 +170,12 @@ async def test_update_slots(slot_manager):
}
@mark.asyncio
async def test_delete_slots(slot_manager):
slot_manager.redis_repository.delete = AsyncMock(
def test_delete_slots(slot_manager):
slot_manager.redis_repository.delete = MagicMock(
side_effect=[None, Exception('Test exception')]
)
response = await slot_manager.delete_slots({'to_delete': ['1', '2']})
response = slot_manager.delete_slots({'to_delete': ['1', '2']})
slot_manager.redis_repository.delete.assert_has_calls(
[call('slot:opc_tags:1'), call('slot:opc_tags:2')]
@@ -196,8 +187,7 @@ async def test_delete_slots(slot_manager):
}
@mark.asyncio
async def test_get_last_data_timestamp_none(slot_manager):
def test_get_last_data_timestamp_none(slot_manager):
"""Test get_last_data_timestamp"""
test_data = {
**metadata,
@@ -206,15 +196,14 @@ async def test_get_last_data_timestamp_none(slot_manager):
'mail_type': 'test_mail_type',
}
slot_manager.redis_repository.get = AsyncMock(return_value=None)
slot_manager.redis_repository.get = MagicMock(return_value=None)
result = await slot_manager.get_last_data_timestamp(test_data)
result = slot_manager.get_last_data_timestamp(test_data)
assert result is None
@mark.asyncio
async def test_get_last_data_timestamp_not_none(slot_manager):
def test_get_last_data_timestamp_not_none(slot_manager):
"""Test get_last_data_timestamp"""
test_data = {
**metadata,
@@ -223,9 +212,9 @@ async def test_get_last_data_timestamp_not_none(slot_manager):
'mail_type': 'test_mail_type',
}
slot_manager.redis_repository.get = AsyncMock(return_value='2023-01-01 12:00:00')
slot_manager.redis_repository.get = MagicMock(return_value='2023-01-01 12:00:00')
result = await slot_manager.get_last_data_timestamp(test_data)
result = slot_manager.get_last_data_timestamp(test_data)
slot_manager.redis_repository.get.assert_called_once_with(
'notification_last_timestamp:test_mail_type'
@@ -234,8 +223,7 @@ async def test_get_last_data_timestamp_not_none(slot_manager):
assert result == '2023-01-01 12:00:00'
@mark.asyncio
async def test_get_last_data_timestamp_error(slot_manager):
def test_get_last_data_timestamp_error(slot_manager):
"""Test get_last_data_timestamp error"""
test_data = {
**metadata,
@@ -244,16 +232,15 @@ async def test_get_last_data_timestamp_error(slot_manager):
'mail_type': 'test_mail_type',
}
slot_manager.send_notification_async = AsyncMock()
slot_manager.redis_repository.get = AsyncMock(side_effect=Exception('test'))
slot_manager.redis_repository.get = MagicMock(side_effect=Exception('test'))
try:
await slot_manager.get_last_data_timestamp(test_data)
slot_manager.get_last_data_timestamp(test_data)
except Exception as e:
assert str(e) == 'test'
slot_manager.send_notification_async.assert_called_once_with(
slot_manager.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='REDIS_GET_ERROR',
message='Error getting last data timestamp: test',
@@ -266,8 +253,7 @@ async def test_get_last_data_timestamp_error(slot_manager):
raise AssertionError('Expected exception')
@mark.asyncio
async def test_put_last_data_timestamp_empty_dataframe(slot_manager):
def test_put_last_data_timestamp_empty_dataframe(slot_manager):
"""Test put_last_data_timestamp with empty dataframe"""
test_data = {
**metadata,
@@ -279,15 +265,14 @@ async def test_put_last_data_timestamp_empty_dataframe(slot_manager):
slot_manager.set = MagicMock()
result = await slot_manager.put_last_data_timestamp(test_data)
result = slot_manager.put_last_data_timestamp(test_data)
assert result is None
slot_manager.set.assert_not_called()
@mark.asyncio
async def test_put_last_data_timestamp_not_empty_dataframe(slot_manager):
def test_put_last_data_timestamp_not_empty_dataframe(slot_manager):
"""Test put_last_data_timestamp with not empty dataframe"""
data = DataFrame(
@@ -305,9 +290,9 @@ async def test_put_last_data_timestamp_not_empty_dataframe(slot_manager):
'mail_type': 'test_mail_type',
}
slot_manager.redis_repository.set = AsyncMock()
slot_manager.redis_repository.set = MagicMock()
result = await slot_manager.put_last_data_timestamp(test_data)
result = slot_manager.put_last_data_timestamp(test_data)
assert result == '2023-01-01 12:00:01'
@@ -316,8 +301,7 @@ async def test_put_last_data_timestamp_not_empty_dataframe(slot_manager):
)
@mark.asyncio
async def test_put_last_data_timestamp_error(slot_manager):
def test_put_last_data_timestamp_error(slot_manager):
"""Test put_last_data_timestamp error"""
test_data = {
**metadata,
@@ -333,16 +317,15 @@ async def test_put_last_data_timestamp_error(slot_manager):
'mail_type': 'test_mail_type',
}
slot_manager.send_notification_async = AsyncMock()
slot_manager.redis_repository.set = AsyncMock(side_effect=Exception('test'))
slot_manager.redis_repository.set = MagicMock(side_effect=Exception('test'))
try:
await slot_manager.put_last_data_timestamp(test_data)
slot_manager.put_last_data_timestamp(test_data)
except Exception as e:
assert str(e) == 'test'
slot_manager.send_notification_async.assert_called_once_with(
slot_manager.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='REDIS_SET_ERROR',
message='Error setting last data timestamp: test',
@@ -355,9 +338,8 @@ async def test_put_last_data_timestamp_error(slot_manager):
raise AssertionError('Expected exception')
@mark.asyncio
async def test_filter_notification_alerts(slot_manager):
slot_manager.redis_repository.get = AsyncMock(
def test_filter_notification_alerts(slot_manager):
slot_manager.redis_repository.get = MagicMock(
side_effect=[
None,
(now() - timedelta(seconds=600)).strftime(DATETIME_FORMAT_MS_WITH_TZ),
@@ -383,7 +365,7 @@ async def test_filter_notification_alerts(slot_manager):
'mail_type': 'test_mail_type',
}
response = await slot_manager.filter_notification_alerts(input_data)
response = slot_manager.filter_notification_alerts(input_data)
assert response == {
'test_group_1': {
@@ -404,8 +386,7 @@ async def test_filter_notification_alerts(slot_manager):
}
@mark.asyncio
async def test_store_notification_cache(slot_manager):
def test_store_notification_cache(slot_manager):
"""Test store_notification_cache"""
test_data = {
**metadata,
@@ -419,9 +400,9 @@ async def test_store_notification_cache(slot_manager):
'sent_ttl': 600,
}
slot_manager.redis_repository.set = AsyncMock()
slot_manager.redis_repository.set = MagicMock()
await slot_manager.store_notification_cache(test_data)
slot_manager.store_notification_cache(test_data)
slot_manager.redis_repository.set.assert_called_once_with(
'test_schedule_1:test_notification_id_1', ANY, ttl=600

103
tests/orchestrator/activities/test_temporal_manager.py
View File

@@ -206,7 +206,7 @@ async def test_create_schedule(
'test-workflow',
input_data['schedules']['scouter']['test-schedule'],
id='test-schedule',
task_queue='test-workflow-queue',
task_queue='test-workflow-legacy-queue',
execution_timeout=timedelta(seconds=100),
run_timeout=timedelta(seconds=100),
task_timeout=timedelta(seconds=100),
@@ -216,7 +216,7 @@ async def test_create_schedule(
'test-workflow',
input_data['schedules']['scouter']['test-schedule-invalid-frequency'],
id='test-schedule-invalid-frequency',
task_queue='test-workflow-queue',
task_queue='test-workflow-legacy-queue',
execution_timeout=timedelta(seconds=400),
run_timeout=timedelta(seconds=400),
task_timeout=timedelta(seconds=400),
@@ -226,7 +226,7 @@ async def test_create_schedule(
'test-workflow',
input_data['schedules']['laborious']['test-schedule-laborious'],
id='test-schedule-laborious',
task_queue='test-workflow-queue',
task_queue='test-workflow-legacy-queue',
execution_timeout=timedelta(seconds=500),
run_timeout=timedelta(seconds=500),
task_timeout=timedelta(seconds=500),
@@ -488,3 +488,100 @@ async def test_delete_schedules_with_no_client(temporal_manager):
str(e)
== f'Temporal client for abc not found, clients: {temporal_manager.temporal_clients}'
)
@mark.asyncio
@patch('orchestrator.activities.temporal_manager.parse_frequency', side_effect=parse_frequency)
@patch('orchestrator.activities.temporal_manager.Schedule')
@patch('orchestrator.activities.temporal_manager.ScheduleActionStartWorkflow')
@patch('orchestrator.activities.temporal_manager.ScheduleIntervalSpec')
@patch('orchestrator.activities.temporal_manager.ScheduleSpec')
@patch('orchestrator.activities.temporal_manager.TypedSearchAttributes')
@patch('orchestrator.activities.temporal_manager.SearchAttributePair')
async def test_create_schedules_default_runtime_legacy_queue(
_mock_search_attribute_pair,
_mock_typed_search_attributes,
_mock_schedule_spec,
_mock_schedule_interval_spec,
mock_schedule_action_start_workflow,
_mock_schedule,
_mock_parse_frequency,
temporal_manager,
):
input_data = {
'schedules': {
'scouter': {
'test-schedule': {
'model_id': 1,
'model_name': 'test-model-name',
'workflow_type': 'scouter',
'frequency': '1m',
'data': {'test': 'test'},
}
}
}
}
temporal_manager.temporal_clients['scouter'].create_schedule = AsyncMock()
await temporal_manager.create_schedules(input_data)
mock_schedule_action_start_workflow.assert_called_once_with(
'scouter',
input_data['schedules']['scouter']['test-schedule'],
id='test-schedule',
task_queue='scouter-legacy-queue',
execution_timeout=timedelta(seconds=300),
run_timeout=timedelta(seconds=300),
task_timeout=timedelta(seconds=300),
typed_search_attributes=_mock_typed_search_attributes.return_value,
)
@mark.asyncio
@patch('orchestrator.activities.temporal_manager.parse_frequency', side_effect=parse_frequency)
@patch('orchestrator.activities.temporal_manager.Schedule')
@patch('orchestrator.activities.temporal_manager.ScheduleActionStartWorkflow')
@patch('orchestrator.activities.temporal_manager.ScheduleIntervalSpec')
@patch('orchestrator.activities.temporal_manager.ScheduleSpec')
@patch('orchestrator.activities.temporal_manager.TypedSearchAttributes')
@patch('orchestrator.activities.temporal_manager.SearchAttributePair')
async def test_create_schedules_tenant_runtime_queue(
_mock_search_attribute_pair,
_mock_typed_search_attributes,
_mock_schedule_spec,
_mock_schedule_interval_spec,
mock_schedule_action_start_workflow,
_mock_schedule,
_mock_parse_frequency,
temporal_manager,
):
input_data = {
'schedules': {
'scouter': {
'test-schedule': {
'model_id': 1,
'model_name': 'test-model-name',
'workflow_type': 'scouter',
'frequency': '1m',
'runtime': 'tenant-x',
'data': {'test': 'test'},
}
}
}
}
temporal_manager.temporal_clients['scouter'].create_schedule = AsyncMock()
await temporal_manager.create_schedules(input_data)
mock_schedule_action_start_workflow.assert_called_once_with(
'scouter',
input_data['schedules']['scouter']['test-schedule'],
id='test-schedule',
task_queue='scouter-tenant-x-queue',
execution_timeout=timedelta(seconds=300),
run_timeout=timedelta(seconds=300),
task_timeout=timedelta(seconds=300),
typed_search_attributes=_mock_typed_search_attributes.return_value,
)

22
tests/orchestrator/utils/test_orchestrator_functions.py
View File

@@ -36,10 +36,22 @@ def test_common_config():
'on_conflict': 'error',
'execution_timeout_seconds': 300,
'task_timeout_seconds': 300,
'runtime': 'legacy',
}
assert result == expected
def test_common_config_preserves_explicit_runtime():
config = {
'workflow_type': 'scouter',
'schedule_name': 'test_schedule',
'model_id': 'test_model_id',
'model': {'name': 'test_model_name'},
'runtime': 'tenant-x',
}
assert common_config(config)['runtime'] == 'tenant-x'
def test_drift():
config = {
'workflow_type': 'drift',
@@ -67,6 +79,7 @@ def test_drift():
'drift_metrics': ['kolmogorov_smirnov', 'jensen_shannon'],
'execution_timeout_seconds': 300,
'task_timeout_seconds': 300,
'runtime': 'legacy',
}
assert result == expected
@@ -99,6 +112,7 @@ def test_simple_metrics():
'metrics': ['rmse', 'mse'],
'execution_timeout_seconds': 300,
'task_timeout_seconds': 300,
'runtime': 'legacy',
}
assert result == expected
@@ -129,6 +143,7 @@ def test_minimal_retrain():
'datetime_columns': ['timestamp'],
'execution_timeout_seconds': 300,
'task_timeout_seconds': 300,
'runtime': 'legacy',
}
assert result == expected
@@ -146,6 +161,7 @@ def test_scouter():
'tag_retention_minutes': 10,
'execution_timeout_seconds': 300,
'task_timeout_seconds': 300,
'runtime': 'legacy',
}
result = scouter(config)
expected = {
@@ -169,6 +185,7 @@ def test_scouter():
'execution_timeout_seconds': 300,
'task_timeout_seconds': 300,
'fill_missing_tags': False,
'runtime': 'legacy',
}
assert result == expected
@@ -289,6 +306,7 @@ def test_predictions_batch(mock_process_path_priority, mock_overlap_filter_confi
'predictions_storage_policy': 'erl:1',
'execution_timeout_seconds': 300,
'task_timeout_seconds': 300,
'runtime': 'legacy',
}
assert result == expected
@@ -575,6 +593,7 @@ def test_base_scouter():
'execution_timeout_seconds': 300,
'task_timeout_seconds': 300,
'fill_missing_tags': True,
'runtime': 'legacy',
}
assert result == expected
@@ -635,6 +654,7 @@ def test_pi_web_api_scouter():
'max_count': 5,
'api_timeout': 30,
},
'runtime': 'legacy',
}
assert result == expected
@@ -687,6 +707,7 @@ def test_pi_web_api_scouter_with_timeout_greater_than_frequency():
'max_count': 1,
'api_timeout': 60,
},
'runtime': 'legacy',
}
assert result == expected
@@ -738,5 +759,6 @@ def test_pi_web_api_scouter_with_no_timeout():
'max_count': 1,
'api_timeout': 30,
},
'runtime': 'legacy',
}
assert result == expected

49
tests/orchestrator/worker/test_worker.py Normal file
View File

@@ -0,0 +1,49 @@
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
from pytest import mark
from orchestrator.worker.worker import main
from orchestrator.workflows.alerts import Alerts
from orchestrator.workflows.orchestrator import Orchestrator
from orchestrator.workflows.reports import Reports
def test_main_is_coroutine():
assert asyncio.iscoroutinefunction(main)
@mark.asyncio
@patch('orchestrator.worker.worker.sys')
@patch('orchestrator.worker.worker.start_http_server')
@patch('orchestrator.worker.worker.NotificationHandler')
@patch('orchestrator.worker.worker.Activities')
@patch('orchestrator.worker.worker.prepare_worker')
@patch('orchestrator.worker.worker.client.Client.connect', new_callable=AsyncMock)
async def test_main_starts_three_workers_for_expected_workflows(
_connect_mock,
prepare_worker_mock,
activities_mock,
_notification_handler_mock,
_start_http_server,
_sys_mock,
):
"""
main() must spin up exactly three Temporal workers, one per main workflow
(Orchestrator, Alerts, Reports), and call run() on each.
"""
activities_instance = activities_mock.return_value
activities_instance.connect_to_temporal = AsyncMock()
worker_mock = MagicMock()
worker_mock.run = AsyncMock()
prepare_worker_mock.return_value = worker_mock
await main()
assert prepare_worker_mock.call_count == 3
main_workflows = [call.kwargs['main_workflow'] for call in prepare_worker_mock.call_args_list]
assert main_workflows == [Orchestrator, Alerts, Reports]
assert worker_mock.run.call_count == 3

4
tests/orchestrator/workflows/subworkflows/test_process_notifications.py
View File

@@ -1,7 +1,7 @@
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from orchestrator.activities.activities import Activities
from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications
@@ -92,7 +92,7 @@ async def test_run(workflow_mock, process_notifications):
'data': workflow_mock.execute_local_activity_method.return_value,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_MS_WITH_TZ,
'format': DATETIME_FORMAT_WITH_TZ,
},
},
schedule_to_close_timeout=ANY,