From ab657d0fda40c39388b5e0d5863ce9ddd8b75827 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 28 Jul 2025 15:07:09 -0300 Subject: [PATCH] SIENTIAPDE-1172 feat: enhance email and slot manager activities with new filtering and caching functionalities - Updated Email class to improve attachment handling and logging. - Refactored Formatters class to process notifications more effectively by grouping them based on their configurations. - Introduced new methods in SlotManager for filtering notification alerts and storing notification cache. - Enhanced Alerts workflow to integrate new filtering and caching activities. - Added tests for new functionalities in SlotManager and Formatters classes to ensure reliability. --- orchestrator/activities/email.py | 16 +- orchestrator/activities/formatters.py | 44 +-- orchestrator/activities/slot_manager.py | 76 ++++- orchestrator/workflows/alerts.py | 70 ++++- .../subworkflows/load_notification_package.py | 11 +- tests/orchestrator/activities/test_email.py | 277 ++++++++++++++++++ .../activities/test_formatters.py | 69 +++++ .../activities/test_slot_manager.py | 97 ++++++ .../test_process_notifications.py | 79 +++++ tests/orchestrator/workflows/test_alerts.py | 80 +++++ 10 files changed, 786 insertions(+), 33 deletions(-) create mode 100644 tests/orchestrator/activities/test_email.py create mode 100644 tests/orchestrator/workflows/subworkflows/test_process_notifications.py create mode 100644 tests/orchestrator/workflows/test_alerts.py diff --git a/orchestrator/activities/email.py b/orchestrator/activities/email.py index 3355ebb..4ec30cf 100644 --- a/orchestrator/activities/email.py +++ b/orchestrator/activities/email.py @@ -1,9 +1,8 @@ -from email import encoders -from email.mime.base import MIMEBase -import traceback + from temporalio import workflow, activity with workflow.unsafe.imports_passed_through(): + import traceback import smtplib from typing import Any from sientia_do.temporal.activities.base import BaseActivity @@ -12,6 +11,8 @@ with workflow.unsafe.imports_passed_through(): from orchestrator.utils.email_builder import EmailBuilder from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText + from email.mime.base import MIMEBase + from email import encoders class Email(BaseActivity): @@ -24,7 +25,6 @@ class Email(BaseActivity): self.sender_email = sender_email self.sender_password = sender_password self.port = port - self.logger = logger if self.sender_password: self.server = smtplib.SMTP_SSL(smpt_server, port) @@ -68,7 +68,7 @@ class Email(BaseActivity): Attaches a list of attachments to an email message. Args: attachments (List[Dict]): A list of dictionaries where each dictionary contains - the keys 'attachment_id', 'trigger', and 'timestamp' representing the attachment details. + the keys 'filename' and 'content' representing the attachment details. msg (MIMEMultipart): The email message object to which the attachments will be added. Returns: MIMEMultipart: The email message object with the attachments added. @@ -124,7 +124,11 @@ class Email(BaseActivity): msg['Subject'] = f"SIENTIA™ {mail_type}" msg = self.handle_attachments( - [notification['attachment_content'] + [ + { + "filename": f"{notification['trigger']}_{notification['notification_id']}.txt", + "content": notification['attachment_content'] + } for notification in group_config['notifications'] if notification['attachment_content']], msg) diff --git a/orchestrator/activities/formatters.py b/orchestrator/activities/formatters.py index 9557b31..635c613 100644 --- a/orchestrator/activities/formatters.py +++ b/orchestrator/activities/formatters.py @@ -508,28 +508,30 @@ class Formatters(BaseActivity): for group_name, group_config in receiver_groups.items(): - notification_id = group_config['notification_id'] - trigger = group_config['trigger'] + for notification in group_config['notifications']: - key = f"{notification_id}:{trigger}" + notification_id = notification['notification_id'] + trigger = notification['trigger'] - if key not in data: - data[key] = { - 'status': group_config['status'], - 'timestamp': group_config['timestamp'], - 'groups': [group_name], - 'message': group_config['message'], - 'level': group_config['level'], - 'notification_id': notification_id, - 'block': group_config['block'], - 'schedule': trigger, - 'pipeline': group_config['pipeline'], - 'project': group_config['project'], - 'model_name': group_config['model_name'], - 'model_id': group_config['model_id'], - 'mail_type': mail_type - } - else: - data[key]['groups'].append(group_name) + key = f"{notification_id}:{trigger}" + + if key not in data: + data[key] = { + 'status': group_config['status'], + 'timestamp': notification['timestamp'], + 'groups': [group_name], + 'message': notification['message'], + 'level': notification['level'], + 'notification_id': notification_id, + 'block': notification['block'], + 'schedule': trigger, + 'pipeline': notification['pipeline'], + 'project': notification['project'], + 'model_name': notification['model_name'], + 'model_id': notification['model_id'], + 'mail_type': mail_type + } + else: + data[key]['groups'].append(group_name) return DataFrame(list(data.values())).to_dict() diff --git a/orchestrator/activities/slot_manager.py b/orchestrator/activities/slot_manager.py index b5cf4b7..eb3554f 100644 --- a/orchestrator/activities/slot_manager.py +++ b/orchestrator/activities/slot_manager.py @@ -1,4 +1,4 @@ -from pandas import DataFrame + from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): @@ -9,6 +9,9 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.temporal.activities.redis_base import Redis from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.models import NotificationLevel + from orchestrator.utils.patterns import DEFAULT_DATE_FORMAT + from pandas import DataFrame + from datetime import datetime, timedelta class SlotManager(Redis): @@ -263,3 +266,74 @@ class SlotManager(Redis): raise e return last_data_timestamp + + @activity.defn(name="filter_notification_alerts") + async def filter_notification_alerts(self, input_data: dict[str, Any]) -> dict[str, Any]: + """ + Filter notification alerts + """ + metadata = input_data['metadata'] + notification_package = input_data['notification_package'] + sending_configs = input_data['sending_configs'] + notification_ttl = input_data['notification_ttl'] + + 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]['notifications'] = [] + + ignore_list = receiver_group.get('ignore', []) + + for notification in notification_package: + alert_type = "do_nothing" + notification_id = notification['notification_id'] + # Check if notification was recently sent + key = f"{notification['trigger']}:{notification_id}" + + last_sent = self.get(key) + + if last_sent is None: + alert_type = "core_alerts" + + else: + last_sent = datetime.strptime( + last_sent, DEFAULT_DATE_FORMAT) + + # Check if "notification_ttl" seconds has passed since last sent + if (datetime.now() - last_sent) > timedelta(seconds=notification_ttl): + 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: + receiver_groups[group_name]["notifications"].append( + notification) + + return receiver_groups + + @activity.defn(name="store_notification_cache") + async def store_notification_cache(self, input_data: dict[str, Any]) -> dict[str, Any]: + """ + Store notification cache + """ + metadata = input_data['metadata'] + log_report = DataFrame(input_data['log_report']) + sent_ttl = input_data['sent_ttl'] + + self.info("Storing notification cache...", metadata=metadata) + + now = datetime.now().strftime(DEFAULT_DATE_FORMAT) + + for index, row in log_report.iterrows(): + status = row['status'] + if status == 'sent': + key = f"{row['schedule']}:{row['notification_id']}" + self.set(key, now, ttl=sent_ttl) + + return log_report diff --git a/orchestrator/workflows/alerts.py b/orchestrator/workflows/alerts.py index eee7f1c..fd46f57 100644 --- a/orchestrator/workflows/alerts.py +++ b/orchestrator/workflows/alerts.py @@ -9,14 +9,80 @@ with workflow.unsafe.imports_passed_through(): class Alerts: @workflow.run async def run(self, input_data: dict[str, Any]): + """ + Workflow to send alerts to the users + + Args: + input_data (dict[str, Any]): Input data. It contains the following keys: + - schedule_name: str - Name of the schedule + - notification_ttl: int - Period before consider some notification persistent + - sent_ttl: int - Time to live for the sent notification + + Returns: + None + + Raises: + Exception: If the workflow fails + """ + metadata = { + 'metadata': { + 'schedule_name': input_data['schedule_name'], + 'workflow_name': 'alerts', + 'model_name': '-', + 'model_id': '-' + } + } + + mail_type = "Alerts" + + input_data['metadata'] = metadata + + 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 + ) + + if not package['notification_package'] or not package['sending_configs']: + return + # Filter notification package by groups custom configs, levels and # timestamp cached + receiver_groups = await workflow.execute_local_activity_method( + Activities.filter_notification_alerts, + { + **metadata, + 'notification_package': package['notification_package'], + 'sending_configs': package['sending_configs'], + 'notification_ttl': input_data['notification_ttl'] + } + ) + # Call subworkflow "process_notifications" passing the notification package + log_report = await workflow.execute_child_workflow( + 'process_notifications', + { + 'metadata': metadata, + 'mail_type': mail_type, + 'notification_package': receiver_groups, + 'schema': 'sientia_data', + 'table_name': 'log_report' + } + ) # Store the notification_id sendings to avoid sending them again - - pass + await workflow.execute_activity_method( + Activities.store_notification_cache, + { + **metadata, + 'log_report': log_report, + 'sent_ttl': input_data['sent_ttl'] + } + ) diff --git a/orchestrator/workflows/subworkflows/load_notification_package.py b/orchestrator/workflows/subworkflows/load_notification_package.py index 165a778..8cbc76b 100644 --- a/orchestrator/workflows/subworkflows/load_notification_package.py +++ b/orchestrator/workflows/subworkflows/load_notification_package.py @@ -62,14 +62,19 @@ class LoadNotificationPackage: **metadata, 'collection_name': 'notification_queue', 'last_data_timestamp': last_timestamp, - 'base_data_filter': { - 'level': 'ERROR' - } + 'base_data_filter': input_data['base_data_filter'] }, start_to_close_timeout=timedelta(seconds=60), retry_policy=retry_policy ) + if not notification_package: + return { + 'last_timestamp': last_timestamp, + 'notification_package': [], + 'sending_configs': await sending_configs_handler + } + # Put last collected timestamp in redis "notification_last_timestamp" await workflow.start_activity_method( diff --git a/tests/orchestrator/activities/test_email.py b/tests/orchestrator/activities/test_email.py new file mode 100644 index 0000000..e29c68e --- /dev/null +++ b/tests/orchestrator/activities/test_email.py @@ -0,0 +1,277 @@ +from unittest.mock import MagicMock, call, patch +from pytest import mark, fixture + +from orchestrator.activities.email import Email + + +@fixture +@patch('orchestrator.activities.email.EmailBuilder') +@patch('orchestrator.activities.email.smtplib') +def email(smtplib, email_builder): + email = Email( + sender_email="test@test.com", + sender_password="test", + smpt_server="test", + port=587, + logger=MagicMock(), + notification_handler=MagicMock() + ) + email_builder.send_notification = MagicMock() + + return email + + +@patch('orchestrator.activities.email.EmailBuilder') +@patch('orchestrator.activities.email.smtplib') +def test___init___with_password(smtplib, email_builder): + email = Email( + sender_email="test@test.com", + sender_password="test", + smpt_server="test", + port=587, + logger=MagicMock(), + notification_handler=MagicMock() + ) + + assert email.sender_email == "test@test.com" + assert email.sender_password == "test" + assert email.port == 587 + + smtplib.SMTP_SSL.assert_called_once_with("test", 587) + smtplib.SMTP_SSL.return_value.login.assert_called_once_with( + "test@test.com", "test") + + assert email.server == smtplib.SMTP_SSL.return_value + + +@patch('orchestrator.activities.email.EmailBuilder') +@patch('orchestrator.activities.email.smtplib') +def test___init___without_password(smtplib, email_builder): + email = Email( + sender_email="test@test.com", + sender_password=None, + smpt_server="test", + port=587, + logger=MagicMock(), + notification_handler=MagicMock() + ) + + assert email.sender_email == "test@test.com" + assert email.sender_password is None + assert email.port == 587 + + smtplib.SMTP.assert_called_once_with("test", 587) + assert email.server == smtplib.SMTP.return_value + + +metadata = { + "metadata": { + "schedule_name": "test", + "model_name": "test", + "model_id": "test", + "workflow_name": "test", + } +} + + +@mark.asyncio +async def test_build_email_html(email): + 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" + } + ] + } + }, + "mail_type": "test" + } + + response = await email.build_email_html(input_data) + + assert response == { + "group_1": { + "notifications": [ + { + "type": "test", + "subject": "test", + "body": "test", + }, + { + "type": "test", + "subject": "test", + "body": "test", + } + ], + "html": "test" + } + } + + email.email_builder.build_email.assert_called_once_with( + input_data['receiver_groups']['group_1']['notifications'], + input_data['mail_type'] + ) + + +@patch('orchestrator.activities.email.MIMEBase') +@patch('orchestrator.activities.email.encoders') +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" + } + ] + + response = email.handle_attachments(attachments, message) + + assert response == message + + mime_base.assert_called_with('application', 'octet-stream') + assert mime_base.return_value.set_payload.call_count == 3 + + 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')) + ] + ) + + encoders.encode_base64.assert_called_with(mime_base.return_value) + assert encoders.encode_base64.call_count == 3 + + mime_base.return_value.add_header.assert_has_calls( + [ + call('Content-Disposition', 'attachment; filename="file_1"'), + call('Content-Disposition', 'attachment; filename="file_2"'), + call('Content-Disposition', 'attachment; filename="file_3"') + ] + ) + + message.attach.assert_called_with(mime_base.return_value) + assert message.attach.call_count == 3 + + +@patch('orchestrator.activities.email.MIMEBase') +def test_handle_attachments_failure(mime_base, email): + message = MagicMock() + + mime_base.side_effect = Exception("test") + + attachments = [ + { + "filename": "file_1", + "attachment_content": "test_content_1" + } + ] + + email.handle_attachments(attachments, message) + + assert message.attach.call_count == 0 + + +@mark.asyncio +@patch('orchestrator.activities.email.MIMEText') +@patch('orchestrator.activities.email.MIMEMultipart') +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 + ] + + email.email_builder.send_notification = MagicMock() + + email.server.sendmail = MagicMock( + side_effect=[ + None, + Exception("test") + ] + ) + + input_data = { + **metadata, + "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" + } + ], + "html": "test_html1" + }, + "group_2": { + "members": ["test3@test.com", "test4@test.com"], + "notifications": [], + "html": "test_html2" + } + }, + "mail_type": "test_TYPE" + } + + response = await email.send_email(input_data) + + assert response['group_1']['status'] == 'sent' + assert response['group_2']['status'] == 'failed' + + assert mimemultipart.call_count == 2 + + 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') + ] + ) + + side_effect_2.__setitem__.assert_has_calls( + [ + call('From', 'test@test.com'), + call('To', 'test3@test.com, test4@test.com'), + call('Subject', 'SIENTIA™ test_TYPE') + ] + ) + + email.server.sendmail.assert_has_calls( + [ + call('test@test.com', 'test@test.com, test2@test.com', + side_effect_1.as_string.return_value), + call('test@test.com', 'test3@test.com, test4@test.com', + side_effect_2.as_string.return_value) + ] + ) diff --git a/tests/orchestrator/activities/test_formatters.py b/tests/orchestrator/activities/test_formatters.py index 04ada97..8440e19 100644 --- a/tests/orchestrator/activities/test_formatters.py +++ b/tests/orchestrator/activities/test_formatters.py @@ -1,5 +1,6 @@ from unittest.mock import MagicMock, patch, call, ANY import json +from pandas import DataFrame from pytest import fixture, mark from sientia_do.notifications.models import NotificationLevel from orchestrator.activities.formatters import Formatters @@ -696,3 +697,71 @@ async def test_report_slot_orchestration(formatters): attachment=input_data['deleted_slots'] ) ]) + + +@mark.asyncio +async def test_format_log_report(formatters): + input_data = { + **metadata, + "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" + } + ], + "status": "sent", + }, + "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" + } + ], + "status": "sent", + } + }, + "mail_type": "test_mail_type" + } + + result = await formatters.format_log_report(input_data) + + 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" + } + ] + ) + + assert DataFrame(result).equals(expected_result) diff --git a/tests/orchestrator/activities/test_slot_manager.py b/tests/orchestrator/activities/test_slot_manager.py index ad2e9a5..d43112c 100644 --- a/tests/orchestrator/activities/test_slot_manager.py +++ b/tests/orchestrator/activities/test_slot_manager.py @@ -1,8 +1,10 @@ from unittest.mock import MagicMock, patch, call, ANY +from datetime import datetime, timedelta from pandas import DataFrame from pytest import mark, fixture from orchestrator.activities.slot_manager import SlotManager from sientia_do.notifications.models import NotificationLevel +from orchestrator.utils.patterns import DEFAULT_DATE_FORMAT metadata = { "metadata": { @@ -359,3 +361,98 @@ async def test_put_last_data_timestamp_error(slot_manager): else: assert False, "Expected exception" + + +@mark.asyncio +async def test_filter_notification_alerts(slot_manager): + slot_manager.get = MagicMock(side_effect=[ + None, + (datetime.now() - timedelta(seconds=600) + ).strftime(DEFAULT_DATE_FORMAT), + datetime.now().strftime(DEFAULT_DATE_FORMAT), + None, + (datetime.now() - timedelta(seconds=600) + ).strftime(DEFAULT_DATE_FORMAT), + datetime.now().strftime(DEFAULT_DATE_FORMAT)]) + + 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' + } + ], + 'sending_configs': [ + { + 'group_name': 'test_group_1', + 'contents': ['core_alerts', 'persistent_alerts'] + }, + { + 'group_name': 'test_group_2', + 'contents': ['core_alerts'] + } + ], + 'notification_ttl': 300 + } + + response = await slot_manager.filter_notification_alerts(input_data) + + assert response == { + 'test_group_1': { + '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' + } + ] + }, + 'test_group_2': { + 'group_name': 'test_group_2', + 'contents': ['core_alerts'], + 'notifications': [ + { + 'trigger': 'test_trigger_1', + 'notification_id': 'test_notification_id_1' + } + ] + } + } + + +@mark.asyncio +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 + } + + 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 + ) diff --git a/tests/orchestrator/workflows/subworkflows/test_process_notifications.py b/tests/orchestrator/workflows/subworkflows/test_process_notifications.py new file mode 100644 index 0000000..4dd6913 --- /dev/null +++ b/tests/orchestrator/workflows/subworkflows/test_process_notifications.py @@ -0,0 +1,79 @@ +from unittest.mock import AsyncMock, patch, ANY, call +from pytest import fixture, mark +from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications +from orchestrator.activities.activities import Activities + + +@fixture +def process_notifications(): + return ProcessNotifications() + + +metadata = { + 'metadata': { + 'schedule_name': 'test-schedule-name', + 'workflow_name': 'test-workflow', + 'model_name': '-', + 'model_id': '-', + } +} + + +@mark.asyncio +@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"], + 'mail_type': 'test_mail_type', + 'schema': 'test_schema', + 'table_name': 'test_table_name', + } + + response = await process_notifications.run(input_data) + + 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'] + }, + schedule_to_close_timeout=ANY, + retry_policy=ANY + ), + call( + Activities.format_log_report, + { + **metadata, + 'receiver_groups': workflow_mock.execute_activity_method.return_value, + 'mail_type': input_data['mail_type'] + }, + ) + ]) + + 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 + ), + 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 + }, + schedule_to_close_timeout=ANY, + retry_policy=ANY + ) + ]) diff --git a/tests/orchestrator/workflows/test_alerts.py b/tests/orchestrator/workflows/test_alerts.py new file mode 100644 index 0000000..788e06d --- /dev/null +++ b/tests/orchestrator/workflows/test_alerts.py @@ -0,0 +1,80 @@ +from unittest.mock import AsyncMock, patch, ANY, call +from pytest import fixture, mark +from orchestrator.workflows.alerts import Alerts +from orchestrator.activities.activities import Activities + + +@fixture +def alerts(): + return Alerts() + + +metadata = { + 'metadata': { + 'schedule_name': 'test-schedule-name', + 'workflow_name': 'alerts', + 'model_name': '-', + 'model_id': '-', + } +} + + +@mark.asyncio +@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 + } + + 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( + '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'] + } + ) + ]) + + 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'] + } + ) + ])