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

SIENTIAPDE-1231

Browse Source 
Update .gitignore and refactor metrics.py for improved logging and consistency

- Added coverage.xml to .gitignore to prevent tracking of coverage reports.
- Refactored metric labels in metrics.py for consistency in string formatting and improved readability.
- Enhanced logging messages in various activities to ensure uniformity in message formatting.
This commit is contained in:
vitor-aignosi
2025-10-15 16:00:18 -03:00
parent a5d2b0d3fd
commit ac795c7c53
39 changed files with 4122 additions and 2602 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,18 +1,19 @@
from unittest.mock import ANY, MagicMock, patch
from pytest import mark
from unittest.mock import patch, MagicMock, ANY
from sientia_do.temporal.activities.postgres import Postgres
from laborious.activities.activities import Activities
from laborious.activities.mlflow import MLFlow
from laborious.activities.gates import Gates
from laborious.activities.mlflow import MLFlow
from laborious.activities.opc import OPC
from laborious.activities.storage import Storage
@patch('laborious.activities.activities.Postgres.__init__')
@patch('laborious.activities.activities.Storage.__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):
def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_storage_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
@@ -20,20 +21,23 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
'max_connections': 10,
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
minio_config = {
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
}
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group'
'group_id': 'test-group',
}
logger = MagicMock()
@@ -42,18 +46,19 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
minio_config=minio_config,
opc_config=opc_config,
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
)
assert isinstance(activities, Activities)
assert isinstance(activities, Postgres)
assert isinstance(activities, Storage)
assert isinstance(activities, MLFlow)
assert isinstance(activities, OPC)
assert isinstance(activities, Gates)
mock_postgres_init.assert_called_once_with(
mock_storage_init.assert_called_once_with(
ANY,
host=postgres_config['host'],
port=postgres_config['port'],
@@ -62,8 +67,9 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
dbname=postgres_config['dbname'],
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
minio_config=minio_config,
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
)
mock_mlflow_init.assert_called_once_with(
@@ -72,30 +78,25 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
minio_config=minio_config,
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
)
mock_opc_init.assert_called_once_with(
ANY,
opc_servers=opc_config,
logger=logger,
notification_handler=notification_handler
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
ANY, logger=logger, notification_handler=notification_handler
)
@mark.asyncio
@patch('laborious.activities.activities.Postgres', return_value=MagicMock())
@patch('laborious.activities.activities.Storage', return_value=MagicMock())
@patch('laborious.activities.activities.MLFlow', return_value=MagicMock())
@patch('laborious.activities.activities.OPC', return_value=MagicMock())
async def test_shutdown(mock_opc_init,
_mock_mlflow_init, mock_postgres_init):
async def test_shutdown(mock_opc_init, _mock_mlflow_init, mock_storage_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
@@ -103,20 +104,23 @@ async def test_shutdown(mock_opc_init,
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
'max_connections': 10,
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
minio_config = {
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
}
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group'
'group_id': 'test-group',
}
logger = MagicMock()
@@ -125,11 +129,12 @@ async def test_shutdown(mock_opc_init,
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
minio_config=minio_config,
opc_config=opc_config,
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
)
await activities.shutdown()
mock_opc_init.shutdown.assert_called_once()
mock_postgres_init.close.assert_called_once()
mock_storage_init.close.assert_called_once()

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

@@ -1,6 +1,8 @@
from unittest.mock import MagicMock, ANY, patch
from unittest.mock import ANY, MagicMock, patch
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from laborious.activities.gates import Gates
@@ -20,11 +22,11 @@ def gates_activity():
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@@ -34,20 +36,18 @@ async def test_input_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'value': [1, 2, 3]},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.error.assert_called_once_with(
"Filter INVALID_FILTER not found", metadata['metadata']
'Filter INVALID_FILTER not found', metadata['metadata']
)
@@ -57,28 +57,27 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac
# Arrange
mock_input_filter_functions.__contains__.return_value = True
mock_input_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception("Test error"))
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {
'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
},
'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id="INTPUT_GATE_ERROR__EMPTY_DATA",
notification_id='INTPUT_GATE_ERROR__EMPTY_DATA',
message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \n Test error",
block="input_gate",
block='input_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
@@ -89,14 +88,14 @@ async def test_input_gate_no_filters(gates_activity):
**metadata,
'filters': {},
'data': {'value': [1, 2, 3]},
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@@ -105,18 +104,16 @@ async def test_input_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
},
'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert
assert result == ('STOP', -1, "Input data with bad quality")
assert result == ('STOP', -1, 'Input data with bad quality')
gates_activity.debug.assert_called()
@@ -125,51 +122,49 @@ async def test_mlflow_response_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
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):
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"))
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id="MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER",
notification_id='MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER',
message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error",
block="mlflow_gate",
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
@@ -181,14 +176,14 @@ async def test_mlflow_response_gate_no_filters(gates_activity):
'filters': {},
'data': {'content': {'message': 'success'}},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@@ -197,25 +192,20 @@ async def test_mlflow_response_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'API_ERROR': {'policy': 'STOP'}
},
'filters': {'API_ERROR': {'policy': 'STOP'}},
'data': {
'success': False,
'content': {
'message': 'API error occurred',
'traceback': 'error trace'
}
'content': {'message': 'API error occurred', 'traceback': 'error trace'},
},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert
assert result == ('STOP', -1, "API error occurred")
assert result == ('STOP', -1, 'API error occurred')
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called()
@@ -225,58 +215,53 @@ async def test_mlflow_content_gate_invalid_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
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):
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"))
side_effect=Exception('Test error')
)
input_data = {
**metadata,
'filters': {
'API_ERROR': {'POLICY': 'STOP'}
},
'filters': {'API_ERROR': {'POLICY': 'STOP'}},
'data': {
'success': False,
'content': {
'message': 'API error occurred',
'traceback': 'error trace'
}
'content': {'message': 'API error occurred', 'traceback': 'error trace'},
},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id="MLFLOW_GATE_CONTENT_FILTER__API_ERROR",
notification_id='MLFLOW_GATE_CONTENT_FILTER__API_ERROR',
message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error",
block="mlflow_gate",
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
@@ -288,14 +273,14 @@ async def test_mlflow_content_gate_no_filters(gates_activity):
'filters': {},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert
assert result == (None, 0, "")
assert result == (None, 0, '')
gates_activity.debug.assert_called()
@@ -304,20 +289,17 @@ async def test_mlflow_content_gate_with_filter(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {
'NAN_VALUES': {'policy': 'STOP', 'config': {}}
},
'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}},
'data': {'value': [None, None, None]},
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
'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")
assert result == ('STOP', -1, 'Transformed data not passed the content filter')
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called()
@@ -328,7 +310,8 @@ def test_get_prediction_store_policy_invalid_policy(gates_activity):
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata)
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
@@ -341,7 +324,8 @@ def test_get_prediction_store_policy_invalid_policy_value(gates_activity):
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata)
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
@@ -354,7 +338,8 @@ def test_get_prediction_store_policy_valid_policy_type(gates_activity):
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata)
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'lts'
@@ -367,7 +352,8 @@ def test_get_prediction_store_policy_valid_policy(gates_activity):
# Act
policy_type, policy_value = gates_activity.get_prediction_store_policy(
prediction_store_policy, metadata)
prediction_store_policy, metadata
)
# Assert
assert policy_type == 'erl'
@@ -380,16 +366,12 @@ async def test_format_prediction_no_timestamp(gates_activity):
input_data = {
**metadata,
'data': {
'prediction': {
'2023-05-26 11:12:27': 1
},
'response_time': {
'2023-05-26 11:12:27': 0.1
}
'prediction': {'2023-05-26 11:12:27': 1},
'response_time': {'2023-05-26 11:12:27': 0.1},
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'lts:1'
'prediction_store_policy': 'lts:1',
}
# Act
@@ -402,7 +384,7 @@ async def test_format_prediction_no_timestamp(gates_activity):
assert result['model_id'] == {0: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9}
assert result['prediction_status'] == {0: 'Good'}
assert result['comments'] == {0: ""}
assert result['comments'] == {0: ''}
@mark.asyncio
@@ -420,11 +402,11 @@ async def test_format_prediction_with_timestamp_erl(gates_activity):
'2023-05-26 11:12:27': 0.1,
'2023-05-26 11:12:28': 0.2,
'2023-05-26 11:12:29': 0.3,
}
},
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'erl:2'
'prediction_store_policy': 'erl:2',
}
# Act
@@ -433,12 +415,11 @@ async def test_format_prediction_with_timestamp_erl(gates_activity):
# Assert
assert result['prediction'] == {0: 2, 1: 1}
assert result['response_time'] == {0: 0.2, 1: 0.1}
assert result['timestamp'] == {
0: '2023-05-26 11:12:28', 1: '2023-05-26 11:12:27'}
assert result['timestamp'] == {0: '2023-05-26 11:12:28', 1: '2023-05-26 11:12:27'}
assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
assert result['comments'] == {0: "", 1: ""}
assert result['comments'] == {0: '', 1: ''}
@mark.asyncio
@@ -456,11 +437,11 @@ async def test_format_prediction_with_timestamp_lts(gates_activity):
'2023-05-26 11:12:27': 0.1,
'2023-05-26 11:12:28': 0.2,
'2023-05-26 11:12:29': 0.3,
}
},
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'lts:2'
'prediction_store_policy': 'lts:2',
}
# Act
@@ -469,12 +450,11 @@ async def test_format_prediction_with_timestamp_lts(gates_activity):
# Assert
assert result['prediction'] == {0: 3, 1: 2}
assert result['response_time'] == {0: 0.3, 1: 0.2}
assert result['timestamp'] == {
0: '2023-05-26 11:12:29', 1: '2023-05-26 11:12:28'}
assert result['timestamp'] == {0: '2023-05-26 11:12:29', 1: '2023-05-26 11:12:28'}
assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
assert result['comments'] == {0: "", 1: ""}
assert result['comments'] == {0: '', 1: ''}
@mark.asyncio
@@ -482,22 +462,23 @@ async def test_format_prediction_with_timestamp_invalid_policy(gates_activity):
# Arrange
input_data = {
**metadata,
'data': {'prediction': [1, 2, 3],
'response_time': [0.1, 0.2, 0.3],
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']},
'data': {
'prediction': [1, 2, 3],
'response_time': [0.1, 0.2, 0.3],
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'lts:2'
'prediction_store_policy': 'lts:2',
}
gates_activity.get_prediction_store_policy = MagicMock(
return_value=('invalid', 1))
gates_activity.get_prediction_store_policy = MagicMock(return_value=('invalid', 1))
try:
result = await gates_activity.format_prediction(input_data)
await gates_activity.format_prediction(input_data)
except ValueError as e:
assert str(e) == "Invalid policy type: invalid"
assert str(e) == 'Invalid policy type: invalid'
else:
assert False, "Expected ValueError"
raise AssertionError('Expected ValueError')
@mark.asyncio
@@ -508,7 +489,7 @@ async def test_format_default_prediction(gates_activity):
'timestamp': '2023-05-26 11:12:27',
'model_id': 'test_model',
'prediction_confidence': 0.1,
'comment': 'Test comment'
'comment': 'Test comment',
}
# Act
@@ -528,12 +509,7 @@ async def test_format_default_prediction(gates_activity):
@mark.asyncio
async def test_get_last_timestamp_with_data(gates_activity):
# Arrange
input_data = {
**metadata,
'data': {
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']
}
}
input_data = {**metadata, 'data': {'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']}}
# Act
result = await gates_activity.get_last_timestamp(input_data)
@@ -545,10 +521,7 @@ async def test_get_last_timestamp_with_data(gates_activity):
@mark.asyncio
async def test_get_last_timestamp_no_data(gates_activity):
# Arrange
input_data = {
'data': {},
**metadata
}
input_data = {'data': {}, **metadata}
# Act
result = await gates_activity.get_last_timestamp(input_data)
@@ -567,30 +540,28 @@ async def test_write_metrics(mock_metrics, gates_activity):
'prediction': {
'prediction': [1, 2, 3],
'prediction_confidence': [0.9, 0.8, 0.7],
'response_time': [0.1, 0.2, 0.3]
}
'response_time': [0.1, 0.2, 0.3],
},
}
await gates_activity.write_metrics(input_data)
mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.assert_called_once_with(
pod_id=gates_activity.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name']
pipeline_name=metadata['metadata']['workflow_name'],
)
mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.return_value.inc.assert_called_once_with()
mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.assert_called_once_with(
pod_id=gates_activity.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name']
)
mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.return_value.set.assert_called_once_with(
0.9
pipeline_name=metadata['metadata']['workflow_name'],
)
mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.return_value.set.assert_called_once_with(0.9)
mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.assert_called_once_with(
pod_id=gates_activity.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name']
pipeline_name=metadata['metadata']['workflow_name'],
)
mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
0.1

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

