SIENTIAPDE-1030
Add comprehensive tests and utility functions for orchestrator activities - Introduced tests for the new Formatters class, covering methods for processing schedules and slots. - Enhanced SlotManager tests with update and delete slot functionalities. - Added TemporalManager tests for creating, updating, and deleting schedules, including frequency parsing. - Implemented utility functions for orchestrator operations, including frequency parsing and tag configuration building. - Created tests for utility functions to ensure correct behavior and integration with orchestrator activities. - Established a new converters module for parsing frequency strings into seconds.
This commit is contained in:
vitor-aignosi
@@ -4,12 +4,15 @@ from orchestrator.activities.activities import Activities
|
||||
from orchestrator.activities.couchbase import Couchbase
|
||||
from orchestrator.activities.temporal_manager import TemporalManager
|
||||
from orchestrator.activities.slot_manager import SlotManager
|
||||
from orchestrator.activities.formatters import Formatters
|
||||
|
||||
|
||||
@patch('orchestrator.activities.couchbase.Couchbase.__init__')
|
||||
@patch('orchestrator.activities.temporal_manager.TemporalManager.__init__')
|
||||
@patch('orchestrator.activities.slot_manager.SlotManager.__init__')
|
||||
def test___init__(mock_slot_manager_init, mock_temporal_manager_init,
|
||||
@patch('orchestrator.activities.formatters.Formatters.__init__')
|
||||
def test___init__(mock_formatters_init, mock_slot_manager_init,
|
||||
mock_temporal_manager_init,
|
||||
mock_couchbase_init):
|
||||
|
||||
couchbase_config = {
|
||||
@@ -41,6 +44,7 @@ def test___init__(mock_slot_manager_init, mock_temporal_manager_init,
|
||||
assert isinstance(activities, Couchbase)
|
||||
assert isinstance(activities, TemporalManager)
|
||||
assert isinstance(activities, SlotManager)
|
||||
assert isinstance(activities, Formatters)
|
||||
|
||||
mock_slot_manager_init.assert_called_once_with(
|
||||
ANY,
|
||||
@@ -68,6 +72,12 @@ def test___init__(mock_slot_manager_init, mock_temporal_manager_init,
|
||||
notification_handler=notification_handler
|
||||
)
|
||||
|
||||
mock_formatters_init.assert_called_once_with(
|
||||
ANY,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.activities.couchbase.Cluster')
|
||||
|
||||
480
tests/orchestrator/activities/test_formatters.py
Normal file
480
tests/orchestrator/activities/test_formatters.py
Normal file
@@ -0,0 +1,480 @@
|
||||
from unittest.mock import MagicMock, patch, call, ANY
|
||||
import json
|
||||
from pytest import fixture, mark
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from orchestrator.activities.formatters import Formatters
|
||||
from orchestrator.utils.orchestrator_functions import build_tag_config
|
||||
|
||||
|
||||
@fixture
|
||||
def formatters():
|
||||
return Formatters(
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock()
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("orchestrator.activities.formatters.scouter",
|
||||
return_value="test_scouter")
|
||||
@patch("orchestrator.activities.formatters.predictions_batch",
|
||||
return_value="test_predictions_batch")
|
||||
async def test_process_schedules(mock_predictions_batch, mock_scouter, formatters):
|
||||
input_data = {
|
||||
"pipelines": [
|
||||
{
|
||||
"schedule_name": "test_schedule_name",
|
||||
"workflow_type": "scouter",
|
||||
"model_name": "test_model_name",
|
||||
"model_id": "test_model_id"
|
||||
},
|
||||
{
|
||||
"schedule_name": "test_schedule_name2",
|
||||
"workflow_type": "predictions_batch",
|
||||
"model_name": "test_model_name",
|
||||
"model_id": "test_model_id"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
result = await formatters.process_schedules(input_data)
|
||||
|
||||
assert result == {
|
||||
"test_schedule_name": "test_scouter",
|
||||
"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])
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("orchestrator.activities.formatters.gather_read_tags",
|
||||
return_value={
|
||||
"test_server_name:test_tag_address": {
|
||||
"server_name": "test_server_name",
|
||||
"tag_address": "test_tag_address",
|
||||
"topics": ["raw_test_schedule"]
|
||||
},
|
||||
"test_server_name2:test_tag_address2": {
|
||||
"server_name": "test_server_name2",
|
||||
"tag_address": "test_tag_address2",
|
||||
"topics": ["raw_test_schedule2"]
|
||||
},
|
||||
"test_server_name2:test_tag_address3": {
|
||||
"server_name": "test_server_name2",
|
||||
"tag_address": "test_tag_address3",
|
||||
"topics": ["raw_test_schedule2"]
|
||||
}
|
||||
})
|
||||
@patch("orchestrator.activities.formatters.build_tag_config", side_effect=build_tag_config)
|
||||
async def test_process_slots(mock_build_tag_config, mock_gather_read_tags, formatters):
|
||||
input_data = {
|
||||
"opc_servers": [
|
||||
{
|
||||
"server_name": "test_server_name",
|
||||
"url": "test_url",
|
||||
"uri": "test_uri",
|
||||
"security_spec": {
|
||||
"test_name": "test_spec"
|
||||
}
|
||||
},
|
||||
{
|
||||
"server_name": "test_server_name2",
|
||||
"url": "test_url2",
|
||||
"uri": "test_uri2"
|
||||
}
|
||||
],
|
||||
"active_ingestors": [
|
||||
"test_active_ingestor1",
|
||||
"test_active_ingestor2"
|
||||
],
|
||||
"pipelines": "test_gather_read_tags"
|
||||
}
|
||||
|
||||
result = await formatters.process_slots(input_data)
|
||||
|
||||
mock_gather_read_tags.assert_called_once_with(input_data['pipelines'])
|
||||
mock_build_tag_config.assert_has_calls([
|
||||
call(
|
||||
{
|
||||
"server_name": "test_server_name",
|
||||
"tag_address": "test_tag_address",
|
||||
"topics": ["raw_test_schedule"]
|
||||
},
|
||||
ANY,
|
||||
{
|
||||
"test_server_name": {
|
||||
"server_name": "test_server_name",
|
||||
"url": "test_url",
|
||||
"uri": "test_uri",
|
||||
"security_spec": {
|
||||
"test_name": "test_spec"
|
||||
}
|
||||
},
|
||||
"test_server_name2": {
|
||||
"server_name": "test_server_name2",
|
||||
"url": "test_url2",
|
||||
"uri": "test_uri2"
|
||||
}
|
||||
},
|
||||
1
|
||||
)
|
||||
])
|
||||
mock_build_tag_config.assert_has_calls([
|
||||
call(
|
||||
{
|
||||
"server_name": "test_server_name2",
|
||||
"tag_address": "test_tag_address2",
|
||||
"topics": ["raw_test_schedule2"]
|
||||
},
|
||||
ANY,
|
||||
{
|
||||
"test_server_name": {
|
||||
"server_name": "test_server_name",
|
||||
"url": "test_url",
|
||||
"uri": "test_uri",
|
||||
"security_spec": {
|
||||
"test_name": "test_spec"
|
||||
}
|
||||
},
|
||||
"test_server_name2": {
|
||||
"server_name": "test_server_name2",
|
||||
"url": "test_url2",
|
||||
"uri": "test_uri2"
|
||||
}
|
||||
},
|
||||
1
|
||||
)
|
||||
])
|
||||
mock_build_tag_config.assert_has_calls([
|
||||
call(
|
||||
{
|
||||
"server_name": "test_server_name2",
|
||||
"tag_address": "test_tag_address3",
|
||||
"topics": ["raw_test_schedule2"]
|
||||
},
|
||||
ANY,
|
||||
{
|
||||
"test_server_name": {
|
||||
"server_name": "test_server_name",
|
||||
"url": "test_url",
|
||||
"uri": "test_uri",
|
||||
"security_spec": {
|
||||
"test_name": "test_spec"
|
||||
}
|
||||
},
|
||||
"test_server_name2": {
|
||||
"server_name": "test_server_name2",
|
||||
"url": "test_url2",
|
||||
"uri": "test_uri2"
|
||||
}
|
||||
},
|
||||
2
|
||||
)
|
||||
])
|
||||
|
||||
assert result == {
|
||||
"1": {
|
||||
"test_server_name": {
|
||||
"name": "test_server_name",
|
||||
"url": "test_url",
|
||||
"server_uri": "test_uri",
|
||||
"test_name": "test_spec",
|
||||
"tags": {
|
||||
"test_tag_address": {
|
||||
"server_name": "test_server_name",
|
||||
"tag_address": "test_tag_address",
|
||||
"topics": ["raw_test_schedule"]
|
||||
}
|
||||
}
|
||||
},
|
||||
"test_server_name2": {
|
||||
"name": "test_server_name2",
|
||||
"url": "test_url2",
|
||||
"server_uri": "test_uri2",
|
||||
"tags": {
|
||||
"test_tag_address2": {
|
||||
"server_name": "test_server_name2",
|
||||
"tag_address": "test_tag_address2",
|
||||
"topics": ["raw_test_schedule2"]
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"2": {
|
||||
"test_server_name2": {
|
||||
"name": "test_server_name2",
|
||||
"url": "test_url2",
|
||||
"server_uri": "test_uri2",
|
||||
"tags": {
|
||||
"test_tag_address3": {
|
||||
"server_name": "test_server_name2",
|
||||
"tag_address": "test_tag_address3",
|
||||
"topics": ["raw_test_schedule2"]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
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"}
|
||||
}
|
||||
},
|
||||
"schedule_config": {
|
||||
"test_schedule_name_to_create": {
|
||||
"frequency": 60,
|
||||
"data": {"test": "test"}
|
||||
},
|
||||
"test_schedule_name_to_update": {
|
||||
"frequency": 60,
|
||||
"data": {"test": "test2"}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
result = await formatters.create_schedule_config(input_data)
|
||||
|
||||
assert result == {
|
||||
"to_create": {
|
||||
"test_schedule_name_to_create": {
|
||||
"frequency": 60,
|
||||
"data": {"test": "test"}
|
||||
}
|
||||
},
|
||||
"to_update": {
|
||||
"test_schedule_name_to_update": {
|
||||
"frequency": 60,
|
||||
"data": {"test": "test2"}
|
||||
}
|
||||
},
|
||||
"to_delete": [
|
||||
"test_schedule_name_to_delete"
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_create_slot_config(formatters):
|
||||
input_data = {
|
||||
"current_slot_config": {
|
||||
"1": {
|
||||
"frequency": 60,
|
||||
"data": {"test": "test"}
|
||||
},
|
||||
"2": {
|
||||
"frequency": 60,
|
||||
"data": {"test": "test"}
|
||||
}
|
||||
},
|
||||
"slot_config": {
|
||||
"1": {
|
||||
"frequency": 60,
|
||||
"data": {"test": "test2"}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
result = await formatters.create_slot_config(input_data)
|
||||
|
||||
assert result == {
|
||||
"to_delete": [
|
||||
"2"
|
||||
],
|
||||
"to_insert": {
|
||||
"1": {
|
||||
"frequency": 60,
|
||||
"data": {"test": "test2"}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_send_success_report(formatters):
|
||||
formatters.send_success_report("test_message", "test_notification_id")
|
||||
formatters.notification_handler.build_and_send_notification.assert_called_once_with(
|
||||
"test_notification_id",
|
||||
"test_message",
|
||||
"report_orchestration",
|
||||
NotificationLevel.INFO
|
||||
)
|
||||
|
||||
|
||||
def test_send_error_report(formatters):
|
||||
formatters.send_error_report(
|
||||
"test_message", "test_notification_id", {"test": "test"})
|
||||
formatters.notification_handler.build_and_send_notification.assert_called_once_with(
|
||||
"test_notification_id",
|
||||
"test_message",
|
||||
"report_orchestration",
|
||||
NotificationLevel.ERROR,
|
||||
attachment_content=json.dumps(
|
||||
{"test": "test"}, indent=4, sort_keys=True)
|
||||
)
|
||||
|
||||
|
||||
def test_parse_report(formatters):
|
||||
input_data = {
|
||||
"test_schedule_name_to_create": {
|
||||
"success": True
|
||||
},
|
||||
"test_schedule_name_to_create_error": {
|
||||
"success": False,
|
||||
"error": "test_error"
|
||||
}
|
||||
}
|
||||
|
||||
result = formatters.parse_report(input_data)
|
||||
|
||||
assert result == (
|
||||
["test_schedule_name_to_create"],
|
||||
["test_schedule_name_to_create_error"]
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_report_schedule_orchestration(formatters):
|
||||
formatters.parse_report = MagicMock(
|
||||
side_effect=formatters.parse_report
|
||||
)
|
||||
formatters.send_success_report = MagicMock()
|
||||
formatters.send_error_report = MagicMock()
|
||||
|
||||
input_data = {
|
||||
"created_schedules": {
|
||||
"test_schedule_name_to_create": {
|
||||
"success": True
|
||||
},
|
||||
"test_schedule_name_to_create_error": {
|
||||
"success": False,
|
||||
"error": "test_error"
|
||||
}
|
||||
},
|
||||
"updated_schedules": {
|
||||
"test_schedule_name_to_update": {
|
||||
"success": True
|
||||
},
|
||||
"test_schedule_name_to_update_error": {
|
||||
"success": False,
|
||||
"error": "test_error"
|
||||
}
|
||||
},
|
||||
"deleted_schedules": {
|
||||
"test_schedule_name_to_delete": {
|
||||
"success": True
|
||||
},
|
||||
"test_schedule_name_to_delete_error": {
|
||||
"success": False,
|
||||
"error": "test_error"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
await formatters.report_schedule_orchestration(input_data)
|
||||
|
||||
formatters.parse_report.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",
|
||||
"REPORT_ORCHESTRATION_CREATED_SCHEDULES"
|
||||
),
|
||||
call(
|
||||
"Updated schedules: \n test_schedule_name_to_update",
|
||||
"REPORT_ORCHESTRATION_UPDATED_SCHEDULES"
|
||||
),
|
||||
call(
|
||||
"Deleted schedules: \n 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",
|
||||
"REPORT_ORCHESTRATION_CREATED_SCHEDULES",
|
||||
input_data['created_schedules']
|
||||
),
|
||||
call(
|
||||
"Failed to update schedules: \n test_schedule_name_to_update_error",
|
||||
"REPORT_ORCHESTRATION_UPDATED_SCHEDULES",
|
||||
input_data['updated_schedules']
|
||||
),
|
||||
call(
|
||||
"Failed to delete schedules: \n test_schedule_name_to_delete_error",
|
||||
"REPORT_ORCHESTRATION_DELETED_SCHEDULES",
|
||||
input_data['deleted_schedules']
|
||||
)
|
||||
])
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_report_slot_orchestration(formatters):
|
||||
formatters.parse_report = MagicMock(
|
||||
side_effect=formatters.parse_report
|
||||
)
|
||||
formatters.send_success_report = MagicMock()
|
||||
formatters.send_error_report = MagicMock()
|
||||
|
||||
input_data = {
|
||||
"inserted_slots": {
|
||||
"test_slot_name_to_create": {
|
||||
"success": True
|
||||
},
|
||||
"test_slot_name_to_create_error": {
|
||||
"success": False,
|
||||
"error": "test_error"
|
||||
}
|
||||
},
|
||||
"deleted_slots": {
|
||||
"test_slot_name_to_delete": {
|
||||
"success": True
|
||||
},
|
||||
"test_slot_name_to_delete_error": {
|
||||
"success": False,
|
||||
"error": "test_error"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
await formatters.report_slot_orchestration(input_data)
|
||||
|
||||
formatters.parse_report.assert_has_calls([
|
||||
call(input_data['inserted_slots']),
|
||||
call(input_data['deleted_slots'])
|
||||
])
|
||||
formatters.send_success_report.assert_has_calls([
|
||||
call(
|
||||
"Inserted slots: \n test_slot_name_to_create",
|
||||
"REPORT_ORCHESTRATION_INSERTED_SLOTS"
|
||||
),
|
||||
call(
|
||||
"Deleted slots: \n test_slot_name_to_delete",
|
||||
"REPORT_ORCHESTRATION_DELETED_SLOTS"
|
||||
)
|
||||
])
|
||||
formatters.send_error_report.assert_has_calls([
|
||||
call(
|
||||
"Failed to insert slots: \n test_slot_name_to_create_error",
|
||||
"REPORT_ORCHESTRATION_INSERTED_SLOTS",
|
||||
input_data['inserted_slots']
|
||||
),
|
||||
call(
|
||||
"Failed to delete slots: \n test_slot_name_to_delete_error",
|
||||
"REPORT_ORCHESTRATION_DELETED_SLOTS",
|
||||
input_data['deleted_slots']
|
||||
)
|
||||
])
|
||||
@@ -1,4 +1,4 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import MagicMock, patch, call
|
||||
from pytest import mark, fixture
|
||||
from orchestrator.activities.slot_manager import SlotManager
|
||||
|
||||
@@ -55,3 +55,66 @@ async def test_load_active_ingestors(slot_manager):
|
||||
|
||||
assert response == ["heartbeat:ingestor:1",
|
||||
"heartbeat:ingestor:2", "heartbeat:ingestor:3"]
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_update_slots(slot_manager):
|
||||
slot_manager.set = MagicMock(
|
||||
side_effect=[
|
||||
None,
|
||||
Exception("Test exception")
|
||||
]
|
||||
)
|
||||
|
||||
response = await slot_manager.update_slots({
|
||||
"to_insert": {
|
||||
"1": "value1",
|
||||
"2": "value2"
|
||||
}
|
||||
})
|
||||
|
||||
slot_manager.set.assert_has_calls([
|
||||
call("slot:opc_tags:1", "value1", ttl=None),
|
||||
call("slot:opc_tags:2", "value2", ttl=None)
|
||||
])
|
||||
|
||||
assert response == {
|
||||
"1": {
|
||||
"success": True,
|
||||
"message": "Slot updated successfully"
|
||||
},
|
||||
"2": {
|
||||
"success": False,
|
||||
"message": "Test exception"
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_delete_slots(slot_manager):
|
||||
slot_manager.redis_client.delete = MagicMock(
|
||||
side_effect=[
|
||||
None,
|
||||
Exception("Test exception")
|
||||
]
|
||||
)
|
||||
|
||||
response = await slot_manager.delete_slots({
|
||||
"to_delete": ["1", "2"]
|
||||
})
|
||||
|
||||
slot_manager.redis_client.delete.assert_has_calls([
|
||||
call("slot:opc_tags:1"),
|
||||
call("slot:opc_tags:2")
|
||||
])
|
||||
|
||||
assert response == {
|
||||
"1": {
|
||||
"success": True,
|
||||
"message": "Slot deleted successfully"
|
||||
},
|
||||
"2": {
|
||||
"success": False,
|
||||
"message": "Test exception"
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
from unittest.mock import MagicMock, patch, AsyncMock
|
||||
from unittest.mock import MagicMock, patch, AsyncMock, call
|
||||
from datetime import timedelta
|
||||
import base64
|
||||
import json
|
||||
from pytest import fixture, mark
|
||||
from orchestrator.activities.temporal_manager import TemporalManager
|
||||
from orchestrator.utils.converters import parse_frequency
|
||||
|
||||
|
||||
@fixture
|
||||
@@ -41,7 +43,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.return_value = MagicMock(
|
||||
temporal_manager.temporal_client.get_schedule_handle.return_value = MagicMock(
|
||||
describe=AsyncMock(
|
||||
return_value=MagicMock(
|
||||
schedule=MagicMock(
|
||||
@@ -57,7 +59,7 @@ async def test_load_schedule(_mock_message_to_dict, temporal_manager):
|
||||
)
|
||||
)
|
||||
)
|
||||
temporal_manager.temporal_client.get_schedule.return_value.describe \
|
||||
temporal_manager.temporal_client.get_schedule_handle.return_value.describe \
|
||||
.return_value.schedule.spec = MagicMock(
|
||||
intervals=[
|
||||
MagicMock(
|
||||
@@ -74,7 +76,205 @@ async def test_load_schedule(_mock_message_to_dict, temporal_manager):
|
||||
assert response == {
|
||||
"test-schedule-id": {
|
||||
"frequency": 60,
|
||||
"data": {"test": "test"},
|
||||
"handle": temporal_manager.temporal_client.get_schedule.return_value
|
||||
"data": {"test": "test"}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("orchestrator.activities.temporal_manager.parse_frequency",
|
||||
side_effect=parse_frequency)
|
||||
@patch("orchestrator.activities.temporal_manager.Schedule")
|
||||
@patch("orchestrator.activities.temporal_manager.ScheduleActionStartWorkflow")
|
||||
@patch("orchestrator.activities.temporal_manager.ScheduleIntervalSpec")
|
||||
@patch("orchestrator.activities.temporal_manager.ScheduleSpec")
|
||||
@patch("orchestrator.activities.temporal_manager.TypedSearchAttributes")
|
||||
@patch("orchestrator.activities.temporal_manager.SearchAttributePair")
|
||||
async def test_create_schedule(
|
||||
mock_search_attribute_pair,
|
||||
mock_typed_search_attributes,
|
||||
mock_schedule_spec,
|
||||
mock_schedule_interval_spec,
|
||||
mock_schedule_action_start_workflow,
|
||||
mock_schedule,
|
||||
mock_parse_frequency,
|
||||
temporal_manager):
|
||||
|
||||
input_data = {
|
||||
"schedules": {
|
||||
"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"}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
temporal_manager.temporal_client.create_schedule = AsyncMock()
|
||||
|
||||
report = await temporal_manager.create_schedules(input_data)
|
||||
|
||||
temporal_manager.temporal_client.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
|
||||
)
|
||||
|
||||
mock_schedule_action_start_workflow.assert_has_calls([
|
||||
call(
|
||||
workflow="test-workflow",
|
||||
args=input_data['schedules']['test-schedule'],
|
||||
id="test-schedule",
|
||||
task_queue="test-workflow-queue"
|
||||
),
|
||||
call(
|
||||
workflow="test-workflow",
|
||||
args=input_data['schedules']['test-schedule-invalid-frequency'],
|
||||
id="test-schedule-invalid-frequency",
|
||||
task_queue="test-workflow-queue"
|
||||
)
|
||||
])
|
||||
|
||||
mock_schedule_spec.assert_called_once_with(
|
||||
intervals=[
|
||||
mock_schedule_interval_spec.return_value
|
||||
]
|
||||
)
|
||||
|
||||
mock_schedule_interval_spec.assert_called_once_with(
|
||||
every=timedelta(seconds=60)
|
||||
)
|
||||
|
||||
mock_parse_frequency.assert_has_calls([
|
||||
call("1m"),
|
||||
call("10y")
|
||||
])
|
||||
|
||||
mock_typed_search_attributes.assert_has_calls([
|
||||
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,
|
||||
mock_search_attribute_pair.return_value
|
||||
])
|
||||
])
|
||||
|
||||
mock_search_attribute_pair.assert_has_calls([
|
||||
call(
|
||||
key=temporal_manager.model_id_id_key,
|
||||
value=1
|
||||
),
|
||||
call(
|
||||
key=temporal_manager.model_name_id_key,
|
||||
value="test-model-name"
|
||||
),
|
||||
call(
|
||||
key=temporal_manager.orchestrated_id_key,
|
||||
value="true"
|
||||
)
|
||||
])
|
||||
|
||||
assert report == {
|
||||
"test-schedule": {
|
||||
"success": True,
|
||||
"message": "Schedule created successfully"
|
||||
},
|
||||
"test-schedule-invalid-frequency": {
|
||||
"success": False,
|
||||
"message": "Invalid frequency"
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("orchestrator.activities.temporal_manager.parse_frequency",
|
||||
side_effect=parse_frequency)
|
||||
@patch("orchestrator.activities.temporal_manager.ScheduleIntervalSpec")
|
||||
async def test_update_schedules(
|
||||
_mock_schedule_interval_spec,
|
||||
_mock_parse_frequency,
|
||||
temporal_manager):
|
||||
input_mock = MagicMock(
|
||||
args=MagicMock()
|
||||
)
|
||||
temporal_manager.schedule_handles = {
|
||||
"test-schedule": MagicMock(
|
||||
update=AsyncMock(
|
||||
update=AsyncMock(
|
||||
side_effect=lambda f: f(input_mock)
|
||||
)
|
||||
)
|
||||
)
|
||||
}
|
||||
input_data = {
|
||||
"schedules": {
|
||||
"test-schedule": {
|
||||
"frequency": "1m",
|
||||
"data": {"test": "test"}
|
||||
},
|
||||
"test-schedule_no_handler": {
|
||||
"frequency": "1m",
|
||||
"data": {"test": "test"}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
report = await temporal_manager.update_schedules(input_data)
|
||||
|
||||
temporal_manager.schedule_handles['test-schedule'].update.assert_called_once()
|
||||
|
||||
assert report == {
|
||||
"test-schedule": {
|
||||
"success": True,
|
||||
"message": "Schedule updated successfully"
|
||||
},
|
||||
"test-schedule_no_handler": {
|
||||
"success": False,
|
||||
"message": "Schedule test-schedule_no_handler not found"
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_delete_schedules(temporal_manager):
|
||||
temporal_manager.schedule_handles = {
|
||||
"test-schedule": MagicMock(
|
||||
delete=AsyncMock()
|
||||
)
|
||||
}
|
||||
input_data = {
|
||||
"schedules": [
|
||||
"test-schedule", "test-schedule_no_handler"
|
||||
]
|
||||
}
|
||||
|
||||
report = await temporal_manager.delete_schedules(input_data)
|
||||
|
||||
assert report == {
|
||||
"test-schedule": {
|
||||
"success": True,
|
||||
"message": "Schedule deleted successfully"
|
||||
},
|
||||
"test-schedule_no_handler": {
|
||||
"success": False,
|
||||
"message": "Schedule test-schedule_no_handler not found"
|
||||
}
|
||||
}
|
||||
|
||||
325
tests/orchestrator/utils/test_orchestrator_functions.py
Normal file
325
tests/orchestrator/utils/test_orchestrator_functions.py
Normal file
@@ -0,0 +1,325 @@
|
||||
from unittest.mock import patch, call
|
||||
from orchestrator.utils.orchestrator_functions import (
|
||||
common_config,
|
||||
scouter,
|
||||
predictions_batch,
|
||||
overlap_filter_config,
|
||||
process_path_priority,
|
||||
gather_read_tags,
|
||||
build_tag_config
|
||||
)
|
||||
|
||||
|
||||
def test_common_config():
|
||||
config = {
|
||||
"schedule_name": "test_schedule",
|
||||
"model_id": "test_model_id",
|
||||
"model_name": "test_model_name"
|
||||
}
|
||||
result = common_config(config)
|
||||
expected = {
|
||||
"workflow_type": "scouter",
|
||||
"schedule_name": "test_schedule",
|
||||
"frequency": "1m",
|
||||
"max_retry_policy": 1,
|
||||
"model_id": "test_model_id",
|
||||
"model_name": "test_model_name"
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_scouter():
|
||||
config = {
|
||||
"schedule_name": "test_schedule",
|
||||
"model_id": "test_model_id",
|
||||
"model_name": "test_model_name",
|
||||
"filters": [
|
||||
{
|
||||
"filter_name": "test_filter_name",
|
||||
"policy": "test_policy"
|
||||
}
|
||||
],
|
||||
"read_tags": [
|
||||
{
|
||||
"tag_name": "test_tag_name",
|
||||
"aggr_func": "test_aggr_func",
|
||||
"data_range": [1, 2]
|
||||
}
|
||||
],
|
||||
"tag_retention_minutes": 10
|
||||
}
|
||||
result = scouter(config)
|
||||
expected = {
|
||||
"workflow_type": "scouter",
|
||||
"schedule_name": "test_schedule",
|
||||
"frequency": "1m",
|
||||
"max_retry_policy": 1,
|
||||
"model_id": "test_model_id",
|
||||
"model_name": "test_model_name",
|
||||
"topic": "raw_test_schedule",
|
||||
"trigger_laborious": False,
|
||||
"filters": {
|
||||
"test_filter_name": {
|
||||
"policy": "test_policy"
|
||||
}
|
||||
},
|
||||
"schema": "sientia_data",
|
||||
"table_name": "laborious_data",
|
||||
"retention_time": 10 * 60,
|
||||
"model_tags": {
|
||||
"test_tag_name": {
|
||||
"aggr_func": "test_aggr_func",
|
||||
"data_range": [1, 2]
|
||||
}
|
||||
}
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_overlap_filter_config():
|
||||
config = [
|
||||
{
|
||||
"filter_name": "test_filter_name",
|
||||
"policy": "test_policy"
|
||||
},
|
||||
{
|
||||
"filter_name": "test_filter_name2",
|
||||
"policy": "test_policy2"
|
||||
}
|
||||
]
|
||||
result = overlap_filter_config({
|
||||
"test_filter_name": {
|
||||
"policy": "test_policy"
|
||||
}
|
||||
}, config)
|
||||
expected = {
|
||||
"test_filter_name": {
|
||||
"policy": "test_policy",
|
||||
"config": {}
|
||||
},
|
||||
"test_filter_name2": {
|
||||
"policy": "test_policy2",
|
||||
"config": {}
|
||||
}
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_process_path_priority():
|
||||
config = ["OTHER", "STOP", "CONTINUE"]
|
||||
result = process_path_priority(config)
|
||||
expected = ["STOP", "CONTINUE", "REPEAT"]
|
||||
assert result == expected
|
||||
|
||||
|
||||
@patch('orchestrator.utils.orchestrator_functions.overlap_filter_config',
|
||||
return_value={
|
||||
"test_filter_name": {
|
||||
"policy": "test_policy",
|
||||
"config": {}
|
||||
}
|
||||
})
|
||||
@patch('orchestrator.utils.orchestrator_functions.process_path_priority',
|
||||
return_value=["STOP", "CONTINUE", "REPEAT"])
|
||||
def test_predictions_batch(mock_process_path_priority,
|
||||
mock_overlap_filter_config):
|
||||
config = {
|
||||
"schedule_name": "test_schedule",
|
||||
"workflow_type": "predictions_batch",
|
||||
"model_id": "test_model_id",
|
||||
"model_name": "test_model_name",
|
||||
"query": "test_query",
|
||||
"write_tags": [
|
||||
{
|
||||
"server_name": "test_server_name",
|
||||
"type": "prediction",
|
||||
"addr": "test_addr"
|
||||
},
|
||||
{
|
||||
"server_name": "test_server_name",
|
||||
"type": "confidence",
|
||||
"addr": "test_addr"
|
||||
}
|
||||
],
|
||||
"input_filters": [
|
||||
{
|
||||
"filter_name": "test_filter_name",
|
||||
"policy": "test_policy"
|
||||
}
|
||||
],
|
||||
"mlflow_transform_filters": [
|
||||
{
|
||||
"filter_name": "test_filter_name",
|
||||
"policy": "test_policy"
|
||||
}
|
||||
],
|
||||
"mlflow_predict_filters": [
|
||||
{
|
||||
"filter_name": "test_filter_name",
|
||||
"policy": "test_policy"
|
||||
}
|
||||
],
|
||||
"path_priority": ["STOP", "CONTINUE", "REPEAT"],
|
||||
}
|
||||
|
||||
result = predictions_batch(config)
|
||||
|
||||
mock_overlap_filter_config.assert_has_calls([
|
||||
call({
|
||||
"EMPTY_DATA": {
|
||||
"policy": "STOP",
|
||||
"config": {}
|
||||
}
|
||||
}, config['input_filters'])
|
||||
])
|
||||
mock_overlap_filter_config.assert_has_calls([
|
||||
call({
|
||||
"API_ERROR": {
|
||||
"policy": "STOP",
|
||||
"config": {}
|
||||
}
|
||||
}, config['mlflow_transform_filters'])
|
||||
])
|
||||
mock_overlap_filter_config.assert_has_calls([
|
||||
call({
|
||||
"API_ERROR": {
|
||||
"policy": "STOP",
|
||||
"config": {}
|
||||
}
|
||||
}, config['mlflow_predict_filters'])
|
||||
])
|
||||
mock_process_path_priority.assert_called_once_with(config['path_priority'])
|
||||
|
||||
expected = {
|
||||
"workflow_type": "predictions_batch",
|
||||
"schedule_name": "test_schedule",
|
||||
"frequency": "1m",
|
||||
"max_retry_policy": 1,
|
||||
"model_id": "test_model_id",
|
||||
"model_name": "test_model_name",
|
||||
"query": "test_query",
|
||||
"schema": "sientia_data",
|
||||
"table_name": "predictions",
|
||||
"retention_time": 60 * 60,
|
||||
"opc_output_config": {
|
||||
"test_server_name": {
|
||||
"prediction_tags": {
|
||||
"test_addr": {
|
||||
"data_type": "float"
|
||||
}
|
||||
},
|
||||
"confidence_tags": {
|
||||
"test_addr": {
|
||||
"data_type": "float"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"input_filters": {
|
||||
"test_filter_name": {
|
||||
"policy": "test_policy",
|
||||
"config": {}
|
||||
}
|
||||
},
|
||||
"mlflow_transform_filters": {
|
||||
"test_filter_name": {
|
||||
"policy": "test_policy",
|
||||
"config": {}
|
||||
}
|
||||
},
|
||||
"mlflow_predict_filters": {
|
||||
"test_filter_name": {
|
||||
"policy": "test_policy",
|
||||
"config": {}
|
||||
}
|
||||
},
|
||||
"path_priority": ["STOP", "CONTINUE", "REPEAT"]
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_gather_read_tags():
|
||||
pipelines = [
|
||||
{
|
||||
"schedule_name": "test_schedule",
|
||||
"read_tags": [
|
||||
{
|
||||
"server_name": "test_server_name",
|
||||
"tag_address": "test_tag_address"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"schedule_name": "test_schedule2",
|
||||
"read_tags": [
|
||||
{
|
||||
"server_name": "test_server_name2",
|
||||
"tag_address": "test_tag_address2"
|
||||
},
|
||||
{
|
||||
"server_name": "test_server_name2",
|
||||
"tag_address": "test_tag_address3"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
result = gather_read_tags(pipelines)
|
||||
|
||||
expected = {
|
||||
"test_server_name:test_tag_address": {
|
||||
"server_name": "test_server_name",
|
||||
"tag_address": "test_tag_address",
|
||||
"topics": ["raw_test_schedule"]
|
||||
},
|
||||
"test_server_name2:test_tag_address2": {
|
||||
"server_name": "test_server_name2",
|
||||
"tag_address": "test_tag_address2",
|
||||
"topics": ["raw_test_schedule2"]
|
||||
},
|
||||
"test_server_name2:test_tag_address3": {
|
||||
"server_name": "test_server_name2",
|
||||
"tag_address": "test_tag_address3",
|
||||
"topics": ["raw_test_schedule2"]
|
||||
}
|
||||
}
|
||||
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_build_tag_config():
|
||||
tag = {
|
||||
"server_name": "test_server_name",
|
||||
"tag_address": "test_tag_address"
|
||||
}
|
||||
opc_servers = {
|
||||
"test_server_name": {
|
||||
"url": "test_url",
|
||||
"uri": "test_uri",
|
||||
"security_spec": {
|
||||
"test_name": "test_spec"
|
||||
}
|
||||
}
|
||||
}
|
||||
slot_config = {
|
||||
"1": {}
|
||||
}
|
||||
i = 1
|
||||
result = build_tag_config(tag, slot_config, opc_servers, i)
|
||||
expected = {
|
||||
"1": {
|
||||
"test_server_name": {
|
||||
"name": "test_server_name",
|
||||
"url": "test_url",
|
||||
"server_uri": "test_uri",
|
||||
"tags": {
|
||||
"test_tag_address": {
|
||||
"server_name": "test_server_name",
|
||||
"tag_address": "test_tag_address"
|
||||
}
|
||||
},
|
||||
"test_name": "test_spec"
|
||||
}
|
||||
}
|
||||
}
|
||||
assert result == expected
|
||||
@@ -79,3 +79,111 @@ async def test_run(workflow_mock, orchestrator):
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.process_schedules,
|
||||
{
|
||||
'pipelines': 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_slots,
|
||||
{
|
||||
'opc_servers': workflow_mock.execute_local_activity_method.return_value,
|
||||
'active_ingestors': workflow_mock.execute_local_activity_method.return_value,
|
||||
'pipelines': 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.create_schedule_config,
|
||||
{
|
||||
'current_schedule_config': workflow_mock.execute_local_activity_method.return_value,
|
||||
'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.create_slot_config,
|
||||
{
|
||||
'current_slot_config': workflow_mock.execute_local_activity_method.return_value,
|
||||
'slot_config': 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,
|
||||
{
|
||||
'to_delete':
|
||||
workflow_mock.execute_local_activity_method.return_value['to_delete']
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.update_slots,
|
||||
{
|
||||
'to_insert':
|
||||
workflow_mock.execute_local_activity_method.return_value['to_insert']
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.delete_schedules,
|
||||
{
|
||||
'schedules':
|
||||
workflow_mock.execute_local_activity_method.return_value['to_delete']
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.create_schedules,
|
||||
{
|
||||
'schedules':
|
||||
workflow_mock.execute_local_activity_method.return_value['to_create']
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.update_schedules,
|
||||
{
|
||||
'schedules':
|
||||
workflow_mock.execute_local_activity_method.return_value['to_update']
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
|
||||
Reference in New Issue
Block a user