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

SIENTIAPDE-1231

Browse Source 
Refactor orchestrator activities and configuration files

- Removed unused Temporal timeout environment variables from values.yaml.
- Reorganized import statements in various activity files for better readability.
- Updated logging messages to use consistent formatting across activities.
- Enhanced test cases to ensure proper initialization and shutdown of orchestrator activities.
- Improved overall code structure and readability by applying consistent formatting and style adjustments.
This commit is contained in:
vitor-aignosi
2025-10-15 17:20:52 -03:00
parent 310cae4c2f
commit f3a6af5603
39 changed files with 4988 additions and 3826 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1399
coverage.xml Normal file
View File

File diff suppressed because it is too large Load Diff

139
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,16 +35,17 @@ class Activities( # Couchbase,
notification_handler (NotificationHandler): Notification management handler notification_handler (NotificationHandler): Notification management handler
""" """
def __init__(self, def __init__(
temporal_config: dict[str, Any], self,
# couchbase_config: dict[str, Any], temporal_config: dict[str, Any],
redis_config: dict[str, Any], # couchbase_config: dict[str, Any],
mongodb_config: dict[str, Any], redis_config: dict[str, Any],
email_config: dict[str, Any], mongodb_config: dict[str, Any],
postgres_config: dict[str, Any], email_config: dict[str, Any],
logger: Logger, postgres_config: dict[str, Any],
notification_handler: NotificationHandler): logger: Logger,
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,55 +53,64 @@ class Activities( # Couchbase,
# logger=logger, # logger=logger,
# notification_handler=notification_handler) # notification_handler=notification_handler)
TemporalManager.__init__(self, TemporalManager.__init__(
host=temporal_config['temporal_host'], self,
scouter_namespace=temporal_config['temporal_scouter_namespace'], host=temporal_config['temporal_host'],
laborious_namespace=temporal_config['temporal_laborious_namespace'], scouter_namespace=temporal_config['temporal_scouter_namespace'],
task_timeout=temporal_config['temporal_task_timeout_minutes'], laborious_namespace=temporal_config['temporal_laborious_namespace'],
run_timeout=temporal_config['temporal_run_timeout_minutes'], logger=logger,
execution_timeout=temporal_config['temporal_execution_timeout_minutes'], notification_handler=notification_handler,
logger=logger, )
notification_handler=notification_handler)
SlotManager.__init__(self, SlotManager.__init__(
host=redis_config['host'], self,
port=redis_config['port'], host=redis_config['host'],
username=redis_config['username'], port=redis_config['port'],
password=redis_config['password'], username=redis_config['username'],
logger=logger, password=redis_config['password'],
notification_handler=notification_handler) logger=logger,
notification_handler=notification_handler,
)
Formatters.__init__(self, Formatters.__init__(
scouter_namespace=temporal_config['temporal_scouter_namespace'], self,
laborious_namespace=temporal_config['temporal_laborious_namespace'], scouter_namespace=temporal_config['temporal_scouter_namespace'],
logger=logger, laborious_namespace=temporal_config['temporal_laborious_namespace'],
notification_handler=notification_handler) logger=logger,
notification_handler=notification_handler,
)
MongoDB.__init__(self, MongoDB.__init__(
connection_string=mongodb_config['connection_string'], self,
database_name=mongodb_config['database_name'], connection_string=mongodb_config['connection_string'],
ttl_index_seconds=mongodb_config['ttl_index_seconds'], database_name=mongodb_config['database_name'],
logger=logger, ttl_index_seconds=mongodb_config['ttl_index_seconds'],
notification_handler=notification_handler) logger=logger,
notification_handler=notification_handler,
)
Email.__init__(self, Email.__init__(
sender_email=email_config['sender_email'], self,
sender_password=email_config['sender_password'], sender_email=email_config['sender_email'],
smtp_server=email_config['smtp_server'], sender_password=email_config['sender_password'],
smtp_port=email_config['smtp_port'], smtp_server=email_config['smtp_server'],
logger=logger, smtp_port=email_config['smtp_port'],
notification_handler=notification_handler) logger=logger,
notification_handler=notification_handler,
)
Postgres.__init__(self, Postgres.__init__(
host=postgres_config['host'], self,
port=postgres_config['port'], host=postgres_config['host'],
user=postgres_config['user'], port=postgres_config['port'],
password=postgres_config['password'], user=postgres_config['user'],
dbname=postgres_config['dbname'], password=postgres_config['password'],
min_connections=postgres_config['min_connections'], dbname=postgres_config['dbname'],
max_connections=postgres_config['max_connections'], min_connections=postgres_config['min_connections'],
logger=logger, max_connections=postgres_config['max_connections'],
notification_handler=notification_handler) logger=logger,
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

354
orchestrator/activities/formatters.py
View File

@@ -2,19 +2,25 @@ 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):
@@ -41,16 +47,18 @@ class Formatters(BaseActivity):
notification_handler (NotificationHandler): Notification management handler notification_handler (NotificationHandler): Notification management handler
""" """
def __init__(self, def __init__(
scouter_namespace: str, self,
laborious_namespace: str, scouter_namespace: str,
logger: Logger, notification_handler: NotificationHandler): laborious_namespace: str,
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.
@@ -61,7 +69,7 @@ class Formatters(BaseActivity):
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:
@@ -73,46 +81,43 @@ class Formatters(BaseActivity):
- dict[str, Any]: The schedule config dictionary - dict[str, Any]: The schedule config dictionary
""" """
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 = {self.scouter_namespace: {}, self.laborious_namespace: {}}
self.scouter_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
@@ -130,9 +135,9 @@ class Formatters(BaseActivity):
- dict[str, Any]: The slot config dictionary - dict[str, Any]: The slot config dictionary
""" """
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']
@@ -154,43 +159,42 @@ class Formatters(BaseActivity):
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,9 +206,9 @@ 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']
@@ -219,16 +223,20 @@ 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 Compares the timestamps of the schedule and the current schedule to determine
which schedules need to be updated or created. which schedules need to be updated or created.
@@ -243,23 +251,22 @@ class Formatters(BaseActivity):
""" """
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,
@@ -276,51 +283,37 @@ class Formatters(BaseActivity):
- dict[str, Any]: The schedule config dictionary - dict[str, Any]: The schedule config dictionary
""" """
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 = {self.scouter_namespace: {}, self.laborious_namespace: {}}
self.scouter_namespace: {}, to_create = {self.scouter_namespace: {}, self.laborious_namespace: {}}
self.laborious_namespace: {} to_delete = {self.scouter_namespace: [], self.laborious_namespace: []}
}
to_create = {
self.scouter_namespace: {},
self.laborious_namespace: {}
}
to_delete = {
self.scouter_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,
@@ -337,9 +330,9 @@ class Formatters(BaseActivity):
- dict[str, Any]: The slot config dictionary - dict[str, Any]: The slot config dictionary
""" """
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,21 +342,18 @@ 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: str = None
) -> None:
""" """
Sends a success notification report. Sends a success notification report.
@@ -377,13 +367,14 @@ 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.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,9 +388,9 @@ 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: dict[str, Any]) -> tuple[list[str], dict[str, Any]]:
@@ -408,7 +399,7 @@ class Formatters(BaseActivity):
Args: Args:
input_data (dict[str, Any]): The input data containing schedule reports. input_data (dict[str, Any]): The input data containing schedule reports.
Each item should have 'namespace', 'schedule_name', 'success', 'message', Each item should have 'namespace', 'schedule_name', 'success', 'message',
and optionally 'attachment' fields. and optionally 'attachment' fields.
Returns: Returns:
@@ -416,14 +407,20 @@ class Formatters(BaseActivity):
- List of successful schedule keys in format "namespace/schedule_name" - List of successful schedule keys in format "namespace/schedule_name"
- Dictionary of error keys mapped to their error details - Dictionary of error keys mapped to their 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 = {
'message': value['message'], f'{value["namespace"]}/{value["schedule_name"]}': {
'attachment': value.get('attachment', None) 'message': value['message'],
'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
@@ -440,15 +437,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 +464,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 +496,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 +507,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 +543,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 +556,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,19 +574,19 @@ 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[str, 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 +600,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,7 +630,7 @@ 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']:
@@ -642,7 +638,7 @@ class Formatters(BaseActivity):
return DataFrame(list(data.values())).to_dict() return DataFrame(list(data.values())).to_dict()
@activity.defn(name="filter_notification_reports") @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 +674,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 +688,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

327
orchestrator/activities/mongo_db.py
View File

@@ -1,16 +1,18 @@
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:
@@ -28,10 +30,10 @@ def clear_mongo_id(docs: list) -> 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):
@@ -57,14 +59,18 @@ 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__(
logger: Logger, self,
notification_handler: NotificationHandler): connection_string: str,
database_name: str,
ttl_index_seconds: int,
logger: Logger,
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(self.connection_string, serverSelectionTimeoutMS=5000)
self.connection_string, serverSelectionTimeoutMS=5000)
self.client.server_info() # Trigger an exception if connection fails self.client.server_info() # Trigger an exception if connection fails
self.database = self.client[self.database_name] self.database = self.client[self.database_name]
@@ -72,11 +78,9 @@ class MongoDB(BaseActivity):
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):
""" """
@@ -84,11 +88,11 @@ class MongoDB(BaseActivity):
""" """
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):
""" """
@@ -109,13 +113,15 @@ class MongoDB(BaseActivity):
""" """
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.
@@ -128,33 +134,39 @@ class MongoDB(BaseActivity):
list[dict]: List of documents matching the query. list[dict]: List of documents matching the query.
""" """
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,19 +174,20 @@ 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.
@@ -186,40 +199,45 @@ class MongoDB(BaseActivity):
list[dict]: List of aggregated documents. list[dict]: List of aggregated documents.
""" """
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 +245,85 @@ 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 the timestamps of the pipelines in the MongoDB collection.
input_data: input_data:
- updated_pipelines (list): List of updated pipelines. - updated_pipelines (list): List of updated pipelines.
""" """
updated_pipelines = input_data.get("updated_pipelines", []) updated_pipelines = input_data.get('updated_pipelines', [])
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
date_now = now() date_now = now()
collection = self.database["orchestrated_schedules"] collection = self.database['orchestrated_schedules']
self.info("Updating pipelines timestamps...", metadata=metadata) 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. Create the timestamps of the pipelines in the MongoDB collection.
input_data: input_data:
- created_pipelines (list): List of created pipelines. - created_pipelines (list): List of created pipelines.
""" """
created_pipelines = input_data.get("created_pipelines", []) created_pipelines = input_data.get('created_pipelines', [])
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
collection = self.database["orchestrated_schedules"] collection = self.database['orchestrated_schedules']
self.info("Creating pipelines timestamps...", metadata=metadata) 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 +335,41 @@ 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 the timestamps of the pipelines in the MongoDB collection.
input_data: input_data:
- deleted_pipelines (list): List of deleted pipelines. - deleted_pipelines (list): List of deleted pipelines.
""" """
deleted_pipelines = input_data.get("deleted_pipelines", []) deleted_pipelines = input_data.get('deleted_pipelines', [])
metadata = input_data.get("metadata", {}) metadata = input_data.get('metadata', {})
collection = self.database["orchestrated_schedules"] collection = self.database['orchestrated_schedules']
self.info("Deleting pipelines timestamps...", metadata=metadata) 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 +378,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 +400,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 +427,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 +445,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 +490,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( self.info(f'Loaded {len(data)} documents from MongoDB', metadata=metadata)
f"Loaded {len(data)} documents from MongoDB",
metadata=metadata
)
self.debug( self.debug(f'Loaded data: {data}', metadata=metadata)
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)

308
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
@@ -30,69 +36,65 @@ class TemporalManager(BaseActivity):
Args: Args:
host (str): Temporal server host address host (str): Temporal server host address
scouter_namespace (str): Scouter workflow namespace scouter_namespace (str): Scouter workflow namespace
laborious_namespace (str): Laborious workflow namespace laborious_namespace (str): Laborious workflow namespace
logger (Logger): Application logger instance logger (Logger): Application logger instance
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__(
task_timeout: int, run_timeout: int, execution_timeout: int, 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.task_timeout = task_timeout
self.run_timeout = run_timeout
self.execution_timeout = execution_timeout
self.temporal_clients = {} self.temporal_clients = {}
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 for scouter and laborious workflows.
Creates client connections to both namespaces and stores them for later use. Creates client connections to both namespaces and stores them 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]) -> dict[str, Any]:
""" """
Normalize schedules. Removes schedules with no update time in mongo db collection "orchestrated_schedules". Normalize schedules. Removes schedules with no update time in mongo db collection "orchestrated_schedules".
input_data: input_data:
- orchestrated_schedules (dict[str, Any]): The orchestrated schedules to compare. - orchestrated_schedules (dict[str, Any]): The orchestrated schedules to compare.
""" """
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', {})
@@ -100,20 +102,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()
@@ -123,19 +125,18 @@ 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
@@ -150,47 +151,41 @@ class TemporalManager(BaseActivity):
""" """
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, SearchAttributePair(key=self.model_id_id_key, value=schedule['model_id']),
value=schedule['model_id'] SearchAttributePair(
), key=self.model_name_id_key, value=schedule['model_name']
SearchAttributePair( ),
key=self.model_name_id_key, SearchAttributePair(key=self.orchestrated_id_key, value='true'),
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 = schedule.get('execution_timeout_seconds', 300)
'execution_timeout_seconds', 300) task_timeout_seconds = schedule.get('task_timeout_seconds', 300)
task_timeout_seconds = schedule.get( self.debug(f'Creating schedule {schedule_name}:', metadata=metadata)
'task_timeout_seconds', 300)
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,
@@ -199,54 +194,57 @@ 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( execution_timeout=timedelta(seconds=execution_timeout_seconds),
seconds=execution_timeout_seconds), run_timeout=timedelta(seconds=execution_timeout_seconds),
run_timeout=timedelta( task_timeout=timedelta(seconds=task_timeout_seconds),
seconds=execution_timeout_seconds), typed_search_attributes=search_attributes,
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
@@ -261,27 +259,27 @@ class TemporalManager(BaseActivity):
""" """
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
@@ -310,33 +308,38 @@ 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
@@ -351,54 +354,59 @@ class TemporalManager(BaseActivity):
""" """
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

