gitea_admin/sientia-dataops-orchestrator_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1150

Browse Source 
fix: implement schedule normalization and timestamp management in MongoDB activities, enhance orchestration workflow with new formatting and normalization methods
This commit is contained in:
vitor-aignosi
2025-07-14 12:59:52 -03:00
parent e3a88eb1d5
commit 43206b578d
10 changed files with 648 additions and 135 deletions
Download Patch File Download Diff File Expand all files Collapse all files

77
orchestrator/activities/formatters.py
View File

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

56
orchestrator/activities/mongo_db.py
View File

@@ -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
]
})

91
orchestrator/activities/temporal_manager.py
View File

@@ -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")

1
orchestrator/utils/patterns.py Normal file
View File

@@ -0,0 +1 @@
DEFAULT_DATE_FORMAT = '%Y-%m-%d %H:%M:%S.%f'

29
orchestrator/workflows/orchestrator.py
View File

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

109
test.ipynb
View File

@@ -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': <temporalio.client.Client object at 0x770086544f90>, '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=<class 'temporalio.converter.DefaultPayloadConverter'>, payload_codec=None, failure_converter_class=<class 'temporalio.converter.DefaultFailureConverter'>, payload_converter=<temporalio.converter.DefaultPayloadConverter object at 0x787d51390750>, failure_converter=<temporalio.converter.DefaultFailureConverter object at 0x787d5277b590>), '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': <temporalio.client.Client object at 0x787d53dd3490>, '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",

137
tests/orchestrator/activities/test_formatters.py
View File

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

55
tests/orchestrator/activities/test_mongo_db.py
View File

@@ -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"}
]}
)

199
tests/orchestrator/activities/test_temporal_manager.py
View File

@@ -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}"

29
tests/orchestrator/workflows/test_orchestrator.py
View File

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