diff --git a/orchestrator/activities/mongo_db.py b/orchestrator/activities/mongo_db.py index 0450ad3..fadc2de 100644 --- a/orchestrator/activities/mongo_db.py +++ b/orchestrator/activities/mongo_db.py @@ -220,7 +220,7 @@ class MongoDB(BaseActivity): trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="MONGODB_UPDATE_ERROR", + notification_id="MONGODB_UPDATE_PIPELINES_ERROR", message=f"Failed to update pipelines timestamps: {e}", block="update_pipelines_timestamps", level=NotificationLevel.ERROR, @@ -256,7 +256,7 @@ class MongoDB(BaseActivity): trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="MONGODB_INSERT_ERROR", + notification_id="MONGODB_CREATE_PIPELINES_ERROR", message=f"Failed to create pipelines timestamps: {e}", block="create_pipelines_timestamps", level=NotificationLevel.ERROR, @@ -289,7 +289,7 @@ class MongoDB(BaseActivity): trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="MONGODB_DELETE_ERROR", + notification_id="MONGODB_DELETE_PIPELINES_ERROR", message=f"Failed to delete pipelines timestamps: {e}", block="delete_pipelines_timestamps", level=NotificationLevel.ERROR, diff --git a/orchestrator/activities/temporal_manager.py b/orchestrator/activities/temporal_manager.py index d41eba6..bdca7dd 100644 --- a/orchestrator/activities/temporal_manager.py +++ b/orchestrator/activities/temporal_manager.py @@ -63,90 +63,6 @@ class TemporalManager(BaseActivity): self.laborious_namespace: laborious_client } - @activity.defn(name="load_schedule") - async def load_schedule(self, input_data: dict[str, Any]) -> dict[str, Any]: - """ - Load all orchestrated schedules from Temporal. Filters by search attribute - "Orchestrated" set to "true" and returns a dictionary of schedule_id: - {frequency, data, handle} - - Returns: - dict[str, Any]: A dictionary of orchestrated schedules - """ - - metadata = input_data.get("metadata", {}) - - self.logger.info("Getting orchestrated schedules...") - - orchestrated_schedules = {} - - try: - - for namespace, client in self.temporal_clients.items(): - orchestrated_schedules[namespace] = {} - - self.logger.info( - f"Getting orchestrated schedules for {namespace}") - - async for schedule in await client.list_schedules(): - search_attrs = getattr(schedule, "search_attributes", {}) - if search_attrs.get("orchestrated", ["false"]) == ["true"]: - schedule_id = schedule.id - - self.logger.debug(f"Schedule id: {schedule_id}") - - handle = client.get_schedule_handle( - schedule_id) - - self.logger.debug("Handle acquired") - - self.schedule_handles[namespace][schedule_id] = handle - - self.logger.debug("Describing schedule...") - - desc = await handle.describe( - rpc_timeout=timedelta(seconds=60) - ) - - self.logger.debug("Parsing args...") - - for arg in desc.schedule.action.args: - data = MessageToDict(arg)['data'] - data = base64.b64decode(data).decode('utf-8') - - frequency = desc.schedule.spec.intervals[0].every.seconds - - orchestrated_schedules[namespace][schedule_id] = { - 'frequency': frequency, - 'data': json.loads(data), - } - - await sleep(0.1) - - except Exception as e: - trace = traceback.format_exc() - self.send_notification( - metadata=metadata, - notification_id="TEMPORAL_LOAD_SCHEDULE_ERROR", - message=f"Failed to load orchestrated schedules: {e}", - block="load_schedule", - level=NotificationLevel.ERROR, - attachment_content=trace - ) - self.logger.error(trace) - raise e - - self.logger.info( - f"Found {len(orchestrated_schedules[self.scouter_namespace]) + len(orchestrated_schedules[self.laborious_namespace])} orchestrated schedules") - - self.logger.debug( - f"Orchestrated schedules: {orchestrated_schedules}") - - self.logger.debug( - f"Schedule handles: {self.schedule_handles}") - - return orchestrated_schedules - @activity.defn(name="normalize_schedules") async def normalize_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]: """ diff --git a/tests/orchestrator/activities/test_mongo_db.py b/tests/orchestrator/activities/test_mongo_db.py index 9f82963..355070b 100644 --- a/tests/orchestrator/activities/test_mongo_db.py +++ b/tests/orchestrator/activities/test_mongo_db.py @@ -281,6 +281,35 @@ async def test_update_pipelines_timestamps_success(datetime_mock, mongo_db): ) +@mark.asyncio +@patch("orchestrator.activities.mongo_db.datetime") +async def test_update_pipelines_timestamps_failure(datetime_mock, mongo_db): + input_data = { + "updated_pipelines": [ + {"schedule_name": "test1", "namespace": "test1", "success": True}, + {"schedule_name": "test2", "namespace": "test2", "success": True} + ], + **metadata + } + + 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" + 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", + level=NotificationLevel.ERROR, + attachment_content=ANY + ) + + else: + assert False, "Expected an exception to be raised" + + @mark.asyncio @patch("orchestrator.activities.mongo_db.datetime") async def test_create_pipelines_timestamps_success(datetime_mock, mongo_db): @@ -300,6 +329,33 @@ async def test_create_pipelines_timestamps_success(datetime_mock, mongo_db): ) +@mark.asyncio +@patch("orchestrator.activities.mongo_db.datetime") +async def test_create_pipelines_timestamps_failure(datetime_mock, mongo_db): + 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") + try: + await mongo_db.create_pipelines_timestamps(input_data) + except Exception as e: + 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", + level=NotificationLevel.ERROR, + attachment_content=ANY + ) + + else: + assert False, "Expected an exception to be raised" + + @mark.asyncio @patch("orchestrator.activities.mongo_db.datetime") async def test_delete_pipelines_timestamps_success(datetime_mock, mongo_db): @@ -315,3 +371,30 @@ async def test_delete_pipelines_timestamps_success(datetime_mock, mongo_db): {"schedule_name": "test2", "namespace": "test2"} ]} ) + + +@mark.asyncio +@patch("orchestrator.activities.mongo_db.datetime") +async def test_delete_pipelines_timestamps_failure(datetime_mock, mongo_db): + 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") + try: + await mongo_db.delete_pipelines_timestamps(input_data) + except Exception as e: + 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", + level=NotificationLevel.ERROR, + attachment_content=ANY + ) + + else: + assert False, "Expected an exception to be raised" diff --git a/tests/orchestrator/activities/test_slot_manager.py b/tests/orchestrator/activities/test_slot_manager.py index d61cbb7..9be7686 100644 --- a/tests/orchestrator/activities/test_slot_manager.py +++ b/tests/orchestrator/activities/test_slot_manager.py @@ -1,6 +1,7 @@ -from unittest.mock import MagicMock, patch, call +from unittest.mock import MagicMock, patch, call, ANY from pytest import mark, fixture from orchestrator.activities.slot_manager import SlotManager +from sientia_do.notifications.models import NotificationLevel metadata = { "metadata": { @@ -28,6 +29,7 @@ def slot_manager(_redis_mock): slot_manager.redis_client = MagicMock() slot_manager.logger = MagicMock() slot_manager.notification_handler = MagicMock() + slot_manager.send_notification = MagicMock() return slot_manager @@ -82,6 +84,32 @@ async def test_load_opc_slots_no_decode(slot_manager): } +@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_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" + 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", + level=NotificationLevel.ERROR, + attachment_content=ANY + ) + + else: + assert False, "Expected an exception to be raised" + + @mark.asyncio async def test_load_active_ingestors(slot_manager): slot_manager.redis_client.keys.return_value = [ @@ -93,6 +121,30 @@ async def test_load_active_ingestors(slot_manager): "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"] + + 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" + 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", + level=NotificationLevel.ERROR, + attachment_content=ANY + ) + + else: + assert False, "Expected an exception to be raised" + + @mark.asyncio async def test_update_slots(slot_manager): slot_manager.set = MagicMock( diff --git a/tests/orchestrator/activities/test_temporal_manager.py b/tests/orchestrator/activities/test_temporal_manager.py index 5422d1a..058aa70 100644 --- a/tests/orchestrator/activities/test_temporal_manager.py +++ b/tests/orchestrator/activities/test_temporal_manager.py @@ -1,9 +1,7 @@ from unittest.mock import MagicMock, patch, AsyncMock, call, ANY from datetime import timedelta -import base64 -import json +from sientia_do.notifications.models import NotificationLevel from pytest import fixture, mark -import pytest_asyncio from orchestrator.activities.temporal_manager import TemporalManager from orchestrator.utils.converters import parse_frequency @@ -30,6 +28,7 @@ def temporal_manager(connect_mock): temporal_manager.temporal_clients['scouter'] = MagicMock() temporal_manager.temporal_clients['laborious'] = MagicMock() + temporal_manager.send_notification = MagicMock() return temporal_manager @@ -71,80 +70,6 @@ async def async_iter(): ) -@mark.asyncio -@patch("orchestrator.activities.temporal_manager.MessageToDict", - return_value={"data": base64.b64encode(json.dumps({"test": "test"}).encode('utf-8'))}) -async def test_load_schedule(_mock_message_to_dict, temporal_manager): - # Create async iterator mock - - handle = MagicMock( - describe=AsyncMock( - return_value=MagicMock( - schedule=MagicMock( - action=MagicMock( - args=[ - MagicMock( - data=base64.b64encode(json.dumps( - {"test": "test"}).encode('utf-8')) - ) - ] - ) - ) - ) - ) - ) - - temporal_manager.temporal_clients['scouter'].list_schedules = AsyncMock( - return_value=async_iter() - ) - temporal_manager.temporal_clients['laborious'].list_schedules = AsyncMock( - return_value=async_iter() - ) - - temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock( - return_value=handle - ) - temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock( - return_value=handle - ) - - describe_mock = MagicMock( - intervals=[ - MagicMock( - every=MagicMock( - seconds=60 - ) - ) - ] - ) - temporal_manager.temporal_clients['scouter'].get_schedule_handle.return_value.describe \ - .return_value.schedule.spec = describe_mock - temporal_manager.temporal_clients['laborious'].get_schedule_handle.return_value.describe \ - .return_value.schedule.spec = describe_mock - - response = await temporal_manager.load_schedule(metadata) - - temporal_manager.temporal_clients['scouter'].list_schedules.assert_awaited_once( - ) - temporal_manager.temporal_clients['laborious'].list_schedules.assert_awaited_once( - ) - - assert response == { - "scouter": { - "test-schedule-id": { - "frequency": 60, - "data": {"test": "test"} - } - }, - "laborious": { - "test-schedule-id": { - "frequency": 60, - "data": {"test": "test"} - } - } - } - - @mark.asyncio async def test_normalize_schedules(temporal_manager): input_data = { @@ -182,6 +107,29 @@ async def test_normalize_schedules(temporal_manager): ) +@mark.asyncio +async def test_normalize_schedules_error(temporal_manager): + temporal_manager.temporal_clients['scouter'].list_schedules = AsyncMock( + side_effect=Exception("Test exception") + ) + + try: + await temporal_manager.normalize_schedules(metadata) + except Exception as e: + 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", + level=NotificationLevel.ERROR, + attachment_content=ANY + ) + + else: + assert False, "Expected an exception to be raised" + + @mark.asyncio @patch("orchestrator.activities.temporal_manager.parse_frequency", side_effect=parse_frequency)