vitor-aignosi
ad00661516
Refactor tests for Postgres activities and improve error handling - Updated test_postgres.py to enhance the testing of load_custom_query method, including cases for None data and date conversion. - Refactored repeat_last_prediction tests to use mocks for SQLAlchemy session execution. - Added tests for export_data_to_postgres method, covering both success and error scenarios. - Improved the initialization tests for Activities class to ensure proper instantiation of dependencies. - Enhanced test coverage for OPC repository connection validation. - Updated tests for prediction workflows to streamline input handling and improve clarity. - Introduced tests for connectors configuration to validate environment variable handling for MLFlow, OPC, and Postgres. - Added tests for logger utility to ensure default settings are correctly applied.
370 lines
10 KiB
Python
370 lines
10 KiB
Python
from unittest.mock import MagicMock, ANY, patch
|
|
from pytest import fixture, mark
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from laborious.activities.gates import Gates
|
|
|
|
|
|
@fixture
|
|
def gates_activity():
|
|
return Gates(
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_input_gate_invalid_filter(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
'filters': {
|
|
'INVALID_FILTER': {'POLICY': 'STOP'}
|
|
},
|
|
'data': {'value': [1, 2, 3]},
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.input_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
gates_activity.logger.error.assert_called_once_with(
|
|
"Filter INVALID_FILTER not found"
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.gates.input_filter_functions')
|
|
async def test_input_gate_filter_exception(mock_input_filter_functions, gates_activity):
|
|
# Arrange
|
|
mock_input_filter_functions.__contains__.return_value = True
|
|
mock_input_filter_functions.__getitem__.return_value = MagicMock(
|
|
side_effect=Exception("Test error"))
|
|
input_data = {
|
|
'filters': {
|
|
'EMPTY_DATA': {'POLICY': 'STOP'}
|
|
},
|
|
'data': {'value': []},
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.input_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
|
|
notification_id="INTPUT_GATE_ERROR__EMPTY_DATA",
|
|
message="Error in filter EMPTY_DATA:{'POLICY': 'STOP'}: \n Test error",
|
|
block="input_gate",
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_input_gate_no_filters(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
'filters': {},
|
|
'data': {'value': [1, 2, 3]},
|
|
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.input_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
gates_activity.logger.debug.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_input_gate_with_filter(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
'filters': {
|
|
'EMPTY_DATA': {'POLICY': 'STOP'}
|
|
},
|
|
'data': {'value': []},
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.input_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == ('STOP', -1, "Input data with bad quality")
|
|
gates_activity.logger.debug.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_response_gate_invalid_filter(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
'filters': {
|
|
'INVALID_FILTER': {'POLICY': 'STOP'}
|
|
},
|
|
'data': {'content': {'message': 'success'}},
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_response_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.gates.mlflow_response_filter_functions')
|
|
async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter_functions,
|
|
gates_activity):
|
|
# Arrange
|
|
mock_mlflow_response_filter_functions.__contains__.return_value = True
|
|
mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock(
|
|
side_effect=Exception("Test error"))
|
|
input_data = {
|
|
'filters': {
|
|
'INVALID_FILTER': {'POLICY': 'STOP'}
|
|
},
|
|
'data': {'content': {'message': 'success'}},
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_response_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
|
|
notification_id="MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER",
|
|
message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error",
|
|
block="mlflow_gate",
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_response_gate_no_filters(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
'filters': {},
|
|
'data': {'content': {'message': 'success'}},
|
|
'type': 'test',
|
|
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_response_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
gates_activity.logger.debug.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_response_gate_with_filter(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
'filters': {
|
|
'API_ERROR': {'POLICY': 'STOP'}
|
|
},
|
|
'data': {
|
|
'success': False,
|
|
'content': {
|
|
'message': 'API error occurred',
|
|
'traceback': 'error trace'
|
|
}
|
|
},
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_response_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == ('STOP', -1, "API error occurred")
|
|
gates_activity.logger.debug.assert_called()
|
|
gates_activity.notification_handler.build_and_send_notification.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_content_gate_invalid_filter(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
'filters': {
|
|
'INVALID_FILTER': {'POLICY': 'STOP'}
|
|
},
|
|
'data': {'value': [1, 2, 3]},
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_content_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.gates.mlflow_content_filter_functions')
|
|
async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_functions,
|
|
gates_activity):
|
|
# Arrange
|
|
mock_mlflow_content_filter_functions.__contains__.return_value = True
|
|
mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock(
|
|
side_effect=Exception("Test error"))
|
|
input_data = {
|
|
'filters': {
|
|
'API_ERROR': {'POLICY': 'STOP'}
|
|
},
|
|
'data': {
|
|
'success': False,
|
|
'content': {
|
|
'message': 'API error occurred',
|
|
'traceback': 'error trace'
|
|
}
|
|
},
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_content_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
gates_activity.logger.debug.assert_called()
|
|
gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
|
|
notification_id="MLFLOW_GATE_CONTENT_FILTER__API_ERROR",
|
|
message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error",
|
|
block="mlflow_gate",
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_content_gate_no_filters(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
'filters': {},
|
|
'data': {'value': [1, 2, 3]},
|
|
'type': 'test',
|
|
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_content_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
gates_activity.logger.debug.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_content_gate_with_filter(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
'filters': {
|
|
'NAN_VALUES': {'POLICY': 'STOP'}
|
|
},
|
|
'data': {'value': [None, None, None]},
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_content_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (
|
|
'STOP', -1, "Transformed data not passed the content filter")
|
|
gates_activity.logger.debug.assert_called()
|
|
gates_activity.notification_handler.build_and_send_notification.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_format_prediction(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
'data': {'prediction': [1], 'response_time': [0.1]},
|
|
'timestamp': '2023-05-26 11:12:27',
|
|
'model_id': 'test_model',
|
|
'prediction_confidence': 0.9
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.format_prediction(input_data)
|
|
|
|
# Assert
|
|
assert result['prediction'] == {0: 1}
|
|
assert result['response_time'] == {0: ANY}
|
|
assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
|
|
assert result['model_id'] == {0: 'test_model'}
|
|
assert result['prediction_confidence'] == {0: 0.9}
|
|
assert result['prediction_status'] == {0: 'Good'}
|
|
assert result['comments'] == {0: ""}
|
|
gates_activity.logger.debug.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_format_default_prediction(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
'timestamp': '2023-05-26 11:12:27',
|
|
'model_id': 'test_model',
|
|
'prediction_confidence': 0.1,
|
|
'comment': 'Test comment'
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.format_default_prediction(input_data)
|
|
|
|
# Assert
|
|
assert result['prediction'] == {0: 0}
|
|
assert result['response_time'] == {0: 0}
|
|
assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
|
|
assert result['model_id'] == {0: 'test_model'}
|
|
assert result['prediction_confidence'] == {0: 0.1}
|
|
assert result['prediction_status'] == {0: 'Bad'}
|
|
assert result['comments'] == {0: 'Test comment'}
|
|
gates_activity.logger.debug.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_get_last_timestamp_with_data(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
'data': {
|
|
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']
|
|
}
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.get_last_timestamp(input_data)
|
|
|
|
# Assert
|
|
assert result == '2023-05-26 11:12:28'
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_get_last_timestamp_no_data(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
'data': {}
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.get_last_timestamp(input_data)
|
|
|
|
# Assert
|
|
assert isinstance(result, str) # Should be a timestamp string
|
|
assert len(result) > 0
|