@@ -1,45 +1,68 @@
from datetime import datetime
from unittest.mock import ANY, MagicMock, patch
from unittest.mock import ANY, MagicMock, call, patch
import numpy as np
from pandas import DataFrame, Timestamp
from pytest import fixture, mark, raises
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from laborious.activities.mlflow import MLFlow
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from laborious.activities.mlflow import MLFlow
@patch("laborious.activities.mlflow.MLFlowRepository")
def test___init__(mock_mlflow_repository):
@patch('laborious.activities.mlflow.MLFlowRepository')
@patch('laborious.activities.mlflow.MinioRepository')
def test___init__(mock_minio_repository, mock_mlflow_repository):
mlflow = MLFlow(
mlflow_host="http://localhost",
mlflow_host='http://localhost',
mlflow_port=5000,
mlflow_username="admin",
mlflow_password="admin",
mlflow_username='admin',
mlflow_password='admin',
minio_config={
'endpoint_url': 'http://localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
},
logger=MagicMock(),
notification_handler=MagicMock()
notification_handler=MagicMock(),
)
assert mlflow.mlflow_host == "http://localhost"
assert mlflow.mlflow_host == 'http://localhost'
assert mlflow.mlflow_port == 5000
assert mlflow.mlflow_username == "admin"
assert mlflow.mlflow_password == "admin"
assert mlflow.mlflow_username == 'admin'
assert mlflow.mlflow_password == 'admin'
mock_mlflow_repository.assert_called_once_with(
"http://localhost:5000", "admin", "admin", ANY
mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY)
mock_minio_repository.assert_called_once_with(
logger=ANY,
notification_handler=ANY,
minio_endpoint_url='http://localhost:9000',
minio_access_key='minio',
minio_secret_key='minio123',
minio_region_name='us-east-1',
minio_default_bucket='test',
)
@fixture
@patch("laborious.activities.mlflow.MLFlowRepository")
def mlflow(mock_mlflow_repository):
@patch('laborious.activities.mlflow.MLFlowRepository')
@patch('laborious.activities.mlflow.MinioRepository')
def mlflow(mock_minio_repository, mock_mlflow_repository):
mlflow = MLFlow(
mlflow_host="http://localhost:5000",
mlflow_host='http://localhost:5000',
mlflow_port=5000,
mlflow_username="admin",
mlflow_password="admin",
mlflow_username='admin',
mlflow_password='admin',
minio_config={
'endpoint_url': 'http://localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
},
logger=MagicMock(),
notification_handler=MagicMock()
notification_handler=MagicMock(),
)
mlflow.send_notification = MagicMock()
@@ -48,44 +71,67 @@ def mlflow(mock_mlflow_repository):
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@mark.asyncio
@patch("laborious.activities.mlflow.DataFrame")
@patch("laborious.activities.mlflow.max")
@patch('laborious.activities.mlflow.DataFrame')
@patch('laborious.activities.mlflow.max')
async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02'
# Mock input data
input_data = {
**metadata,
'data': [
{'timestamp': '2024-01-01', 'variable': 'var1',
'value': 1.0, 'created_at': '2024-01-01 12:00:00'},
{'timestamp': '2024-01-01', 'variable': 'var2',
'value': 2.0, 'created_at': '2024-01-01 12:00:00'},
{'timestamp': '2024-01-02', 'variable': 'var1',
'value': 3.0, 'created_at': '2024-01-02 12:00:00'},
{'timestamp': '2024-01-02', 'variable': 'var2',
'value': 4.0, 'created_at': '2024-01-02 12:00:00'},
{'timestamp': '2024-01-02', 'variable': 'var1',
'value': 1.0, 'created_at': '2024-01-01 12:00:00'},
{'timestamp': '2024-01-02', 'variable': 'var2',
'value': 1.0, 'created_at': '2024-01-01 12:00:00'}
{
'timestamp': '2024-01-01',
'variable': 'var1',
'value': 1.0,
'created_at': '2024-01-01 12:00:00',
},
{
'timestamp': '2024-01-01',
'variable': 'var2',
'value': 2.0,
'created_at': '2024-01-01 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var1',
'value': 3.0,
'created_at': '2024-01-02 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var2',
'value': 4.0,
'created_at': '2024-01-02 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var1',
'value': 1.0,
'created_at': '2024-01-01 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var2',
'value': 1.0,
'created_at': '2024-01-01 12:00:00',
},
],
'model_name': 'test_model',
'model_config': {}
'model_config': {},
}
# Mock the transform response
expected_response = {'prediction': [0.5, 0.6], 'timestamp': [
'2024-01-01', '2024-01-02']}
expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']}
mlflow.model_monitoring_repository.transform.return_value = expected_response
mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value
@@ -114,30 +160,25 @@ async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
@mark.asyncio
@patch("laborious.activities.mlflow.DataFrame")
@patch("laborious.activities.mlflow.to_datetime")
@patch("laborious.activities.mlflow.max")
@patch('laborious.activities.mlflow.DataFrame')
@patch('laborious.activities.mlflow.to_datetime')
@patch('laborious.activities.mlflow.max')
async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02'
# Mock input data
input_data = {
**metadata,
'data': {
"variable": {
"2024-01-01": "var1",
"2024-01-02": "var2",
"2024-01-03": "var1",
"2024-01-04": "var2"
'variable': {
'2024-01-01': 'var1',
'2024-01-02': 'var2',
'2024-01-03': 'var1',
'2024-01-04': 'var2',
},
"value": {
"2024-01-01": 1.0,
"2024-01-02": 2.0,
"2024-01-03": 3.0,
"2024-01-04": 4.0
}
'value': {'2024-01-01': 1.0, '2024-01-02': 2.0, '2024-01-03': 3.0, '2024-01-04': 4.0},
},
'model_name': 'test_model',
'model_config': {}
'model_config': {},
}
# Mock the predict response
@@ -148,9 +189,7 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
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
)
mock_dataframe.return_value.replace.assert_called_once_with(np.nan, None, inplace=True)
mock_dataframe.return_value.__setitem__.assert_any_call(
'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
)
@@ -158,9 +197,7 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
mock_to_datetime.assert_called_once_with(
mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
)
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(
DATETIME_FORMAT
)
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
# Verify the response
assert response_data == expected_response
@@ -172,98 +209,211 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
@mark.asyncio
async def test_retrain_model(mlflow):
data = {
"model_id": [4, 5, 6, 7],
"created_at": [1, 2, 3, 4],
"timestamp": [1, 1, 2, 2],
"variable": ["var1", "var2", "var1", "var2"],
"value": [1, 2, 3, 4]
@patch('laborious.activities.mlflow.to_datetime')
async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlflow):
mlflow.model_monitoring_repository.retrain_model.return_value = {
'success': True,
'experiment': 'test_experiment',
'message': 'Model retrained successfully.',
}
mlflow.model_monitoring_repository.retrain_model.return_value = (
'Model retrained successfully', 'test')
response = await mlflow.retrain_model(
{
**metadata,
'object_key': 'test_object_key',
'model_name': 'test_model',
'model_config': {
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
}
)
response = await mlflow.retrain_model({
**metadata,
'data': data,
'model_name': 'test_model'
})
raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value
mlflow.model_monitoring_repository.retrain_model.assert_called_once()
timestamp = raw_data.__getitem__.return_value.max.return_value
raw_data.sort_values.assert_called_once_with('created_at', ascending=False)
raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with(
subset=['variable', 'timestamp'], keep='first'
)
raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value
raw_data.drop.assert_has_calls(
[
call(columns=['model_id'], inplace=True, errors='ignore'),
call(columns=['created_at'], inplace=True, errors='ignore'),
]
)
raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
raw_data = raw_data.pivot.return_value
raw_data.__setitem__.assert_has_calls(
[
call('timestamp', raw_data.index),
call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
call('timestamp', mock_to_datetime.return_value),
]
)
mock_to_datetime.assert_has_calls(
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
)
mock_to_datetime.assert_has_calls(
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
)
mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
data=raw_data,
model_name='test_model',
model_config={
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
metadata=metadata['metadata'],
)
assert response == {
"status": 'Model retrained successfully',
"timestamp": 2,
"experiment": 'test'
'success': True,
'experiment': 'test_experiment',
'message': 'Model retrained successfully.',
'timestamp': timestamp,
}
@mark.asyncio
async def test_retrain_model_error(mlflow):
mlflow.model_monitoring_repository.retrain_model.side_effect = Exception(
'Error retraining model'
)
data = {
"model_id": [4, 5, 6, 7],
"created_at": [1, 2, 3, 4],
"timestamp": [1, 1, 2, 2],
"variable": ["var1", "var2", "var1", "var2"],
"value": [1, 2, 3, 4]
@patch('laborious.activities.mlflow.to_datetime')
async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow):
mlflow.model_monitoring_repository.retrain_model.return_value = {
'success': False,
'traceback': 'test_traceback',
'message': 'Model retrained failed.',
}
try:
await mlflow.retrain_model({
response = await mlflow.retrain_model(
{
**metadata,
'data': data,
'model_name': 'test_model'
})
except Exception as e:
assert str(e) == 'Error retraining model'
mlflow.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='RETRAIN_MODEL_ERROR',
message='Error retraining model test_model: Error retraining model',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=ANY
)
else:
assert False, "No exception raised"
'object_key': 'test_object_key',
'model_name': 'test_model',
'model_config': {
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
}
)
raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value
timestamp = raw_data.__getitem__.return_value.max.return_value
raw_data.sort_values.assert_called_once_with('created_at', ascending=False)
raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with(
subset=['variable', 'timestamp'], keep='first'
)
raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value
raw_data.drop.assert_has_calls(
[
call(columns=['model_id'], inplace=True, errors='ignore'),
call(columns=['created_at'], inplace=True, errors='ignore'),
]
)
raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
raw_data = raw_data.pivot.return_value
raw_data.__setitem__.assert_has_calls(
[
call('timestamp', raw_data.index),
call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
call('timestamp', mock_to_datetime.return_value),
]
)
mock_to_datetime.assert_has_calls(
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
)
mock_to_datetime.assert_has_calls(
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
)
mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
data=raw_data,
model_name='test_model',
model_config={
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
metadata=metadata['metadata'],
)
mlflow.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='RETRAIN_MODEL_ERROR',
message='Error retraining model test_model: Model retrained failed.',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
assert response == {
'success': False,
'traceback': 'test_traceback',
'message': 'Model retrained failed.',
'timestamp': timestamp,
}
@mark.asyncio
async def test_retrain_model_data_error(mlflow):
mlflow.minio_repository.get_parquet_as_dataframe.side_effect = Exception(
'Error loading retrain data'
)
response = await mlflow.retrain_model(
{
**metadata,
'object_key': 'test_object_key',
'model_name': 'test_model',
'model_config': {
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
}
)
assert response == {
'success': False,
'message': 'Error loading retrain data: Error loading retrain data',
'traceback': ANY,
'timestamp': ANY,
}
@mark.asyncio
async def test_update_production_model(mlflow):
mlflow.model_monitoring_repository.update_production_model.return_value = (
{
"data1": 1,
"data2": 2
}
)
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 1,
'experiment': 'test',
'timestamp': 2,
'status': 'success'
'status': 'success',
}
response = await mlflow.update_production_model(input_data)
mlflow.model_monitoring_repository.update_production_model.assert_called_once_with(
experiment='test', model_name='test_model')
experiment='test', model_name='test_model', metadata=metadata['metadata']
)
assert response == {
'data1': {0: 1},
'data2': {0: 2},
'model_id': {0: 1},
'model_name': {0: 'test_model'},
'timestamp': {0: 2},
'status': {0: 'success'}
}
assert response == mlflow.model_monitoring_repository.update_production_model.return_value
@mark.asyncio
@@ -278,7 +428,7 @@ async def test_update_production_model_error(mlflow):
'model_id': 1,
'experiment': 'test',
'timestamp': 2,
'status': 'success'
'status': 'success',
}
try:
@@ -291,7 +441,7 @@ async def test_update_production_model_error(mlflow):
message='Error updating production model test_model: Error updating production model',
block='update_production_model',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
else:
assert False, "No exception raised"
raise AssertionError('No exception raised')

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

