SIENTIAPDE-1243: Initial commit of the model manager project, adding core files and configurations.
This commit introduces the initial project structure, including: - .env.example: Example environment configuration. - .github/workflows/quality-gate.yml: CI workflow for quality checks. - .gitignore: Specifies intentionally untracked files that Git should ignore. - Makefile: Automation of tasks like docker builds. - README.md: Project documentation. - Source code for model management, activities, utils, worker and workflows. - Test suite. - Dockerfile for the simulator. - sonar-project.properties: SonarQube configuration file. - values.yaml: Helm chart values for deployment.
This commit is contained in:
Bruno Domingues
0
tests/laborious/activities/__init__.py
Normal file
0
tests/laborious/activities/__init__.py
Normal file
135
tests/laborious/activities/test_activities.py
Normal file
135
tests/laborious/activities/test_activities.py
Normal file
@@ -0,0 +1,135 @@
|
||||
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.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', 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):
|
||||
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
|
||||
)
|
||||
|
||||
await activities.shutdown()
|
||||
mock_opc_init.shutdown.assert_called_once()
|
||||
mock_postgres_init.close.assert_called_once()
|
||||
597
tests/laborious/activities/test_gates.py
Normal file
597
tests/laborious/activities/test_gates.py
Normal file
@@ -0,0 +1,597 @@
|
||||
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():
|
||||
gates = Gates(
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
)
|
||||
gates.error = MagicMock()
|
||||
gates.debug = MagicMock()
|
||||
gates.info = MagicMock()
|
||||
gates.warning = MagicMock()
|
||||
gates.critical = MagicMock()
|
||||
gates.send_notification = MagicMock()
|
||||
return gates
|
||||
|
||||
|
||||
metadata = {
|
||||
"metadata": {
|
||||
"model_id": "test_model",
|
||||
"model_name": "test_model",
|
||||
"workflow_name": "test_workflow",
|
||||
"schema_name": "test_schedule",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_input_gate_invalid_filter(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'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.error.assert_called_once_with(
|
||||
"Filter INVALID_FILTER not found", metadata['metadata']
|
||||
)
|
||||
|
||||
|
||||
@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 = {
|
||||
**metadata,
|
||||
'filters': {
|
||||
'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
|
||||
},
|
||||
'data': {'value': []},
|
||||
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await gates_activity.input_gate(input_data)
|
||||
|
||||
# Assert
|
||||
assert result == (None, 0, "")
|
||||
gates_activity.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id="INTPUT_GATE_ERROR__EMPTY_DATA",
|
||||
message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \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 = {
|
||||
**metadata,
|
||||
'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.debug.assert_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_input_gate_with_filter(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'filters': {
|
||||
'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
|
||||
},
|
||||
'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.debug.assert_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_mlflow_response_gate_invalid_filter(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'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 = {
|
||||
**metadata,
|
||||
'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.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
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 = {
|
||||
**metadata,
|
||||
'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.debug.assert_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_mlflow_response_gate_with_filter(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'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.debug.assert_called()
|
||||
gates_activity.send_notification.assert_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_mlflow_content_gate_invalid_filter(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'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 = {
|
||||
**metadata,
|
||||
'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.debug.assert_called()
|
||||
gates_activity.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
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 = {
|
||||
**metadata,
|
||||
'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.debug.assert_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_mlflow_content_gate_with_filter(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'filters': {
|
||||
'NAN_VALUES': {'policy': 'STOP', 'config': {}}
|
||||
},
|
||||
'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.debug.assert_called()
|
||||
gates_activity.send_notification.assert_called()
|
||||
|
||||
|
||||
def test_get_prediction_store_policy_invalid_policy(gates_activity):
|
||||
# Arrange
|
||||
prediction_store_policy = 'INVALID_POLICY'
|
||||
|
||||
# Act
|
||||
policy_type, policy_value = gates_activity.get_prediction_store_policy(
|
||||
prediction_store_policy, metadata)
|
||||
|
||||
# Assert
|
||||
assert policy_type == 'lts'
|
||||
assert policy_value == 1
|
||||
|
||||
|
||||
def test_get_prediction_store_policy_invalid_policy_value(gates_activity):
|
||||
# Arrange
|
||||
prediction_store_policy = 'abc:INVALID_VALUE'
|
||||
|
||||
# Act
|
||||
policy_type, policy_value = gates_activity.get_prediction_store_policy(
|
||||
prediction_store_policy, metadata)
|
||||
|
||||
# Assert
|
||||
assert policy_type == 'lts'
|
||||
assert policy_value == 1
|
||||
|
||||
|
||||
def test_get_prediction_store_policy_valid_policy_type(gates_activity):
|
||||
# Arrange
|
||||
prediction_store_policy = 'abc:1'
|
||||
|
||||
# Act
|
||||
policy_type, policy_value = gates_activity.get_prediction_store_policy(
|
||||
prediction_store_policy, metadata)
|
||||
|
||||
# Assert
|
||||
assert policy_type == 'lts'
|
||||
assert policy_value == 1
|
||||
|
||||
|
||||
def test_get_prediction_store_policy_valid_policy(gates_activity):
|
||||
# Arrange
|
||||
prediction_store_policy = 'erl:1'
|
||||
|
||||
# Act
|
||||
policy_type, policy_value = gates_activity.get_prediction_store_policy(
|
||||
prediction_store_policy, metadata)
|
||||
|
||||
# Assert
|
||||
assert policy_type == 'erl'
|
||||
assert policy_value == 1
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_format_prediction_no_timestamp(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'data': {
|
||||
'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'
|
||||
}
|
||||
|
||||
# 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: ""}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_format_prediction_with_timestamp_erl(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'data': {
|
||||
'prediction': {
|
||||
'2023-05-26 11:12:27': 1,
|
||||
'2023-05-26 11:12:28': 2,
|
||||
'2023-05-26 11:12:29': 3,
|
||||
},
|
||||
'response_time': {
|
||||
'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'
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await gates_activity.format_prediction(input_data)
|
||||
|
||||
# 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['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: ""}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_format_prediction_with_timestamp_lts(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'data': {
|
||||
'prediction': {
|
||||
'2023-05-26 11:12:27': 1,
|
||||
'2023-05-26 11:12:28': 2,
|
||||
'2023-05-26 11:12:29': 3,
|
||||
},
|
||||
'response_time': {
|
||||
'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'
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await gates_activity.format_prediction(input_data)
|
||||
|
||||
# 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['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: ""}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
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']},
|
||||
'model_id': 'test_model',
|
||||
'prediction_confidence': 0.9,
|
||||
'prediction_store_policy': 'lts:2'
|
||||
}
|
||||
gates_activity.get_prediction_store_policy = MagicMock(
|
||||
return_value=('invalid', 1))
|
||||
|
||||
try:
|
||||
result = await gates_activity.format_prediction(input_data)
|
||||
except ValueError as e:
|
||||
assert str(e) == "Invalid policy type: invalid"
|
||||
else:
|
||||
assert False, "Expected ValueError"
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_format_default_prediction(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'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.debug.assert_called()
|
||||
|
||||
|
||||
@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']
|
||||
}
|
||||
}
|
||||
|
||||
# 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': {},
|
||||
**metadata
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await gates_activity.get_last_timestamp(input_data)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, str) # Should be a timestamp string
|
||||
assert len(result) > 0
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.gates.metrics')
|
||||
async def test_write_metrics(mock_metrics, gates_activity):
|
||||
"""Test write_metrics method."""
|
||||
input_data = {
|
||||
**metadata,
|
||||
'prediction': {
|
||||
'prediction': [1, 2, 3],
|
||||
'prediction_confidence': [0.9, 0.8, 0.7],
|
||||
'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']
|
||||
)
|
||||
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
|
||||
)
|
||||
|
||||
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']
|
||||
)
|
||||
mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
|
||||
0.1
|
||||
)
|
||||
297
tests/laborious/activities/test_mlflow.py
Normal file
297
tests/laborious/activities/test_mlflow.py
Normal file
@@ -0,0 +1,297 @@
|
||||
from datetime import datetime
|
||||
from unittest.mock import ANY, MagicMock, 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 sientia_do.notifications.models import NotificationLevel
|
||||
|
||||
|
||||
@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", ANY
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
@patch("laborious.activities.mlflow.MLFlowRepository")
|
||||
def mlflow(mock_mlflow_repository):
|
||||
mlflow = MLFlow(
|
||||
mlflow_host="http://localhost:5000",
|
||||
mlflow_port=5000,
|
||||
mlflow_username="admin",
|
||||
mlflow_password="admin",
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock()
|
||||
)
|
||||
|
||||
mlflow.send_notification = MagicMock()
|
||||
|
||||
return mlflow
|
||||
|
||||
|
||||
metadata = {
|
||||
"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")
|
||||
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'}
|
||||
],
|
||||
'model_name': 'test_model',
|
||||
'model_config': {}
|
||||
}
|
||||
|
||||
# Mock the transform response
|
||||
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
|
||||
mock_dataframe.return_value.drop_duplicates.return_value = mock_dataframe.return_value
|
||||
|
||||
# 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, {}, metadata['metadata']
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@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"
|
||||
},
|
||||
"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': {}
|
||||
}
|
||||
|
||||
# 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
|
||||
)
|
||||
mock_dataframe.return_value.__setitem__.assert_any_call(
|
||||
'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
# 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, {}, metadata['metadata']
|
||||
)
|
||||
|
||||
|
||||
@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]
|
||||
}
|
||||
|
||||
mlflow.model_monitoring_repository.retrain_model.return_value = (
|
||||
'Model retrained successfully', 'test')
|
||||
|
||||
response = await mlflow.retrain_model({
|
||||
**metadata,
|
||||
'data': data,
|
||||
'model_name': 'test_model'
|
||||
})
|
||||
|
||||
mlflow.model_monitoring_repository.retrain_model.assert_called_once()
|
||||
|
||||
assert response == {
|
||||
"status": 'Model retrained successfully',
|
||||
"timestamp": 2,
|
||||
"experiment": 'test'
|
||||
}
|
||||
|
||||
|
||||
@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]
|
||||
}
|
||||
|
||||
try:
|
||||
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"
|
||||
|
||||
|
||||
@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'
|
||||
}
|
||||
|
||||
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')
|
||||
|
||||
assert response == {
|
||||
'data1': {0: 1},
|
||||
'data2': {0: 2},
|
||||
'model_id': {0: 1},
|
||||
'model_name': {0: 'test_model'},
|
||||
'timestamp': {0: 2},
|
||||
'status': {0: 'success'}
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_update_production_model_error(mlflow):
|
||||
mlflow.model_monitoring_repository.update_production_model.side_effect = Exception(
|
||||
'Error updating production model'
|
||||
)
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'model_name': 'test_model',
|
||||
'model_id': 1,
|
||||
'experiment': 'test',
|
||||
'timestamp': 2,
|
||||
'status': 'success'
|
||||
}
|
||||
|
||||
try:
|
||||
await mlflow.update_production_model(input_data)
|
||||
except Exception as e:
|
||||
assert str(e) == 'Error updating production model'
|
||||
mlflow.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
|
||||
message='Error updating production model test_model: Error updating production model',
|
||||
block='update_production_model',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY
|
||||
)
|
||||
else:
|
||||
assert False, "No exception raised"
|
||||
369
tests/laborious/activities/test_opc.py
Normal file
369
tests/laborious/activities/test_opc.py
Normal file
@@ -0,0 +1,369 @@
|
||||
from unittest.mock import patch, MagicMock, ANY, call, AsyncMock
|
||||
from pandas import DataFrame
|
||||
from pytest import fixture, mark
|
||||
import pytest_asyncio
|
||||
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",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test__init__():
|
||||
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")
|
||||
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, {}))
|
||||
)
|
||||
server2 = MagicMock(
|
||||
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, {}))
|
||||
)
|
||||
mock_opc_repository.side_effect = [server1, server2, server3]
|
||||
mock_notification_handler = MagicMock()
|
||||
servers = {
|
||||
'server1': {
|
||||
'id': 'server1',
|
||||
'url': 'http://localhost:8080',
|
||||
'server_uri': 'opc.tcp://localhost:4840',
|
||||
'cert_path': '',
|
||||
'private_key_path': '',
|
||||
'server_cert_path': '',
|
||||
'reconnection_interval': 60,
|
||||
},
|
||||
'server2': {
|
||||
'id': 'server2',
|
||||
'url': 'http://localhost:8080',
|
||||
'server_uri': 'opc.tcp://localhost:4840',
|
||||
'cert_path': '',
|
||||
'private_key_path': '',
|
||||
'server_cert_path': '',
|
||||
'reconnection_interval': 60,
|
||||
},
|
||||
'server3': {
|
||||
'id': 'server3',
|
||||
'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
|
||||
)
|
||||
await opc.init_opc()
|
||||
|
||||
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(
|
||||
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'
|
||||
)
|
||||
])
|
||||
|
||||
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
|
||||
)
|
||||
])
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
@patch("laborious.activities.opc.OpcRepository")
|
||||
async def opc(mock_opc_repository):
|
||||
servers = {
|
||||
'server1': {
|
||||
'id': 'server1',
|
||||
'url': 'http://localhost:8080',
|
||||
'server_uri': 'opc.tcp://localhost:4840',
|
||||
'cert_path': '',
|
||||
'private_key_path': '',
|
||||
'server_cert_path': '',
|
||||
'reconnection_interval': 60,
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
|
||||
|
||||
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)
|
||||
@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)
|
||||
assert result is True
|
||||
opc.opc_repository['server1'].write_data.assert_called_once_with(
|
||||
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'
|
||||
})
|
||||
|
||||
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",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_write_data_exception(opc):
|
||||
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)
|
||||
|
||||
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",
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY
|
||||
)
|
||||
|
||||
else:
|
||||
assert False, "Expected an exception to be raised"
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_write_opc_data_success(opc):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'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 = AsyncMock(return_value=True)
|
||||
opc.process_confidence = MagicMock(return_value={'data': 'data'})
|
||||
output = await opc.write_opc_data(input_data)
|
||||
|
||||
# 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']
|
||||
)
|
||||
])
|
||||
assert opc.write_data.call_count == 2
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_write_opc_data_empty_config(opc):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'data': {
|
||||
'prediction': [0.75],
|
||||
'prediction_confidence': [0.95]
|
||||
},
|
||||
'opc_servers': ['server1'],
|
||||
'opc_output_config': {
|
||||
'server1': {
|
||||
'prediction_tags': {},
|
||||
'confidence_tags': {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Act
|
||||
await opc.write_opc_data(input_data)
|
||||
|
||||
# Assert
|
||||
opc.opc_repository['server1'].write_data.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
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]
|
||||
},
|
||||
'opc_output_config': {
|
||||
'server1': {
|
||||
'prediction_tags': {
|
||||
'tag1': {'data_type': 'float'}
|
||||
},
|
||||
'confidence_tags': {
|
||||
'tag2': {'data_type': 'float'}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Act
|
||||
await opc.write_opc_data(input_data)
|
||||
|
||||
# Assert
|
||||
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),
|
||||
])
|
||||
def test_process_confidence(opc, data, success, expected):
|
||||
# Act
|
||||
result = opc.process_confidence(data, success, metadata)
|
||||
|
||||
# Assert
|
||||
assert result['prediction_confidence'][0] == expected
|
||||
|
||||
|
||||
def test_validate_server(opc):
|
||||
assert opc.validate_server('server1', metadata) is True
|
||||
assert opc.validate_server('server2', metadata) is False
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_shutdown(opc):
|
||||
opc.opc_repository['server1'].disconnect = AsyncMock(return_value=True)
|
||||
await opc.shutdown()
|
||||
opc.opc_repository['server1'].disconnect.assert_called_once()
|
||||
Reference in New Issue
Block a user