16
orchestrator/metrics.py
View File

@@ -6,19 +6,19 @@ orchestrator system including application health, email delivery,
and workflow execution 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", 'Number of emails sent',
[*CORE_LABELS, "email_group"], [*CORE_LABELS, 'email_group'],
) )

13
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'),
} }
@@ -61,9 +61,6 @@ def build_temporal_config():
'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'),
'temporal_task_timeout_minutes': int(getenv('TEMPORAL_TASK_TIMEOUT_MINUTES', '5')),
'temporal_run_timeout_minutes': int(getenv('TEMPORAL_RUN_TIMEOUT_MINUTES', '5')),
'temporal_execution_timeout_minutes': int(getenv('TEMPORAL_EXECUTION_TIMEOUT_MINUTES', '5'))
} }
@@ -81,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')),
} }
@@ -96,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')),
} }

12
orchestrator/utils/converters.py
View File

@@ -9,7 +9,7 @@ def parse_frequency(frequency: str) -> int:
Args: Args:
frequency (str): Frequency string with suffix: frequency (str): Frequency string with suffix:
- 's' for seconds (e.g., '30s') - 's' for seconds (e.g., '30s')
- 'm' for minutes (e.g., '5m') - 'm' for minutes (e.g., '5m')
- 'h' for hours (e.g., '2h') - 'h' for hours (e.g., '2h')
- 'd' for days (e.g., '1d') - 'd' for days (e.g., '1d')
@@ -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')

