SIENTIAPDE-1030
Add unit tests for connectors configuration, logger, workflows, and predictions batch - Implement tests for MLflow, OPC, and Postgres configuration builders to validate environment variable handling and default values. - Create tests for the logger to ensure default settings and handler configurations are correct. - Add comprehensive tests for the FormatAndExportPrediction and PredictionProcess workflows, covering various scenarios including path flags and activity execution. - Introduce tests for the PredictionsBatch workflow to verify the execution of local activities and child workflows. - Include a values.yaml file for Kubernetes deployment configuration, specifying image details, service account settings, environment variables, and resource limits.
This commit is contained in:
vitor-aignosi
0
tests/laborious/activities/__init__.py
Normal file
0
tests/laborious/activities/__init__.py
Normal file
193
tests/laborious/activities/test_activities.py
Normal file
193
tests/laborious/activities/test_activities.py
Normal file
@@ -0,0 +1,193 @@
|
||||
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
Normal file
35
tests/laborious/activities/test_base.py
Normal file
@@ -0,0 +1,35 @@
|
||||
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
Normal file
369
tests/laborious/activities/test_gates.py
Normal file
@@ -0,0 +1,369 @@
|
||||
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
Normal file
120
tests/laborious/activities/test_mlflow.py
Normal file
@@ -0,0 +1,120 @@
|
||||
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
Normal file
196
tests/laborious/activities/test_opc.py
Normal file
@@ -0,0 +1,196 @@
|
||||
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
Normal file
159
tests/laborious/activities/test_postgres.py
Normal file
@@ -0,0 +1,159 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user