@@ -1,57 +1,55 @@
from unittest.mock import patch, MagicMock, ANY, call, AsyncMock
from pandas import DataFrame
from pytest import fixture, mark
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
import pytest_asyncio
from pandas import DataFrame
from pytest import mark
from sientia_do.notifications.models import NotificationLevel
from laborious.activities.opc import OPC
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
def test__init__():
servers = {
'server1': 'config'
}
opc = OPC(
opc_servers=servers,
logger=MagicMock(),
notification_handler=MagicMock()
)
servers = {'server1': 'config'}
opc = OPC(opc_servers=servers, logger=MagicMock(), notification_handler=MagicMock())
assert opc.opc_servers == servers
assert opc.opc_repository == {}
@mark.asyncio
@patch("laborious.activities.opc.OpcRepository")
@patch("laborious.activities.opc.OPC.send_notification")
@patch('laborious.activities.opc.OpcRepository')
@patch('laborious.activities.opc.OPC.send_notification')
async def test_init_opc(mock_send_notification, mock_opc_repository):
mock_logger = MagicMock()
server1 = MagicMock(
connect=AsyncMock(return_value=(True, {})),
write_data=AsyncMock(return_value=(True, {}))
connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {}))
)
server2 = MagicMock(
connect=AsyncMock(return_value=(True, {})),
write_data=AsyncMock(return_value=(True, {}))
connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {}))
)
server3 = MagicMock(
connect=AsyncMock(return_value=(False, {
'notification_id': 'OPC_CONNECTION_ERROR_server3',
'message': 'Failed to connect to OPC server: Test error',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': 'Test error'
})),
write_data=AsyncMock(return_value=(True, {}))
connect=AsyncMock(
return_value=(
False,
{
'notification_id': 'OPC_CONNECTION_ERROR_server3',
'message': 'Failed to connect to OPC server: Test error',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': 'Test error',
},
)
),
write_data=AsyncMock(return_value=(True, {})),
)
mock_opc_repository.side_effect = [server1, server2, server3]
mock_notification_handler = MagicMock()
@@ -82,12 +80,10 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
}
},
}
opc = OPC(
opc_servers=servers,
logger=mock_logger,
notification_handler=mock_notification_handler
opc_servers=servers, logger=mock_logger, notification_handler=mock_notification_handler
)
await opc.init_opc()
@@ -97,57 +93,63 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
assert opc.opc_repository['server1'] == server1
assert opc.opc_repository['server2'] == server2
mock_opc_repository.assert_has_calls([
call(
id="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,
pod_id='localhost'
),
])
mock_opc_repository.assert_has_calls([
call(
id="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,
pod_id='localhost'
)
])
mock_opc_repository.assert_has_calls(
[
call(
opc_id='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,
pod_id='localhost',
),
]
)
mock_opc_repository.assert_has_calls(
[
call(
opc_id='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,
pod_id='localhost',
)
]
)
server1.connect.assert_called_once()
server2.connect.assert_called_once()
mock_send_notification.assert_has_calls([
call(
metadata={
'model_id': '-',
'model_name': '-',
'workflow_name': '-',
'schedule_name': 'INITIALIZATION'
},
notification_id="OPC_CONNECTION_ERROR_server3",
message="Failed to connect to OPC server: Test error",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
])
mock_send_notification.assert_has_calls(
[
call(
metadata={
'model_id': '-',
'model_name': '-',
'workflow_name': '-',
'schedule_name': 'INITIALIZATION',
},
notification_id='OPC_CONNECTION_ERROR_server3',
message='Failed to connect to OPC server: Test error',
block='opc_repository',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
]
)
@pytest_asyncio.fixture
@patch("laborious.activities.opc.OpcRepository")
@patch('laborious.activities.opc.OpcRepository')
async def opc(mock_opc_repository):
servers = {
'server1': {
@@ -161,17 +163,9 @@ async def opc(mock_opc_repository):
}
}
mock_opc_repository.return_value.write_data = AsyncMock(
return_value=(True, {})
)
mock_opc_repository.return_value.connect = AsyncMock(
return_value=(True, {})
)
opc = OPC(
opc_servers=servers,
logger=MagicMock(),
notification_handler=MagicMock()
)
mock_opc_repository.return_value.write_data = AsyncMock(return_value=(True, {}))
mock_opc_repository.return_value.connect = AsyncMock(return_value=(True, {}))
opc = OPC(opc_servers=servers, logger=MagicMock(), notification_handler=MagicMock())
await opc.init_opc()
opc.send_notification = MagicMock()
return opc
@@ -188,58 +182,79 @@ WRITE_DATA_CASES = [
@mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
@mark.asyncio
async def test_write_data_success(opc, tag, data_type, data):
result = await opc.write_data(server_id='server1', tag=tag, data=data,
data_type=data_type, tag_type='prediction', metadata=metadata)
result = await opc.write_data(
server_id='server1',
tag=tag,
data=data,
data_type=data_type,
tag_type='prediction',
metadata=metadata,
)
assert result is True
opc.opc_repository['server1'].write_data.assert_called_once_with(
tag, data, data_type, opc.logger, metadata)
tag, data, data_type, opc.logger, metadata
)
@mark.asyncio
async def test_write_data_failed(opc):
opc.opc_repository['server1'].write_data.return_value = (False, {
'notification_id': 'OPC_WRITE_DATA_ERROR_server1',
'message': 'Failed to write data to OPC server: Test error',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': 'Test error'
})
opc.opc_repository['server1'].write_data.return_value = (
False,
{
'notification_id': 'OPC_WRITE_DATA_ERROR_server1',
'message': 'Failed to write data to OPC server: Test error',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': 'Test error',
},
)
result = await opc.write_data(server_id='server1', tag='tag1', data=50,
data_type='int', tag_type='prediction', metadata=metadata)
result = await opc.write_data(
server_id='server1',
tag='tag1',
data=50,
data_type='int',
tag_type='prediction',
metadata=metadata,
)
assert result is False
opc.send_notification.assert_called_once_with(
metadata=metadata,
notification_id="OPC_WRITE_DATA_ERROR_server1",
message="Failed to write data to OPC server: Test error",
block="opc_repository",
notification_id='OPC_WRITE_DATA_ERROR_server1',
message='Failed to write data to OPC server: Test error',
block='opc_repository',
level=NotificationLevel.ERROR,
attachment_content=ANY
attachment_content=ANY,
)
@mark.asyncio
async def test_write_data_exception(opc):
opc.opc_repository['server1'].write_data.side_effect = Exception(
"Test error")
opc.opc_repository['server1'].write_data.side_effect = Exception('Test error')
try:
await opc.write_data(server_id='server1', tag='tag1', data=50,
data_type='int', tag_type='prediction', metadata=metadata)
await opc.write_data(
server_id='server1',
tag='tag1',
data=50,
data_type='int',
tag_type='prediction',
metadata=metadata,
)
except Exception:
opc.send_notification.assert_called_once_with(
metadata=metadata,
notification_id="WRITE_OPC_PREDICTION_ERROR",
message="Error writing data to OPC server: Test error",
block="write_opc_data",
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
attachment_content=ANY,
)
else:
assert False, "Expected an exception to be raised"
raise AssertionError('Expected an exception to be raised')
@mark.asyncio
@@ -247,20 +262,13 @@ async def test_write_opc_data_success(opc):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
},
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
'opc_output_config': {
'server1': {
'prediction_tags': {
'tag1': {'data_type': 'float'}
},
'confidence_tags': {
'tag2': {'data_type': 'float'}
}
'prediction_tags': {'tag1': {'data_type': 'float'}},
'confidence_tags': {'tag2': {'data_type': 'float'}},
}
}
},
}
# Act
@@ -270,25 +278,30 @@ async def test_write_opc_data_success(opc):
# Assert
assert output == {'data': 'data'}
opc.write_data.assert_has_calls([
call(
server_id='server1',
tag='tag1',
data=0.75,
data_type='float',
tag_type='prediction',
metadata=metadata['metadata']
)])
opc.write_data.assert_has_calls([
call(
server_id='server1',
tag='tag2',
data=0.95,
data_type='float',
tag_type='confidence',
metadata=metadata['metadata']
)
])
opc.write_data.assert_has_calls(
[
call(
server_id='server1',
tag='tag1',
data=0.75,
data_type='float',
tag_type='prediction',
metadata=metadata['metadata'],
)
]
)
opc.write_data.assert_has_calls(
[
call(
server_id='server1',
tag='tag2',
data=0.95,
data_type='float',
tag_type='confidence',
metadata=metadata['metadata'],
)
]
)
assert opc.write_data.call_count == 2
@@ -297,17 +310,9 @@ async def test_write_opc_data_empty_config(opc):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
},
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
'opc_servers': ['server1'],
'opc_output_config': {
'server1': {
'prediction_tags': {},
'confidence_tags': {}
}
}
'opc_output_config': {'server1': {'prediction_tags': {}, 'confidence_tags': {}}},
}
# Act
@@ -322,20 +327,13 @@ async def test_write_opc_data_no_validate_server(opc):
opc.validate_server = MagicMock(return_value=False)
input_data = {
**metadata,
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
},
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
'opc_output_config': {
'server1': {
'prediction_tags': {
'tag1': {'data_type': 'float'}
},
'confidence_tags': {
'tag2': {'data_type': 'float'}
}
'prediction_tags': {'tag1': {'data_type': 'float'}},
'confidence_tags': {'tag2': {'data_type': 'float'}},
}
}
},
}
# Act
@@ -345,10 +343,13 @@ async def test_write_opc_data_no_validate_server(opc):
opc.opc_repository['server1'].write_data.assert_not_called()
@mark.parametrize('data,success,expected', [
(DataFrame({'prediction_confidence': [0]}), True, 0),
(DataFrame({'prediction_confidence': [0]}), False, 12),
])
@mark.parametrize(
'data,success,expected',
[
(DataFrame({'prediction_confidence': [0]}), True, 0),
(DataFrame({'prediction_confidence': [0]}), False, 12),
],
)
def test_process_confidence(opc, data, success, expected):
# Act
result = opc.process_confidence(data, success, metadata)

