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.
This commit is contained in:
vitor-aignosi
277
tests/orchestrator/activities/test_email.py
Normal file
277
tests/orchestrator/activities/test_email.py
Normal file
@@ -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)
|
||||
]
|
||||
)
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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
|
||||
)
|
||||
])
|
||||
80
tests/orchestrator/workflows/test_alerts.py
Normal file
80
tests/orchestrator/workflows/test_alerts.py
Normal file
@@ -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']
|
||||
}
|
||||
)
|
||||
])
|
||||
Reference in New Issue
Block a user