SIENTIAPDE-1148

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

View File

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

View File

@@ -41,7 +41,7 @@ class TemporalManager(BaseActivity):
async def connect_to_temporal(self):
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}")
scouter_client = await Client.connect(

View File

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

View File

@@ -9,6 +9,8 @@ from orchestrator.utils.orchestrator_functions import build_tag_config
@fixture
def formatters():
return Formatters(
scouter_namespace="scouter",
laborious_namespace="laborious",
logger=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)
assert result == {
"test_schedule_name": "test_scouter",
"test_schedule_name2": "test_predictions_batch"
"scouter": {
"test_schedule_name": "test_scouter"
},
"laborious": {
"test_schedule_name2": "test_predictions_batch"
}
}
mock_scouter.assert_called_once_with(input_data['pipelines'][0])
mock_predictions_batch.assert_called_once_with(input_data['pipelines'][1])
mock_scouter.assert_called_once_with(
input_data['pipelines'][0])
mock_predictions_batch.assert_called_once_with(
input_data['pipelines'][1])
@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):
input_data = {
"current_schedule_config": {
"test_schedule_name_to_delete": {
"frequency": 60,
"data": {"test": "test"}
},
"test_schedule_name_to_update": {
"frequency": 60,
"data": {"test": "test"}
"scouter": {
"test_schedule_name_to_delete": {
"frequency": 60,
"data": {"test": "test"}
},
"test_schedule_name_to_update": {
"frequency": 60,
"data": {"test": "test"}
}
}
},
"schedule_config": {
"test_schedule_name_to_create": {
"frequency": 60,
"data": {"test": "test"}
"laborious": {
"test_schedule_name_to_create": {
"frequency": 60,
"data": {"test": "test"}
}
},
"test_schedule_name_to_update": {
"frequency": 60,
"data": {"test": "test2"}
"scouter": {
"test_schedule_name_to_update": {
"frequency": 60,
"data": {"test": "test2"}
}
}
}
}
@@ -268,20 +282,29 @@ async def test_create_schedule_config(formatters):
assert result == {
"to_create": {
"test_schedule_name_to_create": {
"frequency": 60,
"data": {"test": "test"}
}
"laborious": {
"test_schedule_name_to_create": {
"frequency": 60,
"data": {"test": "test"}
}
},
"scouter": {}
},
"to_update": {
"test_schedule_name_to_update": {
"frequency": 60,
"data": {"test": "test2"}
}
"scouter": {
"test_schedule_name_to_update": {
"frequency": 60,
"data": {"test": "test2"}
}
},
"laborious": {}
},
"to_delete": [
"test_schedule_name_to_delete"
]
"to_delete": {
"scouter": [
"test_schedule_name_to_delete"
],
"laborious": []
}
}

View File

@@ -3,18 +3,43 @@ from datetime import timedelta
import base64
import json
from pytest import fixture, mark
import pytest_asyncio
from orchestrator.activities.temporal_manager import TemporalManager
from orchestrator.utils.converters import parse_frequency
@fixture
def temporal_manager():
return TemporalManager(
temporal_client=MagicMock(),
@patch("orchestrator.activities.temporal_manager.Client.connect")
def temporal_manager(connect_mock):
temporal_manager = TemporalManager(
host='localhost:7233',
scouter_namespace='scouter',
laborious_namespace='laborious',
logger=MagicMock(),
notification_handler=MagicMock()
)
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
@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(
return_value=async_iter())
temporal_manager.temporal_client.get_schedule_handle.return_value = MagicMock(
handle = MagicMock(
describe=AsyncMock(
return_value=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(
intervals=[
MagicMock(
every=MagicMock(
seconds=60
)
temporal_manager.temporal_clients['scouter'].list_schedules = AsyncMock(
return_value=async_iter()
)
temporal_manager.temporal_clients['laborious'].list_schedules = AsyncMock(
return_value=async_iter()
)
temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock(
return_value=handle
)
temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock(
return_value=handle
)
describe_mock = MagicMock(
intervals=[
MagicMock(
every=MagicMock(
seconds=60
)
]
)
)
]
)
temporal_manager.temporal_clients['scouter'].get_schedule_handle.return_value.describe \
.return_value.schedule.spec = describe_mock
temporal_manager.temporal_clients['laborious'].get_schedule_handle.return_value.describe \
.return_value.schedule.spec = describe_mock
response = await temporal_manager.load_schedule()
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 == {
"test-schedule-id": {
"frequency": 60,
"data": {"test": "test"}
"scouter": {
"test-schedule-id": {
"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 = {
"schedules": {
"test-schedule": {
"model_id": 1,
"model_name": "test-model-name",
"workflow_type": "test-workflow",
"frequency": "1m",
"data": {"test": "test"}
"scouter": {
"test-schedule": {
"model_id": 1,
"model_name": "test-model-name",
"workflow_type": "test-workflow",
"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": {
"model_id": 2,
"model_name": "test-model-name",
"workflow_type": "test-workflow",
"frequency": "10y",
"data": {"test": "test"}
"laborious": {
"test-schedule-laborious": {
"model_id": 1,
"model_name": "test-model-name",
"workflow_type": "test-workflow",
"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)
temporal_manager.temporal_client.create_schedule.assert_called_once_with(
temporal_manager.temporal_clients['scouter'].create_schedule.assert_called_once_with(
"test-schedule",
mock_schedule.return_value,
search_attributes=mock_typed_search_attributes.return_value
)
mock_schedule.assert_called_once_with(
action=mock_schedule_action_start_workflow.return_value,
spec=mock_schedule_spec.return_value
temporal_manager.temporal_clients['laborious'].create_schedule.assert_called_once_with(
"test-schedule-laborious",
mock_schedule.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([
call(
"test-workflow",
input_data['schedules']['test-schedule'],
input_data['schedules']['scouter']['test-schedule'],
id="test-schedule",
task_queue="test-workflow-queue",
execution_timeout=ANY
),
call(
"test-workflow",
input_data['schedules']['test-schedule-invalid-frequency'],
input_data['schedules']['scouter']['test-schedule-invalid-frequency'],
id="test-schedule-invalid-frequency",
task_queue="test-workflow-queue",
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(
intervals=[
mock_schedule_interval_spec.return_value
]
)
mock_schedule_spec.assert_has_calls([
call(
intervals=[
mock_schedule_interval_spec.return_value
]
),
call(
intervals=[
mock_schedule_interval_spec.return_value
]
)
])
mock_schedule_interval_spec.assert_called_once_with(
every=timedelta(seconds=60)
)
mock_schedule_interval_spec.assert_has_calls([
call(
every=timedelta(seconds=60)
),
call(
every=timedelta(seconds=120)
)
])
mock_parse_frequency.assert_has_calls([
call("1m"),
call("10y")
call("10y"),
call("2m")
])
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
]),
call([
mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value
]),
call([
mock_search_attribute_pair.return_value,
mock_search_attribute_pair.return_value,
@@ -188,6 +290,18 @@ async def test_create_schedule(
key=temporal_manager.model_name_id_key,
value="test-model-name"
),
call(
key=temporal_manager.orchestrated_id_key,
value="true"
),
call(
key=temporal_manager.model_id_id_key,
value=2
),
call(
key=temporal_manager.model_name_id_key,
value="test-model-name"
),
call(
key=temporal_manager.orchestrated_id_key,
value="true"
@@ -202,10 +316,35 @@ async def test_create_schedule(
"test-schedule-invalid-frequency": {
"success": False,
"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
@patch("orchestrator.activities.temporal_manager.parse_frequency",
side_effect=parse_frequency)
@@ -218,30 +357,50 @@ async def test_update_schedules(
args=MagicMock()
)
temporal_manager.schedule_handles = {
"test-schedule": MagicMock(
update=AsyncMock(
"scouter": {
"test-schedule": MagicMock(
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 = {
"schedules": {
"test-schedule": {
"frequency": "1m",
"data": {"test": "test"}
"scouter": {
"test-schedule": {
"frequency": "1m",
"data": {"test": "test"}
},
"test-schedule_no_handler": {
"frequency": "1m",
"data": {"test": "test"}
}
},
"test-schedule_no_handler": {
"frequency": "1m",
"data": {"test": "test"}
"laborious": {
"test-schedule-laborious": {
"frequency": "2m",
"data": {"test": "test"}
}
}
}
}
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 == {
"test-schedule": {
@@ -251,21 +410,59 @@ async def test_update_schedules(
"test-schedule_no_handler": {
"success": False,
"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
async def test_delete_schedules(temporal_manager):
temporal_manager.schedule_handles = {
"test-schedule": MagicMock(
delete=AsyncMock()
)
"scouter": {
"test-schedule": MagicMock(
delete=AsyncMock()
)
},
"laborious": {
"test-schedule-laborious": MagicMock(
delete=AsyncMock()
)
}
}
input_data = {
"schedules": [
"test-schedule", "test-schedule_no_handler"
]
"schedules": {
"scouter": [
"test-schedule", "test-schedule_no_handler"
],
"laborious": [
"test-schedule-laborious"
]
}
}
report = await temporal_manager.delete_schedules(input_data)
@@ -278,5 +475,27 @@ async def test_delete_schedules(temporal_manager):
"test-schedule_no_handler": {
"success": False,
"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}"

View File

@@ -1,7 +1,7 @@
from os import environ
from orchestrator.utils.connectors_config import (build_redis_config,
build_couchbase_config,
build_mongodb_config)
build_mongodb_config, build_temporal_config)
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',
'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'
}