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.

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()

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,
     )

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'),
         ]
     )
 

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,
         )

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)

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}'
+        )

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,
     }

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')

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',
     )

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'

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()

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

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)
 

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,
+            )
+        ]
+    )

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)