From f3a6af5603bc916f2b715ebbe56877826c259cca Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 15 Oct 2025 17:20:52 -0300 Subject: [PATCH] SIENTIAPDE-1231 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. --- coverage.xml | 1399 +++++++++++++++++ orchestrator/activities/activities.py | 139 +- orchestrator/activities/couchbase.py | 53 +- orchestrator/activities/email.py | 117 +- orchestrator/activities/formatters.py | 354 +++-- orchestrator/activities/mongo_db.py | 327 ++-- orchestrator/activities/slot_manager.py | 206 ++- orchestrator/activities/temporal_manager.py | 308 ++-- orchestrator/metrics.py | 16 +- orchestrator/utils/connectors_config.py | 13 +- orchestrator/utils/converters.py | 12 +- orchestrator/utils/email_builder.py | 54 +- orchestrator/utils/orchestrator_functions.py | 166 +- orchestrator/worker/worker.py | 89 +- orchestrator/workflows/alerts.py | 37 +- orchestrator/workflows/orchestrator.py | 144 +- orchestrator/workflows/reports.py | 25 +- .../subworkflows/load_notification_package.py | 45 +- .../subworkflows/process_notifications.py | 46 +- pyproject.toml | 160 ++ requirements-dev.txt | 19 + .../activities/test_activities.py | 79 +- .../orchestrator/activities/test_couchbase.py | 55 +- tests/orchestrator/activities/test_email.py | 243 ++- .../activities/test_formatters.py | 1134 +++++++------ .../orchestrator/activities/test_mongo_db.py | 587 +++---- .../activities/test_slot_manager.py | 332 ++-- .../activities/test_temporal_manager.py | 647 ++++---- .../utils/test_connectors_config.py | 39 +- tests/orchestrator/utils/test_converters.py | 14 +- .../orchestrator/utils/test_email_builder.py | 147 +- .../utils/test_orchestrator_functions.py | 493 +++--- .../test_load_notification_package.py | 170 +- .../test_process_notifications.py | 182 ++- tests/orchestrator/workflows/test_alerts.py | 169 +- .../workflows/test_orchestrator.py | 565 +++---- tests/orchestrator/workflows/test_reports.py | 124 +- validate.sh | 99 ++ values.yaml | 6 - 39 files changed, 4988 insertions(+), 3826 deletions(-) create mode 100644 coverage.xml create mode 100644 pyproject.toml create mode 100644 requirements-dev.txt create mode 100755 validate.sh diff --git a/coverage.xml b/coverage.xml new file mode 100644 index 0000000..dcf2359 --- /dev/null +++ b/coverage.xml @@ -0,0 +1,1399 @@ + + + + + + /home/grezewave/Documents/projects/sientia/sientia-dataops-orchestrator_temporal/orchestrator + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/orchestrator/activities/activities.py b/orchestrator/activities/activities.py index 9da4cd8..3385f39 100644 --- a/orchestrator/activities/activities.py +++ b/orchestrator/activities/activities.py @@ -1,21 +1,22 @@ from temporalio import workflow - 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 sientia_do.temporal.activities.postgres import Postgres + from typing import Any + 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, - TemporalManager, SlotManager, Formatters, MongoDB, Email, - Postgres): + TemporalManager, SlotManager, Formatters, MongoDB, Email, Postgres +): """ Central activities orchestrator for Temporal workflow operations. @@ -34,16 +35,17 @@ class Activities( # Couchbase, notification_handler (NotificationHandler): Notification management handler """ - def __init__(self, - temporal_config: dict[str, Any], - # couchbase_config: dict[str, Any], - redis_config: dict[str, Any], - mongodb_config: dict[str, Any], - email_config: dict[str, Any], - postgres_config: dict[str, Any], - logger: Logger, - notification_handler: NotificationHandler): - + def __init__( + self, + temporal_config: dict[str, Any], + # couchbase_config: dict[str, Any], + redis_config: dict[str, Any], + mongodb_config: dict[str, Any], + email_config: dict[str, Any], + postgres_config: dict[str, Any], + logger: Logger, + notification_handler: NotificationHandler, + ): # Initialize parent classes # Couchbase.__init__(self, connection_string=couchbase_config['connection_string'], # username=couchbase_config['username'], @@ -51,55 +53,64 @@ class Activities( # Couchbase, # logger=logger, # notification_handler=notification_handler) - TemporalManager.__init__(self, - host=temporal_config['temporal_host'], - scouter_namespace=temporal_config['temporal_scouter_namespace'], - laborious_namespace=temporal_config['temporal_laborious_namespace'], - task_timeout=temporal_config['temporal_task_timeout_minutes'], - run_timeout=temporal_config['temporal_run_timeout_minutes'], - execution_timeout=temporal_config['temporal_execution_timeout_minutes'], - logger=logger, - notification_handler=notification_handler) + TemporalManager.__init__( + self, + host=temporal_config['temporal_host'], + scouter_namespace=temporal_config['temporal_scouter_namespace'], + laborious_namespace=temporal_config['temporal_laborious_namespace'], + logger=logger, + notification_handler=notification_handler, + ) - SlotManager.__init__(self, - host=redis_config['host'], - port=redis_config['port'], - username=redis_config['username'], - password=redis_config['password'], - logger=logger, - notification_handler=notification_handler) + SlotManager.__init__( + self, + host=redis_config['host'], + port=redis_config['port'], + username=redis_config['username'], + password=redis_config['password'], + logger=logger, + notification_handler=notification_handler, + ) - Formatters.__init__(self, - scouter_namespace=temporal_config['temporal_scouter_namespace'], - laborious_namespace=temporal_config['temporal_laborious_namespace'], - logger=logger, - notification_handler=notification_handler) + Formatters.__init__( + self, + scouter_namespace=temporal_config['temporal_scouter_namespace'], + laborious_namespace=temporal_config['temporal_laborious_namespace'], + logger=logger, + notification_handler=notification_handler, + ) - MongoDB.__init__(self, - connection_string=mongodb_config['connection_string'], - database_name=mongodb_config['database_name'], - ttl_index_seconds=mongodb_config['ttl_index_seconds'], - logger=logger, - notification_handler=notification_handler) + MongoDB.__init__( + self, + connection_string=mongodb_config['connection_string'], + database_name=mongodb_config['database_name'], + ttl_index_seconds=mongodb_config['ttl_index_seconds'], + logger=logger, + notification_handler=notification_handler, + ) - Email.__init__(self, - sender_email=email_config['sender_email'], - sender_password=email_config['sender_password'], - smtp_server=email_config['smtp_server'], - smtp_port=email_config['smtp_port'], - logger=logger, - notification_handler=notification_handler) + Email.__init__( + self, + sender_email=email_config['sender_email'], + sender_password=email_config['sender_password'], + smtp_server=email_config['smtp_server'], + smtp_port=email_config['smtp_port'], + logger=logger, + notification_handler=notification_handler, + ) - Postgres.__init__(self, - host=postgres_config['host'], - port=postgres_config['port'], - user=postgres_config['user'], - password=postgres_config['password'], - dbname=postgres_config['dbname'], - min_connections=postgres_config['min_connections'], - max_connections=postgres_config['max_connections'], - logger=logger, - notification_handler=notification_handler) + Postgres.__init__( + self, + host=postgres_config['host'], + port=postgres_config['port'], + user=postgres_config['user'], + password=postgres_config['password'], + dbname=postgres_config['dbname'], + min_connections=postgres_config['min_connections'], + max_connections=postgres_config['max_connections'], + logger=logger, + notification_handler=notification_handler, + ) def shutdown(self): """ diff --git a/orchestrator/activities/couchbase.py b/orchestrator/activities/couchbase.py index f7a8488..00c17e5 100644 --- a/orchestrator/activities/couchbase.py +++ b/orchestrator/activities/couchbase.py @@ -1,11 +1,12 @@ from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): - from typing import Any - from logging import Logger - from datetime import timedelta import json import traceback + from datetime import timedelta + from logging import Logger + from typing import Any + from couchbase.auth import PasswordAuthenticator from couchbase.cluster import Cluster from couchbase.options import ClusterOptions @@ -34,44 +35,43 @@ class Couchbase(BaseActivity): but maintained for potential future use. """ - def __init__(self, connection_string: str, username: str, - password: str, logger: Logger, - notification_handler: NotificationHandler): - + def __init__( + self, + connection_string: str, + username: str, + password: str, + logger: Logger, + notification_handler: NotificationHandler, + ): self.connection_string = connection_string self.username = username self.password = password - logger.info("Initializing Couchbase connection...") + logger.info('Initializing Couchbase connection...') self.cluster = Cluster( connection_string, ClusterOptions( - authenticator=PasswordAuthenticator( - username=username, - password=password - ) - ) + authenticator=PasswordAuthenticator(username=username, password=password) + ), ) - logger.info("Awaiting Couchbase connection...") + logger.info('Awaiting Couchbase connection...') self.cluster.wait_until_ready(timeout=timedelta(seconds=10)) - logger.info("Couchbase connection ready") + logger.info('Couchbase connection ready') - BaseActivity.__init__(self, - logger=logger, - notification_handler=notification_handler) + BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler) def shutdown(self): try: self.cluster.close() 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): 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]]: """ Load a query from couchbase @@ -84,16 +84,16 @@ class Couchbase(BaseActivity): """ query = input_data['query'] - self.logger.info(f"Executing couchbase query: {query}") + self.logger.info(f'Executing couchbase query: {query}') try: result = self.cluster.query(query) except Exception as e: trace = traceback.format_exc() self.notification_handler.build_and_send_notification( - notification_id="COUCHBASE_LOAD_QUERY_ERROR", - message=f"Failed to execute couchbase query: {e}", - block="load_query_from_couchbase", + notification_id='COUCHBASE_LOAD_QUERY_ERROR', + message=f'Failed to execute couchbase query: {e}', + block='load_query_from_couchbase', level=NotificationLevel.ERROR, attachment_content=trace, ) @@ -106,8 +106,7 @@ class Couchbase(BaseActivity): for row in result.rows(): rows.append(row) - self.logger.info("Fetched %d rows from couchbase", len(rows)) - self.logger.debug("Rows: \n %s", - json.dumps(rows, indent=4, sort_keys=True)) + self.logger.info('Fetched %d rows from couchbase', len(rows)) + self.logger.debug('Rows: \n %s', json.dumps(rows, indent=4, sort_keys=True)) return rows diff --git a/orchestrator/activities/email.py b/orchestrator/activities/email.py index 5a61207..04e715f 100644 --- a/orchestrator/activities/email.py +++ b/orchestrator/activities/email.py @@ -1,21 +1,23 @@ - from smtplib import SMTPServerDisconnected -from temporalio import workflow, activity + +from temporalio import activity, workflow from orchestrator import metrics with workflow.unsafe.imports_passed_through(): - import traceback import smtplib - from typing import Any - from sientia_do.temporal.activities.base import BaseActivity - from sientia_do.observability.logger import Logger - from sientia_do.notifications.handlers import NotificationHandler - from orchestrator.utils.email_builder import EmailBuilder + import traceback + from email import encoders + from email.mime.base import MIMEBase from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText - from email.mime.base import MIMEBase - from email import encoders + from typing import Any + + 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): @@ -35,10 +37,15 @@ class Email(BaseActivity): notification_handler (NotificationHandler): Notification management handler """ - def __init__(self, sender_email: str, sender_password: str, - smtp_server: str, smtp_port: int, - logger: Logger, notification_handler: NotificationHandler): - + def __init__( + self, + sender_email: str, + sender_password: str, + smtp_server: str, + smtp_port: int, + logger: Logger, + notification_handler: NotificationHandler, + ): self.email_builder = EmailBuilder(logger=logger) self.sender_email = sender_email @@ -46,7 +53,7 @@ class Email(BaseActivity): self.smtp_port = smtp_port 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: self.server = smtplib.SMTP(smtp_server, smtp_port, timeout=20) @@ -55,9 +62,7 @@ class Email(BaseActivity): self.server.starttls() self.server.login(self.sender_email, self.sender_password) - BaseActivity.__init__(self, - logger=logger, - notification_handler=notification_handler) + BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler) def shutdown(self): """ @@ -65,7 +70,7 @@ class Email(BaseActivity): """ 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]: """ Build HTML email content for configured receiver groups. @@ -91,18 +96,14 @@ class Email(BaseActivity): receiver_groups = input_data['receiver_groups'] mail_type = input_data['mail_type'] - self.info(f"Building email html for {mail_type} mail type.", - metadata=metadata) + self.info(f'Building email html for {mail_type} mail type.', metadata=metadata) - for group_name, group_config in receiver_groups.items(): - - html = self.email_builder.build_email( - group_config['notifications'], mail_type) + for _group_name, group_config in receiver_groups.items(): + html = self.email_builder.build_email(group_config['notifications'], mail_type) group_config['html'] = html - self.info(f"Email html built for {mail_type} mail type.", - metadata=metadata) + self.info(f'Email html built for {mail_type} mail type.', metadata=metadata) return receiver_groups @@ -124,17 +125,12 @@ class Email(BaseActivity): try: # Create the attachment as a MIMEBase object part = MIMEBase('application', 'octet-stream') - part.set_payload( - attachment['attachment_content'].encode('utf-8')) + part.set_payload(attachment['attachment_content'].encode('utf-8')) encoders.encode_base64(part) - part.add_header( - 'Content-Disposition', - f'attachment; filename="{att_name}"' - ) + part.add_header('Content-Disposition', f'attachment; filename="{att_name}"') msg.attach(part) except Exception as e: - self.logger.error( - f"Failed to attach content of {att_name}: {e}") + self.logger.error(f'Failed to attach content of {att_name}: {e}') raise e @@ -152,30 +148,27 @@ class Email(BaseActivity): Exception: If email sending fails after reconnection attempts. """ try: - self.server.sendmail( - self.sender_email, receivers, msg.as_string()) + self.server.sendmail(self.sender_email, receivers, msg.as_string()) except SMTPServerDisconnected as e: - self.logger.error(f"SMTP server disconnected: {e}") - self.logger.info( - f"Reconnecting to {self.smtp_server}:{self.smtp_port}") + self.logger.error(f'SMTP server disconnected: {e}') + self.logger.info(f'Reconnecting to {self.smtp_server}:{self.smtp_port}') if self.server: try: self.server.quit() except SMTPServerDisconnected as e: - self.logger.info(f"Server already disconnected: {e}") + self.logger.info(f'Server already disconnected: {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 - self.server = smtplib.SMTP( - self.smtp_server, self.smtp_port, timeout=20) + self.server = smtplib.SMTP(self.smtp_server, self.smtp_port, timeout=20) if self.sender_password: self.server.starttls() self.server.login(self.sender_email, self.sender_password) 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]: """ Send email notifications to configured receiver groups. @@ -202,40 +195,38 @@ class Email(BaseActivity): mail_type = input_data['mail_type'] if self.smtp_server is None: - self.info(f"Skipping email sending for {mail_type} mail type.", - metadata=metadata) + self.info(f'Skipping email sending for {mail_type} mail type.', metadata=metadata) return {} - self.info(f"Sending email for {mail_type} mail type.", - metadata=metadata) + self.info(f'Sending email for {mail_type} mail type.', metadata=metadata) for group_name, group_config in receiver_groups.items(): try: - receivers = ", ".join(group_config['members']) + receivers = ', '.join(group_config['members']) - self.info(f"Sending email to {group_name}: {receivers}", - metadata=metadata) + self.info(f'Sending email to {group_name}: {receivers}', metadata=metadata) msg = MIMEMultipart() msg.attach(MIMEText(group_config['html'], 'html')) msg['From'] = self.sender_email msg['To'] = receivers - msg['Subject'] = f"SIENTIA™ {mail_type}" + msg['Subject'] = f'SIENTIA™ {mail_type}' msg = self.handle_attachments( [ { - "filename": f"{notification['trigger']}_{notification['notification_id']}.txt", - "attachment_content": notification['attachment_content'] + 'filename': f'{notification["trigger"]}_{notification["notification_id"]}.txt', + 'attachment_content': notification['attachment_content'], } for notification in group_config['notifications'] - if notification.get('attachment_content') is not None], - msg) + if notification.get('attachment_content') is not None + ], + msg, + ) self.try_send_email(msg, receivers) except Exception as e: - self.error(f"Failed to send email to {group_name}: {e}", - metadata=metadata) + self.error(f'Failed to send email to {group_name}: {e}', metadata=metadata) traceback.print_exc() group_config['status'] = 'failed' else: @@ -244,13 +235,11 @@ class Email(BaseActivity): pod_id=self.pod_id, model_name=metadata['model_name'], pipeline_name=metadata['workflow_name'], - email_group=group_name + email_group=group_name, ).inc() - self.info(f"Email sent to {group_name}: {receivers}", - metadata=metadata) + self.info(f'Email sent to {group_name}: {receivers}', metadata=metadata) - self.info(f"Email sent for {mail_type} mail type.", - metadata=metadata) + self.info(f'Email sent for {mail_type} mail type.', metadata=metadata) return receiver_groups diff --git a/orchestrator/activities/formatters.py b/orchestrator/activities/formatters.py index 6928bd4..d32b4c0 100644 --- a/orchestrator/activities/formatters.py +++ b/orchestrator/activities/formatters.py @@ -2,19 +2,25 @@ from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): import json - from typing import Any from logging import Logger + from math import ceil + from typing import Any + from pandas import DataFrame 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 orchestrator.utils.orchestrator_functions import ( - scouter, predictions_batch, gather_read_tags, build_tag_config, minimal_retrain - ) + from sientia_do.temporal.activities.base import BaseActivity 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): @@ -41,16 +47,18 @@ class Formatters(BaseActivity): notification_handler (NotificationHandler): Notification management handler """ - def __init__(self, - scouter_namespace: str, - laborious_namespace: str, - logger: Logger, notification_handler: NotificationHandler): + def __init__( + self, + scouter_namespace: str, + laborious_namespace: str, + logger: Logger, + notification_handler: NotificationHandler, + ): self.scouter_namespace = scouter_namespace self.laborious_namespace = laborious_namespace - BaseActivity.__init__(self, logger=logger, - notification_handler=notification_handler) + BaseActivity.__init__(self, logger=logger, 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]: """ Process pipeline configurations into Temporal-compatible schedule configurations. @@ -61,7 +69,7 @@ class Formatters(BaseActivity): Pipeline Types Supported: - 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 Args: @@ -73,46 +81,43 @@ class Formatters(BaseActivity): - 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'] - schedule_config = { - self.scouter_namespace: {}, - self.laborious_namespace: {} - } + schedule_config = {self.scouter_namespace: {}, self.laborious_namespace: {}} for pipeline in pipelines: if pipeline['workflow_type'] == 'scouter': schedule_config[self.scouter_namespace][pipeline['schedule_name']] = { **scouter(pipeline), - "updated_at": pipeline.get( - "updated_at", now().strftime(DATETIME_FORMAT_MS_WITH_TZ)) + 'updated_at': pipeline.get( + 'updated_at', now().strftime(DATETIME_FORMAT_MS_WITH_TZ) + ), } 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), - "updated_at": pipeline.get( - "updated_at", now().strftime(DATETIME_FORMAT_MS_WITH_TZ)) + 'updated_at': pipeline.get( + 'updated_at', now().strftime(DATETIME_FORMAT_MS_WITH_TZ) + ), } 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), - "updated_at": pipeline.get( - "updated_at", now().strftime(DATETIME_FORMAT_MS_WITH_TZ)) + 'updated_at': pipeline.get( + 'updated_at', now().strftime(DATETIME_FORMAT_MS_WITH_TZ) + ), } - self.info("Processed schedules", metadata=metadata) - self.debug(json.dumps( - schedule_config, indent=4, sort_keys=True), metadata=metadata) + self.info('Processed schedules', metadata=metadata) + self.debug(json.dumps(schedule_config, indent=4, sort_keys=True), metadata=metadata) 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]: """ 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 """ - 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'] opc_servers_list = input_data['opc_servers'] @@ -154,43 +159,42 @@ class Formatters(BaseActivity): last_index = 0 for i in range(1, number_of_slots): - slot_config[f"{i}"] = {} - for tag in tags[last_index:last_index + tags_per_slot]: + slot_config[f'{i}'] = {} + for tag in tags[last_index : last_index + tags_per_slot]: try: - slot_config = build_tag_config( - tag, slot_config.copy(), opc_servers, i) + slot_config = build_tag_config(tag, slot_config.copy(), opc_servers, i) except ValueError as e: self.send_notification( metadata=metadata, - notification_id="ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR", + notification_id='ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR', message=str(e), - block="orchestrator", - level=NotificationLevel.ERROR + block='orchestrator', + level=NotificationLevel.ERROR, ) last_index += tags_per_slot - slot_config[f"{number_of_slots}"] = {} + slot_config[f'{number_of_slots}'] = {} for tag in tags[last_index:]: try: 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: self.send_notification( metadata=metadata, - notification_id="ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR", + notification_id='ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR', message=str(e), - block="orchestrator", - level=NotificationLevel.ERROR + block='orchestrator', + level=NotificationLevel.ERROR, ) - self.info("Processed slots", metadata=metadata) - self.debug(json.dumps( - slot_config, indent=4, sort_keys=True), metadata=metadata) + self.info('Processed slots', metadata=metadata) + self.debug(json.dumps(slot_config, indent=4, sort_keys=True), metadata=metadata) 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]: """ Formats the schedule config to a dictionary with the schedule name as the key. @@ -202,9 +206,9 @@ class Formatters(BaseActivity): Returns: 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'] @@ -219,16 +223,20 @@ class Formatters(BaseActivity): config[namespace][schedule_name] = updated_at - self.info("Formatted schedule config", metadata=metadata) - self.debug(json.dumps( - config, indent=4, sort_keys=True), metadata=metadata) + self.info('Formatted schedule config', metadata=metadata) + self.debug(json.dumps(config, indent=4, sort_keys=True), metadata=metadata) return config - def compare_config_timestamps(self, - schedules: dict[str, Any], current_schedules: dict[str, Any], - to_update: dict[str, Any], to_create: dict[str, Any], - namespace: str, metadata: dict[str, Any]): + def compare_config_timestamps( + self, + schedules: 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 which schedules need to be updated or created. @@ -243,23 +251,22 @@ class Formatters(BaseActivity): """ for schedule_name, schedule in schedules.items(): 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] 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: to_update[namespace][schedule_name] = schedule else: to_create[namespace][schedule_name] = schedule - @activity.defn(name="create_schedule_config") - async def create_schedule_config(self, - input_data: dict[str, Any]) -> dict[str, Any]: + @activity.defn(name='create_schedule_config') + async def create_schedule_config(self, input_data: dict[str, Any]) -> dict[str, Any]: """ Creates a schedule config dictionary based on the input data. 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 """ - 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'] schedule_config = input_data['schedule_config'] - to_update = { - self.scouter_namespace: {}, - self.laborious_namespace: {} - } - to_create = { - self.scouter_namespace: {}, - self.laborious_namespace: {} - } - to_delete = { - self.scouter_namespace: [], - self.laborious_namespace: [] - } + to_update = {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(): current_schedules = current_schedule_config.get(namespace, {}) 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 schedule_name in schedules: if schedule_name not in schedule_config[namespace]: to_delete[namespace].append(schedule_name) - output = { - "to_update": to_update, - "to_create": to_create, - "to_delete": to_delete - } + output = {'to_update': to_update, 'to_create': to_create, 'to_delete': to_delete} - self.info("Created schedule config", metadata=metadata) - self.debug(json.dumps( - output, indent=4, sort_keys=True), metadata=metadata) + self.info('Created schedule config', metadata=metadata) + self.debug(json.dumps(output, indent=4, sort_keys=True), metadata=metadata) return output - @activity.defn(name="create_slot_config") - async def create_slot_config(self, - input_data: dict[str, Any]) -> dict[str, Any]: + @activity.defn(name='create_slot_config') + async def create_slot_config(self, input_data: dict[str, Any]) -> dict[str, Any]: """ Creates a slot config dictionary based on the input data. 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 """ - 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'] slot_config = input_data['slot_config'] @@ -349,21 +342,18 @@ class Formatters(BaseActivity): number_of_slots = len(slot_config) if number_of_current_slots > number_of_slots: - to_delete = [str(i) for i in range( - number_of_slots + 1, number_of_current_slots + 1)] + to_delete = [str(i) for i in range(number_of_slots + 1, number_of_current_slots + 1)] - output = { - "to_delete": to_delete, - "to_insert": slot_config - } + output = {'to_delete': to_delete, 'to_insert': slot_config} - self.info("Created slot config", metadata=metadata) - self.debug(json.dumps( - output, indent=4, sort_keys=True), metadata=metadata) + self.info('Created slot config', metadata=metadata) + self.debug(json.dumps(output, indent=4, sort_keys=True), metadata=metadata) 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. @@ -377,13 +367,14 @@ class Formatters(BaseActivity): metadata=metadata, notification_id=notification_id, message=message, - block="report_orchestration", + block='report_orchestration', 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, - attachment: str) -> None: + def send_error_report( + self, metadata: dict[str, Any], message: str, notification_id: str, attachment: str + ) -> None: """ Sends an error notification report. @@ -397,9 +388,9 @@ class Formatters(BaseActivity): metadata=metadata, notification_id=notification_id, message=message, - block="report_orchestration", + block='report_orchestration', 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]]: @@ -408,7 +399,7 @@ class Formatters(BaseActivity): Args: 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. Returns: @@ -416,14 +407,20 @@ class Formatters(BaseActivity): - List of successful schedule keys in format "namespace/schedule_name" - Dictionary of error keys mapped to their error details """ - success_keys = [f"{value['namespace']}/{value['schedule_name']}" - for value in input_data if value['success']] + success_keys = [ + f'{value["namespace"]}/{value["schedule_name"]}' + for value in input_data + if value['success'] + ] - error_keys = {f"{value['namespace']}/{value['schedule_name']}": { - 'message': value['message'], - 'attachment': value.get('attachment', None) + error_keys = { + f'{value["namespace"]}/{value["schedule_name"]}': { + '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 @@ -440,15 +437,20 @@ class Formatters(BaseActivity): - List of successful keys - List of error keys """ - success_keys = [key for key, value - in input_data.items() if value['success']] + success_keys = [key for key, value in input_data.items() if value['success']] - error_keys = [key for key, value - in input_data.items() if not value['success']] + error_keys = [key for key, value in input_data.items() if not value['success']] 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. @@ -462,30 +464,28 @@ class Formatters(BaseActivity): if len(success_keys) > 0: self.send_success_report( 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'], - attachment=schedule_data['items'] + attachment=schedule_data['items'], ) if len(error_keys) > 0: attachment = [] for key, value in error_keys.items(): if value['attachment'] is not None: - attachment.append( - f"{key}:\n{value['message']}\n{value['attachment']}") + attachment.append(f'{key}:\n{value["message"]}\n{value["attachment"]}') else: - attachment.append(f"{key}:\n{value['message']}") + attachment.append(f'{key}:\n{value["message"]}') self.send_error_report( metadata=metadata, - message=f"Fails on {schedule_type}: \n {', '.join(error_keys)}", - notification_id=f"{schedule_data['id']}_ERROR", - attachment=topic_separator.join(attachment) + message=f'Fails on {schedule_type}: \n {", ".join(error_keys)}', + notification_id=f'{schedule_data["id"]}_ERROR', + attachment=topic_separator.join(attachment), ) - @activity.defn(name="report_schedule_orchestration") - async def report_schedule_orchestration(self, - input_data: dict[str, Any]) -> None: + @activity.defn(name='report_schedule_orchestration') + async def report_schedule_orchestration(self, input_data: dict[str, Any]) -> None: """ Reports the orchestration result to the notification handler. @@ -496,9 +496,9 @@ class Formatters(BaseActivity): - updated_schedules (dict[str, Any]): The updated 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'] updated_schedules = input_data['updated_schedules'] @@ -507,34 +507,32 @@ class Formatters(BaseActivity): schedules_report = { 'created schedules': { 'items': created_schedules, - 'id': 'REPORT_ORCHESTRATION_CREATED_SCHEDULES' + 'id': 'REPORT_ORCHESTRATION_CREATED_SCHEDULES', }, 'updated schedules': { 'items': updated_schedules, - 'id': 'REPORT_ORCHESTRATION_UPDATED_SCHEDULES' + 'id': 'REPORT_ORCHESTRATION_UPDATED_SCHEDULES', }, 'deleted schedules': { 'items': deleted_schedules, - 'id': 'REPORT_ORCHESTRATION_DELETED_SCHEDULES' - } + 'id': 'REPORT_ORCHESTRATION_DELETED_SCHEDULES', + }, } for schedule_type, schedule_data in schedules_report.items(): if len(schedule_data['items']) > 0: - success_keys, error_keys = self.parse_report_schedule( - schedule_data['items']) + success_keys, error_keys = self.parse_report_schedule(schedule_data['items']) self.manage_and_send_report( metadata=metadata, success_keys=success_keys, error_keys=error_keys, schedule_type=schedule_type, - schedule_data=schedule_data + schedule_data=schedule_data, ) - @activity.defn(name="report_slot_orchestration") - async def report_slot_orchestration(self, - input_data: dict[str, Any]) -> None: + @activity.defn(name='report_slot_orchestration') + async def report_slot_orchestration(self, input_data: dict[str, Any]) -> None: """ Reports the orchestration result to the notification handler. @@ -545,9 +543,9 @@ class Formatters(BaseActivity): - 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'] deleted_slots = input_data['deleted_slots'] @@ -558,16 +556,16 @@ class Formatters(BaseActivity): if len(success_keys) > 0: self.send_success_report( metadata=metadata, - message=f"Inserted slots: \n {', '.join(success_keys)}", - notification_id="REPORT_ORCHESTRATION_INSERTED_SLOTS" + message=f'Inserted slots: \n {", ".join(success_keys)}', + notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS', ) if len(error_keys) > 0: self.send_error_report( metadata=metadata, - message=f"Failed to insert slots: \n {', '.join(error_keys)}", - notification_id="REPORT_ORCHESTRATION_INSERTED_SLOTS", - attachment=inserted_slots + message=f'Failed to insert slots: \n {", ".join(error_keys)}', + notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS', + attachment=inserted_slots, ) if len(deleted_slots) > 0: @@ -576,19 +574,19 @@ class Formatters(BaseActivity): if len(success_keys) > 0: self.send_success_report( metadata=metadata, - message=f"Deleted slots: \n {', '.join(success_keys)}", - notification_id="REPORT_ORCHESTRATION_DELETED_SLOTS" + message=f'Deleted slots: \n {", ".join(success_keys)}', + notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS', ) if len(error_keys) > 0: self.send_error_report( metadata=metadata, - message=f"Failed to delete slots: \n {', '.join(error_keys)}", - notification_id="REPORT_ORCHESTRATION_DELETED_SLOTS", - attachment=deleted_slots + message=f'Failed to delete slots: \n {", ".join(error_keys)}', + notification_id='REPORT_ORCHESTRATION_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]: """ Formats the receiver_groups status to a dataframe to be stored in the database. @@ -602,23 +600,21 @@ class Formatters(BaseActivity): Returns: dict[str, Any]: The formatted log report as a dictionary representation of a DataFrame. """ - metadata = input_data["metadata"] - mail_type = input_data["mail_type"] + metadata = input_data['metadata'] + 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'] data = {} for group_name, group_config in receiver_groups.items(): - for notification in group_config['notifications']: - notification_id = notification['notification_id'] trigger = notification['trigger'] - key = f"{notification_id}:{trigger}" + key = f'{notification_id}:{trigger}' if key not in data: data[key] = { @@ -634,7 +630,7 @@ class Formatters(BaseActivity): 'project': notification['project'], 'model_name': notification['model_name'], 'model_id': notification['model_id'], - 'mail_type': mail_type + 'mail_type': mail_type, } else: if group_name not in data[key]['groups']: @@ -642,7 +638,7 @@ class Formatters(BaseActivity): 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]: """ Filter notifications for comprehensive scheduled reports. @@ -678,16 +674,13 @@ class Formatters(BaseActivity): notification_package = input_data['notification_package'] sending_configs = input_data['sending_configs'] - self.info("Filtering notification reports...", metadata=metadata) + self.info('Filtering notification reports...', metadata=metadata) receiver_groups = {} for receiver_group in sending_configs: group_name = receiver_group['group_name'] - receiver_groups[group_name] = { - **receiver_group, - "notifications": [] - } + receiver_groups[group_name] = {**receiver_group, 'notifications': []} receiver_groups[group_name]['notifications'] = [] already_added_keys = [] @@ -695,15 +688,18 @@ class Formatters(BaseActivity): ignore_list = receiver_group.get('ignore', []) for notification in notification_package: - alert_type = "reports" + alert_type = 'reports' notification_id = notification['notification_id'] - key = f"{notification['trigger']}:{notification_id}" + key = f'{notification["trigger"]}:{notification_id}' # 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: - receiver_groups[group_name]["notifications"].append( - notification) + if ( + alert_type in receiver_group['contents'] + 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) # Remove groups with no notifications diff --git a/orchestrator/activities/mongo_db.py b/orchestrator/activities/mongo_db.py index 0be4e73..1dfba8a 100644 --- a/orchestrator/activities/mongo_db.py +++ b/orchestrator/activities/mongo_db.py @@ -1,16 +1,18 @@ -from temporalio import workflow, activity +from datetime import UTC +from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): - from typing import Any import traceback + from datetime import datetime from logging import Logger + from typing import Any + from pymongo import MongoClient 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 sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now - from datetime import datetime, timezone def clear_mongo_id(docs: list) -> list: @@ -28,10 +30,10 @@ def clear_mongo_id(docs: list) -> list: clear_mongo_id(doc) elif isinstance(doc, dict): - if "_id" in doc: - del doc["_id"] + if '_id' in doc: + del doc['_id'] - for key, value in doc.items(): + for _key, value in doc.items(): if isinstance(value, list): clear_mongo_id(value) elif isinstance(value, dict): @@ -57,14 +59,18 @@ class MongoDB(BaseActivity): notification_handler (NotificationHandler): Notification management handler """ - def __init__(self, connection_string: str, database_name: str, ttl_index_seconds: int, - logger: Logger, - notification_handler: NotificationHandler): + def __init__( + self, + connection_string: str, + database_name: str, + ttl_index_seconds: int, + logger: Logger, + notification_handler: NotificationHandler, + ): self.connection_string = connection_string self.database_name = database_name - self.client = MongoClient( - self.connection_string, serverSelectionTimeoutMS=5000) + self.client = MongoClient(self.connection_string, serverSelectionTimeoutMS=5000) self.client.server_info() # Trigger an exception if connection fails self.database = self.client[self.database_name] @@ -72,11 +78,9 @@ class MongoDB(BaseActivity): self.ttl_index_seconds = ttl_index_seconds # Initialize MongoDB client here (omitted for brevity) - logger.info("MongoDB connection initialized") + logger.info('MongoDB connection initialized') - BaseActivity.__init__(self, - logger=logger, - notification_handler=notification_handler) + BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler) def shutdown(self): """ @@ -84,11 +88,11 @@ class MongoDB(BaseActivity): """ try: if self.client: - self.logger.info("Closing MongoDB connection...") + self.logger.info('Closing MongoDB connection...') self.client.close() - self.logger.info("MongoDB connection closed successfully") + self.logger.info('MongoDB connection closed successfully') 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): """ @@ -109,13 +113,15 @@ class MongoDB(BaseActivity): """ collection = self.database[collection_name] - documents = list(collection.find(filters, {"_id": 0})) + documents = list(collection.find(filters, {'_id': 0})) documents = clear_mongo_id(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]]: """ 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. """ - query = input_data.get("query", {}) - metadata = input_data.get("metadata", {}) - timestamp_fields = input_data.get("timestamp_fields", []) + query = input_data.get('query', {}) + metadata = input_data.get('metadata', {}) + timestamp_fields = input_data.get('timestamp_fields', []) - collection_name = query.get("collection") + collection_name = query.get('collection') 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( - 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: documents = self.find(collection_name, filters) 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 timestamp_field in timestamp_fields: if timestamp_field in document: - document[timestamp_field] = document[timestamp_field].replace(tzinfo=timezone.utc).strftime( - DATETIME_FORMAT_MS_WITH_TZ) + document[timestamp_field] = ( + document[timestamp_field] + .replace(tzinfo=UTC) + .strftime(DATETIME_FORMAT_MS_WITH_TZ) + ) - self.debug( - f"Documents loaded: {documents}", metadata=metadata) + self.debug(f'Documents loaded: {documents}', metadata=metadata) return documents @@ -162,19 +174,20 @@ class MongoDB(BaseActivity): trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="MONGODB_QUERY_ERROR", - message=f"Failed to execute MongoDB query: {e}", - block="load_query_from_mongodb", + notification_id='MONGODB_QUERY_ERROR', + message=f'Failed to execute MongoDB query: {e}', + block='load_query_from_mongodb', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) self.error(trace, metadata=metadata) raise e - @activity.defn(name="aggregate_documents_in_mongodb") - async def aggregate_documents_in_mongodb(self, - input_data: dict[str, Any]) -> list[dict[str, Any]]: + @activity.defn(name='aggregate_documents_in_mongodb') + async def aggregate_documents_in_mongodb( + self, input_data: dict[str, Any] + ) -> list[dict[str, Any]]: """ 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. """ - query = input_data.get("query", {}) - metadata = input_data.get("metadata", {}) - timestamp_fields = input_data.get("timestamp_fields", []) + query = input_data.get('query', {}) + metadata = input_data.get('metadata', {}) + timestamp_fields = input_data.get('timestamp_fields', []) - collection_name = query.get("collection") + collection_name = query.get('collection') if not collection_name: - raise ValueError("Collection name must be provided in the query.") - aggregation = query.get("aggregation") + raise ValueError('Collection name must be provided in the query.') + aggregation = query.get('aggregation') if not aggregation: - raise ValueError("Aggregation must be provided.") - aggregation.append({"$project": {"_id": 0}}) + raise ValueError('Aggregation must be provided.') + aggregation.append({'$project': {'_id': 0}}) 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: collection = self.database[collection_name] - aggregated_documents = list( - collection.aggregate(aggregation)) + aggregated_documents = list(collection.aggregate(aggregation)) aggregated_documents = clear_mongo_id(aggregated_documents) 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 timestamp_field in timestamp_fields: if timestamp_field in document: - document[timestamp_field] = document[timestamp_field].replace(tzinfo=timezone.utc).strftime( - DATETIME_FORMAT_MS_WITH_TZ) + document[timestamp_field] = ( + document[timestamp_field] + .replace(tzinfo=UTC) + .strftime(DATETIME_FORMAT_MS_WITH_TZ) + ) - self.debug( - f"Aggregation result: {aggregated_documents}", metadata=metadata) + self.debug(f'Aggregation result: {aggregated_documents}', metadata=metadata) return aggregated_documents @@ -227,83 +245,85 @@ class MongoDB(BaseActivity): trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="MONGODB_AGGREGATION_ERROR", - message=f"Failed to execute MongoDB aggregation: {e}", - block="aggregate_documents_in_mongodb", + notification_id='MONGODB_AGGREGATION_ERROR', + message=f'Failed to execute MongoDB aggregation: {e}', + block='aggregate_documents_in_mongodb', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) self.error(trace, metadata=metadata) 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: """ Update the timestamps of the pipelines in the MongoDB collection. input_data: - updated_pipelines (list): List of updated pipelines. """ - updated_pipelines = input_data.get("updated_pipelines", []) - metadata = input_data.get("metadata", {}) + updated_pipelines = input_data.get('updated_pipelines', []) + metadata = input_data.get('metadata', {}) 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 argument = [ - {"schedule_name": pipeline["schedule_name"], - "namespace": pipeline["namespace"]} - for pipeline in updated_pipelines if pipeline["success"] + {'schedule_name': pipeline['schedule_name'], 'namespace': pipeline['namespace']} + for pipeline in updated_pipelines + if pipeline['success'] ] - data_filter = {"$or": argument} if argument else {} + data_filter = {'$or': argument} if argument else {} try: - collection.update_many( - data_filter, - {"$set": {"updated_at": date_now}} - ) + collection.update_many(data_filter, {'$set': {'updated_at': date_now}}) success_count += 1 except Exception as e: trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="MONGODB_UPDATE_PIPELINES_ERROR", - message=f"Failed to update pipelines timestamps: {e}", - block="update_pipelines_timestamps", + notification_id='MONGODB_UPDATE_PIPELINES_ERROR', + message=f'Failed to update pipelines timestamps: {e}', + block='update_pipelines_timestamps', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) self.error(trace, metadata=metadata) raise e 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: """ Create the timestamps of the pipelines in the MongoDB collection. input_data: - created_pipelines (list): List of created pipelines. """ - created_pipelines = input_data.get("created_pipelines", []) - metadata = input_data.get("metadata", {}) - collection = self.database["orchestrated_schedules"] + created_pipelines = input_data.get('created_pipelines', []) + metadata = input_data.get('metadata', {}) + collection = self.database['orchestrated_schedules'] - self.info("Creating pipelines timestamps...", metadata=metadata) + self.info('Creating pipelines timestamps...', metadata=metadata) success_count = 0 date_now = now() argument = [ - {"schedule_name": pipeline["schedule_name"], - "namespace": pipeline["namespace"], - "updated_at": date_now} - for pipeline in created_pipelines if pipeline["success"] + { + 'schedule_name': pipeline['schedule_name'], + 'namespace': pipeline['namespace'], + 'updated_at': date_now, + } + for pipeline in created_pipelines + if pipeline['success'] ] data_filter = argument if argument else {} @@ -315,39 +335,41 @@ class MongoDB(BaseActivity): trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="MONGODB_CREATE_PIPELINES_ERROR", - message=f"Failed to create pipelines timestamps: {e}", - block="create_pipelines_timestamps", + notification_id='MONGODB_CREATE_PIPELINES_ERROR', + message=f'Failed to create pipelines timestamps: {e}', + block='create_pipelines_timestamps', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) self.error(trace, metadata=metadata) raise e 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: """ Delete the timestamps of the pipelines in the MongoDB collection. input_data: - deleted_pipelines (list): List of deleted pipelines. """ - deleted_pipelines = input_data.get("deleted_pipelines", []) - metadata = input_data.get("metadata", {}) - collection = self.database["orchestrated_schedules"] + deleted_pipelines = input_data.get('deleted_pipelines', []) + metadata = input_data.get('metadata', {}) + collection = self.database['orchestrated_schedules'] - self.info("Deleting pipelines timestamps...", metadata=metadata) + self.info('Deleting pipelines timestamps...', metadata=metadata) success_count = 0 argument = [ - {"schedule_name": pipeline["schedule_name"], - "namespace": pipeline["namespace"]} - for pipeline in deleted_pipelines if pipeline["success"] + {'schedule_name': pipeline['schedule_name'], 'namespace': pipeline['namespace']} + for pipeline in deleted_pipelines + if pipeline['success'] ] - data_filter = {"$or": argument} if argument else {} + data_filter = {'$or': argument} if argument else {} try: collection.delete_many(data_filter) @@ -356,19 +378,21 @@ class MongoDB(BaseActivity): trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="MONGODB_DELETE_PIPELINES_ERROR", - message=f"Failed to delete pipelines timestamps: {e}", - block="delete_pipelines_timestamps", + notification_id='MONGODB_DELETE_PIPELINES_ERROR', + message=f'Failed to delete pipelines timestamps: {e}', + block='delete_pipelines_timestamps', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) self.error(trace, metadata=metadata) raise e 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: """ Create a collection with a TTL index. @@ -376,20 +400,21 @@ class MongoDB(BaseActivity): - collection_name (str): The name of the collection to create. - ttl_index (str): The name of the TTL index to create. """ - pipelines = input_data.get("pipelines", {}) - metadata = input_data.get("metadata", {}) + pipelines = input_data.get('pipelines', {}) + metadata = input_data.get('metadata', {}) self.info( - f"Creating collection with TTL index for pipelines: {list(pipelines.keys())}", - metadata=metadata) + f'Creating collection with TTL index for pipelines: {list(pipelines.keys())}', + metadata=metadata, + ) collection_names = self.database.list_collection_names() created_collections = [] created_indexes = [] - for pipeline_name, pipeline_config in pipelines.items(): - collection = pipeline_config["topic"] + for _pipeline_name, pipeline_config in pipelines.items(): + collection = pipeline_config['topic'] try: # Check if collection exists @@ -402,16 +427,17 @@ class MongoDB(BaseActivity): existing_indexes = collection.list_indexes() ttl_index_exists = False 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 break # Create TTL index if it doesn't exist if not ttl_index_exists: collection.create_index( - "inserted_at", - expireAfterSeconds=self.ttl_index_seconds, - background=True + 'inserted_at', expireAfterSeconds=self.ttl_index_seconds, background=True ) created_indexes.append(collection) @@ -419,30 +445,24 @@ class MongoDB(BaseActivity): trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="MONGODB_CREATE_COLLECTION_ERROR", - message=f"Failed to create collection {collection} with TTL index: {e}", - block="create_collection_with_ttl_index", + notification_id='MONGODB_CREATE_COLLECTION_ERROR', + message=f'Failed to create collection {collection} with TTL index: {e}', + block='create_collection_with_ttl_index', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) self.error(trace, metadata=metadata) raise e self.info( - f"Created {len(created_collections)} collections and {len(created_indexes)} indexes", - metadata=metadata + f'Created {len(created_collections)} collections and {len(created_indexes)} indexes', + metadata=metadata, ) - self.debug( - f"Created collections: {created_collections}", - metadata=metadata - ) - self.debug( - f"Created indexes: {created_indexes}", - metadata=metadata - ) + self.debug(f'Created collections: {created_collections}', 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]]: """ 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'] base_data_filter = input_data['base_data_filter'] - self.debug( - f"Loading data from MongoDB: {input_data}", - metadata=metadata - ) + self.debug(f'Loading data from MongoDB: {input_data}', metadata=metadata) try: - if last_data_timestamp is None: data_filter = base_data_filter else: data_filter = { **base_data_filter, - "timestamp": { - "$gt": datetime.strptime(last_data_timestamp, DATETIME_FORMAT_MS_WITH_TZ) - } + 'timestamp': { + '$gt': datetime.strptime(last_data_timestamp, DATETIME_FORMAT_MS_WITH_TZ) + }, } - self.debug( - f"Data filter: {data_filter}", - metadata=metadata - ) + self.debug(f'Data filter: {data_filter}', metadata=metadata) data = self.find(collection_name, data_filter) - self.debug( - f"Collected: {data}", - metadata=metadata - ) + self.debug(f'Collected: {data}', metadata=metadata) for item in data: - item['timestamp'] = item['timestamp'].replace(tzinfo=timezone.utc).strftime( - DATETIME_FORMAT_MS_WITH_TZ) + item['timestamp'] = ( + item['timestamp'].replace(tzinfo=UTC).strftime(DATETIME_FORMAT_MS_WITH_TZ) + ) - self.info( - f"Loaded {len(data)} documents from MongoDB", - metadata=metadata - ) + self.info(f'Loaded {len(data)} documents from MongoDB', metadata=metadata) - self.debug( - f"Loaded data: {data}", - metadata=metadata - ) + self.debug(f'Loaded data: {data}', metadata=metadata) return data except Exception as e: trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="MONGO_LOAD_ERROR", - message=f"Error loading data from MongoDB: {e}", - block="load_latest_data", + notification_id='MONGO_LOAD_ERROR', + message=f'Error loading data from MongoDB: {e}', + block='load_latest_data', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) raise e diff --git a/orchestrator/activities/slot_manager.py b/orchestrator/activities/slot_manager.py index 9accfd1..e41562b 100644 --- a/orchestrator/activities/slot_manager.py +++ b/orchestrator/activities/slot_manager.py @@ -1,17 +1,17 @@ - from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): - from typing import Any - import traceback import json + import traceback + from datetime import datetime, timedelta 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.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 pandas import DataFrame - from datetime import datetime, timedelta class SlotManager(Redis): @@ -40,14 +40,18 @@ class SlotManager(Redis): notification_handler (NotificationHandler): Notification management handler """ - def __init__(self, host: str, port: int, - username: str, password: str, - logger: Logger, notification_handler: NotificationHandler): + def __init__( + self, + 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, - password, logger, notification_handler) - - @activity.defn(name="load_opc_slots") + @activity.defn(name='load_opc_slots') 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. @@ -69,17 +73,16 @@ class SlotManager(Redis): 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 = {} 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 isinstance(slot_keys[0], bytes): @@ -93,20 +96,20 @@ class SlotManager(Redis): trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="REDIS_GET_ERROR", - message=f"Failed to load OPC slots: {e}", - block="load_opc_slots", + notification_id='REDIS_GET_ERROR', + message=f'Failed to load OPC slots: {e}', + block='load_opc_slots', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) self.error(trace, metadata=metadata) 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 - @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]: """ Load all active ingestors from Redis @@ -115,19 +118,16 @@ class SlotManager(Redis): 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: + 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( - f"Loaded {len(active_ingestors)} active ingestors", metadata=metadata) - - self.debug( - f"Active ingestors: \n {active_ingestors}", metadata=metadata) + self.debug(f'Active ingestors: \n {active_ingestors}', metadata=metadata) ingestors = [] @@ -142,16 +142,16 @@ class SlotManager(Redis): trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="REDIS_GET_ERROR", - message=f"Failed to load active ingestors: {e}", - block="load_active_ingestors", + notification_id='REDIS_GET_ERROR', + message=f'Failed to load active ingestors: {e}', + block='load_active_ingestors', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) self.error(trace, metadata=metadata) 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]: """ Update OPC slots in Redis @@ -166,8 +166,8 @@ class SlotManager(Redis): """ to_insert = input_data['to_insert'] - metadata = input_data.get("metadata", {}) - self.info("Updating OPC slots...", metadata=metadata) + metadata = input_data.get('metadata', {}) + self.info('Updating OPC slots...', metadata=metadata) report = {} @@ -175,30 +175,20 @@ class SlotManager(Redis): for slot in to_insert: try: - self.set(f"slot:opc_tags:{slot}", - to_insert[slot], ttl=None) - report[slot] = { - "success": True, - "message": "Slot updated successfully" - } + self.set(f'slot:opc_tags:{slot}', to_insert[slot], ttl=None) + report[slot] = {'success': True, 'message': 'Slot updated successfully'} success_count += 1 except Exception as e: - self.error( - f"Failed to update slot {slot}: {str(e)}", metadata=metadata) - report[slot] = { - "success": False, - "message": str(e) - } + self.error(f'Failed to update slot {slot}: {str(e)}', metadata=metadata) + report[slot] = {'success': False, 'message': str(e)} - self.info( - f"Updated {success_count} of {len(to_insert)} OPC slots", metadata=metadata) + self.info(f'Updated {success_count} of {len(to_insert)} OPC slots', metadata=metadata) - self.debug( - f"Report: \n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata) + self.debug(f'Report: \n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata) 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]: """ Delete OPC slots from Redis @@ -213,9 +203,9 @@ class SlotManager(Redis): """ 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 = {} @@ -223,29 +213,20 @@ class SlotManager(Redis): for slot in to_delete: try: - self.redis_client.delete(f"slot:opc_tags:{slot}") - report[slot] = { - "success": True, - "message": "Slot deleted successfully" - } + self.redis_client.delete(f'slot:opc_tags:{slot}') + report[slot] = {'success': True, 'message': 'Slot deleted successfully'} success_count += 1 except Exception as e: - self.error( - f"Failed to delete slot {slot}: {str(e)}", metadata=metadata) - report[slot] = { - "success": False, - "message": str(e) - } + self.error(f'Failed to delete slot {slot}: {str(e)}', metadata=metadata) + report[slot] = {'success': False, 'message': str(e)} - self.info( - f"Deleted {success_count} of {len(to_delete)} OPC slots", metadata=metadata) + self.info(f'Deleted {success_count} of {len(to_delete)} OPC slots', metadata=metadata) - self.debug( - f"Report: \n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata) + self.debug(f'Report: \n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata) 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: """ 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. """ metadata = input_data['metadata'] - key = f"notification_last_timestamp:{input_data['mail_type']}" + key = f'notification_last_timestamp:{input_data["mail_type"]}' try: data_hold = self.get(key) except Exception as e: self.send_notification( metadata=metadata, - notification_id="REDIS_GET_ERROR", - message=f"Error getting last data timestamp: {e}", - block="get_last_data_timestamp", + notification_id='REDIS_GET_ERROR', + message=f'Error getting last data timestamp: {e}', + block='get_last_data_timestamp', level=NotificationLevel.ERROR, - attachment_content=traceback.format_exc() + attachment_content=traceback.format_exc(), ) raise e - self.debug( - f"Last collected timestamp: {data_hold}", - metadata=metadata - ) + self.debug(f'Last collected timestamp: {data_hold}', metadata=metadata) if not data_hold: return None 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]): """ 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. """ 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']) if data.empty: - self.warning("No data to insert", - metadata=metadata - ) + self.warning('No data to insert', metadata=metadata) return None last_data_timestamp = data['timestamp'].max() - self.debug( - f"Last collected timestamp to insert: {last_data_timestamp}", - metadata=metadata - ) + self.debug(f'Last collected timestamp to insert: {last_data_timestamp}', metadata=metadata) 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: self.send_notification( metadata=metadata, - notification_id="REDIS_SET_ERROR", - message=f"Error setting last data timestamp: {e}", - block="put_last_data_timestamp", + notification_id='REDIS_SET_ERROR', + message=f'Error setting last data timestamp: {e}', + block='put_last_data_timestamp', level=NotificationLevel.ERROR, - attachment_content=traceback.format_exc() + attachment_content=traceback.format_exc(), ) raise e 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]: """ Filter notification alerts with intelligent TTL-based duplicate prevention. @@ -369,16 +342,13 @@ class SlotManager(Redis): sending_configs = input_data['sending_configs'] notification_ttl = input_data['notification_ttl'] - self.info("Filtering notification alerts...", metadata=metadata) + self.info('Filtering notification alerts...', metadata=metadata) receiver_groups = {} for receiver_group in sending_configs: group_name = receiver_group['group_name'] - receiver_groups[group_name] = { - **receiver_group, - "notifications": [] - } + receiver_groups[group_name] = {**receiver_group, 'notifications': []} receiver_groups[group_name]['notifications'] = [] already_added_keys = [] @@ -386,28 +356,30 @@ class SlotManager(Redis): ignore_list = receiver_group.get('ignore', []) for notification in notification_package: - alert_type = "do_nothing" + alert_type = 'do_nothing' notification_id = notification['notification_id'] # Check if notification was recently sent - key = f"{notification['trigger']}:{notification_id}" + key = f'{notification["trigger"]}:{notification_id}' last_sent = self.get(key) if last_sent is None: - alert_type = "core_alerts" + alert_type = 'core_alerts' else: - last_sent = datetime.strptime( - last_sent, DATETIME_FORMAT_MS_WITH_TZ) + last_sent = datetime.strptime(last_sent, DATETIME_FORMAT_MS_WITH_TZ) # Check if "notification_ttl" seconds has passed since last sent if (now() - last_sent) > timedelta(seconds=notification_ttl): - alert_type = "persistent_alerts" + alert_type = 'persistent_alerts' # 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: - receiver_groups[group_name]["notifications"].append( - notification) + if ( + alert_type in receiver_group['contents'] + 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) # Remove groups with no notifications @@ -419,7 +391,7 @@ class SlotManager(Redis): 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: """ Store notification cache in Redis to track recently sent notifications. @@ -434,14 +406,14 @@ class SlotManager(Redis): log_report = DataFrame(input_data['log_report']) 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) - for index, row in log_report.iterrows(): + for _index, row in log_report.iterrows(): status = row['status'] 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.info("Notification cache stored...", metadata=metadata) + self.info('Notification cache stored...', metadata=metadata) diff --git a/orchestrator/activities/temporal_manager.py b/orchestrator/activities/temporal_manager.py index c54251d..6e7f75e 100644 --- a/orchestrator/activities/temporal_manager.py +++ b/orchestrator/activities/temporal_manager.py @@ -1,20 +1,26 @@ from temporalio import activity, workflow 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 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 - 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 @@ -30,69 +36,65 @@ class TemporalManager(BaseActivity): Args: host (str): Temporal server host address scouter_namespace (str): Scouter workflow namespace - laborious_namespace (str): Laborious workflow namespace + laborious_namespace (str): Laborious workflow namespace logger (Logger): Application logger instance notification_handler (NotificationHandler): Notification management handler """ - def __init__(self, host: str, scouter_namespace: str, laborious_namespace: str, - task_timeout: int, run_timeout: int, execution_timeout: int, logger: Logger, notification_handler: NotificationHandler): - + def __init__( + self, + host: str, + scouter_namespace: str, + laborious_namespace: str, + logger: Logger, + notification_handler: NotificationHandler, + ): self.temporal_host = host self.scouter_namespace = scouter_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.model_id_id_key = SearchAttributeKey.for_keyword("model_id") - self.model_name_id_key = SearchAttributeKey.for_keyword("model_name") - self.orchestrated_id_key = SearchAttributeKey.for_keyword( - "orchestrated") + self.model_id_id_key = SearchAttributeKey.for_keyword('model_id') + self.model_name_id_key = SearchAttributeKey.for_keyword('model_name') + self.orchestrated_id_key = SearchAttributeKey.for_keyword('orchestrated') - BaseActivity.__init__(self, - logger=logger, - notification_handler=notification_handler) + BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler) async def connect_to_temporal(self): """ Connect to Temporal server namespaces for scouter and laborious workflows. Creates client connections to both namespaces and stores them for later use. """ - self.logger.info( - f"Connecting to Temporal side namespaces at {self.temporal_host}") - self.logger.info(f"Scouter namespace: {self.scouter_namespace}") + self.logger.info(f'Connecting to Temporal side namespaces at {self.temporal_host}') + self.logger.info(f'Scouter namespace: {self.scouter_namespace}') scouter_client = await Client.connect( - target_host=self.temporal_host, - namespace=self.scouter_namespace + target_host=self.temporal_host, 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( - target_host=self.temporal_host, - namespace=self.laborious_namespace + target_host=self.temporal_host, namespace=self.laborious_namespace ) self.temporal_clients = { 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]: """ Normalize schedules. Removes schedules with no update time in mongo db collection "orchestrated_schedules". input_data: - orchestrated_schedules (dict[str, Any]): The orchestrated schedules to compare. """ - metadata = input_data["metadata"] + metadata = input_data['metadata'] remove_count = 0 - self.info("Getting orchestrated schedules...", metadata=metadata) + self.info('Getting orchestrated schedules...', metadata=metadata) orchestrated_schedules = input_data.get('orchestrated_schedules', {}) @@ -100,20 +102,20 @@ class TemporalManager(BaseActivity): try: schedules = orchestrated_schedules.get(namespace, {}) - self.info( - f"Getting orchestrated schedules for {namespace}", metadata=metadata) + self.info(f'Getting orchestrated schedules for {namespace}', metadata=metadata) async for schedule in await client.list_schedules(): - search_attrs = getattr(schedule, "search_attributes", {}) - if search_attrs.get("orchestrated", ["false"]) == ["true"]: + search_attrs = getattr(schedule, 'search_attributes', {}) + if search_attrs.get('orchestrated', ['false']) == ['true']: schedule_id = schedule.id if schedule_id not in schedules: 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( - schedule_id) + handle = client.get_schedule_handle(schedule_id) await handle.delete() @@ -123,19 +125,18 @@ class TemporalManager(BaseActivity): trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="TEMPORAL_NORMALIZE_SCHEDULES_ERROR", - message=f"Failed to normalize schedules: {e}", - block="normalize_schedules", + notification_id='TEMPORAL_NORMALIZE_SCHEDULES_ERROR', + message=f'Failed to normalize schedules: {e}', + block='normalize_schedules', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) self.error(trace, metadata=metadata) raise e - self.info( - f"Removed {remove_count} schedules", metadata=metadata) + self.info(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]]: """ Create schedules in Temporal @@ -150,47 +151,41 @@ class TemporalManager(BaseActivity): """ schedules_to_create = input_data['schedules'] - metadata = input_data.get("metadata", {}) + metadata = input_data.get('metadata', {}) report = [] success_count = 0 - self.info("Creating schedules...", metadata=metadata) + self.info('Creating schedules...', metadata=metadata) for namespace, schedules in schedules_to_create.items(): client = self.temporal_clients.get(namespace) if not client: 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(): - search_attributes = TypedSearchAttributes([ - SearchAttributePair( - key=self.model_id_id_key, - value=schedule['model_id'] - ), - SearchAttributePair( - key=self.model_name_id_key, - value=schedule['model_name'] - ), - SearchAttributePair( - key=self.orchestrated_id_key, - value="true" - ) - ]) + search_attributes = TypedSearchAttributes( + [ + SearchAttributePair(key=self.model_id_id_key, value=schedule['model_id']), + SearchAttributePair( + key=self.model_name_id_key, value=schedule['model_name'] + ), + SearchAttributePair(key=self.orchestrated_id_key, value='true'), + ] + ) workflow_type = schedule['workflow_type'] try: - execution_timeout_seconds = schedule.get( - 'execution_timeout_seconds', 300) - task_timeout_seconds = schedule.get( - 'task_timeout_seconds', 300) + execution_timeout_seconds = schedule.get('execution_timeout_seconds', 300) + task_timeout_seconds = schedule.get('task_timeout_seconds', 300) + self.debug(f'Creating schedule {schedule_name}:', metadata=metadata) self.debug( - f"Creating schedule {schedule_name}:", metadata=metadata) - self.debug( - f"{json.dumps(schedule, indent=4, sort_keys=True)}", metadata=metadata) + f'{json.dumps(schedule, indent=4, sort_keys=True)}', metadata=metadata + ) await client.create_schedule( schedule_name, @@ -199,54 +194,57 @@ class TemporalManager(BaseActivity): workflow_type, schedule, id=schedule_name, - task_queue=f"{workflow_type}-queue", - execution_timeout=timedelta( - seconds=execution_timeout_seconds), - run_timeout=timedelta( - seconds=execution_timeout_seconds), - task_timeout=timedelta( - seconds=task_timeout_seconds), - typed_search_attributes=search_attributes + task_queue=f'{workflow_type}-queue', + execution_timeout=timedelta(seconds=execution_timeout_seconds), + run_timeout=timedelta(seconds=execution_timeout_seconds), + task_timeout=timedelta(seconds=task_timeout_seconds), + typed_search_attributes=search_attributes, ), spec=ScheduleSpec( intervals=[ ScheduleIntervalSpec( - every=timedelta(seconds=parse_frequency( - schedule.get('frequency', '1m'))) + every=timedelta( + seconds=parse_frequency(schedule.get('frequency', '1m')) + ) ) ] - ) + ), ), - search_attributes=search_attributes + search_attributes=search_attributes, ) - report.append({ - "namespace": namespace, - "schedule_name": schedule_name, - "success": True, - "message": "Schedule created successfully" - }) + report.append( + { + 'namespace': namespace, + 'schedule_name': schedule_name, + 'success': True, + 'message': 'Schedule created successfully', + } + ) success_count += 1 except Exception as e: self.error( - f"Failed to create schedule {schedule_name}: {str(e)}", metadata=metadata) - report.append({ - "namespace": namespace, - "schedule_name": schedule_name, - "success": False, - "message": str(e) - }) + f'Failed to create schedule {schedule_name}: {str(e)}', metadata=metadata + ) + report.append( + { + 'namespace': namespace, + 'schedule_name': schedule_name, + 'success': False, + 'message': str(e), + } + ) 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( - f"\n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata) + self.debug(f'\n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata) 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]]: """ Update schedules in Temporal @@ -261,27 +259,27 @@ class TemporalManager(BaseActivity): """ schedules_to_update = input_data['schedules'] - metadata = input_data.get("metadata", {}) + metadata = input_data.get('metadata', {}) report = [] success_count = 0 - self.info("Updating schedules...", metadata=metadata) + self.info('Updating schedules...', metadata=metadata) for namespace, schedules in schedules_to_update.items(): client = self.temporal_clients.get(namespace) if not client: 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(): try: - handler = client.get_schedule_handle( - schedule_name) + handler = client.get_schedule_handle(schedule_name) if not handler: - raise ValueError(f"Schedule {schedule_name} not found") + raise ValueError(f'Schedule {schedule_name} not found') # fmt: off async def update_schedule(input_data: ScheduleUpdateInput) -> ScheduleUpdate: # NOSONAR @@ -310,33 +308,38 @@ class TemporalManager(BaseActivity): del update_schedule - report.append({ - "namespace": namespace, - "schedule_name": schedule_name, - "success": True, - "message": "Schedule updated successfully" - }) + report.append( + { + 'namespace': namespace, + 'schedule_name': schedule_name, + 'success': True, + 'message': 'Schedule updated successfully', + } + ) success_count += 1 except Exception as e: self.error( - f"Failed to update schedule {schedule_name}: {str(e)}", metadata=metadata) - report.append({ - "namespace": namespace, - "schedule_name": schedule_name, - "success": False, - "message": str(e) - }) + f'Failed to update schedule {schedule_name}: {str(e)}', metadata=metadata + ) + report.append( + { + 'namespace': namespace, + 'schedule_name': schedule_name, + 'success': False, + 'message': str(e), + } + ) 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( - f"\n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata) + self.debug(f'\n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata) 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]]: """ Delete schedules in Temporal @@ -351,54 +354,59 @@ class TemporalManager(BaseActivity): """ schedules_to_delete = input_data['schedules'] - metadata = input_data.get("metadata", {}) + metadata = input_data.get('metadata', {}) report = [] success_count = 0 - self.info("Deleting schedules...", metadata=metadata) + self.info('Deleting schedules...', metadata=metadata) for namespace, schedules in schedules_to_delete.items(): client = self.temporal_clients.get(namespace) if not client: 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: try: - handler = client.get_schedule_handle( - schedule_name) + handler = client.get_schedule_handle(schedule_name) if not handler: - raise ValueError(f"Schedule {schedule_name} not found") + raise ValueError(f'Schedule {schedule_name} not found') await handler.delete() - report.append({ - "namespace": namespace, - "schedule_name": schedule_name, - "success": True, - "message": "Schedule deleted successfully" - }) + report.append( + { + 'namespace': namespace, + 'schedule_name': schedule_name, + 'success': True, + 'message': 'Schedule deleted successfully', + } + ) success_count += 1 except Exception as e: trace = traceback.format_exc() self.error( - f"Failed to delete schedule {schedule_name}: {str(e)}", metadata=metadata) - report.append({ - "namespace": namespace, - "schedule_name": schedule_name, - "success": False, - "message": str(e), - "attachment": trace - }) + f'Failed to delete schedule {schedule_name}: {str(e)}', metadata=metadata + ) + report.append( + { + 'namespace': namespace, + 'schedule_name': schedule_name, + 'success': False, + 'message': str(e), + 'attachment': trace, + } + ) 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( - f"\n {json.dumps(report, indent=4, sort_keys=True)}", metadata=metadata) + self.debug(f'\n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata) return report diff --git a/orchestrator/metrics.py b/orchestrator/metrics.py index 2a6a56b..cd7379e 100644 --- a/orchestrator/metrics.py +++ b/orchestrator/metrics.py @@ -6,19 +6,19 @@ orchestrator system including application health, email delivery, and workflow execution metrics. """ -from prometheus_client import Gauge, Counter +from prometheus_client import Counter, Gauge APP_UP = Gauge( - "app_up", - "Indicates if the application is running (1) or shutting down (0)", - ["pod_id"], + 'app_up', + 'Indicates if the application is running (1) or shutting down (0)', + ['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", - "Number of emails sent", - [*CORE_LABELS, "email_group"], + 'email_sent_count', + 'Number of emails sent', + [*CORE_LABELS, 'email_group'], ) diff --git a/orchestrator/utils/connectors_config.py b/orchestrator/utils/connectors_config.py index 2493bab..2e1466c 100644 --- a/orchestrator/utils/connectors_config.py +++ b/orchestrator/utils/connectors_config.py @@ -12,7 +12,7 @@ def build_redis_config(): 'host': getenv('REDIS_HOST', 'localhost'), 'port': int(getenv('REDIS_PORT', '6379')), 'username': getenv('REDIS_USERNAME', 'default'), - 'password': getenv('REDIS_PASSWORD', 'bdnZOpcyiL') + 'password': getenv('REDIS_PASSWORD', 'bdnZOpcyiL'), } @@ -31,7 +31,7 @@ def build_mongodb_config(): return { 'connection_string': connection_string, '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 { 'connection_string': getenv('COUCHBASE_CONNECTION_STRING', 'couchbase://localhost'), '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_scouter_namespace': getenv('TEMPORAL_SCOUTER_NAMESPACE', 'scouter'), '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'), 'dbname': getenv('POSTGRES_DBNAME', 'sientia'), '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_password': getenv('EMAIL_SENDER_PASSWORD', 'sientia'), 'smtp_server': getenv('EMAIL_SMTP_SERVER', None), - 'smtp_port': int(getenv('EMAIL_SMTP_PORT', '587')) + 'smtp_port': int(getenv('EMAIL_SMTP_PORT', '587')), } diff --git a/orchestrator/utils/converters.py b/orchestrator/utils/converters.py index eff3678..424dd24 100644 --- a/orchestrator/utils/converters.py +++ b/orchestrator/utils/converters.py @@ -9,7 +9,7 @@ def parse_frequency(frequency: str) -> int: Args: frequency (str): Frequency string with suffix: - 's' for seconds (e.g., '30s') - - 'm' for minutes (e.g., '5m') + - 'm' for minutes (e.g., '5m') - 'h' for hours (e.g., '2h') - 'd' for days (e.g., '1d') @@ -19,13 +19,13 @@ def parse_frequency(frequency: str) -> int: Raises: ValueError: If frequency format is invalid """ - if frequency.endswith("s"): + if frequency.endswith('s'): return int(frequency[:-1]) - elif frequency.endswith("m"): + elif frequency.endswith('m'): return int(frequency[:-1]) * 60 - elif frequency.endswith("h"): + elif frequency.endswith('h'): return int(frequency[:-1]) * 60 * 60 - elif frequency.endswith("d"): + elif frequency.endswith('d'): return int(frequency[:-1]) * 60 * 60 * 24 else: - raise ValueError("Invalid frequency") + raise ValueError('Invalid frequency') diff --git a/orchestrator/utils/email_builder.py b/orchestrator/utils/email_builder.py index 1d74ece..8c22203 100644 --- a/orchestrator/utils/email_builder.py +++ b/orchestrator/utils/email_builder.py @@ -1,8 +1,5 @@ -import json -from sientia_do.observability.logger import Logger -from sientia_do.notifications.models import NotificationLevel from jinja2 import Template -import re +from sientia_do.observability.logger import Logger class EmailBuilder: @@ -24,9 +21,9 @@ class EmailBuilder: self.report_template_file = './orchestrator/utils/templates/email_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() - with open(self.general_template_file, 'r') as file: + with open(self.general_template_file) as file: self.general_template = file.read() def replace_parameters(self, template: str, parameters: dict) -> str: @@ -41,9 +38,9 @@ class EmailBuilder: str: The rendered template with parameters replaced. """ # 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: """ @@ -57,21 +54,25 @@ class EmailBuilder: Returns: dict: Dictionary with mail_type and rendered event sections for each notification level. """ - error_models = general_events.get('ERROR', {}).get('models', []) - warning_models = general_events.get('WARNING', {}).get('models', []) - info_models = general_events.get('INFO', {}).get('models', []) + error_events = general_events.get('ERROR', {}) + warning_events = general_events.get('WARNING', {}) + 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 { 'mail_type': mail_type, - 'error_events': self.replace_parameters(self.general_template, - general_events.get( - 'ERROR')) if error_models else '', - 'warning_events': self.replace_parameters(self.general_template, - general_events.get( - 'WARNING')) if warning_models else '', - 'info_events': self.replace_parameters(self.general_template, - general_events.get( - 'INFO')) if info_models else '', + 'error_events': self.replace_parameters(self.general_template, error_events) + if error_models + else '', + 'warning_events': self.replace_parameters(self.general_template, warning_events) + if warning_models + else '', + 'info_events': self.replace_parameters(self.general_template, info_events) + if info_models + else '', } def build_email(self, report_data: list[dict], mail_type: str) -> str: @@ -90,29 +91,26 @@ class EmailBuilder: general_events = {} for report in report_data: - level = report['level'] model_name = report['model_name'] if level not in general_events: general_events[level] = { 'section_name': f'{level.capitalize()}s detected:', - 'models': {} + 'models': {}, } if model_name not in general_events[level]['models']: general_events[level]['models'][model_name] = { 'model_name': model_name, - 'events': [] + 'events': [], } - general_events[level]['models'][model_name]['events'].append( - report) + general_events[level]['models'][model_name]['events'].append(report) for _type, content in general_events.items(): content['models'] = list(content['models'].values()) return self.replace_parameters( - self.report_template, self.parameters( - general_events, mail_type - )) + self.report_template, self.parameters(general_events, mail_type) + ) diff --git a/orchestrator/utils/orchestrator_functions.py b/orchestrator/utils/orchestrator_functions.py index 38ca1f4..c6e5cc0 100644 --- a/orchestrator/utils/orchestrator_functions.py +++ b/orchestrator/utils/orchestrator_functions.py @@ -19,17 +19,15 @@ def common_config(config: dict[str, Any]): """ model = config['model'] return { - "workflow_type": config['workflow_type'], - "schedule_name": config['schedule_name'], - "frequency": config.get('frequency', '1m'), - "max_retry_policy": config.get('max_retry_policy', 1), - - "model_id": config['model_id'], - "model_name": model['name'], - "model_config": model.get('model_config', {}), - - "execution_timeout_seconds": config.get('execution_timeout_seconds', 300), - "task_timeout_seconds": config.get('task_timeout_seconds', 300), + 'workflow_type': config['workflow_type'], + 'schedule_name': config['schedule_name'], + 'frequency': config.get('frequency', '1m'), + 'max_retry_policy': config.get('max_retry_policy', 1), + 'model_id': config['model_id'], + 'model_name': model['name'], + 'model_config': model.get('model_config', {}), + '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 { **common_config(config), - "workflow_type": "minimal_retrain", - "schedule_name": config['schedule_name'], - "query": config['query'], - "schema": "sientia_data", - "table_name": "log_retrain", - "datetime_columns": config.get('datetime_columns', []), + 'workflow_type': 'minimal_retrain', + 'schedule_name': config['schedule_name'], + 'query': config['query'], + 'schema': 'sientia_data', + 'table_name': 'log_retrain', + 'datetime_columns': config.get('datetime_columns', []), } @@ -79,28 +77,25 @@ def scouter(config: dict[str, Any]): """ filters = {} for f in config.get('filters', []): - filters[f['filter_name']] = { - "policy": f['policy'] - } + filters[f['filter_name']] = {'policy': f['policy']} tags = {} for tag in config['read_tags']: tags[tag['tag_name']] = { - "aggr_func": tag.get('aggr_func', 'lts'), - "data_range": tag.get('data_range', [-100, 100]) + 'aggr_func': tag.get('aggr_func', 'lts'), + 'data_range': tag.get('data_range', [-100, 100]), } return { **common_config(config), - - "topic": f"raw_{config['schedule_name']}", - "trigger_laborious": False, - "filters": filters, - "schema": "sientia_data", - "table_name": "laborious_data", - "retention_time": config.get('tag_retention_minutes', 60) * 60, - "model_tags": tags, - "debug_data_package": config.get('debug_data_package', False) + 'topic': f'raw_{config["schedule_name"]}', + 'trigger_laborious': False, + 'filters': filters, + 'schema': 'sientia_data', + 'table_name': 'laborious_data', + 'retention_time': config.get('tag_retention_minutes', 60) * 60, + 'model_tags': tags, + '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: base_filter_config[fil['filter_name']] = { - "policy": fil['policy'], - "config": fil.get('config', {}) + 'policy': fil['policy'], + 'config': fil.get('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"]. """ for priority in path_priority[:]: - if priority not in ["STOP", "CONTINUE", "REPEAT"]: + if priority not in ['STOP', 'CONTINUE', 'REPEAT']: path_priority.remove(priority) - for priority in ["STOP", "CONTINUE", "REPEAT"]: + for priority in ['STOP', 'CONTINUE', 'REPEAT']: if priority not in path_priority: path_priority.append(priority) @@ -169,7 +164,7 @@ def predictions_batch(config: dict[str, Any]): Returns: 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', []): if tag['server_id'] not in tags: tags[tag['server_id']] = {} @@ -177,51 +172,43 @@ def predictions_batch(config: dict[str, Any]): tag_type = tag['type'] 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']]: tags[tag['server_id']][tag_type_str] = {} 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', ["STOP", "CONTINUE", "REPEAT"])) + path_priority = process_path_priority( + config.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']) + ) return { **common_config(config), - - "query": config['query'], - "datetime_columns": config.get('datetime_columns', []), - "schema": "sientia_data", - "table_name": "predictions", - "retention_time": config.get('model_retention_minutes', 60) * 60, - "opc_output_config": tags, - "input_filters": overlap_filter_config({ - "EMPTY_DATA": { - "policy": "STOP", - "config": {} - } - }, config.get('input_filters', [])), - "mlflow_transform_filters": overlap_filter_config({ - "EMPTY_DATA": { - "policy": "STOP", - "config": {} + 'query': config['query'], + 'datetime_columns': config.get('datetime_columns', []), + 'schema': 'sientia_data', + 'table_name': 'predictions', + 'retention_time': config.get('model_retention_minutes', 60) * 60, + 'opc_output_config': tags, + 'input_filters': overlap_filter_config( + {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, config.get('input_filters', []) + ), + 'mlflow_transform_filters': overlap_filter_config( + { + 'EMPTY_DATA': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'policy': 'STOP', 'config': {}}, }, - "API_ERROR": { - "policy": "STOP", - "config": {} - } - }, config.get('mlflow_transform_filters', [])), - "mlflow_predict_filters": overlap_filter_config({ - "API_ERROR": { - "policy": "STOP", - "config": {} - } - }, config.get('mlflow_predict_filters', [])), - "path_priority": path_priority, - "predictions_storage_policy": config.get('predictions_storage_policy', 'lts:1') + config.get('mlflow_transform_filters', []), + ), + 'mlflow_predict_filters': overlap_filter_config( + {'API_ERROR': {'policy': 'STOP', '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 for pipeline in pipelines: 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: - tags[tag_string] = { - **tag, - "topics": [] - } + tags[tag_string] = {**tag, 'topics': []} - tags[tag_string]['topics'].append( - f"raw_{pipeline['schedule_name']}") + tags[tag_string]['topics'].append(f'raw_{pipeline["schedule_name"]}') return tags -def build_tag_config(tag: dict[str, Any], slot_config: dict[str, Any], - opc_servers: dict[str, Any], i: int): +def build_tag_config( + 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. @@ -276,22 +260,22 @@ def build_tag_config(tag: dict[str, Any], slot_config: dict[str, Any], server_id = tag['server_id'] 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'] - if server_name not in slot_config[f"{i}"]: - slot_config[f"{i}"][server_name] = { - "server_id": server_id, - "name": server_name, - "url": opc_servers[server_id]['url'], - "server_uri": opc_servers[server_id]['uri'], - "cert_path": opc_servers[server_id].get('cert_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), - "tags": {} + if server_name not in slot_config[f'{i}']: + slot_config[f'{i}'][server_name] = { + 'server_id': server_id, + 'name': server_name, + 'url': opc_servers[server_id]['url'], + 'server_uri': opc_servers[server_id]['uri'], + 'cert_path': opc_servers[server_id].get('cert_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), + 'tags': {}, } - slot_config[f"{i}"][server_name]["tags"][tag['tag_address']] = { + slot_config[f'{i}'][server_name]['tags'][tag['tag_address']] = { **tag, } diff --git a/orchestrator/worker/worker.py b/orchestrator/worker/worker.py index 9439a38..6283cbd 100644 --- a/orchestrator/worker/worker.py +++ b/orchestrator/worker/worker.py @@ -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.runtime import Runtime, TelemetryConfig, PrometheusConfig with workflow.unsafe.imports_passed_through(): + import asyncio import os import sys - import asyncio - from orchestrator.workflows.alerts import Alerts - 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 prometheus_client import start_http_server from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler 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") -SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091")) + from orchestrator import metrics + 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(): @@ -49,10 +53,9 @@ async def main(): 'schedule_name': '-', } - logger.custom_info( - f'Starting Worker with POD_ID: {POD_ID}', metadata=metadata) + logger.custom_info(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() logger.custom_info('Starting Notification Handler...', metadata=metadata) @@ -66,22 +69,21 @@ async def main(): ) 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( telemetry=TelemetryConfig( - metrics=PrometheusConfig( - bind_address=f"0.0.0.0:{SDK_METRICS_PORT}") + metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}') ) ) - logger.custom_info( - f'Starting Temporal Client at {host}:{namespace}', metadata=metadata) + logger.custom_info(f'Starting Temporal Client at {host}:{namespace}', metadata=metadata) temporal_client = await client.Client.connect( target_host=host, namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'), - runtime=new_runtime + runtime=new_runtime, ) logger.custom_info('Starting Activities...', metadata=metadata) @@ -93,7 +95,7 @@ async def main(): email_config=build_email_config(), postgres_config=build_postgres_config(), logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) await activities.connect_to_temporal() @@ -133,7 +135,7 @@ async def main(): activities.report_schedule_orchestration, activities.report_slot_orchestration, activities.format_schedule_config, - ] + ], ), Worker( temporal_client, @@ -145,19 +147,16 @@ async def main(): activities.find_documents_in_mongodb, activities.load_latest_data, activities.put_last_data_timestamp, - # Format and filter notifications activities.filter_notification_alerts, - # Send email and export data to postgres activities.build_email_html, activities.send_email, activities.format_log_report, activities.export_data_to_postgres, - # Store notification cache - activities.store_notification_cache - ] + activities.store_notification_cache, + ], ), Worker( temporal_client, @@ -169,17 +168,15 @@ async def main(): activities.find_documents_in_mongodb, activities.load_latest_data, activities.put_last_data_timestamp, - # Format and filter notifications activities.filter_notification_reports, - # Send email and export data to postgres activities.build_email_html, activities.send_email, activities.format_log_report, - activities.export_data_to_postgres - ] - ) + activities.export_data_to_postgres, + ], + ), ] handlers = [] @@ -193,7 +190,7 @@ async def main(): # If an exception occurs in any of the worker handlers, it will be propagated here. await asyncio.gather(*handlers) except BaseException as e: - logger.error(f"An unhandled exception occurred: {e}", exc_info=True) + logger.error(f'An unhandled exception occurred: {e}', exc_info=True) finally: if notification_handler: notification_handler.shutdown() @@ -212,12 +209,12 @@ def start_prometheus_server(): Exits the application if the server fails to start. """ try: - port = int(os.getenv("HTTP_METRICS_PORT", 9090)) + port = int(os.getenv('HTTP_METRICS_PORT', 9090)) 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) except Exception as e: - print(f"Failed to start Prometheus server: {e}") + print(f'Failed to start Prometheus server: {e}') os._exit(1) diff --git a/orchestrator/workflows/alerts.py b/orchestrator/workflows/alerts.py index 5f2ec31..436ebfa 100644 --- a/orchestrator/workflows/alerts.py +++ b/orchestrator/workflows/alerts.py @@ -1,13 +1,15 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from orchestrator.activities.activities import Activities - from typing import Any 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="alerts") + +@workflow.defn(name='alerts') class Alerts: """ Alerts workflow for real-time error notification delivery. @@ -51,26 +53,21 @@ class Alerts: 'schedule_name': input_data['schedule_name'], 'workflow_name': 'alerts', 'model_name': '-', - 'model_id': '-' + 'model_id': '-', } } - mail_type = "Alerts" + mail_type = 'Alerts' input_data['metadata'] = metadata input_data['mail_type'] = mail_type - input_data['base_data_filter'] = { - 'level': 'ERROR' - } + input_data['base_data_filter'] = {'level': 'ERROR'} # Call subworkflow "load_notification_package" passing the static filters # (level = "ERROR" and timestamp > last timestamp) - package = await workflow.execute_child_workflow( - 'load_notification_package', - input_data - ) + package = await workflow.execute_child_workflow('load_notification_package', input_data) if not package['notification_package'] or not package['sending_configs']: return @@ -84,10 +81,10 @@ class Alerts: **metadata, 'notification_package': package['notification_package'], 'sending_configs': package['sending_configs'], - 'notification_ttl': input_data['notification_ttl'] + 'notification_ttl': input_data['notification_ttl'], }, schedule_to_close_timeout=timedelta(seconds=60), - retry_policy=retry_policy + retry_policy=retry_policy, ) if not receiver_groups: @@ -101,8 +98,8 @@ class Alerts: 'mail_type': mail_type, 'notification_package': receiver_groups, 'schema': 'sientia_data', - 'table_name': 'log_report' - } + 'table_name': 'log_report', + }, ) if not log_report: @@ -111,11 +108,7 @@ class Alerts: # Store the notification_id sendings to avoid sending them again await workflow.execute_activity_method( 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), - retry_policy=retry_policy + retry_policy=retry_policy, ) diff --git a/orchestrator/workflows/orchestrator.py b/orchestrator/workflows/orchestrator.py index d6b8ede..a0fe690 100644 --- a/orchestrator/workflows/orchestrator.py +++ b/orchestrator/workflows/orchestrator.py @@ -1,13 +1,15 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from orchestrator.activities.activities import Activities - from typing import Any 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="orchestrator") + +@workflow.defn(name='orchestrator') class Orchestrator: """ Main orchestrator workflow for pipeline and resource management. @@ -49,7 +51,7 @@ class Orchestrator: 'schedule_name': input_data.get('schedule_name', 'orchestrator'), 'model_name': '-', 'model_id': '-', - 'workflow_name': input_data['workflow_name'] + 'workflow_name': input_data['workflow_name'], } } @@ -58,33 +60,28 @@ class Orchestrator: { **metadata, 'query': input_data['pipelines_query'], - "timestamp_fields": ["updated_at"] + 'timestamp_fields': ['updated_at'], }, 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( Activities.find_documents_in_mongodb, - { - **metadata, - 'query': input_data['opc_servers_query'] - }, + {**metadata, 'query': input_data['opc_servers_query']}, 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( Activities.find_documents_in_mongodb, { **metadata, - 'query': { - 'collection': 'orchestrated_schedules' - }, - "timestamp_fields": ["updated_at"] + 'query': {'collection': 'orchestrated_schedules'}, + 'timestamp_fields': ['updated_at'], }, 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( @@ -93,7 +90,7 @@ class Orchestrator: **metadata, }, 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( @@ -102,7 +99,7 @@ class Orchestrator: **metadata, }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) pipeline_config = await pipeline_config_handler @@ -113,22 +110,16 @@ class Orchestrator: formatted_orchestrated_schedules_handler = workflow.start_local_activity_method( Activities.format_schedule_config, - { - **metadata, - 'schedule_config': orchestrated_schedules - }, + {**metadata, 'schedule_config': orchestrated_schedules}, 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( Activities.process_schedules, - { - **metadata, - 'pipelines': pipeline_config - }, + {**metadata, 'pipelines': pipeline_config}, 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( @@ -140,7 +131,7 @@ class Orchestrator: 'pipelines': pipeline_config, }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) schedules_config = await schedules_config_handler @@ -152,41 +143,31 @@ class Orchestrator: { **metadata, 'current_schedule_config': formatted_orchestrated_schedules, - 'schedule_config': schedules_config + 'schedule_config': schedules_config, }, 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( 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, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) normalize_schedules_handler = workflow.start_activity_method( Activities.normalize_schedules, - { - **metadata, - 'orchestrated_schedules': formatted_orchestrated_schedules - }, + {**metadata, 'orchestrated_schedules': formatted_orchestrated_schedules}, 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( Activities.create_collection_with_ttl_index, - { - **metadata, - 'pipelines': schedules_config['scouter'] - }, + {**metadata, 'pipelines': schedules_config['scouter']}, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) schedule_actions = await schedule_actions_handler @@ -196,52 +177,37 @@ class Orchestrator: slot_deletion_report_handler = workflow.start_activity_method( Activities.delete_slots, - { - **metadata, - 'to_delete': slot_actions['to_delete'] - }, + {**metadata, 'to_delete': slot_actions['to_delete']}, 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( Activities.update_slots, - { - **metadata, - 'to_insert': slot_actions['to_insert'] - }, + {**metadata, 'to_insert': slot_actions['to_insert']}, 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( Activities.delete_schedules, - { - **metadata, - 'schedules': schedule_actions['to_delete'] - }, + {**metadata, 'schedules': schedule_actions['to_delete']}, 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( Activities.create_schedules, - { - **metadata, - 'schedules': schedule_actions['to_create'] - }, + {**metadata, 'schedules': schedule_actions['to_create']}, 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( Activities.update_schedules, - { - **metadata, - 'schedules': schedule_actions['to_update'] - }, + {**metadata, 'schedules': schedule_actions['to_update']}, 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 @@ -251,66 +217,52 @@ class Orchestrator: schedule_update_report = await schedule_update_report_handler if schedule_insertion_report or schedule_update_report or schedule_deletion_report: - schedule_report_handler = workflow.start_activity_method( Activities.report_schedule_orchestration, { **metadata, 'created_schedules': schedule_insertion_report, 'updated_schedules': schedule_update_report, - 'deleted_schedules': schedule_deletion_report + 'deleted_schedules': schedule_deletion_report, }, 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: - slot_report_handler = workflow.start_activity_method( Activities.report_slot_orchestration, { **metadata, 'inserted_slots': slot_insertion_report, - 'deleted_slots': slot_deletion_report + 'deleted_slots': slot_deletion_report, }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) if schedule_update_report: - update_pipelines_timestamps_handler = workflow.start_activity_method( Activities.update_pipelines_timestamps, - { - **metadata, - 'updated_pipelines': schedule_update_report - }, + {**metadata, 'updated_pipelines': schedule_update_report}, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) if schedule_insertion_report: - create_pipelines_timestamps_handler = workflow.start_activity_method( Activities.create_pipelines_timestamps, - { - **metadata, - 'created_pipelines': schedule_insertion_report - }, + {**metadata, 'created_pipelines': schedule_insertion_report}, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) if schedule_deletion_report: - delete_pipelines_timestamps_handler = workflow.start_activity_method( Activities.delete_pipelines_timestamps, - { - **metadata, - 'deleted_pipelines': schedule_deletion_report - }, + {**metadata, 'deleted_pipelines': schedule_deletion_report}, 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: diff --git a/orchestrator/workflows/reports.py b/orchestrator/workflows/reports.py index e942296..20be848 100644 --- a/orchestrator/workflows/reports.py +++ b/orchestrator/workflows/reports.py @@ -1,13 +1,15 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from orchestrator.activities.activities import Activities - from typing import Any 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="reports") + +@workflow.defn(name='reports') class Reports: """ Reports workflow for sending scheduled notification summaries. @@ -43,11 +45,11 @@ class Reports: 'schedule_name': input_data['schedule_name'], 'workflow_name': 'reports', 'model_name': '-', - 'model_id': '-' + 'model_id': '-', } } - mail_type = "Reports" + mail_type = 'Reports' input_data['metadata'] = metadata input_data['mail_type'] = mail_type @@ -57,10 +59,7 @@ class Reports: # Call subworkflow "load_notification_package" passing the static filters # (timestamp > last timestamp) - package = await workflow.execute_child_workflow( - 'load_notification_package', - input_data - ) + package = await workflow.execute_child_workflow('load_notification_package', input_data) if not package['notification_package'] or not package['sending_configs']: return @@ -73,10 +72,10 @@ class Reports: { **metadata, 'notification_package': package['notification_package'], - 'sending_configs': package['sending_configs'] + 'sending_configs': package['sending_configs'], }, schedule_to_close_timeout=timedelta(seconds=60), - retry_policy=retry_policy + retry_policy=retry_policy, ) if not receiver_groups: @@ -90,6 +89,6 @@ class Reports: 'mail_type': mail_type, 'notification_package': receiver_groups, 'schema': 'sientia_data', - 'table_name': 'log_report' - } + 'table_name': 'log_report', + }, ) diff --git a/orchestrator/workflows/subworkflows/load_notification_package.py b/orchestrator/workflows/subworkflows/load_notification_package.py index 9e2d35b..f16592e 100644 --- a/orchestrator/workflows/subworkflows/load_notification_package.py +++ b/orchestrator/workflows/subworkflows/load_notification_package.py @@ -1,13 +1,15 @@ from temporalio import workflow 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 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: """ Subworkflow for loading notification data and configuration. @@ -48,28 +50,17 @@ class LoadNotificationPackage: # Load last timestamp from redis "notification_last_timestamp" last_timestamp_handler = workflow.start_local_activity_method( 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), - retry_policy=retry_policy + retry_policy=retry_policy, ) # In parallel, load sending configs from collection "receiver_groups" sending_configs_handler = workflow.start_local_activity_method( 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), - retry_policy=retry_policy + retry_policy=retry_policy, ) last_timestamp = await last_timestamp_handler @@ -83,10 +74,10 @@ class LoadNotificationPackage: **metadata, 'collection_name': 'notification_queue', '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), - retry_policy=retry_policy + retry_policy=retry_policy, ) sending_configs = await sending_configs_handler @@ -94,20 +85,16 @@ class LoadNotificationPackage: return { 'last_timestamp': last_timestamp, 'notification_package': notification_package, - 'sending_configs': sending_configs + 'sending_configs': sending_configs, } # Put last collected timestamp in redis "notification_last_timestamp" await workflow.start_activity_method( 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), - retry_policy=retry_policy + retry_policy=retry_policy, ) # Return a dict with the following keys: @@ -117,5 +104,5 @@ class LoadNotificationPackage: return { 'last_timestamp': last_timestamp, 'notification_package': notification_package, - 'sending_configs': sending_configs + 'sending_configs': sending_configs, } diff --git a/orchestrator/workflows/subworkflows/process_notifications.py b/orchestrator/workflows/subworkflows/process_notifications.py index 7809ac1..7e25641 100644 --- a/orchestrator/workflows/subworkflows/process_notifications.py +++ b/orchestrator/workflows/subworkflows/process_notifications.py @@ -1,14 +1,16 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from orchestrator.activities.activities import Activities - from typing import Any 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.policies import retry_policy + + from orchestrator.activities.activities import Activities -@workflow.defn(name="process_notifications") +@workflow.defn(name='process_notifications') class ProcessNotifications: """ Subworkflow for processing and sending notification emails. @@ -43,30 +45,26 @@ class ProcessNotifications: 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 data_to_sent = await workflow.execute_local_activity_method( Activities.build_email_html, { **metadata, - "receiver_groups": input_data["notification_package"], - "mail_type": input_data["mail_type"] + 'receiver_groups': input_data['notification_package'], + 'mail_type': input_data['mail_type'], }, 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 log_report = await workflow.execute_activity_method( 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), - retry_policy=retry_policy + retry_policy=retry_policy, ) if not log_report: @@ -75,13 +73,9 @@ class ProcessNotifications: # Format the log report to a dataframe to be stored in the database log_report = await workflow.execute_local_activity_method( 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), - retry_policy=retry_policy + retry_policy=retry_policy, ) # Store sending log in postgres database "log_report" @@ -89,16 +83,16 @@ class ProcessNotifications: Activities.export_data_to_postgres, { **metadata, - "schema": input_data["schema"], - "table_name": input_data["table_name"], - "data": log_report, + 'schema': input_data['schema'], + 'table_name': input_data['table_name'], + 'data': log_report, 'timestamp_conversion': { 'column': 'timestamp', - 'format': DATETIME_FORMAT_MS_WITH_TZ - } + 'format': DATETIME_FORMAT_MS_WITH_TZ, + }, }, schedule_to_close_timeout=timedelta(seconds=60), - retry_policy=retry_policy + retry_policy=retry_policy, ) # Return the log report to the caller diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..471e7c4 --- /dev/null +++ b/pyproject.toml @@ -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 \ No newline at end of file diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..56ab376 --- /dev/null +++ b/requirements-dev.txt @@ -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 \ No newline at end of file diff --git a/tests/orchestrator/activities/test_activities.py b/tests/orchestrator/activities/test_activities.py index 9a0c1ff..fb5cb3b 100644 --- a/tests/orchestrator/activities/test_activities.py +++ b/tests/orchestrator/activities/test_activities.py @@ -1,11 +1,10 @@ -from unittest.mock import patch, MagicMock, ANY -from pytest import mark -from orchestrator.activities import mongo_db +from unittest.mock import ANY, MagicMock, patch + 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.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__') @@ -14,37 +13,33 @@ from orchestrator.activities.formatters import Formatters @patch('orchestrator.activities.formatters.Formatters.__init__') @patch('orchestrator.activities.email.Email.__init__') @patch('sientia_do.temporal.activities.postgres.Postgres.__init__') -def test___init__(mock_postgres_init, - mock_email_init, - mock_formatters_init, - mock_slot_manager_init, - mock_temporal_manager_init, - mock_mongodb_init): - +def test___init__( + mock_postgres_init, + mock_email_init, + mock_formatters_init, + mock_slot_manager_init, + mock_temporal_manager_init, + mock_mongodb_init, +): mongo_db_config = { 'connection_string': 'mongodb://localhost:27017', 'database_name': 'test_db', - 'ttl_index_seconds': 3600 + 'ttl_index_seconds': 3600, } - redis_config = { - 'host': 'localhost', - 'port': 6379, - 'username': 'admin', - 'password': 'password' - } + redis_config = {'host': 'localhost', 'port': 6379, 'username': 'admin', 'password': 'password'} temporal_config = { 'temporal_host': 'localhost', 'temporal_scouter_namespace': 'scouter', - 'temporal_laborious_namespace': 'laborious' + 'temporal_laborious_namespace': 'laborious', } email_config = { 'sender_email': 'test@test.com', 'sender_password': 'test', 'smtp_server': 'test', - 'smtp_port': 587 + 'smtp_port': 587, } postgres_config = { @@ -67,7 +62,7 @@ def test___init__(mock_postgres_init, email_config=email_config, postgres_config=postgres_config, logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) assert isinstance(activities, Activities) @@ -78,12 +73,12 @@ def test___init__(mock_postgres_init, mock_slot_manager_init.assert_called_once_with( ANY, - host="localhost", + host='localhost', port=6379, - username="admin", - password="password", + username='admin', + password='password', logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) mock_mongodb_init.assert_called_once_with( @@ -92,7 +87,7 @@ def test___init__(mock_postgres_init, database_name='test_db', ttl_index_seconds=3600, logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) mock_temporal_manager_init.assert_called_once_with( @@ -101,7 +96,7 @@ def test___init__(mock_postgres_init, scouter_namespace='scouter', laborious_namespace='laborious', logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) mock_formatters_init.assert_called_once_with( @@ -109,7 +104,7 @@ def test___init__(mock_postgres_init, scouter_namespace='scouter', laborious_namespace='laborious', 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('sientia_do.temporal.activities.postgres.Postgres.close') @patch('orchestrator.activities.mongo_db.MongoDB.shutdown') -def test_shutdown(mock_mongodb_close, - mock_postgres_shutdown, - mock_email_close, - mock_postgres_init, - mock_email_init, - mock_formatters_init, - mock_slot_manager_init, - mock_temporal_manager_init, - mock_mongodb_init, - ): - +def test_shutdown( + mock_mongodb_close, + mock_postgres_shutdown, + mock_email_close, + mock_postgres_init, + mock_email_init, + mock_formatters_init, + mock_slot_manager_init, + mock_temporal_manager_init, + mock_mongodb_init, +): activities = Activities( temporal_config=MagicMock(), redis_config=MagicMock(), @@ -140,7 +135,7 @@ def test_shutdown(mock_mongodb_close, email_config=MagicMock(), postgres_config=MagicMock(), logger=MagicMock(), - notification_handler=MagicMock() + notification_handler=MagicMock(), ) activities.shutdown() diff --git a/tests/orchestrator/activities/test_couchbase.py b/tests/orchestrator/activities/test_couchbase.py index 7f78ec5..8fd0098 100644 --- a/tests/orchestrator/activities/test_couchbase.py +++ b/tests/orchestrator/activities/test_couchbase.py @@ -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 sientia_do.notifications.models import NotificationLevel + from orchestrator.activities.couchbase import Couchbase @fixture -@patch("orchestrator.activities.couchbase.Cluster") +@patch('orchestrator.activities.couchbase.Cluster') def couchbase(_cluster_mock): return Couchbase( - connection_string="couchbase://localhost", - username="admin", - password="password", + connection_string='couchbase://localhost', + username='admin', + password='password', logger=MagicMock(), notification_handler=MagicMock(), ) @@ -22,27 +24,30 @@ def test_shutdown_success(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.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 async def test_load_query_from_couchbase_success(couchbase): couchbase.cluster.query.return_value.rows.return_value = [ - {"id": "1", "name": "test"}, - {"id": "2", "name": "test2"}, + {'id': '1', 'name': 'test'}, + {'id': '2', 'name': 'test2'}, ] - query = "SELECT * FROM bucket" + query = 'SELECT * FROM bucket' - result = await couchbase.load_query_from_couchbase({ - "query": query, - }) + result = await couchbase.load_query_from_couchbase( + { + 'query': query, + } + ) assert result == [ - {"id": "1", "name": "test"}, - {"id": "2", "name": "test2"}, + {'id': '1', 'name': 'test'}, + {'id': '2', 'name': 'test2'}, ] couchbase.cluster.query.assert_called_once_with(query) 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 async def test_load_query_from_couchbase_failure(couchbase): - couchbase.cluster.query.side_effect = Exception("Test error") - query = "SELECT * FROM bucket" + couchbase.cluster.query.side_effect = ValueError('Test error') + query = 'SELECT * FROM bucket' - with raises(Exception): - await couchbase.load_query_from_couchbase({ - "query": query, - }) + with raises(ValueError): + await couchbase.load_query_from_couchbase( + { + 'query': query, + } + ) couchbase.cluster.query.assert_called_once_with(query) couchbase.notification_handler.build_and_send_notification.assert_called_once_with( - notification_id="COUCHBASE_LOAD_QUERY_ERROR", - message="Failed to execute couchbase query: Test error", - block="load_query_from_couchbase", + notification_id='COUCHBASE_LOAD_QUERY_ERROR', + message='Failed to execute couchbase query: Test error', + block='load_query_from_couchbase', level=NotificationLevel.ERROR, attachment_content=ANY, ) diff --git a/tests/orchestrator/activities/test_email.py b/tests/orchestrator/activities/test_email.py index 19c3408..59b3701 100644 --- a/tests/orchestrator/activities/test_email.py +++ b/tests/orchestrator/activities/test_email.py @@ -1,6 +1,7 @@ from smtplib import SMTPServerDisconnected from unittest.mock import MagicMock, call, patch -from pytest import mark, fixture + +from pytest import fixture, mark from orchestrator.activities.email import Email @@ -10,12 +11,12 @@ from orchestrator.activities.email import Email @patch('orchestrator.activities.email.smtplib') def email(smtplib, email_builder): email = Email( - sender_email="test@test.com", - sender_password="test", - smtp_server="test", + sender_email='test@test.com', + sender_password='test', + smtp_server='test', smtp_port=587, logger=MagicMock(), - notification_handler=MagicMock() + notification_handler=MagicMock(), ) email_builder.send_notification = MagicMock() @@ -26,22 +27,21 @@ def email(smtplib, email_builder): @patch('orchestrator.activities.email.smtplib') def test___init___with_password(smtplib, email_builder): email = Email( - sender_email="test@test.com", - sender_password="test", - smtp_server="test", + sender_email='test@test.com', + sender_password='test', + smtp_server='test', smtp_port=587, logger=MagicMock(), - notification_handler=MagicMock() + notification_handler=MagicMock(), ) - assert email.sender_email == "test@test.com" - assert email.sender_password == "test" + assert email.sender_email == 'test@test.com' + assert email.sender_password == 'test' 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.login.assert_called_once_with( - "test@test.com", "test") + smtplib.SMTP.return_value.login.assert_called_once_with('test@test.com', 'test') assert email.server == smtplib.SMTP.return_value @@ -50,28 +50,28 @@ def test___init___with_password(smtplib, email_builder): @patch('orchestrator.activities.email.smtplib') def test___init___without_password(smtplib, email_builder): email = Email( - sender_email="test@test.com", + sender_email='test@test.com', sender_password=None, - smtp_server="test", + smtp_server='test', smtp_port=587, 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.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 metadata = { - "metadata": { - "schedule_name": "test", - "model_name": "test", - "model_id": "test", - "workflow_name": "test", + 'metadata': { + 'schedule_name': 'test', + 'model_name': 'test', + 'model_id': 'test', + 'workflow_name': 'test', } } @@ -84,53 +84,42 @@ def test_shutdown(email): @mark.asyncio async def test_build_email_html(email): - email.email_builder.build_email = MagicMock( - return_value="test" - ) + email.email_builder.build_email = MagicMock(return_value='test') input_data = { **metadata, - "receiver_groups": { - "group_1": { - "notifications": [ - { - "type": "test", - "subject": "test", - "body": "test" - }, - { - "type": "test", - "subject": "test", - "body": "test" - } + 'receiver_groups': { + 'group_1': { + 'notifications': [ + {'type': '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) assert response == { - "group_1": { - "notifications": [ + 'group_1': { + 'notifications': [ { - "type": "test", - "subject": "test", - "body": "test", + 'type': 'test', + 'subject': 'test', + 'body': 'test', }, { - "type": "test", - "subject": "test", - "body": "test", - } + 'type': 'test', + 'subject': 'test', + 'body': 'test', + }, ], - "html": "test" + 'html': 'test', } } email.email_builder.build_email.assert_called_once_with( - input_data['receiver_groups']['group_1']['notifications'], - input_data['mail_type'] + input_data['receiver_groups']['group_1']['notifications'], input_data['mail_type'] ) @@ -140,18 +129,9 @@ def test_handle_attachments_success(encoders, mime_base, email): message = MagicMock() attachments = [ - { - "filename": "file_1", - "attachment_content": "test_content_1" - }, - { - "filename": "file_2", - "attachment_content": "test_content_2" - }, - { - "filename": "file_3", - "attachment_content": "test_content_3" - } + {'filename': 'file_1', 'attachment_content': 'test_content_1'}, + {'filename': 'file_2', 'attachment_content': 'test_content_2'}, + {'filename': 'file_3', 'attachment_content': 'test_content_3'}, ] 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( [ - call("test_content_1".encode('utf-8')), - call("test_content_2".encode('utf-8')), - call("test_content_3".encode('utf-8')) + call(b'test_content_1'), + call(b'test_content_2'), + 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_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): message = MagicMock() - mime_base.side_effect = Exception("test") + mime_base.side_effect = Exception('test') - attachments = [ - { - "filename": "file_1", - "attachment_content": "test_content_1" - } - ] + attachments = [{'filename': 'file_1', 'attachment_content': 'test_content_1'}] try: email.handle_attachments(attachments, message) except Exception as e: - assert str(e) == "test" + assert str(e) == 'test' assert message.attach.call_count == 0 @@ -210,86 +185,69 @@ def test_try_send_email_success(email): msg = MagicMock() - email.try_send_email(msg, "test") + email.try_send_email(msg, 'test') 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') def test_try_send_email_reconnect_quit_success(smtp, email): - email.server.sendmail = MagicMock( - side_effect=SMTPServerDisconnected("test") - ) + email.server.sendmail = MagicMock(side_effect=SMTPServerDisconnected('test')) email.server.quit = MagicMock() msg = MagicMock() - email.try_send_email(msg, "test") + email.try_send_email(msg, 'test') - smtp.assert_has_calls([ - call("test", 587, timeout=20) - ]) + smtp.assert_has_calls([call('test', 587, timeout=20)]) smtp.return_value.starttls.assert_called_once() - smtp.return_value.login.assert_called_once_with( - "test@test.com", "test") + smtp.return_value.login.assert_called_once_with('test@test.com', 'test') 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') def test_try_send_email_reconnect_quit_failure_disconnect(smtp, email): - email.server.sendmail = MagicMock( - side_effect=SMTPServerDisconnected("test") - ) - email.server.quit = MagicMock( - side_effect=SMTPServerDisconnected("test") - ) + email.server.sendmail = MagicMock(side_effect=SMTPServerDisconnected('test')) + email.server.quit = MagicMock(side_effect=SMTPServerDisconnected('test')) msg = MagicMock() - email.try_send_email(msg, "test") + email.try_send_email(msg, 'test') - smtp.assert_has_calls([ - call("test", 587, timeout=20) - ]) + smtp.assert_has_calls([call('test', 587, timeout=20)]) smtp.return_value.starttls.assert_called_once() - smtp.return_value.login.assert_called_once_with( - "test@test.com", "test") + smtp.return_value.login.assert_called_once_with('test@test.com', 'test') 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') def test_try_send_email_reconnect_quit_failure(smtp, email): - email.server.sendmail = MagicMock( - side_effect=SMTPServerDisconnected("test") - ) - email.server.quit = MagicMock( - side_effect=Exception("test") - ) + email.server.sendmail = MagicMock(side_effect=SMTPServerDisconnected('test')) + email.server.quit = MagicMock(side_effect=Exception('test')) msg = MagicMock() try: - email.try_send_email(msg, "test") + email.try_send_email(msg, 'test') except Exception as e: - assert str(e) == "test" + assert str(e) == 'test' else: - assert False, "Expected exception" + raise AssertionError('Expected exception') @mark.asyncio async def test_send_email_without_smtp_server(email): email.smtp_server = None - input_data = { - **metadata, - "receiver_groups": {}, - "mail_type": "test_TYPE" - } + input_data = {**metadata, 'receiver_groups': {}, 'mail_type': 'test_TYPE'} response = await email.send_email(input_data) assert response == {} @@ -301,41 +259,33 @@ async def test_send_email_without_smtp_server(email): async def test_send_email(mimemultipart, mimetext, email): side_effect_1 = MagicMock() side_effect_2 = MagicMock() - mimemultipart.side_effect = [ - side_effect_1, - side_effect_2 - ] + mimemultipart.side_effect = [side_effect_1, side_effect_2] email.email_builder.send_notification = MagicMock() - email.try_send_email = MagicMock( - side_effect=[ - None, - Exception("test") - ] - ) + email.try_send_email = MagicMock(side_effect=[None, Exception('test')]) input_data = { **metadata, - "receiver_groups": { - "group_1": { - "members": ["test@test.com", "test2@test.com"], - "notifications": [ + 'receiver_groups': { + 'group_1': { + 'members': ['test@test.com', 'test2@test.com'], + 'notifications': [ { - "attachment_content": "test_content_1", - "trigger": "test_trigger", - "notification_id": "test_notification_id" + 'attachment_content': 'test_content_1', + 'trigger': 'test_trigger', + '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) @@ -345,18 +295,13 @@ async def test_send_email(mimemultipart, mimetext, email): assert mimemultipart.call_count == 2 - mimetext.assert_has_calls( - [ - call("test_html1", "html"), - call("test_html2", "html") - ] - ) + mimetext.assert_has_calls([call('test_html1', 'html'), call('test_html2', 'html')]) side_effect_1.__setitem__.assert_has_calls( [ call('From', 'test@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('To', 'test3@test.com, test4@test.com'), - call('Subject', 'SIENTIA™ test_TYPE') + call('Subject', 'SIENTIA™ test_TYPE'), ] ) email.try_send_email.assert_has_calls( [ 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'), ] ) diff --git a/tests/orchestrator/activities/test_formatters.py b/tests/orchestrator/activities/test_formatters.py index aec0e12..f111123 100644 --- a/tests/orchestrator/activities/test_formatters.py +++ b/tests/orchestrator/activities/test_formatters.py @@ -1,8 +1,10 @@ -from unittest.mock import MagicMock, patch, call, ANY import json +from unittest.mock import ANY, MagicMock, call, patch + from pandas import DataFrame from pytest import fixture, mark from sientia_do.notifications.models import NotificationLevel + from orchestrator.activities.formatters import Formatters from orchestrator.utils.orchestrator_functions import build_tag_config @@ -10,10 +12,10 @@ from orchestrator.utils.orchestrator_functions import build_tag_config @fixture def formatters(): formatters = Formatters( - scouter_namespace="scouter", - laborious_namespace="laborious", + scouter_namespace='scouter', + laborious_namespace='laborious', logger=MagicMock(), - notification_handler=MagicMock() + notification_handler=MagicMock(), ) formatters.send_notification = MagicMock() @@ -21,762 +23,717 @@ def formatters(): metadata = { - "metadata": { - "schedule_name": "test_schedule_name", - "workflow_name": "test_workflow_name", - "model_name": "test_model_name", - "model_id": "test_model_id" + 'metadata': { + 'schedule_name': 'test_schedule_name', + 'workflow_name': 'test_workflow_name', + 'model_name': 'test_model_name', + 'model_id': 'test_model_id', } } @mark.asyncio -@patch("orchestrator.activities.formatters.scouter", - return_value={"test_scouter": "test_scouter"}) -@patch("orchestrator.activities.formatters.predictions_batch", - return_value={"test_predictions_batch": "test_predictions_batch"}) -@patch("orchestrator.activities.formatters.minimal_retrain", - return_value={"test_minimal_retrain": "test_minimal_retrain"}) -async def test_process_schedules(mock_minimal_retrain, mock_predictions_batch, mock_scouter, formatters): +@patch('orchestrator.activities.formatters.scouter', return_value={'test_scouter': 'test_scouter'}) +@patch( + 'orchestrator.activities.formatters.predictions_batch', + return_value={'test_predictions_batch': 'test_predictions_batch'}, +) +@patch( + 'orchestrator.activities.formatters.minimal_retrain', + return_value={'test_minimal_retrain': 'test_minimal_retrain'}, +) +async def test_process_schedules( + mock_minimal_retrain, mock_predictions_batch, mock_scouter, formatters +): input_data = { - "pipelines": [ + 'pipelines': [ { - "schedule_name": "test_schedule_name", - "workflow_type": "scouter", - "model_name": "test_model_name", - "model_id": "test_model_id", - "updated_at": "2021-01-01" + 'schedule_name': 'test_schedule_name', + 'workflow_type': 'scouter', + 'model_name': 'test_model_name', + 'model_id': 'test_model_id', + 'updated_at': '2021-01-01', }, { - "schedule_name": "test_schedule_name2", - "workflow_type": "predictions_batch", - "model_name": "test_model_name", - "model_id": "test_model_id", - "updated_at": "2021-01-02" + 'schedule_name': 'test_schedule_name2', + 'workflow_type': 'predictions_batch', + 'model_name': 'test_model_name', + 'model_id': 'test_model_id', + 'updated_at': '2021-01-02', }, { - "schedule_name": "test_schedule_name3", - "workflow_type": "minimal_retrain", - "model_name": "test_model_name", - "model_id": "test_model_id", - "updated_at": "2021-01-03" - } + 'schedule_name': 'test_schedule_name3', + 'workflow_type': 'minimal_retrain', + 'model_name': 'test_model_name', + 'model_id': 'test_model_id', + 'updated_at': '2021-01-03', + }, ] } result = await formatters.process_schedules(input_data) assert result == { - "scouter": { - "test_schedule_name": { - "test_scouter": "test_scouter", - "updated_at": "2021-01-01" - } + 'scouter': { + 'test_schedule_name': {'test_scouter': 'test_scouter', 'updated_at': '2021-01-01'} }, - "laborious": { - "test_schedule_name2": { - "test_predictions_batch": "test_predictions_batch", - "updated_at": "2021-01-02" + 'laborious': { + 'test_schedule_name2': { + 'test_predictions_batch': 'test_predictions_batch', + 'updated_at': '2021-01-02', }, - "test_schedule_name3": { - "test_minimal_retrain": "test_minimal_retrain", - "updated_at": "2021-01-03" - } - } + 'test_schedule_name3': { + 'test_minimal_retrain': 'test_minimal_retrain', + 'updated_at': '2021-01-03', + }, + }, } - mock_scouter.assert_called_once_with( - input_data['pipelines'][0]) - mock_predictions_batch.assert_called_once_with( - input_data['pipelines'][1]) + mock_scouter.assert_called_once_with(input_data['pipelines'][0]) + mock_predictions_batch.assert_called_once_with(input_data['pipelines'][1]) @mark.asyncio -@patch("orchestrator.activities.formatters.gather_read_tags", - return_value={ - "1:test_tag_address": { - "server_id": "1", - "server_name": "test_server_name", - "tag_address": "test_tag_address", - "topics": ["raw_test_schedule"] - }, - "2:test_tag_address2": { - "server_id": "2", - "server_name": "test_server_name2", - "tag_address": "test_tag_address2", - "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"] - } - }) -@patch("orchestrator.activities.formatters.build_tag_config", side_effect=build_tag_config) +@patch( + 'orchestrator.activities.formatters.gather_read_tags', + return_value={ + '1:test_tag_address': { + 'server_id': '1', + 'server_name': 'test_server_name', + 'tag_address': 'test_tag_address', + 'topics': ['raw_test_schedule'], + }, + '2:test_tag_address2': { + 'server_id': '2', + 'server_name': 'test_server_name2', + 'tag_address': 'test_tag_address2', + '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'], + }, + }, +) +@patch('orchestrator.activities.formatters.build_tag_config', side_effect=build_tag_config) async def test_process_slots(mock_build_tag_config, mock_gather_read_tags, formatters): input_data = { - "opc_servers": [ - { - "id": "1", - "server_name": "test_server_name", - "url": "test_url", - "uri": "test_uri" - }, - { - "id": "2", - "server_name": "test_server_name2", - "url": "test_url2", - "uri": "test_uri2" - } + 'opc_servers': [ + {'id': '1', 'server_name': 'test_server_name', 'url': 'test_url', 'uri': 'test_uri'}, + {'id': '2', 'server_name': 'test_server_name2', 'url': 'test_url2', 'uri': 'test_uri2'}, ], - "active_ingestors": [ - "test_active_ingestor1", - "test_active_ingestor2" - ], - "pipelines": "test_gather_read_tags" + 'active_ingestors': ['test_active_ingestor1', 'test_active_ingestor2'], + 'pipelines': 'test_gather_read_tags', } result = await formatters.process_slots(input_data) mock_gather_read_tags.assert_called_once_with(input_data['pipelines']) - mock_build_tag_config.assert_has_calls([ - call( - { - "server_id": "1", - "server_name": "test_server_name", - "tag_address": "test_tag_address", - "topics": ["raw_test_schedule"] - }, - ANY, - { - "1": { - "id": "1", - "server_name": "test_server_name", - "url": "test_url", - "uri": "test_uri" + mock_build_tag_config.assert_has_calls( + [ + call( + { + 'server_id': '1', + 'server_name': 'test_server_name', + 'tag_address': 'test_tag_address', + 'topics': ['raw_test_schedule'], }, - "2": { - "id": "2", - "server_name": "test_server_name2", - "url": "test_url2", - "uri": "test_uri2" - } - }, - 1 - ) - ]) - mock_build_tag_config.assert_has_calls([ - call( - { - "server_id": "2", - "server_name": "test_server_name2", - "tag_address": "test_tag_address2", - "topics": ["raw_test_schedule2"] - }, - ANY, - { - "1": { - "id": "1", - "server_name": "test_server_name", - "url": "test_url", - "uri": "test_uri" + ANY, + { + '1': { + 'id': '1', + 'server_name': 'test_server_name', + 'url': 'test_url', + 'uri': 'test_uri', + }, + '2': { + 'id': '2', + 'server_name': 'test_server_name2', + 'url': 'test_url2', + 'uri': 'test_uri2', + }, }, - "2": { - "id": "2", - "server_name": "test_server_name2", - "url": "test_url2", - "uri": "test_uri2" - } - }, - 1 - ) - ]) - mock_build_tag_config.assert_has_calls([ - call( - { - "server_id": "2", - "server_name": "test_server_name2", - "tag_address": "test_tag_address3", - "topics": ["raw_test_schedule2"] - }, - ANY, - { - "1": { - "id": "1", - "server_name": "test_server_name", - "url": "test_url", - "uri": "test_uri" + 1, + ) + ] + ) + mock_build_tag_config.assert_has_calls( + [ + call( + { + 'server_id': '2', + 'server_name': 'test_server_name2', + 'tag_address': 'test_tag_address2', + 'topics': ['raw_test_schedule2'], }, - "2": { - "id": "2", - "server_name": "test_server_name2", - "url": "test_url2", - "uri": "test_uri2" - } - }, - 2 - ) - ]) + ANY, + { + '1': { + 'id': '1', + 'server_name': 'test_server_name', + 'url': 'test_url', + 'uri': 'test_uri', + }, + '2': { + 'id': '2', + 'server_name': 'test_server_name2', + 'url': 'test_url2', + 'uri': 'test_uri2', + }, + }, + 1, + ) + ] + ) + mock_build_tag_config.assert_has_calls( + [ + call( + { + 'server_id': '2', + 'server_name': 'test_server_name2', + 'tag_address': 'test_tag_address3', + 'topics': ['raw_test_schedule2'], + }, + ANY, + { + '1': { + 'id': '1', + 'server_name': 'test_server_name', + 'url': 'test_url', + 'uri': 'test_uri', + }, + '2': { + 'id': '2', + 'server_name': 'test_server_name2', + 'url': 'test_url2', + 'uri': 'test_uri2', + }, + }, + 2, + ) + ] + ) assert result == { - "1": { - "test_server_name": { - "server_id": "1", - "name": "test_server_name", - "url": "test_url", - "server_uri": "test_uri", - "cert_path": None, - "private_key_path": None, - "server_cert_path": None, - "tags": { - "test_tag_address": { - "server_id": "1", - "server_name": "test_server_name", - "tag_address": "test_tag_address", - "topics": ["raw_test_schedule"] + '1': { + 'test_server_name': { + 'server_id': '1', + 'name': 'test_server_name', + 'url': 'test_url', + 'server_uri': 'test_uri', + 'cert_path': None, + 'private_key_path': None, + 'server_cert_path': None, + 'tags': { + 'test_tag_address': { + 'server_id': '1', + 'server_name': 'test_server_name', + 'tag_address': 'test_tag_address', + 'topics': ['raw_test_schedule'], } - } + }, }, - "test_server_name2": { - "server_id": "2", - "name": "test_server_name2", - "url": "test_url2", - "server_uri": "test_uri2", - "cert_path": None, - "private_key_path": None, - "server_cert_path": None, - "tags": { - "test_tag_address2": { - "server_id": "2", - "server_name": "test_server_name2", - "tag_address": "test_tag_address2", - "topics": ["raw_test_schedule2"] + 'test_server_name2': { + 'server_id': '2', + 'name': 'test_server_name2', + 'url': 'test_url2', + 'server_uri': 'test_uri2', + 'cert_path': None, + 'private_key_path': None, + 'server_cert_path': None, + 'tags': { + 'test_tag_address2': { + 'server_id': '2', + 'server_name': 'test_server_name2', + 'tag_address': 'test_tag_address2', + 'topics': ['raw_test_schedule2'], } - } + }, + }, + }, + '2': { + 'test_server_name2': { + 'server_id': '2', + 'name': 'test_server_name2', + 'url': 'test_url2', + 'server_uri': 'test_uri2', + 'cert_path': None, + 'private_key_path': None, + 'server_cert_path': None, + 'tags': { + 'test_tag_address3': { + 'server_id': '2', + 'server_name': 'test_server_name2', + 'tag_address': 'test_tag_address3', + 'topics': ['raw_test_schedule2'], + } + }, } }, - "2": { - "test_server_name2": { - "server_id": "2", - "name": "test_server_name2", - "url": "test_url2", - "server_uri": "test_uri2", - "cert_path": None, - "private_key_path": None, - "server_cert_path": None, - "tags": { - "test_tag_address3": { - "server_id": "2", - "server_name": "test_server_name2", - "tag_address": "test_tag_address3", - "topics": ["raw_test_schedule2"] - } - } - } - } } @mark.asyncio -@patch("orchestrator.activities.formatters.gather_read_tags", - return_value={ - "1:test_tag_address": { - "server_id": "1", - "server_name": "test_server_name", - "tag_address": "test_tag_address", - "topics": ["raw_test_schedule"] - }, - "2:test_tag_address2": { - "server_id": "2", - "server_name": "test_server_name2", - "tag_address": "test_tag_address2", - "topics": ["raw_test_schedule2"] - } - }) -@patch("orchestrator.activities.formatters.build_tag_config", side_effect=ValueError("test_error")) +@patch( + 'orchestrator.activities.formatters.gather_read_tags', + return_value={ + '1:test_tag_address': { + 'server_id': '1', + 'server_name': 'test_server_name', + 'tag_address': 'test_tag_address', + 'topics': ['raw_test_schedule'], + }, + '2:test_tag_address2': { + 'server_id': '2', + 'server_name': 'test_server_name2', + 'tag_address': 'test_tag_address2', + 'topics': ['raw_test_schedule2'], + }, + }, +) +@patch('orchestrator.activities.formatters.build_tag_config', side_effect=ValueError('test_error')) async def test_process_slots_exception(mock_build_tag_config, mock_gather_read_tags, formatters): input_data = { **metadata, - "opc_servers": [ + 'opc_servers': [ { - "id": "1", - "server_name": "test_server_name", - "url": "test_url", - "uri": "test_uri", - "cert_path": 'test_cert_path', - "private_key_path": 'test_private_key_path', - "server_cert_path": 'test_server_cert_path' + 'id': '1', + 'server_name': 'test_server_name', + 'url': 'test_url', + 'uri': 'test_uri', + 'cert_path': 'test_cert_path', + 'private_key_path': 'test_private_key_path', + 'server_cert_path': 'test_server_cert_path', }, { - "id": "2", - "server_name": "test_server_name2", - "url": "test_url2", - "uri": "test_uri2", - "cert_path": 'test_cert_path2', - "private_key_path": 'test_private_key_path2', - "server_cert_path": 'test_server_cert_path2' - } + 'id': '2', + 'server_name': 'test_server_name2', + 'url': 'test_url2', + 'uri': 'test_uri2', + 'cert_path': 'test_cert_path2', + 'private_key_path': 'test_private_key_path2', + 'server_cert_path': 'test_server_cert_path2', + }, ], - "active_ingestors": [ - "test_active_ingestor1", - "test_active_ingestor2" - ], - "pipelines": "test_gather_read_tags" + 'active_ingestors': ['test_active_ingestor1', 'test_active_ingestor2'], + 'pipelines': 'test_gather_read_tags', } await formatters.process_slots(input_data) - formatters.send_notification.assert_has_calls([ - call( - metadata=metadata['metadata'], - notification_id="ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR", - message="test_error", - block="orchestrator", - level=NotificationLevel.ERROR - ), - call( - metadata=metadata['metadata'], - notification_id="ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR", - message="test_error", - block="orchestrator", - level=NotificationLevel.ERROR - ) - ]) + formatters.send_notification.assert_has_calls( + [ + call( + metadata=metadata['metadata'], + notification_id='ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR', + message='test_error', + block='orchestrator', + level=NotificationLevel.ERROR, + ), + call( + metadata=metadata['metadata'], + notification_id='ORCHESTRATOR_BUILD_TAG_CONFIG_ERROR', + message='test_error', + block='orchestrator', + level=NotificationLevel.ERROR, + ), + ] + ) @mark.asyncio async def test_format_schedule_config(formatters): input_data = { - "schedule_config": [ - {"namespace": "test_namespace1", "schedule_name": "test1", - "updated_at": "2021-01-01"}, - {"namespace": "test_namespace2", "schedule_name": "test2", - "updated_at": "2021-01-02"} + 'schedule_config': [ + {'namespace': 'test_namespace1', 'schedule_name': 'test1', 'updated_at': '2021-01-01'}, + {'namespace': 'test_namespace2', 'schedule_name': 'test2', 'updated_at': '2021-01-02'}, ], - **metadata + **metadata, } result = await formatters.format_schedule_config(input_data) assert result == { - "test_namespace1": { - "test1": "2021-01-01" - }, - "test_namespace2": { - "test2": "2021-01-02" - } + 'test_namespace1': {'test1': '2021-01-01'}, + 'test_namespace2': {'test2': '2021-01-02'}, } @mark.asyncio async def test_create_schedule_config(formatters): input_data = { - "current_schedule_config": { - "scouter": { - "test_schedule_name_to_delete": '2021-01-01', - "test_schedule_name_to_update": '2021-01-02' + 'current_schedule_config': { + 'scouter': { + 'test_schedule_name_to_delete': '2021-01-01', + 'test_schedule_name_to_update': '2021-01-02', } }, - "schedule_config": { - "laborious": { - "test_schedule_name_to_create": { - "frequency": 60, - "data": {"test": "test"}, - "updated_at": "2021-01-03" + 'schedule_config': { + 'laborious': { + 'test_schedule_name_to_create': { + 'frequency': 60, + 'data': {'test': 'test'}, + 'updated_at': '2021-01-03', } }, - "scouter": { - "test_schedule_name_to_update": { - "frequency": 60, - "data": {"test": "test2"}, - "updated_at": "2021-01-04" + 'scouter': { + 'test_schedule_name_to_update': { + 'frequency': 60, + 'data': {'test': 'test2'}, + 'updated_at': '2021-01-04', } - } - } + }, + }, } result = await formatters.create_schedule_config(input_data) assert result == { - "to_create": { - "laborious": { - "test_schedule_name_to_create": { - "frequency": 60, - "data": {"test": "test"}, - "updated_at": "2021-01-03" + 'to_create': { + 'laborious': { + 'test_schedule_name_to_create': { + 'frequency': 60, + 'data': {'test': 'test'}, + 'updated_at': '2021-01-03', } }, - "scouter": {} + 'scouter': {}, }, - "to_update": { - "scouter": { - "test_schedule_name_to_update": { - "frequency": 60, - "data": {"test": "test2"}, - "updated_at": "2021-01-04" + 'to_update': { + 'scouter': { + 'test_schedule_name_to_update': { + 'frequency': 60, + 'data': {'test': 'test2'}, + 'updated_at': '2021-01-04', } }, - "laborious": {} + 'laborious': {}, }, - "to_delete": { - "scouter": [ - "test_schedule_name_to_delete" - ], - "laborious": [] - } + 'to_delete': {'scouter': ['test_schedule_name_to_delete'], 'laborious': []}, } @mark.asyncio async def test_create_slot_config(formatters): input_data = { - "current_slot_config": { - "1": { - "frequency": 60, - "data": {"test": "test"} - }, - "2": { - "frequency": 60, - "data": {"test": "test"} - } + 'current_slot_config': { + '1': {'frequency': 60, 'data': {'test': 'test'}}, + '2': {'frequency': 60, 'data': {'test': 'test'}}, }, - "slot_config": { - "1": { - "frequency": 60, - "data": {"test": "test2"} - } - } + 'slot_config': {'1': {'frequency': 60, 'data': {'test': 'test2'}}}, } result = await formatters.create_slot_config(input_data) assert result == { - "to_delete": [ - "2" - ], - "to_insert": { - "1": { - "frequency": 60, - "data": {"test": "test2"} - } - } + 'to_delete': ['2'], + 'to_insert': {'1': {'frequency': 60, 'data': {'test': 'test2'}}}, } def test_send_success_report(formatters): formatters.send_success_report( metadata=metadata, - message="test_message", - notification_id="test_notification_id", - attachment={"test": "test"} + message='test_message', + notification_id='test_notification_id', + attachment={'test': 'test'}, ) formatters.send_notification.assert_called_once_with( metadata=metadata, - notification_id="test_notification_id", - message="test_message", - block="report_orchestration", + notification_id='test_notification_id', + message='test_message', + block='report_orchestration', level=NotificationLevel.INFO, - attachment_content=json.dumps( - {"test": "test"}, indent=4, sort_keys=True) + attachment_content=json.dumps({'test': 'test'}, indent=4, sort_keys=True), ) def test_send_error_report(formatters): formatters.send_error_report( metadata=metadata, - message="test_message", - notification_id="test_notification_id", - attachment="test_attachment" + message='test_message', + notification_id='test_notification_id', + attachment='test_attachment', ) formatters.send_notification.assert_called_once_with( metadata=metadata, - notification_id="test_notification_id", - message="test_message", - block="report_orchestration", + notification_id='test_notification_id', + message='test_message', + block='report_orchestration', level=NotificationLevel.ERROR, - attachment_content="test_attachment" + attachment_content='test_attachment', ) def test_parse_report(formatters): input_data = { - "test_key": { - "success": True - }, - "test_key2": { - "success": False, - "message": "test_error" - } + 'test_key': {'success': True}, + 'test_key2': {'success': False, 'message': 'test_error'}, } result = formatters.parse_report(input_data) - assert result == ( - ["test_key"], - ["test_key2"] - ) + assert result == (['test_key'], ['test_key2']) def test_parse_report_schedule(formatters): input_data = [ { - "namespace": "test_namespace", - "schedule_name": "test_schedule_name_to_create", - "success": True + 'namespace': 'test_namespace', + 'schedule_name': 'test_schedule_name_to_create', + 'success': True, }, { - "namespace": "test_namespace", - "schedule_name": "test_schedule_name_to_create_error", - "success": False, - "message": "test_error", - "attachment": "test_attachment" - } + 'namespace': 'test_namespace', + 'schedule_name': 'test_schedule_name_to_create_error', + 'success': False, + 'message': 'test_error', + 'attachment': 'test_attachment', + }, ] result = formatters.parse_report_schedule(input_data) assert result == ( - ["test_namespace/test_schedule_name_to_create"], - {"test_namespace/test_schedule_name_to_create_error": { - 'message': 'test_error', - 'attachment': 'test_attachment' - }} + ['test_namespace/test_schedule_name_to_create'], + { + 'test_namespace/test_schedule_name_to_create_error': { + 'message': 'test_error', + 'attachment': 'test_attachment', + } + }, ) @mark.asyncio async def test_report_schedule_orchestration(formatters): - formatters.parse_report_schedule = MagicMock( - side_effect=formatters.parse_report_schedule - ) + formatters.parse_report_schedule = MagicMock(side_effect=formatters.parse_report_schedule) formatters.send_success_report = MagicMock() formatters.send_error_report = MagicMock() input_data = { **metadata, - "created_schedules": [ + 'created_schedules': [ { - "namespace": "test_namespace", - "schedule_name": "test_schedule_name_to_create", - "success": True + 'namespace': 'test_namespace', + 'schedule_name': 'test_schedule_name_to_create', + 'success': True, }, { - "namespace": "test_namespace", - "schedule_name": "test_schedule_name_to_create_error", - "success": False, - "message": "test_error1", - "attachment": "test_attachment1" + 'namespace': 'test_namespace', + 'schedule_name': 'test_schedule_name_to_create_error', + 'success': False, + 'message': 'test_error1', + 'attachment': 'test_attachment1', }, { - "namespace": "test_namespace", - "schedule_name": "test_schedule_name_to_create_error2", - "success": False, - "message": "test_error2" - } + 'namespace': 'test_namespace', + 'schedule_name': 'test_schedule_name_to_create_error2', + 'success': False, + 'message': 'test_error2', + }, ], - "updated_schedules": [ + 'updated_schedules': [ { - "namespace": "test_namespace", - "schedule_name": "test_schedule_name_to_update", - "success": True + 'namespace': 'test_namespace', + 'schedule_name': 'test_schedule_name_to_update', + 'success': True, }, { - "namespace": "test_namespace", - "schedule_name": "test_schedule_name_to_update_error", - "success": False, - "message": "test_error2", - "attachment": "test_attachment2" + 'namespace': 'test_namespace', + 'schedule_name': 'test_schedule_name_to_update_error', + 'success': False, + 'message': 'test_error2', + 'attachment': 'test_attachment2', }, { - "namespace": "test_namespace", - "schedule_name": "test_schedule_name_to_update_error2", - "success": False, - "message": "test_error3", - "attachment": "test_attachment3" + 'namespace': 'test_namespace', + 'schedule_name': 'test_schedule_name_to_update_error2', + 'success': False, + 'message': 'test_error3', + 'attachment': 'test_attachment3', }, { - "namespace": "test_namespace", - "schedule_name": "test_schedule_name_to_update_error3", - "success": False, - "message": "test_error4" - } + 'namespace': 'test_namespace', + 'schedule_name': 'test_schedule_name_to_update_error3', + 'success': False, + 'message': 'test_error4', + }, ], - "deleted_schedules": [ + 'deleted_schedules': [ { - "namespace": "test_namespace", - "schedule_name": "test_schedule_name_to_delete", - "success": True + 'namespace': 'test_namespace', + 'schedule_name': 'test_schedule_name_to_delete', + 'success': True, }, { - "namespace": "test_namespace", - "schedule_name": "test_schedule_name_to_delete_error", - "success": False, - "message": "test_error4", - "attachment": "test_attachment4" + 'namespace': 'test_namespace', + 'schedule_name': 'test_schedule_name_to_delete_error', + 'success': False, + 'message': 'test_error4', + 'attachment': 'test_attachment4', }, { - "namespace": "test_namespace", - "schedule_name": "test_schedule_name_to_delete_error2", - "success": False, - "message": "test_error5" - } - ] + 'namespace': 'test_namespace', + 'schedule_name': 'test_schedule_name_to_delete_error2', + 'success': False, + 'message': 'test_error5', + }, + ], } await formatters.report_schedule_orchestration(input_data) - formatters.parse_report_schedule.assert_has_calls([ - call(input_data['created_schedules']), - call(input_data['updated_schedules']), - call(input_data['deleted_schedules']) - ]) - formatters.send_success_report.assert_has_calls([ - call( - metadata=metadata['metadata'], - message="Successfully created schedules: \n test_namespace/test_schedule_name_to_create", - notification_id="REPORT_ORCHESTRATION_CREATED_SCHEDULES", - attachment=input_data['created_schedules'] - ), - call( - metadata=metadata['metadata'], - message="Successfully updated schedules: \n test_namespace/test_schedule_name_to_update", - notification_id="REPORT_ORCHESTRATION_UPDATED_SCHEDULES", - attachment=input_data['updated_schedules'] - ), - call( - metadata=metadata['metadata'], - message="Successfully deleted schedules: \n test_namespace/test_schedule_name_to_delete", - notification_id="REPORT_ORCHESTRATION_DELETED_SCHEDULES", - attachment=input_data['deleted_schedules'] - ) - ]) - formatters.send_error_report.assert_has_calls([ - call( - metadata=metadata['metadata'], - message="Fails on created schedules: \n test_namespace/test_schedule_name_to_create_error, test_namespace/test_schedule_name_to_create_error2", - notification_id="REPORT_ORCHESTRATION_CREATED_SCHEDULES_ERROR", - attachment="test_namespace/test_schedule_name_to_create_error:\ntest_error1\ntest_attachment1\n ========== \ntest_namespace/test_schedule_name_to_create_error2:\ntest_error2" - ), - call( - metadata=metadata['metadata'], - message="Fails on updated schedules: \n test_namespace/test_schedule_name_to_update_error, test_namespace/test_schedule_name_to_update_error2, test_namespace/test_schedule_name_to_update_error3", - notification_id="REPORT_ORCHESTRATION_UPDATED_SCHEDULES_ERROR", - attachment="test_namespace/test_schedule_name_to_update_error:\ntest_error2\ntest_attachment2\n ========== \ntest_namespace/test_schedule_name_to_update_error2:\ntest_error3\ntest_attachment3\n ========== \ntest_namespace/test_schedule_name_to_update_error3:\ntest_error4" - ), - call( - metadata=metadata['metadata'], - message="Fails on deleted schedules: \n test_namespace/test_schedule_name_to_delete_error, test_namespace/test_schedule_name_to_delete_error2", - notification_id="REPORT_ORCHESTRATION_DELETED_SCHEDULES_ERROR", - attachment="test_namespace/test_schedule_name_to_delete_error:\ntest_error4\ntest_attachment4\n ========== \ntest_namespace/test_schedule_name_to_delete_error2:\ntest_error5" - ) - ]) + formatters.parse_report_schedule.assert_has_calls( + [ + call(input_data['created_schedules']), + call(input_data['updated_schedules']), + call(input_data['deleted_schedules']), + ] + ) + formatters.send_success_report.assert_has_calls( + [ + call( + metadata=metadata['metadata'], + message='Successfully created schedules: \n test_namespace/test_schedule_name_to_create', + notification_id='REPORT_ORCHESTRATION_CREATED_SCHEDULES', + attachment=input_data['created_schedules'], + ), + call( + metadata=metadata['metadata'], + message='Successfully updated schedules: \n test_namespace/test_schedule_name_to_update', + notification_id='REPORT_ORCHESTRATION_UPDATED_SCHEDULES', + attachment=input_data['updated_schedules'], + ), + call( + metadata=metadata['metadata'], + message='Successfully deleted schedules: \n test_namespace/test_schedule_name_to_delete', + notification_id='REPORT_ORCHESTRATION_DELETED_SCHEDULES', + attachment=input_data['deleted_schedules'], + ), + ] + ) + formatters.send_error_report.assert_has_calls( + [ + call( + metadata=metadata['metadata'], + message='Fails on created schedules: \n test_namespace/test_schedule_name_to_create_error, test_namespace/test_schedule_name_to_create_error2', + notification_id='REPORT_ORCHESTRATION_CREATED_SCHEDULES_ERROR', + attachment='test_namespace/test_schedule_name_to_create_error:\ntest_error1\ntest_attachment1\n ========== \ntest_namespace/test_schedule_name_to_create_error2:\ntest_error2', + ), + call( + metadata=metadata['metadata'], + message='Fails on updated schedules: \n test_namespace/test_schedule_name_to_update_error, test_namespace/test_schedule_name_to_update_error2, test_namespace/test_schedule_name_to_update_error3', + notification_id='REPORT_ORCHESTRATION_UPDATED_SCHEDULES_ERROR', + attachment='test_namespace/test_schedule_name_to_update_error:\ntest_error2\ntest_attachment2\n ========== \ntest_namespace/test_schedule_name_to_update_error2:\ntest_error3\ntest_attachment3\n ========== \ntest_namespace/test_schedule_name_to_update_error3:\ntest_error4', + ), + call( + metadata=metadata['metadata'], + message='Fails on deleted schedules: \n test_namespace/test_schedule_name_to_delete_error, test_namespace/test_schedule_name_to_delete_error2', + notification_id='REPORT_ORCHESTRATION_DELETED_SCHEDULES_ERROR', + attachment='test_namespace/test_schedule_name_to_delete_error:\ntest_error4\ntest_attachment4\n ========== \ntest_namespace/test_schedule_name_to_delete_error2:\ntest_error5', + ), + ] + ) @mark.asyncio async def test_report_slot_orchestration(formatters): - formatters.parse_report = MagicMock( - side_effect=formatters.parse_report - ) + formatters.parse_report = MagicMock(side_effect=formatters.parse_report) formatters.send_success_report = MagicMock() formatters.send_error_report = MagicMock() input_data = { **metadata, - "inserted_slots": { - "test_slot_name_to_create": { - "success": True - }, - "test_slot_name_to_create_error": { - "success": False, - "error": "test_error" - } + 'inserted_slots': { + 'test_slot_name_to_create': {'success': True}, + 'test_slot_name_to_create_error': {'success': False, 'error': 'test_error'}, + }, + 'deleted_slots': { + 'test_slot_name_to_delete': {'success': True}, + 'test_slot_name_to_delete_error': {'success': False, 'error': 'test_error'}, }, - "deleted_slots": { - "test_slot_name_to_delete": { - "success": True - }, - "test_slot_name_to_delete_error": { - "success": False, - "error": "test_error" - } - } } await formatters.report_slot_orchestration(input_data) - formatters.parse_report.assert_has_calls([ - call(input_data['inserted_slots']), - call(input_data['deleted_slots']) - ]) - formatters.send_success_report.assert_has_calls([ - call( - metadata=metadata['metadata'], - message="Inserted slots: \n test_slot_name_to_create", - notification_id="REPORT_ORCHESTRATION_INSERTED_SLOTS" - ), - call( - metadata=metadata['metadata'], - message="Deleted slots: \n test_slot_name_to_delete", - notification_id="REPORT_ORCHESTRATION_DELETED_SLOTS" - ) - ]) - formatters.send_error_report.assert_has_calls([ - call( - metadata=metadata['metadata'], - message="Failed to insert slots: \n test_slot_name_to_create_error", - notification_id="REPORT_ORCHESTRATION_INSERTED_SLOTS", - attachment=input_data['inserted_slots'] - ), - call( - metadata=metadata['metadata'], - message="Failed to delete slots: \n test_slot_name_to_delete_error", - notification_id="REPORT_ORCHESTRATION_DELETED_SLOTS", - attachment=input_data['deleted_slots'] - ) - ]) + formatters.parse_report.assert_has_calls( + [call(input_data['inserted_slots']), call(input_data['deleted_slots'])] + ) + formatters.send_success_report.assert_has_calls( + [ + call( + metadata=metadata['metadata'], + message='Inserted slots: \n test_slot_name_to_create', + notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS', + ), + call( + metadata=metadata['metadata'], + message='Deleted slots: \n test_slot_name_to_delete', + notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS', + ), + ] + ) + formatters.send_error_report.assert_has_calls( + [ + call( + metadata=metadata['metadata'], + message='Failed to insert slots: \n test_slot_name_to_create_error', + notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS', + attachment=input_data['inserted_slots'], + ), + call( + metadata=metadata['metadata'], + message='Failed to delete slots: \n test_slot_name_to_delete_error', + notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS', + attachment=input_data['deleted_slots'], + ), + ] + ) @mark.asyncio async def test_format_log_report(formatters): input_data = { **metadata, - "receiver_groups": { - "test_receiver_group": { - "notifications": [ + 'receiver_groups': { + 'test_receiver_group': { + 'notifications': [ { - "notification_id": "test_notification_id", - "trigger": "test_trigger", - "timestamp": "2021-01-01", - "message": "test_message", - "level": "test_level", - "block": "test_block", - "pipeline": "test_pipeline", - "project": "test_project", - "model_name": "test_model_name", - "model_id": "test_model_id" + 'notification_id': 'test_notification_id', + 'trigger': 'test_trigger', + 'timestamp': '2021-01-01', + 'message': 'test_message', + 'level': 'test_level', + 'block': 'test_block', + 'pipeline': 'test_pipeline', + 'project': 'test_project', + 'model_name': 'test_model_name', + 'model_id': 'test_model_id', } ], - "status": "sent", + 'status': 'sent', }, - "test_receiver_group2": { - "notifications": [ + 'test_receiver_group2': { + 'notifications': [ { - "notification_id": "test_notification_id", - "trigger": "test_trigger", - "timestamp": "2021-01-01", - "message": "test_message", - "level": "test_level", - "block": "test_block", - "pipeline": "test_pipeline", - "project": "test_project", - "model_name": "test_model_name", - "model_id": "test_model_id" + 'notification_id': 'test_notification_id', + 'trigger': 'test_trigger', + 'timestamp': '2021-01-01', + 'message': 'test_message', + 'level': 'test_level', + 'block': 'test_block', + 'pipeline': 'test_pipeline', + 'project': 'test_project', + 'model_name': 'test_model_name', + 'model_id': 'test_model_id', } ], - "status": "sent", - } + 'status': 'sent', + }, }, - "mail_type": "test_mail_type" + 'mail_type': 'test_mail_type', } result = await formatters.format_log_report(input_data) @@ -784,19 +741,19 @@ async def test_format_log_report(formatters): expected_result = DataFrame( [ { - "status": "sent", - "timestamp": "2021-01-01", - "groups": ["test_receiver_group", "test_receiver_group2"], - "message": "test_message", - "level": "test_level", - "notification_id": "test_notification_id", - "block": "test_block", - "schedule": "test_trigger", - "pipeline": "test_pipeline", - "project": "test_project", - "model_name": "test_model_name", - "model_id": "test_model_id", - "mail_type": "test_mail_type" + 'status': 'sent', + 'timestamp': '2021-01-01', + 'groups': ['test_receiver_group', 'test_receiver_group2'], + 'message': 'test_message', + 'level': 'test_level', + 'notification_id': 'test_notification_id', + 'block': 'test_block', + 'schedule': 'test_trigger', + 'pipeline': 'test_pipeline', + 'project': 'test_project', + 'model_name': 'test_model_name', + 'model_id': 'test_model_id', + 'mail_type': 'test_mail_type', } ] ) @@ -806,34 +763,21 @@ async def test_format_log_report(formatters): @mark.asyncio async def test_filter_notification_reports(formatters): - input_data = { **metadata, 'notification_package': [ - { - 'trigger': 'test_trigger_1', - 'notification_id': 'test_notification_id_1' - }, - { - 'trigger': 'test_trigger_2', - 'notification_id': 'test_notification_id_2' - }, - { - 'trigger': 'test_trigger_3', - 'notification_id': 'test_notification_id_3' - } + {'trigger': 'test_trigger_1', 'notification_id': 'test_notification_id_1'}, + {'trigger': 'test_trigger_2', 'notification_id': 'test_notification_id_2'}, + {'trigger': 'test_trigger_3', 'notification_id': 'test_notification_id_3'}, ], 'sending_configs': [ { 'group_name': 'test_group_1', 'contents': ['reports'], - 'ignore': ['test_notification_id_1'] + 'ignore': ['test_notification_id_1'], }, - { - 'group_name': 'test_group_2', - 'contents': ['core_alerts'] - } - ] + {'group_name': 'test_group_2', 'contents': ['core_alerts']}, + ], } response = await formatters.filter_notification_reports(input_data) @@ -844,14 +788,8 @@ async def test_filter_notification_reports(formatters): 'contents': ['reports'], 'ignore': ['test_notification_id_1'], 'notifications': [ - { - 'trigger': 'test_trigger_2', - 'notification_id': 'test_notification_id_2' - }, - { - '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'}, + ], } } diff --git a/tests/orchestrator/activities/test_mongo_db.py b/tests/orchestrator/activities/test_mongo_db.py index bce033d..53e757f 100644 --- a/tests/orchestrator/activities/test_mongo_db.py +++ b/tests/orchestrator/activities/test_mongo_db.py @@ -1,512 +1,461 @@ -from curses import meta from datetime import datetime -from unittest.mock import MagicMock, patch, ANY +from unittest.mock import ANY, MagicMock, patch + 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.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ +from orchestrator.activities.mongo_db import MongoDB, clear_mongo_id + def test_clear_mongo_id(): input_data = [ [ { - "name": "test", - "_id": "12345", + 'name': 'test', + '_id': '12345', } ], { - "name": "test", - "_id": "12345", - "nested": { - "_id": "67890", - "value": [1, 2, 3], - "list": [{"_id": "abcde", "item": "value"}] + 'name': 'test', + '_id': '12345', + 'nested': { + '_id': '67890', + 'value': [1, 2, 3], + 'list': [{'_id': 'abcde', 'item': 'value'}], }, - "nested_list": [ - {"_id": "fghij", "item": "value1"}, - {"_id": "klmno", "item": "value2"} - ] - } + 'nested_list': [{'_id': 'fghij', 'item': 'value1'}, {'_id': 'klmno', 'item': 'value2'}], + }, ] output = clear_mongo_id(input_data) assert output == [ - [ - {"name": "test"} - ], + [{'name': 'test'}], { - "name": "test", - "nested": { - "value": [1, 2, 3], - "list": [{"item": "value"}] - }, - "nested_list": [ - {"item": "value1"}, - {"item": "value2"} - ] - }] + 'name': 'test', + 'nested': {'value': [1, 2, 3], 'list': [{'item': 'value'}]}, + 'nested_list': [{'item': 'value1'}, {'item': 'value2'}], + }, + ] @fixture -@patch("orchestrator.activities.mongo_db.MongoClient") +@patch('orchestrator.activities.mongo_db.MongoClient') def mongo_db(mongo_mock): - mongo = ( - MongoDB( - connection_string="mongodb://localhost:27017", - database_name="test_db", - ttl_index_seconds=3600, - logger=MagicMock(), - notification_handler=MagicMock() - ) + mongo = MongoDB( + connection_string='mongodb://localhost:27017', + database_name='test_db', + ttl_index_seconds=3600, + logger=MagicMock(), + notification_handler=MagicMock(), ) mongo.send_notification = MagicMock() return mongo -@patch("orchestrator.activities.mongo_db.MongoClient") +@patch('orchestrator.activities.mongo_db.MongoClient') def test___init__(mongo_mock): mongo_db = MongoDB( - connection_string="mongodb://localhost:27017", - database_name="test_db", + connection_string='mongodb://localhost:27017', + database_name='test_db', ttl_index_seconds=3600, logger=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 + 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) 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): mongo_db.shutdown() mongo_db.client.close.assert_called_once() - mongo_db.logger.info.assert_any_call("Closing MongoDB connection...") - mongo_db.logger.info.assert_any_call( - "MongoDB connection closed successfully") + mongo_db.logger.info.assert_any_call('Closing MongoDB connection...') + mongo_db.logger.info.assert_any_call('MongoDB connection closed successfully') 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.logger.error.assert_called_once_with( - "Failed to close MongoDB connection: Close failed" + 'Failed to close MongoDB connection: Close failed' ) @mark.asyncio async def test_find_documents_in_mongodb_success(mongo_db): - input_data = {"collection": "test_collection", "filters": { - "name": {"$exists": True} - }} + input_data = {'collection': 'test_collection', 'filters': {'name': {'$exists': True}}} mock_collection = MagicMock() mock_collection.find.return_value = [ { - "_id": "12345", - "name": "test1", - "timestamp": datetime.strptime( - "2023-01-01 12:00:00.000000+0000", DATETIME_FORMAT_MS_WITH_TZ)}, + '_id': '12345', + 'name': 'test1', + 'timestamp': datetime.strptime( + '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ + ), + }, { - "_id": "67890", - "name": "test2", - "timestamp": datetime.strptime( - "2023-01-01 12:00:00.000000+0000", DATETIME_FORMAT_MS_WITH_TZ)} + '_id': '67890', + 'name': 'test2', + 'timestamp': datetime.strptime( + '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ + ), + }, ] mongo_db.database.__getitem__.return_value = mock_collection result = await mongo_db.find_documents_in_mongodb( - { - "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} + {'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}) + + metadata = { - "metadata": { - "schedule_name": "test_schedule_name", - "workflow_name": "test_workflow_name", - "model_name": "test_model_name", - "model_id": "test_model_id" + 'metadata': { + 'schedule_name': 'test_schedule_name', + 'workflow_name': 'test_workflow_name', + 'model_name': 'test_model_name', + 'model_id': 'test_model_id', } } @mark.asyncio async def test_find_documents_in_mongodb_failure(mongo_db): - input_data = {"collection": "test_collection", "filters": { - "name": {"$exists": True} - }} + input_data = {'collection': 'test_collection', 'filters': {'name': {'$exists': True}}} mongo_db.database.__getitem__.return_value = MagicMock( - find=MagicMock(side_effect=Exception("Error")) + find=MagicMock(side_effect=Exception('Error')) ) try: - await mongo_db.find_documents_in_mongodb( - { - "query": input_data, - **metadata - }) + await mongo_db.find_documents_in_mongodb({'query': input_data, **metadata}) except Exception as e: - assert str(e) == "Error" + assert str(e) == 'Error' mongo_db.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="MONGODB_QUERY_ERROR", - message="Failed to execute MongoDB query: Error", + notification_id='MONGODB_QUERY_ERROR', + message='Failed to execute MongoDB query: Error', level=NotificationLevel.ERROR, - block="load_query_from_mongodb", - attachment_content=ANY + block='load_query_from_mongodb', + attachment_content=ANY, ) else: - assert False, "Expected an exception to be raised" + raise AssertionError('Expected an exception to be raised') @mark.asyncio async def test_find_documents_in_mongodb_missing_collection(mongo_db): - input_data = {"query": {"filters": {}}} + input_data = {'query': {'filters': {}}} try: await mongo_db.find_documents_in_mongodb(input_data) 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: - assert False, "Expected a ValueError to be raised" + raise AssertionError('Expected a ValueError to be raised') @mark.asyncio async def test_aggregate_documents_in_mongodb_success(mongo_db): - input_data = {"collection": "test_collection", "aggregation": [ - {"$match": {"name": {"$exists": True}}}, - {"$project": {"name": 1}} - ]} + input_data = { + 'collection': 'test_collection', + 'aggregation': [{'$match': {'name': {'$exists': True}}}, {'$project': {'name': 1}}], + } mock_collection = MagicMock() mock_collection.aggregate.return_value = [ { - "_id": "asdad", - "name": "test1", - "timestamp": datetime.strptime( - "2023-01-01 12:00:00.000000+0000", DATETIME_FORMAT_MS_WITH_TZ)}, + '_id': 'asdad', + 'name': 'test1', + 'timestamp': datetime.strptime( + '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ + ), + }, { - "_id": "adzx", - "name": "test2", - "timestamp": datetime.strptime( - "2023-01-01 12:00:00.000000+0000", DATETIME_FORMAT_MS_WITH_TZ)} + '_id': 'adzx', + 'name': 'test2', + 'timestamp': datetime.strptime( + '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ + ), + }, ] mongo_db.database.__getitem__.return_value = mock_collection 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 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"} - expected_pipeline = input_data["aggregation"] - expected_pipeline.append({"$project": {"_id": 0}}) + 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'} + expected_pipeline = input_data['aggregation'] + expected_pipeline.append({'$project': {'_id': 0}}) - mock_collection.aggregate.assert_called_once_with( - expected_pipeline - ) + mock_collection.aggregate.assert_called_once_with(expected_pipeline) @mark.asyncio async def test_aggregate_documents_in_mongodb_failure(mongo_db): - input_data = {"collection": "test_collection", "aggregation": [ - {"$match": {"name": {"$exists": True}}}, - {"$project": {"name": 1}} - ]} + input_data = { + 'collection': 'test_collection', + 'aggregation': [{'$match': {'name': {'$exists': True}}}, {'$project': {'name': 1}}], + } mongo_db.database.__getitem__.return_value = MagicMock( - aggregate=MagicMock(side_effect=Exception("Error")) + aggregate=MagicMock(side_effect=Exception('Error')) ) try: - await mongo_db.aggregate_documents_in_mongodb( - { - "query": input_data, - **metadata - }) + await mongo_db.aggregate_documents_in_mongodb({'query': input_data, **metadata}) except Exception as e: - assert str(e) == "Error" + assert str(e) == 'Error' mongo_db.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="MONGODB_AGGREGATION_ERROR", - message="Failed to execute MongoDB aggregation: Error", - block="aggregate_documents_in_mongodb", + notification_id='MONGODB_AGGREGATION_ERROR', + message='Failed to execute MongoDB aggregation: Error', + block='aggregate_documents_in_mongodb', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "Expected an exception to be raised" + raise AssertionError('Expected an exception to be raised') @mark.asyncio async def test_aggregate_documents_in_mongodb_missing_collection(mongo_db): - input_data = {"query": {"aggregation": []}} + input_data = {'query': {'aggregation': []}} try: await mongo_db.aggregate_documents_in_mongodb(input_data) 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: - assert False, "Expected a ValueError to be raised" + raise AssertionError('Expected a ValueError to be raised') @mark.asyncio async def test_aggregate_documents_in_mongodb_missing_aggregation(mongo_db): - input_data = {"query": {"collection": "test_collection"}} + input_data = {'query': {'collection': 'test_collection'}} try: await mongo_db.aggregate_documents_in_mongodb(input_data) except ValueError as e: - assert str(e) == "Aggregation must be provided." + assert str(e) == 'Aggregation must be provided.' else: - assert False, "Expected a ValueError to be raised" + raise AssertionError('Expected a ValueError to be raised') @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): - input_data = {"updated_pipelines": [ - {"schedule_name": "test1", "namespace": "test1", "success": True}, - {"schedule_name": "test2", "namespace": "test2", "success": True} - ]} - mongo_db.database["pipelines"].update_many.return_value = MagicMock() + input_data = { + 'updated_pipelines': [ + {'schedule_name': 'test1', 'namespace': 'test1', 'success': True}, + {'schedule_name': 'test2', 'namespace': 'test2', 'success': True}, + ] + } + mongo_db.database['pipelines'].update_many.return_value = MagicMock() await mongo_db.update_pipelines_timestamps(input_data) - mongo_db.database["pipelines"].update_many.assert_called_once_with( - {"$or": [ - {"schedule_name": "test1", "namespace": "test1"}, - {"schedule_name": "test2", "namespace": "test2"} - ]}, - {"$set": { - "updated_at": now_mock.return_value}} + mongo_db.database['pipelines'].update_many.assert_called_once_with( + { + '$or': [ + {'schedule_name': 'test1', 'namespace': 'test1'}, + {'schedule_name': 'test2', 'namespace': 'test2'}, + ] + }, + {'$set': {'updated_at': now_mock.return_value}}, ) @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): input_data = { - "updated_pipelines": [ - {"schedule_name": "test1", "namespace": "test1", "success": True}, - {"schedule_name": "test2", "namespace": "test2", "success": True} + 'updated_pipelines': [ + {'schedule_name': 'test1', 'namespace': 'test1', '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: await mongo_db.update_pipelines_timestamps(input_data) except Exception as e: - assert str(e) == "Error" + assert str(e) == 'Error' mongo_db.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="MONGODB_UPDATE_PIPELINES_ERROR", - message="Failed to update pipelines timestamps: Error", - block="update_pipelines_timestamps", + notification_id='MONGODB_UPDATE_PIPELINES_ERROR', + message='Failed to update pipelines timestamps: Error', + block='update_pipelines_timestamps', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "Expected an exception to be raised" + raise AssertionError('Expected an exception to be raised') @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): - input_data = {"created_pipelines": [ - {"schedule_name": "test1", "namespace": "test1", "success": True}, - {"schedule_name": "test2", "namespace": "test2", "success": True} - ]} - mongo_db.database["pipelines"].insert_many.return_value = MagicMock() + input_data = { + 'created_pipelines': [ + {'schedule_name': 'test1', 'namespace': 'test1', 'success': True}, + {'schedule_name': 'test2', 'namespace': 'test2', 'success': True}, + ] + } + mongo_db.database['pipelines'].insert_many.return_value = MagicMock() 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", - "updated_at": now_mock.return_value}, - {"schedule_name": "test2", "namespace": "test2", - "updated_at": now_mock.return_value} + {'schedule_name': 'test1', 'namespace': 'test1', 'updated_at': now_mock.return_value}, + {'schedule_name': 'test2', 'namespace': 'test2', 'updated_at': now_mock.return_value}, ] ) @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): - input_data = {"created_pipelines": [ - {"schedule_name": "test1", "namespace": "test1", "success": True}, - {"schedule_name": "test2", "namespace": "test2", "success": True} - ], - **metadata + input_data = { + 'created_pipelines': [ + {'schedule_name': 'test1', 'namespace': 'test1', 'success': True}, + {'schedule_name': 'test2', 'namespace': 'test2', 'success': True}, + ], + **metadata, } - mongo_db.database["pipelines"].insert_many.side_effect = Exception("Error") + mongo_db.database['pipelines'].insert_many.side_effect = Exception('Error') try: await mongo_db.create_pipelines_timestamps(input_data) except Exception as e: - assert str(e) == "Error" + assert str(e) == 'Error' mongo_db.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="MONGODB_CREATE_PIPELINES_ERROR", - message="Failed to create pipelines timestamps: Error", - block="create_pipelines_timestamps", + notification_id='MONGODB_CREATE_PIPELINES_ERROR', + message='Failed to create pipelines timestamps: Error', + block='create_pipelines_timestamps', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "Expected an exception to be raised" + raise AssertionError('Expected an exception to be raised') @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): - input_data = {"deleted_pipelines": [ - {"schedule_name": "test1", "namespace": "test1", "success": True}, - {"schedule_name": "test2", "namespace": "test2", "success": True} - ]} - mongo_db.database["pipelines"].delete_many.return_value = MagicMock() + input_data = { + 'deleted_pipelines': [ + {'schedule_name': 'test1', 'namespace': 'test1', 'success': True}, + {'schedule_name': 'test2', 'namespace': 'test2', 'success': True}, + ] + } + mongo_db.database['pipelines'].delete_many.return_value = MagicMock() await mongo_db.delete_pipelines_timestamps(input_data) - mongo_db.database["pipelines"].delete_many.assert_called_once_with( - {"$or": [ - {"schedule_name": "test1", "namespace": "test1"}, - {"schedule_name": "test2", "namespace": "test2"} - ]} + mongo_db.database['pipelines'].delete_many.assert_called_once_with( + { + '$or': [ + {'schedule_name': 'test1', 'namespace': 'test1'}, + {'schedule_name': 'test2', 'namespace': 'test2'}, + ] + } ) @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): - input_data = {"deleted_pipelines": [ - {"schedule_name": "test1", "namespace": "test1", "success": True}, - {"schedule_name": "test2", "namespace": "test2", "success": True} - ], - **metadata + input_data = { + 'deleted_pipelines': [ + {'schedule_name': 'test1', 'namespace': 'test1', 'success': True}, + {'schedule_name': 'test2', 'namespace': 'test2', 'success': True}, + ], + **metadata, } - mongo_db.database["pipelines"].delete_many.side_effect = Exception("Error") + mongo_db.database['pipelines'].delete_many.side_effect = Exception('Error') try: await mongo_db.delete_pipelines_timestamps(input_data) except Exception as e: - assert str(e) == "Error" + assert str(e) == 'Error' mongo_db.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="MONGODB_DELETE_PIPELINES_ERROR", - message="Failed to delete pipelines timestamps: Error", - block="delete_pipelines_timestamps", + notification_id='MONGODB_DELETE_PIPELINES_ERROR', + message='Failed to delete pipelines timestamps: Error', + block='delete_pipelines_timestamps', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "Expected an exception to be raised" + raise AssertionError('Expected an exception to be raised') @mark.asyncio async def test_create_collection_with_ttl_index_success(mongo_db): input_data = { **metadata, - "pipelines": { - "scouter-pipeline": { - "topic": "raw_scouter_pipeline" - }, - "scouter-pipeline-2": { - "topic": "raw_scouter_pipeline_2" - }, - "scouter-pipeline-3": { - "topic": "raw_scouter_pipeline_3" - } - } + 'pipelines': { + 'scouter-pipeline': {'topic': 'raw_scouter_pipeline'}, + 'scouter-pipeline-2': {'topic': 'raw_scouter_pipeline_2'}, + 'scouter-pipeline-3': {'topic': 'raw_scouter_pipeline_3'}, + }, } mongo_db.database.list_collection_names.return_value = [ - "raw_scouter_pipeline_2", - "raw_scouter_pipeline_3" + 'raw_scouter_pipeline_2', + 'raw_scouter_pipeline_3', ] collection_1 = MagicMock( list_indexes=MagicMock( return_value=[ { - "key": "asdad", + 'key': 'asdad', } ] ) ) collection_2 = MagicMock( - list_indexes=MagicMock( - return_value=[ - { - "key": "inserted_at", - "expireAfterSeconds": None - } - ] - ) + list_indexes=MagicMock(return_value=[{'key': 'inserted_at', 'expireAfterSeconds': None}]) ) collection_3 = MagicMock( - list_indexes=MagicMock( - return_value=[ - { - "key": "inserted_at", - "expireAfterSeconds": 3600 - } - ] - ) + list_indexes=MagicMock(return_value=[{'key': 'inserted_at', 'expireAfterSeconds': 3600}]) ) mongo_db.database.__getitem__ = MagicMock( - side_effect=[ - collection_1, - collection_2, - collection_3 - ] + side_effect=[collection_1, collection_2, collection_3] ) await mongo_db.create_collection_with_ttl_index(input_data) mongo_db.database.list_collection_names.assert_called_once_with() - mongo_db.database.create_collection.assert_called_once_with( - "raw_scouter_pipeline" - ) + mongo_db.database.create_collection.assert_called_once_with('raw_scouter_pipeline') collection_1.list_indexes.assert_called_once() collection_1.create_index.assert_called_once_with( - "inserted_at", - expireAfterSeconds=3600, - background=True + 'inserted_at', expireAfterSeconds=3600, background=True ) collection_2.list_indexes.assert_called_once() collection_2.create_index.assert_called_once_with( - "inserted_at", - expireAfterSeconds=3600, - background=True + 'inserted_at', expireAfterSeconds=3600, background=True ) collection_3.list_indexes.assert_called_once() @@ -515,33 +464,26 @@ async def test_create_collection_with_ttl_index_success(mongo_db): @mark.asyncio async def test_create_collection_with_ttl_index_failure(mongo_db): - input_data = { - **metadata, - "pipelines": { - "scouter-pipeline": { - "topic": "raw_scouter_pipeline" - } - } - } + input_data = {**metadata, 'pipelines': {'scouter-pipeline': {'topic': 'raw_scouter_pipeline'}}} 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: await mongo_db.create_collection_with_ttl_index(input_data) except Exception as e: - assert str(e) == "Error" + assert str(e) == 'Error' mongo_db.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="MONGODB_CREATE_COLLECTION_ERROR", - message="Failed to create collection raw_scouter_pipeline with TTL index: Error", - block="create_collection_with_ttl_index", + notification_id='MONGODB_CREATE_COLLECTION_ERROR', + message='Failed to create collection raw_scouter_pipeline with TTL index: Error', + block='create_collection_with_ttl_index', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "Expected an exception to be raised" + raise AssertionError('Expected an exception to be raised') @mark.asyncio @@ -555,34 +497,25 @@ async def test_load_latest_data_none_last_data_timestamp(mongo_db): 'name': 'test1', 'value': 1, '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({ - '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( + result = await mongo_db.load_latest_data( { - 'level': 'ERROR' - }, - {"_id": 0} + 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, + 'collection_name': 'test_collection', + 'last_data_timestamp': None, + 'base_data_filter': {'level': 'ERROR'}, + } ) - assert result == [{ - 'name': 'test1', - 'value': 1, - 'timestamp': '2023-01-01 12:00:00.000000+0000' - }] + mongo_db.database.__getitem__.assert_called_once_with('test_collection') + + collection.find.assert_called_once_with({'level': 'ERROR'}, {'_id': 0}) + + assert result == [{'name': 'test1', 'value': 1, 'timestamp': '2023-01-01 12:00:00.000000+0000'}] @mark.asyncio @@ -596,38 +529,35 @@ async def test_load_latest_data_not_none_last_data_timestamp(mongo_db): 'name': 'test1', 'value': 1, '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({ - 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, - 'collection_name': 'test_collection', - 'last_data_timestamp': '2023-01-01 12:00:00.000000+0000', - 'base_data_filter': { - 'level': 'ERROR' + result = await mongo_db.load_latest_data( + { + 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, + 'collection_name': 'test_collection', + 'last_data_timestamp': '2023-01-01 12:00:00.000000+0000', + 'base_data_filter': {'level': 'ERROR'}, } - }) + ) - mongo_db.database.__getitem__.assert_called_once_with( - 'test_collection') + mongo_db.database.__getitem__.assert_called_once_with('test_collection') collection.find.assert_called_once_with( { 'level': 'ERROR', 'timestamp': { '$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 == [{ - 'name': 'test1', - 'value': 1, - '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 @@ -640,23 +570,22 @@ async def test_load_latest_data_error(mongo_db): collection.find.side_effect = Exception('test') try: - await mongo_db.load_latest_data({ - 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, - 'collection_name': 'test_collection', - 'last_data_timestamp': '2023-01-01 12:00:00.000000+0000', - 'base_data_filter': { - 'level': 'ERROR' + await mongo_db.load_latest_data( + { + 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, + 'collection_name': 'test_collection', + 'last_data_timestamp': '2023-01-01 12:00:00.000000+0000', + 'base_data_filter': {'level': 'ERROR'}, } - }) + ) except Exception as e: assert str(e) == 'test' mongo_db.send_notification.assert_called_once_with( - metadata={'workflow_name': 'test_pipeline', - 'schedule_name': 'test_schedule'}, + metadata={'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, notification_id='MONGO_LOAD_ERROR', message='Error loading data from MongoDB: test', block='load_latest_data', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) diff --git a/tests/orchestrator/activities/test_slot_manager.py b/tests/orchestrator/activities/test_slot_manager.py index 80cd384..fe6be1b 100644 --- a/tests/orchestrator/activities/test_slot_manager.py +++ b/tests/orchestrator/activities/test_slot_manager.py @@ -1,32 +1,33 @@ -from unittest.mock import MagicMock, patch, call, ANY -from datetime import datetime, timedelta +from datetime import timedelta +from unittest.mock import ANY, MagicMock, call, patch + from pandas import DataFrame -from pytest import mark, fixture -from orchestrator.activities.slot_manager import SlotManager +from pytest import fixture, mark from sientia_do.notifications.models import NotificationLevel from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now +from orchestrator.activities.slot_manager import SlotManager + metadata = { - "metadata": { - "schedule_name": "test_schedule_name", - "workflow_name": "test_workflow_name", - "model_name": "test_model_name", - "model_id": "test_model_id" + 'metadata': { + 'schedule_name': 'test_schedule_name', + 'workflow_name': 'test_workflow_name', + 'model_name': 'test_model_name', + 'model_id': 'test_model_id', } } @fixture -@patch("orchestrator.activities.slot_manager.Redis.__init__") +@patch('orchestrator.activities.slot_manager.Redis.__init__') def slot_manager(_redis_mock): - slot_manager = SlotManager( - host="localhost", + host='localhost', port=6379, - username="admin", - password="password", + username='admin', + password='password', logger=MagicMock(), - notification_handler=MagicMock() + notification_handler=MagicMock(), ) slot_manager.redis_client = MagicMock() @@ -46,168 +47,137 @@ async def test_load_opc_slots_no_slot_keys(slot_manager): @mark.asyncio async def test_load_opc_slots(slot_manager): 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( - side_effect=[ - "value1", - "value2", - None - ] - ) + slot_manager.get = MagicMock(side_effect=['value1', 'value2', None]) response = await slot_manager.load_opc_slots(metadata) assert response == { - "slot:opc_tags:1": "value1", - "slot:opc_tags:2": "value2", - "slot:opc_tags:3": None + 'slot:opc_tags:1': 'value1', + 'slot:opc_tags:2': 'value2', + 'slot:opc_tags:3': None, } @mark.asyncio async def test_load_opc_slots_no_decode(slot_manager): 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( - side_effect=[ - "value1", - "value2", - None - ] - ) + slot_manager.get = MagicMock(side_effect=['value1', 'value2', None]) response = await slot_manager.load_opc_slots(metadata) assert response == { - "slot:opc_tags:1": "value1", - "slot:opc_tags:2": "value2", - "slot:opc_tags:3": None + 'slot:opc_tags:1': 'value1', + 'slot:opc_tags:2': 'value2', + 'slot:opc_tags:3': None, } @mark.asyncio async def test_load_opc_slots_error(slot_manager): 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( - side_effect=Exception("Test exception") - ) + slot_manager.get = MagicMock(side_effect=Exception('Test exception')) try: await slot_manager.load_opc_slots(metadata) except Exception as e: - assert str(e) == "Test exception" + assert str(e) == 'Test exception' slot_manager.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="REDIS_GET_ERROR", - message="Failed to load OPC slots: Test exception", - block="load_opc_slots", + notification_id='REDIS_GET_ERROR', + message='Failed to load OPC slots: Test exception', + block='load_opc_slots', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "Expected an exception to be raised" + raise AssertionError('Expected an exception to be raised') @mark.asyncio async def test_load_active_ingestors(slot_manager): 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) - assert response == ["heartbeat:ingestor:1", - "heartbeat:ingestor:2", "heartbeat:ingestor:3"] + assert response == ['heartbeat:ingestor:1', 'heartbeat:ingestor:2', 'heartbeat:ingestor:3'] @mark.asyncio async def test_load_active_ingestors_error(slot_manager): slot_manager.redis_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: await slot_manager.load_active_ingestors(metadata) except Exception as e: - assert str(e) == "Test exception" + assert str(e) == 'Test exception' slot_manager.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="REDIS_GET_ERROR", - message="Failed to load active ingestors: Test exception", - block="load_active_ingestors", + notification_id='REDIS_GET_ERROR', + message='Failed to load active ingestors: Test exception', + block='load_active_ingestors', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "Expected an exception to be raised" + raise AssertionError('Expected an exception to be raised') @mark.asyncio async def test_update_slots(slot_manager): - slot_manager.set = MagicMock( - side_effect=[ - None, - Exception("Test exception") - ] + slot_manager.set = MagicMock(side_effect=[None, Exception('Test exception')]) + + 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)] ) - 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 == { - "1": { - "success": True, - "message": "Slot updated successfully" - }, - "2": { - "success": False, - "message": "Test exception" - } + '1': {'success': True, 'message': 'Slot updated successfully'}, + '2': {'success': False, 'message': 'Test exception'}, } @mark.asyncio async def test_delete_slots(slot_manager): - slot_manager.redis_client.delete = MagicMock( - side_effect=[ - None, - Exception("Test exception") - ] + slot_manager.redis_client.delete = MagicMock(side_effect=[None, Exception('Test exception')]) + + 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')] ) - 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 == { - "1": { - "success": True, - "message": "Slot deleted successfully" - }, - "2": { - "success": False, - "message": "Test exception" - } + '1': {'success': True, '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, 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule', - 'mail_type': 'test_mail_type' + 'mail_type': 'test_mail_type', } slot_manager.get = MagicMock(return_value=None) @@ -235,16 +205,14 @@ async def test_get_last_data_timestamp_not_none(slot_manager): **metadata, 'workflow_name': 'test_pipeline', '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') result = await slot_manager.get_last_data_timestamp(test_data) - slot_manager.get.assert_called_once_with( - 'notification_last_timestamp:test_mail_type' - ) + slot_manager.get.assert_called_once_with('notification_last_timestamp:test_mail_type') assert result == '2023-01-01 12:00:00' @@ -256,14 +224,13 @@ async def test_get_last_data_timestamp_error(slot_manager): **metadata, 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule', - 'mail_type': 'test_mail_type' + 'mail_type': 'test_mail_type', } slot_manager.send_notification = MagicMock() slot_manager.get = MagicMock(side_effect=Exception('test')) try: - await slot_manager.get_last_data_timestamp(test_data) 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( metadata=metadata['metadata'], - notification_id="REDIS_GET_ERROR", - message="Error getting last data timestamp: test", - block="get_last_data_timestamp", + notification_id='REDIS_GET_ERROR', + message='Error getting last data timestamp: test', + block='get_last_data_timestamp', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "Expected exception" + raise AssertionError('Expected exception') @mark.asyncio @@ -290,7 +257,7 @@ async def test_put_last_data_timestamp_empty_dataframe(slot_manager): 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule', 'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records'), - 'mail_type': 'test_mail_type' + 'mail_type': 'test_mail_type', } 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): """Test put_last_data_timestamp with not empty dataframe""" - data = DataFrame({ - 'name': ['sensor1', 'sensor2'], - 'value': [25.5, 30.0], - 'timestamp': ['2023-01-01 12:00:00', '2023-01-01 12:00:01'] - }) + data = DataFrame( + { + 'name': ['sensor1', 'sensor2'], + 'value': [25.5, 30.0], + 'timestamp': ['2023-01-01 12:00:00', '2023-01-01 12:00:01'], + } + ) test_data = { **metadata, 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule', 'data': data.to_dict('records'), - 'mail_type': 'test_mail_type' + 'mail_type': 'test_mail_type', } 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' slot_manager.set.assert_called_once_with( - 'notification_last_timestamp:test_mail_type', - '2023-01-01 12:00:01', - ttl=18000 + 'notification_last_timestamp:test_mail_type', '2023-01-01 12:00:01', ttl=18000 ) @@ -339,12 +306,14 @@ async def test_put_last_data_timestamp_error(slot_manager): **metadata, 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule', - 'data': DataFrame({ - 'name': ['sensor1', 'sensor2'], - 'value': [25.5, 30.0], - 'timestamp': ['2023-01-01 12:00:00'] * 2 - }).to_dict('records'), - 'mail_type': 'test_mail_type' + 'data': DataFrame( + { + 'name': ['sensor1', 'sensor2'], + 'value': [25.5, 30.0], + 'timestamp': ['2023-01-01 12:00:00'] * 2, + } + ).to_dict('records'), + 'mail_type': 'test_mail_type', } 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( metadata=metadata['metadata'], - notification_id="REDIS_SET_ERROR", - message="Error setting last data timestamp: test", - block="put_last_data_timestamp", + notification_id='REDIS_SET_ERROR', + message='Error setting last data timestamp: test', + block='put_last_data_timestamp', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "Expected exception" + raise AssertionError('Expected exception') @mark.asyncio async def test_filter_notification_alerts(slot_manager): - slot_manager.get = MagicMock(side_effect=[ - None, - (now() - timedelta(seconds=600) - ).strftime(DATETIME_FORMAT_MS_WITH_TZ), - now().strftime(DATETIME_FORMAT_MS_WITH_TZ), - None, - (now() - timedelta(seconds=600) - ).strftime(DATETIME_FORMAT_MS_WITH_TZ), - now().strftime(DATETIME_FORMAT_MS_WITH_TZ)]) + slot_manager.get = MagicMock( + side_effect=[ + None, + (now() - timedelta(seconds=600)).strftime(DATETIME_FORMAT_MS_WITH_TZ), + now().strftime(DATETIME_FORMAT_MS_WITH_TZ), + None, + (now() - timedelta(seconds=600)).strftime(DATETIME_FORMAT_MS_WITH_TZ), + now().strftime(DATETIME_FORMAT_MS_WITH_TZ), + ] + ) input_data = { **metadata, 'notification_package': [ - { - 'trigger': 'test_trigger_1', - 'notification_id': 'test_notification_id_1' - }, - { - 'trigger': 'test_trigger_2', - 'notification_id': 'test_notification_id_2' - }, - { - 'trigger': 'test_trigger_3', - 'notification_id': 'test_notification_id_3' - } + {'trigger': 'test_trigger_1', 'notification_id': 'test_notification_id_1'}, + {'trigger': 'test_trigger_2', 'notification_id': 'test_notification_id_2'}, + {'trigger': 'test_trigger_3', 'notification_id': 'test_notification_id_3'}, ], 'sending_configs': [ - { - 'group_name': 'test_group_1', - 'contents': ['core_alerts', 'persistent_alerts'] - }, - { - 'group_name': 'test_group_2', - 'contents': ['core_alerts'] - } + {'group_name': 'test_group_1', 'contents': ['core_alerts', 'persistent_alerts']}, + {'group_name': 'test_group_2', 'contents': ['core_alerts']}, ], 'notification_ttl': 300, - 'mail_type': 'test_mail_type' + 'mail_type': 'test_mail_type', } 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', 'contents': ['core_alerts', 'persistent_alerts'], 'notifications': [ - { - 'trigger': 'test_trigger_1', - 'notification_id': 'test_notification_id_1' - }, - { - 'trigger': 'test_trigger_2', - 'notification_id': 'test_notification_id_2' - } - ] + {'trigger': 'test_trigger_1', 'notification_id': 'test_notification_id_1'}, + {'trigger': 'test_trigger_2', 'notification_id': 'test_notification_id_2'}, + ], }, 'test_group_2': { 'group_name': 'test_group_2', 'contents': ['core_alerts'], '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_data = { **metadata, - 'log_report': DataFrame({ - 'status': ['sent', 'error'], - 'schedule': ['test_schedule_1', 'test_schedule_2'], - 'notification_id': ['test_notification_id_1', 'test_notification_id_2'] - }).to_dict(), - 'sent_ttl': 600 + 'log_report': DataFrame( + { + 'status': ['sent', 'error'], + 'schedule': ['test_schedule_1', 'test_schedule_2'], + 'notification_id': ['test_notification_id_1', 'test_notification_id_2'], + } + ).to_dict(), + 'sent_ttl': 600, } slot_manager.set = MagicMock() await slot_manager.store_notification_cache(test_data) - slot_manager.set.assert_called_once_with( - "test_schedule_1:test_notification_id_1", - ANY, - ttl=600 - ) + slot_manager.set.assert_called_once_with('test_schedule_1:test_notification_id_1', ANY, ttl=600) diff --git a/tests/orchestrator/activities/test_temporal_manager.py b/tests/orchestrator/activities/test_temporal_manager.py index ab658ae..89ee56b 100644 --- a/tests/orchestrator/activities/test_temporal_manager.py +++ b/tests/orchestrator/activities/test_temporal_manager.py @@ -1,29 +1,31 @@ -from unittest.mock import MagicMock, patch, AsyncMock, call, ANY 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 sientia_do.notifications.models import NotificationLevel + from orchestrator.activities.temporal_manager import TemporalManager from orchestrator.utils.converters import parse_frequency metadata = { - "metadata": { - "schedule_name": "test_schedule_name", - "workflow_name": "test_workflow_name", - "model_name": "test_model_name", - "model_id": "test_model_id" + 'metadata': { + 'schedule_name': 'test_schedule_name', + 'workflow_name': 'test_workflow_name', + 'model_name': 'test_model_name', + 'model_id': 'test_model_id', } } @fixture -@patch("orchestrator.activities.temporal_manager.Client.connect") +@patch('orchestrator.activities.temporal_manager.Client.connect') def temporal_manager(connect_mock): temporal_manager = TemporalManager( host='localhost:7233', scouter_namespace='scouter', laborious_namespace='laborious', logger=MagicMock(), - notification_handler=MagicMock() + notification_handler=MagicMock(), ) temporal_manager.temporal_clients['scouter'] = MagicMock() @@ -34,50 +36,31 @@ def temporal_manager(connect_mock): @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): await temporal_manager.connect_to_temporal() - connect_mock.assert_has_calls([ - call( - target_host='localhost:7233', - namespace='scouter' - ), - call( - target_host='localhost:7233', - namespace='laborious' - ) - ]) + connect_mock.assert_has_calls( + [ + call(target_host='localhost:7233', namespace='scouter'), + call(target_host='localhost:7233', namespace='laborious'), + ] + ) async def async_iter(): - yield MagicMock( - id="test-schedule-id", - search_attributes={ - "orchestrated": ["true"] - } - ) - yield MagicMock( - id="test-schedule-id-2", - search_attributes={ - "Attr": ["false"] - } - ) - yield MagicMock( - id="test-schedule-id-3", - search_attributes={ - "Attr": ["false"] - } - ) + yield MagicMock(id='test-schedule-id', search_attributes={'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 async def test_normalize_schedules(temporal_manager): input_data = { - "orchestrated_schedules": { - "scouter": {"test-scouter": "2021-01-01"}, - "laborious": {"test-schedule-id1": "2021-01-01"} + 'orchestrated_schedules': { + 'scouter': {'test-scouter': '2021-01-01'}, + 'laborious': {'test-schedule-id1': '2021-01-01'}, }, - **metadata + **metadata, } 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( - return_value=MagicMock( - delete=AsyncMock() - ) + return_value=MagicMock(delete=AsyncMock()) ) temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock( - return_value=MagicMock( - delete=AsyncMock() - ) + return_value=MagicMock(delete=AsyncMock()) ) await temporal_manager.normalize_schedules(input_data) - temporal_manager.temporal_clients['scouter'].get_schedule_handle.assert_has_calls([ - 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.assert_has_calls( + [ + call('test-schedule-id'), + ] ) + temporal_manager.temporal_clients[ + 'scouter' + ].get_schedule_handle.return_value.delete.assert_awaited_once() + @mark.asyncio async def test_normalize_schedules_error(temporal_manager): temporal_manager.temporal_clients['scouter'].list_schedules = AsyncMock( - side_effect=Exception("Test exception") + side_effect=Exception('Test exception') ) try: await temporal_manager.normalize_schedules(metadata) except Exception as e: - assert str(e) == "Test exception" + assert str(e) == 'Test exception' temporal_manager.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="TEMPORAL_NORMALIZE_SCHEDULES_ERROR", - message="Failed to normalize schedules: Test exception", - block="normalize_schedules", + notification_id='TEMPORAL_NORMALIZE_SCHEDULES_ERROR', + message='Failed to normalize schedules: Test exception', + block='normalize_schedules', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "Expected an exception to be raised" + raise AssertionError('Expected an exception to be raised') @mark.asyncio -@patch("orchestrator.activities.temporal_manager.parse_frequency", - side_effect=parse_frequency) -@patch("orchestrator.activities.temporal_manager.Schedule") -@patch("orchestrator.activities.temporal_manager.ScheduleActionStartWorkflow") -@patch("orchestrator.activities.temporal_manager.ScheduleIntervalSpec") -@patch("orchestrator.activities.temporal_manager.ScheduleSpec") -@patch("orchestrator.activities.temporal_manager.TypedSearchAttributes") -@patch("orchestrator.activities.temporal_manager.SearchAttributePair") +@patch('orchestrator.activities.temporal_manager.parse_frequency', side_effect=parse_frequency) +@patch('orchestrator.activities.temporal_manager.Schedule') +@patch('orchestrator.activities.temporal_manager.ScheduleActionStartWorkflow') +@patch('orchestrator.activities.temporal_manager.ScheduleIntervalSpec') +@patch('orchestrator.activities.temporal_manager.ScheduleSpec') +@patch('orchestrator.activities.temporal_manager.TypedSearchAttributes') +@patch('orchestrator.activities.temporal_manager.SearchAttributePair') async def test_create_schedule( - mock_search_attribute_pair, - mock_typed_search_attributes, - mock_schedule_spec, - mock_schedule_interval_spec, - mock_schedule_action_start_workflow, - mock_schedule, - mock_parse_frequency, - temporal_manager): - + mock_search_attribute_pair, + mock_typed_search_attributes, + mock_schedule_spec, + mock_schedule_interval_spec, + mock_schedule_action_start_workflow, + mock_schedule, + mock_parse_frequency, + temporal_manager, +): input_data = { - "schedules": { - "scouter": { - "test-schedule": { - "model_id": 1, - "model_name": "test-model-name", - "workflow_type": "test-workflow", - "frequency": "1m", - "data": {"test": "test"} + 'schedules': { + 'scouter': { + 'test-schedule': { + 'model_id': 1, + 'model_name': 'test-model-name', + 'workflow_type': 'test-workflow', + 'frequency': '1m', + 'data': {'test': 'test'}, + 'execution_timeout_seconds': 100, + 'task_timeout_seconds': 100, }, - "test-schedule-invalid-frequency": { - "model_id": 2, - "model_name": "test-model-name", - "workflow_type": "test-workflow", - "frequency": "10y", - "data": {"test": "test"} + 'test-schedule-invalid-frequency': { + 'model_id': 2, + 'model_name': 'test-model-name', + 'workflow_type': 'test-workflow', + 'frequency': '10y', + '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['laborious'].create_schedule = AsyncMock( - ) + temporal_manager.temporal_clients['laborious'].create_schedule = AsyncMock() report = await temporal_manager.create_schedules(input_data) temporal_manager.temporal_clients['scouter'].create_schedule.assert_called_once_with( - "test-schedule", + 'test-schedule', 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( - "test-schedule-laborious", + 'test-schedule-laborious', 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([ - call( - action=mock_schedule_action_start_workflow.return_value, - spec=mock_schedule_spec.return_value - ), - call( - action=mock_schedule_action_start_workflow.return_value, - spec=mock_schedule_spec.return_value - ) - ]) + mock_schedule.assert_has_calls( + [ + call( + action=mock_schedule_action_start_workflow.return_value, + spec=mock_schedule_spec.return_value, + ), + call( + action=mock_schedule_action_start_workflow.return_value, + spec=mock_schedule_spec.return_value, + ), + ] + ) - mock_schedule_action_start_workflow.assert_has_calls([ - call( - "test-workflow", - input_data['schedules']['scouter']['test-schedule'], - id="test-schedule", - task_queue="test-workflow-queue", - execution_timeout=ANY, - typed_search_attributes=mock_typed_search_attributes.return_value, - ), - call( - "test-workflow", - input_data['schedules']['scouter']['test-schedule-invalid-frequency'], - id="test-schedule-invalid-frequency", - task_queue="test-workflow-queue", - execution_timeout=ANY, - typed_search_attributes=mock_typed_search_attributes.return_value, - ), - call( - "test-workflow", - input_data['schedules']['laborious']['test-schedule-laborious'], - id="test-schedule-laborious", - task_queue="test-workflow-queue", - execution_timeout=ANY, - typed_search_attributes=mock_typed_search_attributes.return_value, - ) - ]) + mock_schedule_action_start_workflow.assert_has_calls( + [ + call( + 'test-workflow', + input_data['schedules']['scouter']['test-schedule'], + id='test-schedule', + task_queue='test-workflow-queue', + execution_timeout=timedelta(seconds=100), + run_timeout=timedelta(seconds=100), + task_timeout=timedelta(seconds=100), + typed_search_attributes=mock_typed_search_attributes.return_value, + ), + call( + 'test-workflow', + input_data['schedules']['scouter']['test-schedule-invalid-frequency'], + id='test-schedule-invalid-frequency', + task_queue='test-workflow-queue', + execution_timeout=timedelta(seconds=400), + run_timeout=timedelta(seconds=400), + task_timeout=timedelta(seconds=400), + typed_search_attributes=mock_typed_search_attributes.return_value, + ), + call( + '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([ - call( - intervals=[ - mock_schedule_interval_spec.return_value - ] - ), - call( - intervals=[ - mock_schedule_interval_spec.return_value - ] - ) - ]) + mock_schedule_spec.assert_has_calls( + [ + call(intervals=[mock_schedule_interval_spec.return_value]), + call(intervals=[mock_schedule_interval_spec.return_value]), + ] + ) - mock_schedule_interval_spec.assert_has_calls([ - call( - every=timedelta(seconds=60) - ), - call( - every=timedelta(seconds=120) - ) - ]) + mock_schedule_interval_spec.assert_has_calls( + [call(every=timedelta(seconds=60)), call(every=timedelta(seconds=120))] + ) - mock_parse_frequency.assert_has_calls([ - call("1m"), - call("10y"), - call("2m") - ]) + mock_parse_frequency.assert_has_calls([call('1m'), call('10y'), call('2m')]) - mock_typed_search_attributes.assert_has_calls([ - call([ - mock_search_attribute_pair.return_value, - mock_search_attribute_pair.return_value, - mock_search_attribute_pair.return_value - ]), - call([ - mock_search_attribute_pair.return_value, - mock_search_attribute_pair.return_value, - mock_search_attribute_pair.return_value - ]), - call([ - mock_search_attribute_pair.return_value, - mock_search_attribute_pair.return_value, - mock_search_attribute_pair.return_value - ]) - ]) + mock_typed_search_attributes.assert_has_calls( + [ + call( + [ + mock_search_attribute_pair.return_value, + mock_search_attribute_pair.return_value, + mock_search_attribute_pair.return_value, + ] + ), + call( + [ + mock_search_attribute_pair.return_value, + mock_search_attribute_pair.return_value, + mock_search_attribute_pair.return_value, + ] + ), + call( + [ + mock_search_attribute_pair.return_value, + mock_search_attribute_pair.return_value, + mock_search_attribute_pair.return_value, + ] + ), + ] + ) - mock_search_attribute_pair.assert_has_calls([ - call( - key=temporal_manager.model_id_id_key, - value=1 - ), - call( - key=temporal_manager.model_name_id_key, - value="test-model-name" - ), - call( - key=temporal_manager.orchestrated_id_key, - value="true" - ), - call( - key=temporal_manager.model_id_id_key, - value=2 - ), - call( - key=temporal_manager.model_name_id_key, - value="test-model-name" - ), - call( - key=temporal_manager.orchestrated_id_key, - value="true" - ) - ]) + mock_search_attribute_pair.assert_has_calls( + [ + call(key=temporal_manager.model_id_id_key, value=1), + call(key=temporal_manager.model_name_id_key, value='test-model-name'), + call(key=temporal_manager.orchestrated_id_key, value='true'), + call(key=temporal_manager.model_id_id_key, value=2), + call(key=temporal_manager.model_name_id_key, value='test-model-name'), + call(key=temporal_manager.orchestrated_id_key, value='true'), + ] + ) assert report == [ { - "schedule_name": "test-schedule", - "namespace": "scouter", - "success": True, - "message": "Schedule created successfully" + 'schedule_name': 'test-schedule', + 'namespace': 'scouter', + 'success': True, + 'message': 'Schedule created successfully', }, { - "schedule_name": "test-schedule-invalid-frequency", - "namespace": "scouter", - "success": False, - "message": "Invalid frequency" + 'schedule_name': 'test-schedule-invalid-frequency', + 'namespace': 'scouter', + 'success': False, + 'message': 'Invalid frequency', }, { - "schedule_name": "test-schedule-laborious", - "namespace": "laborious", - "success": True, - "message": "Schedule created successfully" - } + 'schedule_name': 'test-schedule-laborious', + 'namespace': 'laborious', + 'success': True, + 'message': 'Schedule created successfully', + }, ] @@ -334,136 +307,94 @@ async def test_create_schedule( async def test_create_schedules_with_no_client(temporal_manager): temporal_manager.temporal_clients = {} input_data = { - "schedules": { - "abc": { - "test-schedule": { - "frequency": "1m", - "data": {"test": "test"} - } - } - } + 'schedules': {'abc': {'test-schedule': {'frequency': '1m', 'data': {'test': 'test'}}}} } try: await temporal_manager.create_schedules(input_data) except Exception as e: - assert str( - e) == f"Temporal client for abc not found, clients: {temporal_manager.temporal_clients}" + assert ( + str(e) + == f'Temporal client for abc not found, clients: {temporal_manager.temporal_clients}' + ) @mark.asyncio -@patch("orchestrator.activities.temporal_manager.parse_frequency", - side_effect=parse_frequency) -@patch("orchestrator.activities.temporal_manager.ScheduleIntervalSpec") +@patch('orchestrator.activities.temporal_manager.parse_frequency', side_effect=parse_frequency) +@patch('orchestrator.activities.temporal_manager.ScheduleIntervalSpec') async def test_update_schedules( - _mock_schedule_interval_spec, - _mock_parse_frequency, - temporal_manager): - input_mock = MagicMock( - args=MagicMock() - ) + _mock_schedule_interval_spec, _mock_parse_frequency, temporal_manager +): + input_mock = MagicMock(args=MagicMock()) temporal_manager.schedule_handles = { - "scouter": { - "test-schedule": MagicMock( - update=AsyncMock( - update=AsyncMock( - side_effect=lambda f: f(input_mock) - ) - ) + 'scouter': { + 'test-schedule': MagicMock( + update=AsyncMock(update=AsyncMock(side_effect=lambda f: f(input_mock))) ) }, - "laborious": { - "test-schedule-laborious": MagicMock( - update=AsyncMock( - update=AsyncMock( - side_effect=lambda f: f(input_mock) - ) - ) + 'laborious': { + 'test-schedule-laborious': MagicMock( + update=AsyncMock(update=AsyncMock(side_effect=lambda f: f(input_mock))) ) - } + }, } input_data = { - "schedules": { - "scouter": { - "test-schedule": { - "frequency": "1m", - "data": {"test": "test"} - }, - "test-schedule_no_handler": { - "frequency": "1m", - "data": {"test": "test"} - } + 'schedules': { + 'scouter': { + 'test-schedule': {'frequency': '1m', 'data': {'test': 'test'}}, + 'test-schedule_no_handler': {'frequency': '1m', 'data': {'test': 'test'}}, }, - "laborious": { - "test-schedule-laborious": { - "frequency": "2m", - "data": {"test": "test"} - } - } + 'laborious': {'test-schedule-laborious': {'frequency': '2m', 'data': {'test': 'test'}}}, } } handler_scouter = MagicMock( - update=AsyncMock( - update=AsyncMock( - side_effect=lambda f: f(input_mock) - ) - ) + update=AsyncMock(update=AsyncMock(side_effect=lambda f: f(input_mock))) ) handler_laborious = MagicMock( - update=AsyncMock( - update=AsyncMock( - side_effect=lambda f: f(input_mock) - ) - ) + update=AsyncMock(update=AsyncMock(side_effect=lambda f: f(input_mock))) ) temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock( - side_effect=[ - handler_scouter, - None - ] + side_effect=[handler_scouter, None] ) temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock( - side_effect=[ - handler_laborious - ] + side_effect=[handler_laborious] ) report = await temporal_manager.update_schedules(input_data) - temporal_manager.temporal_clients['scouter'].get_schedule_handle.assert_has_calls([ - call("test-schedule"), - call("test-schedule_no_handler") - ]) - temporal_manager.temporal_clients['laborious'].get_schedule_handle.assert_has_calls([ - call("test-schedule-laborious") - ]) + temporal_manager.temporal_clients['scouter'].get_schedule_handle.assert_has_calls( + [call('test-schedule'), call('test-schedule_no_handler')] + ) + temporal_manager.temporal_clients['laborious'].get_schedule_handle.assert_has_calls( + [call('test-schedule-laborious')] + ) handler_scouter.update.assert_called_once() handler_laborious.update.assert_called_once() assert report == [ { - "schedule_name": "test-schedule", - "namespace": "scouter", - "success": True, - "message": "Schedule updated successfully" + 'schedule_name': 'test-schedule', + 'namespace': 'scouter', + 'success': True, + 'message': 'Schedule updated successfully', }, { - "schedule_name": "test-schedule_no_handler", - "namespace": "scouter", - "success": False, - "message": "Schedule test-schedule_no_handler not found" + 'schedule_name': 'test-schedule_no_handler', + 'namespace': 'scouter', + 'success': False, + 'message': 'Schedule test-schedule_no_handler not found', }, { - "schedule_name": "test-schedule-laborious", - "namespace": "laborious", - "success": True, - "message": "Schedule updated successfully" - } + 'schedule_name': 'test-schedule-laborious', + 'namespace': 'laborious', + 'success': True, + 'message': 'Schedule updated successfully', + }, ] @@ -471,67 +402,41 @@ async def test_update_schedules( async def test_update_schedules_with_no_client(temporal_manager): temporal_manager.temporal_clients = {} input_data = { - "schedules": { - "abc": { - "test-schedule": { - "frequency": "1m", - "data": {"test": "test"} - } - } - } + 'schedules': {'abc': {'test-schedule': {'frequency': '1m', 'data': {'test': 'test'}}}} } try: await temporal_manager.update_schedules(input_data) except Exception as e: - assert str( - e) == f"Temporal client for abc not found, clients: {temporal_manager.temporal_clients}" + assert ( + str(e) + == f'Temporal client for abc not found, clients: {temporal_manager.temporal_clients}' + ) @mark.asyncio async def test_delete_schedules(temporal_manager): temporal_manager.schedule_handles = { - "scouter": { - "test-schedule": MagicMock( - delete=AsyncMock() - ) - }, - "laborious": { - "test-schedule-laborious": MagicMock( - delete=AsyncMock() - ) - } + 'scouter': {'test-schedule': MagicMock(delete=AsyncMock())}, + 'laborious': {'test-schedule-laborious': MagicMock(delete=AsyncMock())}, } input_data = { - "schedules": { - "scouter": [ - "test-schedule", "test-schedule_no_handler" - ], - "laborious": [ - "test-schedule-laborious" - ] + 'schedules': { + 'scouter': ['test-schedule', 'test-schedule_no_handler'], + 'laborious': ['test-schedule-laborious'], } } - handler_scouter = MagicMock( - delete=AsyncMock() - ) + handler_scouter = MagicMock(delete=AsyncMock()) - handler_laborious = MagicMock( - delete=AsyncMock() - ) + handler_laborious = MagicMock(delete=AsyncMock()) temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock( - side_effect=[ - handler_scouter, - None - ] + side_effect=[handler_scouter, None] ) temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock( - side_effect=[ - handler_laborious - ] + side_effect=[handler_laborious] ) report = await temporal_manager.delete_schedules(input_data) @@ -541,40 +446,36 @@ async def test_delete_schedules(temporal_manager): assert report == [ { - "schedule_name": "test-schedule", - "namespace": "scouter", - "success": True, - "message": "Schedule deleted successfully" + 'schedule_name': 'test-schedule', + 'namespace': 'scouter', + 'success': True, + 'message': 'Schedule deleted successfully', }, { - "schedule_name": "test-schedule_no_handler", - "namespace": "scouter", - "success": False, - "message": "Schedule test-schedule_no_handler not found", - "attachment": ANY + 'schedule_name': 'test-schedule_no_handler', + 'namespace': 'scouter', + 'success': False, + 'message': 'Schedule test-schedule_no_handler not found', + 'attachment': ANY, }, { - "schedule_name": "test-schedule-laborious", - "namespace": "laborious", - "success": True, - "message": "Schedule deleted successfully" - } + 'schedule_name': 'test-schedule-laborious', + 'namespace': 'laborious', + 'success': True, + 'message': 'Schedule deleted successfully', + }, ] @mark.asyncio async def test_delete_schedules_with_no_client(temporal_manager): temporal_manager.temporal_clients = {} - input_data = { - "schedules": { - "abc": [ - "test-schedule" - ] - } - } + input_data = {'schedules': {'abc': ['test-schedule']}} try: await temporal_manager.delete_schedules(input_data) except Exception as e: - assert str( - e) == f"Temporal client for abc not found, clients: {temporal_manager.temporal_clients}" + assert ( + str(e) + == f'Temporal client for abc not found, clients: {temporal_manager.temporal_clients}' + ) diff --git a/tests/orchestrator/utils/test_connectors_config.py b/tests/orchestrator/utils/test_connectors_config.py index 13977fa..33d982b 100644 --- a/tests/orchestrator/utils/test_connectors_config.py +++ b/tests/orchestrator/utils/test_connectors_config.py @@ -1,10 +1,13 @@ from os import environ -from orchestrator.utils.connectors_config import (build_redis_config, - build_couchbase_config, - build_mongodb_config, - build_temporal_config, - build_email_config, - build_postgres_config) + +from orchestrator.utils.connectors_config import ( + build_couchbase_config, + build_email_config, + build_mongodb_config, + build_postgres_config, + build_redis_config, + build_temporal_config, +) def test_build_redis_config_with_env_vars(): @@ -16,7 +19,7 @@ def test_build_redis_config_with_env_vars(): 'host': 'localhost', 'port': 6379, 'username': 'sientia', - 'password': 'sientia' + 'password': 'sientia', } @@ -27,7 +30,7 @@ def test_build_couchbase_config_with_env_vars(): assert build_couchbase_config() == { 'connection_string': 'couchbase://localhost', 'username': 'sientia', - 'password': 'sientia' + 'password': 'sientia', } @@ -40,7 +43,7 @@ def test_build_redis_config_with_defaults(): 'host': 'localhost', 'port': 6379, 'username': 'default', - 'password': 'bdnZOpcyiL' + 'password': 'bdnZOpcyiL', } @@ -51,7 +54,7 @@ def test_build_couchbase_config_with_defaults(): assert build_couchbase_config() == { 'connection_string': 'couchbase://localhost', 'username': 'sientia', - 'password': 'sientia' + 'password': 'sientia', } @@ -64,7 +67,7 @@ def test_build_mongo_db_config_with_env_vars(): assert build_mongodb_config() == { 'connection_string': 'mongodb://sientia1:sientia1@localhost:27018', '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() == { 'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018', '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_namespace': 'default', '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_namespace': 'default', '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_password': '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_password': 'sientia', 'smtp_server': None, - 'smtp_port': 587 + 'smtp_port': 587, } @@ -147,7 +150,7 @@ def test_build_postgres_config_with_env_vars(): 'password': 'sientia', 'dbname': 'sientia', 'min_connections': 5, - 'max_connections': 20 + 'max_connections': 20, } @@ -167,5 +170,5 @@ def test_build_postgres_config_with_defaults(): 'password': 'sientia', 'dbname': 'sientia', 'min_connections': 5, - 'max_connections': 20 + 'max_connections': 20, } diff --git a/tests/orchestrator/utils/test_converters.py b/tests/orchestrator/utils/test_converters.py index 430316e..d307860 100644 --- a/tests/orchestrator/utils/test_converters.py +++ b/tests/orchestrator/utils/test_converters.py @@ -2,14 +2,14 @@ from orchestrator.utils.converters import parse_frequency def test_parse_frequency(): - assert parse_frequency("1s") == 1 - assert parse_frequency("1m") == 60 - assert parse_frequency("1h") == 60 * 60 - assert parse_frequency("1d") == 60 * 60 * 24 + assert parse_frequency('1s') == 1 + assert parse_frequency('1m') == 60 + assert parse_frequency('1h') == 60 * 60 + assert parse_frequency('1d') == 60 * 60 * 24 try: - parse_frequency("1") + parse_frequency('1') except ValueError as e: - assert str(e) == "Invalid frequency" + assert str(e) == 'Invalid frequency' else: - assert False + raise AssertionError('Expected an exception to be raised') diff --git a/tests/orchestrator/utils/test_email_builder.py b/tests/orchestrator/utils/test_email_builder.py index d2145be..0d600d2 100644 --- a/tests/orchestrator/utils/test_email_builder.py +++ b/tests/orchestrator/utils/test_email_builder.py @@ -1,5 +1,5 @@ -import json from unittest.mock import MagicMock, patch + from pytest import fixture from orchestrator.utils.email_builder import EmailBuilder @@ -7,7 +7,7 @@ from orchestrator.utils.email_builder import EmailBuilder @fixture @patch('orchestrator.utils.email_builder.open') -def report_builder(open): +def report_builder(open_mock): return EmailBuilder(MagicMock()) @@ -26,22 +26,30 @@ def test_parameters(report_builder): general_events = { 'ERROR': { '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': { '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': { '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') @@ -55,24 +63,38 @@ def test_parameters(report_builder): report_builder.replace_parameters.assert_any_call( report_builder.general_template, - {'models': [ - {'model_name': 'model_name', 'events': [ - {'notification_id': 'ID_1', 'level': 'ERROR', 'project': 'project'}]}, - ]} + { + 'models': [ + { + 'model_name': 'model_name', + 'events': [{'notification_id': 'ID_1', 'level': 'ERROR', 'project': 'project'}], + }, + ] + }, ) report_builder.replace_parameters.assert_any_call( report_builder.general_template, - {'models': [ - {'model_name': 'model_name', 'events': [ - {'notification_id': 'ID_2', 'level': 'WARNING', 'project': 'project'}]}, - ]} + { + 'models': [ + { + 'model_name': 'model_name', + 'events': [ + {'notification_id': 'ID_2', 'level': 'WARNING', 'project': 'project'} + ], + }, + ] + }, ) report_builder.replace_parameters.assert_any_call( report_builder.general_template, - {'models': [ - {'model_name': 'model_name', 'events': [ - {'notification_id': 'ID_3', 'level': 'INFO', 'project': 'project'}]}, - ]} + { + 'models': [ + { + '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_data = [ - {'notification_id': 'ID_1', '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'} + { + 'notification_id': 'ID_1', + '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') report_builder.replace_parameters.assert_called_once_with( - report_builder.report_template, - report_builder.parameters.return_value + report_builder.report_template, report_builder.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', 'events': [ - {'notification_id': 'ID_1', 'level': 'ERROR', - 'project': 'project', 'model_name': 'model_name'}, - {'notification_id': 'ID_3', 'level': 'ERROR', - 'project': 'project', 'model_name': 'model_name'} - ] + { + 'notification_id': 'ID_1', + 'level': 'ERROR', + 'project': 'project', + 'model_name': 'model_name', + }, + { + 'notification_id': 'ID_3', + 'level': 'ERROR', + 'project': 'project', + 'model_name': 'model_name', + }, + ], } - ] + ], }, 'WARNING': { 'section_name': 'Warnings detected:', @@ -122,11 +167,15 @@ def test_build_email(report_builder): { 'model_name': 'model_name', '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': { 'section_name': 'Infos detected:', @@ -134,12 +183,16 @@ def test_build_email(report_builder): { 'model_name': 'model_name', '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', ) diff --git a/tests/orchestrator/utils/test_orchestrator_functions.py b/tests/orchestrator/utils/test_orchestrator_functions.py index 28edc73..7e35124 100644 --- a/tests/orchestrator/utils/test_orchestrator_functions.py +++ b/tests/orchestrator/utils/test_orchestrator_functions.py @@ -1,304 +1,190 @@ -from unittest.mock import patch, call +from unittest.mock import call, patch + from orchestrator.utils.orchestrator_functions import ( + build_tag_config, common_config, - minimal_retrain, - scouter, - predictions_batch, - overlap_filter_config, - process_path_priority, gather_read_tags, - build_tag_config + minimal_retrain, + overlap_filter_config, + predictions_batch, + process_path_priority, + scouter, ) def test_common_config(): config = { - "workflow_type": "scouter", - "schedule_name": "test_schedule", - "model_id": "test_model_id", - "model": { - "name": "test_model_name", - "model_config": { - "test_config": "test_config" - } - } + 'workflow_type': 'scouter', + 'schedule_name': 'test_schedule', + 'model_id': 'test_model_id', + 'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}}, } result = common_config(config) expected = { - "workflow_type": "scouter", - "schedule_name": "test_schedule", - "frequency": "1m", - "max_retry_policy": 1, - "model_id": "test_model_id", - "model_name": "test_model_name", - "model_config": { - "test_config": "test_config" - } + 'workflow_type': 'scouter', + 'schedule_name': 'test_schedule', + 'frequency': '1m', + 'max_retry_policy': 1, + 'model_id': 'test_model_id', + 'model_name': 'test_model_name', + 'model_config': {'test_config': 'test_config'}, + 'execution_timeout_seconds': 300, + 'task_timeout_seconds': 300, } assert result == expected def test_minimal_retrain(): config = { - "workflow_type": "minimal_retrain", - "schedule_name": "test_schedule", - "model_id": "test_model_id", - "model": { - "name": "test_model_name", - "model_config": { - "test_config": "test_config" - } - }, - "query": "select * from sientia_data.laborious_data order by \"timestamp\" desc limit 30;", - "datetime_columns": ["timestamp"] + 'workflow_type': 'minimal_retrain', + 'schedule_name': 'test_schedule', + 'model_id': 'test_model_id', + 'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}}, + 'query': 'select * from sientia_data.laborious_data order by "timestamp" desc limit 30;', + 'datetime_columns': ['timestamp'], } result = minimal_retrain(config) expected = { - "workflow_type": "minimal_retrain", - "schedule_name": "test_schedule", - "frequency": "1m", - "max_retry_policy": 1, - "model_id": "test_model_id", - "model_name": "test_model_name", - "model_config": { - "test_config": "test_config" - }, - "query": "select * from sientia_data.laborious_data order by \"timestamp\" desc limit 30;", - "schema": "sientia_data", - "table_name": "log_retrain", - "datetime_columns": ["timestamp"] + 'workflow_type': 'minimal_retrain', + 'schedule_name': 'test_schedule', + 'frequency': '1m', + 'max_retry_policy': 1, + 'model_id': 'test_model_id', + 'model_name': 'test_model_name', + 'model_config': {'test_config': 'test_config'}, + 'query': 'select * from sientia_data.laborious_data order by "timestamp" desc limit 30;', + 'schema': 'sientia_data', + 'table_name': 'log_retrain', + 'datetime_columns': ['timestamp'], + 'execution_timeout_seconds': 300, + 'task_timeout_seconds': 300, } assert result == expected def test_scouter(): config = { - "workflow_type": "scouter", - "schedule_name": "test_schedule", - "model_id": "test_model_id", - "model": { - "name": "test_model_name", - "model_config": { - "test_config": "test_config" - } - }, - "filters": [ - { - "filter_name": "test_filter_name", - "policy": "test_policy" - } + 'workflow_type': 'scouter', + 'schedule_name': 'test_schedule', + 'model_id': 'test_model_id', + 'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}}, + 'filters': [{'filter_name': 'test_filter_name', 'policy': 'test_policy'}], + 'read_tags': [ + {'tag_name': 'test_tag_name', 'aggr_func': 'test_aggr_func', 'data_range': [1, 2]} ], - "read_tags": [ - { - "tag_name": "test_tag_name", - "aggr_func": "test_aggr_func", - "data_range": [1, 2] - } - ], - "tag_retention_minutes": 10 + 'tag_retention_minutes': 10, + 'execution_timeout_seconds': 300, + 'task_timeout_seconds': 300, } result = scouter(config) expected = { - "workflow_type": "scouter", - "schedule_name": "test_schedule", - "frequency": "1m", - "max_retry_policy": 1, - "model_id": "test_model_id", - "model_name": "test_model_name", - "model_config": { - "test_config": "test_config" - }, - "topic": "raw_test_schedule", - "trigger_laborious": False, - "filters": { - "test_filter_name": { - "policy": "test_policy" - } - }, - "schema": "sientia_data", - "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 + 'workflow_type': 'scouter', + 'schedule_name': 'test_schedule', + 'frequency': '1m', + 'max_retry_policy': 1, + 'model_id': 'test_model_id', + 'model_name': 'test_model_name', + 'model_config': {'test_config': 'test_config'}, + 'topic': 'raw_test_schedule', + 'trigger_laborious': False, + 'filters': {'test_filter_name': {'policy': 'test_policy'}}, + 'schema': 'sientia_data', + '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, + 'execution_timeout_seconds': 300, + 'task_timeout_seconds': 300, } assert result == expected def test_overlap_filter_config(): config = [ - { - "filter_name": "test_filter_name", - "policy": "test_policy" - }, - { - "filter_name": "test_filter_name2", - "policy": "test_policy2" - } + {'filter_name': 'test_filter_name', 'policy': 'test_policy'}, + {'filter_name': 'test_filter_name2', 'policy': 'test_policy2'}, ] - result = overlap_filter_config({ - "test_filter_name": { - "policy": "test_policy" - } - }, config) + result = overlap_filter_config({'test_filter_name': {'policy': 'test_policy'}}, config) expected = { - "test_filter_name": { - "policy": "test_policy", - "config": {} - }, - "test_filter_name2": { - "policy": "test_policy2", - "config": {} - } + 'test_filter_name': {'policy': 'test_policy', 'config': {}}, + 'test_filter_name2': {'policy': 'test_policy2', 'config': {}}, } assert result == expected def test_process_path_priority(): - config = ["OTHER", "STOP", "CONTINUE"] + config = ['OTHER', 'STOP', 'CONTINUE'] result = process_path_priority(config) - expected = ["STOP", "CONTINUE", "REPEAT"] + expected = ['STOP', 'CONTINUE', 'REPEAT'] assert result == expected -@patch('orchestrator.utils.orchestrator_functions.overlap_filter_config', - return_value={ - "test_filter_name": { - "policy": "test_policy", - "config": {} - } - }) -@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): +@patch( + 'orchestrator.utils.orchestrator_functions.overlap_filter_config', + return_value={'test_filter_name': {'policy': 'test_policy', 'config': {}}}, +) +@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): config = { - "schedule_name": "test_schedule", - "workflow_type": "predictions_batch", - "model_id": "test_model_id", - "model": { - "name": "test_model_name", - "model_config": { - "test_config": "test_config" - } - }, - "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" - } + 'schedule_name': 'test_schedule', + 'workflow_type': 'predictions_batch', + 'model_id': 'test_model_id', + 'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}}, + '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": [ - { - "filter_name": "test_filter_name", - "policy": "test_policy" - } - ], - "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" + 'input_filters': [{'filter_name': 'test_filter_name', 'policy': 'test_policy'}], + '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) - mock_overlap_filter_config.assert_has_calls([ - call({ - "EMPTY_DATA": { - "policy": "STOP", - "config": {} - } - }, config['input_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_overlap_filter_config.assert_has_calls( + [call({'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, config['input_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']) expected = { - "workflow_type": "predictions_batch", - "schedule_name": "test_schedule", - "frequency": "1m", - "max_retry_policy": 1, - "model_id": "test_model_id", - "model_name": "test_model_name", - "model_config": { - "test_config": "test_config" - }, - "query": "test_query", - "schema": "sientia_data", - "table_name": "predictions", - "retention_time": 60 * 60, - "opc_output_config": { - "test_server_id": { - "prediction_tags": { - "test_addr": { - "data_type": "float" - } - }, - "confidence_tags": { - "test_addr": { - "data_type": "float" - } - } + 'workflow_type': 'predictions_batch', + 'schedule_name': 'test_schedule', + 'frequency': '1m', + 'max_retry_policy': 1, + 'model_id': 'test_model_id', + 'model_name': 'test_model_name', + 'model_config': {'test_config': 'test_config'}, + 'query': 'test_query', + 'schema': 'sientia_data', + 'table_name': 'predictions', + 'retention_time': 60 * 60, + 'opc_output_config': { + 'test_server_id': { + 'prediction_tags': {'test_addr': {'data_type': 'float'}}, + 'confidence_tags': {'test_addr': {'data_type': 'float'}}, } }, - "input_filters": { - "test_filter_name": { - "policy": "test_policy", - "config": {} - } - }, - "mlflow_transform_filters": { - "test_filter_name": { - "policy": "test_policy", - "config": {} - } - }, - "mlflow_predict_filters": { - "test_filter_name": { - "policy": "test_policy", - "config": {} - } - }, - "path_priority": ["STOP", "CONTINUE", "REPEAT"], - "datetime_columns": ["timestamp"], - "predictions_storage_policy": "erl:1" + 'input_filters': {'test_filter_name': {'policy': 'test_policy', 'config': {}}}, + 'mlflow_transform_filters': {'test_filter_name': {'policy': 'test_policy', 'config': {}}}, + 'mlflow_predict_filters': {'test_filter_name': {'policy': 'test_policy', 'config': {}}}, + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], + 'datetime_columns': ['timestamp'], + 'predictions_storage_policy': 'erl:1', + 'execution_timeout_seconds': 300, + 'task_timeout_seconds': 300, } assert result == expected @@ -306,93 +192,81 @@ def test_predictions_batch(mock_process_path_priority, def test_gather_read_tags(): pipelines = [ { - "schedule_name": "test_schedule", - "read_tags": [ + 'schedule_name': 'test_schedule', + 'read_tags': [ { - "server_id": "1", - "server_name": "test_server_name", - "tag_address": "test_tag_address" + 'server_id': '1', + 'server_name': 'test_server_name', + 'tag_address': 'test_tag_address', } - ] + ], }, { - "schedule_name": "test_schedule2", - "read_tags": [ + 'schedule_name': 'test_schedule2', + 'read_tags': [ { - "server_id": "2", - "server_name": "test_server_name2", - "tag_address": "test_tag_address2" + 'server_id': '2', + 'server_name': 'test_server_name2', + 'tag_address': 'test_tag_address2', }, { - "server_id": "2", - "server_name": "test_server_name2", - "tag_address": "test_tag_address3" - } - ] - } + 'server_id': '2', + 'server_name': 'test_server_name2', + 'tag_address': 'test_tag_address3', + }, + ], + }, ] result = gather_read_tags(pipelines) expected = { - "1:test_tag_address": { - "server_id": "1", - "server_name": "test_server_name", - "tag_address": "test_tag_address", - "topics": ["raw_test_schedule"] + '1:test_tag_address': { + 'server_id': '1', + 'server_name': 'test_server_name', + 'tag_address': 'test_tag_address', + 'topics': ['raw_test_schedule'], }, - "2:test_tag_address2": { - "server_id": "2", - "server_name": "test_server_name2", - "tag_address": "test_tag_address2", - "topics": ["raw_test_schedule2"] + '2:test_tag_address2': { + 'server_id': '2', + 'server_name': 'test_server_name2', + 'tag_address': 'test_tag_address2', + '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 def test_build_tag_config(): - tag = { - "server_id": "1", - "server_name": "test_server_name", - "tag_address": "test_tag_address" - } - opc_servers = { - "1": { - "server_name": "test_server_name", - "url": "test_url", - "uri": "test_uri" - } - } - slot_config = { - "1": {} - } + tag = {'server_id': '1', 'server_name': 'test_server_name', 'tag_address': 'test_tag_address'} + opc_servers = {'1': {'server_name': 'test_server_name', 'url': 'test_url', 'uri': 'test_uri'}} + slot_config = {'1': {}} i = 1 result = build_tag_config(tag, slot_config, opc_servers, i) expected = { - "1": { - "test_server_name": { - "server_id": "1", - "name": "test_server_name", - "url": "test_url", - "server_uri": "test_uri", - "cert_path": None, - "private_key_path": None, - "server_cert_path": None, - "tags": { - "test_tag_address": { - "server_id": "1", - "server_name": "test_server_name", - "tag_address": "test_tag_address" + '1': { + 'test_server_name': { + 'server_id': '1', + 'name': 'test_server_name', + 'url': 'test_url', + 'server_uri': 'test_uri', + 'cert_path': None, + 'private_key_path': None, + 'server_cert_path': None, + 'tags': { + 'test_tag_address': { + 'server_id': '1', + 'server_name': 'test_server_name', + 'tag_address': 'test_tag_address', } - } + }, } } } @@ -400,13 +274,10 @@ def test_build_tag_config(): def test_build_tag_config_no_server_id(): - tag = { - "server_id": "1", - "tag_address": "test_tag_address" - } + tag = {'server_id': '1', 'tag_address': 'test_tag_address'} opc_servers = {} try: build_tag_config(tag, {}, opc_servers, 1) except ValueError as e: - assert str(e) == "Server 1 not found in opc_servers" + assert str(e) == 'Server 1 not found in opc_servers' diff --git a/tests/orchestrator/workflows/subworkflows/test_load_notification_package.py b/tests/orchestrator/workflows/subworkflows/test_load_notification_package.py index dc77662..faeac7f 100644 --- a/tests/orchestrator/workflows/subworkflows/test_load_notification_package.py +++ b/tests/orchestrator/workflows/subworkflows/test_load_notification_package.py @@ -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 orchestrator.workflows.subworkflows.load_notification_package import LoadNotificationPackage + from orchestrator.activities.activities import Activities +from orchestrator.workflows.subworkflows.load_notification_package import LoadNotificationPackage @fixture @@ -20,14 +22,14 @@ metadata = { @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): input_data = { 'metadata': metadata, 'mail_type': 'test_mail_type', - 'base_data_filter': { - 'level': 'ERROR' - } + 'base_data_filter': {'level': 'ERROR'}, } workflow_mock.start_local_activity_method.side_effect = [ @@ -41,7 +43,7 @@ async def test_run(workflow_mock, load_notification_package): { 'id_r': '1', } - ] + ], ] output = await load_notification_package.run(input_data) @@ -57,123 +59,113 @@ async def test_run(workflow_mock, load_notification_package): { 'id_r': '1', } - ] + ], } - workflow_mock.start_local_activity_method.assert_has_calls([ - call( - Activities.get_last_data_timestamp, - { - **input_data['metadata'], - 'mail_type': 'test_mail_type' - }, - start_to_close_timeout=ANY, - retry_policy=ANY - ) - ]) + workflow_mock.start_local_activity_method.assert_has_calls( + [ + call( + Activities.get_last_data_timestamp, + {**input_data['metadata'], 'mail_type': 'test_mail_type'}, + start_to_close_timeout=ANY, + retry_policy=ANY, + ) + ] + ) - workflow_mock.start_local_activity_method.assert_has_calls([ - call( - Activities.find_documents_in_mongodb, - { - **input_data['metadata'], - 'query': { - 'collection': 'receiver_groups', - 'filters': { - 'active': True - } - } - }, - start_to_close_timeout=ANY, - retry_policy=ANY - ) - ]) + workflow_mock.start_local_activity_method.assert_has_calls( + [ + call( + Activities.find_documents_in_mongodb, + { + **input_data['metadata'], + 'query': {'collection': 'receiver_groups', 'filters': {'active': True}}, + }, + start_to_close_timeout=ANY, + retry_policy=ANY, + ) + ] + ) - workflow_mock.start_local_activity_method.assert_has_calls([ - call( - Activities.load_latest_data, - { - **input_data['metadata'], - 'collection_name': 'notification_queue', - 'last_data_timestamp': '2023-01-01 12:00:00', - 'base_data_filter': { - 'level': 'ERROR' - } - }, - start_to_close_timeout=ANY, - retry_policy=ANY - ) - ]) + workflow_mock.start_local_activity_method.assert_has_calls( + [ + call( + Activities.load_latest_data, + { + **input_data['metadata'], + 'collection_name': 'notification_queue', + 'last_data_timestamp': '2023-01-01 12:00:00', + 'base_data_filter': {'level': 'ERROR'}, + }, + start_to_close_timeout=ANY, + retry_policy=ANY, + ) + ] + ) - workflow_mock.start_activity_method.assert_has_calls([ - call( - Activities.put_last_data_timestamp, - { - **input_data['metadata'], - 'data': [ - { - 'id': '1', - } - ], - 'mail_type': 'test_mail_type' - }, - start_to_close_timeout=ANY, - retry_policy=ANY - ) - ]) + workflow_mock.start_activity_method.assert_has_calls( + [ + call( + Activities.put_last_data_timestamp, + { + **input_data['metadata'], + 'data': [ + { + 'id': '1', + } + ], + 'mail_type': 'test_mail_type', + }, + start_to_close_timeout=ANY, + retry_policy=ANY, + ) + ] + ) @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): input_data = { 'metadata': metadata, 'mail_type': 'test_mail_type', - 'base_data_filter': { - 'level': 'ERROR' - } + 'base_data_filter': {'level': 'ERROR'}, } - workflow_mock.start_local_activity_method.side_effect = [ - '2023-01-01 12:00:00', - [], - [] - ] + workflow_mock.start_local_activity_method.side_effect = ['2023-01-01 12:00:00', [], []] output = await load_notification_package.run(input_data) assert output == { 'last_timestamp': '2023-01-01 12:00:00', 'notification_package': [], - 'sending_configs': [] + 'sending_configs': [], } workflow_mock.start_activity_method.assert_not_called() @mark.asyncio -@patch("orchestrator.workflows.subworkflows.load_notification_package.workflow", new_callable=AsyncMock) -async def test_run_no_data(workflow_mock, load_notification_package): +@patch( + 'orchestrator.workflows.subworkflows.load_notification_package.workflow', new_callable=AsyncMock +) +async def test_run_has_data(workflow_mock, load_notification_package): input_data = { 'metadata': metadata, - 'base_data_filter': { - 'level': 'ERROR' - }, - 'mail_type': 'test_mail_type' + 'base_data_filter': {'level': 'ERROR'}, + 'mail_type': 'test_mail_type', } - workflow_mock.start_local_activity_method.side_effect = [ - '2023-01-01 12:00:00', - ["data"], - [] - ] + workflow_mock.start_local_activity_method.side_effect = ['2023-01-01 12:00:00', ['data'], []] output = await load_notification_package.run(input_data) assert output == { 'last_timestamp': '2023-01-01 12:00:00', - 'notification_package': ["data"], - 'sending_configs': [] + 'notification_package': ['data'], + 'sending_configs': [], } workflow_mock.start_activity_method.assert_not_called() diff --git a/tests/orchestrator/workflows/subworkflows/test_process_notifications.py b/tests/orchestrator/workflows/subworkflows/test_process_notifications.py index 3ffc3fe..0fb6734 100644 --- a/tests/orchestrator/workflows/subworkflows/test_process_notifications.py +++ b/tests/orchestrator/workflows/subworkflows/test_process_notifications.py @@ -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 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 orchestrator.activities.activities import Activities +from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications + @fixture def process_notifications(): @@ -21,11 +23,11 @@ metadata = { @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): input_data = { 'metadata': metadata, - 'notification_package': ["content"], + 'notification_package': ['content'], 'mail_type': 'test_mail_type', 'schema': 'test_schema', '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 - workflow_mock.execute_local_activity_method.assert_has_calls([ - call( - Activities.build_email_html, - { - **metadata, - 'receiver_groups': input_data['notification_package'], - 'mail_type': input_data['mail_type'] - }, - schedule_to_close_timeout=ANY, - retry_policy=ANY - ) - ]) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call( - Activities.format_log_report, - { - **metadata, - 'receiver_groups': workflow_mock.execute_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([ - 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_local_activity_method.assert_has_calls( + [ + call( + Activities.build_email_html, + { + **metadata, + 'receiver_groups': input_data['notification_package'], + 'mail_type': input_data['mail_type'], + }, + schedule_to_close_timeout=ANY, + retry_policy=ANY, + ) + ] + ) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.format_log_report, + { + **metadata, + 'receiver_groups': workflow_mock.execute_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([ - 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 - ) - ]) + 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( + [ + 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 -@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): input_data = { 'metadata': metadata, - 'notification_package': ["content"], + 'notification_package': ['content'], 'mail_type': 'test_mail_type', 'schema': 'test_schema', '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 = [] - response = await process_notifications.run(input_data) + await process_notifications.run(input_data) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call( - Activities.build_email_html, - { - **metadata, - 'receiver_groups': input_data['notification_package'], - 'mail_type': input_data['mail_type'] - }, - schedule_to_close_timeout=ANY, - retry_policy=ANY - ) - ]) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.build_email_html, + { + **metadata, + 'receiver_groups': input_data['notification_package'], + '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( + [ + 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, + ) + ] ) assert workflow_mock.execute_activity_method.call_count == 1 diff --git a/tests/orchestrator/workflows/test_alerts.py b/tests/orchestrator/workflows/test_alerts.py index 37fa1b3..766b80c 100644 --- a/tests/orchestrator/workflows/test_alerts.py +++ b/tests/orchestrator/workflows/test_alerts.py @@ -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 orchestrator.workflows.alerts import Alerts + from orchestrator.activities.activities import Activities +from orchestrator.workflows.alerts import Alerts @fixture @@ -20,113 +22,103 @@ metadata = { @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): - input_data = { - 'schedule_name': 'test-schedule-name', - 'notification_ttl': 300, - 'sent_ttl': 600 - } + input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600} await alerts.run(input_data) - workflow_mock.execute_child_workflow.assert_has_calls([ - call( - 'load_notification_package', - { - **input_data, - 'metadata': metadata, - 'base_data_filter': { - 'level': 'ERROR' - } - } - ) - ]) + workflow_mock.execute_child_workflow.assert_has_calls( + [ + call( + 'load_notification_package', + {**input_data, 'metadata': metadata, 'base_data_filter': {'level': 'ERROR'}}, + ) + ] + ) - workflow_mock.execute_child_workflow.assert_has_calls([ - call( - 'process_notifications', - { - 'metadata': metadata, - 'mail_type': 'Alerts', - 'notification_package': workflow_mock.execute_local_activity_method.return_value, - 'schema': 'sientia_data', - 'table_name': 'log_report' - } - ) - ]) + workflow_mock.execute_child_workflow.assert_has_calls( + [ + call( + 'process_notifications', + { + 'metadata': metadata, + 'mail_type': 'Alerts', + 'notification_package': workflow_mock.execute_local_activity_method.return_value, + 'schema': 'sientia_data', + 'table_name': 'log_report', + }, + ) + ] + ) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call( - Activities.filter_notification_alerts, - { - **metadata, - 'notification_package': workflow_mock.execute_child_workflow.return_value['notification_package'], - 'sending_configs': workflow_mock.execute_child_workflow.return_value['sending_configs'], - 'notification_ttl': input_data['notification_ttl'] - }, - schedule_to_close_timeout=ANY, - retry_policy=ANY - ) - ]) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.filter_notification_alerts, + { + **metadata, + 'notification_package': workflow_mock.execute_child_workflow.return_value[ + 'notification_package' + ], + 'sending_configs': workflow_mock.execute_child_workflow.return_value[ + 'sending_configs' + ], + 'notification_ttl': input_data['notification_ttl'], + }, + schedule_to_close_timeout=ANY, + retry_policy=ANY, + ) + ] + ) - workflow_mock.execute_activity_method.assert_has_calls([ - call( - Activities.store_notification_cache, - { - **metadata, - 'log_report': workflow_mock.execute_child_workflow.return_value, - 'sent_ttl': input_data['sent_ttl'] - }, - schedule_to_close_timeout=ANY, - retry_policy=ANY - ) - ]) + workflow_mock.execute_activity_method.assert_has_calls( + [ + call( + Activities.store_notification_cache, + { + **metadata, + 'log_report': workflow_mock.execute_child_workflow.return_value, + 'sent_ttl': input_data['sent_ttl'], + }, + schedule_to_close_timeout=ANY, + retry_policy=ANY, + ) + ] + ) @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): workflow_mock.execute_child_workflow.return_value = { 'last_timestamp': '2023-01-01 12:00:00.000000', 'notification_package': [], - 'sending_configs': [] + 'sending_configs': [], } - input_data = { - 'schedule_name': 'test-schedule-name', - 'notification_ttl': 300, - 'sent_ttl': 600 - } + input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600} await alerts.run(input_data) - workflow_mock.execute_child_workflow.assert_has_calls([ - call( - 'load_notification_package', - { - **input_data, - 'metadata': metadata, - 'base_data_filter': { - 'level': 'ERROR' - } - } - ) - ]) + workflow_mock.execute_child_workflow.assert_has_calls( + [ + call( + 'load_notification_package', + {**input_data, 'metadata': metadata, 'base_data_filter': {'level': 'ERROR'}}, + ) + ] + ) workflow_mock.execute_local_activity_method.assert_not_called() @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): workflow_mock.execute_local_activity_method.return_value = {} - input_data = { - 'schedule_name': 'test-schedule-name', - 'notification_ttl': 300, - 'sent_ttl': 600 - } + input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600} await alerts.run(input_data) @@ -134,18 +126,11 @@ async def test_run_no_receiver_groups(workflow_mock, alerts): @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): - workflow_mock.execute_child_workflow.side_effect = [ - MagicMock(), - [] - ] + workflow_mock.execute_child_workflow.side_effect = [MagicMock(), []] - input_data = { - 'schedule_name': 'test-schedule-name', - 'notification_ttl': 300, - 'sent_ttl': 600 - } + input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600} await alerts.run(input_data) diff --git a/tests/orchestrator/workflows/test_orchestrator.py b/tests/orchestrator/workflows/test_orchestrator.py index e89cf73..61243b0 100644 --- a/tests/orchestrator/workflows/test_orchestrator.py +++ b/tests/orchestrator/workflows/test_orchestrator.py @@ -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 orchestrator.workflows.orchestrator import Orchestrator + from orchestrator.activities.activities import Activities +from orchestrator.workflows.orchestrator import Orchestrator @fixture @@ -20,290 +22,327 @@ metadata = { @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): input_data = { - "pipelines_query": "SELECT * FROM bucket", - "opc_servers_query": "SELECT * FROM servers", - "schedule_name": "test-schedule-name", + 'pipelines_query': 'SELECT * FROM bucket', + 'opc_servers_query': 'SELECT * FROM servers', + 'schedule_name': 'test-schedule-name', } await orchestrator.run(input_data) - workflow_mock.start_local_activity_method.assert_has_calls([ - call( - Activities.aggregate_documents_in_mongodb, - { - **metadata, - "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" + workflow_mock.start_local_activity_method.assert_has_calls( + [ + call( + Activities.aggregate_documents_in_mongodb, + { + **metadata, + 'query': input_data['pipelines_query'], + 'timestamp_fields': ['updated_at'], }, - "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([ - call( - Activities.load_opc_slots, - { - **metadata - }, - 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.load_active_ingestors, - { - **metadata - }, - 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, + ) + ] + ) - workflow_mock.start_local_activity_method.assert_has_calls([ - call( - Activities.format_schedule_config, - { - **metadata, - 'schedule_config': workflow_mock.start_local_activity_method.return_value - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.start_local_activity_method.assert_has_calls( + [ + call( + Activities.load_opc_slots, + {**metadata}, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.start_local_activity_method.assert_has_calls([ - call( - Activities.process_schedules, - { - **metadata, - 'pipelines': workflow_mock.start_local_activity_method.return_value - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.start_local_activity_method.assert_has_calls( + [ + call( + Activities.load_active_ingestors, + {**metadata}, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.start_local_activity_method.assert_has_calls([ - call( - Activities.process_slots, - { - **metadata, - 'opc_servers': workflow_mock.start_local_activity_method.return_value, - 'active_ingestors': workflow_mock.start_local_activity_method.return_value, - 'pipelines': workflow_mock.start_local_activity_method.return_value, - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.start_local_activity_method.assert_has_calls( + [ + call( + Activities.format_schedule_config, + { + **metadata, + 'schedule_config': workflow_mock.start_local_activity_method.return_value, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.start_local_activity_method.assert_has_calls([ - call( - Activities.create_schedule_config, - { - **metadata, - 'current_schedule_config': workflow_mock.start_local_activity_method.return_value, - 'schedule_config': workflow_mock.start_local_activity_method.return_value - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.start_local_activity_method.assert_has_calls( + [ + call( + Activities.process_schedules, + {**metadata, 'pipelines': workflow_mock.start_local_activity_method.return_value}, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.start_local_activity_method.assert_has_calls([ - call( - Activities.create_slot_config, - { - **metadata, - 'current_slot_config': workflow_mock.start_local_activity_method.return_value, - 'slot_config': workflow_mock.start_local_activity_method.return_value - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.start_local_activity_method.assert_has_calls( + [ + call( + Activities.process_slots, + { + **metadata, + 'opc_servers': workflow_mock.start_local_activity_method.return_value, + 'active_ingestors': workflow_mock.start_local_activity_method.return_value, + 'pipelines': workflow_mock.start_local_activity_method.return_value, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.start_activity_method.assert_has_calls([ - call( - Activities.normalize_schedules, - { - **metadata, - 'orchestrated_schedules': workflow_mock.start_local_activity_method.return_value - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.start_local_activity_method.assert_has_calls( + [ + call( + Activities.create_schedule_config, + { + **metadata, + 'current_schedule_config': workflow_mock.start_local_activity_method.return_value, + 'schedule_config': workflow_mock.start_local_activity_method.return_value, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.start_activity_method.assert_has_calls([ - call( - Activities.create_collection_with_ttl_index, - { - **metadata, - 'pipelines': workflow_mock.start_local_activity_method.return_value['scouter'] - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.start_local_activity_method.assert_has_calls( + [ + call( + Activities.create_slot_config, + { + **metadata, + 'current_slot_config': workflow_mock.start_local_activity_method.return_value, + 'slot_config': workflow_mock.start_local_activity_method.return_value, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.start_activity_method.assert_has_calls([ - call( - Activities.delete_slots, - { - **metadata, - 'to_delete': - workflow_mock.start_local_activity_method.return_value['to_delete'] - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.start_activity_method.assert_has_calls( + [ + call( + Activities.normalize_schedules, + { + **metadata, + 'orchestrated_schedules': workflow_mock.start_local_activity_method.return_value, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.start_activity_method.assert_has_calls([ - call( - Activities.update_slots, - { - **metadata, - 'to_insert': - workflow_mock.start_local_activity_method.return_value['to_insert'] - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.start_activity_method.assert_has_calls( + [ + call( + Activities.create_collection_with_ttl_index, + { + **metadata, + 'pipelines': workflow_mock.start_local_activity_method.return_value['scouter'], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.start_activity_method.assert_has_calls([ - call( - Activities.delete_schedules, - { - **metadata, - 'schedules': - workflow_mock.start_local_activity_method.return_value['to_delete'] - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.start_activity_method.assert_has_calls( + [ + call( + Activities.delete_slots, + { + **metadata, + 'to_delete': workflow_mock.start_local_activity_method.return_value[ + 'to_delete' + ], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.start_activity_method.assert_has_calls([ - call( - Activities.create_schedules, - { - **metadata, - 'schedules': - workflow_mock.start_local_activity_method.return_value['to_create'] - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.start_activity_method.assert_has_calls( + [ + call( + Activities.update_slots, + { + **metadata, + 'to_insert': workflow_mock.start_local_activity_method.return_value[ + 'to_insert' + ], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.start_activity_method.assert_has_calls([ - call( - Activities.update_schedules, - { - **metadata, - 'schedules': - workflow_mock.start_local_activity_method.return_value['to_update'] - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + workflow_mock.start_activity_method.assert_has_calls( + [ + call( + Activities.delete_schedules, + { + **metadata, + 'schedules': workflow_mock.start_local_activity_method.return_value[ + 'to_delete' + ], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.start_activity_method.assert_has_calls([ - call( - Activities.report_schedule_orchestration, - { - **metadata, - 'created_schedules': workflow_mock.start_activity_method.return_value, - 'updated_schedules': workflow_mock.start_activity_method.return_value, - '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( + [ + call( + Activities.create_schedules, + { + **metadata, + 'schedules': workflow_mock.start_local_activity_method.return_value[ + 'to_create' + ], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.start_activity_method.assert_has_calls([ - call( - Activities.report_slot_orchestration, - { - **metadata, - 'inserted_slots': workflow_mock.start_activity_method.return_value, - 'deleted_slots': 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.update_schedules, + { + **metadata, + 'schedules': workflow_mock.start_local_activity_method.return_value[ + 'to_update' + ], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - workflow_mock.start_activity_method.assert_has_calls([ - call( - Activities.update_pipelines_timestamps, - { - **metadata, - 'updated_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.report_schedule_orchestration, + { + **metadata, + 'created_schedules': workflow_mock.start_activity_method.return_value, + 'updated_schedules': workflow_mock.start_activity_method.return_value, + '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([ - 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.report_slot_orchestration, + { + **metadata, + 'inserted_slots': workflow_mock.start_activity_method.return_value, + 'deleted_slots': 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 - ) - ]) + workflow_mock.start_activity_method.assert_has_calls( + [ + call( + Activities.update_pipelines_timestamps, + { + **metadata, + 'updated_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.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, + ) + ] + ) diff --git a/tests/orchestrator/workflows/test_reports.py b/tests/orchestrator/workflows/test_reports.py index d0134d8..732d44d 100644 --- a/tests/orchestrator/workflows/test_reports.py +++ b/tests/orchestrator/workflows/test_reports.py @@ -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 orchestrator.workflows.reports import Reports + from orchestrator.activities.activities import Activities +from orchestrator.workflows.reports import Reports @fixture @@ -20,95 +22,87 @@ metadata = { @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): - input_data = { - 'schedule_name': 'test-schedule-name', - 'notification_ttl': 300, - 'sent_ttl': 600 - } + input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600} await reports.run(input_data) - workflow_mock.execute_child_workflow.assert_has_calls([ - call( - 'load_notification_package', - { - **input_data, - 'metadata': metadata, - 'base_data_filter': {} - } - ) - ]) + workflow_mock.execute_child_workflow.assert_has_calls( + [ + call( + 'load_notification_package', + {**input_data, 'metadata': metadata, 'base_data_filter': {}}, + ) + ] + ) - workflow_mock.execute_child_workflow.assert_has_calls([ - call( - 'process_notifications', - { - 'metadata': metadata, - 'mail_type': 'Reports', - 'notification_package': workflow_mock.execute_local_activity_method.return_value, - 'schema': 'sientia_data', - 'table_name': 'log_report' - } - ) - ]) + workflow_mock.execute_child_workflow.assert_has_calls( + [ + call( + 'process_notifications', + { + 'metadata': metadata, + 'mail_type': 'Reports', + 'notification_package': workflow_mock.execute_local_activity_method.return_value, + 'schema': 'sientia_data', + 'table_name': 'log_report', + }, + ) + ] + ) - workflow_mock.execute_local_activity_method.assert_has_calls([ - call( - Activities.filter_notification_reports, - { - **metadata, - 'notification_package': workflow_mock.execute_child_workflow.return_value['notification_package'], - 'sending_configs': workflow_mock.execute_child_workflow.return_value['sending_configs'] - }, - schedule_to_close_timeout=ANY, - retry_policy=ANY - ) - ]) + workflow_mock.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.filter_notification_reports, + { + **metadata, + 'notification_package': workflow_mock.execute_child_workflow.return_value[ + 'notification_package' + ], + 'sending_configs': workflow_mock.execute_child_workflow.return_value[ + 'sending_configs' + ], + }, + schedule_to_close_timeout=ANY, + retry_policy=ANY, + ) + ] + ) @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): workflow_mock.execute_child_workflow.return_value = { 'last_timestamp': '2023-01-01 12:00:00.000000', 'notification_package': [], - 'sending_configs': [] + 'sending_configs': [], } - input_data = { - 'schedule_name': 'test-schedule-name', - 'notification_ttl': 300, - 'sent_ttl': 600 - } + input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600} await reports.run(input_data) - workflow_mock.execute_child_workflow.assert_has_calls([ - call( - 'load_notification_package', - { - **input_data, - 'metadata': metadata, - 'base_data_filter': {} - } - ) - ]) + workflow_mock.execute_child_workflow.assert_has_calls( + [ + call( + 'load_notification_package', + {**input_data, 'metadata': metadata, 'base_data_filter': {}}, + ) + ] + ) workflow_mock.execute_local_activity_method.assert_not_called() @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): workflow_mock.execute_local_activity_method.return_value = [] - input_data = { - 'schedule_name': 'test-schedule-name', - 'notification_ttl': 300, - 'sent_ttl': 600 - } + input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600} await reports.run(input_data) diff --git a/validate.sh b/validate.sh new file mode 100755 index 0000000..1cd4474 --- /dev/null +++ b/validate.sh @@ -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 \ No newline at end of file diff --git a/values.yaml b/values.yaml index 96d900d..2a8dc9c 100644 --- a/values.yaml +++ b/values.yaml @@ -235,12 +235,6 @@ env: value: "scouter" - name: TEMPORAL_LABORIOUS_NAMESPACE value: "laborious" - - name: TEMPORAL_TASK_TIMEOUT_MINUTES - value: "10" - - name: TEMPORAL_RUN_TIMEOUT_MINUTES - value: "20" - - name: TEMPORAL_EXECUTION_TIMEOUT_MINUTES - value: "20" ssh: enabled: true