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

SIENTIAPDE-1030

Browse Source 
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
2025-05-30 16:24:52 -03:00
parent f6584314b2
commit 4f37b2be02
59 changed files with 1827 additions and 4377 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
tests/laborious/activities/__init__.py
View File

193
tests/laborious/activities/test_activities.py
View File

@@ -1,193 +0,0 @@
from pytest import mark
from unittest.mock import patch, MagicMock, ANY
from laborious.activities.activities import Activities
from laborious.activities.postgres import Postgres
from laborious.activities.mlflow import MLFlow
from laborious.activities.gates import Gates
from laborious.activities.opc import OPC
@patch('laborious.activities.activities.Postgres.__init__')
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
@patch('laborious.activities.activities.Gates.__init__')
def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgres_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
'user': 'postgres',
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
}
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group'
}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
opc_config=opc_config,
logger=logger,
notification_handler=notification_handler
)
assert isinstance(activities, Activities)
assert isinstance(activities, Postgres)
assert isinstance(activities, MLFlow)
assert isinstance(activities, OPC)
assert isinstance(activities, Gates)
mock_postgres_init.assert_called_once_with(
ANY,
host=postgres_config['host'],
port=postgres_config['port'],
user=postgres_config['user'],
password=postgres_config['password'],
dbname=postgres_config['dbname'],
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
logger=logger,
notification_handler=notification_handler
)
mock_mlflow_init.assert_called_once_with(
ANY,
mlflow_host=mlflow_config['host'],
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
logger=logger,
notification_handler=notification_handler
)
mock_opc_init.assert_called_once_with(
ANY,
opc_servers=opc_config,
logger=logger,
notification_handler=notification_handler
)
mock_gates_init.assert_called_once_with(
ANY,
logger=logger,
notification_handler=notification_handler
)
@mark.asyncio
@patch('laborious.activities.activities.Postgres.__init__')
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
async def test_prepare_activity(_mock_opc_init,
_mock_mlflow_init, _mock_postgres_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
'user': 'postgres',
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
}
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group'
}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
opc_config=opc_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_name == input_data[
'workflow_name']
assert activities.notification_handler.base_notification.schedule_name == 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']
@patch('laborious.activities.activities.Postgres', return_value=MagicMock())
@patch('laborious.activities.activities.MLFlow', return_value=MagicMock())
@patch('laborious.activities.activities.OPC', return_value=MagicMock())
def test_shutdown(mock_opc_init,
_mock_mlflow_init, mock_postgres_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
'user': 'postgres',
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
}
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group'
}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
opc_config=opc_config,
logger=logger,
notification_handler=notification_handler
)
activities.shutdown()
mock_opc_init.shutdown.assert_called_once()
mock_postgres_init.close.assert_called_once()

35
tests/laborious/activities/test_base.py
View File

@@ -1,35 +0,0 @@
from unittest.mock import MagicMock
from laborious.activities.base import BaseActivity
from pytest import fixture, mark
from sientia_do.notifications.models import Notification
@fixture
def base_activity():
return BaseActivity(
logger=MagicMock(),
notification_handler=MagicMock(),
)
@mark.asyncio
async def test_prepare_activity(base_activity):
base_activity.notification_handler.base_notification = Notification(
project="project",
pipeline="pipeline",
trigger="-",
model_name="-",
model_id="-",
)
await base_activity.prepare_activity({
'workflow_name': 'test_workflow',
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id'
})
assert base_activity.notification_handler.base_notification.schedule_name == "test_schedule"
assert base_activity.notification_handler.base_notification.model_name == "test_model"
assert base_activity.notification_handler.base_notification.model_id == "test_model_id"
assert base_activity.notification_handler.base_notification.pipeline_name == "test_workflow"

369
tests/laborious/activities/test_gates.py
View File

