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

Merge pull request #24 from Aignosi/SIENTIAPDE-1231-ajustar-o-retreino-do-courier-no-laborious

Browse Source 
SIENTIAPDE-1231: Refactor Orchestrator Activities and Add Timeout Configuration
This commit is contained in:
Matheus Demoner
2025-10-16 12:22:18 -03:00
committed by GitHub
parent 6091d2ce6b 9031b4c88f
commit a834671be4
41 changed files with 3812 additions and 3902 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@@ -39,6 +39,7 @@ __pycache__/
# Ignorar coverage # Ignorar coverage
htmlcov/ htmlcov/
.coverage .coverage
coverage.xml
# git keys # git keys
git_key* git_key*

97
README.md
View File

@@ -2,6 +2,36 @@
A high-performance, scalable workflow orchestration system built on Temporal.io for automated pipeline management, notification delivery, and resource coordination. The Orchestrator provides enterprise-grade workflow automation, real-time alerting, and comprehensive monitoring capabilities for the SIENTIA platform. A high-performance, scalable workflow orchestration system built on Temporal.io for automated pipeline management, notification delivery, and resource coordination. The Orchestrator provides enterprise-grade workflow automation, real-time alerting, and comprehensive monitoring capabilities for the SIENTIA platform.
## 📑 Table of Contents
- [Features](#features)
- [Core Functionality](#core-functionality)
- [Advanced Capabilities](#advanced-capabilities)
- [Architecture](#architecture)
- [Architecture Principles](#architecture-principles)
- [Task Queue Isolation](#2-task-queue-isolation)
- [Workflows](#-workflows)
- [Orchestrator Workflow](#1-orchestrator-workflow-orchestratorpy)
- [Alerts Workflow](#2-alerts-workflow-alertspy)
- [Reports Workflow](#3-reports-workflow-reportspy)
- [Subworkflows](#subworkflows)
- [Load Notification Package](#1-load-notification-package-load_notification_packagepy)
- [Process Notifications](#2-process-notifications-process_notificationspy)
- [Notification Filtering System](#-notification-filtering-system)
- [Prerequisites](#-prerequisites)
- [Installation](#-installation)
- [How to Run](#-how-to-run)
- [Configuration](#-configuration)
- [Monitoring and Metrics](#-monitoring-and-metrics)
- [Testing](#-testing)
- [Code Quality & Validation](#-code-quality--validation)
- [Development](#-development)
- [Troubleshooting](#-troubleshooting)
- [Performance Tuning](#-performance-tuning)
- [Contributing](#-contributing)
- [License](#-license)
- [Support](#-support)
## Features ## Features
### Core Functionality ### Core Functionality
@@ -646,6 +676,73 @@ pytest tests/activities/test_mongo_db.py
pytest tests/workflows/test_orchestrator.py pytest tests/workflows/test_orchestrator.py
``` ```
## 🛡️ Code Quality & Validation
### Overview
Since Python is not compiled, this project ships a validation workflow to catch issues early. Use the `validate.sh` script to run formatting, linting, type checks, security analysis, and tests in one command.
### Validation Tools
- **Ruff**: formatting and linting (fast, replaces Black/Flake8)
- **mypy**: static typing checks
- **Bandit**: security static analysis
- **pytest**: unit/integration tests with coverage
### Tools Installation
```bash
pip install -r requirements-dev.txt
```
### Complete Validation (recommended)
```bash
./validate.sh
```
What `validate.sh` does:
1. Checks formatting with Ruff
2. Lints code with Ruff
3. Runs mypy type checking
4. Runs Bandit security analysis
5. Executes pytest with coverage (generates HTML report)
Exit codes are propagated so CI can fail fast when quality gates are not met.
### Individual Commands
```bash
# 1) Format check
ruff format --check orchestrator/ tests/
# 2) Lint
ruff check orchestrator/ tests/
# 3) Type check
mypy orchestrator/
# 4) Security
bandit -r orchestrator/ -ll
# 5) Tests with coverage
pytest tests/ --cov=orchestrator --cov-report=html
```
### Automatic Fixes
```bash
# Apply formatting
ruff format orchestrator/ tests/
# Autofix common lint issues
ruff check --fix orchestrator/ tests/
```
### Configuration
Tooling is configured in `pyproject.toml` (lint rules, formatting, typing). Adjust thresholds and rules there as needed.
## 🔧 Development ## 🔧 Development
### Project Structure ### Project Structure

46
init_orchestration.ipynb
View File

File diff suppressed because one or more lines are too long

64
orchestrator/activities/activities.py
View File

@@ -1,21 +1,22 @@
from temporalio import workflow from temporalio import workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from orchestrator.activities.temporal_manager import TemporalManager
from orchestrator.activities.slot_manager import SlotManager
from orchestrator.activities.formatters import Formatters
from orchestrator.activities.email import Email
from orchestrator.activities.mongo_db import MongoDB
from typing import Any
from logging import Logger from logging import Logger
from sientia_do.temporal.activities.postgres import Postgres from typing import Any
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.temporal.activities.postgres import Postgres
from orchestrator.activities.email import Email
from orchestrator.activities.formatters import Formatters
from orchestrator.activities.mongo_db import MongoDB
from orchestrator.activities.slot_manager import SlotManager
from orchestrator.activities.temporal_manager import TemporalManager
class Activities( # Couchbase, class Activities( # Couchbase,
TemporalManager, SlotManager, Formatters, MongoDB, Email, TemporalManager, SlotManager, Formatters, MongoDB, Email, Postgres
Postgres): ):
""" """
Central activities orchestrator for Temporal workflow operations. Central activities orchestrator for Temporal workflow operations.
@@ -34,7 +35,8 @@ class Activities( # Couchbase,
notification_handler (NotificationHandler): Notification management handler notification_handler (NotificationHandler): Notification management handler
""" """
def __init__(self, def __init__(
self,
temporal_config: dict[str, Any], temporal_config: dict[str, Any],
# couchbase_config: dict[str, Any], # couchbase_config: dict[str, Any],
redis_config: dict[str, Any], redis_config: dict[str, Any],
@@ -42,8 +44,8 @@ class Activities( # Couchbase,
email_config: dict[str, Any], email_config: dict[str, Any],
postgres_config: dict[str, Any], postgres_config: dict[str, Any],
logger: Logger, logger: Logger,
notification_handler: NotificationHandler): notification_handler: NotificationHandler,
):
# Initialize parent classes # Initialize parent classes
# Couchbase.__init__(self, connection_string=couchbase_config['connection_string'], # Couchbase.__init__(self, connection_string=couchbase_config['connection_string'],
# username=couchbase_config['username'], # username=couchbase_config['username'],
@@ -51,43 +53,54 @@ class Activities( # Couchbase,
# logger=logger, # logger=logger,
# notification_handler=notification_handler) # notification_handler=notification_handler)
TemporalManager.__init__(self, TemporalManager.__init__(
self,
host=temporal_config['temporal_host'], host=temporal_config['temporal_host'],
scouter_namespace=temporal_config['temporal_scouter_namespace'], scouter_namespace=temporal_config['temporal_scouter_namespace'],
laborious_namespace=temporal_config['temporal_laborious_namespace'], laborious_namespace=temporal_config['temporal_laborious_namespace'],
logger=logger, logger=logger,
notification_handler=notification_handler) notification_handler=notification_handler,
)
SlotManager.__init__(self, SlotManager.__init__(
self,
host=redis_config['host'], host=redis_config['host'],
port=redis_config['port'], port=redis_config['port'],
username=redis_config['username'], username=redis_config['username'],
password=redis_config['password'], password=redis_config['password'],
logger=logger, logger=logger,
notification_handler=notification_handler) notification_handler=notification_handler,
)
Formatters.__init__(self, Formatters.__init__(
self,
scouter_namespace=temporal_config['temporal_scouter_namespace'], scouter_namespace=temporal_config['temporal_scouter_namespace'],
laborious_namespace=temporal_config['temporal_laborious_namespace'], laborious_namespace=temporal_config['temporal_laborious_namespace'],
logger=logger, logger=logger,
notification_handler=notification_handler) notification_handler=notification_handler,
)
MongoDB.__init__(self, MongoDB.__init__(
self,
connection_string=mongodb_config['connection_string'], connection_string=mongodb_config['connection_string'],
database_name=mongodb_config['database_name'], database_name=mongodb_config['database_name'],
ttl_index_seconds=mongodb_config['ttl_index_seconds'], ttl_index_seconds=mongodb_config['ttl_index_seconds'],
logger=logger, logger=logger,
notification_handler=notification_handler) notification_handler=notification_handler,
)
Email.__init__(self, Email.__init__(
self,
sender_email=email_config['sender_email'], sender_email=email_config['sender_email'],
sender_password=email_config['sender_password'], sender_password=email_config['sender_password'],
smtp_server=email_config['smtp_server'], smtp_server=email_config['smtp_server'],
smtp_port=email_config['smtp_port'], smtp_port=email_config['smtp_port'],
logger=logger, logger=logger,
notification_handler=notification_handler) notification_handler=notification_handler,
)
Postgres.__init__(self, Postgres.__init__(
self,
host=postgres_config['host'], host=postgres_config['host'],
port=postgres_config['port'], port=postgres_config['port'],
user=postgres_config['user'], user=postgres_config['user'],
@@ -96,7 +109,8 @@ class Activities( # Couchbase,
min_connections=postgres_config['min_connections'], min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'], max_connections=postgres_config['max_connections'],
logger=logger, logger=logger,
notification_handler=notification_handler) notification_handler=notification_handler,
)
def shutdown(self): def shutdown(self):
""" """

53
orchestrator/activities/couchbase.py
View File

@@ -1,11 +1,12 @@
from temporalio import activity, workflow from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from typing import Any
from logging import Logger
from datetime import timedelta
import json import json
import traceback import traceback
from datetime import timedelta
from logging import Logger
from typing import Any
from couchbase.auth import PasswordAuthenticator from couchbase.auth import PasswordAuthenticator
from couchbase.cluster import Cluster from couchbase.cluster import Cluster
from couchbase.options import ClusterOptions from couchbase.options import ClusterOptions
@@ -34,44 +35,43 @@ class Couchbase(BaseActivity):
but maintained for potential future use. but maintained for potential future use.
""" """
def __init__(self, connection_string: str, username: str, def __init__(
password: str, logger: Logger, self,
notification_handler: NotificationHandler): connection_string: str,
username: str,
password: str,
logger: Logger,
notification_handler: NotificationHandler,
):
self.connection_string = connection_string self.connection_string = connection_string
self.username = username self.username = username
self.password = password self.password = password
logger.info("Initializing Couchbase connection...") logger.info('Initializing Couchbase connection...')
self.cluster = Cluster( self.cluster = Cluster(
connection_string, connection_string,
ClusterOptions( ClusterOptions(
authenticator=PasswordAuthenticator( authenticator=PasswordAuthenticator(username=username, password=password)
username=username, ),
password=password
)
)
) )
logger.info("Awaiting Couchbase connection...") logger.info('Awaiting Couchbase connection...')
self.cluster.wait_until_ready(timeout=timedelta(seconds=10)) self.cluster.wait_until_ready(timeout=timedelta(seconds=10))
logger.info("Couchbase connection ready") logger.info('Couchbase connection ready')
BaseActivity.__init__(self, BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler)
logger=logger,
notification_handler=notification_handler)
def shutdown(self): def shutdown(self):
try: try:
self.cluster.close() self.cluster.close()
except Exception as e: except Exception as e:
self.logger.error(f"Failed to close Couchbase connection: {e}") self.logger.error(f'Failed to close Couchbase connection: {e}')
def __del__(self): def __del__(self):
self.shutdown() self.shutdown()
@activity.defn(name="load_query_from_couchbase") @activity.defn(name='load_query_from_couchbase')
async def load_query_from_couchbase(self, input_data: dict[str, Any]) -> list[dict[str, Any]]: async def load_query_from_couchbase(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
""" """
Load a query from couchbase Load a query from couchbase
@@ -84,16 +84,16 @@ class Couchbase(BaseActivity):
""" """
query = input_data['query'] query = input_data['query']
self.logger.info(f"Executing couchbase query: {query}") self.logger.info(f'Executing couchbase query: {query}')
try: try:
result = self.cluster.query(query) result = self.cluster.query(query)
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.notification_handler.build_and_send_notification( self.notification_handler.build_and_send_notification(
notification_id="COUCHBASE_LOAD_QUERY_ERROR", notification_id='COUCHBASE_LOAD_QUERY_ERROR',
message=f"Failed to execute couchbase query: {e}", message=f'Failed to execute couchbase query: {e}',
block="load_query_from_couchbase", block='load_query_from_couchbase',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace, attachment_content=trace,
) )
@@ -106,8 +106,7 @@ class Couchbase(BaseActivity):
for row in result.rows(): for row in result.rows():
rows.append(row) rows.append(row)
self.logger.info("Fetched %d rows from couchbase", len(rows)) self.logger.info('Fetched %d rows from couchbase', len(rows))
self.logger.debug("Rows: \n %s", self.logger.debug('Rows: \n %s', json.dumps(rows, indent=4, sort_keys=True))
json.dumps(rows, indent=4, sort_keys=True))
return rows return rows

117
orchestrator/activities/email.py
View File

@@ -1,21 +1,23 @@
from smtplib import SMTPServerDisconnected from smtplib import SMTPServerDisconnected
from temporalio import workflow, activity
from temporalio import activity, workflow
from orchestrator import metrics from orchestrator import metrics
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
import traceback
import smtplib import smtplib
from typing import Any import traceback
from sientia_do.temporal.activities.base import BaseActivity from email import encoders
from sientia_do.observability.logger import Logger from email.mime.base import MIMEBase
from sientia_do.notifications.handlers import NotificationHandler
from orchestrator.utils.email_builder import EmailBuilder
from email.mime.multipart import MIMEMultipart from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText from email.mime.text import MIMEText
from email.mime.base import MIMEBase from typing import Any
from email import encoders
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.temporal.activities.base import BaseActivity
from orchestrator.utils.email_builder import EmailBuilder
class Email(BaseActivity): class Email(BaseActivity):
@@ -35,10 +37,15 @@ class Email(BaseActivity):
notification_handler (NotificationHandler): Notification management handler notification_handler (NotificationHandler): Notification management handler
""" """
def __init__(self, sender_email: str, sender_password: str, def __init__(
smtp_server: str, smtp_port: int, self,
logger: Logger, notification_handler: NotificationHandler): sender_email: str,
sender_password: str,
smtp_server: str,
smtp_port: int,
logger: Logger,
notification_handler: NotificationHandler,
):
self.email_builder = EmailBuilder(logger=logger) self.email_builder = EmailBuilder(logger=logger)
self.sender_email = sender_email self.sender_email = sender_email
@@ -46,7 +53,7 @@ class Email(BaseActivity):
self.smtp_port = smtp_port self.smtp_port = smtp_port
self.smtp_server = smtp_server self.smtp_server = smtp_server
logger.info(f"Initializing Email with {smtp_server}:{smtp_port}") logger.info(f'Initializing Email with {smtp_server}:{smtp_port}')
if smtp_server is not None: if smtp_server is not None:
self.server = smtplib.SMTP(smtp_server, smtp_port, timeout=20) self.server = smtplib.SMTP(smtp_server, smtp_port, timeout=20)
@@ -55,9 +62,7 @@ class Email(BaseActivity):
self.server.starttls() self.server.starttls()
self.server.login(self.sender_email, self.sender_password) self.server.login(self.sender_email, self.sender_password)
BaseActivity.__init__(self, BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler)
logger=logger,
notification_handler=notification_handler)
def shutdown(self): def shutdown(self):
""" """
@@ -65,7 +70,7 @@ class Email(BaseActivity):
""" """
self.server.quit() self.server.quit()
@activity.defn(name="build_email_html") @activity.defn(name='build_email_html')
async def build_email_html(self, input_data: dict[str, Any]) -> dict[str, Any]: async def build_email_html(self, input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Build HTML email content for configured receiver groups. Build HTML email content for configured receiver groups.
@@ -91,18 +96,14 @@ class Email(BaseActivity):
receiver_groups = input_data['receiver_groups'] receiver_groups = input_data['receiver_groups']
mail_type = input_data['mail_type'] mail_type = input_data['mail_type']
self.info(f"Building email html for {mail_type} mail type.", self.info(f'Building email html for {mail_type} mail type.', metadata=metadata)
metadata=metadata)
for group_name, group_config in receiver_groups.items(): for _group_name, group_config in receiver_groups.items():
html = self.email_builder.build_email(group_config['notifications'], mail_type)
html = self.email_builder.build_email(
group_config['notifications'], mail_type)
group_config['html'] = html group_config['html'] = html
self.info(f"Email html built for {mail_type} mail type.", self.info(f'Email html built for {mail_type} mail type.', metadata=metadata)
metadata=metadata)
return receiver_groups return receiver_groups
@@ -124,17 +125,12 @@ class Email(BaseActivity):
try: try:
# Create the attachment as a MIMEBase object # Create the attachment as a MIMEBase object
part = MIMEBase('application', 'octet-stream') part = MIMEBase('application', 'octet-stream')
part.set_payload( part.set_payload(attachment['attachment_content'].encode('utf-8'))
attachment['attachment_content'].encode('utf-8'))
encoders.encode_base64(part) encoders.encode_base64(part)
part.add_header( part.add_header('Content-Disposition', f'attachment; filename="{att_name}"')
'Content-Disposition',
f'attachment; filename="{att_name}"'
)
msg.attach(part) msg.attach(part)
except Exception as e: except Exception as e:
self.logger.error( self.logger.error(f'Failed to attach content of {att_name}: {e}')
f"Failed to attach content of {att_name}: {e}")
raise e raise e
@@ -152,30 +148,27 @@ class Email(BaseActivity):
Exception: If email sending fails after reconnection attempts. Exception: If email sending fails after reconnection attempts.
""" """
try: try:
self.server.sendmail( self.server.sendmail(self.sender_email, receivers, msg.as_string())
self.sender_email, receivers, msg.as_string())
except SMTPServerDisconnected as e: except SMTPServerDisconnected as e:
self.logger.error(f"SMTP server disconnected: {e}") self.logger.error(f'SMTP server disconnected: {e}')
self.logger.info( self.logger.info(f'Reconnecting to {self.smtp_server}:{self.smtp_port}')
f"Reconnecting to {self.smtp_server}:{self.smtp_port}")
if self.server: if self.server:
try: try:
self.server.quit() self.server.quit()
except SMTPServerDisconnected as e: except SMTPServerDisconnected as e:
self.logger.info(f"Server already disconnected: {e}") self.logger.info(f'Server already disconnected: {e}')
except Exception as e: except Exception as e:
self.logger.error(f"Failed to quit server: {e}") self.logger.error(f'Failed to quit server: {e}')
raise e raise e
self.server = smtplib.SMTP( self.server = smtplib.SMTP(self.smtp_server, self.smtp_port, timeout=20)
self.smtp_server, self.smtp_port, timeout=20)
if self.sender_password: if self.sender_password:
self.server.starttls() self.server.starttls()
self.server.login(self.sender_email, self.sender_password) self.server.login(self.sender_email, self.sender_password)
self.server.sendmail(self.sender_email, receivers, msg.as_string()) self.server.sendmail(self.sender_email, receivers, msg.as_string())
@activity.defn(name="send_email") @activity.defn(name='send_email')
async def send_email(self, input_data: dict[str, Any]) -> dict[str, Any]: async def send_email(self, input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Send email notifications to configured receiver groups. Send email notifications to configured receiver groups.
@@ -202,40 +195,38 @@ class Email(BaseActivity):
mail_type = input_data['mail_type'] mail_type = input_data['mail_type']
if self.smtp_server is None: if self.smtp_server is None:
self.info(f"Skipping email sending for {mail_type} mail type.", self.info(f'Skipping email sending for {mail_type} mail type.', metadata=metadata)
metadata=metadata)
return {} return {}
self.info(f"Sending email for {mail_type} mail type.", self.info(f'Sending email for {mail_type} mail type.', metadata=metadata)
metadata=metadata)
for group_name, group_config in receiver_groups.items(): for group_name, group_config in receiver_groups.items():
try: try:
receivers = ", ".join(group_config['members']) receivers = ', '.join(group_config['members'])
self.info(f"Sending email to {group_name}: {receivers}", self.info(f'Sending email to {group_name}: {receivers}', metadata=metadata)
metadata=metadata)
msg = MIMEMultipart() msg = MIMEMultipart()
msg.attach(MIMEText(group_config['html'], 'html')) msg.attach(MIMEText(group_config['html'], 'html'))
msg['From'] = self.sender_email msg['From'] = self.sender_email
msg['To'] = receivers msg['To'] = receivers
msg['Subject'] = f"SIENTIA™ {mail_type}" msg['Subject'] = f'SIENTIA™ {mail_type}'
msg = self.handle_attachments( msg = self.handle_attachments(
[ [
{ {
"filename": f"{notification['trigger']}_{notification['notification_id']}.txt", 'filename': f'{notification["trigger"]}_{notification["notification_id"]}.txt',
"attachment_content": notification['attachment_content'] 'attachment_content': notification['attachment_content'],
} }
for notification in group_config['notifications'] for notification in group_config['notifications']
if notification.get('attachment_content') is not None], if notification.get('attachment_content') is not None
msg) ],
msg,
)
self.try_send_email(msg, receivers) self.try_send_email(msg, receivers)
except Exception as e: except Exception as e:
self.error(f"Failed to send email to {group_name}: {e}", self.error(f'Failed to send email to {group_name}: {e}', metadata=metadata)
metadata=metadata)
traceback.print_exc() traceback.print_exc()
group_config['status'] = 'failed' group_config['status'] = 'failed'
else: else:
@@ -244,13 +235,11 @@ class Email(BaseActivity):
pod_id=self.pod_id, pod_id=self.pod_id,
model_name=metadata['model_name'], model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name'], pipeline_name=metadata['workflow_name'],
email_group=group_name email_group=group_name,
).inc() ).inc()
self.info(f"Email sent to {group_name}: {receivers}", self.info(f'Email sent to {group_name}: {receivers}', metadata=metadata)
metadata=metadata)
self.info(f"Email sent for {mail_type} mail type.", self.info(f'Email sent for {mail_type} mail type.', metadata=metadata)
metadata=metadata)
return receiver_groups return receiver_groups

414
orchestrator/activities/formatters.py
View File

@@ -1,20 +1,28 @@
from collections.abc import Hashable
from temporalio import activity, workflow from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
import json import json
from typing import Any
from logging import Logger from logging import Logger
from math import ceil
from typing import Any
from pandas import DataFrame from pandas import DataFrame
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.notifications.models import NotificationLevel from sientia_do.notifications.models import NotificationLevel
from orchestrator.utils.orchestrator_functions import ( from sientia_do.temporal.activities.base import BaseActivity
scouter, predictions_batch, gather_read_tags, build_tag_config, minimal_retrain
)
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
from math import ceil
topic_separator = "\n ========== \n" from orchestrator.utils.orchestrator_functions import (
build_tag_config,
gather_read_tags,
minimal_retrain,
predictions_batch,
scouter,
)
topic_separator = '\n ========== \n'
class Formatters(BaseActivity): class Formatters(BaseActivity):
@@ -26,12 +34,11 @@ class Formatters(BaseActivity):
distribution across active ingestors, and implementing notification filtering for distribution across active ingestors, and implementing notification filtering for
scheduled reports. scheduled reports.
Key Features: Key features:
- Pipeline schedule configuration formatting (scouter, predictions_batch, minimal_retrain) - Pipeline schedule configuration formatting ("scouter", "predictions_batch", "minimal_retrain")
- OPC slot distribution across active ingestors - OPC slot distribution across active ingestors
- Notification filtering for comprehensive reports - Notification filtering for comprehensive scheduled reports
- Group-based report filtering with ignore list support - Group-based report filtering with ignore list support
- Resource optimization algorithms
- Configuration validation and transformation - Configuration validation and transformation
Args: Args:
@@ -41,16 +48,18 @@ class Formatters(BaseActivity):
notification_handler (NotificationHandler): Notification management handler notification_handler (NotificationHandler): Notification management handler
""" """
def __init__(self, def __init__(
self,
scouter_namespace: str, scouter_namespace: str,
laborious_namespace: str, laborious_namespace: str,
logger: Logger, notification_handler: NotificationHandler): logger: Logger,
notification_handler: NotificationHandler,
):
self.scouter_namespace = scouter_namespace self.scouter_namespace = scouter_namespace
self.laborious_namespace = laborious_namespace self.laborious_namespace = laborious_namespace
BaseActivity.__init__(self, logger=logger, BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler)
notification_handler=notification_handler)
@activity.defn(name="process_schedules") @activity.defn(name='process_schedules')
async def process_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]: async def process_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Process pipeline configurations into Temporal-compatible schedule configurations. Process pipeline configurations into Temporal-compatible schedule configurations.
@@ -59,10 +68,10 @@ class Formatters(BaseActivity):
Temporal schedule configurations, organizing them by workflow type (scouter and Temporal schedule configurations, organizing them by workflow type (scouter and
laborious) and applying the appropriate configuration builders for each pipeline type. laborious) and applying the appropriate configuration builders for each pipeline type.
Pipeline Types Supported: Pipeline types supported:
- scouter: Data collection workflows with OPC tag configurations - "scouter": Data collection workflows with OPC tag configurations
- predictions_batch: ML prediction workflows with OPC write configurations - "predictions_batch": ML prediction workflows with OPC write configurations
- minimal_retrain: Model retraining workflows with SQL query configurations - "minimal_retrain": Model retraining workflows with SQL query configurations
Args: Args:
- input_data (dict[str, Any]): The input data containing - input_data (dict[str, Any]): The input data containing
@@ -70,49 +79,49 @@ class Formatters(BaseActivity):
- pipelines (list[dict[str, Any]]): The schedules to process. - pipelines (list[dict[str, Any]]): The schedules to process.
Returns: Returns:
- dict[str, Any]: The schedule config dictionary - dict[str, Any]: The schedule configuration dictionary keyed by namespace
""" """
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
self.info("Processing schedules...", metadata=metadata) self.info('Processing schedules...', metadata=metadata)
pipelines = input_data['pipelines'] pipelines = input_data['pipelines']
schedule_config = { schedule_config: dict[str, dict[str, Any]] = {
self.scouter_namespace: {}, self.scouter_namespace: {},
self.laborious_namespace: {} self.laborious_namespace: {},
} }
for pipeline in pipelines: for pipeline in pipelines:
if pipeline['workflow_type'] == 'scouter': if pipeline['workflow_type'] == 'scouter':
schedule_config[self.scouter_namespace][pipeline['schedule_name']] = { schedule_config[self.scouter_namespace][pipeline['schedule_name']] = {
**scouter(pipeline), **scouter(pipeline),
"updated_at": pipeline.get( 'updated_at': pipeline.get(
"updated_at", now().strftime(DATETIME_FORMAT_MS_WITH_TZ)) 'updated_at', now().strftime(DATETIME_FORMAT_MS_WITH_TZ)
),
} }
elif pipeline['workflow_type'] == 'predictions_batch': elif pipeline['workflow_type'] == 'predictions_batch':
schedule_config[self.laborious_namespace][pipeline['schedule_name'] schedule_config[self.laborious_namespace][pipeline['schedule_name']] = {
] = {
**predictions_batch(pipeline), **predictions_batch(pipeline),
"updated_at": pipeline.get( 'updated_at': pipeline.get(
"updated_at", now().strftime(DATETIME_FORMAT_MS_WITH_TZ)) 'updated_at', now().strftime(DATETIME_FORMAT_MS_WITH_TZ)
),
} }
elif pipeline['workflow_type'] == 'minimal_retrain': elif pipeline['workflow_type'] == 'minimal_retrain':
schedule_config[self.laborious_namespace][pipeline['schedule_name'] schedule_config[self.laborious_namespace][pipeline['schedule_name']] = {
] = {
**minimal_retrain(pipeline), **minimal_retrain(pipeline),
"updated_at": pipeline.get( 'updated_at': pipeline.get(
"updated_at", now().strftime(DATETIME_FORMAT_MS_WITH_TZ)) 'updated_at', now().strftime(DATETIME_FORMAT_MS_WITH_TZ)
),
} }
self.info("Processed schedules", metadata=metadata) self.info('Processed schedules', metadata=metadata)
self.debug(json.dumps( self.debug(json.dumps(schedule_config, indent=4, sort_keys=True), metadata=metadata)
schedule_config, indent=4, sort_keys=True), metadata=metadata)
return schedule_config return schedule_config
@activity.defn(name="process_slots") @activity.defn(name='process_slots')
async def process_slots(self, input_data: dict[str, Any]) -> dict[str, Any]: async def process_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Extracts all read tags from input pipelines, divides them into slots and Extracts all read tags from input pipelines, divides them into slots and
@@ -127,12 +136,12 @@ class Formatters(BaseActivity):
- active_ingestors (list[str]): The active ingestors to divide into slots. - active_ingestors (list[str]): The active ingestors to divide into slots.
Returns: Returns:
- dict[str, Any]: The slot config dictionary - dict[str, Any]: The slot configuration dictionary keyed by slot id (as string)
""" """
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
self.info("Processing slots...", metadata=metadata) self.info('Processing slots...', metadata=metadata)
pipelines = input_data['pipelines'] pipelines = input_data['pipelines']
opc_servers_list = input_data['opc_servers'] opc_servers_list = input_data['opc_servers']
@@ -150,47 +159,46 @@ class Formatters(BaseActivity):
number_of_slots = len(active_ingestors) if active_ingestors else 1 number_of_slots = len(active_ingestors) if active_ingestors else 1
tags_per_slot = ceil(number_of_tags / number_of_slots) tags_per_slot = ceil(number_of_tags / number_of_slots)
slot_config = {} slot_config: dict[str, Any] = {}
last_index = 0 last_index = 0
for i in range(1, number_of_slots): for i in range(1, number_of_slots):
slot_config[f"{i}"] = {} slot_config[f'{i}'] = {}
for tag in tags[last_index:last_index + tags_per_slot]: for tag in tags[last_index : last_index + tags_per_slot]:
try: try:
slot_config = build_tag_config( slot_config = build_tag_config(tag, slot_config.copy(), opc_servers, i)
tag, slot_config.copy(), opc_servers, i)
except ValueError as e: except ValueError as e:
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id="ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR", notification_id='ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR',
message=str(e), message=str(e),
block="orchestrator", block='orchestrator',
level=NotificationLevel.ERROR level=NotificationLevel.ERROR,
) )
last_index += tags_per_slot last_index += tags_per_slot
slot_config[f"{number_of_slots}"] = {} slot_config[f'{number_of_slots}'] = {}
for tag in tags[last_index:]: for tag in tags[last_index:]:
try: try:
slot_config = build_tag_config( slot_config = build_tag_config(
tag, slot_config.copy(), opc_servers, number_of_slots) tag, slot_config.copy(), opc_servers, number_of_slots
)
except ValueError as e: except ValueError as e:
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id="ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR", notification_id='ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR',
message=str(e), message=str(e),
block="orchestrator", block='orchestrator',
level=NotificationLevel.ERROR level=NotificationLevel.ERROR,
) )
self.info("Processed slots", metadata=metadata) self.info('Processed slots', metadata=metadata)
self.debug(json.dumps( self.debug(json.dumps(slot_config, indent=4, sort_keys=True), metadata=metadata)
slot_config, indent=4, sort_keys=True), metadata=metadata)
return slot_config return slot_config
@activity.defn(name="format_schedule_config") @activity.defn(name='format_schedule_config')
async def format_schedule_config(self, input_data: dict[str, Any]) -> dict[str, Any]: async def format_schedule_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Formats the schedule config to a dictionary with the schedule name as the key. Formats the schedule config to a dictionary with the schedule name as the key.
@@ -202,13 +210,13 @@ class Formatters(BaseActivity):
Returns: Returns:
dict[str, Any]: The formatted schedule config. dict[str, Any]: The formatted schedule config.
""" """
metadata = input_data["metadata"] metadata = input_data['metadata']
self.info("Formatting schedule config...", metadata=metadata) self.info('Formatting schedule config...', metadata=metadata)
schedule_config = input_data['schedule_config'] schedule_config = input_data['schedule_config']
config = {} config: dict[str, dict[str, str]] = {}
for schedule in schedule_config: for schedule in schedule_config:
namespace = schedule['namespace'] namespace = schedule['namespace']
schedule_name = schedule['schedule_name'] schedule_name = schedule['schedule_name']
@@ -219,47 +227,49 @@ class Formatters(BaseActivity):
config[namespace][schedule_name] = updated_at config[namespace][schedule_name] = updated_at
self.info("Formatted schedule config", metadata=metadata) self.info('Formatted schedule config', metadata=metadata)
self.debug(json.dumps( self.debug(json.dumps(config, indent=4, sort_keys=True), metadata=metadata)
config, indent=4, sort_keys=True), metadata=metadata)
return config return config
def compare_config_timestamps(self, def compare_config_timestamps(
schedules: dict[str, Any], current_schedules: dict[str, Any], self,
to_update: dict[str, Any], to_create: dict[str, Any], schedules: dict[str, Any],
namespace: str, metadata: dict[str, Any]): current_schedules: dict[str, Any],
to_update: dict[str, Any],
to_create: dict[str, Any],
namespace: str,
metadata: dict[str, Any],
):
""" """
Compares the timestamps of the schedule and the current schedule to determine Compare new and current schedules to determine which should be updated or created.
which schedules need to be updated or created.
Args: Args:
schedules (dict[str, Any]): The new schedules to compare. schedules (dict[str, Any]): New schedules to compare.
current_schedules (dict[str, Any]): The existing schedules to compare against. current_schedules (dict[str, Any]): Existing schedules to compare against.
to_update (dict[str, Any]): Dictionary to populate with schedules that need updating. to_update (dict[str, Any]): Output accumulator for schedules that need updating.
to_create (dict[str, Any]): Dictionary to populate with schedules that need creating. to_create (dict[str, Any]): Output accumulator for schedules that need creating.
namespace (str): The namespace for the schedules. namespace (str): Namespace for the schedules being compared.
metadata (dict[str, Any]): Metadata for logging purposes. metadata (dict[str, Any]): Metadata for logging purposes.
""" """
for schedule_name, schedule in schedules.items(): for schedule_name, schedule in schedules.items():
if schedule_name in current_schedules: if schedule_name in current_schedules:
update_timestamp = schedule.get('updated_at', now())
update_timestamp = schedule.get(
'updated_at', now())
old_timestamp = current_schedules[schedule_name] old_timestamp = current_schedules[schedule_name]
self.debug( self.debug(
f"Comparing schedule {schedule_name}:{update_timestamp} vs {old_timestamp}", metadata=metadata) f'Comparing schedule {schedule_name}:{update_timestamp} vs {old_timestamp}',
metadata=metadata,
)
if update_timestamp > old_timestamp: if update_timestamp > old_timestamp:
to_update[namespace][schedule_name] = schedule to_update[namespace][schedule_name] = schedule
else: else:
to_create[namespace][schedule_name] = schedule to_create[namespace][schedule_name] = schedule
@activity.defn(name="create_schedule_config") @activity.defn(name='create_schedule_config')
async def create_schedule_config(self, async def create_schedule_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Creates a schedule config dictionary based on the input data. Creates a schedule config dictionary based on the input data.
Checks the existing schedule config and updates it with the new schedule config, Checks the existing schedule config and updates it with the new schedule config,
@@ -273,54 +283,49 @@ class Formatters(BaseActivity):
- schedule_config (dict[str, Any]): The schedule config to process. - schedule_config (dict[str, Any]): The schedule config to process.
Returns: Returns:
- dict[str, Any]: The schedule config dictionary - dict[str, Any]: A dictionary with keys 'to_update', 'to_create', and 'to_delete'
""" """
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
self.info("Creating schedule config...", metadata=metadata) self.info('Creating schedule config...', metadata=metadata)
current_schedule_config = input_data['current_schedule_config'] current_schedule_config = input_data['current_schedule_config']
schedule_config = input_data['schedule_config'] schedule_config = input_data['schedule_config']
to_update = { to_update: dict[str, dict[str, Any]] = {
self.scouter_namespace: {}, self.scouter_namespace: {},
self.laborious_namespace: {} self.laborious_namespace: {},
} }
to_create = { to_create: dict[str, dict[str, Any]] = {
self.scouter_namespace: {}, self.scouter_namespace: {},
self.laborious_namespace: {} self.laborious_namespace: {},
} }
to_delete = { to_delete: dict[str, list[str]] = {
self.scouter_namespace: [], self.scouter_namespace: [],
self.laborious_namespace: [] self.laborious_namespace: [],
} }
for namespace, schedules in schedule_config.items(): for namespace, schedules in schedule_config.items():
current_schedules = current_schedule_config.get(namespace, {}) current_schedules = current_schedule_config.get(namespace, {})
self.compare_config_timestamps( self.compare_config_timestamps(
schedules, current_schedules, to_update, to_create, namespace, metadata) schedules, current_schedules, to_update, to_create, namespace, metadata
)
for namespace, schedules in current_schedule_config.items(): for namespace, schedules in current_schedule_config.items():
for schedule_name in schedules: for schedule_name in schedules:
if schedule_name not in schedule_config[namespace]: if schedule_name not in schedule_config[namespace]:
to_delete[namespace].append(schedule_name) to_delete[namespace].append(schedule_name)
output = { output = {'to_update': to_update, 'to_create': to_create, 'to_delete': to_delete}
"to_update": to_update,
"to_create": to_create,
"to_delete": to_delete
}
self.info("Created schedule config", metadata=metadata) self.info('Created schedule config', metadata=metadata)
self.debug(json.dumps( self.debug(json.dumps(output, indent=4, sort_keys=True), metadata=metadata)
output, indent=4, sort_keys=True), metadata=metadata)
return output return output
@activity.defn(name="create_slot_config") @activity.defn(name='create_slot_config')
async def create_slot_config(self, async def create_slot_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Creates a slot config dictionary based on the input data. Creates a slot config dictionary based on the input data.
Checks the existing slot config and updates it with the new slot config, Checks the existing slot config and updates it with the new slot config,
@@ -334,12 +339,12 @@ class Formatters(BaseActivity):
- slot_config (dict[str, Any]): The slot config to process. - slot_config (dict[str, Any]): The slot config to process.
Returns: Returns:
- dict[str, Any]: The slot config dictionary - dict[str, Any]: A dictionary containing 'to_insert' and 'to_delete'
""" """
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
self.info("Creating slot config...", metadata=metadata) self.info('Creating slot config...', metadata=metadata)
current_slot_config = input_data['current_slot_config'] current_slot_config = input_data['current_slot_config']
slot_config = input_data['slot_config'] slot_config = input_data['slot_config']
@@ -349,41 +354,43 @@ class Formatters(BaseActivity):
number_of_slots = len(slot_config) number_of_slots = len(slot_config)
if number_of_current_slots > number_of_slots: if number_of_current_slots > number_of_slots:
to_delete = [str(i) for i in range( to_delete = [str(i) for i in range(number_of_slots + 1, number_of_current_slots + 1)]
number_of_slots + 1, number_of_current_slots + 1)]
output = { output = {'to_delete': to_delete, 'to_insert': slot_config}
"to_delete": to_delete,
"to_insert": slot_config
}
self.info("Created slot config", metadata=metadata) self.info('Created slot config', metadata=metadata)
self.debug(json.dumps( self.debug(json.dumps(output, indent=4, sort_keys=True), metadata=metadata)
output, indent=4, sort_keys=True), metadata=metadata)
return output return output
def send_success_report(self, metadata: dict[str, Any], message: str, notification_id: str, attachment: str = None) -> None: def send_success_report(
self,
metadata: dict[str, Any],
message: str,
notification_id: str,
attachment: Any | None = None,
) -> None:
""" """
Sends a success notification report. Sends a success notification report.
Args: Args:
metadata (dict[str, Any]): Metadata for the notification. metadata (dict[str, Any]): Metadata for the notification.
message (str): The success message to send. message (str): The success message to send.
notification_id (str): The ID of the notification. notification_id (str): The ID of the notification to send.
attachment (str, optional): Optional attachment content for the notification. attachment (Any | None, optional): Optional attachment content to include.
""" """
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id=notification_id, notification_id=notification_id,
message=message, message=message,
block="report_orchestration", block='report_orchestration',
level=NotificationLevel.INFO, level=NotificationLevel.INFO,
attachment_content=json.dumps(attachment, indent=4, sort_keys=True) attachment_content=json.dumps(attachment, indent=4, sort_keys=True),
) )
def send_error_report(self, metadata: dict[str, Any], message: str, notification_id: str, def send_error_report(
attachment: str) -> None: self, metadata: dict[str, Any], message: str, notification_id: str, attachment: str
) -> None:
""" """
Sends an error notification report. Sends an error notification report.
@@ -397,37 +404,45 @@ class Formatters(BaseActivity):
metadata=metadata, metadata=metadata,
notification_id=notification_id, notification_id=notification_id,
message=message, message=message,
block="report_orchestration", block='report_orchestration',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=attachment attachment_content=attachment,
) )
def parse_report_schedule(self, input_data: dict[str, Any]) -> tuple[list[str], dict[str, Any]]: def parse_report_schedule(
self, input_data: list[dict[str, Any]]
) -> tuple[list[str], dict[str, Any]]:
""" """
Parses the report schedule data to extract success and error information. Parses the report schedule data to extract success and error information.
Args: Args:
input_data (dict[str, Any]): The input data containing schedule reports. input_data (list[dict[str, Any]]): The schedule reports. Each item must
Each item should have 'namespace', 'schedule_name', 'success', 'message', contain 'namespace', 'schedule_name', 'success', 'message', and optionally
and optionally 'attachment' fields. 'attachment'.
Returns: Returns:
tuple[list[str], dict[str, Any]]: A tuple containing: tuple[list[str], dict[str, Any]]: A tuple of:
- List of successful schedule keys in format "namespace/schedule_name" - Successful schedule keys in the form "namespace/schedule_name"
- Dictionary of error keys mapped to their error details - Error map keyed by the same string to error details
""" """
success_keys = [f"{value['namespace']}/{value['schedule_name']}" success_keys = [
for value in input_data if value['success']] f'{value["namespace"]}/{value["schedule_name"]}'
for value in input_data
if value['success']
]
error_keys = {f"{value['namespace']}/{value['schedule_name']}": { error_keys = {
f'{value["namespace"]}/{value["schedule_name"]}': {
'message': value['message'], 'message': value['message'],
'attachment': value.get('attachment', None) 'attachment': value.get('attachment', None),
}
for value in input_data
if not value['success']
} }
for value in input_data if not value['success']}
return success_keys, error_keys return success_keys, error_keys
def parse_report(self, input_data: dict[str, Any]) -> tuple[list[str], list[str]]: def parse_report(self, input_data: dict[str, dict[str, Any]]) -> tuple[list[str], list[str]]:
""" """
Parses the report data to extract success and error keys. Parses the report data to extract success and error keys.
@@ -440,15 +455,20 @@ class Formatters(BaseActivity):
- List of successful keys - List of successful keys
- List of error keys - List of error keys
""" """
success_keys = [key for key, value success_keys = [key for key, value in input_data.items() if value['success']]
in input_data.items() if value['success']]
error_keys = [key for key, value error_keys = [key for key, value in input_data.items() if not value['success']]
in input_data.items() if not value['success']]
return success_keys, error_keys return success_keys, error_keys
def manage_and_send_report(self, metadata: dict[str, Any], success_keys: list[str], error_keys: dict[str, Any], schedule_type: str, schedule_data: dict[str, Any]): def manage_and_send_report(
self,
metadata: dict[str, Any],
success_keys: list[str],
error_keys: dict[str, Any],
schedule_type: str,
schedule_data: dict[str, Any],
):
""" """
Manages and sends success and error reports based on the provided keys and data. Manages and sends success and error reports based on the provided keys and data.
@@ -462,30 +482,28 @@ class Formatters(BaseActivity):
if len(success_keys) > 0: if len(success_keys) > 0:
self.send_success_report( self.send_success_report(
metadata=metadata, metadata=metadata,
message=f"Successfully {schedule_type}: \n {', '.join(success_keys)}", message=f'Successfully {schedule_type}: \n {", ".join(success_keys)}',
notification_id=schedule_data['id'], notification_id=schedule_data['id'],
attachment=schedule_data['items'] attachment=schedule_data['items'],
) )
if len(error_keys) > 0: if len(error_keys) > 0:
attachment = [] attachment = []
for key, value in error_keys.items(): for key, value in error_keys.items():
if value['attachment'] is not None: if value['attachment'] is not None:
attachment.append( attachment.append(f'{key}:\n{value["message"]}\n{value["attachment"]}')
f"{key}:\n{value['message']}\n{value['attachment']}")
else: else:
attachment.append(f"{key}:\n{value['message']}") attachment.append(f'{key}:\n{value["message"]}')
self.send_error_report( self.send_error_report(
metadata=metadata, metadata=metadata,
message=f"Fails on {schedule_type}: \n {', '.join(error_keys)}", message=f'Fails on {schedule_type}: \n {", ".join(error_keys)}',
notification_id=f"{schedule_data['id']}_ERROR", notification_id=f'{schedule_data["id"]}_ERROR',
attachment=topic_separator.join(attachment) attachment=topic_separator.join(attachment),
) )
@activity.defn(name="report_schedule_orchestration") @activity.defn(name='report_schedule_orchestration')
async def report_schedule_orchestration(self, async def report_schedule_orchestration(self, input_data: dict[str, Any]) -> None:
input_data: dict[str, Any]) -> None:
""" """
Reports the orchestration result to the notification handler. Reports the orchestration result to the notification handler.
@@ -496,9 +514,9 @@ class Formatters(BaseActivity):
- updated_schedules (dict[str, Any]): The updated schedules. - updated_schedules (dict[str, Any]): The updated schedules.
- deleted_schedules (list[str]): The deleted schedules. - deleted_schedules (list[str]): The deleted schedules.
""" """
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
self.info("Reporting orchestration...", metadata=metadata) self.info('Reporting orchestration...', metadata=metadata)
created_schedules = input_data['created_schedules'] created_schedules = input_data['created_schedules']
updated_schedules = input_data['updated_schedules'] updated_schedules = input_data['updated_schedules']
@@ -507,34 +525,32 @@ class Formatters(BaseActivity):
schedules_report = { schedules_report = {
'created schedules': { 'created schedules': {
'items': created_schedules, 'items': created_schedules,
'id': 'REPORT_ORCHESTRATION_CREATED_SCHEDULES' 'id': 'REPORT_ORCHESTRATION_CREATED_SCHEDULES',
}, },
'updated schedules': { 'updated schedules': {
'items': updated_schedules, 'items': updated_schedules,
'id': 'REPORT_ORCHESTRATION_UPDATED_SCHEDULES' 'id': 'REPORT_ORCHESTRATION_UPDATED_SCHEDULES',
}, },
'deleted schedules': { 'deleted schedules': {
'items': deleted_schedules, 'items': deleted_schedules,
'id': 'REPORT_ORCHESTRATION_DELETED_SCHEDULES' 'id': 'REPORT_ORCHESTRATION_DELETED_SCHEDULES',
} },
} }
for schedule_type, schedule_data in schedules_report.items(): for schedule_type, schedule_data in schedules_report.items():
if len(schedule_data['items']) > 0: if len(schedule_data['items']) > 0:
success_keys, error_keys = self.parse_report_schedule( success_keys, error_keys = self.parse_report_schedule(schedule_data['items'])
schedule_data['items'])
self.manage_and_send_report( self.manage_and_send_report(
metadata=metadata, metadata=metadata,
success_keys=success_keys, success_keys=success_keys,
error_keys=error_keys, error_keys=error_keys,
schedule_type=schedule_type, schedule_type=schedule_type,
schedule_data=schedule_data schedule_data=schedule_data,
) )
@activity.defn(name="report_slot_orchestration") @activity.defn(name='report_slot_orchestration')
async def report_slot_orchestration(self, async def report_slot_orchestration(self, input_data: dict[str, Any]) -> None:
input_data: dict[str, Any]) -> None:
""" """
Reports the orchestration result to the notification handler. Reports the orchestration result to the notification handler.
@@ -545,9 +561,9 @@ class Formatters(BaseActivity):
- deleted_slots (list[str]): The deleted slots. - deleted_slots (list[str]): The deleted slots.
""" """
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
self.info("Reporting orchestration...", metadata=metadata) self.info('Reporting orchestration...', metadata=metadata)
inserted_slots = input_data['inserted_slots'] inserted_slots = input_data['inserted_slots']
deleted_slots = input_data['deleted_slots'] deleted_slots = input_data['deleted_slots']
@@ -558,16 +574,16 @@ class Formatters(BaseActivity):
if len(success_keys) > 0: if len(success_keys) > 0:
self.send_success_report( self.send_success_report(
metadata=metadata, metadata=metadata,
message=f"Inserted slots: \n {', '.join(success_keys)}", message=f'Inserted slots: \n {", ".join(success_keys)}',
notification_id="REPORT_ORCHESTRATION_INSERTED_SLOTS" notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS',
) )
if len(error_keys) > 0: if len(error_keys) > 0:
self.send_error_report( self.send_error_report(
metadata=metadata, metadata=metadata,
message=f"Failed to insert slots: \n {', '.join(error_keys)}", message=f'Failed to insert slots: \n {", ".join(error_keys)}',
notification_id="REPORT_ORCHESTRATION_INSERTED_SLOTS", notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS',
attachment=inserted_slots attachment=inserted_slots,
) )
if len(deleted_slots) > 0: if len(deleted_slots) > 0:
@@ -576,20 +592,20 @@ class Formatters(BaseActivity):
if len(success_keys) > 0: if len(success_keys) > 0:
self.send_success_report( self.send_success_report(
metadata=metadata, metadata=metadata,
message=f"Deleted slots: \n {', '.join(success_keys)}", message=f'Deleted slots: \n {", ".join(success_keys)}',
notification_id="REPORT_ORCHESTRATION_DELETED_SLOTS" notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS',
) )
if len(error_keys) > 0: if len(error_keys) > 0:
self.send_error_report( self.send_error_report(
metadata=metadata, metadata=metadata,
message=f"Failed to delete slots: \n {', '.join(error_keys)}", message=f'Failed to delete slots: \n {", ".join(error_keys)}',
notification_id="REPORT_ORCHESTRATION_DELETED_SLOTS", notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS',
attachment=deleted_slots attachment=deleted_slots,
) )
@activity.defn(name="format_log_report") @activity.defn(name='format_log_report')
async def format_log_report(self, input_data: dict[str, Any]) -> dict[str, Any]: async def format_log_report(self, input_data: dict[str, Any]) -> dict[Hashable, Any]:
""" """
Formats the receiver_groups status to a dataframe to be stored in the database. Formats the receiver_groups status to a dataframe to be stored in the database.
@@ -602,23 +618,21 @@ class Formatters(BaseActivity):
Returns: Returns:
dict[str, 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"] metadata = input_data['metadata']
mail_type = input_data["mail_type"] mail_type = input_data['mail_type']
self.info("Formatting log report...", metadata=metadata) self.info('Formatting log report...', metadata=metadata)
receiver_groups = input_data['receiver_groups'] receiver_groups = input_data['receiver_groups']
data = {} data = {}
for group_name, group_config in receiver_groups.items(): for group_name, group_config in receiver_groups.items():
for notification in group_config['notifications']: for notification in group_config['notifications']:
notification_id = notification['notification_id'] notification_id = notification['notification_id']
trigger = notification['trigger'] trigger = notification['trigger']
key = f"{notification_id}:{trigger}" key = f'{notification_id}:{trigger}'
if key not in data: if key not in data:
data[key] = { data[key] = {
@@ -634,15 +648,17 @@ class Formatters(BaseActivity):
'project': notification['project'], 'project': notification['project'],
'model_name': notification['model_name'], 'model_name': notification['model_name'],
'model_id': notification['model_id'], 'model_id': notification['model_id'],
'mail_type': mail_type 'mail_type': mail_type,
} }
else: else:
if group_name not in data[key]['groups']: if group_name not in data[key]['groups']:
data[key]['groups'].append(group_name) data[key]['groups'].append(group_name)
return DataFrame(list(data.values())).to_dict() data_values: DataFrame = DataFrame(list(data.values()))
@activity.defn(name="filter_notification_reports") return data_values.to_dict()
@activity.defn(name='filter_notification_reports')
async def filter_notification_reports(self, input_data: dict[str, Any]) -> dict[str, Any]: async def filter_notification_reports(self, input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Filter notifications for comprehensive scheduled reports. Filter notifications for comprehensive scheduled reports.
@@ -678,16 +694,13 @@ class Formatters(BaseActivity):
notification_package = input_data['notification_package'] notification_package = input_data['notification_package']
sending_configs = input_data['sending_configs'] sending_configs = input_data['sending_configs']
self.info("Filtering notification reports...", metadata=metadata) self.info('Filtering notification reports...', metadata=metadata)
receiver_groups = {} receiver_groups = {}
for receiver_group in sending_configs: for receiver_group in sending_configs:
group_name = receiver_group['group_name'] group_name = receiver_group['group_name']
receiver_groups[group_name] = { receiver_groups[group_name] = {**receiver_group, 'notifications': []}
**receiver_group,
"notifications": []
}
receiver_groups[group_name]['notifications'] = [] receiver_groups[group_name]['notifications'] = []
already_added_keys = [] already_added_keys = []
@@ -695,15 +708,18 @@ class Formatters(BaseActivity):
ignore_list = receiver_group.get('ignore', []) ignore_list = receiver_group.get('ignore', [])
for notification in notification_package: for notification in notification_package:
alert_type = "reports" alert_type = 'reports'
notification_id = notification['notification_id'] notification_id = notification['notification_id']
key = f"{notification['trigger']}:{notification_id}" key = f'{notification["trigger"]}:{notification_id}'
# Check if this group must be notified # Check if this group must be notified
if alert_type in receiver_group['contents'] and notification_id not in ignore_list and key not in already_added_keys: if (
receiver_groups[group_name]["notifications"].append( alert_type in receiver_group['contents']
notification) and notification_id not in ignore_list
and key not in already_added_keys
):
receiver_groups[group_name]['notifications'].append(notification)
already_added_keys.append(key) already_added_keys.append(key)
# Remove groups with no notifications # Remove groups with no notifications

382
orchestrator/activities/mongo_db.py
View File

@@ -1,37 +1,39 @@
from temporalio import workflow, activity from datetime import UTC
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from typing import Any
import traceback import traceback
from datetime import datetime
from logging import Logger from logging import Logger
from typing import Any
from pymongo import MongoClient from pymongo import MongoClient
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.base import BaseActivity from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
from datetime import datetime, timezone
def clear_mongo_id(docs: list) -> list: def clear_mongo_id(docs: list) -> list:
""" """
Remove the MongoDB internal `_id` field from the document. Remove MongoDB internal `_id` fields from nested structures.
Args: Args:
docs (list): The document to clear. docs (list): The list of documents or nested structures to clean.
Returns: Returns:
list: The documents without the `_id` field. list: The cleaned documents with `_id` fields removed wherever present.
""" """
for doc in docs: for doc in docs:
if isinstance(doc, list): if isinstance(doc, list):
clear_mongo_id(doc) clear_mongo_id(doc)
elif isinstance(doc, dict): elif isinstance(doc, dict):
if "_id" in doc: if '_id' in doc:
del doc["_id"] del doc['_id']
for key, value in doc.items(): for _key, value in doc.items():
if isinstance(value, list): if isinstance(value, list):
clear_mongo_id(value) clear_mongo_id(value)
elif isinstance(value, dict): elif isinstance(value, dict):
@@ -46,8 +48,8 @@ class MongoDB(BaseActivity):
This class provides MongoDB database operations including document This class provides MongoDB database operations including document
querying, aggregation, timestamp management, and collection management querying, aggregation, timestamp management, and collection management
with TTL indexes. It handles all MongoDB interactions required by with TTL indexes. It centralizes all MongoDB interactions required by
the orchestration system. the orchestration system lifecycle.
Args: Args:
connection_string (str): MongoDB connection string connection_string (str): MongoDB connection string
@@ -57,42 +59,46 @@ class MongoDB(BaseActivity):
notification_handler (NotificationHandler): Notification management handler notification_handler (NotificationHandler): Notification management handler
""" """
def __init__(self, connection_string: str, database_name: str, ttl_index_seconds: int, def __init__(
self,
connection_string: str,
database_name: str,
ttl_index_seconds: int,
logger: Logger, logger: Logger,
notification_handler: NotificationHandler): notification_handler: NotificationHandler,
):
self.connection_string = connection_string self.connection_string = connection_string
self.database_name = database_name self.database_name = database_name
self.client = MongoClient( self.client: MongoClient = MongoClient(
self.connection_string, serverSelectionTimeoutMS=5000) self.connection_string, serverSelectionTimeoutMS=5000
self.client.server_info() # Trigger an exception if connection fails )
self.client.server_info() # Force early failure if connection is invalid
self.database = self.client[self.database_name] self.database = self.client[self.database_name]
self.ttl_index_seconds = ttl_index_seconds self.ttl_index_seconds = ttl_index_seconds
# Initialize MongoDB client here (omitted for brevity) # Initialize MongoDB client here (omitted for brevity)
logger.info("MongoDB connection initialized") logger.info('MongoDB connection initialized')
BaseActivity.__init__(self, BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler)
logger=logger,
notification_handler=notification_handler)
def shutdown(self): def shutdown(self):
""" """
Shutdown the MongoDB connection and clean up resources. Shutdown the MongoDB client and clean up resources.
""" """
try: try:
if self.client: if self.client:
self.logger.info("Closing MongoDB connection...") self.logger.info('Closing MongoDB connection...')
self.client.close() self.client.close()
self.logger.info("MongoDB connection closed successfully") self.logger.info('MongoDB connection closed successfully')
except Exception as e: except Exception as e:
self.logger.error(f"Failed to close MongoDB connection: {e}") self.logger.error(f'Failed to close MongoDB connection: {e}')
def __del__(self): def __del__(self):
""" """
Destructor to ensure MongoDB client is closed when the object is deleted. Ensure the MongoDB client is closed when the object is garbage-collected.
""" """
self.shutdown() self.shutdown()
@@ -105,17 +111,19 @@ class MongoDB(BaseActivity):
filters (dict[str, Any]): The query filters to apply. filters (dict[str, Any]): The query filters to apply.
Returns: Returns:
list[dict[str, Any]]: List of documents matching the filters, with _id fields removed. list[dict[str, Any]]: Documents matching the filters (with `_id` removed).
""" """
collection = self.database[collection_name] collection = self.database[collection_name]
documents = list(collection.find(filters, {"_id": 0})) documents = list(collection.find(filters, {'_id': 0}))
documents = clear_mongo_id(documents) documents = clear_mongo_id(documents)
return documents return documents
@activity.defn(name="find_documents_in_mongodb",) @activity.defn(
name='find_documents_in_mongodb',
)
async def find_documents_in_mongodb(self, input_data: dict[str, Any]) -> list[dict[str, Any]]: async 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. Find documents in a MongoDB collection based on the provided query parameters.
@@ -123,38 +131,45 @@ class MongoDB(BaseActivity):
Args: Args:
- input_data (dict): Input data containing query parameters. Contains: - input_data (dict): Input data containing query parameters. Contains:
- query (dict): Query parameters to filter documents. - query (dict): Query parameters to filter documents.
- timestamp_fields (list[str], optional): Fields to format as RFC3339 with TZ.
Returns: Returns:
list[dict]: List of documents matching the query. list[dict]: Documents matching the query with timestamp fields normalized.
""" """
query = input_data.get("query", {}) query = input_data.get('query', {})
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
timestamp_fields = input_data.get("timestamp_fields", []) timestamp_fields = input_data.get('timestamp_fields', [])
collection_name = query.get("collection") collection_name = query.get('collection')
if not collection_name: if not collection_name:
raise ValueError("Collection name must be provided in the query.") raise ValueError('Collection name must be provided in the query.')
filters = query.get("filters", {}) filters = query.get('filters', {})
self.info( self.info(
f"Loading documents from collection '{collection_name}' with filters: {filters}", metadata=metadata) f"Loading documents from collection '{collection_name}' with filters: {filters}",
metadata=metadata,
)
try: try:
documents = self.find(collection_name, filters) documents = self.find(collection_name, filters)
self.info( self.info(
f"Loaded {len(documents)} documents from collection '{collection_name}'", metadata=metadata) f"Loaded {len(documents)} documents from collection '{collection_name}'",
metadata=metadata,
)
for document in documents: for document in documents:
for timestamp_field in timestamp_fields: for timestamp_field in timestamp_fields:
if timestamp_field in document: if timestamp_field in document:
document[timestamp_field] = document[timestamp_field].replace(tzinfo=timezone.utc).strftime( document[timestamp_field] = (
DATETIME_FORMAT_MS_WITH_TZ) document[timestamp_field]
.replace(tzinfo=UTC)
.strftime(DATETIME_FORMAT_MS_WITH_TZ)
)
self.debug( self.debug(f'Documents loaded: {documents}', metadata=metadata)
f"Documents loaded: {documents}", metadata=metadata)
return documents return documents
@@ -162,64 +177,71 @@ class MongoDB(BaseActivity):
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id="MONGODB_QUERY_ERROR", notification_id='MONGODB_QUERY_ERROR',
message=f"Failed to execute MongoDB query: {e}", message=f'Failed to execute MongoDB query: {e}',
block="load_query_from_mongodb", block='load_query_from_mongodb',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
self.error(trace, metadata=metadata) self.error(trace, metadata=metadata)
raise e raise e
@activity.defn(name="aggregate_documents_in_mongodb") @activity.defn(name='aggregate_documents_in_mongodb')
async def aggregate_documents_in_mongodb(self, async def aggregate_documents_in_mongodb(
input_data: dict[str, Any]) -> list[dict[str, Any]]: self, input_data: dict[str, Any]
) -> list[dict[str, Any]]:
""" """
Aggregate documents in a MongoDB collection based on the provided aggregation pipeline. Aggregate documents in a MongoDB collection based on the provided aggregation pipeline.
Args: Args:
- input_data (dict): Input data containing aggregation parameters. Contains: - input_data (dict): Input data containing aggregation parameters. Contains:
- query (dict): Query parameters to filter documents. - query (dict): Aggregation parameters including 'collection' and 'aggregation'.
- timestamp_fields (list[str], optional): Fields to format as RFC3339 with TZ.
Returns: Returns:
list[dict]: List of aggregated documents. list[dict]: Aggregated documents with timestamp fields normalized.
""" """
query = input_data.get("query", {}) query = input_data.get('query', {})
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
timestamp_fields = input_data.get("timestamp_fields", []) timestamp_fields = input_data.get('timestamp_fields', [])
collection_name = query.get("collection") collection_name = query.get('collection')
if not collection_name: if not collection_name:
raise ValueError("Collection name must be provided in the query.") raise ValueError('Collection name must be provided in the query.')
aggregation = query.get("aggregation") aggregation = query.get('aggregation')
if not aggregation: if not aggregation:
raise ValueError("Aggregation must be provided.") raise ValueError('Aggregation must be provided.')
aggregation.append({"$project": {"_id": 0}}) aggregation.append({'$project': {'_id': 0}})
self.info( self.info(
f"Aggregating documents from collection '{collection_name}' with aggregation: {aggregation}", metadata=metadata) f"Aggregating documents from collection '{collection_name}' with aggregation: {aggregation}",
metadata=metadata,
)
try: try:
collection = self.database[collection_name] collection = self.database[collection_name]
aggregated_documents = list( aggregated_documents = list(collection.aggregate(aggregation))
collection.aggregate(aggregation))
aggregated_documents = clear_mongo_id(aggregated_documents) aggregated_documents = clear_mongo_id(aggregated_documents)
self.info( self.info(
f"Aggregated {len(aggregated_documents)} documents from collection '{collection_name}'", metadata=metadata) f"Aggregated {len(aggregated_documents)} documents from collection '{collection_name}'",
metadata=metadata,
)
for document in aggregated_documents: for document in aggregated_documents:
for timestamp_field in timestamp_fields: for timestamp_field in timestamp_fields:
if timestamp_field in document: if timestamp_field in document:
document[timestamp_field] = document[timestamp_field].replace(tzinfo=timezone.utc).strftime( document[timestamp_field] = (
DATETIME_FORMAT_MS_WITH_TZ) document[timestamp_field]
.replace(tzinfo=UTC)
.strftime(DATETIME_FORMAT_MS_WITH_TZ)
)
self.debug( self.debug(f'Aggregation result: {aggregated_documents}', metadata=metadata)
f"Aggregation result: {aggregated_documents}", metadata=metadata)
return aggregated_documents return aggregated_documents
@@ -227,83 +249,87 @@ class MongoDB(BaseActivity):
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id="MONGODB_AGGREGATION_ERROR", notification_id='MONGODB_AGGREGATION_ERROR',
message=f"Failed to execute MongoDB aggregation: {e}", message=f'Failed to execute MongoDB aggregation: {e}',
block="aggregate_documents_in_mongodb", block='aggregate_documents_in_mongodb',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
self.error(trace, metadata=metadata) self.error(trace, metadata=metadata)
raise e raise e
@activity.defn(name="update_pipelines_timestamps") @activity.defn(name='update_pipelines_timestamps')
async def update_pipelines_timestamps(self, input_data: dict[str, Any]) -> None: async def update_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
""" """
Update the timestamps of the pipelines in the MongoDB collection. Update `updated_at` timestamps for successfully updated pipelines.
input_data:
- updated_pipelines (list): List of updated pipelines.
"""
updated_pipelines = input_data.get("updated_pipelines", [])
metadata = input_data.get("metadata", {})
date_now = now()
collection = self.database["orchestrated_schedules"]
self.info("Updating pipelines timestamps...", metadata=metadata) input_data:
- updated_pipelines (list[dict]): Pipelines with success flags to consider.
"""
updated_pipelines = input_data.get('updated_pipelines', [])
metadata = input_data.get('metadata', {})
date_now = now()
collection = self.database['orchestrated_schedules']
self.info('Updating pipelines timestamps...', metadata=metadata)
success_count = 0 success_count = 0
argument = [ argument = [
{"schedule_name": pipeline["schedule_name"], {'schedule_name': pipeline['schedule_name'], 'namespace': pipeline['namespace']}
"namespace": pipeline["namespace"]} for pipeline in updated_pipelines
for pipeline in updated_pipelines if pipeline["success"] if pipeline['success']
] ]
data_filter = {"$or": argument} if argument else {} data_filter = {'$or': argument} if argument else {}
try: try:
collection.update_many( collection.update_many(data_filter, {'$set': {'updated_at': date_now}})
data_filter,
{"$set": {"updated_at": date_now}}
)
success_count += 1 success_count += 1
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id="MONGODB_UPDATE_PIPELINES_ERROR", notification_id='MONGODB_UPDATE_PIPELINES_ERROR',
message=f"Failed to update pipelines timestamps: {e}", message=f'Failed to update pipelines timestamps: {e}',
block="update_pipelines_timestamps", block='update_pipelines_timestamps',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
self.error(trace, metadata=metadata) self.error(trace, metadata=metadata)
raise e raise e
self.info( self.info(
f"Updated {success_count} of {len(updated_pipelines)} pipelines timestamps", metadata=metadata) f'Updated {success_count} of {len(updated_pipelines)} pipelines timestamps',
metadata=metadata,
)
@activity.defn(name="create_pipelines_timestamps") @activity.defn(name='create_pipelines_timestamps')
async def create_pipelines_timestamps(self, input_data: dict[str, Any]) -> None: async def create_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
""" """
Create the timestamps of the pipelines in the MongoDB collection. Insert `updated_at` timestamps for newly created pipelines.
input_data:
- created_pipelines (list): List of created pipelines.
"""
created_pipelines = input_data.get("created_pipelines", [])
metadata = input_data.get("metadata", {})
collection = self.database["orchestrated_schedules"]
self.info("Creating pipelines timestamps...", metadata=metadata) input_data:
- created_pipelines (list[dict]): Pipelines with success flags to consider.
"""
created_pipelines = input_data.get('created_pipelines', [])
metadata = input_data.get('metadata', {})
collection = self.database['orchestrated_schedules']
self.info('Creating pipelines timestamps...', metadata=metadata)
success_count = 0 success_count = 0
date_now = now() date_now = now()
argument = [ argument = [
{"schedule_name": pipeline["schedule_name"], {
"namespace": pipeline["namespace"], 'schedule_name': pipeline['schedule_name'],
"updated_at": date_now} 'namespace': pipeline['namespace'],
for pipeline in created_pipelines if pipeline["success"] 'updated_at': date_now,
}
for pipeline in created_pipelines
if pipeline['success']
] ]
data_filter = argument if argument else {} data_filter = argument if argument else {}
@@ -315,39 +341,42 @@ class MongoDB(BaseActivity):
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id="MONGODB_CREATE_PIPELINES_ERROR", notification_id='MONGODB_CREATE_PIPELINES_ERROR',
message=f"Failed to create pipelines timestamps: {e}", message=f'Failed to create pipelines timestamps: {e}',
block="create_pipelines_timestamps", block='create_pipelines_timestamps',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
self.error(trace, metadata=metadata) self.error(trace, metadata=metadata)
raise e raise e
self.info( self.info(
f"Created {success_count} of {len(created_pipelines)} pipelines timestamps", metadata=metadata) f'Created {success_count} of {len(created_pipelines)} pipelines timestamps',
metadata=metadata,
)
@activity.defn(name="delete_pipelines_timestamps") @activity.defn(name='delete_pipelines_timestamps')
async def delete_pipelines_timestamps(self, input_data: dict[str, Any]) -> None: async def delete_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
""" """
Delete the timestamps of the pipelines in the MongoDB collection. Delete timestamp rows for successfully deleted pipelines.
input_data:
- deleted_pipelines (list): List of deleted pipelines.
"""
deleted_pipelines = input_data.get("deleted_pipelines", [])
metadata = input_data.get("metadata", {})
collection = self.database["orchestrated_schedules"]
self.info("Deleting pipelines timestamps...", metadata=metadata) input_data:
- deleted_pipelines (list[dict]): Pipelines with success flags to consider.
"""
deleted_pipelines = input_data.get('deleted_pipelines', [])
metadata = input_data.get('metadata', {})
collection = self.database['orchestrated_schedules']
self.info('Deleting pipelines timestamps...', metadata=metadata)
success_count = 0 success_count = 0
argument = [ argument = [
{"schedule_name": pipeline["schedule_name"], {'schedule_name': pipeline['schedule_name'], 'namespace': pipeline['namespace']}
"namespace": pipeline["namespace"]} for pipeline in deleted_pipelines
for pipeline in deleted_pipelines if pipeline["success"] if pipeline['success']
] ]
data_filter = {"$or": argument} if argument else {} data_filter = {'$or': argument} if argument else {}
try: try:
collection.delete_many(data_filter) collection.delete_many(data_filter)
@@ -356,19 +385,21 @@ class MongoDB(BaseActivity):
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id="MONGODB_DELETE_PIPELINES_ERROR", notification_id='MONGODB_DELETE_PIPELINES_ERROR',
message=f"Failed to delete pipelines timestamps: {e}", message=f'Failed to delete pipelines timestamps: {e}',
block="delete_pipelines_timestamps", block='delete_pipelines_timestamps',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
self.error(trace, metadata=metadata) self.error(trace, metadata=metadata)
raise e raise e
self.info( self.info(
f"Deleted {success_count} of {len(deleted_pipelines)} pipelines timestamps", metadata=metadata) f'Deleted {success_count} of {len(deleted_pipelines)} pipelines timestamps',
metadata=metadata,
)
@activity.defn(name="create_collection_with_ttl_index") @activity.defn(name='create_collection_with_ttl_index')
async def create_collection_with_ttl_index(self, input_data: dict[str, Any]) -> None: async def create_collection_with_ttl_index(self, input_data: dict[str, Any]) -> None:
""" """
Create a collection with a TTL index. Create a collection with a TTL index.
@@ -376,20 +407,21 @@ class MongoDB(BaseActivity):
- collection_name (str): The name of the collection to create. - collection_name (str): The name of the collection to create.
- ttl_index (str): The name of the TTL index to create. - ttl_index (str): The name of the TTL index to create.
""" """
pipelines = input_data.get("pipelines", {}) pipelines = input_data.get('pipelines', {})
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
self.info( self.info(
f"Creating collection with TTL index for pipelines: {list(pipelines.keys())}", f'Creating collection with TTL index for pipelines: {list(pipelines.keys())}',
metadata=metadata) metadata=metadata,
)
collection_names = self.database.list_collection_names() collection_names = self.database.list_collection_names()
created_collections = [] created_collections = []
created_indexes = [] created_indexes = []
for pipeline_name, pipeline_config in pipelines.items(): for _pipeline_name, pipeline_config in pipelines.items():
collection = pipeline_config["topic"] collection = pipeline_config['topic']
try: try:
# Check if collection exists # Check if collection exists
@@ -402,16 +434,17 @@ class MongoDB(BaseActivity):
existing_indexes = collection.list_indexes() existing_indexes = collection.list_indexes()
ttl_index_exists = False ttl_index_exists = False
for index in existing_indexes: for index in existing_indexes:
if "inserted_at" in index["key"] and index.get("expireAfterSeconds") is not None: if (
'inserted_at' in index['key']
and index.get('expireAfterSeconds') is not None
):
ttl_index_exists = True ttl_index_exists = True
break break
# Create TTL index if it doesn't exist # Create TTL index if it doesn't exist
if not ttl_index_exists: if not ttl_index_exists:
collection.create_index( collection.create_index(
"inserted_at", 'inserted_at', expireAfterSeconds=self.ttl_index_seconds, background=True
expireAfterSeconds=self.ttl_index_seconds,
background=True
) )
created_indexes.append(collection) created_indexes.append(collection)
@@ -419,30 +452,24 @@ class MongoDB(BaseActivity):
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id="MONGODB_CREATE_COLLECTION_ERROR", notification_id='MONGODB_CREATE_COLLECTION_ERROR',
message=f"Failed to create collection {collection} with TTL index: {e}", message=f'Failed to create collection {collection} with TTL index: {e}',
block="create_collection_with_ttl_index", block='create_collection_with_ttl_index',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
self.error(trace, metadata=metadata) self.error(trace, metadata=metadata)
raise e raise e
self.info( self.info(
f"Created {len(created_collections)} collections and {len(created_indexes)} indexes", f'Created {len(created_collections)} collections and {len(created_indexes)} indexes',
metadata=metadata metadata=metadata,
) )
self.debug( self.debug(f'Created collections: {created_collections}', metadata=metadata)
f"Created collections: {created_collections}", self.debug(f'Created indexes: {created_indexes}', metadata=metadata)
metadata=metadata
)
self.debug(
f"Created indexes: {created_indexes}",
metadata=metadata
)
@activity.defn(name="load_latest_data") @activity.defn(name='load_latest_data')
async def load_latest_data(self, input_data: dict[str, Any]) -> list[dict[str, Any]]: async 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. Load the latest data from MongoDB collection since a specified timestamp.
@@ -470,58 +497,43 @@ class MongoDB(BaseActivity):
last_data_timestamp = input_data['last_data_timestamp'] last_data_timestamp = input_data['last_data_timestamp']
base_data_filter = input_data['base_data_filter'] base_data_filter = input_data['base_data_filter']
self.debug( self.debug(f'Loading data from MongoDB: {input_data}', metadata=metadata)
f"Loading data from MongoDB: {input_data}",
metadata=metadata
)
try: try:
if last_data_timestamp is None: if last_data_timestamp is None:
data_filter = base_data_filter data_filter = base_data_filter
else: else:
data_filter = { data_filter = {
**base_data_filter, **base_data_filter,
"timestamp": { 'timestamp': {
"$gt": datetime.strptime(last_data_timestamp, DATETIME_FORMAT_MS_WITH_TZ) '$gt': datetime.strptime(last_data_timestamp, DATETIME_FORMAT_MS_WITH_TZ)
} },
} }
self.debug( self.debug(f'Data filter: {data_filter}', metadata=metadata)
f"Data filter: {data_filter}",
metadata=metadata
)
data = self.find(collection_name, data_filter) data = self.find(collection_name, data_filter)
self.debug( self.debug(f'Collected: {data}', metadata=metadata)
f"Collected: {data}",
metadata=metadata
)
for item in data: for item in data:
item['timestamp'] = item['timestamp'].replace(tzinfo=timezone.utc).strftime( item['timestamp'] = (
DATETIME_FORMAT_MS_WITH_TZ) item['timestamp'].replace(tzinfo=UTC).strftime(DATETIME_FORMAT_MS_WITH_TZ)
self.info(
f"Loaded {len(data)} documents from MongoDB",
metadata=metadata
) )
self.debug( self.info(f'Loaded {len(data)} documents from MongoDB', metadata=metadata)
f"Loaded data: {data}",
metadata=metadata self.debug(f'Loaded data: {data}', metadata=metadata)
)
return data return data
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id="MONGO_LOAD_ERROR", notification_id='MONGO_LOAD_ERROR',
message=f"Error loading data from MongoDB: {e}", message=f'Error loading data from MongoDB: {e}',
block="load_latest_data", block='load_latest_data',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
raise e raise e

206
orchestrator/activities/slot_manager.py
View File

@@ -1,17 +1,17 @@
from temporalio import activity, workflow from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from typing import Any
import traceback
import json import json
import traceback
from datetime import datetime, timedelta
from logging import Logger from logging import Logger
from sientia_do.temporal.activities.redis_base import Redis from typing import Any
from pandas import DataFrame
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.redis_base import Redis
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
from pandas import DataFrame
from datetime import datetime, timedelta
class SlotManager(Redis): class SlotManager(Redis):
@@ -40,14 +40,18 @@ class SlotManager(Redis):
notification_handler (NotificationHandler): Notification management handler notification_handler (NotificationHandler): Notification management handler
""" """
def __init__(self, host: str, port: int, def __init__(
username: str, password: str, self,
logger: Logger, notification_handler: NotificationHandler): host: str,
port: int,
username: str,
password: str,
logger: Logger,
notification_handler: NotificationHandler,
):
Redis.__init__(self, host, port, username, password, logger, notification_handler)
Redis.__init__(self, host, port, username, @activity.defn(name='load_opc_slots')
password, logger, notification_handler)
@activity.defn(name="load_opc_slots")
async def load_opc_slots(self, input_data: dict[str, Any]) -> dict[str, Any]: async def load_opc_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Load all OPC slots from Redis for current system state assessment. Load all OPC slots from Redis for current system state assessment.
@@ -69,17 +73,16 @@ class SlotManager(Redis):
Exception: If Redis connection fails or data retrieval errors occur Exception: If Redis connection fails or data retrieval errors occur
""" """
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
self.info("Loading OPC slots...", metadata=metadata) self.info('Loading OPC slots...', metadata=metadata)
opc_slots = {} opc_slots = {}
try: try:
slot_keys = self.redis_client.keys('slot:opc_tags:*')
slot_keys = self.redis_client.keys("slot:opc_tags:*") self.debug(f'Slot keys: {slot_keys}', metadata=metadata)
self.debug(f"Slot keys: {slot_keys}", metadata=metadata)
if slot_keys: if slot_keys:
if isinstance(slot_keys[0], bytes): if isinstance(slot_keys[0], bytes):
@@ -93,20 +96,20 @@ class SlotManager(Redis):
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id="REDIS_GET_ERROR", notification_id='REDIS_GET_ERROR',
message=f"Failed to load OPC slots: {e}", message=f'Failed to load OPC slots: {e}',
block="load_opc_slots", block='load_opc_slots',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
self.error(trace, metadata=metadata) self.error(trace, metadata=metadata)
raise e raise e
self.info(f"Loaded {len(opc_slots)} OPC slots", metadata=metadata) self.info(f'Loaded {len(opc_slots)} OPC slots', metadata=metadata)
return opc_slots return opc_slots
@activity.defn(name="load_active_ingestors") @activity.defn(name='load_active_ingestors')
async def load_active_ingestors(self, input_data: dict[str, Any]) -> list[str]: async def load_active_ingestors(self, input_data: dict[str, Any]) -> list[str]:
""" """
Load all active ingestors from Redis Load all active ingestors from Redis
@@ -115,19 +118,16 @@ class SlotManager(Redis):
list[str]: A list of active ingestors list[str]: A list of active ingestors
""" """
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
self.info("Loading active ingestors...", metadata=metadata) self.info('Loading active ingestors...', metadata=metadata)
try: try:
active_ingestors = self.redis_client.keys('heartbeat:ingestor:*')
active_ingestors = self.redis_client.keys("heartbeat:ingestor:*") self.info(f'Loaded {len(active_ingestors)} active ingestors', metadata=metadata)
self.info( self.debug(f'Active ingestors: \n {active_ingestors}', metadata=metadata)
f"Loaded {len(active_ingestors)} active ingestors", metadata=metadata)
self.debug(
f"Active ingestors: \n {active_ingestors}", metadata=metadata)
ingestors = [] ingestors = []
@@ -142,16 +142,16 @@ class SlotManager(Redis):
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id="REDIS_GET_ERROR", notification_id='REDIS_GET_ERROR',
message=f"Failed to load active ingestors: {e}", message=f'Failed to load active ingestors: {e}',
block="load_active_ingestors", block='load_active_ingestors',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
self.error(trace, metadata=metadata) self.error(trace, metadata=metadata)
raise e raise e
@activity.defn(name="update_slots") @activity.defn(name='update_slots')
async def update_slots(self, input_data: dict[str, Any]) -> dict[str, Any]: async def update_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Update OPC slots in Redis Update OPC slots in Redis
@@ -166,8 +166,8 @@ class SlotManager(Redis):
""" """
to_insert = input_data['to_insert'] to_insert = input_data['to_insert']
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
self.info("Updating OPC slots...", metadata=metadata) self.info('Updating OPC slots...', metadata=metadata)
report = {} report = {}
@@ -175,30 +175,20 @@ class SlotManager(Redis):
for slot in to_insert: for slot in to_insert:
try: try:
self.set(f"slot:opc_tags:{slot}", self.set(f'slot:opc_tags:{slot}', to_insert[slot], ttl=None)
to_insert[slot], ttl=None) report[slot] = {'success': True, 'message': 'Slot updated successfully'}
report[slot] = {
"success": True,
"message": "Slot updated successfully"
}
success_count += 1 success_count += 1
except Exception as e: except Exception as e:
self.error( self.error(f'Failed to update slot {slot}: {str(e)}', metadata=metadata)
f"Failed to update slot {slot}: {str(e)}", metadata=metadata) report[slot] = {'success': False, 'message': str(e)}
report[slot] = {
"success": False,
"message": str(e)
}
self.info( self.info(f'Updated {success_count} of {len(to_insert)} OPC slots', metadata=metadata)
f"Updated {success_count} of {len(to_insert)} OPC slots", metadata=metadata)
self.debug( self.debug(f'Report: \n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
f"Report: \n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata)
return report return report
@activity.defn(name="delete_slots") @activity.defn(name='delete_slots')
async def delete_slots(self, input_data: dict[str, Any]) -> dict[str, Any]: async def delete_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Delete OPC slots from Redis Delete OPC slots from Redis
@@ -213,9 +203,9 @@ class SlotManager(Redis):
""" """
to_delete = input_data['to_delete'] to_delete = input_data['to_delete']
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
self.info("Deleting OPC slots...", metadata=metadata) self.info('Deleting OPC slots...', metadata=metadata)
report = {} report = {}
@@ -223,29 +213,20 @@ class SlotManager(Redis):
for slot in to_delete: for slot in to_delete:
try: try:
self.redis_client.delete(f"slot:opc_tags:{slot}") self.redis_client.delete(f'slot:opc_tags:{slot}')
report[slot] = { report[slot] = {'success': True, 'message': 'Slot deleted successfully'}
"success": True,
"message": "Slot deleted successfully"
}
success_count += 1 success_count += 1
except Exception as e: except Exception as e:
self.error( self.error(f'Failed to delete slot {slot}: {str(e)}', metadata=metadata)
f"Failed to delete slot {slot}: {str(e)}", metadata=metadata) report[slot] = {'success': False, 'message': str(e)}
report[slot] = {
"success": False,
"message": str(e)
}
self.info( self.info(f'Deleted {success_count} of {len(to_delete)} OPC slots', metadata=metadata)
f"Deleted {success_count} of {len(to_delete)} OPC slots", metadata=metadata)
self.debug( self.debug(f'Report: \n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
f"Report: \n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata)
return report return report
@activity.defn(name="get_last_data_timestamp") @activity.defn(name='get_last_data_timestamp')
async def get_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None: async def get_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None:
""" """
Gets the last data timestamp from redis. Gets the last data timestamp from redis.
@@ -259,32 +240,29 @@ class SlotManager(Redis):
str | None: The last data timestamp as a string, or None if no timestamp exists. str | None: The last data timestamp as a string, or None if no timestamp exists.
""" """
metadata = input_data['metadata'] metadata = input_data['metadata']
key = f"notification_last_timestamp:{input_data['mail_type']}" key = f'notification_last_timestamp:{input_data["mail_type"]}'
try: try:
data_hold = self.get(key) data_hold = self.get(key)
except Exception as e: except Exception as e:
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id="REDIS_GET_ERROR", notification_id='REDIS_GET_ERROR',
message=f"Error getting last data timestamp: {e}", message=f'Error getting last data timestamp: {e}',
block="get_last_data_timestamp", block='get_last_data_timestamp',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=traceback.format_exc() attachment_content=traceback.format_exc(),
) )
raise e raise e
self.debug( self.debug(f'Last collected timestamp: {data_hold}', metadata=metadata)
f"Last collected timestamp: {data_hold}",
metadata=metadata
)
if not data_hold: if not data_hold:
return None return None
return data_hold return data_hold
@activity.defn(name="put_last_data_timestamp") @activity.defn(name='put_last_data_timestamp')
async def put_last_data_timestamp(self, input_data: dict[str, Any]): async def put_last_data_timestamp(self, input_data: dict[str, Any]):
""" """
Puts the last data timestamp into redis. Puts the last data timestamp into redis.
@@ -299,39 +277,34 @@ class SlotManager(Redis):
str | None: The last data timestamp that was stored, or None if no data exists. str | None: The last data timestamp that was stored, or None if no data exists.
""" """
metadata = input_data['metadata'] metadata = input_data['metadata']
key = f"notification_last_timestamp:{input_data['mail_type']}" key = f'notification_last_timestamp:{input_data["mail_type"]}'
data = DataFrame(input_data['data']) data = DataFrame(input_data['data'])
if data.empty: if data.empty:
self.warning("No data to insert", self.warning('No data to insert', metadata=metadata)
metadata=metadata
)
return None return None
last_data_timestamp = data['timestamp'].max() last_data_timestamp = data['timestamp'].max()
self.debug( self.debug(f'Last collected timestamp to insert: {last_data_timestamp}', metadata=metadata)
f"Last collected timestamp to insert: {last_data_timestamp}",
metadata=metadata
)
try: try:
self.set(key, last_data_timestamp, ttl=60*60*5) self.set(key, last_data_timestamp, ttl=60 * 60 * 5)
except Exception as e: except Exception as e:
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id="REDIS_SET_ERROR", notification_id='REDIS_SET_ERROR',
message=f"Error setting last data timestamp: {e}", message=f'Error setting last data timestamp: {e}',
block="put_last_data_timestamp", block='put_last_data_timestamp',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=traceback.format_exc() attachment_content=traceback.format_exc(),
) )
raise e raise e
return last_data_timestamp return last_data_timestamp
@activity.defn(name="filter_notification_alerts") @activity.defn(name='filter_notification_alerts')
async def filter_notification_alerts(self, input_data: dict[str, Any]) -> dict[str, Any]: async def filter_notification_alerts(self, input_data: dict[str, Any]) -> dict[str, Any]:
""" """
Filter notification alerts with intelligent TTL-based duplicate prevention. Filter notification alerts with intelligent TTL-based duplicate prevention.
@@ -369,16 +342,13 @@ class SlotManager(Redis):
sending_configs = input_data['sending_configs'] sending_configs = input_data['sending_configs']
notification_ttl = input_data['notification_ttl'] notification_ttl = input_data['notification_ttl']
self.info("Filtering notification alerts...", metadata=metadata) self.info('Filtering notification alerts...', metadata=metadata)
receiver_groups = {} receiver_groups = {}
for receiver_group in sending_configs: for receiver_group in sending_configs:
group_name = receiver_group['group_name'] group_name = receiver_group['group_name']
receiver_groups[group_name] = { receiver_groups[group_name] = {**receiver_group, 'notifications': []}
**receiver_group,
"notifications": []
}
receiver_groups[group_name]['notifications'] = [] receiver_groups[group_name]['notifications'] = []
already_added_keys = [] already_added_keys = []
@@ -386,28 +356,30 @@ class SlotManager(Redis):
ignore_list = receiver_group.get('ignore', []) ignore_list = receiver_group.get('ignore', [])
for notification in notification_package: for notification in notification_package:
alert_type = "do_nothing" alert_type = 'do_nothing'
notification_id = notification['notification_id'] notification_id = notification['notification_id']
# Check if notification was recently sent # Check if notification was recently sent
key = f"{notification['trigger']}:{notification_id}" key = f'{notification["trigger"]}:{notification_id}'
last_sent = self.get(key) last_sent = self.get(key)
if last_sent is None: if last_sent is None:
alert_type = "core_alerts" alert_type = 'core_alerts'
else: else:
last_sent = datetime.strptime( last_sent = datetime.strptime(last_sent, DATETIME_FORMAT_MS_WITH_TZ)
last_sent, DATETIME_FORMAT_MS_WITH_TZ)
# Check if "notification_ttl" seconds has passed since last sent # Check if "notification_ttl" seconds has passed since last sent
if (now() - last_sent) > timedelta(seconds=notification_ttl): if (now() - last_sent) > timedelta(seconds=notification_ttl):
alert_type = "persistent_alerts" alert_type = 'persistent_alerts'
# Check if this group must be notified # Check if this group must be notified
if alert_type in receiver_group['contents'] and notification_id not in ignore_list and key not in already_added_keys: if (
receiver_groups[group_name]["notifications"].append( alert_type in receiver_group['contents']
notification) and notification_id not in ignore_list
and key not in already_added_keys
):
receiver_groups[group_name]['notifications'].append(notification)
already_added_keys.append(key) already_added_keys.append(key)
# Remove groups with no notifications # Remove groups with no notifications
@@ -419,7 +391,7 @@ class SlotManager(Redis):
return receiver_groups return receiver_groups
@activity.defn(name="store_notification_cache") @activity.defn(name='store_notification_cache')
async def store_notification_cache(self, input_data: dict[str, Any]) -> None: async def store_notification_cache(self, input_data: dict[str, Any]) -> None:
""" """
Store notification cache in Redis to track recently sent notifications. Store notification cache in Redis to track recently sent notifications.
@@ -434,14 +406,14 @@ class SlotManager(Redis):
log_report = DataFrame(input_data['log_report']) log_report = DataFrame(input_data['log_report'])
sent_ttl = input_data['sent_ttl'] sent_ttl = input_data['sent_ttl']
self.info("Storing notification cache...", metadata=metadata) self.info('Storing notification cache...', metadata=metadata)
date_now = now().strftime(DATETIME_FORMAT_MS_WITH_TZ) date_now = now().strftime(DATETIME_FORMAT_MS_WITH_TZ)
for index, row in log_report.iterrows(): for _index, row in log_report.iterrows():
status = row['status'] status = row['status']
if status == 'sent': if status == 'sent':
key = f"{row['schedule']}:{row['notification_id']}" key = f'{row["schedule"]}:{row["notification_id"]}'
self.set(key, date_now, ttl=sent_ttl) self.set(key, date_now, ttl=sent_ttl)
self.info("Notification cache stored...", metadata=metadata) self.info('Notification cache stored...', metadata=metadata)

316
orchestrator/activities/temporal_manager.py
View File

@@ -1,20 +1,26 @@
from temporalio import activity, workflow from temporalio import activity, workflow
from temporalio.client import ( from temporalio.client import (
Client, Schedule, ScheduleActionStartWorkflow, ScheduleIntervalSpec, ScheduleSpec, ScheduleUpdate, ScheduleUpdateInput) Client,
Schedule,
ScheduleActionStartWorkflow,
ScheduleIntervalSpec,
ScheduleSpec,
ScheduleUpdate,
ScheduleUpdateInput,
)
from temporalio.common import SearchAttributeKey, SearchAttributePair, TypedSearchAttributes from temporalio.common import SearchAttributeKey, SearchAttributePair, TypedSearchAttributes
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
import traceback
from typing import Any
from logging import Logger
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.temporal.activities.base import BaseActivity
from sientia_do.notifications.models import NotificationLevel
from google.protobuf.json_format import MessageToDict
import base64
from datetime import timedelta
import json import json
from asyncio import sleep import traceback
from datetime import timedelta
from logging import Logger
from typing import Any
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.base import BaseActivity
from orchestrator.utils.converters import parse_frequency from orchestrator.utils.converters import parse_frequency
@@ -35,61 +41,64 @@ class TemporalManager(BaseActivity):
notification_handler (NotificationHandler): Notification management handler notification_handler (NotificationHandler): Notification management handler
""" """
def __init__(self, host: str, scouter_namespace: str, laborious_namespace: str, def __init__(
logger: Logger, notification_handler: NotificationHandler): self,
host: str,
scouter_namespace: str,
laborious_namespace: str,
logger: Logger,
notification_handler: NotificationHandler,
):
self.temporal_host = host self.temporal_host = host
self.scouter_namespace = scouter_namespace self.scouter_namespace = scouter_namespace
self.laborious_namespace = laborious_namespace self.laborious_namespace = laborious_namespace
self.temporal_clients = {} self.temporal_clients: dict[str, Client] = {}
self.model_id_id_key = SearchAttributeKey.for_keyword("model_id") self.model_id_id_key = SearchAttributeKey.for_keyword('model_id')
self.model_name_id_key = SearchAttributeKey.for_keyword("model_name") self.model_name_id_key = SearchAttributeKey.for_keyword('model_name')
self.orchestrated_id_key = SearchAttributeKey.for_keyword( self.orchestrated_id_key = SearchAttributeKey.for_keyword('orchestrated')
"orchestrated")
BaseActivity.__init__(self, BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler)
logger=logger,
notification_handler=notification_handler)
async def connect_to_temporal(self): async def connect_to_temporal(self):
""" """
Connect to Temporal server namespaces for scouter and laborious workflows. Connect to Temporal server namespaces used by scouter and laborious workflows.
Creates client connections to both namespaces and stores them for later use. Creates and caches `Client` connections for both namespaces for later use.
""" """
self.logger.info( self.logger.info(f'Connecting to Temporal side namespaces at {self.temporal_host}')
f"Connecting to Temporal side namespaces at {self.temporal_host}") self.logger.info(f'Scouter namespace: {self.scouter_namespace}')
self.logger.info(f"Scouter namespace: {self.scouter_namespace}")
scouter_client = await Client.connect( scouter_client = await Client.connect(
target_host=self.temporal_host, target_host=self.temporal_host, namespace=self.scouter_namespace
namespace=self.scouter_namespace
) )
self.logger.info(f"Laborious namespace: {self.laborious_namespace}") self.logger.info(f'Laborious namespace: {self.laborious_namespace}')
laborious_client = await Client.connect( laborious_client = await Client.connect(
target_host=self.temporal_host, target_host=self.temporal_host, namespace=self.laborious_namespace
namespace=self.laborious_namespace
) )
self.temporal_clients = { self.temporal_clients = {
self.scouter_namespace: scouter_client, self.scouter_namespace: scouter_client,
self.laborious_namespace: laborious_client self.laborious_namespace: laborious_client,
} }
@activity.defn(name="normalize_schedules") @activity.defn(name='normalize_schedules')
async def normalize_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]: async def normalize_schedules(self, input_data: dict[str, Any]):
""" """
Normalize schedules. Removes schedules with no update time in mongo db collection "orchestrated_schedules". Normalize schedules by removing orphaned schedules from Temporal.
Any schedule marked with search attribute `orchestrated=true` that does not
exist in MongoDB collection `orchestrated_schedules` will be deleted.
input_data: input_data:
- orchestrated_schedules (dict[str, Any]): The orchestrated schedules to compare. - orchestrated_schedules (dict[str, Any]): Current orchestrated schedules from MongoDB.
""" """
metadata = input_data["metadata"] metadata = input_data['metadata']
remove_count = 0 remove_count = 0
self.info("Getting orchestrated schedules...", metadata=metadata) self.info('Getting orchestrated schedules...', metadata=metadata)
orchestrated_schedules = input_data.get('orchestrated_schedules', {}) orchestrated_schedules = input_data.get('orchestrated_schedules', {})
@@ -97,20 +106,20 @@ class TemporalManager(BaseActivity):
try: try:
schedules = orchestrated_schedules.get(namespace, {}) schedules = orchestrated_schedules.get(namespace, {})
self.info( self.info(f'Getting orchestrated schedules for {namespace}', metadata=metadata)
f"Getting orchestrated schedules for {namespace}", metadata=metadata)
async for schedule in await client.list_schedules(): async for schedule in await client.list_schedules():
search_attrs = getattr(schedule, "search_attributes", {}) search_attrs = getattr(schedule, 'search_attributes', {})
if search_attrs.get("orchestrated", ["false"]) == ["true"]: if search_attrs.get('orchestrated', ['false']) == ['true']:
schedule_id = schedule.id schedule_id = schedule.id
if schedule_id not in schedules: if schedule_id not in schedules:
self.info( self.info(
f"Schedule {schedule_id} not found in mongo db, cleaning up", metadata=metadata) f'Schedule {schedule_id} not found in mongo db, cleaning up',
metadata=metadata,
)
handle = client.get_schedule_handle( handle = client.get_schedule_handle(schedule_id)
schedule_id)
await handle.delete() await handle.delete()
@@ -120,22 +129,21 @@ class TemporalManager(BaseActivity):
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( self.send_notification(
metadata=metadata, metadata=metadata,
notification_id="TEMPORAL_NORMALIZE_SCHEDULES_ERROR", notification_id='TEMPORAL_NORMALIZE_SCHEDULES_ERROR',
message=f"Failed to normalize schedules: {e}", message=f'Failed to normalize schedules: {e}',
block="normalize_schedules", block='normalize_schedules',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=trace attachment_content=trace,
) )
self.error(trace, metadata=metadata) self.error(trace, metadata=metadata)
raise e raise e
self.info( self.info(f'Removed {remove_count} schedules', metadata=metadata)
f"Removed {remove_count} schedules", metadata=metadata)
@activity.defn(name="create_schedules") @activity.defn(name='create_schedules')
async def create_schedules(self, input_data: dict[str, Any]) -> list[dict[str, Any]]: async def create_schedules(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
""" """
Create schedules in Temporal Create schedules in Temporal.
Args: Args:
- input_data (dict[str, Any]): The input data containing - input_data (dict[str, Any]): The input data containing
@@ -143,47 +151,45 @@ class TemporalManager(BaseActivity):
- schedules (dict[str, Any]): The schedules to create. - schedules (dict[str, Any]): The schedules to create.
Returns: Returns:
- dict[str, Any]: A report of the created schedules. - list[dict[str, Any]]: Report entries for each attempted schedule creation.
""" """
schedules_to_create = input_data['schedules'] schedules_to_create = input_data['schedules']
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
report = [] report = []
success_count = 0 success_count = 0
self.info("Creating schedules...", metadata=metadata) self.info('Creating schedules...', metadata=metadata)
for namespace, schedules in schedules_to_create.items(): for namespace, schedules in schedules_to_create.items():
client = self.temporal_clients.get(namespace) client = self.temporal_clients.get(namespace)
if not client: if not client:
raise ValueError( raise ValueError(
f"Temporal client for {namespace} not found, clients: {self.temporal_clients}") f'Temporal client for {namespace} not found, clients: {self.temporal_clients}'
)
for schedule_name, schedule in schedules.items(): for schedule_name, schedule in schedules.items():
search_attributes = TypedSearchAttributes([ search_attributes = TypedSearchAttributes(
[
SearchAttributePair(key=self.model_id_id_key, value=schedule['model_id']),
SearchAttributePair( SearchAttributePair(
key=self.model_id_id_key, key=self.model_name_id_key, value=schedule['model_name']
value=schedule['model_id']
), ),
SearchAttributePair( SearchAttributePair(key=self.orchestrated_id_key, value='true'),
key=self.model_name_id_key, ]
value=schedule['model_name']
),
SearchAttributePair(
key=self.orchestrated_id_key,
value="true"
) )
])
workflow_type = schedule['workflow_type'] workflow_type = schedule['workflow_type']
try: try:
execution_timeout_seconds = schedule.get('execution_timeout_seconds', 300)
task_timeout_seconds = schedule.get('task_timeout_seconds', 300)
self.debug(f'Creating schedule {schedule_name}:', metadata=metadata)
self.debug( self.debug(
f"Creating schedule {schedule_name}:", metadata=metadata) f'{json.dumps(schedule, indent=4, sort_keys=True)}', metadata=metadata
self.debug( )
f"{json.dumps(schedule, indent=4, sort_keys=True)}", metadata=metadata)
await client.create_schedule( await client.create_schedule(
schedule_name, schedule_name,
@@ -192,52 +198,60 @@ class TemporalManager(BaseActivity):
workflow_type, workflow_type,
schedule, schedule,
id=schedule_name, id=schedule_name,
task_queue=f"{workflow_type}-queue", task_queue=f'{workflow_type}-queue',
execution_timeout=timedelta(minutes=2), execution_timeout=timedelta(seconds=execution_timeout_seconds),
typed_search_attributes=search_attributes run_timeout=timedelta(seconds=execution_timeout_seconds),
task_timeout=timedelta(seconds=task_timeout_seconds),
typed_search_attributes=search_attributes,
), ),
spec=ScheduleSpec( spec=ScheduleSpec(
intervals=[ intervals=[
ScheduleIntervalSpec( ScheduleIntervalSpec(
every=timedelta(seconds=parse_frequency( every=timedelta(
schedule.get('frequency', '1m'))) seconds=parse_frequency(schedule.get('frequency', '1m'))
)
) )
] ]
)
), ),
search_attributes=search_attributes ),
search_attributes=search_attributes,
) )
report.append({ report.append(
"namespace": namespace, {
"schedule_name": schedule_name, 'namespace': namespace,
"success": True, 'schedule_name': schedule_name,
"message": "Schedule created successfully" 'success': True,
}) 'message': 'Schedule created successfully',
}
)
success_count += 1 success_count += 1
except Exception as e: except Exception as e:
self.error( self.error(
f"Failed to create schedule {schedule_name}: {str(e)}", metadata=metadata) f'Failed to create schedule {schedule_name}: {str(e)}', metadata=metadata
report.append({ )
"namespace": namespace, report.append(
"schedule_name": schedule_name, {
"success": False, 'namespace': namespace,
"message": str(e) 'schedule_name': schedule_name,
}) 'success': False,
'message': str(e),
}
)
self.info( self.info(
f"Created {success_count} of {len(schedules_to_create)} schedules", metadata=metadata) f'Created {success_count} of {len(schedules_to_create)} schedules', metadata=metadata
)
self.debug( self.debug(f'\n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
f"\n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata)
return report return report
@activity.defn(name="update_schedules") @activity.defn(name='update_schedules')
async def update_schedules(self, input_data: dict[str, Any]) -> list[dict[str, Any]]: async def update_schedules(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
""" """
Update schedules in Temporal Update schedules in Temporal.
Args: Args:
- input_data (dict[str, Any]): The input data containing - input_data (dict[str, Any]): The input data containing
@@ -245,31 +259,31 @@ class TemporalManager(BaseActivity):
- schedules (dict[str, Any]): The schedules to update. - schedules (dict[str, Any]): The schedules to update.
Returns: Returns:
- dict[str, Any]: A report of the updated schedules. - list[dict[str, Any]]: Report entries for each attempted schedule update.
""" """
schedules_to_update = input_data['schedules'] schedules_to_update = input_data['schedules']
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
report = [] report = []
success_count = 0 success_count = 0
self.info("Updating schedules...", metadata=metadata) self.info('Updating schedules...', metadata=metadata)
for namespace, schedules in schedules_to_update.items(): for namespace, schedules in schedules_to_update.items():
client = self.temporal_clients.get(namespace) client = self.temporal_clients.get(namespace)
if not client: if not client:
raise ValueError( raise ValueError(
f"Temporal client for {namespace} not found, clients: {self.temporal_clients}") f'Temporal client for {namespace} not found, clients: {self.temporal_clients}'
)
for schedule_name, schedule in schedules.items(): for schedule_name, schedule in schedules.items():
try: try:
handler = client.get_schedule_handle( handler = client.get_schedule_handle(schedule_name)
schedule_name)
if not handler: if not handler:
raise ValueError(f"Schedule {schedule_name} not found") raise ValueError(f'Schedule {schedule_name} not found')
# fmt: off # fmt: off
async def update_schedule(input_data: ScheduleUpdateInput) -> ScheduleUpdate: # NOSONAR async def update_schedule(input_data: ScheduleUpdateInput) -> ScheduleUpdate: # NOSONAR
@@ -298,36 +312,41 @@ class TemporalManager(BaseActivity):
del update_schedule del update_schedule
report.append({ report.append(
"namespace": namespace, {
"schedule_name": schedule_name, 'namespace': namespace,
"success": True, 'schedule_name': schedule_name,
"message": "Schedule updated successfully" 'success': True,
}) 'message': 'Schedule updated successfully',
}
)
success_count += 1 success_count += 1
except Exception as e: except Exception as e:
self.error( self.error(
f"Failed to update schedule {schedule_name}: {str(e)}", metadata=metadata) f'Failed to update schedule {schedule_name}: {str(e)}', metadata=metadata
report.append({ )
"namespace": namespace, report.append(
"schedule_name": schedule_name, {
"success": False, 'namespace': namespace,
"message": str(e) 'schedule_name': schedule_name,
}) 'success': False,
'message': str(e),
}
)
self.info( self.info(
f"Updated {success_count} of {len(schedules_to_update)} schedules", metadata=metadata) f'Updated {success_count} of {len(schedules_to_update)} schedules', metadata=metadata
)
self.debug( self.debug(f'\n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
f"\n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata)
return report return report
@activity.defn(name="delete_schedules") @activity.defn(name='delete_schedules')
async def delete_schedules(self, input_data: dict[str, Any]) -> list[dict[str, Any]]: async def delete_schedules(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
""" """
Delete schedules in Temporal Delete schedules in Temporal.
Args: Args:
- input_data (dict[str, Any]): The input data containing - input_data (dict[str, Any]): The input data containing
@@ -335,58 +354,63 @@ class TemporalManager(BaseActivity):
- schedules (list[str]): The schedules to delete. - schedules (list[str]): The schedules to delete.
Returns: Returns:
- dict[str, Any]: A report of the deleted schedules. - list[dict[str, Any]]: Report entries for each attempted schedule deletion.
""" """
schedules_to_delete = input_data['schedules'] schedules_to_delete = input_data['schedules']
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
report = [] report = []
success_count = 0 success_count = 0
self.info("Deleting schedules...", metadata=metadata) self.info('Deleting schedules...', metadata=metadata)
for namespace, schedules in schedules_to_delete.items(): for namespace, schedules in schedules_to_delete.items():
client = self.temporal_clients.get(namespace) client = self.temporal_clients.get(namespace)
if not client: if not client:
raise ValueError( raise ValueError(
f"Temporal client for {namespace} not found, clients: {self.temporal_clients}") f'Temporal client for {namespace} not found, clients: {self.temporal_clients}'
)
for schedule_name in schedules: for schedule_name in schedules:
try: try:
handler = client.get_schedule_handle( handler = client.get_schedule_handle(schedule_name)
schedule_name)
if not handler: if not handler:
raise ValueError(f"Schedule {schedule_name} not found") raise ValueError(f'Schedule {schedule_name} not found')
await handler.delete() await handler.delete()
report.append({ report.append(
"namespace": namespace, {
"schedule_name": schedule_name, 'namespace': namespace,
"success": True, 'schedule_name': schedule_name,
"message": "Schedule deleted successfully" 'success': True,
}) 'message': 'Schedule deleted successfully',
}
)
success_count += 1 success_count += 1
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.error( self.error(
f"Failed to delete schedule {schedule_name}: {str(e)}", metadata=metadata) f'Failed to delete schedule {schedule_name}: {str(e)}', metadata=metadata
report.append({ )
"namespace": namespace, report.append(
"schedule_name": schedule_name, {
"success": False, 'namespace': namespace,
"message": str(e), 'schedule_name': schedule_name,
"attachment": trace 'success': False,
}) 'message': str(e),
'attachment': trace,
}
)
self.info( self.info(
f"Deleted {success_count} of {len(schedules_to_delete)} schedules", metadata=metadata) f'Deleted {success_count} of {len(schedules_to_delete)} schedules', metadata=metadata
)
self.debug( self.debug(f'\n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
f"\n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata)
return report return report

23
orchestrator/metrics.py
View File

@@ -1,24 +1,23 @@
""" """
Prometheus metrics definitions for the orchestrator application. Prometheus metric definitions for the orchestrator application.
This module defines all Prometheus metrics used for monitoring the This module exposes Prometheus counters and gauges for monitoring the
orchestrator system including application health, email delivery, orchestrator, including application health and email delivery metrics.
and workflow execution metrics.
""" """
from prometheus_client import Gauge, Counter from prometheus_client import Counter, Gauge
APP_UP = Gauge( APP_UP = Gauge(
"app_up", 'app_up',
"Indicates if the application is running (1) or shutting down (0)", 'Indicates if the application is running (1) or shutting down (0)',
["pod_id"], ['pod_id'],
) )
CORE_LABELS = ["pod_id", "model_name", "pipeline_name"] CORE_LABELS = ['pod_id', 'model_name', 'pipeline_name']
EMAIL_SENT_COUNT = Counter( EMAIL_SENT_COUNT = Counter(
"email_sent_count", 'email_sent_count',
"Number of emails sent", 'Total number of emails sent by the orchestrator',
[*CORE_LABELS, "email_group"], [*CORE_LABELS, 'email_group'],
) )

12
orchestrator/utils/connectors_config.py
View File

@@ -12,7 +12,7 @@ def build_redis_config():
'host': getenv('REDIS_HOST', 'localhost'), 'host': getenv('REDIS_HOST', 'localhost'),
'port': int(getenv('REDIS_PORT', '6379')), 'port': int(getenv('REDIS_PORT', '6379')),
'username': getenv('REDIS_USERNAME', 'default'), 'username': getenv('REDIS_USERNAME', 'default'),
'password': getenv('REDIS_PASSWORD', 'bdnZOpcyiL') 'password': getenv('REDIS_PASSWORD', 'bdnZOpcyiL'),
} }
@@ -31,7 +31,7 @@ def build_mongodb_config():
return { return {
'connection_string': connection_string, 'connection_string': connection_string,
'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'), 'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'),
'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600 'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600,
} }
@@ -45,7 +45,7 @@ def build_couchbase_config():
return { return {
'connection_string': getenv('COUCHBASE_CONNECTION_STRING', 'couchbase://localhost'), 'connection_string': getenv('COUCHBASE_CONNECTION_STRING', 'couchbase://localhost'),
'username': getenv('COUCHBASE_USERNAME', 'sientia'), 'username': getenv('COUCHBASE_USERNAME', 'sientia'),
'password': getenv('COUCHBASE_PASSWORD', 'sientia') 'password': getenv('COUCHBASE_PASSWORD', 'sientia'),
} }
@@ -60,7 +60,7 @@ def build_temporal_config():
'temporal_host': getenv('TEMPORAL_HOST', 'localhost:7233'), 'temporal_host': getenv('TEMPORAL_HOST', 'localhost:7233'),
'temporal_namespace': getenv('TEMPORAL_NAMESPACE', 'default'), 'temporal_namespace': getenv('TEMPORAL_NAMESPACE', 'default'),
'temporal_scouter_namespace': getenv('TEMPORAL_SCOUTER_NAMESPACE', 'scouter'), 'temporal_scouter_namespace': getenv('TEMPORAL_SCOUTER_NAMESPACE', 'scouter'),
'temporal_laborious_namespace': getenv('TEMPORAL_LABORIOUS_NAMESPACE', 'laborious') 'temporal_laborious_namespace': getenv('TEMPORAL_LABORIOUS_NAMESPACE', 'laborious'),
} }
@@ -78,7 +78,7 @@ def build_postgres_config():
'password': getenv('POSTGRES_PASSWORD', 'sientia'), 'password': getenv('POSTGRES_PASSWORD', 'sientia'),
'dbname': getenv('POSTGRES_DBNAME', 'sientia'), 'dbname': getenv('POSTGRES_DBNAME', 'sientia'),
'min_connections': int(getenv('POSTGRES_MIN_CONNECTIONS', '5')), 'min_connections': int(getenv('POSTGRES_MIN_CONNECTIONS', '5')),
'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20')) 'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20')),
} }
@@ -93,5 +93,5 @@ def build_email_config():
'sender_email': getenv('EMAIL_SENDER', 'sientia-alerts@aignosi.com'), 'sender_email': getenv('EMAIL_SENDER', 'sientia-alerts@aignosi.com'),
'sender_password': getenv('EMAIL_SENDER_PASSWORD', 'sientia'), 'sender_password': getenv('EMAIL_SENDER_PASSWORD', 'sientia'),
'smtp_server': getenv('EMAIL_SMTP_SERVER', None), 'smtp_server': getenv('EMAIL_SMTP_SERVER', None),
'smtp_port': int(getenv('EMAIL_SMTP_PORT', '587')) 'smtp_port': int(getenv('EMAIL_SMTP_PORT', '587')),
} }

10
orchestrator/utils/converters.py
View File

@@ -19,13 +19,13 @@ def parse_frequency(frequency: str) -> int:
Raises: Raises:
ValueError: If frequency format is invalid ValueError: If frequency format is invalid
""" """
if frequency.endswith("s"): if frequency.endswith('s'):
return int(frequency[:-1]) return int(frequency[:-1])
elif frequency.endswith("m"): elif frequency.endswith('m'):
return int(frequency[:-1]) * 60 return int(frequency[:-1]) * 60
elif frequency.endswith("h"): elif frequency.endswith('h'):
return int(frequency[:-1]) * 60 * 60 return int(frequency[:-1]) * 60 * 60
elif frequency.endswith("d"): elif frequency.endswith('d'):
return int(frequency[:-1]) * 60 * 60 * 24 return int(frequency[:-1]) * 60 * 60 * 24
else: else:
raise ValueError("Invalid frequency") raise ValueError('Invalid frequency')

72
orchestrator/utils/email_builder.py
View File

@@ -1,8 +1,7 @@
import json from typing import Any
from sientia_do.observability.logger import Logger
from sientia_do.notifications.models import NotificationLevel
from jinja2 import Template from jinja2 import Template
import re from sientia_do.observability.logger import Logger
class EmailBuilder: class EmailBuilder:
@@ -24,9 +23,9 @@ class EmailBuilder:
self.report_template_file = './orchestrator/utils/templates/email_template.html' self.report_template_file = './orchestrator/utils/templates/email_template.html'
self.general_template_file = './orchestrator/utils/templates/general_template.html' self.general_template_file = './orchestrator/utils/templates/general_template.html'
with open(self.report_template_file, 'r') as file: with open(self.report_template_file) as file:
self.report_template = file.read() self.report_template = file.read()
with open(self.general_template_file, 'r') as file: with open(self.general_template_file) as file:
self.general_template = file.read() self.general_template = file.read()
def replace_parameters(self, template: str, parameters: dict) -> str: def replace_parameters(self, template: str, parameters: dict) -> str:
@@ -40,10 +39,10 @@ class EmailBuilder:
Returns: Returns:
str: The rendered template with parameters replaced. str: The rendered template with parameters replaced.
""" """
# Criar um template Jinja2 # Create a Jinja2 template from the provided string
template = Template(template) template_obj = Template(template)
return template.render(parameters) return template_obj.render(parameters)
def parameters(self, general_events: dict, mail_type: str) -> dict: def parameters(self, general_events: dict, mail_type: str) -> dict:
""" """
@@ -57,29 +56,33 @@ class EmailBuilder:
Returns: Returns:
dict: Dictionary with mail_type and rendered event sections for each notification level. dict: Dictionary with mail_type and rendered event sections for each notification level.
""" """
error_models = general_events.get('ERROR', {}).get('models', []) error_events = general_events.get('ERROR', {})
warning_models = general_events.get('WARNING', {}).get('models', []) warning_events = general_events.get('WARNING', {})
info_models = general_events.get('INFO', {}).get('models', []) info_events = general_events.get('INFO', {})
error_models = error_events.get('models', [])
warning_models = warning_events.get('models', [])
info_models = info_events.get('models', [])
return { return {
'mail_type': mail_type, 'mail_type': mail_type,
'error_events': self.replace_parameters(self.general_template, 'error_events': self.replace_parameters(self.general_template, error_events)
general_events.get( if error_models
'ERROR')) if error_models else '', else '',
'warning_events': self.replace_parameters(self.general_template, 'warning_events': self.replace_parameters(self.general_template, warning_events)
general_events.get( if warning_models
'WARNING')) if warning_models else '', else '',
'info_events': self.replace_parameters(self.general_template, 'info_events': self.replace_parameters(self.general_template, info_events)
general_events.get( if info_models
'INFO')) if info_models else '', else '',
} }
def build_email(self, report_data: list[dict], mail_type: str) -> str: def build_email(self, report_data: list[dict[str, Any]], mail_type: str) -> str:
""" """
Builds the email HTML by organizing report data by notification level and model. Builds the email HTML by organizing report data by notification level and model.
Args: Args:
report_data (list[dict]): List of notification reports, each containing: report_data (List[Dict[str, Any]]): List of notification reports, each containing:
- level (str): Notification level (ERROR, WARNING, INFO) - level (str): Notification level (ERROR, WARNING, INFO)
- model_name (str): Name of the model - model_name (str): Name of the model
- Additional notification details - Additional notification details
@@ -87,32 +90,33 @@ class EmailBuilder:
Returns: Returns:
str: Complete HTML email content ready for sending. str: Complete HTML email content ready for sending.
""" """
general_events = {} general_events: dict[str, dict[str, Any]] = {}
for report in report_data: for report in report_data:
level = report['level'] level = report['level']
model_name = report['model_name'] model_name = report['model_name']
if level not in general_events: if level not in general_events:
general_events[level] = { general_events[level] = {
'section_name': f'{level.capitalize()}s detected:', 'section_name': f'{level.capitalize()}s detected:',
'models': {} 'models': {},
} }
if model_name not in general_events[level]['models']: # Type assertion to help the type checker understand the structure
general_events[level]['models'][model_name] = { level_data = general_events[level]
models_dict = level_data['models']
if model_name not in models_dict:
models_dict[model_name] = {
'model_name': model_name, 'model_name': model_name,
'events': [] 'events': [],
} }
general_events[level]['models'][model_name]['events'].append( models_dict[model_name]['events'].append(report)
report)
for _type, content in general_events.items(): for _type, content in general_events.items():
content['models'] = list(content['models'].values()) content['models'] = list(content['models'].values())
return self.replace_parameters( return self.replace_parameters(
self.report_template, self.parameters( self.report_template, self.parameters(general_events, mail_type)
general_events, mail_type )
))

163
orchestrator/utils/orchestrator_functions.py
View File

@@ -19,14 +19,15 @@ def common_config(config: dict[str, Any]):
""" """
model = config['model'] model = config['model']
return { return {
"workflow_type": config['workflow_type'], 'workflow_type': config['workflow_type'],
"schedule_name": config['schedule_name'], 'schedule_name': config['schedule_name'],
"frequency": config.get('frequency', '1m'), 'frequency': config.get('frequency', '1m'),
"max_retry_policy": config.get('max_retry_policy', 1), 'max_retry_policy': config.get('max_retry_policy', 1),
'model_id': config['model_id'],
"model_id": config['model_id'], 'model_name': model['name'],
"model_name": model['name'], 'model_config': model.get('model_config', {}),
"model_config": model.get('model_config', {}), 'execution_timeout_seconds': config.get('execution_timeout_seconds', 300),
'task_timeout_seconds': config.get('task_timeout_seconds', 300),
} }
@@ -45,12 +46,12 @@ def minimal_retrain(config: dict[str, Any]):
""" """
return { return {
**common_config(config), **common_config(config),
"workflow_type": "minimal_retrain", 'workflow_type': 'minimal_retrain',
"schedule_name": config['schedule_name'], 'schedule_name': config['schedule_name'],
"query": config['query'], 'query': config['query'],
"schema": "sientia_data", 'schema': 'sientia_data',
"table_name": "log_retrain", 'table_name': 'log_retrain',
"datetime_columns": config.get('datetime_columns', []), 'datetime_columns': config.get('datetime_columns', []),
} }
@@ -76,28 +77,25 @@ def scouter(config: dict[str, Any]):
""" """
filters = {} filters = {}
for f in config.get('filters', []): for f in config.get('filters', []):
filters[f['filter_name']] = { filters[f['filter_name']] = {'policy': f['policy']}
"policy": f['policy']
}
tags = {} tags = {}
for tag in config['read_tags']: for tag in config['read_tags']:
tags[tag['tag_name']] = { tags[tag['tag_name']] = {
"aggr_func": tag.get('aggr_func', 'lts'), 'aggr_func': tag.get('aggr_func', 'lts'),
"data_range": tag.get('data_range', [-100, 100]) 'data_range': tag.get('data_range', [-100, 100]),
} }
return { return {
**common_config(config), **common_config(config),
'topic': f'raw_{config["schedule_name"]}',
"topic": f"raw_{config['schedule_name']}", 'trigger_laborious': False,
"trigger_laborious": False, 'filters': filters,
"filters": filters, 'schema': 'sientia_data',
"schema": "sientia_data", 'table_name': 'laborious_data',
"table_name": "laborious_data", 'retention_time': config.get('tag_retention_minutes', 60) * 60,
"retention_time": config.get('tag_retention_minutes', 60) * 60, 'model_tags': tags,
"model_tags": tags, 'debug_data_package': config.get('debug_data_package', False),
"debug_data_package": config.get('debug_data_package', False)
} }
@@ -117,8 +115,8 @@ def overlap_filter_config(base_filter_config: dict[str, Any], config: list[dict[
""" """
for fil in config: for fil in config:
base_filter_config[fil['filter_name']] = { base_filter_config[fil['filter_name']] = {
"policy": fil['policy'], 'policy': fil['policy'],
"config": fil.get('config', {}) 'config': fil.get('config', {}),
} }
return base_filter_config return base_filter_config
@@ -135,10 +133,10 @@ def process_path_priority(path_priority: list[str]):
list[str]: Normalized path priority list with exactly 3 elements: ["STOP", "CONTINUE", "REPEAT"]. list[str]: Normalized path priority list with exactly 3 elements: ["STOP", "CONTINUE", "REPEAT"].
""" """
for priority in path_priority[:]: for priority in path_priority[:]:
if priority not in ["STOP", "CONTINUE", "REPEAT"]: if priority not in ['STOP', 'CONTINUE', 'REPEAT']:
path_priority.remove(priority) path_priority.remove(priority)
for priority in ["STOP", "CONTINUE", "REPEAT"]: for priority in ['STOP', 'CONTINUE', 'REPEAT']:
if priority not in path_priority: if priority not in path_priority:
path_priority.append(priority) path_priority.append(priority)
@@ -166,7 +164,7 @@ def predictions_batch(config: dict[str, Any]):
Returns: Returns:
dict[str, Any]: Predictions batch configuration with OPC output config, filters, and path priority. dict[str, Any]: Predictions batch configuration with OPC output config, filters, and path priority.
""" """
tags = {} tags: dict[str, Any] = {}
for tag in config.get('write_tags', []): for tag in config.get('write_tags', []):
if tag['server_id'] not in tags: if tag['server_id'] not in tags:
tags[tag['server_id']] = {} tags[tag['server_id']] = {}
@@ -174,51 +172,43 @@ def predictions_batch(config: dict[str, Any]):
tag_type = tag['type'] tag_type = tag['type']
if tag_type == 'prediction' or tag_type == 'confidence': if tag_type == 'prediction' or tag_type == 'confidence':
tag_type_str = f"{tag_type}_tags" tag_type_str = f'{tag_type}_tags'
if tag_type_str not in tags[tag['server_id']]: if tag_type_str not in tags[tag['server_id']]:
tags[tag['server_id']][tag_type_str] = {} tags[tag['server_id']][tag_type_str] = {}
tags[tag['server_id']][tag_type_str][tag['addr']] = { tags[tag['server_id']][tag_type_str][tag['addr']] = {
"data_type": tag.get('data_type', 'float'), 'data_type': tag.get('data_type', 'float'),
} }
path_priority = process_path_priority(config.get( path_priority = process_path_priority(
'path_priority', ["STOP", "CONTINUE", "REPEAT"])) config.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT'])
)
return { return {
**common_config(config), **common_config(config),
'query': config['query'],
"query": config['query'], 'datetime_columns': config.get('datetime_columns', []),
"datetime_columns": config.get('datetime_columns', []), 'schema': 'sientia_data',
"schema": "sientia_data", 'table_name': 'predictions',
"table_name": "predictions", 'retention_time': config.get('model_retention_minutes', 60) * 60,
"retention_time": config.get('model_retention_minutes', 60) * 60, 'opc_output_config': tags,
"opc_output_config": tags, 'input_filters': overlap_filter_config(
"input_filters": overlap_filter_config({ {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, config.get('input_filters', [])
"EMPTY_DATA": { ),
"policy": "STOP", 'mlflow_transform_filters': overlap_filter_config(
"config": {} {
} 'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
}, config.get('input_filters', [])), 'API_ERROR': {'policy': 'STOP', 'config': {}},
"mlflow_transform_filters": overlap_filter_config({
"EMPTY_DATA": {
"policy": "STOP",
"config": {}
}, },
"API_ERROR": { config.get('mlflow_transform_filters', []),
"policy": "STOP", ),
"config": {} 'mlflow_predict_filters': overlap_filter_config(
} {'API_ERROR': {'policy': 'STOP', 'config': {}}},
}, config.get('mlflow_transform_filters', [])), config.get('mlflow_predict_filters', []),
"mlflow_predict_filters": overlap_filter_config({ ),
"API_ERROR": { 'path_priority': path_priority,
"policy": "STOP", 'predictions_storage_policy': config.get('predictions_storage_policy', 'lts:1'),
"config": {}
}
}, config.get('mlflow_predict_filters', [])),
"path_priority": path_priority,
"predictions_storage_policy": config.get('predictions_storage_policy', 'lts:1')
} }
@@ -238,21 +228,18 @@ def gather_read_tags(pipelines: list[dict[str, Any]]) -> dict[str, Any]:
# Get all read tags from pipelines # Get all read tags from pipelines
for pipeline in pipelines: for pipeline in pipelines:
for tag in pipeline.get('read_tags', []): for tag in pipeline.get('read_tags', []):
tag_string = f"{tag['server_id']}:{tag['tag_address']}" tag_string = f'{tag["server_id"]}:{tag["tag_address"]}'
if tag_string not in tags: if tag_string not in tags:
tags[tag_string] = { tags[tag_string] = {**tag, 'topics': []}
**tag,
"topics": []
}
tags[tag_string]['topics'].append( tags[tag_string]['topics'].append(f'raw_{pipeline["schedule_name"]}')
f"raw_{pipeline['schedule_name']}")
return tags return tags
def build_tag_config(tag: dict[str, Any], slot_config: dict[str, Any], def build_tag_config(
opc_servers: dict[str, Any], i: int): tag: dict[str, Any], slot_config: dict[str, Any], opc_servers: dict[str, Any], i: int
):
""" """
Build tag configuration for a specific slot and OPC server. Build tag configuration for a specific slot and OPC server.
@@ -273,22 +260,22 @@ def build_tag_config(tag: dict[str, Any], slot_config: dict[str, Any],
server_id = tag['server_id'] server_id = tag['server_id']
if server_id not in opc_servers: if server_id not in opc_servers:
raise ValueError(f"Server {server_id} not found in opc_servers") raise ValueError(f'Server {server_id} not found in opc_servers')
server_name = opc_servers[server_id]['server_name'] server_name = opc_servers[server_id]['server_name']
if server_name not in slot_config[f"{i}"]: if server_name not in slot_config[f'{i}']:
slot_config[f"{i}"][server_name] = { slot_config[f'{i}'][server_name] = {
"server_id": server_id, 'server_id': server_id,
"name": server_name, 'name': server_name,
"url": opc_servers[server_id]['url'], 'url': opc_servers[server_id]['url'],
"server_uri": opc_servers[server_id]['uri'], 'server_uri': opc_servers[server_id]['uri'],
"cert_path": opc_servers[server_id].get('cert_path', None), 'cert_path': opc_servers[server_id].get('cert_path', None),
"private_key_path": opc_servers[server_id].get('private_key_path', None), 'private_key_path': opc_servers[server_id].get('private_key_path', None),
"server_cert_path": opc_servers[server_id].get('server_cert_path', None), 'server_cert_path': opc_servers[server_id].get('server_cert_path', None),
"tags": {} 'tags': {},
} }
slot_config[f"{i}"][server_name]["tags"][tag['tag_address']] = { slot_config[f'{i}'][server_name]['tags'][tag['tag_address']] = {
**tag, **tag,
} }

89
orchestrator/worker/worker.py
View File

@@ -1,32 +1,36 @@
from temporalio import workflow, client from temporalio import client, workflow
from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
from temporalio.worker import Worker from temporalio.worker import Worker
from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
import asyncio
import os import os
import sys import sys
import asyncio
from orchestrator.workflows.alerts import Alerts from prometheus_client import start_http_server
from orchestrator.workflows.reports import Reports
from orchestrator.workflows.subworkflows.load_notification_package import LoadNotificationPackage
from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications
from orchestrator.workflows.orchestrator import Orchestrator
from orchestrator.activities.activities import Activities
from orchestrator.utils.connectors_config import (
# build_couchbase_config,
build_redis_config,
build_mongodb_config,
build_temporal_config,
build_email_config,
build_postgres_config
)
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import get_logger from sientia_do.observability.logger import get_logger
from prometheus_client import start_http_server
from orchestrator import metrics
POD_ID = os.getenv("POD_ID") from orchestrator import metrics
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091")) from orchestrator.activities.activities import Activities
from orchestrator.utils.connectors_config import (
build_email_config,
build_mongodb_config,
build_postgres_config,
# build_couchbase_config,
build_redis_config,
build_temporal_config,
)
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
POD_ID = os.getenv('POD_ID')
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
async def main(): async def main():
@@ -49,10 +53,9 @@ async def main():
'schedule_name': '-', 'schedule_name': '-',
} }
logger.custom_info( logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata=metadata)
f'Starting Worker with POD_ID: {POD_ID}', metadata=metadata)
logger.custom_info("Starting prometheus client...", metadata=metadata) logger.custom_info('Starting prometheus client...', metadata=metadata)
start_prometheus_server() start_prometheus_server()
logger.custom_info('Starting Notification Handler...', metadata=metadata) logger.custom_info('Starting Notification Handler...', metadata=metadata)
@@ -66,22 +69,21 @@ async def main():
) )
logger.custom_info( logger.custom_info(
f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata=metadata) f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata=metadata
)
new_runtime = Runtime( new_runtime = Runtime(
telemetry=TelemetryConfig( telemetry=TelemetryConfig(
metrics=PrometheusConfig( metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
bind_address=f"0.0.0.0:{SDK_METRICS_PORT}")
) )
) )
logger.custom_info( logger.custom_info(f'Starting Temporal Client at {host}:{namespace}', metadata=metadata)
f'Starting Temporal Client at {host}:{namespace}', metadata=metadata)
temporal_client = await client.Client.connect( temporal_client = await client.Client.connect(
target_host=host, target_host=host,
namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'), namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'),
runtime=new_runtime runtime=new_runtime,
) )
logger.custom_info('Starting Activities...', metadata=metadata) logger.custom_info('Starting Activities...', metadata=metadata)
@@ -93,7 +95,7 @@ async def main():
email_config=build_email_config(), email_config=build_email_config(),
postgres_config=build_postgres_config(), postgres_config=build_postgres_config(),
logger=logger, logger=logger,
notification_handler=notification_handler notification_handler=notification_handler,
) )
await activities.connect_to_temporal() await activities.connect_to_temporal()
@@ -133,7 +135,7 @@ async def main():
activities.report_schedule_orchestration, activities.report_schedule_orchestration,
activities.report_slot_orchestration, activities.report_slot_orchestration,
activities.format_schedule_config, activities.format_schedule_config,
] ],
), ),
Worker( Worker(
temporal_client, temporal_client,
@@ -145,19 +147,16 @@ async def main():
activities.find_documents_in_mongodb, activities.find_documents_in_mongodb,
activities.load_latest_data, activities.load_latest_data,
activities.put_last_data_timestamp, activities.put_last_data_timestamp,
# Format and filter notifications # Format and filter notifications
activities.filter_notification_alerts, activities.filter_notification_alerts,
# Send email and export data to postgres # Send email and export data to postgres
activities.build_email_html, activities.build_email_html,
activities.send_email, activities.send_email,
activities.format_log_report, activities.format_log_report,
activities.export_data_to_postgres, activities.export_data_to_postgres,
# Store notification cache # Store notification cache
activities.store_notification_cache activities.store_notification_cache,
] ],
), ),
Worker( Worker(
temporal_client, temporal_client,
@@ -169,17 +168,15 @@ async def main():
activities.find_documents_in_mongodb, activities.find_documents_in_mongodb,
activities.load_latest_data, activities.load_latest_data,
activities.put_last_data_timestamp, activities.put_last_data_timestamp,
# Format and filter notifications # Format and filter notifications
activities.filter_notification_reports, activities.filter_notification_reports,
# Send email and export data to postgres # Send email and export data to postgres
activities.build_email_html, activities.build_email_html,
activities.send_email, activities.send_email,
activities.format_log_report, activities.format_log_report,
activities.export_data_to_postgres activities.export_data_to_postgres,
] ],
) ),
] ]
handlers = [] handlers = []
@@ -193,7 +190,7 @@ async def main():
# If an exception occurs in any of the worker handlers, it will be propagated here. # If an exception occurs in any of the worker handlers, it will be propagated here.
await asyncio.gather(*handlers) await asyncio.gather(*handlers)
except BaseException as e: except BaseException as e:
logger.error(f"An unhandled exception occurred: {e}", exc_info=True) logger.error(f'An unhandled exception occurred: {e}', exc_info=True)
finally: finally:
if notification_handler: if notification_handler:
notification_handler.shutdown() notification_handler.shutdown()
@@ -212,12 +209,12 @@ def start_prometheus_server():
Exits the application if the server fails to start. Exits the application if the server fails to start.
""" """
try: try:
port = int(os.getenv("HTTP_METRICS_PORT", 9090)) port = int(os.getenv('HTTP_METRICS_PORT', 9090))
start_http_server(port) start_http_server(port)
print(f"Prometheus server started on port {port}.") print(f'Prometheus server started on port {port}.')
metrics.APP_UP.labels(pod_id=POD_ID).set(1) metrics.APP_UP.labels(pod_id=POD_ID).set(1)
except Exception as e: except Exception as e:
print(f"Failed to start Prometheus server: {e}") print(f'Failed to start Prometheus server: {e}')
os._exit(1) os._exit(1)

37
orchestrator/workflows/alerts.py
View File

@@ -1,13 +1,15 @@
from temporalio import workflow from temporalio import workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from orchestrator.activities.activities import Activities
from typing import Any
from datetime import timedelta from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy from sientia_do.temporal.policies import retry_policy
from orchestrator.activities.activities import Activities
@workflow.defn(name="alerts")
@workflow.defn(name='alerts')
class Alerts: class Alerts:
""" """
Alerts workflow for real-time error notification delivery. Alerts workflow for real-time error notification delivery.
@@ -51,26 +53,21 @@ class Alerts:
'schedule_name': input_data['schedule_name'], 'schedule_name': input_data['schedule_name'],
'workflow_name': 'alerts', 'workflow_name': 'alerts',
'model_name': '-', 'model_name': '-',
'model_id': '-' 'model_id': '-',
} }
} }
mail_type = "Alerts" mail_type = 'Alerts'
input_data['metadata'] = metadata input_data['metadata'] = metadata
input_data['mail_type'] = mail_type input_data['mail_type'] = mail_type
input_data['base_data_filter'] = { input_data['base_data_filter'] = {'level': 'ERROR'}
'level': 'ERROR'
}
# Call subworkflow "load_notification_package" passing the static filters # Call subworkflow "load_notification_package" passing the static filters
# (level = "ERROR" and timestamp > last timestamp) # (level = "ERROR" and timestamp > last timestamp)
package = await workflow.execute_child_workflow( package = await workflow.execute_child_workflow('load_notification_package', input_data)
'load_notification_package',
input_data
)
if not package['notification_package'] or not package['sending_configs']: if not package['notification_package'] or not package['sending_configs']:
return return
@@ -84,10 +81,10 @@ class Alerts:
**metadata, **metadata,
'notification_package': package['notification_package'], 'notification_package': package['notification_package'],
'sending_configs': package['sending_configs'], 'sending_configs': package['sending_configs'],
'notification_ttl': input_data['notification_ttl'] 'notification_ttl': input_data['notification_ttl'],
}, },
schedule_to_close_timeout=timedelta(seconds=60), schedule_to_close_timeout=timedelta(seconds=60),
retry_policy=retry_policy retry_policy=retry_policy,
) )
if not receiver_groups: if not receiver_groups:
@@ -101,8 +98,8 @@ class Alerts:
'mail_type': mail_type, 'mail_type': mail_type,
'notification_package': receiver_groups, 'notification_package': receiver_groups,
'schema': 'sientia_data', 'schema': 'sientia_data',
'table_name': 'log_report' 'table_name': 'log_report',
} },
) )
if not log_report: if not log_report:
@@ -111,11 +108,7 @@ class Alerts:
# Store the notification_id sendings to avoid sending them again # Store the notification_id sendings to avoid sending them again
await workflow.execute_activity_method( await workflow.execute_activity_method(
Activities.store_notification_cache, Activities.store_notification_cache,
{ {**metadata, 'log_report': log_report, 'sent_ttl': input_data['sent_ttl']},
**metadata,
'log_report': log_report,
'sent_ttl': input_data['sent_ttl']
},
schedule_to_close_timeout=timedelta(seconds=60), schedule_to_close_timeout=timedelta(seconds=60),
retry_policy=retry_policy retry_policy=retry_policy,
) )

144
orchestrator/workflows/orchestrator.py
View File

@@ -1,13 +1,15 @@
from temporalio import workflow from temporalio import workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from orchestrator.activities.activities import Activities
from typing import Any
from datetime import timedelta from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy from sientia_do.temporal.policies import retry_policy
from orchestrator.activities.activities import Activities
@workflow.defn(name="orchestrator")
@workflow.defn(name='orchestrator')
class Orchestrator: class Orchestrator:
""" """
Main orchestrator workflow for pipeline and resource management. Main orchestrator workflow for pipeline and resource management.
@@ -49,7 +51,7 @@ class Orchestrator:
'schedule_name': input_data.get('schedule_name', 'orchestrator'), 'schedule_name': input_data.get('schedule_name', 'orchestrator'),
'model_name': '-', 'model_name': '-',
'model_id': '-', 'model_id': '-',
'workflow_name': input_data['workflow_name'] 'workflow_name': input_data['workflow_name'],
} }
} }
@@ -58,33 +60,28 @@ class Orchestrator:
{ {
**metadata, **metadata,
'query': input_data['pipelines_query'], 'query': input_data['pipelines_query'],
"timestamp_fields": ["updated_at"] 'timestamp_fields': ['updated_at'],
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
opc_servers_handler = workflow.start_local_activity_method( opc_servers_handler = workflow.start_local_activity_method(
Activities.find_documents_in_mongodb, Activities.find_documents_in_mongodb,
{ {**metadata, 'query': input_data['opc_servers_query']},
**metadata,
'query': input_data['opc_servers_query']
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
orchestrated_schedules_handler = workflow.start_local_activity_method( orchestrated_schedules_handler = workflow.start_local_activity_method(
Activities.find_documents_in_mongodb, Activities.find_documents_in_mongodb,
{ {
**metadata, **metadata,
'query': { 'query': {'collection': 'orchestrated_schedules'},
'collection': 'orchestrated_schedules' 'timestamp_fields': ['updated_at'],
},
"timestamp_fields": ["updated_at"]
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
current_slot_config_handler = workflow.start_local_activity_method( current_slot_config_handler = workflow.start_local_activity_method(
@@ -93,7 +90,7 @@ class Orchestrator:
**metadata, **metadata,
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
active_ingestors_handler = workflow.start_local_activity_method( active_ingestors_handler = workflow.start_local_activity_method(
@@ -102,7 +99,7 @@ class Orchestrator:
**metadata, **metadata,
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
pipeline_config = await pipeline_config_handler pipeline_config = await pipeline_config_handler
@@ -113,22 +110,16 @@ class Orchestrator:
formatted_orchestrated_schedules_handler = workflow.start_local_activity_method( formatted_orchestrated_schedules_handler = workflow.start_local_activity_method(
Activities.format_schedule_config, Activities.format_schedule_config,
{ {**metadata, 'schedule_config': orchestrated_schedules},
**metadata,
'schedule_config': orchestrated_schedules
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
schedules_config_handler = workflow.start_local_activity_method( schedules_config_handler = workflow.start_local_activity_method(
Activities.process_schedules, Activities.process_schedules,
{ {**metadata, 'pipelines': pipeline_config},
**metadata,
'pipelines': pipeline_config
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
slot_config_handler = workflow.start_local_activity_method( slot_config_handler = workflow.start_local_activity_method(
@@ -140,7 +131,7 @@ class Orchestrator:
'pipelines': pipeline_config, 'pipelines': pipeline_config,
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
schedules_config = await schedules_config_handler schedules_config = await schedules_config_handler
@@ -152,41 +143,31 @@ class Orchestrator:
{ {
**metadata, **metadata,
'current_schedule_config': formatted_orchestrated_schedules, 'current_schedule_config': formatted_orchestrated_schedules,
'schedule_config': schedules_config 'schedule_config': schedules_config,
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
slot_actions_handler = workflow.start_local_activity_method( slot_actions_handler = workflow.start_local_activity_method(
Activities.create_slot_config, Activities.create_slot_config,
{ {**metadata, 'current_slot_config': current_slot_config, 'slot_config': slot_config},
**metadata,
'current_slot_config': current_slot_config,
'slot_config': slot_config
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
normalize_schedules_handler = workflow.start_activity_method( normalize_schedules_handler = workflow.start_activity_method(
Activities.normalize_schedules, Activities.normalize_schedules,
{ {**metadata, 'orchestrated_schedules': formatted_orchestrated_schedules},
**metadata,
'orchestrated_schedules': formatted_orchestrated_schedules
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
create_collection_with_ttl_index_handler = workflow.start_activity_method( create_collection_with_ttl_index_handler = workflow.start_activity_method(
Activities.create_collection_with_ttl_index, Activities.create_collection_with_ttl_index,
{ {**metadata, 'pipelines': schedules_config['scouter']},
**metadata,
'pipelines': schedules_config['scouter']
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
schedule_actions = await schedule_actions_handler schedule_actions = await schedule_actions_handler
@@ -196,52 +177,37 @@ class Orchestrator:
slot_deletion_report_handler = workflow.start_activity_method( slot_deletion_report_handler = workflow.start_activity_method(
Activities.delete_slots, Activities.delete_slots,
{ {**metadata, 'to_delete': slot_actions['to_delete']},
**metadata,
'to_delete': slot_actions['to_delete']
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
slot_insertion_report_handler = workflow.start_activity_method( slot_insertion_report_handler = workflow.start_activity_method(
Activities.update_slots, Activities.update_slots,
{ {**metadata, 'to_insert': slot_actions['to_insert']},
**metadata,
'to_insert': slot_actions['to_insert']
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
schedule_deletion_report_handler = workflow.start_activity_method( schedule_deletion_report_handler = workflow.start_activity_method(
Activities.delete_schedules, Activities.delete_schedules,
{ {**metadata, 'schedules': schedule_actions['to_delete']},
**metadata,
'schedules': schedule_actions['to_delete']
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
schedule_insertion_report_handler = workflow.start_activity_method( schedule_insertion_report_handler = workflow.start_activity_method(
Activities.create_schedules, Activities.create_schedules,
{ {**metadata, 'schedules': schedule_actions['to_create']},
**metadata,
'schedules': schedule_actions['to_create']
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
schedule_update_report_handler = workflow.start_activity_method( schedule_update_report_handler = workflow.start_activity_method(
Activities.update_schedules, Activities.update_schedules,
{ {**metadata, 'schedules': schedule_actions['to_update']},
**metadata,
'schedules': schedule_actions['to_update']
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
slot_deletion_report = await slot_deletion_report_handler slot_deletion_report = await slot_deletion_report_handler
@@ -251,66 +217,52 @@ class Orchestrator:
schedule_update_report = await schedule_update_report_handler schedule_update_report = await schedule_update_report_handler
if schedule_insertion_report or schedule_update_report or schedule_deletion_report: if schedule_insertion_report or schedule_update_report or schedule_deletion_report:
schedule_report_handler = workflow.start_activity_method( schedule_report_handler = workflow.start_activity_method(
Activities.report_schedule_orchestration, Activities.report_schedule_orchestration,
{ {
**metadata, **metadata,
'created_schedules': schedule_insertion_report, 'created_schedules': schedule_insertion_report,
'updated_schedules': schedule_update_report, 'updated_schedules': schedule_update_report,
'deleted_schedules': schedule_deletion_report 'deleted_schedules': schedule_deletion_report,
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
if slot_insertion_report or slot_deletion_report: if slot_insertion_report or slot_deletion_report:
slot_report_handler = workflow.start_activity_method( slot_report_handler = workflow.start_activity_method(
Activities.report_slot_orchestration, Activities.report_slot_orchestration,
{ {
**metadata, **metadata,
'inserted_slots': slot_insertion_report, 'inserted_slots': slot_insertion_report,
'deleted_slots': slot_deletion_report 'deleted_slots': slot_deletion_report,
}, },
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
if schedule_update_report: if schedule_update_report:
update_pipelines_timestamps_handler = workflow.start_activity_method( update_pipelines_timestamps_handler = workflow.start_activity_method(
Activities.update_pipelines_timestamps, Activities.update_pipelines_timestamps,
{ {**metadata, 'updated_pipelines': schedule_update_report},
**metadata,
'updated_pipelines': schedule_update_report
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
if schedule_insertion_report: if schedule_insertion_report:
create_pipelines_timestamps_handler = workflow.start_activity_method( create_pipelines_timestamps_handler = workflow.start_activity_method(
Activities.create_pipelines_timestamps, Activities.create_pipelines_timestamps,
{ {**metadata, 'created_pipelines': schedule_insertion_report},
**metadata,
'created_pipelines': schedule_insertion_report
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
if schedule_deletion_report: if schedule_deletion_report:
delete_pipelines_timestamps_handler = workflow.start_activity_method( delete_pipelines_timestamps_handler = workflow.start_activity_method(
Activities.delete_pipelines_timestamps, Activities.delete_pipelines_timestamps,
{ {**metadata, 'deleted_pipelines': schedule_deletion_report},
**metadata,
'deleted_pipelines': schedule_deletion_report
},
retry_policy=retry_policy, retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60) start_to_close_timeout=timedelta(seconds=60),
) )
if schedule_insertion_report or schedule_update_report or schedule_deletion_report: if schedule_insertion_report or schedule_update_report or schedule_deletion_report:

25
orchestrator/workflows/reports.py
View File

@@ -1,13 +1,15 @@
from temporalio import workflow from temporalio import workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from orchestrator.activities.activities import Activities
from typing import Any
from datetime import timedelta from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy from sientia_do.temporal.policies import retry_policy
from orchestrator.activities.activities import Activities
@workflow.defn(name="reports")
@workflow.defn(name='reports')
class Reports: class Reports:
""" """
Reports workflow for sending scheduled notification summaries. Reports workflow for sending scheduled notification summaries.
@@ -43,11 +45,11 @@ class Reports:
'schedule_name': input_data['schedule_name'], 'schedule_name': input_data['schedule_name'],
'workflow_name': 'reports', 'workflow_name': 'reports',
'model_name': '-', 'model_name': '-',
'model_id': '-' 'model_id': '-',
} }
} }
mail_type = "Reports" mail_type = 'Reports'
input_data['metadata'] = metadata input_data['metadata'] = metadata
input_data['mail_type'] = mail_type input_data['mail_type'] = mail_type
@@ -57,10 +59,7 @@ class Reports:
# Call subworkflow "load_notification_package" passing the static filters # Call subworkflow "load_notification_package" passing the static filters
# (timestamp > last timestamp) # (timestamp > last timestamp)
package = await workflow.execute_child_workflow( package = await workflow.execute_child_workflow('load_notification_package', input_data)
'load_notification_package',
input_data
)
if not package['notification_package'] or not package['sending_configs']: if not package['notification_package'] or not package['sending_configs']:
return return
@@ -73,10 +72,10 @@ class Reports:
{ {
**metadata, **metadata,
'notification_package': package['notification_package'], 'notification_package': package['notification_package'],
'sending_configs': package['sending_configs'] 'sending_configs': package['sending_configs'],
}, },
schedule_to_close_timeout=timedelta(seconds=60), schedule_to_close_timeout=timedelta(seconds=60),
retry_policy=retry_policy retry_policy=retry_policy,
) )
if not receiver_groups: if not receiver_groups:
@@ -90,6 +89,6 @@ class Reports:
'mail_type': mail_type, 'mail_type': mail_type,
'notification_package': receiver_groups, 'notification_package': receiver_groups,
'schema': 'sientia_data', 'schema': 'sientia_data',
'table_name': 'log_report' 'table_name': 'log_report',
} },
) )

45
orchestrator/workflows/subworkflows/load_notification_package.py
View File

@@ -1,13 +1,15 @@
from temporalio import workflow from temporalio import workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from orchestrator.activities.activities import Activities
from typing import Any
from sientia_do.temporal.policies import retry_policy
from datetime import timedelta from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy
from orchestrator.activities.activities import Activities
@workflow.defn(name="load_notification_package") @workflow.defn(name='load_notification_package')
class LoadNotificationPackage: class LoadNotificationPackage:
""" """
Subworkflow for loading notification data and configuration. Subworkflow for loading notification data and configuration.
@@ -48,28 +50,17 @@ class LoadNotificationPackage:
# Load last timestamp from redis "notification_last_timestamp" # Load last timestamp from redis "notification_last_timestamp"
last_timestamp_handler = workflow.start_local_activity_method( last_timestamp_handler = workflow.start_local_activity_method(
Activities.get_last_data_timestamp, Activities.get_last_data_timestamp,
{ {**metadata, 'mail_type': input_data['mail_type']},
**metadata,
'mail_type': input_data['mail_type']
},
start_to_close_timeout=timedelta(seconds=60), start_to_close_timeout=timedelta(seconds=60),
retry_policy=retry_policy retry_policy=retry_policy,
) )
# In parallel, load sending configs from collection "receiver_groups" # In parallel, load sending configs from collection "receiver_groups"
sending_configs_handler = workflow.start_local_activity_method( sending_configs_handler = workflow.start_local_activity_method(
Activities.find_documents_in_mongodb, Activities.find_documents_in_mongodb,
{ {**metadata, 'query': {'collection': 'receiver_groups', 'filters': {'active': True}}},
**metadata,
'query': {
'collection': 'receiver_groups',
'filters': {
'active': True
}
}
},
start_to_close_timeout=timedelta(seconds=60), start_to_close_timeout=timedelta(seconds=60),
retry_policy=retry_policy retry_policy=retry_policy,
) )
last_timestamp = await last_timestamp_handler last_timestamp = await last_timestamp_handler
@@ -83,10 +74,10 @@ class LoadNotificationPackage:
**metadata, **metadata,
'collection_name': 'notification_queue', 'collection_name': 'notification_queue',
'last_data_timestamp': last_timestamp, 'last_data_timestamp': last_timestamp,
'base_data_filter': input_data['base_data_filter'] 'base_data_filter': input_data['base_data_filter'],
}, },
start_to_close_timeout=timedelta(seconds=60), start_to_close_timeout=timedelta(seconds=60),
retry_policy=retry_policy retry_policy=retry_policy,
) )
sending_configs = await sending_configs_handler sending_configs = await sending_configs_handler
@@ -94,20 +85,16 @@ class LoadNotificationPackage:
return { return {
'last_timestamp': last_timestamp, 'last_timestamp': last_timestamp,
'notification_package': notification_package, 'notification_package': notification_package,
'sending_configs': sending_configs 'sending_configs': sending_configs,
} }
# Put last collected timestamp in redis "notification_last_timestamp" # Put last collected timestamp in redis "notification_last_timestamp"
await workflow.start_activity_method( await workflow.start_activity_method(
Activities.put_last_data_timestamp, Activities.put_last_data_timestamp,
{ {**metadata, 'data': notification_package, 'mail_type': input_data['mail_type']},
**metadata,
'data': notification_package,
'mail_type': input_data['mail_type']
},
start_to_close_timeout=timedelta(seconds=60), start_to_close_timeout=timedelta(seconds=60),
retry_policy=retry_policy retry_policy=retry_policy,
) )
# Return a dict with the following keys: # Return a dict with the following keys:
@@ -117,5 +104,5 @@ class LoadNotificationPackage:
return { return {
'last_timestamp': last_timestamp, 'last_timestamp': last_timestamp,
'notification_package': notification_package, 'notification_package': notification_package,
'sending_configs': sending_configs 'sending_configs': sending_configs,
} }

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

@@ -1,14 +1,16 @@
from temporalio import workflow from temporalio import workflow
with workflow.unsafe.imports_passed_through(): with workflow.unsafe.imports_passed_through():
from orchestrator.activities.activities import Activities
from typing import Any
from datetime import timedelta from datetime import timedelta
from sientia_do.temporal.policies import retry_policy from typing import Any
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ
from sientia_do.temporal.policies import retry_policy
from orchestrator.activities.activities import Activities
@workflow.defn(name="process_notifications") @workflow.defn(name='process_notifications')
class ProcessNotifications: class ProcessNotifications:
""" """
Subworkflow for processing and sending notification emails. Subworkflow for processing and sending notification emails.
@@ -43,30 +45,26 @@ class ProcessNotifications:
Exception: If notification processing fails Exception: If notification processing fails
""" """
metadata = input_data["metadata"] metadata = input_data['metadata']
# Use notification package to create the report html for each group and each model # Use notification package to create the report html for each group and each model
data_to_sent = await workflow.execute_local_activity_method( data_to_sent = await workflow.execute_local_activity_method(
Activities.build_email_html, Activities.build_email_html,
{ {
**metadata, **metadata,
"receiver_groups": input_data["notification_package"], 'receiver_groups': input_data['notification_package'],
"mail_type": input_data["mail_type"] 'mail_type': input_data['mail_type'],
}, },
schedule_to_close_timeout=timedelta(seconds=60), schedule_to_close_timeout=timedelta(seconds=60),
retry_policy=retry_policy retry_policy=retry_policy,
) )
# Send the report html to the receivers of each group # Send the report html to the receivers of each group
log_report = await workflow.execute_activity_method( log_report = await workflow.execute_activity_method(
Activities.send_email, Activities.send_email,
{ {**metadata, 'receiver_groups': data_to_sent, 'mail_type': input_data['mail_type']},
**metadata,
"receiver_groups": data_to_sent,
"mail_type": input_data["mail_type"]
},
schedule_to_close_timeout=timedelta(seconds=60), schedule_to_close_timeout=timedelta(seconds=60),
retry_policy=retry_policy retry_policy=retry_policy,
) )
if not log_report: if not log_report:
@@ -75,13 +73,9 @@ class ProcessNotifications:
# Format the log report to a dataframe to be stored in the database # Format the log report to a dataframe to be stored in the database
log_report = await workflow.execute_local_activity_method( log_report = await workflow.execute_local_activity_method(
Activities.format_log_report, Activities.format_log_report,
{ {**metadata, 'receiver_groups': log_report, 'mail_type': input_data['mail_type']},
**metadata,
"receiver_groups": log_report,
"mail_type": input_data["mail_type"]
},
schedule_to_close_timeout=timedelta(seconds=60), schedule_to_close_timeout=timedelta(seconds=60),
retry_policy=retry_policy retry_policy=retry_policy,
) )
# Store sending log in postgres database "log_report" # Store sending log in postgres database "log_report"
@@ -89,16 +83,16 @@ class ProcessNotifications:
Activities.export_data_to_postgres, Activities.export_data_to_postgres,
{ {
**metadata, **metadata,
"schema": input_data["schema"], 'schema': input_data['schema'],
"table_name": input_data["table_name"], 'table_name': input_data['table_name'],
"data": log_report, 'data': log_report,
'timestamp_conversion': { 'timestamp_conversion': {
'column': 'timestamp', 'column': 'timestamp',
'format': DATETIME_FORMAT_MS_WITH_TZ 'format': DATETIME_FORMAT_MS_WITH_TZ,
} },
}, },
schedule_to_close_timeout=timedelta(seconds=60), schedule_to_close_timeout=timedelta(seconds=60),
retry_policy=retry_policy retry_policy=retry_policy,
) )
# Return the log report to the caller # Return the log report to the caller

160
pyproject.toml Normal file
View File

@@ -0,0 +1,160 @@
[build-system]
requires = ["setuptools>=61.0"]
build-backend = "setuptools.build_meta"
[project]
name = "orchestrator"
version = "0.0.0"
description = "Sientia DataOps Orchestrator - ML Model Orchestration System"
readme = "README.md"
requires-python = ">=3.11"
authors = [
{name = "Aignosi", email = "dev@aignosi.com"}
]
[tool.ruff]
line-length = 100
target-version = "py311"
exclude = [
".git",
".venv",
"venv",
"__pycache__",
"*.pyc",
".pytest_cache",
"htmlcov",
]
[tool.ruff.lint]
select = [
"E", # pycodestyle errors
"W", # pycodestyle warnings
"F", # pyflakes
"I", # isort
"B", # flake8-bugbear
"C4", # flake8-comprehensions
"UP", # pyupgrade
"N", # pep8-naming
"YTT", # flake8-2020
"S", # flake8-bandit
"BLE", # flake8-blind-except
"A", # flake8-builtins
"C90", # mccabe complexity
]
ignore = [
"B023", # ignore blind assignment, we need to assign the schedule to the schedule action
"BLE001",# ignore blind except, we need to send notifications with any error
"E501", # line too long (handled by formatter)
"S101", # use of assert (needed for tests)
"S105", # possible hardcoded password (false positives)
"S106", # possible hardcoded password (false positives)
"N802", # function name should be lowercase (temporal decorators)
"N806", # variable in function should be lowercase
]
[tool.ruff.lint.per-file-ignores]
"tests/**/*.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
[tool.ruff.format]
quote-style = "single"
indent-style = "space"
line-ending = "auto"
[tool.mypy]
python_version = "3.11"
warn_return_any = false
warn_unused_configs = true
disallow_untyped_defs = false
disallow_incomplete_defs = false
check_untyped_defs = true
no_implicit_optional = true
warn_redundant_casts = true
warn_unused_ignores = false
warn_no_return = true
strict_equality = true
ignore_missing_imports = true
# Ignore missing imports for external packages
[[tool.mypy.overrides]]
module = "temporalio.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "sientia_do.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "mlflow.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "prometheus_client.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "sientia.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "pandas.*"
ignore_missing_imports = true
[tool.pytest.ini_options]
testpaths = ["tests"]
python_files = ["test_*.py"]
python_classes = ["Test*"]
python_functions = ["test_*"]
addopts = [
"-v",
"--strict-markers",
"--cov=model_manager",
"--cov-report=term-missing",
"--cov-report=html",
"--cov-report=xml",
]
markers = [
"asyncio: marks tests as async",
"integration: marks tests as integration tests",
"unit: marks tests as unit tests",
]
[tool.coverage.run]
source = ["model_manager"]
omit = [
"*/tests/*",
"*/venv/*",
"*/__pycache__/*",
"*/site-packages/*",
]
branch = true
[tool.coverage.report]
precision = 2
show_missing = true
skip_covered = false
exclude_lines = [
"pragma: no cover",
"def __repr__",
"def __str__",
"raise AssertionError",
"raise NotImplementedError",
"if __name__ == .__main__.:",
"if TYPE_CHECKING:",
"class .*\\bProtocol\\):",
"@(abc\\.)?abstractmethod",
]
[tool.coverage.html]
directory = "htmlcov"
[tool.bandit]
exclude_dirs = ["tests", "venv", ".venv"]
skips = ["B101", "B601"] # Skip assert and shell injection in controlled environments

19
requirements-dev.txt Normal file
View File

@@ -0,0 +1,19 @@
# Development and Testing Dependencies
# These packages are only needed for development, testing, and code quality checks
# Install with: pip install -r requirements-dev.txt
# Code Quality & Linting
ruff>=0.1.0 # Fast Python linter and formatter (replaces flake8, black, isort)
mypy>=1.7.0 # Static type checker
bandit>=1.7.5 # Security vulnerability scanner
pandas-stubs>=2.0.0 # Type stubs for pandas
types-requests>=2.31.0 # Type stubs for requests
# Testing
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)
# Development Tools
ipython>=8.12.0 # Enhanced Python shell
ipdb>=0.13.13 # IPython debugger

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

@@ -1,11 +1,10 @@
from unittest.mock import patch, MagicMock, ANY from unittest.mock import ANY, MagicMock, patch
from pytest import mark
from orchestrator.activities import mongo_db
from orchestrator.activities.activities import Activities from orchestrator.activities.activities import Activities
from orchestrator.activities.mongo_db import MongoDB
from orchestrator.activities.temporal_manager import TemporalManager
from orchestrator.activities.slot_manager import SlotManager
from orchestrator.activities.formatters import Formatters from orchestrator.activities.formatters import Formatters
from orchestrator.activities.mongo_db import MongoDB
from orchestrator.activities.slot_manager import SlotManager
from orchestrator.activities.temporal_manager import TemporalManager
@patch('orchestrator.activities.mongo_db.MongoDB.__init__') @patch('orchestrator.activities.mongo_db.MongoDB.__init__')
@@ -14,37 +13,33 @@ from orchestrator.activities.formatters import Formatters
@patch('orchestrator.activities.formatters.Formatters.__init__') @patch('orchestrator.activities.formatters.Formatters.__init__')
@patch('orchestrator.activities.email.Email.__init__') @patch('orchestrator.activities.email.Email.__init__')
@patch('sientia_do.temporal.activities.postgres.Postgres.__init__') @patch('sientia_do.temporal.activities.postgres.Postgres.__init__')
def test___init__(mock_postgres_init, def test___init__(
mock_postgres_init,
mock_email_init, mock_email_init,
mock_formatters_init, mock_formatters_init,
mock_slot_manager_init, mock_slot_manager_init,
mock_temporal_manager_init, mock_temporal_manager_init,
mock_mongodb_init): mock_mongodb_init,
):
mongo_db_config = { mongo_db_config = {
'connection_string': 'mongodb://localhost:27017', 'connection_string': 'mongodb://localhost:27017',
'database_name': 'test_db', 'database_name': 'test_db',
'ttl_index_seconds': 3600 'ttl_index_seconds': 3600,
} }
redis_config = { redis_config = {'host': 'localhost', 'port': 6379, 'username': 'admin', 'password': 'password'}
'host': 'localhost',
'port': 6379,
'username': 'admin',
'password': 'password'
}
temporal_config = { temporal_config = {
'temporal_host': 'localhost', 'temporal_host': 'localhost',
'temporal_scouter_namespace': 'scouter', 'temporal_scouter_namespace': 'scouter',
'temporal_laborious_namespace': 'laborious' 'temporal_laborious_namespace': 'laborious',
} }
email_config = { email_config = {
'sender_email': 'test@test.com', 'sender_email': 'test@test.com',
'sender_password': 'test', 'sender_password': 'test',
'smtp_server': 'test', 'smtp_server': 'test',
'smtp_port': 587 'smtp_port': 587,
} }
postgres_config = { postgres_config = {
@@ -67,7 +62,7 @@ def test___init__(mock_postgres_init,
email_config=email_config, email_config=email_config,
postgres_config=postgres_config, postgres_config=postgres_config,
logger=logger, logger=logger,
notification_handler=notification_handler notification_handler=notification_handler,
) )
assert isinstance(activities, Activities) assert isinstance(activities, Activities)
@@ -78,12 +73,12 @@ def test___init__(mock_postgres_init,
mock_slot_manager_init.assert_called_once_with( mock_slot_manager_init.assert_called_once_with(
ANY, ANY,
host="localhost", host='localhost',
port=6379, port=6379,
username="admin", username='admin',
password="password", password='password',
logger=logger, logger=logger,
notification_handler=notification_handler notification_handler=notification_handler,
) )
mock_mongodb_init.assert_called_once_with( mock_mongodb_init.assert_called_once_with(
@@ -92,7 +87,7 @@ def test___init__(mock_postgres_init,
database_name='test_db', database_name='test_db',
ttl_index_seconds=3600, ttl_index_seconds=3600,
logger=logger, logger=logger,
notification_handler=notification_handler notification_handler=notification_handler,
) )
mock_temporal_manager_init.assert_called_once_with( mock_temporal_manager_init.assert_called_once_with(
@@ -101,7 +96,7 @@ def test___init__(mock_postgres_init,
scouter_namespace='scouter', scouter_namespace='scouter',
laborious_namespace='laborious', laborious_namespace='laborious',
logger=logger, logger=logger,
notification_handler=notification_handler notification_handler=notification_handler,
) )
mock_formatters_init.assert_called_once_with( mock_formatters_init.assert_called_once_with(
@@ -109,7 +104,7 @@ def test___init__(mock_postgres_init,
scouter_namespace='scouter', scouter_namespace='scouter',
laborious_namespace='laborious', laborious_namespace='laborious',
logger=logger, logger=logger,
notification_handler=notification_handler notification_handler=notification_handler,
) )
@@ -122,7 +117,8 @@ def test___init__(mock_postgres_init,
@patch('orchestrator.activities.email.Email.shutdown') @patch('orchestrator.activities.email.Email.shutdown')
@patch('sientia_do.temporal.activities.postgres.Postgres.close') @patch('sientia_do.temporal.activities.postgres.Postgres.close')
@patch('orchestrator.activities.mongo_db.MongoDB.shutdown') @patch('orchestrator.activities.mongo_db.MongoDB.shutdown')
def test_shutdown(mock_mongodb_close, def test_shutdown(
mock_mongodb_close,
mock_postgres_shutdown, mock_postgres_shutdown,
mock_email_close, mock_email_close,
mock_postgres_init, mock_postgres_init,
@@ -131,8 +127,7 @@ def test_shutdown(mock_mongodb_close,
mock_slot_manager_init, mock_slot_manager_init,
mock_temporal_manager_init, mock_temporal_manager_init,
mock_mongodb_init, mock_mongodb_init,
): ):
activities = Activities( activities = Activities(
temporal_config=MagicMock(), temporal_config=MagicMock(),
redis_config=MagicMock(), redis_config=MagicMock(),
@@ -140,7 +135,7 @@ def test_shutdown(mock_mongodb_close,
email_config=MagicMock(), email_config=MagicMock(),
postgres_config=MagicMock(), postgres_config=MagicMock(),
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock() notification_handler=MagicMock(),
) )
activities.shutdown() activities.shutdown()

55
tests/orchestrator/activities/test_couchbase.py
View File

@@ -1,16 +1,18 @@
from unittest.mock import MagicMock, patch, ANY from unittest.mock import ANY, MagicMock, patch
from pytest import fixture, mark, raises from pytest import fixture, mark, raises
from sientia_do.notifications.models import NotificationLevel from sientia_do.notifications.models import NotificationLevel
from orchestrator.activities.couchbase import Couchbase from orchestrator.activities.couchbase import Couchbase
@fixture @fixture
@patch("orchestrator.activities.couchbase.Cluster") @patch('orchestrator.activities.couchbase.Cluster')
def couchbase(_cluster_mock): def couchbase(_cluster_mock):
return Couchbase( return Couchbase(
connection_string="couchbase://localhost", connection_string='couchbase://localhost',
username="admin", username='admin',
password="password", password='password',
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock(), notification_handler=MagicMock(),
) )
@@ -22,27 +24,30 @@ def test_shutdown_success(couchbase):
def test_shutdown_failure(couchbase): def test_shutdown_failure(couchbase):
couchbase.cluster.close.side_effect = Exception("Test error") couchbase.cluster.close.side_effect = Exception('Test error')
couchbase.shutdown() couchbase.shutdown()
couchbase.logger.error.assert_called_once_with( couchbase.logger.error.assert_called_once_with(
f"Failed to close Couchbase connection: {couchbase.cluster.close.side_effect}") f'Failed to close Couchbase connection: {couchbase.cluster.close.side_effect}'
)
@mark.asyncio @mark.asyncio
async def test_load_query_from_couchbase_success(couchbase): async def test_load_query_from_couchbase_success(couchbase):
couchbase.cluster.query.return_value.rows.return_value = [ couchbase.cluster.query.return_value.rows.return_value = [
{"id": "1", "name": "test"}, {'id': '1', 'name': 'test'},
{"id": "2", "name": "test2"}, {'id': '2', 'name': 'test2'},
] ]
query = "SELECT * FROM bucket" query = 'SELECT * FROM bucket'
result = await couchbase.load_query_from_couchbase({ result = await couchbase.load_query_from_couchbase(
"query": query, {
}) 'query': query,
}
)
assert result == [ assert result == [
{"id": "1", "name": "test"}, {'id': '1', 'name': 'test'},
{"id": "2", "name": "test2"}, {'id': '2', 'name': 'test2'},
] ]
couchbase.cluster.query.assert_called_once_with(query) couchbase.cluster.query.assert_called_once_with(query)
couchbase.notification_handler.build_and_send_notification.assert_not_called() couchbase.notification_handler.build_and_send_notification.assert_not_called()
@@ -50,19 +55,21 @@ async def test_load_query_from_couchbase_success(couchbase):
@mark.asyncio @mark.asyncio
async def test_load_query_from_couchbase_failure(couchbase): async def test_load_query_from_couchbase_failure(couchbase):
couchbase.cluster.query.side_effect = Exception("Test error") couchbase.cluster.query.side_effect = ValueError('Test error')
query = "SELECT * FROM bucket" query = 'SELECT * FROM bucket'
with raises(Exception): with raises(ValueError):
await couchbase.load_query_from_couchbase({ await couchbase.load_query_from_couchbase(
"query": query, {
}) 'query': query,
}
)
couchbase.cluster.query.assert_called_once_with(query) couchbase.cluster.query.assert_called_once_with(query)
couchbase.notification_handler.build_and_send_notification.assert_called_once_with( couchbase.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="COUCHBASE_LOAD_QUERY_ERROR", notification_id='COUCHBASE_LOAD_QUERY_ERROR',
message="Failed to execute couchbase query: Test error", message='Failed to execute couchbase query: Test error',
block="load_query_from_couchbase", block='load_query_from_couchbase',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=ANY, attachment_content=ANY,
) )

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

@@ -1,6 +1,7 @@
from smtplib import SMTPServerDisconnected from smtplib import SMTPServerDisconnected
from unittest.mock import MagicMock, call, patch from unittest.mock import MagicMock, call, patch
from pytest import mark, fixture
from pytest import fixture, mark
from orchestrator.activities.email import Email from orchestrator.activities.email import Email
@@ -10,12 +11,12 @@ from orchestrator.activities.email import Email
@patch('orchestrator.activities.email.smtplib') @patch('orchestrator.activities.email.smtplib')
def email(smtplib, email_builder): def email(smtplib, email_builder):
email = Email( email = Email(
sender_email="test@test.com", sender_email='test@test.com',
sender_password="test", sender_password='test',
smtp_server="test", smtp_server='test',
smtp_port=587, smtp_port=587,
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock() notification_handler=MagicMock(),
) )
email_builder.send_notification = MagicMock() email_builder.send_notification = MagicMock()
@@ -26,22 +27,21 @@ def email(smtplib, email_builder):
@patch('orchestrator.activities.email.smtplib') @patch('orchestrator.activities.email.smtplib')
def test___init___with_password(smtplib, email_builder): def test___init___with_password(smtplib, email_builder):
email = Email( email = Email(
sender_email="test@test.com", sender_email='test@test.com',
sender_password="test", sender_password='test',
smtp_server="test", smtp_server='test',
smtp_port=587, smtp_port=587,
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock() notification_handler=MagicMock(),
) )
assert email.sender_email == "test@test.com" assert email.sender_email == 'test@test.com'
assert email.sender_password == "test" assert email.sender_password == 'test'
assert email.smtp_port == 587 assert email.smtp_port == 587
smtplib.SMTP.assert_called_once_with("test", 587, timeout=20) smtplib.SMTP.assert_called_once_with('test', 587, timeout=20)
smtplib.SMTP.return_value.starttls.assert_called_once() smtplib.SMTP.return_value.starttls.assert_called_once()
smtplib.SMTP.return_value.login.assert_called_once_with( smtplib.SMTP.return_value.login.assert_called_once_with('test@test.com', 'test')
"test@test.com", "test")
assert email.server == smtplib.SMTP.return_value assert email.server == smtplib.SMTP.return_value
@@ -50,28 +50,28 @@ def test___init___with_password(smtplib, email_builder):
@patch('orchestrator.activities.email.smtplib') @patch('orchestrator.activities.email.smtplib')
def test___init___without_password(smtplib, email_builder): def test___init___without_password(smtplib, email_builder):
email = Email( email = Email(
sender_email="test@test.com", sender_email='test@test.com',
sender_password=None, sender_password=None,
smtp_server="test", smtp_server='test',
smtp_port=587, smtp_port=587,
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock() notification_handler=MagicMock(),
) )
assert email.sender_email == "test@test.com" assert email.sender_email == 'test@test.com'
assert email.sender_password is None assert email.sender_password is None
assert email.smtp_port == 587 assert email.smtp_port == 587
smtplib.SMTP.assert_called_once_with("test", 587, timeout=20) smtplib.SMTP.assert_called_once_with('test', 587, timeout=20)
assert email.server == smtplib.SMTP.return_value assert email.server == smtplib.SMTP.return_value
metadata = { metadata = {
"metadata": { 'metadata': {
"schedule_name": "test", 'schedule_name': 'test',
"model_name": "test", 'model_name': 'test',
"model_id": "test", 'model_id': 'test',
"workflow_name": "test", 'workflow_name': 'test',
} }
} }
@@ -84,53 +84,42 @@ def test_shutdown(email):
@mark.asyncio @mark.asyncio
async def test_build_email_html(email): async def test_build_email_html(email):
email.email_builder.build_email = MagicMock( email.email_builder.build_email = MagicMock(return_value='test')
return_value="test"
)
input_data = { input_data = {
**metadata, **metadata,
"receiver_groups": { 'receiver_groups': {
"group_1": { 'group_1': {
"notifications": [ 'notifications': [
{ {'type': 'test', 'subject': 'test', 'body': 'test'},
"type": "test", {'type': 'test', 'subject': 'test', 'body': 'test'},
"subject": "test",
"body": "test"
},
{
"type": "test",
"subject": "test",
"body": "test"
}
] ]
} }
}, },
"mail_type": "test" 'mail_type': 'test',
} }
response = await email.build_email_html(input_data) response = await email.build_email_html(input_data)
assert response == { assert response == {
"group_1": { 'group_1': {
"notifications": [ 'notifications': [
{ {
"type": "test", 'type': 'test',
"subject": "test", 'subject': 'test',
"body": "test", 'body': 'test',
}, },
{ {
"type": "test", 'type': 'test',
"subject": "test", 'subject': 'test',
"body": "test", 'body': 'test',
} },
], ],
"html": "test" 'html': 'test',
} }
} }
email.email_builder.build_email.assert_called_once_with( email.email_builder.build_email.assert_called_once_with(
input_data['receiver_groups']['group_1']['notifications'], input_data['receiver_groups']['group_1']['notifications'], input_data['mail_type']
input_data['mail_type']
) )
@@ -140,18 +129,9 @@ def test_handle_attachments_success(encoders, mime_base, email):
message = MagicMock() message = MagicMock()
attachments = [ attachments = [
{ {'filename': 'file_1', 'attachment_content': 'test_content_1'},
"filename": "file_1", {'filename': 'file_2', 'attachment_content': 'test_content_2'},
"attachment_content": "test_content_1" {'filename': 'file_3', 'attachment_content': 'test_content_3'},
},
{
"filename": "file_2",
"attachment_content": "test_content_2"
},
{
"filename": "file_3",
"attachment_content": "test_content_3"
}
] ]
response = email.handle_attachments(attachments, message) response = email.handle_attachments(attachments, message)
@@ -163,9 +143,9 @@ def test_handle_attachments_success(encoders, mime_base, email):
mime_base.return_value.set_payload.assert_has_calls( mime_base.return_value.set_payload.assert_has_calls(
[ [
call("test_content_1".encode('utf-8')), call(b'test_content_1'),
call("test_content_2".encode('utf-8')), call(b'test_content_2'),
call("test_content_3".encode('utf-8')) call(b'test_content_3'),
] ]
) )
@@ -176,7 +156,7 @@ def test_handle_attachments_success(encoders, mime_base, email):
[ [
call('Content-Disposition', 'attachment; filename="file_1"'), call('Content-Disposition', 'attachment; filename="file_1"'),
call('Content-Disposition', 'attachment; filename="file_2"'), call('Content-Disposition', 'attachment; filename="file_2"'),
call('Content-Disposition', 'attachment; filename="file_3"') call('Content-Disposition', 'attachment; filename="file_3"'),
] ]
) )
@@ -188,19 +168,14 @@ def test_handle_attachments_success(encoders, mime_base, email):
def test_handle_attachments_failure(mime_base, email): def test_handle_attachments_failure(mime_base, email):
message = MagicMock() message = MagicMock()
mime_base.side_effect = Exception("test") mime_base.side_effect = Exception('test')
attachments = [ attachments = [{'filename': 'file_1', 'attachment_content': 'test_content_1'}]
{
"filename": "file_1",
"attachment_content": "test_content_1"
}
]
try: try:
email.handle_attachments(attachments, message) email.handle_attachments(attachments, message)
except Exception as e: except Exception as e:
assert str(e) == "test" assert str(e) == 'test'
assert message.attach.call_count == 0 assert message.attach.call_count == 0
@@ -210,86 +185,69 @@ def test_try_send_email_success(email):
msg = MagicMock() msg = MagicMock()
email.try_send_email(msg, "test") email.try_send_email(msg, 'test')
email.server.sendmail.assert_called_once_with( email.server.sendmail.assert_called_once_with(
"test@test.com", "test", msg.as_string.return_value) 'test@test.com', 'test', msg.as_string.return_value
)
@patch('orchestrator.activities.email.smtplib.SMTP') @patch('orchestrator.activities.email.smtplib.SMTP')
def test_try_send_email_reconnect_quit_success(smtp, email): def test_try_send_email_reconnect_quit_success(smtp, email):
email.server.sendmail = MagicMock( email.server.sendmail = MagicMock(side_effect=SMTPServerDisconnected('test'))
side_effect=SMTPServerDisconnected("test")
)
email.server.quit = MagicMock() email.server.quit = MagicMock()
msg = MagicMock() msg = MagicMock()
email.try_send_email(msg, "test") email.try_send_email(msg, 'test')
smtp.assert_has_calls([ smtp.assert_has_calls([call('test', 587, timeout=20)])
call("test", 587, timeout=20)
])
smtp.return_value.starttls.assert_called_once() smtp.return_value.starttls.assert_called_once()
smtp.return_value.login.assert_called_once_with( smtp.return_value.login.assert_called_once_with('test@test.com', 'test')
"test@test.com", "test")
smtp.return_value.sendmail.assert_called_once_with( smtp.return_value.sendmail.assert_called_once_with(
"test@test.com", "test", msg.as_string.return_value) 'test@test.com', 'test', msg.as_string.return_value
)
@patch('orchestrator.activities.email.smtplib.SMTP') @patch('orchestrator.activities.email.smtplib.SMTP')
def test_try_send_email_reconnect_quit_failure_disconnect(smtp, email): def test_try_send_email_reconnect_quit_failure_disconnect(smtp, email):
email.server.sendmail = MagicMock( email.server.sendmail = MagicMock(side_effect=SMTPServerDisconnected('test'))
side_effect=SMTPServerDisconnected("test") email.server.quit = MagicMock(side_effect=SMTPServerDisconnected('test'))
)
email.server.quit = MagicMock(
side_effect=SMTPServerDisconnected("test")
)
msg = MagicMock() msg = MagicMock()
email.try_send_email(msg, "test") email.try_send_email(msg, 'test')
smtp.assert_has_calls([ smtp.assert_has_calls([call('test', 587, timeout=20)])
call("test", 587, timeout=20)
])
smtp.return_value.starttls.assert_called_once() smtp.return_value.starttls.assert_called_once()
smtp.return_value.login.assert_called_once_with( smtp.return_value.login.assert_called_once_with('test@test.com', 'test')
"test@test.com", "test")
smtp.return_value.sendmail.assert_called_once_with( smtp.return_value.sendmail.assert_called_once_with(
"test@test.com", "test", msg.as_string.return_value) 'test@test.com', 'test', msg.as_string.return_value
)
@patch('orchestrator.activities.email.smtplib.SMTP') @patch('orchestrator.activities.email.smtplib.SMTP')
def test_try_send_email_reconnect_quit_failure(smtp, email): def test_try_send_email_reconnect_quit_failure(smtp, email):
email.server.sendmail = MagicMock( email.server.sendmail = MagicMock(side_effect=SMTPServerDisconnected('test'))
side_effect=SMTPServerDisconnected("test") email.server.quit = MagicMock(side_effect=Exception('test'))
)
email.server.quit = MagicMock(
side_effect=Exception("test")
)
msg = MagicMock() msg = MagicMock()
try: try:
email.try_send_email(msg, "test") email.try_send_email(msg, 'test')
except Exception as e: except Exception as e:
assert str(e) == "test" assert str(e) == 'test'
else: else:
assert False, "Expected exception" raise AssertionError('Expected exception')
@mark.asyncio @mark.asyncio
async def test_send_email_without_smtp_server(email): async def test_send_email_without_smtp_server(email):
email.smtp_server = None email.smtp_server = None
input_data = { input_data = {**metadata, 'receiver_groups': {}, 'mail_type': 'test_TYPE'}
**metadata,
"receiver_groups": {},
"mail_type": "test_TYPE"
}
response = await email.send_email(input_data) response = await email.send_email(input_data)
assert response == {} assert response == {}
@@ -301,41 +259,33 @@ async def test_send_email_without_smtp_server(email):
async def test_send_email(mimemultipart, mimetext, email): async def test_send_email(mimemultipart, mimetext, email):
side_effect_1 = MagicMock() side_effect_1 = MagicMock()
side_effect_2 = MagicMock() side_effect_2 = MagicMock()
mimemultipart.side_effect = [ mimemultipart.side_effect = [side_effect_1, side_effect_2]
side_effect_1,
side_effect_2
]
email.email_builder.send_notification = MagicMock() email.email_builder.send_notification = MagicMock()
email.try_send_email = MagicMock( email.try_send_email = MagicMock(side_effect=[None, Exception('test')])
side_effect=[
None,
Exception("test")
]
)
input_data = { input_data = {
**metadata, **metadata,
"receiver_groups": { 'receiver_groups': {
"group_1": { 'group_1': {
"members": ["test@test.com", "test2@test.com"], 'members': ['test@test.com', 'test2@test.com'],
"notifications": [ 'notifications': [
{ {
"attachment_content": "test_content_1", 'attachment_content': 'test_content_1',
"trigger": "test_trigger", 'trigger': 'test_trigger',
"notification_id": "test_notification_id" 'notification_id': 'test_notification_id',
} }
], ],
"html": "test_html1" 'html': 'test_html1',
}, },
"group_2": { 'group_2': {
"members": ["test3@test.com", "test4@test.com"], 'members': ['test3@test.com', 'test4@test.com'],
"notifications": [], 'notifications': [],
"html": "test_html2" 'html': 'test_html2',
}
}, },
"mail_type": "test_TYPE" },
'mail_type': 'test_TYPE',
} }
response = await email.send_email(input_data) response = await email.send_email(input_data)
@@ -345,18 +295,13 @@ async def test_send_email(mimemultipart, mimetext, email):
assert mimemultipart.call_count == 2 assert mimemultipart.call_count == 2
mimetext.assert_has_calls( mimetext.assert_has_calls([call('test_html1', 'html'), call('test_html2', 'html')])
[
call("test_html1", "html"),
call("test_html2", "html")
]
)
side_effect_1.__setitem__.assert_has_calls( side_effect_1.__setitem__.assert_has_calls(
[ [
call('From', 'test@test.com'), call('From', 'test@test.com'),
call('To', 'test@test.com, test2@test.com'), call('To', 'test@test.com, test2@test.com'),
call('Subject', 'SIENTIA™ test_TYPE') call('Subject', 'SIENTIA™ test_TYPE'),
] ]
) )
@@ -364,14 +309,14 @@ async def test_send_email(mimemultipart, mimetext, email):
[ [
call('From', 'test@test.com'), call('From', 'test@test.com'),
call('To', 'test3@test.com, test4@test.com'), call('To', 'test3@test.com, test4@test.com'),
call('Subject', 'SIENTIA™ test_TYPE') call('Subject', 'SIENTIA™ test_TYPE'),
] ]
) )
email.try_send_email.assert_has_calls( email.try_send_email.assert_has_calls(
[ [
call(side_effect_1, 'test@test.com, test2@test.com'), call(side_effect_1, 'test@test.com, test2@test.com'),
call(side_effect_2, 'test3@test.com, test4@test.com') call(side_effect_2, 'test3@test.com, test4@test.com'),
] ]
) )

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

File diff suppressed because it is too large Load Diff

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

@@ -1,512 +1,461 @@
from curses import meta
from datetime import datetime from datetime import datetime
from unittest.mock import MagicMock, patch, ANY from unittest.mock import ANY, MagicMock, patch
from pytest import fixture, mark from pytest import fixture, mark
from orchestrator.activities.mongo_db import clear_mongo_id
from orchestrator.activities.mongo_db import MongoDB
from sientia_do.notifications.models import NotificationLevel from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ
from orchestrator.activities.mongo_db import MongoDB, clear_mongo_id
def test_clear_mongo_id(): def test_clear_mongo_id():
input_data = [ input_data = [
[ [
{ {
"name": "test", 'name': 'test',
"_id": "12345", '_id': '12345',
} }
], ],
{ {
"name": "test", 'name': 'test',
"_id": "12345", '_id': '12345',
"nested": { 'nested': {
"_id": "67890", '_id': '67890',
"value": [1, 2, 3], 'value': [1, 2, 3],
"list": [{"_id": "abcde", "item": "value"}] 'list': [{'_id': 'abcde', 'item': 'value'}],
},
'nested_list': [{'_id': 'fghij', 'item': 'value1'}, {'_id': 'klmno', 'item': 'value2'}],
}, },
"nested_list": [
{"_id": "fghij", "item": "value1"},
{"_id": "klmno", "item": "value2"}
]
}
] ]
output = clear_mongo_id(input_data) output = clear_mongo_id(input_data)
assert output == [ assert output == [
[ [{'name': 'test'}],
{"name": "test"}
],
{ {
"name": "test", 'name': 'test',
"nested": { 'nested': {'value': [1, 2, 3], 'list': [{'item': 'value'}]},
"value": [1, 2, 3], 'nested_list': [{'item': 'value1'}, {'item': 'value2'}],
"list": [{"item": "value"}]
}, },
"nested_list": [
{"item": "value1"},
{"item": "value2"}
] ]
}]
@fixture @fixture
@patch("orchestrator.activities.mongo_db.MongoClient") @patch('orchestrator.activities.mongo_db.MongoClient')
def mongo_db(mongo_mock): def mongo_db(mongo_mock):
mongo = ( mongo = MongoDB(
MongoDB( connection_string='mongodb://localhost:27017',
connection_string="mongodb://localhost:27017", database_name='test_db',
database_name="test_db",
ttl_index_seconds=3600, ttl_index_seconds=3600,
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock() notification_handler=MagicMock(),
)
) )
mongo.send_notification = MagicMock() mongo.send_notification = MagicMock()
return mongo return mongo
@patch("orchestrator.activities.mongo_db.MongoClient") @patch('orchestrator.activities.mongo_db.MongoClient')
def test___init__(mongo_mock): def test___init__(mongo_mock):
mongo_db = MongoDB( mongo_db = MongoDB(
connection_string="mongodb://localhost:27017", connection_string='mongodb://localhost:27017',
database_name="test_db", database_name='test_db',
ttl_index_seconds=3600, ttl_index_seconds=3600,
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock() notification_handler=MagicMock(),
)
assert mongo_db.connection_string == "mongodb://localhost:27017"
assert mongo_db.database_name == "test_db"
mongo_mock.assert_called_once_with(
"mongodb://localhost:27017", serverSelectionTimeoutMS=5000
) )
assert mongo_db.connection_string == 'mongodb://localhost:27017'
assert mongo_db.database_name == 'test_db'
mongo_mock.assert_called_once_with('mongodb://localhost:27017', serverSelectionTimeoutMS=5000)
mongo_db.client.server_info.assert_called_once() mongo_db.client.server_info.assert_called_once()
mongo_db.client.__getitem__.assert_called_once_with("test_db") mongo_db.client.__getitem__.assert_called_once_with('test_db')
def test_shutdown_success(mongo_db): def test_shutdown_success(mongo_db):
mongo_db.shutdown() mongo_db.shutdown()
mongo_db.client.close.assert_called_once() mongo_db.client.close.assert_called_once()
mongo_db.logger.info.assert_any_call("Closing MongoDB connection...") mongo_db.logger.info.assert_any_call('Closing MongoDB connection...')
mongo_db.logger.info.assert_any_call( mongo_db.logger.info.assert_any_call('MongoDB connection closed successfully')
"MongoDB connection closed successfully")
def test_shutdown_failure(mongo_db): def test_shutdown_failure(mongo_db):
mongo_db.client.close.side_effect = Exception("Close failed") mongo_db.client.close.side_effect = Exception('Close failed')
mongo_db.shutdown() mongo_db.shutdown()
mongo_db.logger.error.assert_called_once_with( mongo_db.logger.error.assert_called_once_with(
"Failed to close MongoDB connection: Close failed" 'Failed to close MongoDB connection: Close failed'
) )
@mark.asyncio @mark.asyncio
async def test_find_documents_in_mongodb_success(mongo_db): async def test_find_documents_in_mongodb_success(mongo_db):
input_data = {"collection": "test_collection", "filters": { input_data = {'collection': 'test_collection', 'filters': {'name': {'$exists': True}}}
"name": {"$exists": True}
}}
mock_collection = MagicMock() mock_collection = MagicMock()
mock_collection.find.return_value = [ mock_collection.find.return_value = [
{ {
"_id": "12345", '_id': '12345',
"name": "test1", 'name': 'test1',
"timestamp": datetime.strptime( 'timestamp': datetime.strptime(
"2023-01-01 12:00:00.000000+0000", DATETIME_FORMAT_MS_WITH_TZ)}, '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
),
},
{ {
"_id": "67890", '_id': '67890',
"name": "test2", 'name': 'test2',
"timestamp": datetime.strptime( 'timestamp': datetime.strptime(
"2023-01-01 12:00:00.000000+0000", DATETIME_FORMAT_MS_WITH_TZ)} '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
),
},
] ]
mongo_db.database.__getitem__.return_value = mock_collection mongo_db.database.__getitem__.return_value = mock_collection
result = await mongo_db.find_documents_in_mongodb( result = await mongo_db.find_documents_in_mongodb(
{ {'query': input_data, 'timestamp_fields': ['timestamp']}
"query": input_data,
"timestamp_fields": ["timestamp"]
})
assert len(result) == 2
assert result[0] == {"name": "test1",
"timestamp": "2023-01-01 12:00:00.000000+0000"}
assert result[1] == {"name": "test2",
"timestamp": "2023-01-01 12:00:00.000000+0000"}
mock_collection.find.assert_called_once_with(
{"name": {"$exists": True}}, {"_id": 0}
) )
assert len(result) == 2
assert result[0] == {'name': 'test1', 'timestamp': '2023-01-01 12:00:00.000000+0000'}
assert result[1] == {'name': 'test2', 'timestamp': '2023-01-01 12:00:00.000000+0000'}
mock_collection.find.assert_called_once_with({'name': {'$exists': True}}, {'_id': 0})
metadata = { metadata = {
"metadata": { 'metadata': {
"schedule_name": "test_schedule_name", 'schedule_name': 'test_schedule_name',
"workflow_name": "test_workflow_name", 'workflow_name': 'test_workflow_name',
"model_name": "test_model_name", 'model_name': 'test_model_name',
"model_id": "test_model_id" 'model_id': 'test_model_id',
} }
} }
@mark.asyncio @mark.asyncio
async def test_find_documents_in_mongodb_failure(mongo_db): async def test_find_documents_in_mongodb_failure(mongo_db):
input_data = {"collection": "test_collection", "filters": { input_data = {'collection': 'test_collection', 'filters': {'name': {'$exists': True}}}
"name": {"$exists": True}
}}
mongo_db.database.__getitem__.return_value = MagicMock( mongo_db.database.__getitem__.return_value = MagicMock(
find=MagicMock(side_effect=Exception("Error")) find=MagicMock(side_effect=Exception('Error'))
) )
try: try:
await mongo_db.find_documents_in_mongodb( await mongo_db.find_documents_in_mongodb({'query': input_data, **metadata})
{
"query": input_data,
**metadata
})
except Exception as e: except Exception as e:
assert str(e) == "Error" assert str(e) == 'Error'
mongo_db.send_notification.assert_called_once_with( mongo_db.send_notification.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id="MONGODB_QUERY_ERROR", notification_id='MONGODB_QUERY_ERROR',
message="Failed to execute MongoDB query: Error", message='Failed to execute MongoDB query: Error',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
block="load_query_from_mongodb", block='load_query_from_mongodb',
attachment_content=ANY attachment_content=ANY,
) )
else: else:
assert False, "Expected an exception to be raised" raise AssertionError('Expected an exception to be raised')
@mark.asyncio @mark.asyncio
async def test_find_documents_in_mongodb_missing_collection(mongo_db): async def test_find_documents_in_mongodb_missing_collection(mongo_db):
input_data = {"query": {"filters": {}}} input_data = {'query': {'filters': {}}}
try: try:
await mongo_db.find_documents_in_mongodb(input_data) await mongo_db.find_documents_in_mongodb(input_data)
except ValueError as e: except ValueError as e:
assert str(e) == "Collection name must be provided in the query." assert str(e) == 'Collection name must be provided in the query.'
else: else:
assert False, "Expected a ValueError to be raised" raise AssertionError('Expected a ValueError to be raised')
@mark.asyncio @mark.asyncio
async def test_aggregate_documents_in_mongodb_success(mongo_db): async def test_aggregate_documents_in_mongodb_success(mongo_db):
input_data = {"collection": "test_collection", "aggregation": [ input_data = {
{"$match": {"name": {"$exists": True}}}, 'collection': 'test_collection',
{"$project": {"name": 1}} 'aggregation': [{'$match': {'name': {'$exists': True}}}, {'$project': {'name': 1}}],
]} }
mock_collection = MagicMock() mock_collection = MagicMock()
mock_collection.aggregate.return_value = [ mock_collection.aggregate.return_value = [
{ {
"_id": "asdad", '_id': 'asdad',
"name": "test1", 'name': 'test1',
"timestamp": datetime.strptime( 'timestamp': datetime.strptime(
"2023-01-01 12:00:00.000000+0000", DATETIME_FORMAT_MS_WITH_TZ)}, '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
),
},
{ {
"_id": "adzx", '_id': 'adzx',
"name": "test2", 'name': 'test2',
"timestamp": datetime.strptime( 'timestamp': datetime.strptime(
"2023-01-01 12:00:00.000000+0000", DATETIME_FORMAT_MS_WITH_TZ)} '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
),
},
] ]
mongo_db.database.__getitem__.return_value = mock_collection mongo_db.database.__getitem__.return_value = mock_collection
result = await mongo_db.aggregate_documents_in_mongodb( result = await mongo_db.aggregate_documents_in_mongodb(
{ {'query': input_data, 'timestamp_fields': ['timestamp']}
"query": input_data, )
"timestamp_fields": ["timestamp"]
})
assert len(result) == 2 assert len(result) == 2
assert result[0] == {"name": "test1", assert result[0] == {'name': 'test1', 'timestamp': '2023-01-01 12:00:00.000000+0000'}
"timestamp": "2023-01-01 12:00:00.000000+0000"} assert result[1] == {'name': 'test2', 'timestamp': '2023-01-01 12:00:00.000000+0000'}
assert result[1] == {"name": "test2", expected_pipeline = input_data['aggregation']
"timestamp": "2023-01-01 12:00:00.000000+0000"} expected_pipeline.append({'$project': {'_id': 0}})
expected_pipeline = input_data["aggregation"]
expected_pipeline.append({"$project": {"_id": 0}})
mock_collection.aggregate.assert_called_once_with( mock_collection.aggregate.assert_called_once_with(expected_pipeline)
expected_pipeline
)
@mark.asyncio @mark.asyncio
async def test_aggregate_documents_in_mongodb_failure(mongo_db): async def test_aggregate_documents_in_mongodb_failure(mongo_db):
input_data = {"collection": "test_collection", "aggregation": [ input_data = {
{"$match": {"name": {"$exists": True}}}, 'collection': 'test_collection',
{"$project": {"name": 1}} 'aggregation': [{'$match': {'name': {'$exists': True}}}, {'$project': {'name': 1}}],
]} }
mongo_db.database.__getitem__.return_value = MagicMock( mongo_db.database.__getitem__.return_value = MagicMock(
aggregate=MagicMock(side_effect=Exception("Error")) aggregate=MagicMock(side_effect=Exception('Error'))
) )
try: try:
await mongo_db.aggregate_documents_in_mongodb( await mongo_db.aggregate_documents_in_mongodb({'query': input_data, **metadata})
{
"query": input_data,
**metadata
})
except Exception as e: except Exception as e:
assert str(e) == "Error" assert str(e) == 'Error'
mongo_db.send_notification.assert_called_once_with( mongo_db.send_notification.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id="MONGODB_AGGREGATION_ERROR", notification_id='MONGODB_AGGREGATION_ERROR',
message="Failed to execute MongoDB aggregation: Error", message='Failed to execute MongoDB aggregation: Error',
block="aggregate_documents_in_mongodb", block='aggregate_documents_in_mongodb',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=ANY attachment_content=ANY,
) )
else: else:
assert False, "Expected an exception to be raised" raise AssertionError('Expected an exception to be raised')
@mark.asyncio @mark.asyncio
async def test_aggregate_documents_in_mongodb_missing_collection(mongo_db): async def test_aggregate_documents_in_mongodb_missing_collection(mongo_db):
input_data = {"query": {"aggregation": []}} input_data = {'query': {'aggregation': []}}
try: try:
await mongo_db.aggregate_documents_in_mongodb(input_data) await mongo_db.aggregate_documents_in_mongodb(input_data)
except ValueError as e: except ValueError as e:
assert str(e) == "Collection name must be provided in the query." assert str(e) == 'Collection name must be provided in the query.'
else: else:
assert False, "Expected a ValueError to be raised" raise AssertionError('Expected a ValueError to be raised')
@mark.asyncio @mark.asyncio
async def test_aggregate_documents_in_mongodb_missing_aggregation(mongo_db): async def test_aggregate_documents_in_mongodb_missing_aggregation(mongo_db):
input_data = {"query": {"collection": "test_collection"}} input_data = {'query': {'collection': 'test_collection'}}
try: try:
await mongo_db.aggregate_documents_in_mongodb(input_data) await mongo_db.aggregate_documents_in_mongodb(input_data)
except ValueError as e: except ValueError as e:
assert str(e) == "Aggregation must be provided." assert str(e) == 'Aggregation must be provided.'
else: else:
assert False, "Expected a ValueError to be raised" raise AssertionError('Expected a ValueError to be raised')
@mark.asyncio @mark.asyncio
@patch("orchestrator.activities.mongo_db.now") @patch('orchestrator.activities.mongo_db.now')
async def test_update_pipelines_timestamps_success(now_mock, mongo_db): async def test_update_pipelines_timestamps_success(now_mock, mongo_db):
input_data = {"updated_pipelines": [ input_data = {
{"schedule_name": "test1", "namespace": "test1", "success": True}, 'updated_pipelines': [
{"schedule_name": "test2", "namespace": "test2", "success": True} {'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
]} {'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
mongo_db.database["pipelines"].update_many.return_value = MagicMock() ]
}
mongo_db.database['pipelines'].update_many.return_value = MagicMock()
await mongo_db.update_pipelines_timestamps(input_data) await mongo_db.update_pipelines_timestamps(input_data)
mongo_db.database["pipelines"].update_many.assert_called_once_with( mongo_db.database['pipelines'].update_many.assert_called_once_with(
{"$or": [ {
{"schedule_name": "test1", "namespace": "test1"}, '$or': [
{"schedule_name": "test2", "namespace": "test2"} {'schedule_name': 'test1', 'namespace': 'test1'},
]}, {'schedule_name': 'test2', 'namespace': 'test2'},
{"$set": { ]
"updated_at": now_mock.return_value}} },
{'$set': {'updated_at': now_mock.return_value}},
) )
@mark.asyncio @mark.asyncio
@patch("orchestrator.activities.mongo_db.now") @patch('orchestrator.activities.mongo_db.now')
async def test_update_pipelines_timestamps_failure(now_mock, mongo_db): async def test_update_pipelines_timestamps_failure(now_mock, mongo_db):
input_data = { input_data = {
"updated_pipelines": [ 'updated_pipelines': [
{"schedule_name": "test1", "namespace": "test1", "success": True}, {'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
{"schedule_name": "test2", "namespace": "test2", "success": True} {'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
], ],
**metadata **metadata,
} }
mongo_db.database["pipelines"].update_many.side_effect = Exception("Error") mongo_db.database['pipelines'].update_many.side_effect = Exception('Error')
try: try:
await mongo_db.update_pipelines_timestamps(input_data) await mongo_db.update_pipelines_timestamps(input_data)
except Exception as e: except Exception as e:
assert str(e) == "Error" assert str(e) == 'Error'
mongo_db.send_notification.assert_called_once_with( mongo_db.send_notification.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id="MONGODB_UPDATE_PIPELINES_ERROR", notification_id='MONGODB_UPDATE_PIPELINES_ERROR',
message="Failed to update pipelines timestamps: Error", message='Failed to update pipelines timestamps: Error',
block="update_pipelines_timestamps", block='update_pipelines_timestamps',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=ANY attachment_content=ANY,
) )
else: else:
assert False, "Expected an exception to be raised" raise AssertionError('Expected an exception to be raised')
@mark.asyncio @mark.asyncio
@patch("orchestrator.activities.mongo_db.now") @patch('orchestrator.activities.mongo_db.now')
async def test_create_pipelines_timestamps_success(now_mock, mongo_db): async def test_create_pipelines_timestamps_success(now_mock, mongo_db):
input_data = {"created_pipelines": [ input_data = {
{"schedule_name": "test1", "namespace": "test1", "success": True}, 'created_pipelines': [
{"schedule_name": "test2", "namespace": "test2", "success": True} {'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
]} {'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
mongo_db.database["pipelines"].insert_many.return_value = MagicMock() ]
}
mongo_db.database['pipelines'].insert_many.return_value = MagicMock()
await mongo_db.create_pipelines_timestamps(input_data) await mongo_db.create_pipelines_timestamps(input_data)
mongo_db.database["pipelines"].insert_many.assert_called_once_with( mongo_db.database['pipelines'].insert_many.assert_called_once_with(
[ [
{"schedule_name": "test1", "namespace": "test1", {'schedule_name': 'test1', 'namespace': 'test1', 'updated_at': now_mock.return_value},
"updated_at": now_mock.return_value}, {'schedule_name': 'test2', 'namespace': 'test2', 'updated_at': now_mock.return_value},
{"schedule_name": "test2", "namespace": "test2",
"updated_at": now_mock.return_value}
] ]
) )
@mark.asyncio @mark.asyncio
@patch("orchestrator.activities.mongo_db.now") @patch('orchestrator.activities.mongo_db.now')
async def test_create_pipelines_timestamps_failure(now_mock, mongo_db): async def test_create_pipelines_timestamps_failure(now_mock, mongo_db):
input_data = {"created_pipelines": [ input_data = {
{"schedule_name": "test1", "namespace": "test1", "success": True}, 'created_pipelines': [
{"schedule_name": "test2", "namespace": "test2", "success": True} {'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
{'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
], ],
**metadata **metadata,
} }
mongo_db.database["pipelines"].insert_many.side_effect = Exception("Error") mongo_db.database['pipelines'].insert_many.side_effect = Exception('Error')
try: try:
await mongo_db.create_pipelines_timestamps(input_data) await mongo_db.create_pipelines_timestamps(input_data)
except Exception as e: except Exception as e:
assert str(e) == "Error" assert str(e) == 'Error'
mongo_db.send_notification.assert_called_once_with( mongo_db.send_notification.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id="MONGODB_CREATE_PIPELINES_ERROR", notification_id='MONGODB_CREATE_PIPELINES_ERROR',
message="Failed to create pipelines timestamps: Error", message='Failed to create pipelines timestamps: Error',
block="create_pipelines_timestamps", block='create_pipelines_timestamps',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=ANY attachment_content=ANY,
) )
else: else:
assert False, "Expected an exception to be raised" raise AssertionError('Expected an exception to be raised')
@mark.asyncio @mark.asyncio
@patch("orchestrator.activities.mongo_db.now") @patch('orchestrator.activities.mongo_db.now')
async def test_delete_pipelines_timestamps_success(now_mock, mongo_db): async def test_delete_pipelines_timestamps_success(now_mock, mongo_db):
input_data = {"deleted_pipelines": [ input_data = {
{"schedule_name": "test1", "namespace": "test1", "success": True}, 'deleted_pipelines': [
{"schedule_name": "test2", "namespace": "test2", "success": True} {'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
]} {'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
mongo_db.database["pipelines"].delete_many.return_value = MagicMock() ]
}
mongo_db.database['pipelines'].delete_many.return_value = MagicMock()
await mongo_db.delete_pipelines_timestamps(input_data) await mongo_db.delete_pipelines_timestamps(input_data)
mongo_db.database["pipelines"].delete_many.assert_called_once_with( mongo_db.database['pipelines'].delete_many.assert_called_once_with(
{"$or": [ {
{"schedule_name": "test1", "namespace": "test1"}, '$or': [
{"schedule_name": "test2", "namespace": "test2"} {'schedule_name': 'test1', 'namespace': 'test1'},
]} {'schedule_name': 'test2', 'namespace': 'test2'},
]
}
) )
@mark.asyncio @mark.asyncio
@patch("orchestrator.activities.mongo_db.now") @patch('orchestrator.activities.mongo_db.now')
async def test_delete_pipelines_timestamps_failure(now_mock, mongo_db): async def test_delete_pipelines_timestamps_failure(now_mock, mongo_db):
input_data = {"deleted_pipelines": [ input_data = {
{"schedule_name": "test1", "namespace": "test1", "success": True}, 'deleted_pipelines': [
{"schedule_name": "test2", "namespace": "test2", "success": True} {'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
{'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
], ],
**metadata **metadata,
} }
mongo_db.database["pipelines"].delete_many.side_effect = Exception("Error") mongo_db.database['pipelines'].delete_many.side_effect = Exception('Error')
try: try:
await mongo_db.delete_pipelines_timestamps(input_data) await mongo_db.delete_pipelines_timestamps(input_data)
except Exception as e: except Exception as e:
assert str(e) == "Error" assert str(e) == 'Error'
mongo_db.send_notification.assert_called_once_with( mongo_db.send_notification.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id="MONGODB_DELETE_PIPELINES_ERROR", notification_id='MONGODB_DELETE_PIPELINES_ERROR',
message="Failed to delete pipelines timestamps: Error", message='Failed to delete pipelines timestamps: Error',
block="delete_pipelines_timestamps", block='delete_pipelines_timestamps',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=ANY attachment_content=ANY,
) )
else: else:
assert False, "Expected an exception to be raised" raise AssertionError('Expected an exception to be raised')
@mark.asyncio @mark.asyncio
async def test_create_collection_with_ttl_index_success(mongo_db): async def test_create_collection_with_ttl_index_success(mongo_db):
input_data = { input_data = {
**metadata, **metadata,
"pipelines": { 'pipelines': {
"scouter-pipeline": { 'scouter-pipeline': {'topic': 'raw_scouter_pipeline'},
"topic": "raw_scouter_pipeline" 'scouter-pipeline-2': {'topic': 'raw_scouter_pipeline_2'},
'scouter-pipeline-3': {'topic': 'raw_scouter_pipeline_3'},
}, },
"scouter-pipeline-2": {
"topic": "raw_scouter_pipeline_2"
},
"scouter-pipeline-3": {
"topic": "raw_scouter_pipeline_3"
}
}
} }
mongo_db.database.list_collection_names.return_value = [ mongo_db.database.list_collection_names.return_value = [
"raw_scouter_pipeline_2", 'raw_scouter_pipeline_2',
"raw_scouter_pipeline_3" 'raw_scouter_pipeline_3',
] ]
collection_1 = MagicMock( collection_1 = MagicMock(
list_indexes=MagicMock( list_indexes=MagicMock(
return_value=[ return_value=[
{ {
"key": "asdad", 'key': 'asdad',
} }
] ]
) )
) )
collection_2 = MagicMock( collection_2 = MagicMock(
list_indexes=MagicMock( list_indexes=MagicMock(return_value=[{'key': 'inserted_at', 'expireAfterSeconds': None}])
return_value=[
{
"key": "inserted_at",
"expireAfterSeconds": None
}
]
)
) )
collection_3 = MagicMock( collection_3 = MagicMock(
list_indexes=MagicMock( list_indexes=MagicMock(return_value=[{'key': 'inserted_at', 'expireAfterSeconds': 3600}])
return_value=[
{
"key": "inserted_at",
"expireAfterSeconds": 3600
}
]
)
) )
mongo_db.database.__getitem__ = MagicMock( mongo_db.database.__getitem__ = MagicMock(
side_effect=[ side_effect=[collection_1, collection_2, collection_3]
collection_1,
collection_2,
collection_3
]
) )
await mongo_db.create_collection_with_ttl_index(input_data) await mongo_db.create_collection_with_ttl_index(input_data)
mongo_db.database.list_collection_names.assert_called_once_with() mongo_db.database.list_collection_names.assert_called_once_with()
mongo_db.database.create_collection.assert_called_once_with( mongo_db.database.create_collection.assert_called_once_with('raw_scouter_pipeline')
"raw_scouter_pipeline"
)
collection_1.list_indexes.assert_called_once() collection_1.list_indexes.assert_called_once()
collection_1.create_index.assert_called_once_with( collection_1.create_index.assert_called_once_with(
"inserted_at", 'inserted_at', expireAfterSeconds=3600, background=True
expireAfterSeconds=3600,
background=True
) )
collection_2.list_indexes.assert_called_once() collection_2.list_indexes.assert_called_once()
collection_2.create_index.assert_called_once_with( collection_2.create_index.assert_called_once_with(
"inserted_at", 'inserted_at', expireAfterSeconds=3600, background=True
expireAfterSeconds=3600,
background=True
) )
collection_3.list_indexes.assert_called_once() collection_3.list_indexes.assert_called_once()
@@ -515,33 +464,26 @@ async def test_create_collection_with_ttl_index_success(mongo_db):
@mark.asyncio @mark.asyncio
async def test_create_collection_with_ttl_index_failure(mongo_db): async def test_create_collection_with_ttl_index_failure(mongo_db):
input_data = { input_data = {**metadata, 'pipelines': {'scouter-pipeline': {'topic': 'raw_scouter_pipeline'}}}
**metadata,
"pipelines": {
"scouter-pipeline": {
"topic": "raw_scouter_pipeline"
}
}
}
mongo_db.database.list_collection_names.return_value = [] mongo_db.database.list_collection_names.return_value = []
mongo_db.database.create_collection.side_effect = Exception("Error") mongo_db.database.create_collection.side_effect = Exception('Error')
try: try:
await mongo_db.create_collection_with_ttl_index(input_data) await mongo_db.create_collection_with_ttl_index(input_data)
except Exception as e: except Exception as e:
assert str(e) == "Error" assert str(e) == 'Error'
mongo_db.send_notification.assert_called_once_with( mongo_db.send_notification.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id="MONGODB_CREATE_COLLECTION_ERROR", notification_id='MONGODB_CREATE_COLLECTION_ERROR',
message="Failed to create collection raw_scouter_pipeline with TTL index: Error", message='Failed to create collection raw_scouter_pipeline with TTL index: Error',
block="create_collection_with_ttl_index", block='create_collection_with_ttl_index',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=ANY attachment_content=ANY,
) )
else: else:
assert False, "Expected an exception to be raised" raise AssertionError('Expected an exception to be raised')
@mark.asyncio @mark.asyncio
@@ -555,34 +497,25 @@ async def test_load_latest_data_none_last_data_timestamp(mongo_db):
'name': 'test1', 'name': 'test1',
'value': 1, 'value': 1,
'timestamp': datetime.strptime( 'timestamp': datetime.strptime(
'2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ) '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
),
} }
] ]
result = await mongo_db.load_latest_data({ result = await mongo_db.load_latest_data(
{
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'collection_name': 'test_collection', 'collection_name': 'test_collection',
'last_data_timestamp': None, 'last_data_timestamp': None,
'base_data_filter': { 'base_data_filter': {'level': 'ERROR'},
'level': 'ERROR'
} }
})
mongo_db.database.__getitem__.assert_called_once_with(
'test_collection')
collection.find.assert_called_once_with(
{
'level': 'ERROR'
},
{"_id": 0}
) )
assert result == [{ mongo_db.database.__getitem__.assert_called_once_with('test_collection')
'name': 'test1',
'value': 1, collection.find.assert_called_once_with({'level': 'ERROR'}, {'_id': 0})
'timestamp': '2023-01-01 12:00:00.000000+0000'
}] assert result == [{'name': 'test1', 'value': 1, 'timestamp': '2023-01-01 12:00:00.000000+0000'}]
@mark.asyncio @mark.asyncio
@@ -596,38 +529,35 @@ async def test_load_latest_data_not_none_last_data_timestamp(mongo_db):
'name': 'test1', 'name': 'test1',
'value': 1, 'value': 1,
'timestamp': datetime.strptime( 'timestamp': datetime.strptime(
'2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ) '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
),
} }
] ]
result = await mongo_db.load_latest_data({ result = await mongo_db.load_latest_data(
{
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'collection_name': 'test_collection', 'collection_name': 'test_collection',
'last_data_timestamp': '2023-01-01 12:00:00.000000+0000', 'last_data_timestamp': '2023-01-01 12:00:00.000000+0000',
'base_data_filter': { 'base_data_filter': {'level': 'ERROR'},
'level': 'ERROR'
} }
}) )
mongo_db.database.__getitem__.assert_called_once_with( mongo_db.database.__getitem__.assert_called_once_with('test_collection')
'test_collection')
collection.find.assert_called_once_with( collection.find.assert_called_once_with(
{ {
'level': 'ERROR', 'level': 'ERROR',
'timestamp': { 'timestamp': {
'$gt': datetime.strptime( '$gt': datetime.strptime(
'2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ) '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
} )
}, },
{"_id": 0} },
{'_id': 0},
) )
assert result == [{ assert result == [{'name': 'test1', 'value': 1, 'timestamp': '2023-01-01 12:00:00.000000+0000'}]
'name': 'test1',
'value': 1,
'timestamp': '2023-01-01 12:00:00.000000+0000'
}]
@mark.asyncio @mark.asyncio
@@ -640,23 +570,22 @@ async def test_load_latest_data_error(mongo_db):
collection.find.side_effect = Exception('test') collection.find.side_effect = Exception('test')
try: try:
await mongo_db.load_latest_data({ await mongo_db.load_latest_data(
{
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'collection_name': 'test_collection', 'collection_name': 'test_collection',
'last_data_timestamp': '2023-01-01 12:00:00.000000+0000', 'last_data_timestamp': '2023-01-01 12:00:00.000000+0000',
'base_data_filter': { 'base_data_filter': {'level': 'ERROR'},
'level': 'ERROR'
} }
}) )
except Exception as e: except Exception as e:
assert str(e) == 'test' assert str(e) == 'test'
mongo_db.send_notification.assert_called_once_with( mongo_db.send_notification.assert_called_once_with(
metadata={'workflow_name': 'test_pipeline', metadata={'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'schedule_name': 'test_schedule'},
notification_id='MONGO_LOAD_ERROR', notification_id='MONGO_LOAD_ERROR',
message='Error loading data from MongoDB: test', message='Error loading data from MongoDB: test',
block='load_latest_data', block='load_latest_data',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=ANY attachment_content=ANY,
) )

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

@@ -1,32 +1,33 @@
from unittest.mock import MagicMock, patch, call, ANY from datetime import timedelta
from datetime import datetime, timedelta from unittest.mock import ANY, MagicMock, call, patch
from pandas import DataFrame from pandas import DataFrame
from pytest import mark, fixture from pytest import fixture, mark
from orchestrator.activities.slot_manager import SlotManager
from sientia_do.notifications.models import NotificationLevel from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
from orchestrator.activities.slot_manager import SlotManager
metadata = { metadata = {
"metadata": { 'metadata': {
"schedule_name": "test_schedule_name", 'schedule_name': 'test_schedule_name',
"workflow_name": "test_workflow_name", 'workflow_name': 'test_workflow_name',
"model_name": "test_model_name", 'model_name': 'test_model_name',
"model_id": "test_model_id" 'model_id': 'test_model_id',
} }
} }
@fixture @fixture
@patch("orchestrator.activities.slot_manager.Redis.__init__") @patch('orchestrator.activities.slot_manager.Redis.__init__')
def slot_manager(_redis_mock): def slot_manager(_redis_mock):
slot_manager = SlotManager( slot_manager = SlotManager(
host="localhost", host='localhost',
port=6379, port=6379,
username="admin", username='admin',
password="password", password='password',
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock() notification_handler=MagicMock(),
) )
slot_manager.redis_client = MagicMock() slot_manager.redis_client = MagicMock()
@@ -46,168 +47,137 @@ async def test_load_opc_slots_no_slot_keys(slot_manager):
@mark.asyncio @mark.asyncio
async def test_load_opc_slots(slot_manager): async def test_load_opc_slots(slot_manager):
slot_manager.redis_client.keys.return_value = [ slot_manager.redis_client.keys.return_value = [
b"slot:opc_tags:1", b"slot:opc_tags:2", b"slot:opc_tags:3"] b'slot:opc_tags:1',
b'slot:opc_tags:2',
slot_manager.get = MagicMock( b'slot:opc_tags:3',
side_effect=[
"value1",
"value2",
None
] ]
)
slot_manager.get = MagicMock(side_effect=['value1', 'value2', None])
response = await slot_manager.load_opc_slots(metadata) response = await slot_manager.load_opc_slots(metadata)
assert response == { assert response == {
"slot:opc_tags:1": "value1", 'slot:opc_tags:1': 'value1',
"slot:opc_tags:2": "value2", 'slot:opc_tags:2': 'value2',
"slot:opc_tags:3": None 'slot:opc_tags:3': None,
} }
@mark.asyncio @mark.asyncio
async def test_load_opc_slots_no_decode(slot_manager): async def test_load_opc_slots_no_decode(slot_manager):
slot_manager.redis_client.keys.return_value = [ slot_manager.redis_client.keys.return_value = [
"slot:opc_tags:1", "slot:opc_tags:2", "slot:opc_tags:3"] 'slot:opc_tags:1',
'slot:opc_tags:2',
slot_manager.get = MagicMock( 'slot:opc_tags:3',
side_effect=[
"value1",
"value2",
None
] ]
)
slot_manager.get = MagicMock(side_effect=['value1', 'value2', None])
response = await slot_manager.load_opc_slots(metadata) response = await slot_manager.load_opc_slots(metadata)
assert response == { assert response == {
"slot:opc_tags:1": "value1", 'slot:opc_tags:1': 'value1',
"slot:opc_tags:2": "value2", 'slot:opc_tags:2': 'value2',
"slot:opc_tags:3": None 'slot:opc_tags:3': None,
} }
@mark.asyncio @mark.asyncio
async def test_load_opc_slots_error(slot_manager): async def test_load_opc_slots_error(slot_manager):
slot_manager.redis_client.keys.return_value = [ slot_manager.redis_client.keys.return_value = [
"slot:opc_tags:1", "slot:opc_tags:2", "slot:opc_tags:3"] 'slot:opc_tags:1',
'slot:opc_tags:2',
'slot:opc_tags:3',
]
slot_manager.get = MagicMock( slot_manager.get = MagicMock(side_effect=Exception('Test exception'))
side_effect=Exception("Test exception")
)
try: try:
await slot_manager.load_opc_slots(metadata) await slot_manager.load_opc_slots(metadata)
except Exception as e: except Exception as e:
assert str(e) == "Test exception" assert str(e) == 'Test exception'
slot_manager.send_notification.assert_called_once_with( slot_manager.send_notification.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id="REDIS_GET_ERROR", notification_id='REDIS_GET_ERROR',
message="Failed to load OPC slots: Test exception", message='Failed to load OPC slots: Test exception',
block="load_opc_slots", block='load_opc_slots',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=ANY attachment_content=ANY,
) )
else: else:
assert False, "Expected an exception to be raised" raise AssertionError('Expected an exception to be raised')
@mark.asyncio @mark.asyncio
async def test_load_active_ingestors(slot_manager): async def test_load_active_ingestors(slot_manager):
slot_manager.redis_client.keys.return_value = [ slot_manager.redis_client.keys.return_value = [
b"heartbeat:ingestor:1", b"heartbeat:ingestor:2", "heartbeat:ingestor:3"] b'heartbeat:ingestor:1',
b'heartbeat:ingestor:2',
'heartbeat:ingestor:3',
]
response = await slot_manager.load_active_ingestors(metadata) response = await slot_manager.load_active_ingestors(metadata)
assert response == ["heartbeat:ingestor:1", assert response == ['heartbeat:ingestor:1', 'heartbeat:ingestor:2', 'heartbeat:ingestor:3']
"heartbeat:ingestor:2", "heartbeat:ingestor:3"]
@mark.asyncio @mark.asyncio
async def test_load_active_ingestors_error(slot_manager): async def test_load_active_ingestors_error(slot_manager):
slot_manager.redis_client.keys.return_value = [ slot_manager.redis_client.keys.return_value = [
"heartbeat:ingestor:1", "heartbeat:ingestor:2", "heartbeat:ingestor:3"] 'heartbeat:ingestor:1',
'heartbeat:ingestor:2',
'heartbeat:ingestor:3',
]
slot_manager.redis_client.keys.side_effect = Exception("Test exception") slot_manager.redis_client.keys.side_effect = Exception('Test exception')
try: try:
await slot_manager.load_active_ingestors(metadata) await slot_manager.load_active_ingestors(metadata)
except Exception as e: except Exception as e:
assert str(e) == "Test exception" assert str(e) == 'Test exception'
slot_manager.send_notification.assert_called_once_with( slot_manager.send_notification.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id="REDIS_GET_ERROR", notification_id='REDIS_GET_ERROR',
message="Failed to load active ingestors: Test exception", message='Failed to load active ingestors: Test exception',
block="load_active_ingestors", block='load_active_ingestors',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=ANY attachment_content=ANY,
) )
else: else:
assert False, "Expected an exception to be raised" raise AssertionError('Expected an exception to be raised')
@mark.asyncio @mark.asyncio
async def test_update_slots(slot_manager): async def test_update_slots(slot_manager):
slot_manager.set = MagicMock( slot_manager.set = MagicMock(side_effect=[None, Exception('Test exception')])
side_effect=[
None, response = await slot_manager.update_slots({'to_insert': {'1': 'value1', '2': 'value2'}})
Exception("Test exception")
] slot_manager.set.assert_has_calls(
[call('slot:opc_tags:1', 'value1', ttl=None), call('slot:opc_tags:2', 'value2', ttl=None)]
) )
response = await slot_manager.update_slots({
"to_insert": {
"1": "value1",
"2": "value2"
}
})
slot_manager.set.assert_has_calls([
call("slot:opc_tags:1", "value1", ttl=None),
call("slot:opc_tags:2", "value2", ttl=None)
])
assert response == { assert response == {
"1": { '1': {'success': True, 'message': 'Slot updated successfully'},
"success": True, '2': {'success': False, 'message': 'Test exception'},
"message": "Slot updated successfully"
},
"2": {
"success": False,
"message": "Test exception"
}
} }
@mark.asyncio @mark.asyncio
async def test_delete_slots(slot_manager): async def test_delete_slots(slot_manager):
slot_manager.redis_client.delete = MagicMock( slot_manager.redis_client.delete = MagicMock(side_effect=[None, Exception('Test exception')])
side_effect=[
None, response = await slot_manager.delete_slots({'to_delete': ['1', '2']})
Exception("Test exception")
] slot_manager.redis_client.delete.assert_has_calls(
[call('slot:opc_tags:1'), call('slot:opc_tags:2')]
) )
response = await slot_manager.delete_slots({
"to_delete": ["1", "2"]
})
slot_manager.redis_client.delete.assert_has_calls([
call("slot:opc_tags:1"),
call("slot:opc_tags:2")
])
assert response == { assert response == {
"1": { '1': {'success': True, 'message': 'Slot deleted successfully'},
"success": True, '2': {'success': False, 'message': 'Test exception'},
"message": "Slot deleted successfully"
},
"2": {
"success": False,
"message": "Test exception"
}
} }
@@ -218,7 +188,7 @@ async def test_get_last_data_timestamp_none(slot_manager):
**metadata, **metadata,
'workflow_name': 'test_pipeline', 'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule', 'schedule_name': 'test_schedule',
'mail_type': 'test_mail_type' 'mail_type': 'test_mail_type',
} }
slot_manager.get = MagicMock(return_value=None) slot_manager.get = MagicMock(return_value=None)
@@ -235,16 +205,14 @@ async def test_get_last_data_timestamp_not_none(slot_manager):
**metadata, **metadata,
'workflow_name': 'test_pipeline', 'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule', 'schedule_name': 'test_schedule',
'mail_type': 'test_mail_type' 'mail_type': 'test_mail_type',
} }
slot_manager.get = MagicMock(return_value='2023-01-01 12:00:00') slot_manager.get = MagicMock(return_value='2023-01-01 12:00:00')
result = await slot_manager.get_last_data_timestamp(test_data) result = await slot_manager.get_last_data_timestamp(test_data)
slot_manager.get.assert_called_once_with( slot_manager.get.assert_called_once_with('notification_last_timestamp:test_mail_type')
'notification_last_timestamp:test_mail_type'
)
assert result == '2023-01-01 12:00:00' assert result == '2023-01-01 12:00:00'
@@ -256,14 +224,13 @@ async def test_get_last_data_timestamp_error(slot_manager):
**metadata, **metadata,
'workflow_name': 'test_pipeline', 'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule', 'schedule_name': 'test_schedule',
'mail_type': 'test_mail_type' 'mail_type': 'test_mail_type',
} }
slot_manager.send_notification = MagicMock() slot_manager.send_notification = MagicMock()
slot_manager.get = MagicMock(side_effect=Exception('test')) slot_manager.get = MagicMock(side_effect=Exception('test'))
try: try:
await slot_manager.get_last_data_timestamp(test_data) await slot_manager.get_last_data_timestamp(test_data)
except Exception as e: except Exception as e:
@@ -271,15 +238,15 @@ async def test_get_last_data_timestamp_error(slot_manager):
slot_manager.send_notification.assert_called_once_with( slot_manager.send_notification.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id="REDIS_GET_ERROR", notification_id='REDIS_GET_ERROR',
message="Error getting last data timestamp: test", message='Error getting last data timestamp: test',
block="get_last_data_timestamp", block='get_last_data_timestamp',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=ANY attachment_content=ANY,
) )
else: else:
assert False, "Expected exception" raise AssertionError('Expected exception')
@mark.asyncio @mark.asyncio
@@ -290,7 +257,7 @@ async def test_put_last_data_timestamp_empty_dataframe(slot_manager):
'workflow_name': 'test_pipeline', 'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule', 'schedule_name': 'test_schedule',
'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records'), 'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records'),
'mail_type': 'test_mail_type' 'mail_type': 'test_mail_type',
} }
slot_manager.set = MagicMock() slot_manager.set = MagicMock()
@@ -306,17 +273,19 @@ async def test_put_last_data_timestamp_empty_dataframe(slot_manager):
async def test_put_last_data_timestamp_not_empty_dataframe(slot_manager): async def test_put_last_data_timestamp_not_empty_dataframe(slot_manager):
"""Test put_last_data_timestamp with not empty dataframe""" """Test put_last_data_timestamp with not empty dataframe"""
data = DataFrame({ data = DataFrame(
{
'name': ['sensor1', 'sensor2'], 'name': ['sensor1', 'sensor2'],
'value': [25.5, 30.0], 'value': [25.5, 30.0],
'timestamp': ['2023-01-01 12:00:00', '2023-01-01 12:00:01'] 'timestamp': ['2023-01-01 12:00:00', '2023-01-01 12:00:01'],
}) }
)
test_data = { test_data = {
**metadata, **metadata,
'workflow_name': 'test_pipeline', 'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule', 'schedule_name': 'test_schedule',
'data': data.to_dict('records'), 'data': data.to_dict('records'),
'mail_type': 'test_mail_type' 'mail_type': 'test_mail_type',
} }
slot_manager.set = MagicMock() slot_manager.set = MagicMock()
@@ -326,9 +295,7 @@ async def test_put_last_data_timestamp_not_empty_dataframe(slot_manager):
assert result == '2023-01-01 12:00:01' assert result == '2023-01-01 12:00:01'
slot_manager.set.assert_called_once_with( slot_manager.set.assert_called_once_with(
'notification_last_timestamp:test_mail_type', 'notification_last_timestamp:test_mail_type', '2023-01-01 12:00:01', ttl=18000
'2023-01-01 12:00:01',
ttl=18000
) )
@@ -339,12 +306,14 @@ async def test_put_last_data_timestamp_error(slot_manager):
**metadata, **metadata,
'workflow_name': 'test_pipeline', 'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule', 'schedule_name': 'test_schedule',
'data': DataFrame({ 'data': DataFrame(
{
'name': ['sensor1', 'sensor2'], 'name': ['sensor1', 'sensor2'],
'value': [25.5, 30.0], 'value': [25.5, 30.0],
'timestamp': ['2023-01-01 12:00:00'] * 2 'timestamp': ['2023-01-01 12:00:00'] * 2,
}).to_dict('records'), }
'mail_type': 'test_mail_type' ).to_dict('records'),
'mail_type': 'test_mail_type',
} }
slot_manager.send_notification = MagicMock() slot_manager.send_notification = MagicMock()
@@ -358,57 +327,43 @@ async def test_put_last_data_timestamp_error(slot_manager):
slot_manager.send_notification.assert_called_once_with( slot_manager.send_notification.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id="REDIS_SET_ERROR", notification_id='REDIS_SET_ERROR',
message="Error setting last data timestamp: test", message='Error setting last data timestamp: test',
block="put_last_data_timestamp", block='put_last_data_timestamp',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=ANY attachment_content=ANY,
) )
else: else:
assert False, "Expected exception" raise AssertionError('Expected exception')
@mark.asyncio @mark.asyncio
async def test_filter_notification_alerts(slot_manager): async def test_filter_notification_alerts(slot_manager):
slot_manager.get = MagicMock(side_effect=[ slot_manager.get = MagicMock(
side_effect=[
None, None,
(now() - timedelta(seconds=600) (now() - timedelta(seconds=600)).strftime(DATETIME_FORMAT_MS_WITH_TZ),
).strftime(DATETIME_FORMAT_MS_WITH_TZ),
now().strftime(DATETIME_FORMAT_MS_WITH_TZ), now().strftime(DATETIME_FORMAT_MS_WITH_TZ),
None, None,
(now() - timedelta(seconds=600) (now() - timedelta(seconds=600)).strftime(DATETIME_FORMAT_MS_WITH_TZ),
).strftime(DATETIME_FORMAT_MS_WITH_TZ), now().strftime(DATETIME_FORMAT_MS_WITH_TZ),
now().strftime(DATETIME_FORMAT_MS_WITH_TZ)]) ]
)
input_data = { input_data = {
**metadata, **metadata,
'notification_package': [ 'notification_package': [
{ {'trigger': 'test_trigger_1', 'notification_id': 'test_notification_id_1'},
'trigger': 'test_trigger_1', {'trigger': 'test_trigger_2', 'notification_id': 'test_notification_id_2'},
'notification_id': 'test_notification_id_1' {'trigger': 'test_trigger_3', 'notification_id': 'test_notification_id_3'},
},
{
'trigger': 'test_trigger_2',
'notification_id': 'test_notification_id_2'
},
{
'trigger': 'test_trigger_3',
'notification_id': 'test_notification_id_3'
}
], ],
'sending_configs': [ 'sending_configs': [
{ {'group_name': 'test_group_1', 'contents': ['core_alerts', 'persistent_alerts']},
'group_name': 'test_group_1', {'group_name': 'test_group_2', 'contents': ['core_alerts']},
'contents': ['core_alerts', 'persistent_alerts']
},
{
'group_name': 'test_group_2',
'contents': ['core_alerts']
}
], ],
'notification_ttl': 300, 'notification_ttl': 300,
'mail_type': 'test_mail_type' 'mail_type': 'test_mail_type',
} }
response = await slot_manager.filter_notification_alerts(input_data) response = await slot_manager.filter_notification_alerts(input_data)
@@ -418,26 +373,17 @@ async def test_filter_notification_alerts(slot_manager):
'group_name': 'test_group_1', 'group_name': 'test_group_1',
'contents': ['core_alerts', 'persistent_alerts'], 'contents': ['core_alerts', 'persistent_alerts'],
'notifications': [ 'notifications': [
{ {'trigger': 'test_trigger_1', 'notification_id': 'test_notification_id_1'},
'trigger': 'test_trigger_1', {'trigger': 'test_trigger_2', 'notification_id': 'test_notification_id_2'},
'notification_id': 'test_notification_id_1' ],
},
{
'trigger': 'test_trigger_2',
'notification_id': 'test_notification_id_2'
}
]
}, },
'test_group_2': { 'test_group_2': {
'group_name': 'test_group_2', 'group_name': 'test_group_2',
'contents': ['core_alerts'], 'contents': ['core_alerts'],
'notifications': [ 'notifications': [
{ {'trigger': 'test_trigger_1', 'notification_id': 'test_notification_id_1'}
'trigger': 'test_trigger_1', ],
'notification_id': 'test_notification_id_1' },
}
]
}
} }
@@ -446,20 +392,18 @@ async def test_store_notification_cache(slot_manager):
"""Test store_notification_cache""" """Test store_notification_cache"""
test_data = { test_data = {
**metadata, **metadata,
'log_report': DataFrame({ 'log_report': DataFrame(
{
'status': ['sent', 'error'], 'status': ['sent', 'error'],
'schedule': ['test_schedule_1', 'test_schedule_2'], 'schedule': ['test_schedule_1', 'test_schedule_2'],
'notification_id': ['test_notification_id_1', 'test_notification_id_2'] 'notification_id': ['test_notification_id_1', 'test_notification_id_2'],
}).to_dict(), }
'sent_ttl': 600 ).to_dict(),
'sent_ttl': 600,
} }
slot_manager.set = MagicMock() slot_manager.set = MagicMock()
await slot_manager.store_notification_cache(test_data) await slot_manager.store_notification_cache(test_data)
slot_manager.set.assert_called_once_with( slot_manager.set.assert_called_once_with('test_schedule_1:test_notification_id_1', ANY, ttl=600)
"test_schedule_1:test_notification_id_1",
ANY,
ttl=600
)

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

@@ -1,29 +1,31 @@
from unittest.mock import MagicMock, patch, AsyncMock, call, ANY
from datetime import timedelta from datetime import timedelta
from sientia_do.notifications.models import NotificationLevel from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
from pytest import fixture, mark from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from orchestrator.activities.temporal_manager import TemporalManager from orchestrator.activities.temporal_manager import TemporalManager
from orchestrator.utils.converters import parse_frequency from orchestrator.utils.converters import parse_frequency
metadata = { metadata = {
"metadata": { 'metadata': {
"schedule_name": "test_schedule_name", 'schedule_name': 'test_schedule_name',
"workflow_name": "test_workflow_name", 'workflow_name': 'test_workflow_name',
"model_name": "test_model_name", 'model_name': 'test_model_name',
"model_id": "test_model_id" 'model_id': 'test_model_id',
} }
} }
@fixture @fixture
@patch("orchestrator.activities.temporal_manager.Client.connect") @patch('orchestrator.activities.temporal_manager.Client.connect')
def temporal_manager(connect_mock): def temporal_manager(connect_mock):
temporal_manager = TemporalManager( temporal_manager = TemporalManager(
host='localhost:7233', host='localhost:7233',
scouter_namespace='scouter', scouter_namespace='scouter',
laborious_namespace='laborious', laborious_namespace='laborious',
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock() notification_handler=MagicMock(),
) )
temporal_manager.temporal_clients['scouter'] = MagicMock() temporal_manager.temporal_clients['scouter'] = MagicMock()
@@ -34,50 +36,31 @@ def temporal_manager(connect_mock):
@mark.asyncio @mark.asyncio
@patch("orchestrator.activities.temporal_manager.Client.connect", new_callable=AsyncMock) @patch('orchestrator.activities.temporal_manager.Client.connect', new_callable=AsyncMock)
async def test_connect_to_temporal(connect_mock, temporal_manager): async def test_connect_to_temporal(connect_mock, temporal_manager):
await temporal_manager.connect_to_temporal() await temporal_manager.connect_to_temporal()
connect_mock.assert_has_calls([ connect_mock.assert_has_calls(
call( [
target_host='localhost:7233', call(target_host='localhost:7233', namespace='scouter'),
namespace='scouter' call(target_host='localhost:7233', namespace='laborious'),
), ]
call(
target_host='localhost:7233',
namespace='laborious'
) )
])
async def async_iter(): async def async_iter():
yield MagicMock( yield MagicMock(id='test-schedule-id', search_attributes={'orchestrated': ['true']})
id="test-schedule-id", yield MagicMock(id='test-schedule-id-2', search_attributes={'Attr': ['false']})
search_attributes={ yield MagicMock(id='test-schedule-id-3', search_attributes={'Attr': ['false']})
"orchestrated": ["true"]
}
)
yield MagicMock(
id="test-schedule-id-2",
search_attributes={
"Attr": ["false"]
}
)
yield MagicMock(
id="test-schedule-id-3",
search_attributes={
"Attr": ["false"]
}
)
@mark.asyncio @mark.asyncio
async def test_normalize_schedules(temporal_manager): async def test_normalize_schedules(temporal_manager):
input_data = { input_data = {
"orchestrated_schedules": { 'orchestrated_schedules': {
"scouter": {"test-scouter": "2021-01-01"}, 'scouter': {'test-scouter': '2021-01-01'},
"laborious": {"test-schedule-id1": "2021-01-01"} 'laborious': {'test-schedule-id1': '2021-01-01'},
}, },
**metadata **metadata,
} }
temporal_manager.temporal_clients['scouter'].list_schedules = AsyncMock( temporal_manager.temporal_clients['scouter'].list_schedules = AsyncMock(
@@ -88,58 +71,56 @@ async def test_normalize_schedules(temporal_manager):
) )
temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock( temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock(
return_value=MagicMock( return_value=MagicMock(delete=AsyncMock())
delete=AsyncMock()
)
) )
temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock( temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock(
return_value=MagicMock( return_value=MagicMock(delete=AsyncMock())
delete=AsyncMock()
)
) )
await temporal_manager.normalize_schedules(input_data) await temporal_manager.normalize_schedules(input_data)
temporal_manager.temporal_clients['scouter'].get_schedule_handle.assert_has_calls([ temporal_manager.temporal_clients['scouter'].get_schedule_handle.assert_has_calls(
call("test-schedule-id"), [
]) call('test-schedule-id'),
]
temporal_manager.temporal_clients['scouter'].get_schedule_handle.return_value.delete.assert_awaited_once(
) )
temporal_manager.temporal_clients[
'scouter'
].get_schedule_handle.return_value.delete.assert_awaited_once()
@mark.asyncio @mark.asyncio
async def test_normalize_schedules_error(temporal_manager): async def test_normalize_schedules_error(temporal_manager):
temporal_manager.temporal_clients['scouter'].list_schedules = AsyncMock( temporal_manager.temporal_clients['scouter'].list_schedules = AsyncMock(
side_effect=Exception("Test exception") side_effect=Exception('Test exception')
) )
try: try:
await temporal_manager.normalize_schedules(metadata) await temporal_manager.normalize_schedules(metadata)
except Exception as e: except Exception as e:
assert str(e) == "Test exception" assert str(e) == 'Test exception'
temporal_manager.send_notification.assert_called_once_with( temporal_manager.send_notification.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id="TEMPORAL_NORMALIZE_SCHEDULES_ERROR", notification_id='TEMPORAL_NORMALIZE_SCHEDULES_ERROR',
message="Failed to normalize schedules: Test exception", message='Failed to normalize schedules: Test exception',
block="normalize_schedules", block='normalize_schedules',
level=NotificationLevel.ERROR, level=NotificationLevel.ERROR,
attachment_content=ANY attachment_content=ANY,
) )
else: else:
assert False, "Expected an exception to be raised" raise AssertionError('Expected an exception to be raised')
@mark.asyncio @mark.asyncio
@patch("orchestrator.activities.temporal_manager.parse_frequency", @patch('orchestrator.activities.temporal_manager.parse_frequency', side_effect=parse_frequency)
side_effect=parse_frequency) @patch('orchestrator.activities.temporal_manager.Schedule')
@patch("orchestrator.activities.temporal_manager.Schedule") @patch('orchestrator.activities.temporal_manager.ScheduleActionStartWorkflow')
@patch("orchestrator.activities.temporal_manager.ScheduleActionStartWorkflow") @patch('orchestrator.activities.temporal_manager.ScheduleIntervalSpec')
@patch("orchestrator.activities.temporal_manager.ScheduleIntervalSpec") @patch('orchestrator.activities.temporal_manager.ScheduleSpec')
@patch("orchestrator.activities.temporal_manager.ScheduleSpec") @patch('orchestrator.activities.temporal_manager.TypedSearchAttributes')
@patch("orchestrator.activities.temporal_manager.TypedSearchAttributes") @patch('orchestrator.activities.temporal_manager.SearchAttributePair')
@patch("orchestrator.activities.temporal_manager.SearchAttributePair")
async def test_create_schedule( async def test_create_schedule(
mock_search_attribute_pair, mock_search_attribute_pair,
mock_typed_search_attributes, mock_typed_search_attributes,
@@ -148,185 +129,177 @@ async def test_create_schedule(
mock_schedule_action_start_workflow, mock_schedule_action_start_workflow,
mock_schedule, mock_schedule,
mock_parse_frequency, mock_parse_frequency,
temporal_manager): temporal_manager,
):
input_data = { input_data = {
"schedules": { 'schedules': {
"scouter": { 'scouter': {
"test-schedule": { 'test-schedule': {
"model_id": 1, 'model_id': 1,
"model_name": "test-model-name", 'model_name': 'test-model-name',
"workflow_type": "test-workflow", 'workflow_type': 'test-workflow',
"frequency": "1m", 'frequency': '1m',
"data": {"test": "test"} 'data': {'test': 'test'},
'execution_timeout_seconds': 100,
'task_timeout_seconds': 100,
}, },
"test-schedule-invalid-frequency": { 'test-schedule-invalid-frequency': {
"model_id": 2, 'model_id': 2,
"model_name": "test-model-name", 'model_name': 'test-model-name',
"workflow_type": "test-workflow", 'workflow_type': 'test-workflow',
"frequency": "10y", 'frequency': '10y',
"data": {"test": "test"} 'data': {'test': 'test'},
'execution_timeout_seconds': 400,
'task_timeout_seconds': 400,
},
},
'laborious': {
'test-schedule-laborious': {
'model_id': 1,
'model_name': 'test-model-name',
'workflow_type': 'test-workflow',
'frequency': '2m',
'data': {'test': 'test'},
'execution_timeout_seconds': 500,
'task_timeout_seconds': 500,
} }
}, },
"laborious": {
"test-schedule-laborious": {
"model_id": 1,
"model_name": "test-model-name",
"workflow_type": "test-workflow",
"frequency": "2m",
"data": {"test": "test"}
}
}
} }
} }
temporal_manager.temporal_clients['scouter'].create_schedule = AsyncMock() temporal_manager.temporal_clients['scouter'].create_schedule = AsyncMock()
temporal_manager.temporal_clients['laborious'].create_schedule = AsyncMock( temporal_manager.temporal_clients['laborious'].create_schedule = AsyncMock()
)
report = await temporal_manager.create_schedules(input_data) report = await temporal_manager.create_schedules(input_data)
temporal_manager.temporal_clients['scouter'].create_schedule.assert_called_once_with( temporal_manager.temporal_clients['scouter'].create_schedule.assert_called_once_with(
"test-schedule", 'test-schedule',
mock_schedule.return_value, mock_schedule.return_value,
search_attributes=mock_typed_search_attributes.return_value search_attributes=mock_typed_search_attributes.return_value,
) )
temporal_manager.temporal_clients['laborious'].create_schedule.assert_called_once_with( temporal_manager.temporal_clients['laborious'].create_schedule.assert_called_once_with(
"test-schedule-laborious", 'test-schedule-laborious',
mock_schedule.return_value, mock_schedule.return_value,
search_attributes=mock_typed_search_attributes.return_value search_attributes=mock_typed_search_attributes.return_value,
) )
mock_schedule.assert_has_calls([ mock_schedule.assert_has_calls(
[
call( call(
action=mock_schedule_action_start_workflow.return_value, action=mock_schedule_action_start_workflow.return_value,
spec=mock_schedule_spec.return_value spec=mock_schedule_spec.return_value,
), ),
call( call(
action=mock_schedule_action_start_workflow.return_value, action=mock_schedule_action_start_workflow.return_value,
spec=mock_schedule_spec.return_value spec=mock_schedule_spec.return_value,
),
]
) )
])
mock_schedule_action_start_workflow.assert_has_calls([ mock_schedule_action_start_workflow.assert_has_calls(
[
call( call(
"test-workflow", 'test-workflow',
input_data['schedules']['scouter']['test-schedule'], input_data['schedules']['scouter']['test-schedule'],
id="test-schedule", id='test-schedule',
task_queue="test-workflow-queue", task_queue='test-workflow-queue',
execution_timeout=ANY, execution_timeout=timedelta(seconds=100),
run_timeout=timedelta(seconds=100),
task_timeout=timedelta(seconds=100),
typed_search_attributes=mock_typed_search_attributes.return_value, typed_search_attributes=mock_typed_search_attributes.return_value,
), ),
call( call(
"test-workflow", 'test-workflow',
input_data['schedules']['scouter']['test-schedule-invalid-frequency'], input_data['schedules']['scouter']['test-schedule-invalid-frequency'],
id="test-schedule-invalid-frequency", id='test-schedule-invalid-frequency',
task_queue="test-workflow-queue", task_queue='test-workflow-queue',
execution_timeout=ANY, execution_timeout=timedelta(seconds=400),
run_timeout=timedelta(seconds=400),
task_timeout=timedelta(seconds=400),
typed_search_attributes=mock_typed_search_attributes.return_value, typed_search_attributes=mock_typed_search_attributes.return_value,
), ),
call( call(
"test-workflow", 'test-workflow',
input_data['schedules']['laborious']['test-schedule-laborious'], input_data['schedules']['laborious']['test-schedule-laborious'],
id="test-schedule-laborious", id='test-schedule-laborious',
task_queue="test-workflow-queue", task_queue='test-workflow-queue',
execution_timeout=ANY, execution_timeout=timedelta(seconds=500),
run_timeout=timedelta(seconds=500),
task_timeout=timedelta(seconds=500),
typed_search_attributes=mock_typed_search_attributes.return_value, typed_search_attributes=mock_typed_search_attributes.return_value,
),
]
) )
])
mock_schedule_spec.assert_has_calls([ mock_schedule_spec.assert_has_calls(
[
call(intervals=[mock_schedule_interval_spec.return_value]),
call(intervals=[mock_schedule_interval_spec.return_value]),
]
)
mock_schedule_interval_spec.assert_has_calls(
[call(every=timedelta(seconds=60)), call(every=timedelta(seconds=120))]
)
mock_parse_frequency.assert_has_calls([call('1m'), call('10y'), call('2m')])
mock_typed_search_attributes.assert_has_calls(
[
call( call(
intervals=[ [
mock_schedule_interval_spec.return_value mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value,
] ]
), ),
call( call(
intervals=[ [
mock_schedule_interval_spec.return_value mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value,
]
),
call(
[
mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value,
]
),
] ]
) )
])
mock_schedule_interval_spec.assert_has_calls([ mock_search_attribute_pair.assert_has_calls(
call( [
every=timedelta(seconds=60) call(key=temporal_manager.model_id_id_key, value=1),
), call(key=temporal_manager.model_name_id_key, value='test-model-name'),
call( call(key=temporal_manager.orchestrated_id_key, value='true'),
every=timedelta(seconds=120) call(key=temporal_manager.model_id_id_key, value=2),
call(key=temporal_manager.model_name_id_key, value='test-model-name'),
call(key=temporal_manager.orchestrated_id_key, value='true'),
]
) )
])
mock_parse_frequency.assert_has_calls([
call("1m"),
call("10y"),
call("2m")
])
mock_typed_search_attributes.assert_has_calls([
call([
mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value
]),
call([
mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value
]),
call([
mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value
])
])
mock_search_attribute_pair.assert_has_calls([
call(
key=temporal_manager.model_id_id_key,
value=1
),
call(
key=temporal_manager.model_name_id_key,
value="test-model-name"
),
call(
key=temporal_manager.orchestrated_id_key,
value="true"
),
call(
key=temporal_manager.model_id_id_key,
value=2
),
call(
key=temporal_manager.model_name_id_key,
value="test-model-name"
),
call(
key=temporal_manager.orchestrated_id_key,
value="true"
)
])
assert report == [ assert report == [
{ {
"schedule_name": "test-schedule", 'schedule_name': 'test-schedule',
"namespace": "scouter", 'namespace': 'scouter',
"success": True, 'success': True,
"message": "Schedule created successfully" 'message': 'Schedule created successfully',
}, },
{ {
"schedule_name": "test-schedule-invalid-frequency", 'schedule_name': 'test-schedule-invalid-frequency',
"namespace": "scouter", 'namespace': 'scouter',
"success": False, 'success': False,
"message": "Invalid frequency" 'message': 'Invalid frequency',
}, },
{ {
"schedule_name": "test-schedule-laborious", 'schedule_name': 'test-schedule-laborious',
"namespace": "laborious", 'namespace': 'laborious',
"success": True, 'success': True,
"message": "Schedule created successfully" 'message': 'Schedule created successfully',
} },
] ]
@@ -334,136 +307,94 @@ async def test_create_schedule(
async def test_create_schedules_with_no_client(temporal_manager): async def test_create_schedules_with_no_client(temporal_manager):
temporal_manager.temporal_clients = {} temporal_manager.temporal_clients = {}
input_data = { input_data = {
"schedules": { 'schedules': {'abc': {'test-schedule': {'frequency': '1m', 'data': {'test': 'test'}}}}
"abc": {
"test-schedule": {
"frequency": "1m",
"data": {"test": "test"}
}
}
}
} }
try: try:
await temporal_manager.create_schedules(input_data) await temporal_manager.create_schedules(input_data)
except Exception as e: except Exception as e:
assert str( assert (
e) == f"Temporal client for abc not found, clients: {temporal_manager.temporal_clients}" str(e)
== f'Temporal client for abc not found, clients: {temporal_manager.temporal_clients}'
)
@mark.asyncio @mark.asyncio
@patch("orchestrator.activities.temporal_manager.parse_frequency", @patch('orchestrator.activities.temporal_manager.parse_frequency', side_effect=parse_frequency)
side_effect=parse_frequency) @patch('orchestrator.activities.temporal_manager.ScheduleIntervalSpec')
@patch("orchestrator.activities.temporal_manager.ScheduleIntervalSpec")
async def test_update_schedules( async def test_update_schedules(
_mock_schedule_interval_spec, _mock_schedule_interval_spec, _mock_parse_frequency, temporal_manager
_mock_parse_frequency, ):
temporal_manager): input_mock = MagicMock(args=MagicMock())
input_mock = MagicMock(
args=MagicMock()
)
temporal_manager.schedule_handles = { temporal_manager.schedule_handles = {
"scouter": { 'scouter': {
"test-schedule": MagicMock( 'test-schedule': MagicMock(
update=AsyncMock( update=AsyncMock(update=AsyncMock(side_effect=lambda f: f(input_mock)))
update=AsyncMock(
side_effect=lambda f: f(input_mock)
)
)
) )
}, },
"laborious": { 'laborious': {
"test-schedule-laborious": MagicMock( 'test-schedule-laborious': MagicMock(
update=AsyncMock( update=AsyncMock(update=AsyncMock(side_effect=lambda f: f(input_mock)))
update=AsyncMock(
side_effect=lambda f: f(input_mock)
) )
) },
)
}
} }
input_data = { input_data = {
"schedules": { 'schedules': {
"scouter": { 'scouter': {
"test-schedule": { 'test-schedule': {'frequency': '1m', 'data': {'test': 'test'}},
"frequency": "1m", 'test-schedule_no_handler': {'frequency': '1m', 'data': {'test': 'test'}},
"data": {"test": "test"}
}, },
"test-schedule_no_handler": { 'laborious': {'test-schedule-laborious': {'frequency': '2m', 'data': {'test': 'test'}}},
"frequency": "1m",
"data": {"test": "test"}
}
},
"laborious": {
"test-schedule-laborious": {
"frequency": "2m",
"data": {"test": "test"}
}
}
} }
} }
handler_scouter = MagicMock( handler_scouter = MagicMock(
update=AsyncMock( update=AsyncMock(update=AsyncMock(side_effect=lambda f: f(input_mock)))
update=AsyncMock(
side_effect=lambda f: f(input_mock)
)
)
) )
handler_laborious = MagicMock( handler_laborious = MagicMock(
update=AsyncMock( update=AsyncMock(update=AsyncMock(side_effect=lambda f: f(input_mock)))
update=AsyncMock(
side_effect=lambda f: f(input_mock)
)
)
) )
temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock( temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock(
side_effect=[ side_effect=[handler_scouter, None]
handler_scouter,
None
]
) )
temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock( temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock(
side_effect=[ side_effect=[handler_laborious]
handler_laborious
]
) )
report = await temporal_manager.update_schedules(input_data) report = await temporal_manager.update_schedules(input_data)
temporal_manager.temporal_clients['scouter'].get_schedule_handle.assert_has_calls([ temporal_manager.temporal_clients['scouter'].get_schedule_handle.assert_has_calls(
call("test-schedule"), [call('test-schedule'), call('test-schedule_no_handler')]
call("test-schedule_no_handler") )
]) temporal_manager.temporal_clients['laborious'].get_schedule_handle.assert_has_calls(
temporal_manager.temporal_clients['laborious'].get_schedule_handle.assert_has_calls([ [call('test-schedule-laborious')]
call("test-schedule-laborious") )
])
handler_scouter.update.assert_called_once() handler_scouter.update.assert_called_once()
handler_laborious.update.assert_called_once() handler_laborious.update.assert_called_once()
assert report == [ assert report == [
{ {
"schedule_name": "test-schedule", 'schedule_name': 'test-schedule',
"namespace": "scouter", 'namespace': 'scouter',
"success": True, 'success': True,
"message": "Schedule updated successfully" 'message': 'Schedule updated successfully',
}, },
{ {
"schedule_name": "test-schedule_no_handler", 'schedule_name': 'test-schedule_no_handler',
"namespace": "scouter", 'namespace': 'scouter',
"success": False, 'success': False,
"message": "Schedule test-schedule_no_handler not found" 'message': 'Schedule test-schedule_no_handler not found',
}, },
{ {
"schedule_name": "test-schedule-laborious", 'schedule_name': 'test-schedule-laborious',
"namespace": "laborious", 'namespace': 'laborious',
"success": True, 'success': True,
"message": "Schedule updated successfully" 'message': 'Schedule updated successfully',
} },
] ]
@@ -471,67 +402,41 @@ async def test_update_schedules(
async def test_update_schedules_with_no_client(temporal_manager): async def test_update_schedules_with_no_client(temporal_manager):
temporal_manager.temporal_clients = {} temporal_manager.temporal_clients = {}
input_data = { input_data = {
"schedules": { 'schedules': {'abc': {'test-schedule': {'frequency': '1m', 'data': {'test': 'test'}}}}
"abc": {
"test-schedule": {
"frequency": "1m",
"data": {"test": "test"}
}
}
}
} }
try: try:
await temporal_manager.update_schedules(input_data) await temporal_manager.update_schedules(input_data)
except Exception as e: except Exception as e:
assert str( assert (
e) == f"Temporal client for abc not found, clients: {temporal_manager.temporal_clients}" str(e)
== f'Temporal client for abc not found, clients: {temporal_manager.temporal_clients}'
)
@mark.asyncio @mark.asyncio
async def test_delete_schedules(temporal_manager): async def test_delete_schedules(temporal_manager):
temporal_manager.schedule_handles = { temporal_manager.schedule_handles = {
"scouter": { 'scouter': {'test-schedule': MagicMock(delete=AsyncMock())},
"test-schedule": MagicMock( 'laborious': {'test-schedule-laborious': MagicMock(delete=AsyncMock())},
delete=AsyncMock()
)
},
"laborious": {
"test-schedule-laborious": MagicMock(
delete=AsyncMock()
)
}
} }
input_data = { input_data = {
"schedules": { 'schedules': {
"scouter": [ 'scouter': ['test-schedule', 'test-schedule_no_handler'],
"test-schedule", "test-schedule_no_handler" 'laborious': ['test-schedule-laborious'],
],
"laborious": [
"test-schedule-laborious"
]
} }
} }
handler_scouter = MagicMock( handler_scouter = MagicMock(delete=AsyncMock())
delete=AsyncMock()
)
handler_laborious = MagicMock( handler_laborious = MagicMock(delete=AsyncMock())
delete=AsyncMock()
)
temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock( temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock(
side_effect=[ side_effect=[handler_scouter, None]
handler_scouter,
None
]
) )
temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock( temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock(
side_effect=[ side_effect=[handler_laborious]
handler_laborious
]
) )
report = await temporal_manager.delete_schedules(input_data) report = await temporal_manager.delete_schedules(input_data)
@@ -541,40 +446,36 @@ async def test_delete_schedules(temporal_manager):
assert report == [ assert report == [
{ {
"schedule_name": "test-schedule", 'schedule_name': 'test-schedule',
"namespace": "scouter", 'namespace': 'scouter',
"success": True, 'success': True,
"message": "Schedule deleted successfully" 'message': 'Schedule deleted successfully',
}, },
{ {
"schedule_name": "test-schedule_no_handler", 'schedule_name': 'test-schedule_no_handler',
"namespace": "scouter", 'namespace': 'scouter',
"success": False, 'success': False,
"message": "Schedule test-schedule_no_handler not found", 'message': 'Schedule test-schedule_no_handler not found',
"attachment": ANY 'attachment': ANY,
}, },
{ {
"schedule_name": "test-schedule-laborious", 'schedule_name': 'test-schedule-laborious',
"namespace": "laborious", 'namespace': 'laborious',
"success": True, 'success': True,
"message": "Schedule deleted successfully" 'message': 'Schedule deleted successfully',
} },
] ]
@mark.asyncio @mark.asyncio
async def test_delete_schedules_with_no_client(temporal_manager): async def test_delete_schedules_with_no_client(temporal_manager):
temporal_manager.temporal_clients = {} temporal_manager.temporal_clients = {}
input_data = { input_data = {'schedules': {'abc': ['test-schedule']}}
"schedules": {
"abc": [
"test-schedule"
]
}
}
try: try:
await temporal_manager.delete_schedules(input_data) await temporal_manager.delete_schedules(input_data)
except Exception as e: except Exception as e:
assert str( assert (
e) == f"Temporal client for abc not found, clients: {temporal_manager.temporal_clients}" str(e)
== f'Temporal client for abc not found, clients: {temporal_manager.temporal_clients}'
)

35
tests/orchestrator/utils/test_connectors_config.py
View File

@@ -1,10 +1,13 @@
from os import environ from os import environ
from orchestrator.utils.connectors_config import (build_redis_config,
from orchestrator.utils.connectors_config import (
build_couchbase_config, build_couchbase_config,
build_mongodb_config,
build_temporal_config,
build_email_config, build_email_config,
build_postgres_config) build_mongodb_config,
build_postgres_config,
build_redis_config,
build_temporal_config,
)
def test_build_redis_config_with_env_vars(): def test_build_redis_config_with_env_vars():
@@ -16,7 +19,7 @@ def test_build_redis_config_with_env_vars():
'host': 'localhost', 'host': 'localhost',
'port': 6379, 'port': 6379,
'username': 'sientia', 'username': 'sientia',
'password': 'sientia' 'password': 'sientia',
} }
@@ -27,7 +30,7 @@ def test_build_couchbase_config_with_env_vars():
assert build_couchbase_config() == { assert build_couchbase_config() == {
'connection_string': 'couchbase://localhost', 'connection_string': 'couchbase://localhost',
'username': 'sientia', 'username': 'sientia',
'password': 'sientia' 'password': 'sientia',
} }
@@ -40,7 +43,7 @@ def test_build_redis_config_with_defaults():
'host': 'localhost', 'host': 'localhost',
'port': 6379, 'port': 6379,
'username': 'default', 'username': 'default',
'password': 'bdnZOpcyiL' 'password': 'bdnZOpcyiL',
} }
@@ -51,7 +54,7 @@ def test_build_couchbase_config_with_defaults():
assert build_couchbase_config() == { assert build_couchbase_config() == {
'connection_string': 'couchbase://localhost', 'connection_string': 'couchbase://localhost',
'username': 'sientia', 'username': 'sientia',
'password': 'sientia' 'password': 'sientia',
} }
@@ -64,7 +67,7 @@ def test_build_mongo_db_config_with_env_vars():
assert build_mongodb_config() == { assert build_mongodb_config() == {
'connection_string': 'mongodb://sientia1:sientia1@localhost:27018', 'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
'database_name': 'test_db', 'database_name': 'test_db',
'ttl_index_seconds': 7200 'ttl_index_seconds': 7200,
} }
@@ -77,7 +80,7 @@ def test_build_mongo_db_config_with_defaults():
assert build_mongodb_config() == { assert build_mongodb_config() == {
'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018', 'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
'database_name': 'sientia', 'database_name': 'sientia',
'ttl_index_seconds': 3600 'ttl_index_seconds': 3600,
} }
@@ -89,7 +92,7 @@ def test_build_temporal_config_with_env_vars():
'temporal_host': 'localhost:7233', 'temporal_host': 'localhost:7233',
'temporal_namespace': 'default', 'temporal_namespace': 'default',
'temporal_scouter_namespace': 'scouter', 'temporal_scouter_namespace': 'scouter',
'temporal_laborious_namespace': 'laborious' 'temporal_laborious_namespace': 'laborious',
} }
@@ -101,7 +104,7 @@ def test_build_temporal_config_with_defaults():
'temporal_host': 'localhost:7233', 'temporal_host': 'localhost:7233',
'temporal_namespace': 'default', 'temporal_namespace': 'default',
'temporal_scouter_namespace': 'scouter', 'temporal_scouter_namespace': 'scouter',
'temporal_laborious_namespace': 'laborious' 'temporal_laborious_namespace': 'laborious',
} }
@@ -114,7 +117,7 @@ def test_build_email_config_with_env_vars():
'sender_email': 'test@test.com', 'sender_email': 'test@test.com',
'sender_password': 'test', 'sender_password': 'test',
'smtp_server': 'test', 'smtp_server': 'test',
'smtp_port': 587 'smtp_port': 587,
} }
@@ -128,7 +131,7 @@ def test_build_email_config_with_defaults():
'sender_email': 'sientia-alerts@aignosi.com', 'sender_email': 'sientia-alerts@aignosi.com',
'sender_password': 'sientia', 'sender_password': 'sientia',
'smtp_server': None, 'smtp_server': None,
'smtp_port': 587 'smtp_port': 587,
} }
@@ -147,7 +150,7 @@ def test_build_postgres_config_with_env_vars():
'password': 'sientia', 'password': 'sientia',
'dbname': 'sientia', 'dbname': 'sientia',
'min_connections': 5, 'min_connections': 5,
'max_connections': 20 'max_connections': 20,
} }
@@ -167,5 +170,5 @@ def test_build_postgres_config_with_defaults():
'password': 'sientia', 'password': 'sientia',
'dbname': 'sientia', 'dbname': 'sientia',
'min_connections': 5, 'min_connections': 5,
'max_connections': 20 'max_connections': 20,
} }

14
tests/orchestrator/utils/test_converters.py
View File

@@ -2,14 +2,14 @@ from orchestrator.utils.converters import parse_frequency
def test_parse_frequency(): def test_parse_frequency():
assert parse_frequency("1s") == 1 assert parse_frequency('1s') == 1
assert parse_frequency("1m") == 60 assert parse_frequency('1m') == 60
assert parse_frequency("1h") == 60 * 60 assert parse_frequency('1h') == 60 * 60
assert parse_frequency("1d") == 60 * 60 * 24 assert parse_frequency('1d') == 60 * 60 * 24
try: try:
parse_frequency("1") parse_frequency('1')
except ValueError as e: except ValueError as e:
assert str(e) == "Invalid frequency" assert str(e) == 'Invalid frequency'
else: else:
assert False raise AssertionError('Expected an exception to be raised')

145
tests/orchestrator/utils/test_email_builder.py
View File

@@ -1,5 +1,5 @@
import json
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock, patch
from pytest import fixture from pytest import fixture
from orchestrator.utils.email_builder import EmailBuilder from orchestrator.utils.email_builder import EmailBuilder
@@ -7,7 +7,7 @@ from orchestrator.utils.email_builder import EmailBuilder
@fixture @fixture
@patch('orchestrator.utils.email_builder.open') @patch('orchestrator.utils.email_builder.open')
def report_builder(open): def report_builder(open_mock):
return EmailBuilder(MagicMock()) return EmailBuilder(MagicMock())
@@ -26,22 +26,30 @@ def test_parameters(report_builder):
general_events = { general_events = {
'ERROR': { 'ERROR': {
'models': [ 'models': [
{'model_name': 'model_name', 'events': [ {
{'notification_id': 'ID_1', 'level': 'ERROR', 'project': 'project'}]}, 'model_name': 'model_name',
'events': [{'notification_id': 'ID_1', 'level': 'ERROR', 'project': 'project'}],
},
] ]
}, },
'WARNING': { 'WARNING': {
'models': [ 'models': [
{'model_name': 'model_name', 'events': [ {
{'notification_id': 'ID_2', 'level': 'WARNING', 'project': 'project'}]}, 'model_name': 'model_name',
'events': [
{'notification_id': 'ID_2', 'level': 'WARNING', 'project': 'project'}
],
},
] ]
}, },
'INFO': { 'INFO': {
'models': [ 'models': [
{'model_name': 'model_name', 'events': [ {
{'notification_id': 'ID_3', 'level': 'INFO', 'project': 'project'}]}, 'model_name': 'model_name',
'events': [{'notification_id': 'ID_3', 'level': 'INFO', 'project': 'project'}],
},
] ]
} },
} }
output = report_builder.parameters(general_events, 'model_name') output = report_builder.parameters(general_events, 'model_name')
@@ -55,24 +63,38 @@ def test_parameters(report_builder):
report_builder.replace_parameters.assert_any_call( report_builder.replace_parameters.assert_any_call(
report_builder.general_template, report_builder.general_template,
{'models': [ {
{'model_name': 'model_name', 'events': [ 'models': [
{'notification_id': 'ID_1', 'level': 'ERROR', 'project': 'project'}]}, {
]} 'model_name': 'model_name',
'events': [{'notification_id': 'ID_1', 'level': 'ERROR', 'project': 'project'}],
},
]
},
) )
report_builder.replace_parameters.assert_any_call( report_builder.replace_parameters.assert_any_call(
report_builder.general_template, report_builder.general_template,
{'models': [ {
{'model_name': 'model_name', 'events': [ 'models': [
{'notification_id': 'ID_2', 'level': 'WARNING', 'project': 'project'}]}, {
]} 'model_name': 'model_name',
'events': [
{'notification_id': 'ID_2', 'level': 'WARNING', 'project': 'project'}
],
},
]
},
) )
report_builder.replace_parameters.assert_any_call( report_builder.replace_parameters.assert_any_call(
report_builder.general_template, report_builder.general_template,
{'models': [ {
{'model_name': 'model_name', 'events': [ 'models': [
{'notification_id': 'ID_3', 'level': 'INFO', 'project': 'project'}]}, {
]} 'model_name': 'model_name',
'events': [{'notification_id': 'ID_3', 'level': 'INFO', 'project': 'project'}],
},
]
},
) )
@@ -81,21 +103,36 @@ def test_build_email(report_builder):
report_builder.replace_parameters = MagicMock() report_builder.replace_parameters = MagicMock()
report_data = [ report_data = [
{'notification_id': 'ID_1', 'level': 'ERROR', {
'project': 'project', 'model_name': 'model_name'}, 'notification_id': 'ID_1',
{'notification_id': 'ID_2', 'level': 'WARNING', 'level': 'ERROR',
'project': 'project', 'model_name': 'model_name'}, 'project': 'project',
{'notification_id': 'ID_2', 'level': 'INFO', 'model_name': 'model_name',
'project': 'project', 'model_name': 'model_name'}, },
{'notification_id': 'ID_3', 'level': 'ERROR', {
'project': 'project', 'model_name': 'model_name'} 'notification_id': 'ID_2',
'level': 'WARNING',
'project': 'project',
'model_name': 'model_name',
},
{
'notification_id': 'ID_2',
'level': 'INFO',
'project': 'project',
'model_name': 'model_name',
},
{
'notification_id': 'ID_3',
'level': 'ERROR',
'project': 'project',
'model_name': 'model_name',
},
] ]
html = report_builder.build_email(report_data, 'type_1') html = report_builder.build_email(report_data, 'type_1')
report_builder.replace_parameters.assert_called_once_with( report_builder.replace_parameters.assert_called_once_with(
report_builder.report_template, report_builder.report_template, report_builder.parameters.return_value
report_builder.parameters.return_value
) )
assert html == report_builder.replace_parameters.return_value assert html == report_builder.replace_parameters.return_value
@@ -108,13 +145,21 @@ def test_build_email(report_builder):
{ {
'model_name': 'model_name', 'model_name': 'model_name',
'events': [ 'events': [
{'notification_id': 'ID_1', 'level': 'ERROR', {
'project': 'project', 'model_name': 'model_name'}, 'notification_id': 'ID_1',
{'notification_id': 'ID_3', 'level': 'ERROR', 'level': 'ERROR',
'project': 'project', 'model_name': 'model_name'} 'project': 'project',
] 'model_name': 'model_name',
},
{
'notification_id': 'ID_3',
'level': 'ERROR',
'project': 'project',
'model_name': 'model_name',
},
],
} }
] ],
}, },
'WARNING': { 'WARNING': {
'section_name': 'Warnings detected:', 'section_name': 'Warnings detected:',
@@ -122,11 +167,15 @@ def test_build_email(report_builder):
{ {
'model_name': 'model_name', 'model_name': 'model_name',
'events': [ 'events': [
{'notification_id': 'ID_2', 'level': 'WARNING', {
'project': 'project', 'model_name': 'model_name'} 'notification_id': 'ID_2',
] 'level': 'WARNING',
'project': 'project',
'model_name': 'model_name',
} }
] ],
}
],
}, },
'INFO': { 'INFO': {
'section_name': 'Infos detected:', 'section_name': 'Infos detected:',
@@ -134,12 +183,16 @@ def test_build_email(report_builder):
{ {
'model_name': 'model_name', 'model_name': 'model_name',
'events': [ 'events': [
{'notification_id': 'ID_2', 'level': 'INFO', {
'project': 'project', 'model_name': 'model_name'} 'notification_id': 'ID_2',
] 'level': 'INFO',
'project': 'project',
'model_name': 'model_name',
} }
] ],
} }
],
}, },
'type_1' },
'type_1',
) )

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

@@ -1,304 +1,190 @@
from unittest.mock import patch, call from unittest.mock import call, patch
from orchestrator.utils.orchestrator_functions import ( from orchestrator.utils.orchestrator_functions import (
build_tag_config,
common_config, common_config,
minimal_retrain,
scouter,
predictions_batch,
overlap_filter_config,
process_path_priority,
gather_read_tags, gather_read_tags,
build_tag_config minimal_retrain,
overlap_filter_config,
predictions_batch,
process_path_priority,
scouter,
) )
def test_common_config(): def test_common_config():
config = { config = {
"workflow_type": "scouter", 'workflow_type': 'scouter',
"schedule_name": "test_schedule", 'schedule_name': 'test_schedule',
"model_id": "test_model_id", 'model_id': 'test_model_id',
"model": { 'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}},
"name": "test_model_name",
"model_config": {
"test_config": "test_config"
}
}
} }
result = common_config(config) result = common_config(config)
expected = { expected = {
"workflow_type": "scouter", 'workflow_type': 'scouter',
"schedule_name": "test_schedule", 'schedule_name': 'test_schedule',
"frequency": "1m", 'frequency': '1m',
"max_retry_policy": 1, 'max_retry_policy': 1,
"model_id": "test_model_id", 'model_id': 'test_model_id',
"model_name": "test_model_name", 'model_name': 'test_model_name',
"model_config": { 'model_config': {'test_config': 'test_config'},
"test_config": "test_config" 'execution_timeout_seconds': 300,
} 'task_timeout_seconds': 300,
} }
assert result == expected assert result == expected
def test_minimal_retrain(): def test_minimal_retrain():
config = { config = {
"workflow_type": "minimal_retrain", 'workflow_type': 'minimal_retrain',
"schedule_name": "test_schedule", 'schedule_name': 'test_schedule',
"model_id": "test_model_id", 'model_id': 'test_model_id',
"model": { 'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}},
"name": "test_model_name", 'query': 'select * from sientia_data.laborious_data order by "timestamp" desc limit 30;',
"model_config": { 'datetime_columns': ['timestamp'],
"test_config": "test_config"
}
},
"query": "select * from sientia_data.laborious_data order by \"timestamp\" desc limit 30;",
"datetime_columns": ["timestamp"]
} }
result = minimal_retrain(config) result = minimal_retrain(config)
expected = { expected = {
"workflow_type": "minimal_retrain", 'workflow_type': 'minimal_retrain',
"schedule_name": "test_schedule", 'schedule_name': 'test_schedule',
"frequency": "1m", 'frequency': '1m',
"max_retry_policy": 1, 'max_retry_policy': 1,
"model_id": "test_model_id", 'model_id': 'test_model_id',
"model_name": "test_model_name", 'model_name': 'test_model_name',
"model_config": { 'model_config': {'test_config': 'test_config'},
"test_config": "test_config" 'query': 'select * from sientia_data.laborious_data order by "timestamp" desc limit 30;',
}, 'schema': 'sientia_data',
"query": "select * from sientia_data.laborious_data order by \"timestamp\" desc limit 30;", 'table_name': 'log_retrain',
"schema": "sientia_data", 'datetime_columns': ['timestamp'],
"table_name": "log_retrain", 'execution_timeout_seconds': 300,
"datetime_columns": ["timestamp"] 'task_timeout_seconds': 300,
} }
assert result == expected assert result == expected
def test_scouter(): def test_scouter():
config = { config = {
"workflow_type": "scouter", 'workflow_type': 'scouter',
"schedule_name": "test_schedule", 'schedule_name': 'test_schedule',
"model_id": "test_model_id", 'model_id': 'test_model_id',
"model": { 'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}},
"name": "test_model_name", 'filters': [{'filter_name': 'test_filter_name', 'policy': 'test_policy'}],
"model_config": { 'read_tags': [
"test_config": "test_config" {'tag_name': 'test_tag_name', 'aggr_func': 'test_aggr_func', 'data_range': [1, 2]}
}
},
"filters": [
{
"filter_name": "test_filter_name",
"policy": "test_policy"
}
], ],
"read_tags": [ 'tag_retention_minutes': 10,
{ 'execution_timeout_seconds': 300,
"tag_name": "test_tag_name", 'task_timeout_seconds': 300,
"aggr_func": "test_aggr_func",
"data_range": [1, 2]
}
],
"tag_retention_minutes": 10
} }
result = scouter(config) result = scouter(config)
expected = { expected = {
"workflow_type": "scouter", 'workflow_type': 'scouter',
"schedule_name": "test_schedule", 'schedule_name': 'test_schedule',
"frequency": "1m", 'frequency': '1m',
"max_retry_policy": 1, 'max_retry_policy': 1,
"model_id": "test_model_id", 'model_id': 'test_model_id',
"model_name": "test_model_name", 'model_name': 'test_model_name',
"model_config": { 'model_config': {'test_config': 'test_config'},
"test_config": "test_config" 'topic': 'raw_test_schedule',
}, 'trigger_laborious': False,
"topic": "raw_test_schedule", 'filters': {'test_filter_name': {'policy': 'test_policy'}},
"trigger_laborious": False, 'schema': 'sientia_data',
"filters": { 'table_name': 'laborious_data',
"test_filter_name": { 'retention_time': 10 * 60,
"policy": "test_policy" 'model_tags': {'test_tag_name': {'aggr_func': 'test_aggr_func', 'data_range': [1, 2]}},
} 'debug_data_package': False,
}, 'execution_timeout_seconds': 300,
"schema": "sientia_data", 'task_timeout_seconds': 300,
"table_name": "laborious_data",
"retention_time": 10 * 60,
"model_tags": {
"test_tag_name": {
"aggr_func": "test_aggr_func",
"data_range": [1, 2]
}
},
"debug_data_package": False
} }
assert result == expected assert result == expected
def test_overlap_filter_config(): def test_overlap_filter_config():
config = [ config = [
{ {'filter_name': 'test_filter_name', 'policy': 'test_policy'},
"filter_name": "test_filter_name", {'filter_name': 'test_filter_name2', 'policy': 'test_policy2'},
"policy": "test_policy"
},
{
"filter_name": "test_filter_name2",
"policy": "test_policy2"
}
] ]
result = overlap_filter_config({ result = overlap_filter_config({'test_filter_name': {'policy': 'test_policy'}}, config)
"test_filter_name": {
"policy": "test_policy"
}
}, config)
expected = { expected = {
"test_filter_name": { 'test_filter_name': {'policy': 'test_policy', 'config': {}},
"policy": "test_policy", 'test_filter_name2': {'policy': 'test_policy2', 'config': {}},
"config": {}
},
"test_filter_name2": {
"policy": "test_policy2",
"config": {}
}
} }
assert result == expected assert result == expected
def test_process_path_priority(): def test_process_path_priority():
config = ["OTHER", "STOP", "CONTINUE"] config = ['OTHER', 'STOP', 'CONTINUE']
result = process_path_priority(config) result = process_path_priority(config)
expected = ["STOP", "CONTINUE", "REPEAT"] expected = ['STOP', 'CONTINUE', 'REPEAT']
assert result == expected assert result == expected
@patch('orchestrator.utils.orchestrator_functions.overlap_filter_config', @patch(
return_value={ 'orchestrator.utils.orchestrator_functions.overlap_filter_config',
"test_filter_name": { return_value={'test_filter_name': {'policy': 'test_policy', 'config': {}}},
"policy": "test_policy", )
"config": {} @patch(
} 'orchestrator.utils.orchestrator_functions.process_path_priority',
}) return_value=['STOP', 'CONTINUE', 'REPEAT'],
@patch('orchestrator.utils.orchestrator_functions.process_path_priority', )
return_value=["STOP", "CONTINUE", "REPEAT"]) def test_predictions_batch(mock_process_path_priority, mock_overlap_filter_config):
def test_predictions_batch(mock_process_path_priority,
mock_overlap_filter_config):
config = { config = {
"schedule_name": "test_schedule", 'schedule_name': 'test_schedule',
"workflow_type": "predictions_batch", 'workflow_type': 'predictions_batch',
"model_id": "test_model_id", 'model_id': 'test_model_id',
"model": { 'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}},
"name": "test_model_name", 'query': 'test_query',
"model_config": { 'write_tags': [
"test_config": "test_config" {'server_id': 'test_server_id', 'type': 'prediction', 'addr': 'test_addr'},
} {'server_id': 'test_server_id', 'type': 'confidence', 'addr': 'test_addr'},
},
"query": "test_query",
"write_tags": [
{
"server_id": "test_server_id",
"type": "prediction",
"addr": "test_addr"
},
{
"server_id": "test_server_id",
"type": "confidence",
"addr": "test_addr"
}
], ],
"input_filters": [ 'input_filters': [{'filter_name': 'test_filter_name', 'policy': 'test_policy'}],
{ 'mlflow_transform_filters': [{'filter_name': 'test_filter_name', 'policy': 'test_policy'}],
"filter_name": "test_filter_name", 'mlflow_predict_filters': [{'filter_name': 'test_filter_name', 'policy': 'test_policy'}],
"policy": "test_policy" 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
} 'datetime_columns': ['timestamp'],
], 'predictions_storage_policy': 'erl:1',
"mlflow_transform_filters": [
{
"filter_name": "test_filter_name",
"policy": "test_policy"
}
],
"mlflow_predict_filters": [
{
"filter_name": "test_filter_name",
"policy": "test_policy"
}
],
"path_priority": ["STOP", "CONTINUE", "REPEAT"],
"datetime_columns": ["timestamp"],
"predictions_storage_policy": "erl:1"
} }
result = predictions_batch(config) result = predictions_batch(config)
mock_overlap_filter_config.assert_has_calls([ mock_overlap_filter_config.assert_has_calls(
call({ [call({'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, config['input_filters'])]
"EMPTY_DATA": { )
"policy": "STOP", mock_overlap_filter_config.assert_has_calls(
"config": {} [call({'API_ERROR': {'policy': 'STOP', 'config': {}}}, config['mlflow_transform_filters'])]
} )
}, config['input_filters']) mock_overlap_filter_config.assert_has_calls(
]) [call({'API_ERROR': {'policy': 'STOP', 'config': {}}}, config['mlflow_predict_filters'])]
mock_overlap_filter_config.assert_has_calls([ )
call({
"API_ERROR": {
"policy": "STOP",
"config": {}
}
}, config['mlflow_transform_filters'])
])
mock_overlap_filter_config.assert_has_calls([
call({
"API_ERROR": {
"policy": "STOP",
"config": {}
}
}, config['mlflow_predict_filters'])
])
mock_process_path_priority.assert_called_once_with(config['path_priority']) mock_process_path_priority.assert_called_once_with(config['path_priority'])
expected = { expected = {
"workflow_type": "predictions_batch", 'workflow_type': 'predictions_batch',
"schedule_name": "test_schedule", 'schedule_name': 'test_schedule',
"frequency": "1m", 'frequency': '1m',
"max_retry_policy": 1, 'max_retry_policy': 1,
"model_id": "test_model_id", 'model_id': 'test_model_id',
"model_name": "test_model_name", 'model_name': 'test_model_name',
"model_config": { 'model_config': {'test_config': 'test_config'},
"test_config": "test_config" 'query': 'test_query',
}, 'schema': 'sientia_data',
"query": "test_query", 'table_name': 'predictions',
"schema": "sientia_data", 'retention_time': 60 * 60,
"table_name": "predictions", 'opc_output_config': {
"retention_time": 60 * 60, 'test_server_id': {
"opc_output_config": { 'prediction_tags': {'test_addr': {'data_type': 'float'}},
"test_server_id": { 'confidence_tags': {'test_addr': {'data_type': 'float'}},
"prediction_tags": {
"test_addr": {
"data_type": "float"
} }
}, },
"confidence_tags": { 'input_filters': {'test_filter_name': {'policy': 'test_policy', 'config': {}}},
"test_addr": { 'mlflow_transform_filters': {'test_filter_name': {'policy': 'test_policy', 'config': {}}},
"data_type": "float" 'mlflow_predict_filters': {'test_filter_name': {'policy': 'test_policy', 'config': {}}},
} 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
} 'datetime_columns': ['timestamp'],
} 'predictions_storage_policy': 'erl:1',
}, 'execution_timeout_seconds': 300,
"input_filters": { 'task_timeout_seconds': 300,
"test_filter_name": {
"policy": "test_policy",
"config": {}
}
},
"mlflow_transform_filters": {
"test_filter_name": {
"policy": "test_policy",
"config": {}
}
},
"mlflow_predict_filters": {
"test_filter_name": {
"policy": "test_policy",
"config": {}
}
},
"path_priority": ["STOP", "CONTINUE", "REPEAT"],
"datetime_columns": ["timestamp"],
"predictions_storage_policy": "erl:1"
} }
assert result == expected assert result == expected
@@ -306,93 +192,81 @@ def test_predictions_batch(mock_process_path_priority,
def test_gather_read_tags(): def test_gather_read_tags():
pipelines = [ pipelines = [
{ {
"schedule_name": "test_schedule", 'schedule_name': 'test_schedule',
"read_tags": [ 'read_tags': [
{ {
"server_id": "1", 'server_id': '1',
"server_name": "test_server_name", 'server_name': 'test_server_name',
"tag_address": "test_tag_address" 'tag_address': 'test_tag_address',
} }
] ],
}, },
{ {
"schedule_name": "test_schedule2", 'schedule_name': 'test_schedule2',
"read_tags": [ 'read_tags': [
{ {
"server_id": "2", 'server_id': '2',
"server_name": "test_server_name2", 'server_name': 'test_server_name2',
"tag_address": "test_tag_address2" 'tag_address': 'test_tag_address2',
}, },
{ {
"server_id": "2", 'server_id': '2',
"server_name": "test_server_name2", 'server_name': 'test_server_name2',
"tag_address": "test_tag_address3" 'tag_address': 'test_tag_address3',
} },
] ],
} },
] ]
result = gather_read_tags(pipelines) result = gather_read_tags(pipelines)
expected = { expected = {
"1:test_tag_address": { '1:test_tag_address': {
"server_id": "1", 'server_id': '1',
"server_name": "test_server_name", 'server_name': 'test_server_name',
"tag_address": "test_tag_address", 'tag_address': 'test_tag_address',
"topics": ["raw_test_schedule"] 'topics': ['raw_test_schedule'],
}, },
"2:test_tag_address2": { '2:test_tag_address2': {
"server_id": "2", 'server_id': '2',
"server_name": "test_server_name2", 'server_name': 'test_server_name2',
"tag_address": "test_tag_address2", 'tag_address': 'test_tag_address2',
"topics": ["raw_test_schedule2"] 'topics': ['raw_test_schedule2'],
},
'2:test_tag_address3': {
'server_id': '2',
'server_name': 'test_server_name2',
'tag_address': 'test_tag_address3',
'topics': ['raw_test_schedule2'],
}, },
"2:test_tag_address3": {
"server_id": "2",
"server_name": "test_server_name2",
"tag_address": "test_tag_address3",
"topics": ["raw_test_schedule2"]
}
} }
assert result == expected assert result == expected
def test_build_tag_config(): def test_build_tag_config():
tag = { tag = {'server_id': '1', 'server_name': 'test_server_name', 'tag_address': 'test_tag_address'}
"server_id": "1", opc_servers = {'1': {'server_name': 'test_server_name', 'url': 'test_url', 'uri': 'test_uri'}}
"server_name": "test_server_name", slot_config = {'1': {}}
"tag_address": "test_tag_address"
}
opc_servers = {
"1": {
"server_name": "test_server_name",
"url": "test_url",
"uri": "test_uri"
}
}
slot_config = {
"1": {}
}
i = 1 i = 1
result = build_tag_config(tag, slot_config, opc_servers, i) result = build_tag_config(tag, slot_config, opc_servers, i)
expected = { expected = {
"1": { '1': {
"test_server_name": { 'test_server_name': {
"server_id": "1", 'server_id': '1',
"name": "test_server_name", 'name': 'test_server_name',
"url": "test_url", 'url': 'test_url',
"server_uri": "test_uri", 'server_uri': 'test_uri',
"cert_path": None, 'cert_path': None,
"private_key_path": None, 'private_key_path': None,
"server_cert_path": None, 'server_cert_path': None,
"tags": { 'tags': {
"test_tag_address": { 'test_tag_address': {
"server_id": "1", 'server_id': '1',
"server_name": "test_server_name", 'server_name': 'test_server_name',
"tag_address": "test_tag_address" 'tag_address': 'test_tag_address',
}
} }
},
} }
} }
} }
@@ -400,13 +274,10 @@ def test_build_tag_config():
def test_build_tag_config_no_server_id(): def test_build_tag_config_no_server_id():
tag = { tag = {'server_id': '1', 'tag_address': 'test_tag_address'}
"server_id": "1",
"tag_address": "test_tag_address"
}
opc_servers = {} opc_servers = {}
try: try:
build_tag_config(tag, {}, opc_servers, 1) build_tag_config(tag, {}, opc_servers, 1)
except ValueError as e: except ValueError as e:
assert str(e) == "Server 1 not found in opc_servers" assert str(e) == 'Server 1 not found in opc_servers'

106
tests/orchestrator/workflows/subworkflows/test_load_notification_package.py
View File

@@ -1,7 +1,9 @@
from unittest.mock import AsyncMock, patch, ANY, call from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark from pytest import fixture, mark
from orchestrator.workflows.subworkflows.load_notification_package import LoadNotificationPackage
from orchestrator.activities.activities import Activities from orchestrator.activities.activities import Activities
from orchestrator.workflows.subworkflows.load_notification_package import LoadNotificationPackage
@fixture @fixture
@@ -20,14 +22,14 @@ metadata = {
@mark.asyncio @mark.asyncio
@patch("orchestrator.workflows.subworkflows.load_notification_package.workflow", new_callable=AsyncMock) @patch(
'orchestrator.workflows.subworkflows.load_notification_package.workflow', new_callable=AsyncMock
)
async def test_run(workflow_mock, load_notification_package): async def test_run(workflow_mock, load_notification_package):
input_data = { input_data = {
'metadata': metadata, 'metadata': metadata,
'mail_type': 'test_mail_type', 'mail_type': 'test_mail_type',
'base_data_filter': { 'base_data_filter': {'level': 'ERROR'},
'level': 'ERROR'
}
} }
workflow_mock.start_local_activity_method.side_effect = [ workflow_mock.start_local_activity_method.side_effect = [
@@ -41,7 +43,7 @@ async def test_run(workflow_mock, load_notification_package):
{ {
'id_r': '1', 'id_r': '1',
} }
] ],
] ]
output = await load_notification_package.run(input_data) output = await load_notification_package.run(input_data)
@@ -57,55 +59,52 @@ async def test_run(workflow_mock, load_notification_package):
{ {
'id_r': '1', 'id_r': '1',
} }
] ],
} }
workflow_mock.start_local_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
[
call( call(
Activities.get_last_data_timestamp, Activities.get_last_data_timestamp,
{ {**input_data['metadata'], 'mail_type': 'test_mail_type'},
**input_data['metadata'],
'mail_type': 'test_mail_type'
},
start_to_close_timeout=ANY, start_to_close_timeout=ANY,
retry_policy=ANY retry_policy=ANY,
)
]
) )
])
workflow_mock.start_local_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
[
call( call(
Activities.find_documents_in_mongodb, Activities.find_documents_in_mongodb,
{ {
**input_data['metadata'], **input_data['metadata'],
'query': { 'query': {'collection': 'receiver_groups', 'filters': {'active': True}},
'collection': 'receiver_groups',
'filters': {
'active': True
}
}
}, },
start_to_close_timeout=ANY, start_to_close_timeout=ANY,
retry_policy=ANY retry_policy=ANY,
)
]
) )
])
workflow_mock.start_local_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
[
call( call(
Activities.load_latest_data, Activities.load_latest_data,
{ {
**input_data['metadata'], **input_data['metadata'],
'collection_name': 'notification_queue', 'collection_name': 'notification_queue',
'last_data_timestamp': '2023-01-01 12:00:00', 'last_data_timestamp': '2023-01-01 12:00:00',
'base_data_filter': { 'base_data_filter': {'level': 'ERROR'},
'level': 'ERROR'
}
}, },
start_to_close_timeout=ANY, start_to_close_timeout=ANY,
retry_policy=ANY retry_policy=ANY,
)
]
) )
])
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
[
call( call(
Activities.put_last_data_timestamp, Activities.put_last_data_timestamp,
{ {
@@ -115,65 +114,58 @@ async def test_run(workflow_mock, load_notification_package):
'id': '1', 'id': '1',
} }
], ],
'mail_type': 'test_mail_type' 'mail_type': 'test_mail_type',
}, },
start_to_close_timeout=ANY, start_to_close_timeout=ANY,
retry_policy=ANY retry_policy=ANY,
)
]
) )
])
@mark.asyncio @mark.asyncio
@patch("orchestrator.workflows.subworkflows.load_notification_package.workflow", new_callable=AsyncMock) @patch(
'orchestrator.workflows.subworkflows.load_notification_package.workflow', new_callable=AsyncMock
)
async def test_run_no_data(workflow_mock, load_notification_package): async def test_run_no_data(workflow_mock, load_notification_package):
input_data = { input_data = {
'metadata': metadata, 'metadata': metadata,
'mail_type': 'test_mail_type', 'mail_type': 'test_mail_type',
'base_data_filter': { 'base_data_filter': {'level': 'ERROR'},
'level': 'ERROR'
}
} }
workflow_mock.start_local_activity_method.side_effect = [ workflow_mock.start_local_activity_method.side_effect = ['2023-01-01 12:00:00', [], []]
'2023-01-01 12:00:00',
[],
[]
]
output = await load_notification_package.run(input_data) output = await load_notification_package.run(input_data)
assert output == { assert output == {
'last_timestamp': '2023-01-01 12:00:00', 'last_timestamp': '2023-01-01 12:00:00',
'notification_package': [], 'notification_package': [],
'sending_configs': [] 'sending_configs': [],
} }
workflow_mock.start_activity_method.assert_not_called() workflow_mock.start_activity_method.assert_not_called()
@mark.asyncio @mark.asyncio
@patch("orchestrator.workflows.subworkflows.load_notification_package.workflow", new_callable=AsyncMock) @patch(
async def test_run_no_data(workflow_mock, load_notification_package): 'orchestrator.workflows.subworkflows.load_notification_package.workflow', new_callable=AsyncMock
)
async def test_run_has_data(workflow_mock, load_notification_package):
input_data = { input_data = {
'metadata': metadata, 'metadata': metadata,
'base_data_filter': { 'base_data_filter': {'level': 'ERROR'},
'level': 'ERROR' 'mail_type': 'test_mail_type',
},
'mail_type': 'test_mail_type'
} }
workflow_mock.start_local_activity_method.side_effect = [ workflow_mock.start_local_activity_method.side_effect = ['2023-01-01 12:00:00', ['data'], []]
'2023-01-01 12:00:00',
["data"],
[]
]
output = await load_notification_package.run(input_data) output = await load_notification_package.run(input_data)
assert output == { assert output == {
'last_timestamp': '2023-01-01 12:00:00', 'last_timestamp': '2023-01-01 12:00:00',
'notification_package': ["data"], 'notification_package': ['data'],
'sending_configs': [] 'sending_configs': [],
} }
workflow_mock.start_activity_method.assert_not_called() workflow_mock.start_activity_method.assert_not_called()

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

@@ -1,9 +1,11 @@
from unittest.mock import AsyncMock, patch, ANY, call from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark from pytest import fixture, mark
from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications
from orchestrator.activities.activities import Activities
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ
from orchestrator.activities.activities import Activities
from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications
@fixture @fixture
def process_notifications(): def process_notifications():
@@ -21,11 +23,11 @@ metadata = {
@mark.asyncio @mark.asyncio
@patch("orchestrator.workflows.subworkflows.process_notifications.workflow", new_callable=AsyncMock) @patch('orchestrator.workflows.subworkflows.process_notifications.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock, process_notifications): async def test_run(workflow_mock, process_notifications):
input_data = { input_data = {
'metadata': metadata, 'metadata': metadata,
'notification_package': ["content"], 'notification_package': ['content'],
'mail_type': 'test_mail_type', 'mail_type': 'test_mail_type',
'schema': 'test_schema', 'schema': 'test_schema',
'table_name': 'test_table_name', 'table_name': 'test_table_name',
@@ -35,45 +37,52 @@ async def test_run(workflow_mock, process_notifications):
assert response == workflow_mock.execute_local_activity_method.return_value assert response == workflow_mock.execute_local_activity_method.return_value
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls(
[
call( call(
Activities.build_email_html, Activities.build_email_html,
{ {
**metadata, **metadata,
'receiver_groups': input_data['notification_package'], 'receiver_groups': input_data['notification_package'],
'mail_type': input_data['mail_type'] 'mail_type': input_data['mail_type'],
}, },
schedule_to_close_timeout=ANY, schedule_to_close_timeout=ANY,
retry_policy=ANY retry_policy=ANY,
) )
]) ]
workflow_mock.execute_local_activity_method.assert_has_calls([ )
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call( call(
Activities.format_log_report, Activities.format_log_report,
{ {
**metadata, **metadata,
'receiver_groups': workflow_mock.execute_activity_method.return_value, 'receiver_groups': workflow_mock.execute_activity_method.return_value,
'mail_type': input_data['mail_type'] 'mail_type': input_data['mail_type'],
}, },
schedule_to_close_timeout=ANY, schedule_to_close_timeout=ANY,
retry_policy=ANY retry_policy=ANY,
)
]
) )
])
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_activity_method.assert_has_calls(
[
call( call(
Activities.send_email, Activities.send_email,
{ {
**metadata, **metadata,
'receiver_groups': workflow_mock.execute_local_activity_method.return_value, 'receiver_groups': workflow_mock.execute_local_activity_method.return_value,
'mail_type': input_data['mail_type'] 'mail_type': input_data['mail_type'],
}, },
schedule_to_close_timeout=ANY, schedule_to_close_timeout=ANY,
retry_policy=ANY retry_policy=ANY,
)] )
]
) )
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_activity_method.assert_has_calls(
[
call( call(
Activities.export_data_to_postgres, Activities.export_data_to_postgres,
{ {
@@ -83,21 +92,22 @@ async def test_run(workflow_mock, process_notifications):
'data': workflow_mock.execute_local_activity_method.return_value, 'data': workflow_mock.execute_local_activity_method.return_value,
'timestamp_conversion': { 'timestamp_conversion': {
'column': 'timestamp', 'column': 'timestamp',
'format': DATETIME_FORMAT_MS_WITH_TZ 'format': DATETIME_FORMAT_MS_WITH_TZ,
} },
}, },
schedule_to_close_timeout=ANY, schedule_to_close_timeout=ANY,
retry_policy=ANY retry_policy=ANY,
)
]
) )
])
@mark.asyncio @mark.asyncio
@patch("orchestrator.workflows.subworkflows.process_notifications.workflow", new_callable=AsyncMock) @patch('orchestrator.workflows.subworkflows.process_notifications.workflow', new_callable=AsyncMock)
async def test_run_send_email_return_empty(workflow_mock, process_notifications): async def test_run_send_email_return_empty(workflow_mock, process_notifications):
input_data = { input_data = {
'metadata': metadata, 'metadata': metadata,
'notification_package': ["content"], 'notification_package': ['content'],
'mail_type': 'test_mail_type', 'mail_type': 'test_mail_type',
'schema': 'test_schema', 'schema': 'test_schema',
'table_name': 'test_table_name', 'table_name': 'test_table_name',
@@ -105,32 +115,36 @@ async def test_run_send_email_return_empty(workflow_mock, process_notifications)
workflow_mock.execute_activity_method.return_value = [] workflow_mock.execute_activity_method.return_value = []
response = await process_notifications.run(input_data) await process_notifications.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls(
[
call( call(
Activities.build_email_html, Activities.build_email_html,
{ {
**metadata, **metadata,
'receiver_groups': input_data['notification_package'], 'receiver_groups': input_data['notification_package'],
'mail_type': input_data['mail_type'] 'mail_type': input_data['mail_type'],
}, },
schedule_to_close_timeout=ANY, schedule_to_close_timeout=ANY,
retry_policy=ANY retry_policy=ANY,
)
]
) )
])
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_activity_method.assert_has_calls(
[
call( call(
Activities.send_email, Activities.send_email,
{ {
**metadata, **metadata,
'receiver_groups': workflow_mock.execute_local_activity_method.return_value, 'receiver_groups': workflow_mock.execute_local_activity_method.return_value,
'mail_type': input_data['mail_type'] 'mail_type': input_data['mail_type'],
}, },
schedule_to_close_timeout=ANY, schedule_to_close_timeout=ANY,
retry_policy=ANY retry_policy=ANY,
)] )
]
) )
assert workflow_mock.execute_activity_method.call_count == 1 assert workflow_mock.execute_activity_method.call_count == 1

111
tests/orchestrator/workflows/test_alerts.py
View File

@@ -1,7 +1,9 @@
from unittest.mock import AsyncMock, MagicMock, patch, ANY, call from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
from pytest import fixture, mark from pytest import fixture, mark
from orchestrator.workflows.alerts import Alerts
from orchestrator.activities.activities import Activities from orchestrator.activities.activities import Activities
from orchestrator.workflows.alerts import Alerts
@fixture @fixture
@@ -20,30 +22,23 @@ metadata = {
@mark.asyncio @mark.asyncio
@patch("orchestrator.workflows.alerts.workflow", new_callable=AsyncMock) @patch('orchestrator.workflows.alerts.workflow', new_callable=AsyncMock)
async def test_run_full_flow(workflow_mock, alerts): async def test_run_full_flow(workflow_mock, alerts):
input_data = { input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600}
'schedule_name': 'test-schedule-name',
'notification_ttl': 300,
'sent_ttl': 600
}
await alerts.run(input_data) await alerts.run(input_data)
workflow_mock.execute_child_workflow.assert_has_calls([ workflow_mock.execute_child_workflow.assert_has_calls(
[
call( call(
'load_notification_package', 'load_notification_package',
{ {**input_data, 'metadata': metadata, 'base_data_filter': {'level': 'ERROR'}},
**input_data, )
'metadata': metadata, ]
'base_data_filter': {
'level': 'ERROR'
}
}
) )
])
workflow_mock.execute_child_workflow.assert_has_calls([ workflow_mock.execute_child_workflow.assert_has_calls(
[
call( call(
'process_notifications', 'process_notifications',
{ {
@@ -51,82 +46,79 @@ async def test_run_full_flow(workflow_mock, alerts):
'mail_type': 'Alerts', 'mail_type': 'Alerts',
'notification_package': workflow_mock.execute_local_activity_method.return_value, 'notification_package': workflow_mock.execute_local_activity_method.return_value,
'schema': 'sientia_data', 'schema': 'sientia_data',
'table_name': 'log_report' 'table_name': 'log_report',
} },
)
]
) )
])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls(
[
call( call(
Activities.filter_notification_alerts, Activities.filter_notification_alerts,
{ {
**metadata, **metadata,
'notification_package': workflow_mock.execute_child_workflow.return_value['notification_package'], 'notification_package': workflow_mock.execute_child_workflow.return_value[
'sending_configs': workflow_mock.execute_child_workflow.return_value['sending_configs'], 'notification_package'
'notification_ttl': input_data['notification_ttl'] ],
'sending_configs': workflow_mock.execute_child_workflow.return_value[
'sending_configs'
],
'notification_ttl': input_data['notification_ttl'],
}, },
schedule_to_close_timeout=ANY, schedule_to_close_timeout=ANY,
retry_policy=ANY retry_policy=ANY,
)
]
) )
])
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_activity_method.assert_has_calls(
[
call( call(
Activities.store_notification_cache, Activities.store_notification_cache,
{ {
**metadata, **metadata,
'log_report': workflow_mock.execute_child_workflow.return_value, 'log_report': workflow_mock.execute_child_workflow.return_value,
'sent_ttl': input_data['sent_ttl'] 'sent_ttl': input_data['sent_ttl'],
}, },
schedule_to_close_timeout=ANY, schedule_to_close_timeout=ANY,
retry_policy=ANY retry_policy=ANY,
)
]
) )
])
@mark.asyncio @mark.asyncio
@patch("orchestrator.workflows.alerts.workflow", new_callable=AsyncMock) @patch('orchestrator.workflows.alerts.workflow', new_callable=AsyncMock)
async def test_run_no_data(workflow_mock, alerts): async def test_run_no_data(workflow_mock, alerts):
workflow_mock.execute_child_workflow.return_value = { workflow_mock.execute_child_workflow.return_value = {
'last_timestamp': '2023-01-01 12:00:00.000000', 'last_timestamp': '2023-01-01 12:00:00.000000',
'notification_package': [], 'notification_package': [],
'sending_configs': [] 'sending_configs': [],
} }
input_data = { input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600}
'schedule_name': 'test-schedule-name',
'notification_ttl': 300,
'sent_ttl': 600
}
await alerts.run(input_data) await alerts.run(input_data)
workflow_mock.execute_child_workflow.assert_has_calls([ workflow_mock.execute_child_workflow.assert_has_calls(
[
call( call(
'load_notification_package', 'load_notification_package',
{ {**input_data, 'metadata': metadata, 'base_data_filter': {'level': 'ERROR'}},
**input_data, )
'metadata': metadata, ]
'base_data_filter': {
'level': 'ERROR'
}
}
) )
])
workflow_mock.execute_local_activity_method.assert_not_called() workflow_mock.execute_local_activity_method.assert_not_called()
@mark.asyncio @mark.asyncio
@patch("orchestrator.workflows.alerts.workflow", new_callable=AsyncMock) @patch('orchestrator.workflows.alerts.workflow', new_callable=AsyncMock)
async def test_run_no_receiver_groups(workflow_mock, alerts): async def test_run_no_receiver_groups(workflow_mock, alerts):
workflow_mock.execute_local_activity_method.return_value = {} workflow_mock.execute_local_activity_method.return_value = {}
input_data = { input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600}
'schedule_name': 'test-schedule-name',
'notification_ttl': 300,
'sent_ttl': 600
}
await alerts.run(input_data) await alerts.run(input_data)
@@ -134,18 +126,11 @@ async def test_run_no_receiver_groups(workflow_mock, alerts):
@mark.asyncio @mark.asyncio
@patch("orchestrator.workflows.alerts.workflow", new_callable=AsyncMock) @patch('orchestrator.workflows.alerts.workflow', new_callable=AsyncMock)
async def test_run_no_log_report(workflow_mock, alerts): async def test_run_no_log_report(workflow_mock, alerts):
workflow_mock.execute_child_workflow.side_effect = [ workflow_mock.execute_child_workflow.side_effect = [MagicMock(), []]
MagicMock(),
[]
]
input_data = { input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600}
'schedule_name': 'test-schedule-name',
'notification_ttl': 300,
'sent_ttl': 600
}
await alerts.run(input_data) await alerts.run(input_data)

261
tests/orchestrator/workflows/test_orchestrator.py
View File

@@ -1,7 +1,9 @@
from unittest.mock import AsyncMock, patch, ANY, call from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark from pytest import fixture, mark
from orchestrator.workflows.orchestrator import Orchestrator
from orchestrator.activities.activities import Activities from orchestrator.activities.activities import Activities
from orchestrator.workflows.orchestrator import Orchestrator
@fixture @fixture
@@ -20,103 +22,106 @@ metadata = {
@mark.asyncio @mark.asyncio
@patch("orchestrator.workflows.orchestrator.workflow", new_callable=AsyncMock) @patch('orchestrator.workflows.orchestrator.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock, orchestrator): async def test_run(workflow_mock, orchestrator):
input_data = { input_data = {
"pipelines_query": "SELECT * FROM bucket", 'pipelines_query': 'SELECT * FROM bucket',
"opc_servers_query": "SELECT * FROM servers", 'opc_servers_query': 'SELECT * FROM servers',
"schedule_name": "test-schedule-name", 'schedule_name': 'test-schedule-name',
} }
await orchestrator.run(input_data) await orchestrator.run(input_data)
workflow_mock.start_local_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
[
call( call(
Activities.aggregate_documents_in_mongodb, Activities.aggregate_documents_in_mongodb,
{ {
**metadata, **metadata,
"query": input_data["pipelines_query"], 'query': input_data['pipelines_query'],
"timestamp_fields": ["updated_at"] 'timestamp_fields': ['updated_at'],
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_local_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
[
call(
Activities.find_documents_in_mongodb,
{**metadata, 'query': input_data['opc_servers_query']},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.start_local_activity_method.assert_has_calls(
[
call( call(
Activities.find_documents_in_mongodb, Activities.find_documents_in_mongodb,
{ {
**metadata, **metadata,
"query": input_data["opc_servers_query"] 'query': {'collection': 'orchestrated_schedules'},
'timestamp_fields': ['updated_at'],
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
) )
]) ]
workflow_mock.start_local_activity_method.assert_has_calls([
call(
Activities.find_documents_in_mongodb,
{
**metadata,
"query": {
"collection": "orchestrated_schedules"
},
"timestamp_fields": ["updated_at"]
},
retry_policy=ANY,
start_to_close_timeout=ANY
) )
])
workflow_mock.start_local_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
[
call( call(
Activities.load_opc_slots, Activities.load_opc_slots,
{ {**metadata},
**metadata
},
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_local_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
[
call( call(
Activities.load_active_ingestors, Activities.load_active_ingestors,
{ {**metadata},
**metadata
},
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_local_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
[
call( call(
Activities.format_schedule_config, Activities.format_schedule_config,
{ {
**metadata, **metadata,
'schedule_config': workflow_mock.start_local_activity_method.return_value 'schedule_config': workflow_mock.start_local_activity_method.return_value,
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_local_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
[
call( call(
Activities.process_schedules, Activities.process_schedules,
{ {**metadata, 'pipelines': workflow_mock.start_local_activity_method.return_value},
**metadata,
'pipelines': workflow_mock.start_local_activity_method.return_value
},
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_local_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
[
call( call(
Activities.process_slots, Activities.process_slots,
{ {
@@ -126,126 +131,151 @@ async def test_run(workflow_mock, orchestrator):
'pipelines': workflow_mock.start_local_activity_method.return_value, 'pipelines': workflow_mock.start_local_activity_method.return_value,
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_local_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
[
call( call(
Activities.create_schedule_config, Activities.create_schedule_config,
{ {
**metadata, **metadata,
'current_schedule_config': workflow_mock.start_local_activity_method.return_value, 'current_schedule_config': workflow_mock.start_local_activity_method.return_value,
'schedule_config': workflow_mock.start_local_activity_method.return_value 'schedule_config': workflow_mock.start_local_activity_method.return_value,
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_local_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
[
call( call(
Activities.create_slot_config, Activities.create_slot_config,
{ {
**metadata, **metadata,
'current_slot_config': workflow_mock.start_local_activity_method.return_value, 'current_slot_config': workflow_mock.start_local_activity_method.return_value,
'slot_config': workflow_mock.start_local_activity_method.return_value 'slot_config': workflow_mock.start_local_activity_method.return_value,
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
[
call( call(
Activities.normalize_schedules, Activities.normalize_schedules,
{ {
**metadata, **metadata,
'orchestrated_schedules': workflow_mock.start_local_activity_method.return_value 'orchestrated_schedules': workflow_mock.start_local_activity_method.return_value,
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
[
call( call(
Activities.create_collection_with_ttl_index, Activities.create_collection_with_ttl_index,
{ {
**metadata, **metadata,
'pipelines': workflow_mock.start_local_activity_method.return_value['scouter'] 'pipelines': workflow_mock.start_local_activity_method.return_value['scouter'],
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
[
call( call(
Activities.delete_slots, Activities.delete_slots,
{ {
**metadata, **metadata,
'to_delete': 'to_delete': workflow_mock.start_local_activity_method.return_value[
workflow_mock.start_local_activity_method.return_value['to_delete'] 'to_delete'
],
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
[
call( call(
Activities.update_slots, Activities.update_slots,
{ {
**metadata, **metadata,
'to_insert': 'to_insert': workflow_mock.start_local_activity_method.return_value[
workflow_mock.start_local_activity_method.return_value['to_insert'] 'to_insert'
],
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
[
call( call(
Activities.delete_schedules, Activities.delete_schedules,
{ {
**metadata, **metadata,
'schedules': 'schedules': workflow_mock.start_local_activity_method.return_value[
workflow_mock.start_local_activity_method.return_value['to_delete'] 'to_delete'
],
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
[
call( call(
Activities.create_schedules, Activities.create_schedules,
{ {
**metadata, **metadata,
'schedules': 'schedules': workflow_mock.start_local_activity_method.return_value[
workflow_mock.start_local_activity_method.return_value['to_create'] 'to_create'
],
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
[
call( call(
Activities.update_schedules, Activities.update_schedules,
{ {
**metadata, **metadata,
'schedules': 'schedules': workflow_mock.start_local_activity_method.return_value[
workflow_mock.start_local_activity_method.return_value['to_update'] 'to_update'
],
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
[
call( call(
Activities.report_schedule_orchestration, Activities.report_schedule_orchestration,
{ {
@@ -255,11 +285,13 @@ async def test_run(workflow_mock, orchestrator):
'deleted_schedules': workflow_mock.start_activity_method.return_value, 'deleted_schedules': workflow_mock.start_activity_method.return_value,
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
[
call( call(
Activities.report_slot_orchestration, Activities.report_slot_orchestration,
{ {
@@ -268,11 +300,13 @@ async def test_run(workflow_mock, orchestrator):
'deleted_slots': workflow_mock.start_activity_method.return_value, 'deleted_slots': workflow_mock.start_activity_method.return_value,
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
[
call( call(
Activities.update_pipelines_timestamps, Activities.update_pipelines_timestamps,
{ {
@@ -280,11 +314,13 @@ async def test_run(workflow_mock, orchestrator):
'updated_pipelines': workflow_mock.start_activity_method.return_value, 'updated_pipelines': workflow_mock.start_activity_method.return_value,
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
[
call( call(
Activities.delete_pipelines_timestamps, Activities.delete_pipelines_timestamps,
{ {
@@ -292,11 +328,13 @@ async def test_run(workflow_mock, orchestrator):
'deleted_pipelines': workflow_mock.start_activity_method.return_value, 'deleted_pipelines': workflow_mock.start_activity_method.return_value,
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
[
call( call(
Activities.create_pipelines_timestamps, Activities.create_pipelines_timestamps,
{ {
@@ -304,6 +342,7 @@ async def test_run(workflow_mock, orchestrator):
'created_pipelines': workflow_mock.start_activity_method.return_value, 'created_pipelines': workflow_mock.start_activity_method.return_value,
}, },
retry_policy=ANY, retry_policy=ANY,
start_to_close_timeout=ANY start_to_close_timeout=ANY,
)
]
) )
])

82
tests/orchestrator/workflows/test_reports.py
View File

@@ -1,7 +1,9 @@
from unittest.mock import AsyncMock, patch, ANY, call from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark from pytest import fixture, mark
from orchestrator.workflows.reports import Reports
from orchestrator.activities.activities import Activities from orchestrator.activities.activities import Activities
from orchestrator.workflows.reports import Reports
@fixture @fixture
@@ -20,28 +22,23 @@ metadata = {
@mark.asyncio @mark.asyncio
@patch("orchestrator.workflows.reports.workflow", new_callable=AsyncMock) @patch('orchestrator.workflows.reports.workflow', new_callable=AsyncMock)
async def test_run_full_flow(workflow_mock, reports): async def test_run_full_flow(workflow_mock, reports):
input_data = { input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600}
'schedule_name': 'test-schedule-name',
'notification_ttl': 300,
'sent_ttl': 600
}
await reports.run(input_data) await reports.run(input_data)
workflow_mock.execute_child_workflow.assert_has_calls([ workflow_mock.execute_child_workflow.assert_has_calls(
[
call( call(
'load_notification_package', 'load_notification_package',
{ {**input_data, 'metadata': metadata, 'base_data_filter': {}},
**input_data, )
'metadata': metadata, ]
'base_data_filter': {}
}
) )
])
workflow_mock.execute_child_workflow.assert_has_calls([ workflow_mock.execute_child_workflow.assert_has_calls(
[
call( call(
'process_notifications', 'process_notifications',
{ {
@@ -49,66 +46,63 @@ async def test_run_full_flow(workflow_mock, reports):
'mail_type': 'Reports', 'mail_type': 'Reports',
'notification_package': workflow_mock.execute_local_activity_method.return_value, 'notification_package': workflow_mock.execute_local_activity_method.return_value,
'schema': 'sientia_data', 'schema': 'sientia_data',
'table_name': 'log_report' 'table_name': 'log_report',
} },
)
]
) )
])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls(
[
call( call(
Activities.filter_notification_reports, Activities.filter_notification_reports,
{ {
**metadata, **metadata,
'notification_package': workflow_mock.execute_child_workflow.return_value['notification_package'], 'notification_package': workflow_mock.execute_child_workflow.return_value[
'sending_configs': workflow_mock.execute_child_workflow.return_value['sending_configs'] 'notification_package'
],
'sending_configs': workflow_mock.execute_child_workflow.return_value[
'sending_configs'
],
}, },
schedule_to_close_timeout=ANY, schedule_to_close_timeout=ANY,
retry_policy=ANY retry_policy=ANY,
)
]
) )
])
@mark.asyncio @mark.asyncio
@patch("orchestrator.workflows.reports.workflow", new_callable=AsyncMock) @patch('orchestrator.workflows.reports.workflow', new_callable=AsyncMock)
async def test_run_no_data(workflow_mock, reports): async def test_run_no_data(workflow_mock, reports):
workflow_mock.execute_child_workflow.return_value = { workflow_mock.execute_child_workflow.return_value = {
'last_timestamp': '2023-01-01 12:00:00.000000', 'last_timestamp': '2023-01-01 12:00:00.000000',
'notification_package': [], 'notification_package': [],
'sending_configs': [] 'sending_configs': [],
} }
input_data = { input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600}
'schedule_name': 'test-schedule-name',
'notification_ttl': 300,
'sent_ttl': 600
}
await reports.run(input_data) await reports.run(input_data)
workflow_mock.execute_child_workflow.assert_has_calls([ workflow_mock.execute_child_workflow.assert_has_calls(
[
call( call(
'load_notification_package', 'load_notification_package',
{ {**input_data, 'metadata': metadata, 'base_data_filter': {}},
**input_data, )
'metadata': metadata, ]
'base_data_filter': {}
}
) )
])
workflow_mock.execute_local_activity_method.assert_not_called() workflow_mock.execute_local_activity_method.assert_not_called()
@mark.asyncio @mark.asyncio
@patch("orchestrator.workflows.reports.workflow", new_callable=AsyncMock) @patch('orchestrator.workflows.reports.workflow', new_callable=AsyncMock)
async def test_run_no_groups(workflow_mock, reports): async def test_run_no_groups(workflow_mock, reports):
workflow_mock.execute_local_activity_method.return_value = [] workflow_mock.execute_local_activity_method.return_value = []
input_data = { input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600}
'schedule_name': 'test-schedule-name',
'notification_ttl': 300,
'sent_ttl': 600
}
await reports.run(input_data) await reports.run(input_data)

99
validate.sh Executable file
View File

@@ -0,0 +1,99 @@
#!/bin/bash
# Model Manager Code Validation Script
# This script runs all code quality checks before committing or deploying
set -e # Exit on any error
# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color
echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║ Model Manager - Code Validation Suite ║${NC}"
echo -e "${BLUE}╚════════════════════════════════════════════════════════╝${NC}"
echo ""
# Check if virtual environment is activated
if [[ -z "${VIRTUAL_ENV}" ]] && [[ -z "${CONDA_DEFAULT_ENV}" ]]; then
echo -e "${YELLOW}⚠️ Warning: No virtual environment detected${NC}"
echo -e "${YELLOW} Consider activating your venv/conda environment${NC}"
echo ""
fi
# Function to run a validation step
run_step() {
local step_name=$1
local step_command=$2
echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
echo -e "${BLUE}▶ ${step_name}${NC}"
echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
if eval "$step_command"; then
echo -e "${GREEN}✅ ${step_name} - PASSED${NC}"
echo ""
return 0
else
echo -e "${RED}❌ ${step_name} - FAILED${NC}"
echo ""
return 1
fi
}
# Track failures
FAILED_STEPS=()
# Step 1: Code Formatting Check (Ruff)
if ! run_step "1. Code Formatting (Ruff)" "ruff format orchestrator/ tests/ && ruff format --check orchestrator/ tests/"; then
FAILED_STEPS+=("Code Formatting")
fi
# Step 2: Linting (Ruff)
if ! run_step "2. Code Linting (Ruff)" "ruff check --fix orchestrator/ tests/"; then
FAILED_STEPS+=("Linting")
fi
# Step 3: Type Checking (mypy)
if ! run_step "3. Type Checking (mypy)" "mypy orchestrator/"; then
FAILED_STEPS+=("Type Checking")
fi
# Step 4: Security Analysis (Bandit)
if ! run_step "4. Security Analysis (Bandit)" "bandit -r orchestrator/ -ll -q"; then
FAILED_STEPS+=("Security Analysis")
fi
# Step 5: Unit Tests (pytest)
if ! run_step "5. Unit Tests (pytest)" "pytest tests/ --cov=orchestrator --cov-report=term-missing --cov-report=xml --cov-report=html --cov-fail-under=80 -q"; then
FAILED_STEPS+=("Unit Tests")
fi
# Summary
echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║ Validation Summary ║${NC}"
echo -e "${BLUE}╚════════════════════════════════════════════════════════╝${NC}"
echo ""
if [ ${#FAILED_STEPS[@]} -eq 0 ]; then
echo -e "${GREEN}✅ All validation checks passed!${NC}"
echo -e "${GREEN} Your code is ready for commit/deployment.${NC}"
echo ""
exit 0
else
echo -e "${RED}❌ Validation failed for the following steps:${NC}"
for step in "${FAILED_STEPS[@]}"; do
echo -e "${RED} • ${step}${NC}"
done
echo ""
echo -e "${YELLOW}💡 Tips:${NC}"
echo -e "${YELLOW} • Run 'ruff format orchestrator/ tests/' to auto-fix formatting${NC}"
echo -e "${YELLOW} • Run 'ruff check --fix orchestrator/ tests/' to auto-fix linting issues${NC}"
echo -e "${YELLOW} • Review mypy errors and add type hints where needed${NC}"
echo -e "${YELLOW} • Check bandit warnings for security issues${NC}"
echo -e "${YELLOW} • Fix failing tests or improve test coverage${NC}"
echo ""
exit 1
fi

2
values.yaml
View File

@@ -151,7 +151,7 @@ env:
- name: GITHUB_REPO_URL - name: GITHUB_REPO_URL
value: "git@github.com:Aignosi/sientia-dataops-orchestrator_temporal.git" value: "git@github.com:Aignosi/sientia-dataops-orchestrator_temporal.git"
- name: GITHUB_BRANCH - name: GITHUB_BRANCH
value: "SIENTIAPDE-1222-ajustar-a-library-para-fazer-o-download-do-courier" value: "SIENTIAPDE-1231-ajustar-o-retreino-do-courier-no-laborious"
- name: PYTHON_APP - name: PYTHON_APP
value: "orchestrator.worker.worker" value: "orchestrator.worker.worker"