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

SIENTIAPDE-1110

Browse Source 
Refactor notification handling in Gates and OPC activities to use send_notification method with metadata integration, enhancing error reporting and traceability.
This commit is contained in:
vitor-aignosi
2025-07-04 15:19:14 -03:00
parent 58216a91ce
commit 50e6479c49
12 changed files with 293 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
laborious/activities/gates.py
View File

@@ -90,7 +90,8 @@ class Gates(BaseActivity):
filter_output.append(config['policy'])
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
self.send_notification(
metadata=metadata,
notification_id=f"INTPUT_GATE_ERROR__{fil}",
message=f"Error in filter {fil}:{config}: \n {e}",
block="input_gate",
@@ -145,7 +146,8 @@ class Gates(BaseActivity):
if mlflow_response_filter_functions[fil](data, config):
filter_output.append(config['policy'])
comments.append(data['content']['message'])
self.notification_handler.build_and_send_notification(
self.send_notification(
metadata=metadata,
notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
message=data['content']['message'],
block="mlflow_gate",
@@ -154,7 +156,8 @@ class Gates(BaseActivity):
)
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
self.send_notification(
metadata=metadata,
notification_id=f"MLFLOW_GATE_RESPONSE_FILTER__{fil}",
message=f"Error in filter {fil}:{config}: \n {e}",
block="mlflow_gate",
@@ -208,7 +211,8 @@ class Gates(BaseActivity):
try:
if mlflow_content_filter_functions[fil](data, config):
filter_output.append(config['policy'])
self.notification_handler.build_and_send_notification(
self.send_notification(
metadata=metadata,
notification_id=f"{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}",
message=f"Data not passed the content filter {fil}:{config}",
block="mlflow_gate",
@@ -217,7 +221,8 @@ class Gates(BaseActivity):
)
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
self.send_notification(
metadata=metadata,
notification_id=f"MLFLOW_GATE_CONTENT_FILTER__{fil}",
message=f"Error in filter {fil}:{config}: \n {e}",
block="mlflow_gate",

3
laborious/activities/opc.py
View File

@@ -60,7 +60,8 @@ class OPC(BaseActivity):
tag, data, data_type, self.logger, metadata)
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
self.send_notification(
metadata=metadata,
notification_id=f"WRITE_OPC_{tag_type.upper()}_ERROR",
message=f"Error writing data to OPC server: {e}",
block="write_opc_data",

6
laborious/utils/repository/opc_repository.py
View File

@@ -217,7 +217,7 @@ class OpcRepository():
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id=f"OPC_WRITE_GET_NODE_ERROR_{self.id}",
message=f"Failed to get node from OPC server: {e}",
message=f"Failed to get node from OPC server: {e} | metadata: {metadata}",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=trace
@@ -229,7 +229,7 @@ class OpcRepository():
if data_type not in data_type_map:
self.notification_handler.build_and_send_notification(
notification_id=f"OPC_WRITE_DATA_TYPE_ERROR_{self.id}",
message=f"Unsupported data type: {data_type}",
message=f"Unsupported data type: {data_type} | metadata: {metadata}",
block="opc_repository",
level=NotificationLevel.ERROR
)
@@ -258,7 +258,7 @@ class OpcRepository():
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id=f"OPC_WRITE_DATA_ERROR_{self.id}",
message=f"Failed to write data to OPC server: {e}",
message=f"Failed to write data to OPC server: {e} | metadata: {metadata}",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=trace

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

@@ -90,65 +90,6 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
)
@mark.asyncio
@patch('laborious.activities.activities.Postgres.__init__')
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
async def test_prepare_activity(_mock_opc_init,
_mock_mlflow_init, _mock_postgres_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
'user': 'postgres',
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
}
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group'
}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
opc_config=opc_config,
logger=logger,
notification_handler=notification_handler
)
input_data = {
'workflow_name': 'test-workflow-name',
'schedule_name': 'test-schedule-name',
'model_name': 'test-model-name',
'model_id': 'test-model-id'
}
await activities.prepare_activity(input_data)
assert activities.notification_handler.base_notification.pipeline == input_data[
'workflow_name']
assert activities.notification_handler.base_notification.trigger == input_data[
'schedule_name']
assert activities.notification_handler.base_notification.model_name == input_data[
'model_name']
assert activities.notification_handler.base_notification.model_id == input_data[
'model_id']
@patch('laborious.activities.activities.Postgres', return_value=MagicMock())
@patch('laborious.activities.activities.MLFlow', return_value=MagicMock())
@patch('laborious.activities.activities.OPC', return_value=MagicMock())

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

@@ -6,16 +6,34 @@ from laborious.activities.gates import Gates
@fixture
def gates_activity():
return Gates(
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'}
},
@@ -28,8 +46,8 @@ async def test_input_gate_invalid_filter(gates_activity):
# Assert
assert result == (None, 0, "")
gates_activity.logger.error.assert_called_once_with(
"Filter INVALID_FILTER not found"
gates_activity.error.assert_called_once_with(
"Filter INVALID_FILTER not found", metadata['metadata']
)
@@ -41,6 +59,7 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac
mock_input_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception("Test error"))
input_data = {
**metadata,
'filters': {
'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
},
@@ -53,7 +72,8 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac
# Assert
assert result == (None, 0, "")
gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
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",
@@ -66,6 +86,7 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac
async def test_input_gate_no_filters(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {},
'data': {'value': [1, 2, 3]},
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
@@ -76,13 +97,14 @@ async def test_input_gate_no_filters(gates_activity):
# Assert
assert result == (None, 0, "")
gates_activity.logger.debug.assert_called()
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': {}}
},
@@ -95,13 +117,14 @@ async def test_input_gate_with_filter(gates_activity):
# Assert
assert result == ('STOP', -1, "Input data with bad quality")
gates_activity.logger.debug.assert_called()
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'}
},
@@ -126,6 +149,7 @@ async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter
mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception("Test error"))
input_data = {
**metadata,
'filters': {
'INVALID_FILTER': {'POLICY': 'STOP'}
},
@@ -139,7 +163,8 @@ async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter
# Assert
assert result == (None, 0, "")
gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
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",
@@ -152,6 +177,7 @@ async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter
async def test_mlflow_response_gate_no_filters(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {},
'data': {'content': {'message': 'success'}},
'type': 'test',
@@ -163,13 +189,14 @@ async def test_mlflow_response_gate_no_filters(gates_activity):
# Assert
assert result == (None, 0, "")
gates_activity.logger.debug.assert_called()
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'}
},
@@ -189,14 +216,15 @@ async def test_mlflow_response_gate_with_filter(gates_activity):
# Assert
assert result == ('STOP', -1, "API error occurred")
gates_activity.logger.debug.assert_called()
gates_activity.notification_handler.build_and_send_notification.assert_called()
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'}
},
@@ -221,6 +249,7 @@ async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_f
mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock(
side_effect=Exception("Test error"))
input_data = {
**metadata,
'filters': {
'API_ERROR': {'POLICY': 'STOP'}
},
@@ -240,8 +269,9 @@ async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_f
# Assert
assert result == (None, 0, "")
gates_activity.logger.debug.assert_called()
gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
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",
@@ -254,6 +284,7 @@ async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_f
async def test_mlflow_content_gate_no_filters(gates_activity):
# Arrange
input_data = {
**metadata,
'filters': {},
'data': {'value': [1, 2, 3]},
'type': 'test',
@@ -265,13 +296,14 @@ async def test_mlflow_content_gate_no_filters(gates_activity):
# Assert
assert result == (None, 0, "")
gates_activity.logger.debug.assert_called()
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': {}}
},
@@ -286,14 +318,15 @@ async def test_mlflow_content_gate_with_filter(gates_activity):
# Assert
assert result == (
'STOP', -1, "Transformed data not passed the content filter")
gates_activity.logger.debug.assert_called()
gates_activity.notification_handler.build_and_send_notification.assert_called()
gates_activity.debug.assert_called()
gates_activity.send_notification.assert_called()
@mark.asyncio
async def test_format_prediction(gates_activity):
# Arrange
input_data = {
**metadata,
'data': {'prediction': [1], 'response_time': [0.1]},
'timestamp': '2023-05-26 11:12:27',
'model_id': 'test_model',
@@ -311,13 +344,14 @@ async def test_format_prediction(gates_activity):
assert result['prediction_confidence'] == {0: 0.9}
assert result['prediction_status'] == {0: 'Good'}
assert result['comments'] == {0: ""}
gates_activity.logger.debug.assert_called()
gates_activity.debug.assert_called()
@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,
@@ -335,13 +369,14 @@ async def test_format_default_prediction(gates_activity):
assert result['prediction_confidence'] == {0: 0.1}
assert result['prediction_status'] == {0: 'Bad'}
assert result['comments'] == {0: 'Test comment'}
gates_activity.logger.debug.assert_called()
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']
}

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

@@ -39,6 +39,16 @@ def mlflow(mock_mlflow_repository):
)
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")
@@ -46,6 +56,7 @@ async def test_request_transform(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'},
@@ -100,6 +111,7 @@ async def test_request_predict(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},
{'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0},

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

@@ -5,6 +5,15 @@ from laborious.activities.opc 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",
},
}
@patch("laborious.activities.opc.OpcRepository")
def test___init__(mock_opc_repository):
@@ -15,6 +24,7 @@ def test___init__(mock_opc_repository):
mock_notification_handler = MagicMock()
servers = {
'server1': {
'id': 'server1',
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
@@ -23,6 +33,7 @@ def test___init__(mock_opc_repository):
'reconnection_interval': 60,
},
'server2': {
'id': 'server2',
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
@@ -45,7 +56,7 @@ def test___init__(mock_opc_repository):
mock_opc_repository.assert_has_calls([
call(
name="server1",
id="server1",
url="http://localhost:8080",
logger=mock_logger,
server_uri="opc.tcp://localhost:4840",
@@ -58,7 +69,7 @@ def test___init__(mock_opc_repository):
])
mock_opc_repository.assert_has_calls([
call(
name="server2",
id="server2",
url="http://localhost:8080",
logger=mock_logger,
server_uri="opc.tcp://localhost:4840",
@@ -79,6 +90,7 @@ def test___init__(mock_opc_repository):
def opc(mock_opc_repository):
servers = {
'server1': {
'id': 'server1',
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
@@ -90,12 +102,15 @@ def opc(mock_opc_repository):
mock_opc_repository.write_data = MagicMock(
return_value=True
)
return OPC(
opc = OPC(
opc_servers=servers,
logger=MagicMock(),
notification_handler=MagicMock()
)
opc.send_notification = MagicMock()
return opc
WRITE_DATA_CASES = [
('tag1', 'int', 50),
@@ -107,10 +122,10 @@ WRITE_DATA_CASES = [
@mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
def test_write_data_success(opc, tag, data_type, data):
assert opc.write_data(server='server1', tag=tag, data=data,
data_type=data_type, tag_type='prediction')
assert opc.write_data(server_id='server1', tag=tag, data=data,
data_type=data_type, tag_type='prediction', metadata=metadata)
opc.opc_repository['server1'].write_data.assert_called_once_with(
tag, data, data_type)
tag, data, data_type, opc.logger, metadata)
def test_write_data_exception(opc):
@@ -118,11 +133,12 @@ def test_write_data_exception(opc):
"Test error")
try:
opc.write_data(server='server1', tag='tag1', data=50,
data_type='int', tag_type='prediction')
opc.write_data(server_id='server1', tag='tag1', data=50,
data_type='int', tag_type='prediction', metadata=metadata)
except Exception:
opc.notification_handler.build_and_send_notification.assert_called_once_with(
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",
@@ -138,6 +154,7 @@ def test_write_data_exception(opc):
async def test_write_opc_data_success(opc):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
@@ -163,19 +180,21 @@ async def test_write_opc_data_success(opc):
assert output == {'data': 'data'}
opc.write_data.assert_has_calls([
call(
server='server1',
server_id='server1',
tag='tag1',
data=0.75,
data_type='float',
tag_type='prediction'
tag_type='prediction',
metadata=metadata['metadata']
)])
opc.write_data.assert_has_calls([
call(
server='server1',
server_id='server1',
tag='tag2',
data=0.95,
data_type='float',
tag_type='confidence'
tag_type='confidence',
metadata=metadata['metadata']
)
])
assert opc.write_data.call_count == 2
@@ -185,6 +204,7 @@ async def test_write_opc_data_success(opc):
async def test_write_opc_data_empty_config(opc):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
@@ -209,7 +229,7 @@ async def test_write_opc_data_empty_config(opc):
])
def test_process_confidence(opc, data, success, expected):
# Act
result = opc.process_confidence(data, success)
result = opc.process_confidence(data, success, metadata)
# Assert
assert result['prediction_confidence'][0] == expected

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

@@ -14,7 +14,7 @@ def mock_logger():
@fixture
def opc_repository(mock_logger):
return OpcRepository(
name="test_repo",
id="test_repo",
url="opc.tcp://localhost:4840",
logger=mock_logger,
notification_handler=Mock(),
@@ -34,8 +34,18 @@ def mock_client():
yield client_instance
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
},
}
def test_init(opc_repository):
assert opc_repository.name == "test_repo"
assert opc_repository.id == "test_repo"
assert opc_repository.url == "opc.tcp://localhost:4840"
assert opc_repository.server_uri == "urn:test:server"
assert opc_repository.cert_path == "/path/to/cert.pem"
@@ -110,7 +120,7 @@ def test_try_connect_fail(opc_repository):
opc_repository.client.connect.assert_called_once()
assert opc_repository.last_reconnection_time is not None
opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id=f"OPC_CONNECTION_ERROR_{opc_repository.name}",
notification_id=f"OPC_CONNECTION_ERROR_{opc_repository.id}",
message="Failed to connect to OPC server: Test error",
block="opc_repository",
level=NotificationLevel.ERROR,
@@ -126,6 +136,17 @@ def test_disconnect(opc_repository, mock_client):
assert opc_repository.client is None
def test_disconnect_error(opc_repository, mock_client):
opc_repository.client = mock_client
mock_client.disconnect.side_effect = Exception("Test error")
opc_repository.disconnect()
opc_repository.logger.error.assert_called_once_with(
"Failed to disconnect from OPC server: Test error"
)
assert opc_repository.client is None
def test_validate_connection_none_client(opc_repository):
opc_repository.client = None
opc_repository.connect = MagicMock()
@@ -174,11 +195,11 @@ def test_validate_connection_lost_time_to_reconect(_mock_datetime, opc_repositor
opc_repository.client = MagicMock()
opc_repository.client.aio_obj.uaclient.protocol = None
opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0)
opc_repository.try_connect = MagicMock()
opc_repository.connect = MagicMock()
response = opc_repository.validate_connection()
opc_repository.try_connect.assert_called_once()
assert response == opc_repository.try_connect.return_value
opc_repository.connect.assert_called_once()
assert response == opc_repository.connect.return_value
def test_validate_connection_failed(opc_repository):
@@ -192,7 +213,8 @@ def test_validate_connection_failed(opc_repository):
def test_write_data_validate_connection_do_nothing(opc_repository):
opc_repository.validate_connection = MagicMock(return_value=True)
opc_repository.client = MagicMock()
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata)
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
@@ -201,7 +223,8 @@ def test_write_data_validate_connection_failed(opc_repository):
opc_repository.validate_connection = MagicMock(return_value=False)
opc_repository.client = MagicMock()
opc_repository.error_count = 0
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata)
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_not_called()
@@ -211,12 +234,13 @@ def test_write_data_get_node_failed(opc_repository):
opc_repository.client = MagicMock()
opc_repository.error_count = 0
opc_repository.client.get_node.side_effect = Exception("Test error")
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata)
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id=f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.name}",
message="Failed to get node from OPC server: Test error",
notification_id=f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}",
message="Failed to get node from OPC server: Test error | metadata: {'metadata': {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}}",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=ANY
@@ -230,13 +254,14 @@ def test_write_data_invalid_data_type(opc_repository, mock_client):
mock_node = MagicMock()
mock_client.get_node.return_value = mock_node
opc_repository.write_data("ns=2;s=TestNode", 42.0, "invalid_type")
opc_repository.write_data("ns=2;s=TestNode", 42.0,
"invalid_type", opc_repository.logger, metadata)
opc_repository.validate_connection.assert_called_once()
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id=f"OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.name}",
message="Unsupported data type: invalid_type",
notification_id=f"OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}",
message="Unsupported data type: invalid_type | metadata: {'metadata': {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}}",
block="opc_repository",
level=NotificationLevel.ERROR
)
@@ -248,12 +273,11 @@ def test_write_data(opc_repository, mock_client):
mock_node = MagicMock()
mock_client.get_node.return_value = mock_node
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata)
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
mock_node.write_value.assert_called_once()
opc_repository.logger.info.assert_called_once_with(
"Writing 42.0 - <class 'float'> to " + str(mock_node))
def test_write_data_write_value_failed(opc_repository, mock_client):
@@ -263,13 +287,14 @@ def test_write_data_write_value_failed(opc_repository, mock_client):
opc_repository.error_count = 0
mock_client.get_node.return_value = mock_node
mock_node.write_value.side_effect = Exception("Test error")
opc_repository.write_data("ns=2;s=TestNode", 42.0, "float")
opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata)
opc_repository.validate_connection.assert_called_once()
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
mock_node.write_value.assert_called_once()
opc_repository.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id=f"OPC_WRITE_DATA_ERROR_{opc_repository.name}",
message="Failed to write data to OPC server: Test error",
notification_id=f"OPC_WRITE_DATA_ERROR_{opc_repository.id}",
message="Failed to write data to OPC server: Test error | metadata: {'metadata': {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}}",
block="opc_repository",
level=NotificationLevel.ERROR,
attachment_content=ANY

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