@@ -1,369 +0,0 @@
from unittest.mock import MagicMock, ANY, patch
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from laborious.activities.gates import Gates
@fixture
def gates_activity():
return Gates(
logger=MagicMock(),
notification_handler=MagicMock(),
)
@mark.asyncio
async def test_input_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'data': {'value': [1, 2, 3]},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
gates_activity.logger.error.assert_called_once_with(
"Filter INVALID_FILTER not found"
)
@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_filter_exception(mock_input_filter_functions, gates_activity):
# Arrange
mock_input_filter_functions.__contains__.return_value = True
mock_input_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception("Test error"))
input_data = {
'filters': {
'EMPTY_DATA': {'POLICY': 'STOP'}
},
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="INTPUT_GATE_ERROR__EMPTY_DATA",
message="Error in filter EMPTY_DATA:{'POLICY': 'STOP'}: \n Test error",
block="input_gate",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
@mark.asyncio
async def test_input_gate_no_filters(gates_activity):
# Arrange
input_data = {
'filters': {},
'data': {'value': [1, 2, 3]},
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
gates_activity.logger.debug.assert_called()
@mark.asyncio
async def test_input_gate_with_filter(gates_activity):
# Arrange
input_data = {
'filters': {
'EMPTY_DATA': {'POLICY': 'STOP'}
},
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == ('STOP', -1, "Input data with bad quality")
gates_activity.logger.debug.assert_called()
@mark.asyncio
async def test_mlflow_response_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
@mark.asyncio
@patch('laborious.activities.gates.mlflow_response_filter_functions')
async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter_functions,
gates_activity):
# Arrange
mock_mlflow_response_filter_functions.__contains__.return_value = True
mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception("Test error"))
input_data = {
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER",
message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error",
block="mlflow_gate",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
@mark.asyncio
async def test_mlflow_response_gate_no_filters(gates_activity):
# Arrange
input_data = {
'filters': {},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
gates_activity.logger.debug.assert_called()
@mark.asyncio
async def test_mlflow_response_gate_with_filter(gates_activity):
# Arrange
input_data = {
'filters': {
'API_ERROR': {'POLICY': 'STOP'}
},
'data': {
'success': False,
'content': {
'message': 'API error occurred',
'traceback': 'error trace'
}
},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == ('STOP', -1, "API error occurred")
gates_activity.logger.debug.assert_called()
gates_activity.notification_handler.build_and_send_notification.assert_called()
@mark.asyncio
async def test_mlflow_content_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
@mark.asyncio
@patch('laborious.activities.gates.mlflow_content_filter_functions')
async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_functions,
gates_activity):
# Arrange
mock_mlflow_content_filter_functions.__contains__.return_value = True
mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception("Test error"))
input_data = {
'filters': {
'API_ERROR': {'POLICY': 'STOP'}
},
'data': {
'success': False,
'content': {
'message': 'API error occurred',
'traceback': 'error trace'
}
},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
gates_activity.logger.debug.assert_called()
gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="MLFLOW_GATE_CONTENT_FILTER__API_ERROR",
message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error",
block="mlflow_gate",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
@mark.asyncio
async def test_mlflow_content_gate_no_filters(gates_activity):
# Arrange
input_data = {
'filters': {},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
gates_activity.logger.debug.assert_called()
@mark.asyncio
async def test_mlflow_content_gate_with_filter(gates_activity):
# Arrange
input_data = {
'filters': {
'NAN_VALUES': {'POLICY': 'STOP'}
},
'data': {'value': [None, None, None]},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (
'STOP', -1, "Transformed data not passed the content filter")
gates_activity.logger.debug.assert_called()
gates_activity.notification_handler.build_and_send_notification.assert_called()
@mark.asyncio
async def test_format_prediction(gates_activity):
# Arrange
input_data = {
'data': {'prediction': [1], 'response_time': [0.1]},
'timestamp': '2023-05-26 11:12:27',
'model_id': 'test_model',
'prediction_confidence': 0.9
}
# Act
result = await gates_activity.format_prediction(input_data)
# Assert
assert result['prediction'] == {0: 1}
assert result['response_time'] == {0: ANY}
assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
assert result['model_id'] == {0: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9}
assert result['prediction_status'] == {0: 'Good'}
assert result['comments'] == {0: ""}
gates_activity.logger.debug.assert_called()
@mark.asyncio
async def test_format_default_prediction(gates_activity):
# Arrange
input_data = {
'timestamp': '2023-05-26 11:12:27',
'model_id': 'test_model',
'prediction_confidence': 0.1,
'comment': 'Test comment'
}
# Act
result = await gates_activity.format_default_prediction(input_data)
# Assert
assert result['prediction'] == {0: 0}
assert result['response_time'] == {0: 0}
assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
assert result['model_id'] == {0: 'test_model'}
assert result['prediction_confidence'] == {0: 0.1}
assert result['prediction_status'] == {0: 'Bad'}
assert result['comments'] == {0: 'Test comment'}
gates_activity.logger.debug.assert_called()
@mark.asyncio
async def test_get_last_timestamp_with_data(gates_activity):
# Arrange
input_data = {
'data': {
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']
}
}
# Act
result = await gates_activity.get_last_timestamp(input_data)
# Assert
assert result == '2023-05-26 11:12:28'
@mark.asyncio
async def test_get_last_timestamp_no_data(gates_activity):
# Arrange
input_data = {
'data': {}
}
# Act
result = await gates_activity.get_last_timestamp(input_data)
# Assert
assert isinstance(result, str) # Should be a timestamp string
assert len(result) > 0

120
tests/laborious/activities/test_mlflow.py
View File

@@ -1,120 +0,0 @@
from unittest.mock import MagicMock, patch
import numpy as np
from pytest import fixture, mark
from laborious.activities.mlflow import MLFlow
@patch("laborious.activities.mlflow.MLFlowRepository")
def test___init__(mock_mlflow_repository):
mlflow = MLFlow(
mlflow_host="http://localhost",
mlflow_port=5000,
mlflow_username="admin",
mlflow_password="admin",
logger=MagicMock(),
notification_handler=MagicMock()
)
assert mlflow.mlflow_host == "http://localhost"
assert mlflow.mlflow_port == 5000
assert mlflow.mlflow_username == "admin"
assert mlflow.mlflow_password == "admin"
mock_mlflow_repository.assert_called_once_with(
"http://localhost:5000", "admin", "admin"
)
@fixture
@patch("laborious.activities.mlflow.MLFlowRepository")
def mlflow(mock_mlflow_repository):
return MLFlow(
mlflow_host="http://localhost:5000",
mlflow_port=5000,
mlflow_username="admin",
mlflow_password="admin",
logger=MagicMock(),
notification_handler=MagicMock()
)
@mark.asyncio
@patch("laborious.activities.mlflow.DataFrame")
@patch("laborious.activities.mlflow.max")
async def test_request_transform(mock_max, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02'
# Mock input data
input_data = {
'data': [
{'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0},
{'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0},
{'timestamp': '2024-01-02', 'variable': 'var1', 'value': 3.0},
{'timestamp': '2024-01-02', 'variable': 'var2', 'value': 4.0}
],
'model_name': 'test_model',
'model_retention': 30
}
# Mock the transform response
expected_response = {'prediction': [0.5, 0.6]}
mlflow.model_monitoring_repository.transform.return_value = expected_response
# Call the method
response_data = await mlflow.request_transform(input_data)
# Verify the data was correctly transformed
mock_dataframe.assert_called_once_with(input_data['data'])
mock_dataframe.return_value.pivot.assert_called_once_with(
index='timestamp', columns='variable', values='value'
)
mock_dataframe = mock_dataframe.return_value.pivot.return_value
mock_dataframe.fillna.assert_called_once_with(np.nan, inplace=True)
mock_dataframe.reset_index.assert_called_once()
mock_dataframe.columns.name = None
# Verify the response
assert response_data == expected_response
# Verify the repository was called with correct arguments
mlflow.model_monitoring_repository.transform.assert_called_once_with(
'test_model', mock_dataframe, 30
)
@mark.asyncio
@patch("laborious.activities.mlflow.DataFrame")
@patch("laborious.activities.mlflow.max")
async def test_request_predict(mock_max, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02'
# Mock input data
input_data = {
'data': [
{'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0},
{'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0},
{'timestamp': '2024-01-02', 'variable': 'var1', 'value': 3.0},
{'timestamp': '2024-01-02', 'variable': 'var2', 'value': 4.0}
],
'model_name': 'test_model',
'model_retention': 30
}
# Mock the predict response
expected_response = {'prediction': [0.5, 0.6]}
mlflow.model_monitoring_repository.predict.return_value = expected_response
# Call the method
response_data = await mlflow.request_predict(input_data)
mock_dataframe.assert_called_once_with(input_data['data'])
mock_dataframe.return_value.replace.assert_called_once_with(
np.nan, None, inplace=True
)
# Verify the response
assert response_data == expected_response
# Verify the repository was called with correct arguments
mlflow.model_monitoring_repository.predict.assert_called_once_with(
'test_model', mock_dataframe.return_value, 30
)

196
tests/laborious/activities/test_opc.py
View File

@@ -1,196 +0,0 @@
from unittest.mock import patch, MagicMock, ANY, call
from pytest import fixture, mark
from laborious.activities.opc import NotificationLevel
from laborious.activities.opc import OPC
@patch("laborious.activities.opc.OpcRepository")
def test___init__(mock_opc_repository):
mock_logger = MagicMock()
server1 = MagicMock()
server2 = MagicMock()
mock_opc_repository.side_effect = [server1, server2]
mock_notification_handler = MagicMock()
servers = {
'server1': {
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
},
'server2': {
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
}
}
opc = OPC(
opc_servers=servers,
logger=mock_logger,
notification_handler=mock_notification_handler
)
assert opc.opc_servers == servers
assert opc.logger == mock_logger
assert opc.notification_handler == mock_notification_handler
assert opc.opc_repository['server1'] == server1
assert opc.opc_repository['server2'] == server2
mock_opc_repository.assert_has_calls([
call(
name="server1",
url="http://localhost:8080",
logger=mock_logger,
server_uri="opc.tcp://localhost:4840",
cert_path="",
private_key_path="",
server_cert_path="",
notification_handler=mock_notification_handler,
reconnection_interval=60,
),
])
mock_opc_repository.assert_has_calls([
call(
name="server2",
url="http://localhost:8080",
logger=mock_logger,
server_uri="opc.tcp://localhost:4840",
cert_path="",
private_key_path="",
server_cert_path="",
notification_handler=mock_notification_handler,
reconnection_interval=60,
)
])
server1.connect.assert_called_once()
server2.connect.assert_called_once()
@fixture
@patch("laborious.activities.opc.OpcRepository")
def opc(_mock_opc_repository):
servers = {
'server1': {
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
}
}
return OPC(
opc_servers=servers,
logger=MagicMock(),
notification_handler=MagicMock()
)
WRITE_DATA_CASES = [
('tag1', 'int', 50),
('tag2', 'float', 50.5),
('tag3', 'bool', True),
('tag4', 'string', 'test'),
]
@mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
def test_write_data_success(opc, tag, data_type, data):
opc.write_data(server='server1', tag=tag, data=data,
data_type=data_type, tag_type='prediction')
opc.opc_repository['server1'].write_data.assert_called_once_with(
tag, data, data_type)
def test_write_data_exception(opc):
opc.opc_repository['server1'].write_data.side_effect = Exception(
"Test error")
opc.write_data(server='server1', tag='tag1', data=50,
data_type='int', tag_type='prediction')
opc.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="WRITE_OPC_PREDICTION_ERROR",
message="Error writing data to OPC server: Test error",
block="write_opc_data",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
opc.logger.error.assert_called_once()
@mark.asyncio
async def test_write_opc_data_success(opc):
# Arrange
input_data = {
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
},
'opc_output_config': {
'server1': {
'prediction_tags': {
'tag1': {'data_type': 'float'}
},
'confidence_tags': {
'tag2': {'data_type': 'float'}
}
}
}
}
# Act
opc.write_data = MagicMock()
await opc.write_opc_data(input_data)
# Assert
opc.write_data.assert_has_calls([
call(
server='server1',
tag='tag1',
data=0.75,
data_type='float',
tag_type='prediction'
)])
opc.write_data.assert_has_calls([
call(
server='server1',
tag='tag2',
data=0.95,
data_type='float',
tag_type='confidence'
)
])
assert opc.write_data.call_count == 2
@mark.asyncio
async def test_write_opc_data_empty_config(opc):
# Arrange
input_data = {
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
},
'opc_servers': ['server1'],
'opc_output_config': {
'prediction_tags': {},
'confidence_tags': {}
}
}
# Act
await opc.write_opc_data(input_data)
# Assert
opc.opc_repository['server1'].write_data.assert_not_called()
def test_shutdown(opc):
opc.shutdown()
opc.opc_repository['server1'].disconnect.assert_called_once()

159
tests/laborious/activities/test_postgres.py
View File

@@ -1,159 +0,0 @@
from unittest.mock import MagicMock, patch
from pytest import fixture, mark
import pandas as pd
from laborious.activities.postgres import Postgres
@fixture
@patch("laborious.activities.postgres.create_engine")
def postgres_activity(_mock_create_engine):
return Postgres(
host="localhost",
port=5432,
user="test_user",
password="test_password",
dbname="test_db",
min_connections=1,
max_connections=5,
logger=MagicMock(),
notification_handler=MagicMock()
)
@mark.asyncio
@patch("laborious.activities.postgres.read_sql_query")
async def test_load_custom_query_none_data(mock_read_sql_query, postgres_activity):
query = "SELECT * FROM test_table LIMIT 1"
mock_read_sql_query.return_value = None
result = await postgres_activity.load_custom_query(query)
assert isinstance(result, dict)
assert len(result) == 0
@mark.asyncio
@patch("laborious.activities.postgres.read_sql_query")
async def test_load_custom_query_date_converted(mock_read_sql_query, postgres_activity):
query = "SELECT * FROM test_table LIMIT 1"
mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]})
mock_data['date'] = pd.to_datetime('2022-01-01')
mock_read_sql_query.return_value = mock_data
result = await postgres_activity.load_custom_query(query)
assert isinstance(result, dict)
assert len(result) == 3
assert "column1" in result
assert "column2" in result
assert "date" in result
assert result['date'] == {0: '2022-01-01 00:00:00'}
@mark.asyncio
@patch("laborious.activities.postgres.read_sql_query")
async def test_load_custom_query_success(mock_read_sql_query, postgres_activity):
query = "SELECT * FROM test_table LIMIT 1"
mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]})
mock_read_sql_query.return_value = mock_data
result = await postgres_activity.load_custom_query(query)
assert isinstance(result, dict)
assert len(result) == 2
assert "column1" in result
assert "column2" in result
postgres_activity.logger.info.assert_called()
@mark.asyncio
async def test_load_custom_query_error(postgres_activity):
query = "SELECT * FROM non_existent_table"
error_msg = "Table not found"
with patch("laborious.activities.postgres.read_sql_query", side_effect=ValueError(error_msg)):
result = await postgres_activity.load_custom_query(query)
assert isinstance(result, dict)
assert len(result) == 0
postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
postgres_activity.logger.error.assert_called()
@mark.asyncio
async def test_repeat_last_prediction_success(postgres_activity):
query_items = {
"schema": "public",
"table_name": "predictions",
"model": 1
}
with patch("sqlalchemy.orm.session.Session.execute") as mock_execute:
await postgres_activity.repeat_last_prediction(query_items)
mock_execute.assert_called_once()
postgres_activity.logger.info.assert_called()
@mark.asyncio
async def test_repeat_last_prediction_error(postgres_activity):
query_items = {
"schema": "public",
"table_name": "predictions",
"model": 1
}
error_msg = "Database error"
with patch("sqlalchemy.orm.session.Session.execute", side_effect=ValueError(error_msg)):
await postgres_activity.repeat_last_prediction(query_items)
postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
postgres_activity.logger.error.assert_called()
@mark.asyncio
async def test_export_data_to_postgres_success(postgres_activity):
input_data = {
"schema": "public",
"table_name": "test_table",
"data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]})
}
with patch("laborious.activities.postgres.DataFrame.to_sql") as mock_to_sql:
await postgres_activity.export_data_to_postgres(input_data)
mock_to_sql.assert_called_once()
postgres_activity.logger.debug.assert_called()
@mark.asyncio
async def test_export_data_to_postgres_error(postgres_activity):
input_data = {
"schema": "public",
"table_name": "test_table",
"data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]})
}
error_msg = "Export failed"
with patch("laborious.activities.postgres.DataFrame.to_sql", side_effect=ValueError(error_msg)):
await postgres_activity.export_data_to_postgres(input_data)
postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
postgres_activity.logger.error.assert_called()
@mark.asyncio
async def test_close(postgres_activity):
postgres_activity.close()
postgres_activity.engine.dispose.assert_called_once()
@mark.asyncio
async def test_del(postgres_activity):
postgres_activity.close = MagicMock()
postgres_activity.__del__()
postgres_activity.close.assert_called_once()

0
tests/laborious/utils/__init__.py
View File

0
tests/laborious/utils/filters/__init__.py
View File

30
tests/laborious/utils/filters/test_conditional_filters.py
View File

@@ -1,30 +0,0 @@
from pandas import DataFrame
from laborious.utils.filters.conditional_filters import (
filter_specific_variables_null_values,
filter_empty_data
)
def test_filter_specific_variables_null_values():
assert filter_specific_variables_null_values(
DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
config={'VARIABLES': ['variable2']}) is False
def test_filter_specific_variables_null_values_with_null_values():
assert filter_specific_variables_null_values(
DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, None]}),
config={'VARIABLES': ['variable2']}) is True
def test_filter_empty_data():
assert filter_empty_data(DataFrame(), {}) is True
def test_filter_empty_data_with_data():
assert filter_empty_data(
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
{}) is False

