From 43206b578d83ee1b872ccb3cbb5df58bd842ca38 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 14 Jul 2025 12:59:52 -0300 Subject: [PATCH] SIENTIAPDE-1150 fix: implement schedule normalization and timestamp management in MongoDB activities, enhance orchestration workflow with new formatting and normalization methods --- orchestrator/activities/formatters.py | 77 +++++-- orchestrator/activities/mongo_db.py | 56 +++++ orchestrator/activities/temporal_manager.py | 91 ++++++-- orchestrator/utils/patterns.py | 1 + orchestrator/workflows/orchestrator.py | 29 ++- test.ipynb | 109 +++++++++- .../activities/test_formatters.py | 137 ++++++++---- .../orchestrator/activities/test_mongo_db.py | 55 +++++ .../activities/test_temporal_manager.py | 199 ++++++++++++++---- .../workflows/test_orchestrator.py | 29 ++- 10 files changed, 648 insertions(+), 135 deletions(-) create mode 100644 orchestrator/utils/patterns.py diff --git a/orchestrator/activities/formatters.py b/orchestrator/activities/formatters.py index 376ec44..834826f 100644 --- a/orchestrator/activities/formatters.py +++ b/orchestrator/activities/formatters.py @@ -1,15 +1,19 @@ -from sientia_do.notifications.models import NotificationLevel + from temporalio import activity, workflow + with workflow.unsafe.imports_passed_through(): import json from typing import Any from logging import Logger + from datetime import datetime from sientia_do.notifications.handlers import NotificationHandler from sientia_do.temporal.activities.base import BaseActivity + from sientia_do.notifications.models import NotificationLevel from orchestrator.utils.orchestrator_functions import ( scouter, predictions_batch, gather_read_tags, build_tag_config ) + from orchestrator.utils.patterns import DEFAULT_DATE_FORMAT from math import ceil @@ -119,29 +123,48 @@ class Formatters(BaseActivity): return slot_config - def compare_configs(self, - schedules: dict[str, Any], current_schedules: dict[str, Any], - to_update: dict[str, Any], to_create: dict[str, Any], - namespace: str): - """ - Compares the schedules and current schedules, and updates the - to_update and to_create dictionaries. + @activity.defn(name="format_schedule_config") + async def format_schedule_config(self, input_data: dict[str, Any]) -> dict[str, Any]: """ + Formats the schedule config to a dictionary with the schedule name as the key. + input_data: + - schedule_config (list[dict[str, Any]]): The schedule config to format. + Returns: + - dict[str, Any]: The formatted schedule config. + """ + schedule_config = input_data['schedule_config'] + + config = {} + for schedule in schedule_config: + namespace = schedule['namespace'] + schedule_name = schedule['schedule_name'] + updated_at = schedule['updated_at'] + + if namespace not in config: + config[namespace] = {} + + config[namespace][schedule_name] = updated_at + + return config + + def compare_config_timestamps(self, + schedules: dict[str, Any], current_schedules: dict[str, Any], + to_update: dict[str, Any], to_create: dict[str, Any], + namespace: str): + """ + Compares the timestamps of the schedule and the current schedule. + """ for schedule_name, schedule in schedules.items(): if schedule_name in current_schedules: - old_config = current_schedules[schedule_name]['data'] + update_timestamp = schedule.get( + 'updated_at', datetime.now().strftime(DEFAULT_DATE_FORMAT)) - self.logger.debug(f"Comparing {schedule_name}:") - self.logger.debug(json.dumps( - old_config, indent=4, sort_keys=True)) - self.logger.debug(json.dumps( - schedule, indent=4, sort_keys=True)) + old_timestamp = current_schedules[schedule_name] - if schedule != old_config: + if update_timestamp > old_timestamp: to_update[namespace][schedule_name] = schedule - - elif schedule_name not in current_schedules: + else: to_create[namespace][schedule_name] = schedule @activity.defn(name="create_schedule_config") @@ -183,7 +206,7 @@ class Formatters(BaseActivity): for namespace, schedules in schedule_config.items(): current_schedules = current_schedule_config.get(namespace, {}) - self.compare_configs( + self.compare_config_timestamps( schedules, current_schedules, to_update, to_create, namespace) for namespace, schedules in current_schedule_config.items(): @@ -264,6 +287,15 @@ class Formatters(BaseActivity): attachment_content=json.dumps(attachment, indent=4, sort_keys=True) ) + def parse_report_schedule(self, input_data: dict[str, Any]) -> tuple[list[str], list[str]]: + success_keys = [f"{value['namespace']}/{value['schedule_name']}" + for value in input_data if value['success']] + + error_keys = [f"{value['namespace']}/{value['schedule_name']}: {value['message']}" + for value in input_data if not value['success']] + + return success_keys, error_keys + def parse_report(self, input_data: dict[str, Any]) -> tuple[list[str], list[str]]: success_keys = [key for key, value in input_data.items() if value['success']] @@ -295,7 +327,8 @@ class Formatters(BaseActivity): # Send report for created schedules if len(created_schedules) > 0: - success_keys, error_keys = self.parse_report(created_schedules) + success_keys, error_keys = self.parse_report_schedule( + created_schedules) if len(success_keys) > 0: self.send_success_report( @@ -312,7 +345,8 @@ class Formatters(BaseActivity): # Send report for updated schedules if len(updated_schedules) > 0: - success_keys, error_keys = self.parse_report(updated_schedules) + success_keys, error_keys = self.parse_report_schedule( + updated_schedules) if len(success_keys) > 0: self.send_success_report( @@ -328,7 +362,8 @@ class Formatters(BaseActivity): ) if len(deleted_schedules) > 0: - success_keys, error_keys = self.parse_report(deleted_schedules) + success_keys, error_keys = self.parse_report_schedule( + deleted_schedules) if len(success_keys) > 0: self.send_success_report( diff --git a/orchestrator/activities/mongo_db.py b/orchestrator/activities/mongo_db.py index e5c663f..bd9c548 100644 --- a/orchestrator/activities/mongo_db.py +++ b/orchestrator/activities/mongo_db.py @@ -1,5 +1,8 @@ +import datetime from temporalio import workflow, activity +from orchestrator.utils.patterns import DEFAULT_DATE_FORMAT + with workflow.unsafe.imports_passed_through(): from typing import Any import traceback @@ -184,3 +187,56 @@ class MongoDB(BaseActivity): self.logger.error(trace) raise e + + @activity.defn(name="update_pipelines_timestamps") + async def update_pipelines_timestamps(self, input_data: dict[str, Any]) -> None: + """ + Update the timestamps of the pipelines in the MongoDB collection. + input_data: + - updated_pipelines (list): List of updated pipelines. + """ + updated_pipelines = input_data.get("updated_pipelines", []) + now = datetime.now().strftime(DEFAULT_DATE_FORMAT) + + self.database["pipelines"].update_many( + {"$or": [ + {"schedule_name": pipeline["schedule_name"], + "namespace": pipeline["namespace"]} + for pipeline in updated_pipelines + ]}, + {"$set": {"updated_at": now}} + ) + + @activity.defn(name="create_pipelines_timestamps") + async def create_pipelines_timestamps(self, input_data: dict[str, Any]) -> None: + """ + Create the timestamps of the pipelines in the MongoDB collection. + input_data: + - created_pipelines (list): List of created pipelines. + """ + created_pipelines = input_data.get("created_pipelines", []) + + now = datetime.now().strftime(DEFAULT_DATE_FORMAT) + + self.database["pipelines"].insert_many([{ + "schedule_name": pipeline["schedule_name"], + "namespace": pipeline["namespace"], + "updated_at": now + } for pipeline in created_pipelines]) + + @activity.defn(name="delete_pipelines_timestamps") + async def delete_pipelines_timestamps(self, input_data: dict[str, Any]) -> None: + """ + Delete the timestamps of the pipelines in the MongoDB collection. + input_data: + - deleted_pipelines (list): List of deleted pipelines. + """ + deleted_pipelines = input_data.get("deleted_pipelines", []) + + self.database["pipelines"].delete_many({ + "$or": [ + {"schedule_name": pipeline["schedule_name"], + "namespace": pipeline["namespace"]} + for pipeline in deleted_pipelines + ] + }) diff --git a/orchestrator/activities/temporal_manager.py b/orchestrator/activities/temporal_manager.py index c07b621..59b6274 100644 --- a/orchestrator/activities/temporal_manager.py +++ b/orchestrator/activities/temporal_manager.py @@ -129,6 +129,37 @@ class TemporalManager(BaseActivity): return orchestrated_schedules + @activity.defn(name="normalize_schedules") + async def normalize_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]: + """ + Normalize schedules. Removes schedules with no update time in mongo db collection "orchestrated_schedules". + input_data: + - orchestrated_schedules (dict[str, Any]): The orchestrated schedules to compare. + """ + self.logger.info("Getting orchestrated schedules...") + + orchestrated_schedules = input_data.get('orchestrated_schedules', {}) + + for namespace, client in self.temporal_clients.items(): + schedules = orchestrated_schedules.get(namespace, {}) + + self.logger.info( + "Getting orchestrated schedules for %s", 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 + + if schedule_id not in schedules: + self.logger.info( + "Schedule %s not found in mongo db, cleaning up", schedule_id) + + handle = client.get_schedule_handle( + schedule_id) + + await handle.delete() + @activity.defn(name="create_schedules") async def create_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]: """ @@ -145,7 +176,7 @@ class TemporalManager(BaseActivity): schedules_to_create = input_data['schedules'] - report = {} + report = [] for namespace, schedules in schedules_to_create.items(): client = self.temporal_clients.get(namespace) @@ -198,17 +229,21 @@ class TemporalManager(BaseActivity): search_attributes=search_attributes ) - report[schedule_name] = { + report.append({ + "namespace": namespace, + "schedule_name": schedule_name, "success": True, "message": "Schedule created successfully" - } + }) except Exception as e: self.logger.error("Failed to create schedule %s: %s", schedule_name, str(e)) - report[schedule_name] = { + report.append({ + "namespace": namespace, + "schedule_name": schedule_name, "success": False, "message": str(e) - } + }) self.logger.info(f"Processed {len(schedules_to_create)} schedules") @@ -233,18 +268,19 @@ class TemporalManager(BaseActivity): schedules_to_update = input_data['schedules'] - report = {} + report = [] for namespace, schedules in schedules_to_update.items(): - schedule_handles = self.schedule_handles.get(namespace, None) + client = self.temporal_clients.get(namespace) - if schedule_handles is None: + if not client: raise ValueError( - f"Schedule handles for {namespace} not found, handles: {self.schedule_handles}") + f"Temporal client for {namespace} not found, clients: {self.temporal_clients}") for schedule_name, schedule in schedules.items(): try: - handler = schedule_handles.get(schedule_name) + handler = client.get_schedule_handle( + schedule_name) if not handler: raise ValueError(f"Schedule {schedule_name} not found") @@ -276,17 +312,21 @@ class TemporalManager(BaseActivity): del update_schedule - report[schedule_name] = { + report.append({ + "namespace": namespace, + "schedule_name": schedule_name, "success": True, "message": "Schedule updated successfully" - } + }) except Exception as e: self.logger.error("Failed to update schedule %s: %s", schedule_name, str(e)) - report[schedule_name] = { + report.append({ + "namespace": namespace, + "schedule_name": schedule_name, "success": False, "message": str(e) - } + }) self.logger.info(f"Processed {len(schedules_to_update)} schedules") @@ -311,18 +351,19 @@ class TemporalManager(BaseActivity): schedules_to_delete = input_data['schedules'] - report = {} + report = [] for namespace, schedules in schedules_to_delete.items(): - schedule_handles = self.schedule_handles.get(namespace, None) + client = self.temporal_clients.get(namespace) - if schedule_handles is None: + if not client: raise ValueError( - f"Schedule handles for {namespace} not found, handles: {self.schedule_handles}") + f"Temporal client for {namespace} not found, clients: {self.temporal_clients}") for schedule_name in schedules: try: - handler = schedule_handles.get(schedule_name) + handler = client.get_schedule_handle( + schedule_name) if not handler: raise ValueError(f"Schedule {schedule_name} not found") @@ -331,17 +372,21 @@ class TemporalManager(BaseActivity): del self.schedule_handles[namespace][schedule_name] - report[schedule_name] = { + report.append({ + "namespace": namespace, + "schedule_name": schedule_name, "success": True, "message": "Schedule deleted successfully" - } + }) except Exception as e: self.logger.error("Failed to delete schedule %s: %s", schedule_name, str(e)) - report[schedule_name] = { + report.append({ + "namespace": namespace, + "schedule_name": schedule_name, "success": False, "message": str(e) - } + }) self.logger.info(f"Processed {len(schedules_to_delete)} schedules") diff --git a/orchestrator/utils/patterns.py b/orchestrator/utils/patterns.py new file mode 100644 index 0000000..fe00beb --- /dev/null +++ b/orchestrator/utils/patterns.py @@ -0,0 +1 @@ +DEFAULT_DATE_FORMAT = '%Y-%m-%d %H:%M:%S.%f' diff --git a/orchestrator/workflows/orchestrator.py b/orchestrator/workflows/orchestrator.py index 7bece5f..d195446 100644 --- a/orchestrator/workflows/orchestrator.py +++ b/orchestrator/workflows/orchestrator.py @@ -33,7 +33,12 @@ class Orchestrator: ) orchestrated_schedules_handler = workflow.execute_local_activity_method( - Activities.load_schedule, + Activities.find_documents_in_mongodb, + { + 'query': { + 'collection': 'orchestrated_schedules' + } + }, retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=600) ) @@ -56,6 +61,15 @@ class Orchestrator: opc_servers = await opc_servers_handler active_ingestors = await active_ingestors_handler + formatted_orchestrated_schedules_handler = workflow.execute_local_activity_method( + Activities.format_schedule_config, + { + 'schedule_config': orchestrated_schedules + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60) + ) + schedules_config_handler = workflow.execute_local_activity_method( Activities.process_schedules, { @@ -78,11 +92,12 @@ class Orchestrator: schedules_config = await schedules_config_handler slot_config = await slot_config_handler + formatted_orchestrated_schedules = await formatted_orchestrated_schedules_handler schedule_actions_handler = workflow.execute_local_activity_method( Activities.create_schedule_config, { - 'current_schedule_config': orchestrated_schedules, + 'current_schedule_config': formatted_orchestrated_schedules, 'schedule_config': schedules_config }, retry_policy=retry_policy, @@ -99,8 +114,18 @@ class Orchestrator: start_to_close_timeout=timedelta(seconds=60) ) + normalize_schedules_handler = workflow.execute_local_activity_method( + Activities.normalize_schedules, + { + 'orchestrated_schedules': formatted_orchestrated_schedules + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60) + ) + schedule_actions = await schedule_actions_handler slot_actions = await slot_actions_handler + await normalize_schedules_handler slot_deletion_report_handler = workflow.execute_activity_method( Activities.delete_slots, diff --git a/test.ipynb b/test.ipynb index 1fcad1d..1843727 100644 --- a/test.ipynb +++ b/test.ipynb @@ -110,7 +110,7 @@ }, { "cell_type": "code", - "execution_count": 14, + "execution_count": 3, "id": "7d01f160", "metadata": {}, "outputs": [], @@ -128,9 +128,6 @@ " namespace=os.getenv('TEMPORAL_NAMESPACE', 'default')\n", ")\n", "\n", - "manager = TemporalManager(temporal_client=temporal_client,\n", - " logger=logger,\n", - " notification_handler=MagicMock())\n", " " ] }, @@ -382,7 +379,7 @@ }, { "cell_type": "code", - "execution_count": 17, + "execution_count": 4, "id": "f81b3728", "metadata": {}, "outputs": [ @@ -390,12 +387,112 @@ "name": "stdout", "output_type": "stream", "text": [ - "{'_client': , 'id': 'orchestrator'}\n" + "{'id': 'orchestrator', 'schedule': ScheduleListSchedule(action=ScheduleListActionStartWorkflow(workflow='orchestrator'), spec=ScheduleSpec(calendars=[], intervals=[ScheduleIntervalSpec(every=datetime.timedelta(seconds=3600), offset=None)], cron_expressions=[], skip=[], start_at=None, end_at=None, jitter=None, time_zone_name=None), state=ScheduleListState(note=None, paused=False)), 'info': ScheduleListInfo(recent_actions=[ScheduleActionResult(scheduled_at=datetime.datetime(2025, 7, 14, 9, 0, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 7, 14, 9, 0, 0, 165083, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='orchestrator-2025-07-14T09:00:00Z', first_execution_run_id='01980829-9700-771a-9c9d-2450eecc9af4')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 7, 14, 10, 0, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 7, 14, 10, 0, 0, 125309, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='orchestrator-2025-07-14T10:00:00Z', first_execution_run_id='01980860-8578-759d-97fa-dbdfc784f88a')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 7, 14, 11, 0, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 7, 14, 11, 0, 0, 124356, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='orchestrator-2025-07-14T11:00:00Z', first_execution_run_id='01980897-73f7-7e83-a7a6-20eaac38e7d4')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 7, 14, 12, 0, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 7, 14, 12, 0, 0, 167329, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='orchestrator-2025-07-14T12:00:00Z', first_execution_run_id='019808ce-62a2-75d6-b11d-c5387cfa258f')), ScheduleActionResult(scheduled_at=datetime.datetime(2025, 7, 14, 13, 0, tzinfo=datetime.timezone.utc), started_at=datetime.datetime(2025, 7, 14, 13, 0, 0, 167575, tzinfo=datetime.timezone.utc), action=ScheduleActionExecutionStartWorkflow(workflow_id='orchestrator-2025-07-14T13:00:00Z', first_execution_run_id='01980905-5122-7807-90ae-3782d0332e79'))], next_action_times=[datetime.datetime(2025, 7, 14, 14, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 7, 14, 15, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 7, 14, 16, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 7, 14, 17, 0, tzinfo=datetime.timezone.utc), datetime.datetime(2025, 7, 14, 18, 0, tzinfo=datetime.timezone.utc)]), 'typed_search_attributes': TypedSearchAttributes(search_attributes=[]), 'search_attributes': {}, 'data_converter': DataConverter(payload_converter_class=, payload_codec=None, failure_converter_class=, payload_converter=, failure_converter=), 'raw_entry': schedule_id: \"orchestrator\"\n", + "info {\n", + " spec {\n", + " interval {\n", + " interval {\n", + " seconds: 3600\n", + " }\n", + " }\n", + " }\n", + " workflow_type {\n", + " name: \"orchestrator\"\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1752483600\n", + " }\n", + " actual_time {\n", + " seconds: 1752483600\n", + " nanos: 165083498\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"orchestrator-2025-07-14T09:00:00Z\"\n", + " run_id: \"01980829-9700-771a-9c9d-2450eecc9af4\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1752487200\n", + " }\n", + " actual_time {\n", + " seconds: 1752487200\n", + " nanos: 125309119\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"orchestrator-2025-07-14T10:00:00Z\"\n", + " run_id: \"01980860-8578-759d-97fa-dbdfc784f88a\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1752490800\n", + " }\n", + " actual_time {\n", + " seconds: 1752490800\n", + " nanos: 124356906\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"orchestrator-2025-07-14T11:00:00Z\"\n", + " run_id: \"01980897-73f7-7e83-a7a6-20eaac38e7d4\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1752494400\n", + " }\n", + " actual_time {\n", + " seconds: 1752494400\n", + " nanos: 167329351\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"orchestrator-2025-07-14T12:00:00Z\"\n", + " run_id: \"019808ce-62a2-75d6-b11d-c5387cfa258f\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_COMPLETED\n", + " }\n", + " recent_actions {\n", + " schedule_time {\n", + " seconds: 1752498000\n", + " }\n", + " actual_time {\n", + " seconds: 1752498000\n", + " nanos: 167575982\n", + " }\n", + " start_workflow_result {\n", + " workflow_id: \"orchestrator-2025-07-14T13:00:00Z\"\n", + " run_id: \"01980905-5122-7807-90ae-3782d0332e79\"\n", + " }\n", + " start_workflow_status: WORKFLOW_EXECUTION_STATUS_RUNNING\n", + " }\n", + " future_action_times {\n", + " seconds: 1752501600\n", + " }\n", + " future_action_times {\n", + " seconds: 1752505200\n", + " }\n", + " future_action_times {\n", + " seconds: 1752508800\n", + " }\n", + " future_action_times {\n", + " seconds: 1752512400\n", + " }\n", + " future_action_times {\n", + " seconds: 1752516000\n", + " }\n", + "}\n", + "}\n", + "{'_client': , 'id': 'orchestrator'}\n" ] } ], "source": [ "async for schedule in await temporal_client.list_schedules():\n", + " print(vars(schedule))\n", " id = schedule.id\n", "\n", " handle = temporal_client.get_schedule_handle(id)\n", diff --git a/tests/orchestrator/activities/test_formatters.py b/tests/orchestrator/activities/test_formatters.py index db63ba5..240a1e5 100644 --- a/tests/orchestrator/activities/test_formatters.py +++ b/tests/orchestrator/activities/test_formatters.py @@ -247,32 +247,51 @@ async def test_process_slots(mock_build_tag_config, mock_gather_read_tags, forma } +@mark.asyncio +async def test_format_schedule_config(formatters): + input_data = { + "schedule_config": [ + {"namespace": "test_namespace1", "schedule_name": "test1", + "updated_at": "2021-01-01"}, + {"namespace": "test_namespace2", "schedule_name": "test2", + "updated_at": "2021-01-02"} + ] + } + + result = await formatters.format_schedule_config(input_data) + + assert result == { + "test_namespace1": { + "test1": "2021-01-01" + }, + "test_namespace2": { + "test2": "2021-01-02" + } + } + + @mark.asyncio async def test_create_schedule_config(formatters): input_data = { "current_schedule_config": { "scouter": { - "test_schedule_name_to_delete": { - "frequency": 60, - "data": {"test": "test"} - }, - "test_schedule_name_to_update": { - "frequency": 60, - "data": {"test": "test"} - } + "test_schedule_name_to_delete": '2021-01-01', + "test_schedule_name_to_update": '2021-01-02' } }, "schedule_config": { "laborious": { "test_schedule_name_to_create": { "frequency": 60, - "data": {"test": "test"} + "data": {"test": "test"}, + "updated_at": "2021-01-03" } }, "scouter": { "test_schedule_name_to_update": { "frequency": 60, - "data": {"test": "test2"} + "data": {"test": "test2"}, + "updated_at": "2021-01-04" } } } @@ -285,7 +304,8 @@ async def test_create_schedule_config(formatters): "laborious": { "test_schedule_name_to_create": { "frequency": 60, - "data": {"test": "test"} + "data": {"test": "test"}, + "updated_at": "2021-01-03" } }, "scouter": {} @@ -294,7 +314,8 @@ async def test_create_schedule_config(formatters): "scouter": { "test_schedule_name_to_update": { "frequency": 60, - "data": {"test": "test2"} + "data": {"test": "test2"}, + "updated_at": "2021-01-04" } }, "laborious": {} @@ -369,95 +390,129 @@ def test_send_error_report(formatters): def test_parse_report(formatters): input_data = { - "test_schedule_name_to_create": { + "test_key": { "success": True }, - "test_schedule_name_to_create_error": { + "test_key2": { "success": False, - "error": "test_error" + "message": "test_error" } } result = formatters.parse_report(input_data) assert result == ( - ["test_schedule_name_to_create"], - ["test_schedule_name_to_create_error"] + ["test_key"], + ["test_key2"] + ) + + +def test_parse_report_schedule(formatters): + input_data = [ + { + "namespace": "test_namespace", + "schedule_name": "test_schedule_name_to_create", + "success": True + }, + { + "namespace": "test_namespace", + "schedule_name": "test_schedule_name_to_create_error", + "success": False, + "message": "test_error" + } + ] + result = formatters.parse_report_schedule(input_data) + + assert result == ( + ["test_namespace/test_schedule_name_to_create"], + ["test_namespace/test_schedule_name_to_create_error: test_error"] ) @mark.asyncio async def test_report_schedule_orchestration(formatters): - formatters.parse_report = MagicMock( - side_effect=formatters.parse_report + formatters.parse_report_schedule = MagicMock( + side_effect=formatters.parse_report_schedule ) formatters.send_success_report = MagicMock() formatters.send_error_report = MagicMock() input_data = { - "created_schedules": { - "test_schedule_name_to_create": { + "created_schedules": [ + { + "namespace": "test_namespace", + "schedule_name": "test_schedule_name_to_create", "success": True }, - "test_schedule_name_to_create_error": { + { + "namespace": "test_namespace", + "schedule_name": "test_schedule_name_to_create_error", "success": False, - "error": "test_error" + "message": "test_error" } - }, - "updated_schedules": { - "test_schedule_name_to_update": { + ], + "updated_schedules": [ + { + "namespace": "test_namespace", + "schedule_name": "test_schedule_name_to_update", "success": True }, - "test_schedule_name_to_update_error": { + { + "namespace": "test_namespace", + "schedule_name": "test_schedule_name_to_update_error", "success": False, - "error": "test_error" + "message": "test_error" } - }, - "deleted_schedules": { - "test_schedule_name_to_delete": { + ], + "deleted_schedules": [ + { + "namespace": "test_namespace", + "schedule_name": "test_schedule_name_to_delete", "success": True }, - "test_schedule_name_to_delete_error": { + { + "namespace": "test_namespace", + "schedule_name": "test_schedule_name_to_delete_error", "success": False, - "error": "test_error" + "message": "test_error" } - } + ] } await formatters.report_schedule_orchestration(input_data) - formatters.parse_report.assert_has_calls([ + formatters.parse_report_schedule.assert_has_calls([ call(input_data['created_schedules']), call(input_data['updated_schedules']), call(input_data['deleted_schedules']) ]) formatters.send_success_report.assert_has_calls([ call( - "Created schedules: \n test_schedule_name_to_create", + "Created schedules: \n test_namespace/test_schedule_name_to_create", "REPORT_ORCHESTRATION_CREATED_SCHEDULES" ), call( - "Updated schedules: \n test_schedule_name_to_update", + "Updated schedules: \n test_namespace/test_schedule_name_to_update", "REPORT_ORCHESTRATION_UPDATED_SCHEDULES" ), call( - "Deleted schedules: \n test_schedule_name_to_delete", + "Deleted schedules: \n test_namespace/test_schedule_name_to_delete", "REPORT_ORCHESTRATION_DELETED_SCHEDULES" ) ]) formatters.send_error_report.assert_has_calls([ call( - "Failed to create schedules: \n test_schedule_name_to_create_error", + "Failed to create schedules: \n test_namespace/test_schedule_name_to_create_error: test_error", "REPORT_ORCHESTRATION_CREATED_SCHEDULES", input_data['created_schedules'] ), call( - "Failed to update schedules: \n test_schedule_name_to_update_error", + "Failed to update schedules: \n test_namespace/test_schedule_name_to_update_error: test_error", "REPORT_ORCHESTRATION_UPDATED_SCHEDULES", input_data['updated_schedules'] ), call( - "Failed to delete schedules: \n test_schedule_name_to_delete_error", + "Failed to delete schedules: \n test_namespace/test_schedule_name_to_delete_error: test_error", "REPORT_ORCHESTRATION_DELETED_SCHEDULES", input_data['deleted_schedules'] ) diff --git a/tests/orchestrator/activities/test_mongo_db.py b/tests/orchestrator/activities/test_mongo_db.py index cb89784..3db578e 100644 --- a/tests/orchestrator/activities/test_mongo_db.py +++ b/tests/orchestrator/activities/test_mongo_db.py @@ -244,3 +244,58 @@ async def test_aggregate_documents_in_mongodb_missing_aggregation(mongo_db): else: assert False, "Expected a ValueError to be raised" + + +@mark.asyncio +@patch("orchestrator.activities.mongo_db.datetime") +async def test_update_pipelines_timestamps_success(datetime_mock, mongo_db): + input_data = {"updated_pipelines": [ + {"schedule_name": "test1", "namespace": "test1"}, + {"schedule_name": "test2", "namespace": "test2"} + ]} + 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": datetime_mock.now.return_value.strftime.return_value}} + ) + + +@mark.asyncio +@patch("orchestrator.activities.mongo_db.datetime") +async def test_create_pipelines_timestamps_success(datetime_mock, mongo_db): + input_data = {"created_pipelines": [ + {"schedule_name": "test1", "namespace": "test1"}, + {"schedule_name": "test2", "namespace": "test2"} + ]} + 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( + [ + {"schedule_name": "test1", "namespace": "test1", + "updated_at": datetime_mock.now.return_value.strftime.return_value}, + {"schedule_name": "test2", "namespace": "test2", + "updated_at": datetime_mock.now.return_value.strftime.return_value} + ] + ) + + +@mark.asyncio +@patch("orchestrator.activities.mongo_db.datetime") +async def test_delete_pipelines_timestamps_success(datetime_mock, mongo_db): + input_data = {"deleted_pipelines": [ + {"schedule_name": "test1", "namespace": "test1"}, + {"schedule_name": "test2", "namespace": "test2"} + ]} + 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"} + ]} + ) diff --git a/tests/orchestrator/activities/test_temporal_manager.py b/tests/orchestrator/activities/test_temporal_manager.py index 5410052..ebdd519 100644 --- a/tests/orchestrator/activities/test_temporal_manager.py +++ b/tests/orchestrator/activities/test_temporal_manager.py @@ -41,30 +41,32 @@ async def test_connect_to_temporal(connect_mock, temporal_manager): ]) +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"] + } + ) + + @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 - 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"] - } - ) handle = MagicMock( describe=AsyncMock( @@ -134,6 +136,43 @@ async def test_load_schedule(_mock_message_to_dict, temporal_manager): } +@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"} + } + } + + 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=MagicMock( + delete=AsyncMock() + ) + ) + temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock( + 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( + ) + + @mark.asyncio @patch("orchestrator.activities.temporal_manager.parse_frequency", side_effect=parse_frequency) @@ -308,20 +347,26 @@ async def test_create_schedule( ) ]) - assert report == { - "test-schedule": { + assert report == [ + { + "schedule_name": "test-schedule", + "namespace": "scouter", "success": True, "message": "Schedule created successfully" }, - "test-schedule-invalid-frequency": { + { + "schedule_name": "test-schedule-invalid-frequency", + "namespace": "scouter", "success": False, "message": "Invalid frequency" }, - "test-schedule-laborious": { + { + "schedule_name": "test-schedule-laborious", + "namespace": "laborious", "success": True, "message": "Schedule created successfully" } - } + ] @mark.asyncio @@ -397,29 +442,72 @@ async def test_update_schedules( } } + handler_scouter = MagicMock( + 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) + ) + ) + ) + + temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock( + side_effect=[ + handler_scouter, + None + ] + ) + + temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock( + side_effect=[ + handler_laborious + ] + ) + report = await temporal_manager.update_schedules(input_data) - temporal_manager.schedule_handles['scouter']['test-schedule'].update.assert_called_once() - temporal_manager.schedule_handles['laborious']['test-schedule-laborious'].update.assert_called_once() + 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") + ]) - assert report == { - "test-schedule": { + 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" }, - "test-schedule_no_handler": { + { + "schedule_name": "test-schedule_no_handler", + "namespace": "scouter", "success": False, "message": "Schedule test-schedule_no_handler not found" }, - "test-schedule-laborious": { + { + "schedule_name": "test-schedule-laborious", + "namespace": "laborious", "success": True, "message": "Schedule updated successfully" } - } + ] @mark.asyncio -async def test_update_schedules_with_no_handle(temporal_manager): +async def test_update_schedules_with_no_client(temporal_manager): temporal_manager.temporal_clients = {} input_data = { "schedules": { @@ -436,7 +524,7 @@ async def test_update_schedules_with_no_handle(temporal_manager): await temporal_manager.update_schedules(input_data) except Exception as e: assert str( - e) == f"Schedule handles for abc not found, handles: {temporal_manager.schedule_handles}" + e) == f"Temporal client for abc not found, clients: {temporal_manager.temporal_clients}" @mark.asyncio @@ -464,28 +552,57 @@ async def test_delete_schedules(temporal_manager): ] } } + handler_scouter = MagicMock( + delete=AsyncMock() + ) + + handler_laborious = MagicMock( + delete=AsyncMock() + ) + + temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock( + side_effect=[ + handler_scouter, + None + ] + ) + + temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock( + side_effect=[ + handler_laborious + ] + ) report = await temporal_manager.delete_schedules(input_data) - assert report == { - "test-schedule": { + handler_scouter.delete.assert_called_once() + handler_laborious.delete.assert_called_once() + + assert report == [ + { + "schedule_name": "test-schedule", + "namespace": "scouter", "success": True, "message": "Schedule deleted successfully" }, - "test-schedule_no_handler": { + { + "schedule_name": "test-schedule_no_handler", + "namespace": "scouter", "success": False, "message": "Schedule test-schedule_no_handler not found" }, - "test-schedule-laborious": { + { + "schedule_name": "test-schedule-laborious", + "namespace": "laborious", "success": True, "message": "Schedule deleted successfully" } - } + ] @mark.asyncio -async def test_delete_schedules_with_no_handle(temporal_manager): - temporal_manager.schedule_handles = {} +async def test_delete_schedules_with_no_client(temporal_manager): + temporal_manager.temporal_clients = {} input_data = { "schedules": { "abc": [ @@ -498,4 +615,4 @@ async def test_delete_schedules_with_no_handle(temporal_manager): await temporal_manager.delete_schedules(input_data) except Exception as e: assert str( - e) == f"Schedule handles for abc not found, handles: {temporal_manager.schedule_handles}" + e) == f"Temporal client for abc not found, clients: {temporal_manager.temporal_clients}" diff --git a/tests/orchestrator/workflows/test_orchestrator.py b/tests/orchestrator/workflows/test_orchestrator.py index c4fd4c3..45f4411 100644 --- a/tests/orchestrator/workflows/test_orchestrator.py +++ b/tests/orchestrator/workflows/test_orchestrator.py @@ -44,7 +44,12 @@ async def test_run(workflow_mock, orchestrator): workflow_mock.execute_local_activity_method.assert_has_calls([ call( - Activities.load_schedule, + Activities.find_documents_in_mongodb, + { + "query": { + "collection": "orchestrated_schedules" + } + }, retry_policy=ANY, start_to_close_timeout=ANY ) @@ -66,6 +71,17 @@ async def test_run(workflow_mock, orchestrator): ) ]) + workflow_mock.execute_local_activity_method.assert_has_calls([ + call( + Activities.format_schedule_config, + { + 'schedule_config': workflow_mock.execute_local_activity_method.return_value + }, + retry_policy=ANY, + start_to_close_timeout=ANY + ) + ]) + workflow_mock.execute_local_activity_method.assert_has_calls([ call( Activities.process_schedules, @@ -114,6 +130,17 @@ async def test_run(workflow_mock, orchestrator): ) ]) + workflow_mock.execute_local_activity_method.assert_has_calls([ + call( + Activities.normalize_schedules, + { + 'orchestrated_schedules': workflow_mock.execute_local_activity_method.return_value + }, + retry_policy=ANY, + start_to_close_timeout=ANY + ) + ]) + workflow_mock.execute_activity_method.assert_has_calls([ call( Activities.delete_slots,