SIENTIAPDE-1030
Add unit tests for orchestrator activities and workflows - Implement tests for Activities class, covering initialization and prepare_activity method. - Create tests for Couchbase class, including successful and failed query loading. - Add tests for SlotManager class, verifying OPC slot loading and active ingestor retrieval. - Develop tests for TemporalManager class, focusing on schedule loading functionality. - Introduce tests for Orchestrator class, ensuring proper execution of workflow activities. - Establish a new test suite for orchestrator activities and workflows in the tests directory.
This commit is contained in:
vitor-aignosi
0
tests/orchestrator/__init__.py
Normal file
0
tests/orchestrator/__init__.py
Normal file
116
tests/orchestrator/activities/test_activities.py
Normal file
116
tests/orchestrator/activities/test_activities.py
Normal file
@@ -0,0 +1,116 @@
|
||||
from unittest.mock import patch, MagicMock, ANY
|
||||
from pytest import mark
|
||||
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
|
||||
|
||||
|
||||
@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,
|
||||
mock_couchbase_init):
|
||||
|
||||
couchbase_config = {
|
||||
'connection_string': 'couchbase://localhost',
|
||||
'username': 'admin',
|
||||
'password': 'password'
|
||||
}
|
||||
|
||||
redis_config = {
|
||||
'host': 'localhost',
|
||||
'port': 6379,
|
||||
'username': 'admin',
|
||||
'password': 'password'
|
||||
}
|
||||
|
||||
temporal_client = MagicMock()
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
|
||||
activities = Activities(
|
||||
temporal_client=temporal_client,
|
||||
couchbase_config=couchbase_config,
|
||||
redis_config=redis_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler
|
||||
)
|
||||
|
||||
assert isinstance(activities, Activities)
|
||||
assert isinstance(activities, Couchbase)
|
||||
assert isinstance(activities, TemporalManager)
|
||||
assert isinstance(activities, SlotManager)
|
||||
|
||||
mock_slot_manager_init.assert_called_once_with(
|
||||
ANY,
|
||||
host="localhost",
|
||||
port=6379,
|
||||
username="admin",
|
||||
password="password",
|
||||
logger=logger,
|
||||
notification_handler=notification_handler
|
||||
)
|
||||
|
||||
mock_couchbase_init.assert_called_once_with(
|
||||
ANY,
|
||||
connection_string=couchbase_config['connection_string'],
|
||||
username=couchbase_config['username'],
|
||||
password=couchbase_config['password'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler
|
||||
)
|
||||
|
||||
mock_temporal_manager_init.assert_called_once_with(
|
||||
ANY,
|
||||
temporal_client=temporal_client,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.activities.couchbase.Cluster')
|
||||
async def test_prepare_activity(_mock_cluster):
|
||||
couchbase_config = {
|
||||
'connection_string': 'couchbase://localhost',
|
||||
'username': 'admin',
|
||||
'password': 'password'
|
||||
}
|
||||
|
||||
redis_config = {
|
||||
'host': 'localhost',
|
||||
'port': 6379,
|
||||
'username': 'admin',
|
||||
'password': 'password'
|
||||
}
|
||||
|
||||
temporal_client = MagicMock()
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
|
||||
activities = Activities(
|
||||
temporal_client=temporal_client,
|
||||
couchbase_config=couchbase_config,
|
||||
redis_config=redis_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler
|
||||
)
|
||||
|
||||
input_data = {
|
||||
'workflow_name': 'test-workflow-name',
|
||||
'schedule_name': 'test-schedule-name',
|
||||
'model_name': 'test-model-name',
|
||||
'model_id': 'test-model-id'
|
||||
}
|
||||
|
||||
await activities.prepare_activity(input_data)
|
||||
|
||||
assert activities.notification_handler.base_notification.pipeline == input_data[
|
||||
'workflow_name']
|
||||
assert activities.notification_handler.base_notification.trigger == input_data[
|
||||
'schedule_name']
|
||||
assert activities.notification_handler.base_notification.model_name == input_data[
|
||||
'model_name']
|
||||
assert activities.notification_handler.base_notification.model_id == input_data[
|
||||
'model_id']
|
||||
56
tests/orchestrator/activities/test_couchbase.py
Normal file
56
tests/orchestrator/activities/test_couchbase.py
Normal file
@@ -0,0 +1,56 @@
|
||||
from unittest.mock import MagicMock, patch, ANY
|
||||
from pytest import fixture, mark, raises
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from orchestrator.activities.couchbase import Couchbase
|
||||
|
||||
|
||||
@fixture
|
||||
@patch("orchestrator.activities.couchbase.Cluster")
|
||||
def couchbase(_cluster_mock):
|
||||
return Couchbase(
|
||||
connection_string="couchbase://localhost",
|
||||
username="admin",
|
||||
password="password",
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_load_query_from_couchbase_success(couchbase):
|
||||
couchbase.cluster.query.return_value.rows.return_value = [
|
||||
{"id": "1", "name": "test"},
|
||||
{"id": "2", "name": "test2"},
|
||||
]
|
||||
query = "SELECT * FROM bucket"
|
||||
|
||||
result = await couchbase.load_query_from_couchbase({
|
||||
"query": query,
|
||||
})
|
||||
|
||||
assert result == [
|
||||
{"id": "1", "name": "test"},
|
||||
{"id": "2", "name": "test2"},
|
||||
]
|
||||
couchbase.cluster.query.assert_called_once_with(query)
|
||||
couchbase.notification_handler.build_and_send_notification.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_load_query_from_couchbase_failure(couchbase):
|
||||
couchbase.cluster.query.side_effect = Exception("Test error")
|
||||
query = "SELECT * FROM bucket"
|
||||
|
||||
with raises(Exception):
|
||||
await couchbase.load_query_from_couchbase({
|
||||
"query": query,
|
||||
})
|
||||
|
||||
couchbase.cluster.query.assert_called_once_with(query)
|
||||
couchbase.notification_handler.build_and_send_notification.assert_called_once_with(
|
||||
notification_id="COUCHBASE_LOAD_QUERY_ERROR",
|
||||
message="Failed to execute couchbase query: Test error",
|
||||
block="load_query_from_couchbase",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
57
tests/orchestrator/activities/test_slot_manager.py
Normal file
57
tests/orchestrator/activities/test_slot_manager.py
Normal file
@@ -0,0 +1,57 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
from pytest import mark, fixture
|
||||
from orchestrator.activities.slot_manager import SlotManager
|
||||
|
||||
|
||||
@fixture
|
||||
@patch("orchestrator.activities.slot_manager.Redis.__init__")
|
||||
def slot_manager(_redis_mock):
|
||||
|
||||
slot_manager = SlotManager(
|
||||
host="localhost",
|
||||
port=6379,
|
||||
username="admin",
|
||||
password="password",
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock()
|
||||
)
|
||||
|
||||
slot_manager.redis_client = MagicMock()
|
||||
slot_manager.logger = MagicMock()
|
||||
slot_manager.notification_handler = MagicMock()
|
||||
|
||||
return slot_manager
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_load_opc_slots_no_slot_keys(slot_manager):
|
||||
slot_manager.redis_client.keys.return_value = []
|
||||
assert await slot_manager.load_opc_slots() == {}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_load_opc_slots(slot_manager):
|
||||
slot_manager.redis_client.keys.return_value = [
|
||||
b"slot:opc_tags:1", b"slot:opc_tags:2", b"slot:opc_tags:3"]
|
||||
|
||||
slot_manager.redis_client.mget.return_value = [
|
||||
b"value1", "value2", None]
|
||||
|
||||
response = await slot_manager.load_opc_slots()
|
||||
|
||||
assert response == {
|
||||
"slot:opc_tags:1": "value1",
|
||||
"slot:opc_tags:2": "value2",
|
||||
"slot:opc_tags:3": None
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_load_active_ingestors(slot_manager):
|
||||
slot_manager.redis_client.keys.return_value = [
|
||||
b"heartbeat:ingestor:1", b"heartbeat:ingestor:2", b"heartbeat:ingestor:3"]
|
||||
|
||||
response = await slot_manager.load_active_ingestors()
|
||||
|
||||
assert response == ["heartbeat:ingestor:1",
|
||||
"heartbeat:ingestor:2", "heartbeat:ingestor:3"]
|
||||
80
tests/orchestrator/activities/test_temporal_manager.py
Normal file
80
tests/orchestrator/activities/test_temporal_manager.py
Normal file
@@ -0,0 +1,80 @@
|
||||
from unittest.mock import MagicMock, patch, AsyncMock
|
||||
import base64
|
||||
import json
|
||||
from pytest import fixture, mark
|
||||
from orchestrator.activities.temporal_manager import TemporalManager
|
||||
|
||||
|
||||
@fixture
|
||||
def temporal_manager():
|
||||
return TemporalManager(
|
||||
temporal_client=MagicMock(),
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock()
|
||||
)
|
||||
|
||||
|
||||
@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"]
|
||||
}
|
||||
)
|
||||
|
||||
temporal_manager.temporal_client.list_schedules = AsyncMock(
|
||||
return_value=async_iter())
|
||||
temporal_manager.temporal_client.get_schedule.return_value = MagicMock(
|
||||
describe=AsyncMock(
|
||||
return_value=MagicMock(
|
||||
schedule=MagicMock(
|
||||
action=MagicMock(
|
||||
args=[
|
||||
MagicMock(
|
||||
data=base64.b64encode(json.dumps(
|
||||
{"test": "test"}).encode('utf-8'))
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
temporal_manager.temporal_client.get_schedule.return_value.describe \
|
||||
.return_value.schedule.spec = MagicMock(
|
||||
intervals=[
|
||||
MagicMock(
|
||||
every=MagicMock(
|
||||
seconds=60
|
||||
)
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
response = await temporal_manager.load_schedule()
|
||||
|
||||
temporal_manager.temporal_client.list_schedules.assert_called_once()
|
||||
assert response == {
|
||||
"test-schedule-id": {
|
||||
"frequency": 60,
|
||||
"data": {"test": "test"},
|
||||
"handle": temporal_manager.temporal_client.get_schedule.return_value
|
||||
}
|
||||
}
|
||||
81
tests/orchestrator/workflows/test_orchestrator.py
Normal file
81
tests/orchestrator/workflows/test_orchestrator.py
Normal file
@@ -0,0 +1,81 @@
|
||||
from unittest.mock import AsyncMock, patch, ANY, call
|
||||
from pytest import fixture, mark
|
||||
from orchestrator.workflows.orchestrator import Orchestrator
|
||||
from orchestrator.activities.activities import Activities
|
||||
|
||||
|
||||
@fixture
|
||||
def orchestrator():
|
||||
return Orchestrator()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch("orchestrator.workflows.orchestrator.workflow", new_callable=AsyncMock)
|
||||
async def test_run(workflow_mock, orchestrator):
|
||||
input_data = {
|
||||
"pipelines_query": "SELECT * FROM bucket",
|
||||
"opc_servers_query": "SELECT * FROM servers",
|
||||
"schedule_name": "test-schedule-name",
|
||||
}
|
||||
|
||||
await orchestrator.run(input_data)
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.prepare_activity,
|
||||
{
|
||||
"workflow_name": "orchestrator",
|
||||
"schedule_name": "test-schedule-name",
|
||||
"model_name": "-",
|
||||
"model_id": "-"
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.load_query_from_couchbase,
|
||||
{
|
||||
"query": input_data["pipelines_query"]
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.load_query_from_couchbase,
|
||||
{
|
||||
"query": input_data["opc_servers_query"]
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.load_schedule,
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.load_opc_slots,
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.load_active_ingestors,
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
Reference in New Issue
Block a user