22
tests/laborious/utils/filters/test_mlflow_filters.py
View File

@@ -1,22 +0,0 @@
from pandas import DataFrame
from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
def test_api_error_filter_invalid_response():
assert api_error_filter(None, {}) == True # NOSONAR
def test_api_error_filter_valid_response_fail():
assert api_error_filter({'success': False}, {}) == True
def test_api_error_filter_valid_response_success():
assert api_error_filter({'success': True}, {}) == False
def test_nan_values_filter_all_nan_values():
assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) == True
def test_nan_values_filter_no_nan_values():
assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) == False

278
tests/laborious/utils/repository/test_model_repository.py
View File

@@ -1,278 +0,0 @@
from unittest.mock import ANY, MagicMock, patch
import numpy as np
from pandas import DataFrame
import pytest
from laborious.utils.repository.model_repository import MLFlowRepository
@pytest.fixture
def mlflow_repository():
with patch('laborious.utils.repository.model_repository.ModelServing', autospec=True) as MockModelServing:
mock_instance = MockModelServing.return_value
mock_instance.get_transformed_data = MagicMock()
repo = MLFlowRepository(
host='http://localhost:5000',
username='admin',
password='admin'
)
return repo
def test_get_current_data_df(mlflow_repository):
current_data = {
'prediction': [1, 3],
'target': [1, 1],
}
mlflow_repository.model_serving.get_transformed_data.return_value = {
'var1': [1, 2],
'var2': [2, np.nan],
}
expected = DataFrame({
'var1': [1],
'var2': [2],
'prediction': [1],
'target': [1],
})
output = mlflow_repository.get_current_data_df(current_data,
'model', 'target')
mlflow_repository.model_serving.get_transformed_data.assert_called_once_with(
'model', current_data, by='model')
diff = output.compare(expected)
assert diff.empty
def test_get_artifact(mlflow_repository):
mlflow_repository.get_artifact(
'destination', 'search_by', 'run_id', 'model', 'artifact'
)
mlflow_repository.model_serving.get_artifact.assert_called_once_with(
destination='destination',
search_by='search_by',
run_id='run_id',
model_name='model',
artifact_name='artifact'
)
def test_calculate_model_metrics(mlflow_repository):
mlflow_repository.model_serving.get_model_metrics.return_value = 'data'
real_data = 'real_data'
predictions = 'predictions'
flag = 'flag'
output = mlflow_repository.calculate_model_metrics(
real_data, predictions, flag
)
mlflow_repository.model_serving.get_model_metrics.assert_called_once_with(
reference_data=None,
real_data=real_data,
predictions=predictions,
type_flag=flag
)
assert output == 'data'
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_by_run_id(mlflow, mlflow_repository):
mlflow.get_run.return_value = MagicMock(
info=MagicMock(
experiment_id='0',
)
)
mlflow.get_experiment.return_value = MagicMock()
mlflow.get_experiment.return_value.name = 'test'
output = mlflow_repository.get_experiment_by_run_id('0')
assert output == 'test'
mlflow.get_run.assert_called_once_with('0')
mlflow.get_experiment.assert_called_once_with('0')
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_next_run_name(mlflow, mlflow_repository):
mlflow.search_runs.return_value = [1, 2, 3]
output = mlflow_repository.get_next_run_name('run')
assert output == 'run-4'
mlflow.search_runs.assert_called_once_with(
experiment_names=['run'],
order_by=['start_time desc'],
)
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_success(mlflow, mlflow_repository):
mlflow.get_experiment_by_name.return_value = MagicMock(
experiment_id='0')
output = mlflow_repository.get_experiment('test')
assert output == 0
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_error(mlflow, mlflow_repository):
mlflow.get_experiment_by_name.return_value = None
try:
mlflow_repository.get_experiment('test')
except ValueError as e:
assert str(e) == 'Experiment test not found'
else:
assert False
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_last_run(mlflow, mlflow_repository):
mlflow.search_runs.return_value = DataFrame({
'params.retrain': ['True', 'False', 'True', 'False'],
'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'],
'run_id': ['0', '1', '2', '3'],
})
output = mlflow_repository.get_experiment_last_run(0)
mlflow.search_runs.assert_called_once_with(
experiment_ids=[0],
filter_string="",
output_format="pandas",
)
assert output == '2'
@patch('laborious.utils.repository.model_repository.mlflow')
def test_update_production_model_by_run_id(mlflow, mlflow_repository):
client_mock = MagicMock()
mlflow.tracking.MlflowClient.return_value = client_mock
client_mock.get_registered_model.return_value = MagicMock(
latest_versions=[
MagicMock(version='1'),
MagicMock(version='2'),
MagicMock(version='3'),
]
)
output = mlflow_repository.update_production_model_by_run_id('0', 'test')
mlflow.register_model.assert_called_once_with(
"runs:/0/prediction_model",
'test',
)
mlflow.tracking.MlflowClient.assert_called_once()
client_mock.get_registered_model.assert_called_once_with('test')
client_mock.transition_model_version_stage.assert_called_once_with(
name='test',
version='3',
stage='Production',
archive_existing_versions=True,
)
assert output == {
'model_name': 'test',
'version': '3',
'mlflow_run_id': '0',
}
def test_update_production_model(mlflow_repository):
connector = mlflow_repository
with patch.object(connector, 'get_experiment',
return_value='0') as get_experiment:
with patch.object(connector, 'get_experiment_last_run',
return_value='2') as get_experiment_last_run:
with patch.object(connector, 'update_production_model_by_run_id',
return_value={'model_name': 'test', 'version': '3',
'mlflow_run_id': '0'}) as update_production_model_by_run_id:
output = connector.update_production_model('0', 'test')
get_experiment.assert_called_once_with('0')
get_experiment_last_run.assert_called_once_with('0')
update_production_model_by_run_id.assert_called_once_with(
'2', 'test')
assert output == {
'model_name': 'test',
'version': '3',
'mlflow_run_id': '0',
'mlflow_experiment_id': '0',
}
def test_transform_success(mlflow_repository):
data = 'data'
model_name = 'model'
output = mlflow_repository.transform(model_name, data, 1)
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
model_name, data, 1)
assert output == {
'success': True,
'content': mlflow_repository.model_serving.get_cached_transform.return_value.to_dict.return_value
}
def test_transform_error(mlflow_repository):
data = 'data'
model_name = 'model'
mlflow_repository.model_serving.get_cached_transform.side_effect = Exception(
'error')
output = mlflow_repository.transform(model_name, data, 1)
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
model_name, data, 1)
assert output == {
'success': False,
'content': {
'message': 'error',
'traceback': ANY
}
}
def test_predict_success(mlflow_repository):
data = 'data'
model_name = 'model'
mlflow_repository.model_serving.get_cached_predict.return_value = np.array(
[2, 3]
)
output = mlflow_repository.predict(model_name, data, 1)
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
model_name, data, 1)
assert output['success'] is True
assert output['content'] == {'prediction': {
0: 3}, 'response_time': ANY}
def test_predict_error(mlflow_repository):
data = 'data'
model_name = 'model'
mlflow_repository.model_serving.get_cached_predict = MagicMock(
side_effect=Exception('error')
)
output = mlflow_repository.predict(model_name, data, 1)
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
model_name, data, 1)
assert output == {
'success': False,
'content': {
'message': 'error',
'traceback': ANY
}
}