54
orchestrator/utils/email_builder.py
View File

@@ -1,8 +1,5 @@
import json
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 +21,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:
@@ -41,9 +38,9 @@ class EmailBuilder:
str: The rendered template with parameters replaced. str: The rendered template with parameters replaced.
""" """
# Criar um template Jinja2 # Criar um template Jinja2
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,21 +54,25 @@ 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], mail_type: str) -> str:
@@ -90,29 +91,26 @@ class EmailBuilder:
general_events = {} general_events = {}
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']: if model_name not in general_events[level]['models']:
general_events[level]['models'][model_name] = { general_events[level]['models'][model_name] = {
'model_name': model_name, 'model_name': model_name,
'events': [] 'events': [],
} }
general_events[level]['models'][model_name]['events'].append( general_events[level]['models'][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 )
))

166
orchestrator/utils/orchestrator_functions.py
View File

@@ -19,17 +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),
"execution_timeout_seconds": config.get('execution_timeout_seconds', 300),
"task_timeout_seconds": config.get('task_timeout_seconds', 300),
} }
@@ -48,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', []),
} }
@@ -79,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)
} }
@@ -120,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
@@ -138,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)
@@ -169,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']] = {}
@@ -177,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')
} }
@@ -241,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.
@@ -276,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

79
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_email_init, mock_postgres_init,
mock_formatters_init, mock_email_init,
mock_slot_manager_init, mock_formatters_init,
mock_temporal_manager_init, mock_slot_manager_init,
mock_mongodb_init): mock_temporal_manager_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,17 +117,17 @@ 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_postgres_shutdown, mock_mongodb_close,
mock_email_close, mock_postgres_shutdown,
mock_postgres_init, mock_email_close,
mock_email_init, mock_postgres_init,
mock_formatters_init, mock_email_init,
mock_slot_manager_init, mock_formatters_init,
mock_temporal_manager_init, mock_slot_manager_init,
mock_mongodb_init, mock_temporal_manager_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': {
'members': ['test3@test.com', 'test4@test.com'],
'notifications': [],
'html': 'test_html2',
}, },
"group_2": {
"members": ["test3@test.com", "test4@test.com"],
"notifications": [],
"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'),
] ]
) )

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