203
tests/laborious/activities/test_storage.py Normal file
View File

@@ -0,0 +1,203 @@
import datetime
from unittest.mock import ANY, AsyncMock, MagicMock, patch
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.postgres import Postgres
from laborious.activities.storage import Storage
metadata = {
'metadata': {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schedule_name': 'test_schedule',
}
}
@fixture
@patch('laborious.activities.storage.MinioRepository')
def storage(mock_minio_repository):
return Storage(
host='localhost',
port=5432,
user='postgres',
password='postgres',
dbname='postgres',
min_connections=1,
max_connections=10,
minio_config={
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
},
logger=MagicMock(),
notification_handler=MagicMock(),
)
@patch('laborious.activities.storage.MinioRepository')
def test___init___not_hasattr(mock_minio_repository):
logger = MagicMock()
notification_handler = MagicMock()
storage = Storage(
host='localhost',
port=5432,
user='postgres',
password='postgres',
dbname='postgres',
min_connections=1,
max_connections=10,
minio_config={
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
},
logger=logger,
notification_handler=notification_handler,
)
assert isinstance(storage, Postgres)
mock_minio_repository.assert_called_once_with(
logger=logger,
notification_handler=notification_handler,
minio_endpoint_url='localhost:9000',
minio_access_key='minio',
minio_secret_key='minio123',
minio_region_name='us-east-1',
minio_default_bucket='test',
)
@patch('laborious.activities.storage.MinioRepository')
def test___init___none_minio_repository(mock_minio_repository, storage):
storage.minio_repository = None
logger = MagicMock()
notification_handler = MagicMock()
storage.__init__(
host='localhost',
port=5432,
user='postgres',
password='postgres',
dbname='postgres',
min_connections=1,
max_connections=10,
minio_config={
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
},
logger=logger,
notification_handler=notification_handler,
)
mock_minio_repository.assert_called_once_with(
logger=logger,
notification_handler=notification_handler,
minio_endpoint_url='localhost:9000',
minio_access_key='minio',
minio_secret_key='minio123',
minio_region_name='us-east-1',
minio_default_bucket='test',
)
@patch('laborious.activities.storage.MinioRepository')
def test___init___done_repository(mock_minio_repository, storage):
storage.__init__(
host='localhost',
port=5432,
user='postgres',
password='postgres',
dbname='postgres',
min_connections=1,
max_connections=10,
minio_config={
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'region_name': 'us-east-1',
'default_bucket': 'test',
},
logger=MagicMock(),
notification_handler=MagicMock(),
)
mock_minio_repository.assert_not_called()
assert storage.minio_repository is not None
@mark.asyncio
async def test_query_to_minio_not_data(storage):
storage.load_custom_query = AsyncMock(return_value=None)
result = await storage.query_to_minio({})
storage.load_custom_query.assert_called_once_with({})
assert result['success'] is False
assert result['message'] == 'No data returned from query'
@mark.asyncio
@patch('laborious.activities.storage.pd.DataFrame')
@patch('laborious.activities.storage.now')
async def test_query_to_minio_success(now, dataframe, storage):
data = [{'a': 1}, {'a': 2}, {'a': 3}]
storage.load_custom_query = AsyncMock(return_value=data)
now.return_value = datetime.datetime(2024, 1, 1, 0, 0, 0)
storage.minio_repository.minio_bucket = 'test'
result = await storage.query_to_minio({'object_prefix': 'test', **metadata})
dataframe.assert_called_once_with(data)
storage.minio_repository.store_dataframe_as_parquet.assert_called_once_with(
dataframe=dataframe.return_value,
uri='s3://test/test_2024-01-01_00-00-00.parquet',
object_name='test_2024-01-01_00-00-00.parquet',
metadata=metadata['metadata'],
)
assert result['success'] is True
assert result['object_key'] == 'test_2024-01-01_00-00-00.parquet'
assert result['uri'] == 's3://test/test_2024-01-01_00-00-00.parquet'
@mark.asyncio
async def test_query_to_minio_error(storage):
storage.send_notification = MagicMock()
storage.load_custom_query = AsyncMock(side_effect=Exception('test'))
result = await storage.query_to_minio({**metadata, 'object_prefix': 'test'})
assert result['success'] is False
assert result['message'] == 'test'
storage.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='ERROR_STORING_QUERY_TO_MINIO',
message='Error storing query to MinIO: test',
block='query_to_minio',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
def test_close(storage):
storage.minio_repository = MagicMock()
storage.close()
assert storage.minio_repository is None
def test___del__(storage):
storage.close = MagicMock()
storage.__del__()
storage.close.assert_called_once()