259
tests/laborious/utils/repository/test_opc_repository.py
View File

@@ -1,259 +0,0 @@
from unittest.mock import Mock, patch, MagicMock, ANY, call
from asyncua.crypto.security_policies import SecurityPolicyBasic256
from pytest import fixture
from laborious.utils.repository.opc_repository import OpcRepository
from sientia_do.notifications.models import NotificationLevel
from datetime import datetime
@fixture
def mock_logger():
return Mock()
@fixture
def opc_repository(mock_logger):
return OpcRepository(
name="test_repo",
url="opc.tcp://localhost:4840",
logger=mock_logger,
notification_handler=Mock(),
reconnection_interval=60,
server_uri="urn:test:server",
cert_path="/path/to/cert.pem",
private_key_path="/path/to/key.pem",
server_cert_path="/path/to/server_cert.pem"
)
@fixture
def mock_client():
with patch('laborious.utils.repository.opc_repository.Client') as mock:
client_instance = MagicMock()
mock.return_value = client_instance
yield client_instance
def test_init(opc_repository):
assert opc_repository.name == "test_repo"
assert opc_repository.url == "opc.tcp://localhost:4840"
assert opc_repository.server_uri == "urn:test:server"
assert opc_repository.cert_path == "/path/to/cert.pem"
assert opc_repository.private_key_path == "/path/to/key.pem"
assert opc_repository.server_cert_path == "/path/to/server_cert.pem"
assert opc_repository.reconnection_interval == 60
assert opc_repository.client is None
assert opc_repository.last_reconnection_time is None
assert opc_repository.error_count == 0
def test_set_security(opc_repository, mock_client):
opc_repository.client = mock_client
opc_repository.set_security()
mock_client.application_uri = "urn:test:server"
mock_client.set_security.assert_called_once_with(
SecurityPolicyBasic256,
certificate="/path/to/cert.pem",
private_key="/path/to/key.pem",
server_certificate="/path/to/server_cert.pem"
)
assert mock_client.secure_channel_timeout == 10000000
assert mock_client.session_timeout == 10000000
def test_set_security_missing_certificates(opc_repository):
opc_repository.cert_path = None
opc_repository.private_key_path = None
try:
opc_repository.set_security()
except ValueError as e:
assert str(
e) == "Certificate and private key paths must be provided for secure connection."
def test_connect_with_security(opc_repository, mock_client):
opc_repository.try_connect = MagicMock()
opc_repository.connect()
opc_repository.try_connect.assert_called_once()
assert opc_repository.client == mock_client
def test_connect_without_security(opc_repository, mock_client):
opc_repository.cert_path = None
opc_repository.try_connect = MagicMock()
opc_repository.set_security = MagicMock()
opc_repository.connect()
opc_repository.try_connect.assert_called_once()
opc_repository.set_security.assert_not_called()
assert opc_repository.client == mock_client
def test_try_connect_sucess(opc_repository):
opc_repository.last_reconnection_time = None
opc_repository.client = MagicMock()
opc_repository.try_connect()
opc_repository.client.connect.assert_called_once()
assert opc_repository.last_reconnection_time is not None
def test_try_connect_fail(opc_repository):
opc_repository.last_reconnection_time = None
opc_repository.client = MagicMock()
opc_repository.client.connect.side_effect = Exception("Test error")
opc_repository.try_connect()
opc_repository.client.connect.assert_called_once()
assert opc_repository.last_reconnection_time is not None
opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id=f"OPC_CONNECTION_ERROR_{opc_repository.name}",
message="Failed to connect to OPC server: Test error",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
def test_disconnect(opc_repository, mock_client):
opc_repository.client = mock_client
opc_repository.disconnect()
mock_client.disconnect.assert_called_once()
assert opc_repository.client is None
def test_validate_connection_none_client(opc_repository):
opc_repository.client = None
opc_repository.connect = MagicMock()
response = opc_repository.validate_connection()
assert response
opc_repository.connect.assert_called_once()
def test_validate_connection_error_count_disconnect_error(opc_repository):
opc_repository.error_count = 6
opc_repository.client = MagicMock()
opc_repository.disconnect = MagicMock(side_effect=Exception("Test error"))
opc_repository.connect = MagicMock()
response = opc_repository.validate_connection()
assert response == opc_repository.connect.return_value
opc_repository.disconnect.assert_called_once()
opc_repository.connect.assert_called_once()
opc_repository.logger.error.assert_has_calls(
[
call("Failed to disconnect from OPC server: Test error"),
]
)
@patch('laborious.utils.repository.opc_repository.hasattr', return_value=True)
@patch('laborious.utils.repository.opc_repository.datetime',
MagicMock(now=MagicMock(return_value=datetime(2025, 1, 1, 0, 0, 0))))
def test_validate_connection_lost_not_time_to_reconect(_mock_datetime, opc_repository):
opc_repository.error_count = 0
opc_repository.client = MagicMock()
opc_repository.client.aio_obj.uaclient.protocol = None
opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0)
opc_repository.try_connect = MagicMock()
response = opc_repository.validate_connection()
opc_repository.try_connect.assert_not_called()
assert response is False
@patch('laborious.utils.repository.opc_repository.hasattr', return_value=True)
@patch('laborious.utils.repository.opc_repository.datetime',
MagicMock(now=MagicMock(return_value=datetime(2025, 1, 1, 1, 0, 0))))
def test_validate_connection_lost_time_to_reconect(_mock_datetime, opc_repository):
opc_repository.error_count = 0
opc_repository.client = MagicMock()
opc_repository.client.aio_obj.uaclient.protocol = None
opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0)
opc_repository.try_connect = MagicMock()
response = opc_repository.validate_connection()
opc_repository.try_connect.assert_called_once()
assert response == opc_repository.try_connect.return_value
def test_validate_connection_failed(opc_repository):
opc_repository.client = MagicMock()
opc_repository.error_count = 0
output = opc_repository.validate_connection()
assert output is True
def test_write_data_validate_connection_do_nothing(opc_repository):
opc_repository.validate_connection = MagicMock(return_value=True)
opc_repository.client = MagicMock()
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
def test_write_data_validate_connection_failed(opc_repository):
opc_repository.validate_connection = MagicMock(return_value=False)
opc_repository.client = MagicMock()
opc_repository.error_count = 0
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_not_called()
def test_write_data_get_node_failed(opc_repository):
opc_repository.validate_connection = MagicMock(return_value=True)
opc_repository.client = MagicMock()
opc_repository.error_count = 0
opc_repository.client.get_node.side_effect = Exception("Test error")
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id=f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.name}",
message="Failed to get node from OPC server: Test error",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
assert opc_repository.error_count == 1
def test_write_data(opc_repository, mock_client):
opc_repository.validate_connection = MagicMock(return_value=True)
opc_repository.client = mock_client
mock_node = MagicMock()
mock_client.get_node.return_value = mock_node
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
mock_node.write_value.assert_called_once()
opc_repository.logger.info.assert_called_once_with(
"Writing 42.0 - <class 'float'> to " + str(mock_node))
def test_write_data_write_value_failed(opc_repository, mock_client):
opc_repository.validate_connection = MagicMock(return_value=True)
opc_repository.client = mock_client
mock_node = MagicMock()
opc_repository.error_count = 0
mock_client.get_node.return_value = mock_node
mock_node.write_value.side_effect = Exception("Test error")
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
opc_repository.validate_connection.assert_called_once()
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
mock_node.write_value.assert_called_once()
opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id=f"OPC_WRITE_DATA_ERROR_{opc_repository.name}",
message="Failed to write data to OPC server: Test error",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
assert opc_repository.error_count == 1