@@ -54,7 +54,7 @@ def test_build_opc_config_with_env_vars():
def test_build_opc_config_with_individual_env_vars():
# Arrange
environ.pop('OPC_CONFIG', None)
environ['OPC_NAME'] = 'test-name'
environ['OPC_ID'] = '1'
environ['OPC_URL'] = 'opc.tcp://test:4840'
environ['OPC_SERVER_URI'] = 'opc.tcp://test:4840'
environ['OPC_RECONNECTION_INTERVAL'] = '300'
@@ -63,16 +63,16 @@ def test_build_opc_config_with_individual_env_vars():
config = build_opc_config()
# Assert
assert config['opc']['name'] == 'test-name'
assert config['opc']['url'] == 'opc.tcp://test:4840'
assert config['opc']['server_uri'] == 'opc.tcp://test:4840'
assert config['opc']['reconnection_interval'] == 300
assert config['1']['id'] == '1'
assert config['1']['url'] == 'opc.tcp://test:4840'
assert config['1']['server_uri'] == 'opc.tcp://test:4840'
assert config['1']['reconnection_interval'] == 300
def test_build_opc_config_with_defaults():
# Arrange
environ.pop('OPC_CONFIG', None)
environ.pop('OPC_NAME', None)
environ.pop('OPC_ID', None)
environ.pop('OPC_URL', None)
environ.pop('OPC_SERVER_URI', None)
environ.pop('OPC_RECONNECTION_INTERVAL', None)
@@ -81,10 +81,10 @@ def test_build_opc_config_with_defaults():
config = build_opc_config()
# Assert
assert config['opc']['name'] == 'opc'
assert config['opc']['url'] == 'opc.tcp://localhost:4840'
assert config['opc']['server_uri'] == 'opc.tcp://localhost:4840'
assert config['opc']['reconnection_interval'] == 120
assert config['1']['id'] == '1'
assert config['1']['url'] == 'opc.tcp://localhost:4840'
assert config['1']['server_uri'] == 'opc.tcp://localhost:4840'
assert config['1']['reconnection_interval'] == 120
def test_build_postgres_config_with_env_vars():

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

