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

SIENTIAPDE-1148

Browse Source 
fix: update schedule handling in temporal_manager and formatters to support multiple namespaces, enhance test coverage, and improve configuration handling
This commit is contained in:
vitor-aignosi
2025-07-11 16:22:19 -03:00
parent 0b0bb4bd45
commit 89c769998f
6 changed files with 377 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
orchestrator/activities/formatters.py
View File

@@ -157,9 +157,10 @@ class Formatters(BaseActivity):
} }
for namespace, schedules in schedule_config.items(): for namespace, schedules in schedule_config.items():
current_schedules = current_schedule_config.get(namespace, {})
for schedule_name, schedule in schedules.items(): for schedule_name, schedule in schedules.items():
if schedule_name in current_schedule_config[namespace]: if schedule_name in current_schedules:
old_config = current_schedule_config[namespace][schedule_name]['data'] old_config = current_schedules[schedule_name]['data']
self.logger.debug(f"Comparing {schedule_name}:") self.logger.debug(f"Comparing {schedule_name}:")
self.logger.debug(json.dumps( self.logger.debug(json.dumps(
@@ -170,7 +171,7 @@ class Formatters(BaseActivity):
if schedule != old_config: if schedule != old_config:
to_update[namespace][schedule_name] = schedule to_update[namespace][schedule_name] = schedule
elif schedule_name not in current_schedule_config: elif schedule_name not in current_schedules:
to_create[namespace][schedule_name] = schedule to_create[namespace][schedule_name] = schedule
for namespace, schedules in current_schedule_config.items(): for namespace, schedules in current_schedule_config.items():

2
orchestrator/activities/temporal_manager.py
View File

@@ -41,7 +41,7 @@ class TemporalManager(BaseActivity):
async def connect_to_temporal(self): async def connect_to_temporal(self):
self.logger.info( self.logger.info(
f"Connecting to Temporal side namespaces at {self.host}") f"Connecting to Temporal side namespaces at {self.temporal_host}")
self.logger.info(f"Scouter namespace: {self.scouter_namespace}") self.logger.info(f"Scouter namespace: {self.scouter_namespace}")
scouter_client = await Client.connect( scouter_client = await Client.connect(

15
tests/orchestrator/activities/test_activities.py
View File

@@ -28,12 +28,17 @@ def test___init__(mock_formatters_init, mock_slot_manager_init,
'password': 'password' 'password': 'password'
} }
temporal_client = MagicMock() temporal_config = {
'temporal_host': 'localhost',
'temporal_scouter_namespace': 'scouter',
'temporal_laborious_namespace': 'laborious'
}
logger = MagicMock() logger = MagicMock()
notification_handler = MagicMock() notification_handler = MagicMock()
activities = Activities( activities = Activities(
temporal_client=temporal_client, temporal_config=temporal_config,
redis_config=redis_config, redis_config=redis_config,
mongodb_config=mongo_db_config, mongodb_config=mongo_db_config,
logger=logger, logger=logger,
@@ -66,13 +71,17 @@ def test___init__(mock_formatters_init, mock_slot_manager_init,
mock_temporal_manager_init.assert_called_once_with( mock_temporal_manager_init.assert_called_once_with(
ANY, ANY,
temporal_client=temporal_client, host='localhost',
scouter_namespace='scouter',
laborious_namespace='laborious',
logger=logger, logger=logger,
notification_handler=notification_handler notification_handler=notification_handler
) )
mock_formatters_init.assert_called_once_with( mock_formatters_init.assert_called_once_with(
ANY, ANY,
scouter_namespace='scouter',
laborious_namespace='laborious',
logger=logger, logger=logger,
notification_handler=notification_handler notification_handler=notification_handler
) )

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

@@ -9,6 +9,8 @@ from orchestrator.utils.orchestrator_functions import build_tag_config
@fixture @fixture
def formatters(): def formatters():
return Formatters( return Formatters(
scouter_namespace="scouter",
laborious_namespace="laborious",
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock() notification_handler=MagicMock()
) )
@@ -40,12 +42,18 @@ async def test_process_schedules(mock_predictions_batch, mock_scouter, formatter
result = await formatters.process_schedules(input_data) result = await formatters.process_schedules(input_data)
assert result == { assert result == {
"test_schedule_name": "test_scouter", "scouter": {
"test_schedule_name2": "test_predictions_batch" "test_schedule_name": "test_scouter"
},
"laborious": {
"test_schedule_name2": "test_predictions_batch"
}
} }
mock_scouter.assert_called_once_with(input_data['pipelines'][0]) mock_scouter.assert_called_once_with(
mock_predictions_batch.assert_called_once_with(input_data['pipelines'][1]) input_data['pipelines'][0])
mock_predictions_batch.assert_called_once_with(
input_data['pipelines'][1])
@mark.asyncio @mark.asyncio
@@ -243,23 +251,29 @@ async def test_process_slots(mock_build_tag_config, mock_gather_read_tags, forma
async def test_create_schedule_config(formatters): async def test_create_schedule_config(formatters):
input_data = { input_data = {
"current_schedule_config": { "current_schedule_config": {
"test_schedule_name_to_delete": { "scouter": {
"frequency": 60, "test_schedule_name_to_delete": {
"data": {"test": "test"} "frequency": 60,
}, "data": {"test": "test"}
"test_schedule_name_to_update": { },
"frequency": 60, "test_schedule_name_to_update": {
"data": {"test": "test"} "frequency": 60,
"data": {"test": "test"}
}
} }
}, },
"schedule_config": { "schedule_config": {
"test_schedule_name_to_create": { "laborious": {
"frequency": 60, "test_schedule_name_to_create": {
"data": {"test": "test"} "frequency": 60,
"data": {"test": "test"}
}
}, },
"test_schedule_name_to_update": { "scouter": {
"frequency": 60, "test_schedule_name_to_update": {
"data": {"test": "test2"} "frequency": 60,
"data": {"test": "test2"}
}
} }
} }
} }
@@ -268,20 +282,29 @@ async def test_create_schedule_config(formatters):
assert result == { assert result == {
"to_create": { "to_create": {
"test_schedule_name_to_create": { "laborious": {
"frequency": 60, "test_schedule_name_to_create": {
"data": {"test": "test"} "frequency": 60,
} "data": {"test": "test"}
}
},
"scouter": {}
}, },
"to_update": { "to_update": {
"test_schedule_name_to_update": { "scouter": {
"frequency": 60, "test_schedule_name_to_update": {
"data": {"test": "test2"} "frequency": 60,
} "data": {"test": "test2"}
}
},
"laborious": {}
}, },
"to_delete": [ "to_delete": {
"test_schedule_name_to_delete" "scouter": [
] "test_schedule_name_to_delete"
],
"laborious": []
}
} }

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

@@ -3,18 +3,43 @@ from datetime import timedelta
import base64 import base64
import json import json
from pytest import fixture, mark from pytest import fixture, mark
import pytest_asyncio
from orchestrator.activities.temporal_manager import TemporalManager from orchestrator.activities.temporal_manager import TemporalManager
from orchestrator.utils.converters import parse_frequency from orchestrator.utils.converters import parse_frequency
@fixture @fixture
def temporal_manager(): @patch("orchestrator.activities.temporal_manager.Client.connect")
return TemporalManager( def temporal_manager(connect_mock):
temporal_client=MagicMock(), temporal_manager = TemporalManager(
host='localhost:7233',
scouter_namespace='scouter',
laborious_namespace='laborious',
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock() notification_handler=MagicMock()
) )
temporal_manager.temporal_clients['scouter'] = MagicMock()
temporal_manager.temporal_clients['laborious'] = MagicMock()
return temporal_manager
@mark.asyncio
@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'
)
])
@mark.asyncio @mark.asyncio
@patch("orchestrator.activities.temporal_manager.MessageToDict", @patch("orchestrator.activities.temporal_manager.MessageToDict",
@@ -41,9 +66,7 @@ async def test_load_schedule(_mock_message_to_dict, temporal_manager):
} }
) )
temporal_manager.temporal_client.list_schedules = AsyncMock( handle = MagicMock(
return_value=async_iter())
temporal_manager.temporal_client.get_schedule_handle.return_value = MagicMock(
describe=AsyncMock( describe=AsyncMock(
return_value=MagicMock( return_value=MagicMock(
schedule=MagicMock( schedule=MagicMock(
@@ -59,24 +82,54 @@ async def test_load_schedule(_mock_message_to_dict, temporal_manager):
) )
) )
) )
temporal_manager.temporal_client.get_schedule_handle.return_value.describe \
.return_value.schedule.spec = MagicMock( temporal_manager.temporal_clients['scouter'].list_schedules = AsyncMock(
intervals=[ return_value=async_iter()
MagicMock( )
every=MagicMock( temporal_manager.temporal_clients['laborious'].list_schedules = AsyncMock(
seconds=60 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() response = await temporal_manager.load_schedule()
temporal_manager.temporal_client.list_schedules.assert_called_once() temporal_manager.temporal_clients['scouter'].list_schedules.assert_awaited_once(
)
temporal_manager.temporal_clients['laborious'].list_schedules.assert_awaited_once(
)
assert response == { assert response == {
"test-schedule-id": { "scouter": {
"frequency": 60, "test-schedule-id": {
"data": {"test": "test"} "frequency": 60,
"data": {"test": "test"}
}
},
"laborious": {
"test-schedule-id": {
"frequency": 60,
"data": {"test": "test"}
}
} }
} }
@@ -102,68 +155,112 @@ async def test_create_schedule(
input_data = { input_data = {
"schedules": { "schedules": {
"test-schedule": { "scouter": {
"model_id": 1, "test-schedule": {
"model_name": "test-model-name", "model_id": 1,
"workflow_type": "test-workflow", "model_name": "test-model-name",
"frequency": "1m", "workflow_type": "test-workflow",
"data": {"test": "test"} "frequency": "1m",
"data": {"test": "test"}
},
"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": { "laborious": {
"model_id": 2, "test-schedule-laborious": {
"model_name": "test-model-name", "model_id": 1,
"workflow_type": "test-workflow", "model_name": "test-model-name",
"frequency": "10y", "workflow_type": "test-workflow",
"data": {"test": "test"} "frequency": "2m",
"data": {"test": "test"}
}
} }
} }
} }
temporal_manager.temporal_client.create_schedule = AsyncMock() temporal_manager.temporal_clients['scouter'].create_schedule = AsyncMock()
temporal_manager.temporal_clients['laborious'].create_schedule = AsyncMock(
)
report = await temporal_manager.create_schedules(input_data) report = await temporal_manager.create_schedules(input_data)
temporal_manager.temporal_client.create_schedule.assert_called_once_with( temporal_manager.temporal_clients['scouter'].create_schedule.assert_called_once_with(
"test-schedule", "test-schedule",
mock_schedule.return_value, 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(
mock_schedule.assert_called_once_with( "test-schedule-laborious",
action=mock_schedule_action_start_workflow.return_value, mock_schedule.return_value,
spec=mock_schedule_spec.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_action_start_workflow.assert_has_calls([ mock_schedule_action_start_workflow.assert_has_calls([
call( call(
"test-workflow", "test-workflow",
input_data['schedules']['test-schedule'], input_data['schedules']['scouter']['test-schedule'],
id="test-schedule", id="test-schedule",
task_queue="test-workflow-queue", task_queue="test-workflow-queue",
execution_timeout=ANY execution_timeout=ANY
), ),
call( call(
"test-workflow", "test-workflow",
input_data['schedules']['test-schedule-invalid-frequency'], input_data['schedules']['scouter']['test-schedule-invalid-frequency'],
id="test-schedule-invalid-frequency", id="test-schedule-invalid-frequency",
task_queue="test-workflow-queue", task_queue="test-workflow-queue",
execution_timeout=ANY execution_timeout=ANY
),
call(
"test-workflow",
input_data['schedules']['laborious']['test-schedule-laborious'],
id="test-schedule-laborious",
task_queue="test-workflow-queue",
execution_timeout=ANY
) )
]) ])
mock_schedule_spec.assert_called_once_with( mock_schedule_spec.assert_has_calls([
intervals=[ call(
mock_schedule_interval_spec.return_value intervals=[
] mock_schedule_interval_spec.return_value
) ]
),
call(
intervals=[
mock_schedule_interval_spec.return_value
]
)
])
mock_schedule_interval_spec.assert_called_once_with( mock_schedule_interval_spec.assert_has_calls([
every=timedelta(seconds=60) call(
) every=timedelta(seconds=60)
),
call(
every=timedelta(seconds=120)
)
])
mock_parse_frequency.assert_has_calls([ mock_parse_frequency.assert_has_calls([
call("1m"), call("1m"),
call("10y") call("10y"),
call("2m")
]) ])
mock_typed_search_attributes.assert_has_calls([ mock_typed_search_attributes.assert_has_calls([
@@ -172,6 +269,11 @@ async def test_create_schedule(
mock_search_attribute_pair.return_value, 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([ call([
mock_search_attribute_pair.return_value, mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value, mock_search_attribute_pair.return_value,
@@ -188,6 +290,18 @@ async def test_create_schedule(
key=temporal_manager.model_name_id_key, key=temporal_manager.model_name_id_key,
value="test-model-name" 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( call(
key=temporal_manager.orchestrated_id_key, key=temporal_manager.orchestrated_id_key,
value="true" value="true"
@@ -202,10 +316,35 @@ async def test_create_schedule(
"test-schedule-invalid-frequency": { "test-schedule-invalid-frequency": {
"success": False, "success": False,
"message": "Invalid frequency" "message": "Invalid frequency"
},
"test-schedule-laborious": {
"success": True,
"message": "Schedule created successfully"
} }
} }
@mark.asyncio
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"}
}
}
}
}
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}"
@mark.asyncio @mark.asyncio
@patch("orchestrator.activities.temporal_manager.parse_frequency", @patch("orchestrator.activities.temporal_manager.parse_frequency",
side_effect=parse_frequency) side_effect=parse_frequency)
@@ -218,30 +357,50 @@ async def test_update_schedules(
args=MagicMock() args=MagicMock()
) )
temporal_manager.schedule_handles = { temporal_manager.schedule_handles = {
"test-schedule": MagicMock( "scouter": {
update=AsyncMock( "test-schedule": MagicMock(
update=AsyncMock( update=AsyncMock(
side_effect=lambda f: f(input_mock) 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 = { input_data = {
"schedules": { "schedules": {
"test-schedule": { "scouter": {
"frequency": "1m", "test-schedule": {
"data": {"test": "test"} "frequency": "1m",
"data": {"test": "test"}
},
"test-schedule_no_handler": {
"frequency": "1m",
"data": {"test": "test"}
}
}, },
"test-schedule_no_handler": { "laborious": {
"frequency": "1m", "test-schedule-laborious": {
"data": {"test": "test"} "frequency": "2m",
"data": {"test": "test"}
}
} }
} }
} }
report = await temporal_manager.update_schedules(input_data) report = await temporal_manager.update_schedules(input_data)
temporal_manager.schedule_handles['test-schedule'].update.assert_called_once() temporal_manager.schedule_handles['scouter']['test-schedule'].update.assert_called_once()
temporal_manager.schedule_handles['laborious']['test-schedule-laborious'].update.assert_called_once()
assert report == { assert report == {
"test-schedule": { "test-schedule": {
@@ -251,21 +410,59 @@ async def test_update_schedules(
"test-schedule_no_handler": { "test-schedule_no_handler": {
"success": False, "success": False,
"message": "Schedule test-schedule_no_handler not found" "message": "Schedule test-schedule_no_handler not found"
},
"test-schedule-laborious": {
"success": True,
"message": "Schedule updated successfully"
} }
} }
@mark.asyncio
async def test_update_schedules_with_no_handle(temporal_manager):
temporal_manager.temporal_clients = {}
input_data = {
"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"Schedule handles for abc not found, handles: {temporal_manager.schedule_handles}"
@mark.asyncio @mark.asyncio
async def test_delete_schedules(temporal_manager): async def test_delete_schedules(temporal_manager):
temporal_manager.schedule_handles = { temporal_manager.schedule_handles = {
"test-schedule": MagicMock( "scouter": {
delete=AsyncMock() "test-schedule": MagicMock(
) delete=AsyncMock()
)
},
"laborious": {
"test-schedule-laborious": MagicMock(
delete=AsyncMock()
)
}
} }
input_data = { input_data = {
"schedules": [ "schedules": {
"test-schedule", "test-schedule_no_handler" "scouter": [
] "test-schedule", "test-schedule_no_handler"
],
"laborious": [
"test-schedule-laborious"
]
}
} }
report = await temporal_manager.delete_schedules(input_data) report = await temporal_manager.delete_schedules(input_data)
@@ -278,5 +475,27 @@ async def test_delete_schedules(temporal_manager):
"test-schedule_no_handler": { "test-schedule_no_handler": {
"success": False, "success": False,
"message": "Schedule test-schedule_no_handler not found" "message": "Schedule test-schedule_no_handler not found"
},
"test-schedule-laborious": {
"success": True,
"message": "Schedule deleted successfully"
} }
} }
@mark.asyncio
async def test_delete_schedules_with_no_handle(temporal_manager):
temporal_manager.schedule_handles = {}
input_data = {
"schedules": {
"abc": [
"test-schedule"
]
}
}
try:
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}"

26
tests/orchestrator/utils/test_connectors_config.py
View File

@@ -1,7 +1,7 @@
from os import environ from os import environ
from orchestrator.utils.connectors_config import (build_redis_config, from orchestrator.utils.connectors_config import (build_redis_config,
build_couchbase_config, build_couchbase_config,
build_mongodb_config) build_mongodb_config, build_temporal_config)
def test_build_redis_config_with_env_vars(): def test_build_redis_config_with_env_vars():
@@ -74,3 +74,27 @@ def test_build_mongo_db_config_with_defaults():
'connection_string': 'mongodb://sientia:sientia@localhost:27017', 'connection_string': 'mongodb://sientia:sientia@localhost:27017',
'database_name': 'sientia' 'database_name': 'sientia'
} }
def test_build_temporal_config_with_env_vars():
environ['TEMPORAL_HOST'] = 'localhost:7233'
environ['TEMPORAL_SCOUTER_NAMESPACE'] = 'scouter'
environ['TEMPORAL_LABORIOUS_NAMESPACE'] = 'laborious'
assert build_temporal_config() == {
'temporal_host': 'localhost:7233',
'temporal_namespace': 'default',
'temporal_scouter_namespace': 'scouter',
'temporal_laborious_namespace': 'laborious'
}
def test_build_temporal_config_with_defaults():
environ.pop('TEMPORAL_HOST', None)
environ.pop('TEMPORAL_SCOUTER_NAMESPACE', None)
environ.pop('TEMPORAL_LABORIOUS_NAMESPACE', None)
assert build_temporal_config() == {
'temporal_host': 'localhost:7233',
'temporal_namespace': 'default',
'temporal_scouter_namespace': 'scouter',
'temporal_laborious_namespace': 'laborious'
}