133
tests/laborious/utils/test_connectors_config.py
View File

@@ -1,133 +0,0 @@
from os import environ
from laborious.utils.connectors_config import (build_mlflow_config,
build_opc_config,
build_postgres_config)
def test_build_mlflow_config_with_env_vars():
# Arrange
environ['MLFLOW_HOST'] = 'http://test-host'
environ['MLFLOW_PORT'] = '8080'
environ['MLFLOW_USERNAME'] = 'test-user'
environ['MLFLOW_PASSWORD'] = 'test-pass'
# Act
config = build_mlflow_config()
# Assert
assert config['host'] == 'http://test-host'
assert config['port'] == 8080
assert config['username'] == 'test-user'
assert config['password'] == 'test-pass'
def test_build_mlflow_config_with_defaults():
# Arrange
# Clear any existing env vars
environ.pop('MLFLOW_HOST', None)
environ.pop('MLFLOW_PORT', None)
environ.pop('MLFLOW_USERNAME', None)
environ.pop('MLFLOW_PASSWORD', None)
# Act
config = build_mlflow_config()
# Assert
assert config['host'] == 'http://localhost'
assert config['port'] == 5080
assert config['username'] == 'aignosi'
assert config['password'] == 'aignosi'
def test_build_opc_config_with_env_vars():
# Arrange
environ['OPC_CONFIG'] = '{"opc": {"name": "test-opc", "url": "opc.tcp://test:4840"}}'
# Act
config = build_opc_config()
# Assert
assert config['opc']['name'] == 'test-opc'
assert config['opc']['url'] == 'opc.tcp://test:4840'
def test_build_opc_config_with_individual_env_vars():
# Arrange
environ.pop('OPC_CONFIG', None)
environ['OPC_NAME'] = 'test-name'
environ['OPC_URL'] = 'opc.tcp://test:4840'
environ['OPC_SERVER_URI'] = 'opc.tcp://test:4840'
environ['OPC_RECONNECTION_INTERVAL'] = '300'
# Act
config = build_opc_config()
# Assert
assert config['opc']['name'] == 'test-name'
assert config['opc']['url'] == 'opc.tcp://test:4840'
assert config['opc']['server_uri'] == 'opc.tcp://test:4840'
assert config['opc']['reconnection_interval'] == 300
def test_build_opc_config_with_defaults():
# Arrange
environ.pop('OPC_CONFIG', None)
environ.pop('OPC_NAME', None)
environ.pop('OPC_URL', None)
environ.pop('OPC_SERVER_URI', None)
environ.pop('OPC_RECONNECTION_INTERVAL', None)
# Act
config = build_opc_config()
# Assert
assert config['opc']['name'] == 'opc'
assert config['opc']['url'] == 'opc.tcp://localhost:4840'
assert config['opc']['server_uri'] == 'opc.tcp://localhost:4840'
assert config['opc']['reconnection_interval'] == 120
def test_build_postgres_config_with_env_vars():
# Arrange
environ['POSTGRES_HOST'] = 'test-host'
environ['POSTGRES_PORT'] = '5433'
environ['POSTGRES_USER'] = 'test-user'
environ['POSTGRES_PASSWORD'] = 'test-pass'
environ['POSTGRES_DBNAME'] = 'test-db'
environ['POSTGRES_MIN_CONNECTIONS'] = '10'
environ['POSTGRES_MAX_CONNECTIONS'] = '30'
# Act
config = build_postgres_config()
# Assert
assert config['host'] == 'test-host'
assert config['port'] == 5433
assert config['user'] == 'test-user'
assert config['password'] == 'test-pass'
assert config['dbname'] == 'test-db'
assert config['min_connections'] == 10
assert config['max_connections'] == 30
def test_build_postgres_config_with_defaults():
# Arrange
environ.pop('POSTGRES_HOST', None)
environ.pop('POSTGRES_PORT', None)
environ.pop('POSTGRES_USER', None)
environ.pop('POSTGRES_PASSWORD', None)
environ.pop('POSTGRES_DBNAME', None)
environ.pop('POSTGRES_MIN_CONNECTIONS', None)
environ.pop('POSTGRES_MAX_CONNECTIONS', None)
# Act
config = build_postgres_config()
# Assert
assert config['host'] == 'localhost'
assert config['port'] == 5432
assert config['user'] == 'sientia'
assert config['password'] == 'sientia'
assert config['dbname'] == 'sientia'
assert config['min_connections'] == 5
assert config['max_connections'] == 20