File diff suppressed because it is too large Load Diff

587
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": [ 'nested_list': [{'_id': 'fghij', 'item': 'value1'}, {'_id': 'klmno', 'item': 'value2'}],
{"_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'},
'collection_name': 'test_collection',
'last_data_timestamp': None,
'base_data_filter': {
'level': 'ERROR'
}
})
mongo_db.database.__getitem__.assert_called_once_with(
'test_collection')
collection.find.assert_called_once_with(
{ {
'level': 'ERROR' 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
}, 'collection_name': 'test_collection',
{"_id": 0} 'last_data_timestamp': None,
'base_data_filter': {'level': 'ERROR'},
}
) )
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'}, {
'collection_name': 'test_collection', 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'last_data_timestamp': '2023-01-01 12:00:00.000000+0000', 'collection_name': 'test_collection',
'base_data_filter': { 'last_data_timestamp': '2023-01-01 12:00:00.000000+0000',
'level': 'ERROR' 'base_data_filter': {'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'}, {
'collection_name': 'test_collection', 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'last_data_timestamp': '2023-01-01 12:00:00.000000+0000', 'collection_name': 'test_collection',
'base_data_filter': { 'last_data_timestamp': '2023-01-01 12:00:00.000000+0000',
'level': 'ERROR' 'base_data_filter': {'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,
) )

332
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',
b'slot:opc_tags:3',
]
slot_manager.get = MagicMock( slot_manager.get = MagicMock(side_effect=['value1', 'value2', None])
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:opc_tags:3',
]
slot_manager.get = MagicMock( slot_manager.get = MagicMock(side_effect=['value1', 'value2', None])
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'], {
'value': [25.5, 30.0], 'name': ['sensor1', 'sensor2'],
'timestamp': ['2023-01-01 12:00:00', '2023-01-01 12:00:01'] 'value': [25.5, 30.0],
}) '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'], {
'value': [25.5, 30.0], 'name': ['sensor1', 'sensor2'],
'timestamp': ['2023-01-01 12:00:00'] * 2 'value': [25.5, 30.0],
}).to_dict('records'), 'timestamp': ['2023-01-01 12:00:00'] * 2,
'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(
None, side_effect=[
(now() - timedelta(seconds=600) None,
).strftime(DATETIME_FORMAT_MS_WITH_TZ), (now() - timedelta(seconds=600)).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'], {
'schedule': ['test_schedule_1', 'test_schedule_2'], 'status': ['sent', 'error'],
'notification_id': ['test_notification_id_1', 'test_notification_id_2'] 'schedule': ['test_schedule_1', 'test_schedule_2'],
}).to_dict(), 'notification_id': ['test_notification_id_1', 'test_notification_id_2'],
'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
)