@@ -10,11 +10,22 @@ def format_and_export_prediction():
return FormatAndExportPrediction()
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
},
}
@mark.asyncio
@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
input_data = {
'metadata': metadata,
"path_flag": None,
"data": {"test": "data"},
"timestamp": "2021-01-01",
@@ -35,7 +46,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
'data': input_data['data'],
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence']
'prediction_confidence': input_data['prediction_confidence'],
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -46,7 +58,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -59,7 +72,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
{
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': workflow_mock.execute_activity_method.return_value
'data': workflow_mock.execute_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -74,6 +88,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
input_data = {
'metadata': metadata,
"path_flag": "default",
"data": {"test": "data"},
"timestamp": "2021-01-01",
@@ -95,7 +110,8 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'comment': input_data['comment']
'comment': input_data['comment'],
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -107,7 +123,8 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -120,7 +137,8 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
{
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': workflow_mock.execute_activity_method.return_value
'data': workflow_mock.execute_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY

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

@@ -9,12 +9,23 @@ def prediction_process():
return PredictionProcess()
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
},
}
@mark.asyncio
@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
async def test_run(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(return_value=False)
# Arrange
input_data = {
'metadata': metadata,
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
@@ -49,51 +60,61 @@ async def test_run(workflow_mock, prediction_process):
assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
retry_policy=ANY, start_to_close_timeout=ANY)])
call(Activities.get_last_timestamp, {
'data': input_data['data'],
**metadata
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
'model_retention': input_data['model_retention'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, {
'data': 'transformed_data',
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
'model_retention': input_data['model_retention'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_called_once_with(
'format_and_export_prediction',
{
'metadata': metadata,
'path_flag': 'continue',
'data': 'predicted_data',
'prediction_confidence': 0.95,
@@ -115,6 +136,7 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(return_value=True)
# Arrange
input_data = {
'metadata': metadata,
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
@@ -141,11 +163,15 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
assert workflow_mock.execute_local_activity_method.call_count == 2
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
'data': input_data['data']}, retry_policy=ANY, start_to_close_timeout=ANY),
'data': input_data['data'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY),
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']}, retry_policy=ANY, start_to_close_timeout=ANY)
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -156,6 +182,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True])
# Arrange
input_data = {
'metadata': metadata,
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
@@ -183,19 +210,26 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 4
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
call(Activities.get_last_timestamp, {
'data': input_data['data'],
**metadata
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']},
'path_priority': input_data['path_priority'],
**metadata
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']},
'model_retention': input_data['model_retention'],
**metadata
},
retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -203,7 +237,8 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -216,6 +251,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
side_effect=[False, False, True])
# Arrange
input_data = {
'metadata': metadata,
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
@@ -247,33 +283,41 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
assert workflow_mock.execute_local_activity_method.call_count == 5
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
call(Activities.get_last_timestamp, {
'data': input_data['data'],
**metadata
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']},
'path_priority': input_data['path_priority'],
**metadata
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
'model_retention': input_data['model_retention'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -285,6 +329,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
side_effect=[False, False, False, True])
# Arrange
input_data = {
'metadata': metadata,
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
@@ -317,46 +362,56 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
call(Activities.get_last_timestamp, {
'data': input_data['data'],
**metadata
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']},
'path_priority': input_data['path_priority'],
**metadata
},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
'model_retention': input_data['model_retention'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, {
'data': 'transformed_data',
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
'model_retention': input_data['model_retention'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -378,6 +433,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
'metadata': metadata,
'schema': schema,
'table_name': table_name,
'model_id': model,
@@ -410,6 +466,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
'metadata': metadata,
'schema': schema,
'table_name': table_name,
'model_id': model,
@@ -424,6 +481,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
workflow_mock.execute_activity_method.assert_called_once_with(
Activities.repeat_last_prediction,
{
**metadata,
'schema': schema,
'table_name': table_name,
'model': model,
@@ -452,6 +510,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
'metadata': metadata,
'schema': schema,
'table_name': table_name,
'model_id': model,
@@ -468,6 +527,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
workflow_mock.execute_child_workflow.assert_called_once_with(
'format_and_export_prediction',
{
'metadata': metadata,
'path_flag': path_flag,
'data': data,
'prediction_confidence': confidence,
@@ -500,6 +560,7 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
**metadata,
'schema': schema,
'table_name': table_name,
'model_id': model,

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

@@ -9,6 +9,16 @@ def predictions_batch() -> PredictionsBatch:
return PredictionsBatch()
metadata = {
"metadata": {
"model_id": "test_model_id",
"model_name": "test_model",
"workflow_name": "predictions_batch",
"schedule_name": "test_schedule",
},
}
@mark.asyncio
@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
@@ -29,27 +39,17 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.prepare_activity,
Activities.load_custom_query,
{
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'predictions_batch'
**metadata,
'query': input_data['query'],
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.load_custom_query,
input_data['query'],
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
prediction_input = {
'metadata': metadata,
'data': {'data': 'test_data'},
'schema': input_data['schema'],
'table_name': input_data['table_name'],