37
tests/laborious/utils/test_logger.py
View File

@@ -1,37 +0,0 @@
import os
from unittest.mock import patch
import logging
import pytest
from laborious.utils.logger import get_logger
@pytest.fixture
def mock_env_vars():
with patch.dict(os.environ, {}, clear=True):
yield
@pytest.mark.usefixtures("mock_env_vars")
@patch('laborious.utils.logger.logging.Formatter')
@patch('laborious.utils.logger.logging.StreamHandler')
def test_get_logger_defaults(mock_stream_handler, mock_formatter):
"""Test logger creation with default settings"""
# Mock the StreamHandler and Formatter
logger = get_logger('test_logger')
# Verify logger settings
assert logger.name == 'test_logger'
assert logger.level == logging.INFO
# Verify handler configuration
mock_stream_handler.return_value.setLevel.assert_called_once_with('INFO')
mock_stream_handler.return_value.setFormatter.assert_called_once()
# Verify formatter configuration
mock_formatter.assert_called_once_with(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
# Verify handler was added to logger
assert len(logger.handlers) == 1

127
tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
View File

@@ -1,127 +0,0 @@
from unittest.mock import call, patch, AsyncMock, ANY
from pytest import mark, fixture
from laborious.activities.activities import Activities
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
@fixture
def format_and_export_prediction():
return FormatAndExportPrediction()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
input_data = {
"path_flag": None,
"data": {"test": "data"},
"timestamp": "2021-01-01",
"model_id": 1,
"prediction_confidence": 0,
"schema": "test_schema",
"table_name": "test_table",
"opc_servers": ["test_server"],
"opc_output_config": {"test": "config"}
}
await format_and_export_prediction.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.format_prediction,
{
'data': input_data['data'],
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence']
},
retry_policy=ANY,
start_to_close_timeout=ANY
)])
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.export_data_to_postgres,
{
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': 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.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
assert workflow_mock.execute_activity_method.call_count == 2
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio
@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
input_data = {
"path_flag": "default",
"data": {"test": "data"},
"timestamp": "2021-01-01",
"model_id": 1,
"prediction_confidence": 0,
"schema": "test_schema",
"table_name": "test_table",
"opc_servers": ["test_server"],
"opc_output_config": {"test": "config"},
"comment": "test_comment"
}
await format_and_export_prediction.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.format_default_prediction,
{
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'comment': input_data['comment']
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.export_data_to_postgres,
{
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': 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.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
assert workflow_mock.execute_activity_method.call_count == 2
assert workflow_mock.execute_local_activity_method.call_count == 1

515
tests/laborious/workflows/subworkflows/test_prediction_process.py
View File

@@ -1,515 +0,0 @@
from unittest.mock import AsyncMock, patch, call, ANY
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
@fixture
def prediction_process():
return PredictionProcess()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_run(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(return_value=False)
# Arrange
input_data = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'},
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
# mlflow_response_gate (transform)
('continue', 0.95, "Error"),
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
# mlflow_response_gate (predict)
('continue', 0.95, "Error"),
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, {
'data': 'transformed_data',
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_called_once_with(
'format_and_export_prediction',
{
'path_flag': 'continue',
'data': 'predicted_data',
'prediction_confidence': 0.95,
'timestamp': '2024-01-01',
'model_id': 1,
'model_name': 'test_model_name',
'model_retention': '30',
'opc_output_config': input_data['opc_output_config'],
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'comment': 'Error'
}
)
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(return_value=True)
# Arrange
input_data = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('stop', 0.95, "Input data with bad quality"), # input_gate
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 2
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
'data': input_data['data']}, retry_policy=ANY, start_to_close_timeout=ANY),
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']}, retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True])
# Arrange
input_data = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('repeat', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95, "Error"), # mlflow_response_gate (transform)
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 4
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']},
retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(
side_effect=[False, False, True])
# Arrange
input_data = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
# mlflow_response_gate (transform)
('continue', 0.95, "Error"),
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 5
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(
side_effect=[False, False, False, True])
# Arrange
input_data = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
# mlflow_response_gate (transform)
('continue', 0.95, "Error"),
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
('continue', 0.95, "Error"), # mlflow_response_gate (predict)
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, {
'data': 'transformed_data',
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_path_flag_handler_stop(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
path_flag = 'STOP'
confidence = 0.95
schema = 'test_schema'
table_name = 'test_table'
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_retention = '30'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
'schema': schema,
'table_name': table_name,
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_retention': model_retention
}, confidence, last_timestamp, ""
)
# Assert
assert result is True
workflow_mock.execute_local_activity_method.assert_not_called()
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
path_flag = 'repeat'
confidence = 0.95
schema = 'test_schema'
table_name = 'test_table'
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_retention = '30'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
'schema': schema,
'table_name': table_name,
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_retention': model_retention
}, confidence, last_timestamp, ""
)
# Assert
assert result is True
workflow_mock.execute_activity_method.assert_called_once_with(
Activities.repeat_last_prediction,
{
'schema': schema,
'table_name': table_name,
'model_id': model
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
workflow_mock.execute_child_workflow.assert_not_called()
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_path_flag_handler_continue(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
path_flag = 'CONTINUE'
confidence = 0.95
schema = 'test_schema'
table_name = 'test_table'
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_retention = '30'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
'schema': schema,
'table_name': table_name,
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_retention': model_retention,
'opc_output_config': {'test': 'config'}
}, confidence, last_timestamp, 'Prediction Process'
)
# Assert
assert result is True
workflow_mock.execute_activity_method.assert_not_called()
workflow_mock.execute_child_workflow.assert_called_once_with(
'format_and_export_prediction',
{
'path_flag': path_flag,
'data': data,
'prediction_confidence': confidence,
'timestamp': last_timestamp,
'model_id': model,
'model_name': model_name,
'model_retention': model_retention,
'schema': schema,
'table_name': table_name,
'comment': 'Prediction Process',
'opc_output_config': {'test': 'config'}
}
)
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
path_flag = 'unknown'
confidence = 0.95
schema = 'test_schema'
table_name = 'test_table'
model = 'test_model'
last_timestamp = '2024-01-01'
model_name = 'test_model_name'
model_retention = '30'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
'schema': schema,
'table_name': table_name,
'model_id': model,
'last_timestamp': last_timestamp,
'model_name': model_name,
'model_retention': model_retention,
'opc_output_config': {'test': 'config'}
}, confidence, last_timestamp, ""
)
# Assert
assert result is False
workflow_mock.execute_activity_method.assert_not_called()
workflow_mock.execute_child_workflow.assert_not_called()