647
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,245 +71,235 @@ 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,
mock_schedule_spec, mock_schedule_spec,
mock_schedule_interval_spec, mock_schedule_interval_spec,
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( [
action=mock_schedule_action_start_workflow.return_value, call(
spec=mock_schedule_spec.return_value action=mock_schedule_action_start_workflow.return_value,
), spec=mock_schedule_spec.return_value,
call( ),
action=mock_schedule_action_start_workflow.return_value, call(
spec=mock_schedule_spec.return_value action=mock_schedule_action_start_workflow.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( [
"test-workflow", call(
input_data['schedules']['scouter']['test-schedule'], 'test-workflow',
id="test-schedule", input_data['schedules']['scouter']['test-schedule'],
task_queue="test-workflow-queue", id='test-schedule',
execution_timeout=ANY, task_queue='test-workflow-queue',
typed_search_attributes=mock_typed_search_attributes.return_value, execution_timeout=timedelta(seconds=100),
), run_timeout=timedelta(seconds=100),
call( task_timeout=timedelta(seconds=100),
"test-workflow", typed_search_attributes=mock_typed_search_attributes.return_value,
input_data['schedules']['scouter']['test-schedule-invalid-frequency'], ),
id="test-schedule-invalid-frequency", call(
task_queue="test-workflow-queue", 'test-workflow',
execution_timeout=ANY, input_data['schedules']['scouter']['test-schedule-invalid-frequency'],
typed_search_attributes=mock_typed_search_attributes.return_value, id='test-schedule-invalid-frequency',
), task_queue='test-workflow-queue',
call( execution_timeout=timedelta(seconds=400),
"test-workflow", run_timeout=timedelta(seconds=400),
input_data['schedules']['laborious']['test-schedule-laborious'], task_timeout=timedelta(seconds=400),
id="test-schedule-laborious", typed_search_attributes=mock_typed_search_attributes.return_value,
task_queue="test-workflow-queue", ),
execution_timeout=ANY, call(
typed_search_attributes=mock_typed_search_attributes.return_value, 'test-workflow',
) input_data['schedules']['laborious']['test-schedule-laborious'],
]) id='test-schedule-laborious',
task_queue='test-workflow-queue',
execution_timeout=timedelta(seconds=500),
run_timeout=timedelta(seconds=500),
task_timeout=timedelta(seconds=500),
typed_search_attributes=mock_typed_search_attributes.return_value,
),
]
)
mock_schedule_spec.assert_has_calls([ mock_schedule_spec.assert_has_calls(
call( [
intervals=[ call(intervals=[mock_schedule_interval_spec.return_value]),
mock_schedule_interval_spec.return_value call(intervals=[mock_schedule_interval_spec.return_value]),
] ]
), )
call(
intervals=[
mock_schedule_interval_spec.return_value
]
)
])
mock_schedule_interval_spec.assert_has_calls([ mock_schedule_interval_spec.assert_has_calls(
call( [call(every=timedelta(seconds=60)), call(every=timedelta(seconds=120))]
every=timedelta(seconds=60) )
),
call(
every=timedelta(seconds=120)
)
])
mock_parse_frequency.assert_has_calls([ mock_parse_frequency.assert_has_calls([call('1m'), call('10y'), call('2m')])
call("1m"),
call("10y"),
call("2m")
])
mock_typed_search_attributes.assert_has_calls([ mock_typed_search_attributes.assert_has_calls(
call([ [
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.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.return_value call(
]), [
call([ mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value, mock_search_attribute_pair.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_search_attribute_pair.assert_has_calls([ mock_search_attribute_pair.assert_has_calls(
call( [
key=temporal_manager.model_id_id_key, call(key=temporal_manager.model_id_id_key, value=1),
value=1 call(key=temporal_manager.model_name_id_key, value='test-model-name'),
), call(key=temporal_manager.orchestrated_id_key, value='true'),
call( call(key=temporal_manager.model_id_id_key, value=2),
key=temporal_manager.model_name_id_key, call(key=temporal_manager.model_name_id_key, value='test-model-name'),
value="test-model-name" call(key=temporal_manager.orchestrated_id_key, value='true'),
), ]
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": {
"frequency": "1m",
"data": {"test": "test"}
}
}, },
"laborious": { 'laborious': {'test-schedule-laborious': {'frequency': '2m', 'data': {'test': 'test'}}},
"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}'
)

39
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,
build_couchbase_config, from orchestrator.utils.connectors_config import (
build_mongodb_config, build_couchbase_config,
build_temporal_config, build_email_config,
build_email_config, build_mongodb_config,
build_postgres_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')

147
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": {
"test_addr": {
"data_type": "float"
}
}
} }
}, },
"input_filters": { 'input_filters': {'test_filter_name': {'policy': 'test_policy', 'config': {}}},
"test_filter_name": { 'mlflow_transform_filters': {'test_filter_name': {'policy': 'test_policy', 'config': {}}},
"policy": "test_policy", 'mlflow_predict_filters': {'test_filter_name': {'policy': 'test_policy', 'config': {}}},
"config": {} 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
} 'datetime_columns': ['timestamp'],
}, 'predictions_storage_policy': 'erl:1',
"mlflow_transform_filters": { 'execution_timeout_seconds': 300,
"test_filter_name": { 'task_timeout_seconds': 300,
"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'

170
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,123 +59,113 @@ 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( [
Activities.get_last_data_timestamp, call(
{ Activities.get_last_data_timestamp,
**input_data['metadata'], {**input_data['metadata'], 'mail_type': 'test_mail_type'},
'mail_type': 'test_mail_type' start_to_close_timeout=ANY,
}, retry_policy=ANY,
start_to_close_timeout=ANY, )
retry_policy=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, call(
{ Activities.find_documents_in_mongodb,
**input_data['metadata'], {
'query': { **input_data['metadata'],
'collection': 'receiver_groups', 'query': {'collection': 'receiver_groups', 'filters': {'active': True}},
'filters': { },
'active': True start_to_close_timeout=ANY,
} retry_policy=ANY,
} )
}, ]
start_to_close_timeout=ANY, )
retry_policy=ANY
)
])
workflow_mock.start_local_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
call( [
Activities.load_latest_data, call(
{ Activities.load_latest_data,
**input_data['metadata'], {
'collection_name': 'notification_queue', **input_data['metadata'],
'last_data_timestamp': '2023-01-01 12:00:00', 'collection_name': 'notification_queue',
'base_data_filter': { 'last_data_timestamp': '2023-01-01 12:00:00',
'level': 'ERROR' 'base_data_filter': {'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( [
Activities.put_last_data_timestamp, call(
{ Activities.put_last_data_timestamp,
**input_data['metadata'], {
'data': [ **input_data['metadata'],
{ 'data': [
'id': '1', {
} 'id': '1',
], }
'mail_type': 'test_mail_type' ],
}, 'mail_type': 'test_mail_type',
start_to_close_timeout=ANY, },
retry_policy=ANY start_to_close_timeout=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()

182
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,69 +37,77 @@ 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( [
Activities.build_email_html, call(
{ Activities.build_email_html,
**metadata, {
'receiver_groups': input_data['notification_package'], **metadata,
'mail_type': input_data['mail_type'] 'receiver_groups': input_data['notification_package'],
}, 'mail_type': input_data['mail_type'],
schedule_to_close_timeout=ANY, },
retry_policy=ANY schedule_to_close_timeout=ANY,
) retry_policy=ANY,
]) )
workflow_mock.execute_local_activity_method.assert_has_calls([ ]
call( )
Activities.format_log_report, workflow_mock.execute_local_activity_method.assert_has_calls(
{ [
**metadata, call(
'receiver_groups': workflow_mock.execute_activity_method.return_value, Activities.format_log_report,
'mail_type': input_data['mail_type'] {
}, **metadata,
schedule_to_close_timeout=ANY, 'receiver_groups': workflow_mock.execute_activity_method.return_value,
retry_policy=ANY 'mail_type': input_data['mail_type'],
) },
]) schedule_to_close_timeout=ANY,
retry_policy=ANY,
workflow_mock.execute_activity_method.assert_has_calls([ )
call( ]
Activities.send_email,
{
**metadata,
'receiver_groups': workflow_mock.execute_local_activity_method.return_value,
'mail_type': input_data['mail_type']
},
schedule_to_close_timeout=ANY,
retry_policy=ANY
)]
) )
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_activity_method.assert_has_calls(
call( [
Activities.export_data_to_postgres, call(
{ Activities.send_email,
**metadata, {
'schema': input_data['schema'], **metadata,
'table_name': input_data['table_name'], 'receiver_groups': workflow_mock.execute_local_activity_method.return_value,
'data': workflow_mock.execute_local_activity_method.return_value, 'mail_type': input_data['mail_type'],
'timestamp_conversion': { },
'column': 'timestamp', schedule_to_close_timeout=ANY,
'format': DATETIME_FORMAT_MS_WITH_TZ retry_policy=ANY,
} )
}, ]
schedule_to_close_timeout=ANY, )
retry_policy=ANY
) workflow_mock.execute_activity_method.assert_has_calls(
]) [
call(
Activities.export_data_to_postgres,
{
**metadata,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': workflow_mock.execute_local_activity_method.return_value,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_MS_WITH_TZ,
},
},
schedule_to_close_timeout=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( [
Activities.build_email_html, call(
{ Activities.build_email_html,
**metadata, {
'receiver_groups': input_data['notification_package'], **metadata,
'mail_type': input_data['mail_type'] 'receiver_groups': input_data['notification_package'],
}, 'mail_type': input_data['mail_type'],
schedule_to_close_timeout=ANY, },
retry_policy=ANY schedule_to_close_timeout=ANY,
) retry_policy=ANY,
]) )
]
)
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_activity_method.assert_has_calls(
call( [
Activities.send_email, call(
{ Activities.send_email,
**metadata, {
'receiver_groups': workflow_mock.execute_local_activity_method.return_value, **metadata,
'mail_type': input_data['mail_type'] 'receiver_groups': workflow_mock.execute_local_activity_method.return_value,
}, 'mail_type': input_data['mail_type'],
schedule_to_close_timeout=ANY, },
retry_policy=ANY schedule_to_close_timeout=ANY,
)] retry_policy=ANY,
)
]
) )
assert workflow_mock.execute_activity_method.call_count == 1 assert workflow_mock.execute_activity_method.call_count == 1

169
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,113 +22,103 @@ 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( [
'load_notification_package', call(
{ 'load_notification_package',
**input_data, {**input_data, 'metadata': metadata, 'base_data_filter': {'level': 'ERROR'}},
'metadata': metadata, )
'base_data_filter': { ]
'level': 'ERROR' )
}
}
)
])
workflow_mock.execute_child_workflow.assert_has_calls([ workflow_mock.execute_child_workflow.assert_has_calls(
call( [
'process_notifications', call(
{ 'process_notifications',
'metadata': metadata, {
'mail_type': 'Alerts', 'metadata': metadata,
'notification_package': workflow_mock.execute_local_activity_method.return_value, 'mail_type': 'Alerts',
'schema': 'sientia_data', 'notification_package': workflow_mock.execute_local_activity_method.return_value,
'table_name': 'log_report' 'schema': 'sientia_data',
} 'table_name': 'log_report',
) },
]) )
]
)
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls(
call( [
Activities.filter_notification_alerts, call(
{ Activities.filter_notification_alerts,
**metadata, {
'notification_package': workflow_mock.execute_child_workflow.return_value['notification_package'], **metadata,
'sending_configs': workflow_mock.execute_child_workflow.return_value['sending_configs'], 'notification_package': workflow_mock.execute_child_workflow.return_value[
'notification_ttl': input_data['notification_ttl'] 'notification_package'
}, ],
schedule_to_close_timeout=ANY, 'sending_configs': workflow_mock.execute_child_workflow.return_value[
retry_policy=ANY 'sending_configs'
) ],
]) 'notification_ttl': input_data['notification_ttl'],
},
schedule_to_close_timeout=ANY,
retry_policy=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls([ workflow_mock.execute_activity_method.assert_has_calls(
call( [
Activities.store_notification_cache, call(
{ Activities.store_notification_cache,
**metadata, {
'log_report': workflow_mock.execute_child_workflow.return_value, **metadata,
'sent_ttl': input_data['sent_ttl'] 'log_report': workflow_mock.execute_child_workflow.return_value,
}, 'sent_ttl': input_data['sent_ttl'],
schedule_to_close_timeout=ANY, },
retry_policy=ANY schedule_to_close_timeout=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( [
'load_notification_package', call(
{ 'load_notification_package',
**input_data, {**input_data, 'metadata': metadata, 'base_data_filter': {'level': 'ERROR'}},
'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)

565
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,290 +22,327 @@ 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( [
Activities.aggregate_documents_in_mongodb, call(
{ Activities.aggregate_documents_in_mongodb,
**metadata, {
"query": input_data["pipelines_query"], **metadata,
"timestamp_fields": ["updated_at"] 'query': input_data['pipelines_query'],
}, 'timestamp_fields': ['updated_at'],
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
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(
Activities.find_documents_in_mongodb,
{
**metadata,
"query": {
"collection": "orchestrated_schedules"
}, },
"timestamp_fields": ["updated_at"] retry_policy=ANY,
}, start_to_close_timeout=ANY,
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( [
Activities.load_opc_slots, call(
{ Activities.find_documents_in_mongodb,
**metadata {**metadata, 'query': input_data['opc_servers_query']},
}, 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.load_active_ingestors, call(
{ Activities.find_documents_in_mongodb,
**metadata {
}, **metadata,
retry_policy=ANY, 'query': {'collection': 'orchestrated_schedules'},
start_to_close_timeout=ANY '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( [
Activities.format_schedule_config, call(
{ Activities.load_opc_slots,
**metadata, {**metadata},
'schedule_config': workflow_mock.start_local_activity_method.return_value retry_policy=ANY,
}, start_to_close_timeout=ANY,
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( [
Activities.process_schedules, call(
{ Activities.load_active_ingestors,
**metadata, {**metadata},
'pipelines': workflow_mock.start_local_activity_method.return_value retry_policy=ANY,
}, start_to_close_timeout=ANY,
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( [
Activities.process_slots, call(
{ Activities.format_schedule_config,
**metadata, {
'opc_servers': workflow_mock.start_local_activity_method.return_value, **metadata,
'active_ingestors': workflow_mock.start_local_activity_method.return_value, 'schedule_config': 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( [
Activities.create_schedule_config, call(
{ Activities.process_schedules,
**metadata, {**metadata, 'pipelines': workflow_mock.start_local_activity_method.return_value},
'current_schedule_config': workflow_mock.start_local_activity_method.return_value, retry_policy=ANY,
'schedule_config': workflow_mock.start_local_activity_method.return_value start_to_close_timeout=ANY,
}, )
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( [
Activities.create_slot_config, call(
{ Activities.process_slots,
**metadata, {
'current_slot_config': workflow_mock.start_local_activity_method.return_value, **metadata,
'slot_config': workflow_mock.start_local_activity_method.return_value 'opc_servers': workflow_mock.start_local_activity_method.return_value,
}, 'active_ingestors': workflow_mock.start_local_activity_method.return_value,
retry_policy=ANY, 'pipelines': workflow_mock.start_local_activity_method.return_value,
start_to_close_timeout=ANY },
) retry_policy=ANY,
]) start_to_close_timeout=ANY,
)
]
)
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
call( [
Activities.normalize_schedules, call(
{ Activities.create_schedule_config,
**metadata, {
'orchestrated_schedules': workflow_mock.start_local_activity_method.return_value **metadata,
}, 'current_schedule_config': workflow_mock.start_local_activity_method.return_value,
retry_policy=ANY, 'schedule_config': workflow_mock.start_local_activity_method.return_value,
start_to_close_timeout=ANY },
) retry_policy=ANY,
]) start_to_close_timeout=ANY,
)
]
)
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_local_activity_method.assert_has_calls(
call( [
Activities.create_collection_with_ttl_index, call(
{ Activities.create_slot_config,
**metadata, {
'pipelines': workflow_mock.start_local_activity_method.return_value['scouter'] **metadata,
}, 'current_slot_config': workflow_mock.start_local_activity_method.return_value,
retry_policy=ANY, 'slot_config': workflow_mock.start_local_activity_method.return_value,
start_to_close_timeout=ANY },
) retry_policy=ANY,
]) start_to_close_timeout=ANY,
)
]
)
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
call( [
Activities.delete_slots, call(
{ Activities.normalize_schedules,
**metadata, {
'to_delete': **metadata,
workflow_mock.start_local_activity_method.return_value['to_delete'] '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( [
Activities.update_slots, call(
{ Activities.create_collection_with_ttl_index,
**metadata, {
'to_insert': **metadata,
workflow_mock.start_local_activity_method.return_value['to_insert'] '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( [
Activities.delete_schedules, call(
{ Activities.delete_slots,
**metadata, {
'schedules': **metadata,
workflow_mock.start_local_activity_method.return_value['to_delete'] 'to_delete': workflow_mock.start_local_activity_method.return_value[
}, 'to_delete'
retry_policy=ANY, ],
start_to_close_timeout=ANY },
) retry_policy=ANY,
]) start_to_close_timeout=ANY,
)
]
)
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
call( [
Activities.create_schedules, call(
{ Activities.update_slots,
**metadata, {
'schedules': **metadata,
workflow_mock.start_local_activity_method.return_value['to_create'] 'to_insert': workflow_mock.start_local_activity_method.return_value[
}, 'to_insert'
retry_policy=ANY, ],
start_to_close_timeout=ANY },
) retry_policy=ANY,
]) start_to_close_timeout=ANY,
)
]
)
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
call( [
Activities.update_schedules, call(
{ Activities.delete_schedules,
**metadata, {
'schedules': **metadata,
workflow_mock.start_local_activity_method.return_value['to_update'] 'schedules': workflow_mock.start_local_activity_method.return_value[
}, 'to_delete'
retry_policy=ANY, ],
start_to_close_timeout=ANY },
) retry_policy=ANY,
]) start_to_close_timeout=ANY,
)
]
)
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
call( [
Activities.report_schedule_orchestration, call(
{ Activities.create_schedules,
**metadata, {
'created_schedules': workflow_mock.start_activity_method.return_value, **metadata,
'updated_schedules': workflow_mock.start_activity_method.return_value, 'schedules': workflow_mock.start_local_activity_method.return_value[
'deleted_schedules': workflow_mock.start_activity_method.return_value, 'to_create'
}, ],
retry_policy=ANY, },
start_to_close_timeout=ANY retry_policy=ANY,
) start_to_close_timeout=ANY,
]) )
]
)
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
call( [
Activities.report_slot_orchestration, call(
{ Activities.update_schedules,
**metadata, {
'inserted_slots': workflow_mock.start_activity_method.return_value, **metadata,
'deleted_slots': workflow_mock.start_activity_method.return_value, 'schedules': workflow_mock.start_local_activity_method.return_value[
}, 'to_update'
retry_policy=ANY, ],
start_to_close_timeout=ANY },
) retry_policy=ANY,
]) start_to_close_timeout=ANY,
)
]
)
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
call( [
Activities.update_pipelines_timestamps, call(
{ Activities.report_schedule_orchestration,
**metadata, {
'updated_pipelines': workflow_mock.start_activity_method.return_value, **metadata,
}, 'created_schedules': workflow_mock.start_activity_method.return_value,
retry_policy=ANY, 'updated_schedules': workflow_mock.start_activity_method.return_value,
start_to_close_timeout=ANY 'deleted_schedules': workflow_mock.start_activity_method.return_value,
) },
]) retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
call( [
Activities.delete_pipelines_timestamps, call(
{ Activities.report_slot_orchestration,
**metadata, {
'deleted_pipelines': workflow_mock.start_activity_method.return_value, **metadata,
}, 'inserted_slots': workflow_mock.start_activity_method.return_value,
retry_policy=ANY, 'deleted_slots': workflow_mock.start_activity_method.return_value,
start_to_close_timeout=ANY },
) retry_policy=ANY,
]) start_to_close_timeout=ANY,
)
]
)
workflow_mock.start_activity_method.assert_has_calls([ workflow_mock.start_activity_method.assert_has_calls(
call( [
Activities.create_pipelines_timestamps, call(
{ Activities.update_pipelines_timestamps,
**metadata, {
'created_pipelines': workflow_mock.start_activity_method.return_value, **metadata,
}, 'updated_pipelines': workflow_mock.start_activity_method.return_value,
retry_policy=ANY, },
start_to_close_timeout=ANY retry_policy=ANY,
) start_to_close_timeout=ANY,
]) )
]
)
workflow_mock.start_activity_method.assert_has_calls(
[
call(
Activities.delete_pipelines_timestamps,
{
**metadata,
'deleted_pipelines': workflow_mock.start_activity_method.return_value,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.start_activity_method.assert_has_calls(
[
call(
Activities.create_pipelines_timestamps,
{
**metadata,
'created_pipelines': workflow_mock.start_activity_method.return_value,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)

124
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,95 +22,87 @@ 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( [
'load_notification_package', call(
{ 'load_notification_package',
**input_data, {**input_data, 'metadata': metadata, 'base_data_filter': {}},
'metadata': metadata, )
'base_data_filter': {} ]
} )
)
])
workflow_mock.execute_child_workflow.assert_has_calls([ workflow_mock.execute_child_workflow.assert_has_calls(
call( [
'process_notifications', call(
{ 'process_notifications',
'metadata': metadata, {
'mail_type': 'Reports', 'metadata': metadata,
'notification_package': workflow_mock.execute_local_activity_method.return_value, 'mail_type': 'Reports',
'schema': 'sientia_data', 'notification_package': workflow_mock.execute_local_activity_method.return_value,
'table_name': 'log_report' 'schema': 'sientia_data',
} 'table_name': 'log_report',
) },
]) )
]
)
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls(
call( [
Activities.filter_notification_reports, call(
{ Activities.filter_notification_reports,
**metadata, {
'notification_package': workflow_mock.execute_child_workflow.return_value['notification_package'], **metadata,
'sending_configs': workflow_mock.execute_child_workflow.return_value['sending_configs'] 'notification_package': workflow_mock.execute_child_workflow.return_value[
}, 'notification_package'
schedule_to_close_timeout=ANY, ],
retry_policy=ANY 'sending_configs': workflow_mock.execute_child_workflow.return_value[
) 'sending_configs'
]) ],
},
schedule_to_close_timeout=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( [
'load_notification_package', call(
{ 'load_notification_package',
**input_data, {**input_data, 'metadata': metadata, 'base_data_filter': {}},
'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

6
values.yaml
View File

@@ -235,12 +235,6 @@ env:
value: "scouter" value: "scouter"
- name: TEMPORAL_LABORIOUS_NAMESPACE - name: TEMPORAL_LABORIOUS_NAMESPACE
value: "laborious" value: "laborious"
- name: TEMPORAL_TASK_TIMEOUT_MINUTES
value: "10"
- name: TEMPORAL_RUN_TIMEOUT_MINUTES
value: "20"
- name: TEMPORAL_EXECUTION_TIMEOUT_MINUTES
value: "20"
ssh: ssh:
enabled: true enabled: true