81
tests/laborious/workflows/test_predictions_batch.py
View File

@@ -1,81 +0,0 @@
from unittest.mock import AsyncMock, call, patch, ANY
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
@fixture
def predictions_batch() -> PredictionsBatch:
return PredictionsBatch()
@mark.asyncio
@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
workflow_mock.execute_local_activity_method.return_value = {
'data': 'test_data'
}
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'query': 'SELECT * FROM test',
'schema': 'test_schema',
'table_name': 'test_table',
'opc_output_config': 'test_opc_output_config'
}
await predictions_batch.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.prepare_activity,
{
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'predictions_batch'
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.load_custom_query,
input_data['query'],
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
prediction_input = {
'data': {'data': 'test_data'},
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'input_filters': input_data.get('input_filters', {
'EMPTY_DATA': {
'POLICY': 'STOP'
}
}),
'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'model_retention': input_data.get('model_retention', 60),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
'opc_output_config': input_data.get('opc_output_config', {})
}
workflow_mock.execute_child_workflow.assert_has_calls([
call(
'prediction_process', prediction_input)
])

0
tests/laborious/__init__.py → tests/orchestrator/__init__.py
View File

116
tests/orchestrator/activities/test_activities.py Normal file
View 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
View 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
View 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
View 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
View 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
)
])