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

SIENTIAPDE-994

Browse Source 
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.
This commit is contained in:
vitor-aignosi
2025-05-26 11:22:50 -03:00
parent 67fe4afaa6
commit ad00661516
13 changed files with 849 additions and 653 deletions
Download Patch File Download Diff File Expand all files Collapse all files

149
tests/laborious/activities/test_activities.py Normal file
View File

@@ -0,0 +1,149 @@
from pytest import mark
from unittest.mock import patch, MagicMock, ANY
from laborious.activities.activities import Activities
from laborious.activities.postgres import Postgres
from laborious.activities.mlflow import MLFlow
from laborious.activities.gates import Gates
from laborious.activities.opc import OPC
@patch('laborious.activities.activities.Postgres.__init__')
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
@patch('laborious.activities.activities.Gates.__init__')
def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgres_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
'user': 'postgres',
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
}
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group'
}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
opc_config=opc_config,
logger=logger,
notification_handler=notification_handler
)
assert isinstance(activities, Activities)
assert isinstance(activities, Postgres)
assert isinstance(activities, MLFlow)
assert isinstance(activities, OPC)
assert isinstance(activities, Gates)
mock_postgres_init.assert_called_once_with(
ANY,
host=postgres_config['host'],
port=postgres_config['port'],
user=postgres_config['user'],
password=postgres_config['password'],
dbname=postgres_config['dbname'],
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
logger=logger,
notification_handler=notification_handler
)
mock_mlflow_init.assert_called_once_with(
ANY,
mlflow_host=mlflow_config['host'],
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
logger=logger,
notification_handler=notification_handler
)
mock_opc_init.assert_called_once_with(
ANY,
opc_servers=opc_config,
logger=logger,
notification_handler=notification_handler
)
mock_gates_init.assert_called_once_with(
ANY,
logger=logger,
notification_handler=notification_handler
)
@mark.asyncio
@patch('laborious.activities.activities.Postgres.__init__')
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
async def test_prepare_activity(_mock_opc_init,
_mock_mlflow_init, _mock_postgres_init):
postgres_config = {
'host': 'localhost',
'port': 5432,
'user': 'postgres',
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
}
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group'
}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
opc_config=opc_config,
logger=logger,
notification_handler=notification_handler
)
input_data = {
'workflow_name': 'test-workflow-name',
'schedule_name': 'test-schedule-name',
'model_name': 'test-model-name',
'model_id': 'test-model-id'
}
await activities.prepare_activity(input_data)
assert activities.notification_handler.base_notification.pipeline_name == input_data[
'workflow_name']
assert activities.notification_handler.base_notification.schedule_name == input_data[
'schedule_name']
assert activities.notification_handler.base_notification.model_name == input_data[
'model_name']
assert activities.notification_handler.base_notification.model_id == input_data[
'model_id']

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

@@ -1,605 +1,369 @@
from unittest.mock import ANY, MagicMock, patch
from pandas import DataFrame
from unittest.mock import MagicMock, ANY, patch
from pytest import fixture, mark
from laborious.activities.gates import Gates
from sientia_do.notifications.models import NotificationLevel
from laborious.activities.gates import Gates
@fixture
def gates():
def gates_activity():
return Gates(
logger=MagicMock(),
notification_handler=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_specific_variables_null_values_with_stop_policy_only(
input_filter_functions_mock,
gates
):
specific_variables_null_values_mock = MagicMock(return_value=True)
empty_data_mock = MagicMock(return_value=False)
def functions_side_effect(x):
if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
return specific_variables_null_values_mock
if x == 'path_confidence':
return {
'stop': -1,
'continue': 2,
'repeat': -1
}
return empty_data_mock
input_filter_functions_mock.__getitem__.side_effect = functions_side_effect
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': {
'SPECIFIC_VARIABLES_NULL_VALUES': {
'POLICY': 'stop',
'VARIABLES': ['variable2']
}
'EMPTY_DATA': {'POLICY': 'STOP'}
},
'data': {
'variable': ['variable1', 'variable2'],
'value': [1, 2]
},
'path_priority': ['stop', 'continue', 'repeat']
'data': {'value': []},
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
result = await gates.input_gate(input_data)
assert result == ('stop', -1, 'Input data with bad quality')
# Act
result = await gates_activity.input_gate(input_data)
input_args = specific_variables_null_values_mock.call_args
assert input_args[0][0].equals(DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
assert input_args[0][1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
empty_data_mock.assert_not_called()
@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_specific_variables_null_values_with_continue_policy_only(
input_filter_functions_mock,
gates
):
specific_variables_null_values_mock = MagicMock(return_value=True)
empty_data_mock = MagicMock(return_value=False)
def functions_side_effect(x):
if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
return specific_variables_null_values_mock
if x == 'path_confidence':
return {
'stop': -1,
'continue': 2,
'repeat': -1
}
return empty_data_mock
input_filter_functions_mock.__getitem__.side_effect = functions_side_effect
input_data = {
'filters': {
'SPECIFIC_VARIABLES_NULL_VALUES': {
'POLICY': 'continue',
'VARIABLES': ['variable2']
}
},
'data': {
'variable': ['variable1', 'variable2'],
'value': [1, 2]
},
'path_priority': ['stop', 'continue', 'repeat']
}
result = await gates.input_gate(input_data)
assert result == ('continue', 2, 'Input data with bad quality')
input_args = specific_variables_null_values_mock.call_args
assert input_args[0][0].equals(DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
assert input_args[0][1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
empty_data_mock.assert_not_called()
@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_specific_variables_null_values_no_filtered(
input_filter_functions_mock,
gates
):
specific_variables_null_values_mock = MagicMock(return_value=False)
empty_data_mock = MagicMock(return_value=False)
def functions_side_effect(x):
if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
return specific_variables_null_values_mock
if x == 'path_confidence':
return {
'stop': -1,
'continue': 2,
'repeat': -1
}
return empty_data_mock
input_filter_functions_mock.__getitem__.side_effect = functions_side_effect
input_data = {
'filters': {
'SPECIFIC_VARIABLES_NULL_VALUES': {
'POLICY': 'stop',
'VARIABLES': ['variable2']
}
},
'data': {
'variable': ['variable1', 'variable2'],
'value': [1, 2]
},
'path_priority': ['stop', 'continue', 'repeat']
}
result = await gates.input_gate(input_data)
assert result == (None, 0, '')
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
assert specific_variables_null_values_input_args[0][
1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
empty_data_mock.assert_not_called()
@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_one_stop_policy(
input_filter_functions_mock,
gates
):
specific_variables_null_values_mock = MagicMock(return_value=True)
empty_data_mock = MagicMock(return_value=True)
def functions_side_effect(x):
if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
return specific_variables_null_values_mock
if x == 'path_confidence':
return {
'stop': -1,
'continue': 2,
'repeat': -1
}
return empty_data_mock
input_filter_functions_mock.__getitem__.side_effect = functions_side_effect
input_data = {
'filters': {
'SPECIFIC_VARIABLES_NULL_VALUES': {
'POLICY': 'stop',
'VARIABLES': ['variable2']
},
'EMPTY_DATA': {
'POLICY': 'continue',
}
},
'data': {
'variable': ['variable1', 'variable2'],
'value': [1, 2]
},
'path_priority': ['stop', 'continue', 'repeat']
}
result = await gates.input_gate(input_data)
assert result == ('stop', -1, 'Input data with bad quality')
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
assert specific_variables_null_values_input_args[0][
1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
empty_data_input_args = empty_data_mock.call_args
assert empty_data_input_args[0][0].equals(DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
assert empty_data_input_args[0][1] == input_data['filters']['EMPTY_DATA']
@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_one_continue_policy(
input_filter_functions_mock,
gates
):
specific_variables_null_values_mock = MagicMock(return_value=False)
empty_data_mock = MagicMock(return_value=True)
def functions_side_effect(x):
if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
return specific_variables_null_values_mock
if x == 'path_confidence':
return {
'stop': -1,
'continue': 2,
'repeat': -1
}
return empty_data_mock
input_filter_functions_mock.__getitem__.side_effect = functions_side_effect
input_data = {
'filters': {
'SPECIFIC_VARIABLES_NULL_VALUES': {
'POLICY': 'stop',
'VARIABLES': ['variable2']
},
'EMPTY_DATA': {
'POLICY': 'continue',
}
},
'data': {
'variable': ['variable1', 'variable2'],
'value': [1, 2]
},
'path_priority': ['stop', 'continue', 'repeat']
}
result = await gates.input_gate(input_data)
assert result == ('continue', 2, 'Input data with bad quality')
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
assert specific_variables_null_values_input_args[0][
1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
empty_data_input_args = empty_data_mock.call_args
assert empty_data_input_args[0][0].equals(DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
assert empty_data_input_args[0][1] == input_data['filters']['EMPTY_DATA']
@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_no_filtered(
input_filter_functions_mock,
gates
):
specific_variables_null_values_mock = MagicMock(return_value=False)
empty_data_mock = MagicMock(return_value=False)
def functions_side_effect(x):
if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
return specific_variables_null_values_mock
if x == 'path_confidence':
return {
'stop': -1,
'continue': 2,
'repeat': -1
}
return empty_data_mock
input_filter_functions_mock.__getitem__.side_effect = functions_side_effect
input_data = {
'filters': {
'SPECIFIC_VARIABLES_NULL_VALUES': {
'POLICY': 'stop',
'VARIABLES': ['variable2']
},
'EMPTY_DATA': {
'POLICY': 'continue',
}
},
'data': {
'variable': ['variable1', 'variable2'],
'value': [1, 2]
},
'path_priority': ['stop', 'continue', 'repeat']
}
result = await gates.input_gate(input_data)
assert result == (None, 0, '')
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
empty_data_input_args = empty_data_mock.call_args
assert empty_data_input_args[0][0].equals(DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
assert empty_data_input_args[0][1] == input_data['filters']['EMPTY_DATA']
@mark.asyncio
@patch('laborious.activities.gates.input_filter_functions')
async def test_input_gate_error(
input_filter_functions_mock,
gates
):
input_filter_functions_mock.__getitem__.side_effect = KeyError('test')
input_data = {
'filters': {
'SPECIFIC_VARIABLES_NULL_VALUES': {
'POLICY': 'stop',
'VARIABLES': ['variable2']
}
},
'data': {
'variable': ['variable1', 'variable2'],
'value': [1, 2]
},
'path_priority': ['stop', 'continue', 'repeat'],
}
result = await gates.input_gate(input_data)
assert result == (None, 0, '')
gates.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id='INTPUT_GATE_ERROR__SPECIFIC_VARIABLES_NULL_VALUES',
message="Error in filter SPECIFIC_VARIABLES_NULL_VALUES:{'POLICY': 'stop', 'VARIABLES': ['variable2']}: \n 'test'",
block='input_gate',
# 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
)
transform_filter_path_confidence = {
'stop': -1,
'continue': 255,
'repeat': -1
}
@mark.asyncio
@patch('laborious.activities.gates.mlflow_response_filter_functions')
async def test_mlflow_response_gate_no_filtered(
mlflow_response_filter_functions_mock,
gates
):
api_error_filter_mock = MagicMock(return_value=False)
def transform_filter_functions_side_effect(x: str):
if x == 'API_ERROR':
return api_error_filter_mock
if x == 'path_confidence':
return transform_filter_path_confidence
mlflow_response_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect
async def test_input_gate_no_filters(gates_activity):
# Arrange
input_data = {
'filters': {
'API_ERROR': {
'POLICY': 'stop',
}
},
'data': {
'variable': ['variable1', 'variable2'],
'value': [1, 2]
},
'path_priority': ['stop', 'continue', 'repeat'],
'type': 'predict'
'filters': {},
'data': {'value': [1, 2, 3]},
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
}
result = await gates.mlflow_response_gate(input_data)
assert result == (None, 0, '')
# Act
result = await gates_activity.input_gate(input_data)
api_error_filter_mock.assert_called_once_with(
input_data['data'],
input_data['filters']['API_ERROR']
)
# Assert
assert result == (None, 0, "")
gates_activity.logger.debug.assert_called()
gates.notification_handler.build_and_send_notification.assert_not_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_filtered(
mlflow_response_filter_functions_mock,
gates
):
api_error_filter_mock = MagicMock(return_value=True)
def transform_filter_functions_side_effect(x: str):
if x == 'API_ERROR':
return api_error_filter_mock
if x == 'path_confidence':
return transform_filter_path_confidence
mlflow_response_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect
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': {
'API_ERROR': {
'POLICY': 'continue',
}
'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': 'Error',
'traceback': 'Error'
'message': 'API error occurred',
'traceback': 'error trace'
}
},
'path_priority': ['stop', 'continue', 'repeat'],
'type': 'predict'
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
result = await gates.mlflow_response_gate(input_data)
assert result == ('continue', 255, "Error")
# Act
result = await gates_activity.mlflow_response_gate(input_data)
api_error_filter_mock.assert_called_once_with(
input_data['data'],
input_data['filters']['API_ERROR']
)
# 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.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id='PREDICT_GATE_RESPONSE_FILTER__API_ERROR',
message=input_data['data']['content']['message'],
block='mlflow_gate',
level=NotificationLevel.WARNING,
attachment_content=input_data['data']['content']['traceback']
)
@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_no_filtered(
mlflow_content_filter_functions_mock,
gates
):
nan_values_filter_mock = MagicMock(return_value=False)
def transform_filter_functions_side_effect(x: str):
if x == 'NAN_VALUES':
return nan_values_filter_mock
if x == 'path_confidence':
return transform_filter_path_confidence
mlflow_content_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect
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': {
'NAN_VALUES': {
'POLICY': 'repeat',
}
'API_ERROR': {'POLICY': 'STOP'}
},
'data': {
'variable': ['variable1', 'variable2'],
'value': [1, 2]
'success': False,
'content': {
'message': 'API error occurred',
'traceback': 'error trace'
}
},
'path_priority': ['stop', 'continue', 'repeat'],
'type': 'predict'
'type': 'test',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
}
result = await gates.mlflow_content_gate(input_data)
assert result == (None, 0, '')
# Act
result = await gates_activity.mlflow_content_gate(input_data)
nan_values_filter_mock_args = nan_values_filter_mock.call_args
assert nan_values_filter_mock_args[0][0].equals(DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
assert nan_values_filter_mock_args[0][1] == input_data['filters']['NAN_VALUES']
gates.notification_handler.build_and_send_notification.assert_not_called()
# 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
@patch('laborious.activities.gates.mlflow_content_filter_functions')
async def test_mlflow_content_gate_filtered(
mlflow_content_filter_functions_mock,
gates
):
nan_values_filter_mock = MagicMock(return_value=True)
def transform_filter_functions_side_effect(x: str):
if x == 'NAN_VALUES':
return nan_values_filter_mock
if x == 'path_confidence':
return transform_filter_path_confidence
mlflow_content_filter_functions_mock.__getitem__.side_effect = \
transform_filter_functions_side_effect
async def test_mlflow_content_gate_no_filters(gates_activity):
# Arrange
input_data = {
'filters': {
'NAN_VALUES': {
'POLICY': 'repeat',
}
},
'data': {
'variable': ['variable1', 'variable2'],
'value': [1, 2]
},
'path_priority': ['stop', 'continue', 'repeat'],
'type': 'predict'
'filters': {},
'data': {'value': [1, 2, 3]},
'type': 'test',
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
}
result = await gates.mlflow_content_gate(input_data)
# 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 == (
'repeat', -1, "Transformed data not passed the content filter")
nan_values_filter_mock_args = nan_values_filter_mock.call_args
assert nan_values_filter_mock_args[0][0].equals(DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
assert nan_values_filter_mock_args[0][1] == input_data['filters']['NAN_VALUES']
gates.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id='PREDICT_GATE_CONTENT_FILTER__NAN_VALUES',
message="Data not passed the content filter NAN_VALUES:{'POLICY': 'repeat'}",
block='mlflow_gate',
level=NotificationLevel.WARNING,
attachment_content=DataFrame(input_data['data']).to_string()
)
'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
):
async def test_format_prediction(gates_activity):
# Arrange
input_data = {
'data': {
'variable': ['variable1', 'variable2'],
'value': [1, 2]
},
'timestamp': '2021-01-01',
'model_id': 'model_id',
'prediction_confidence': 0.95
'data': {'prediction': [1], 'response_time': [0.1]},
'timestamp': '2023-05-26 11:12:27',
'model_id': 'test_model',
'prediction_confidence': 0.9
}
expected_output = DataFrame(input_data['data'])
expected_output['timestamp'] = input_data['timestamp']
expected_output['model_id'] = input_data['model_id']
expected_output['prediction_confidence'] = input_data['prediction_confidence']
expected_output['prediction_status'] = 'Good'
expected_output['comment'] = ''
# Act
result = await gates_activity.format_prediction(input_data)
result = await gates.format_prediction(input_data)
assert result == expected_output.to_dict()
# 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
):
async def test_format_default_prediction(gates_activity):
# Arrange
input_data = {
'timestamp': '2021-01-01',
'model_id': 'model_id',
'prediction_confidence': 0.95,
'comment': 'Comment'
'timestamp': '2023-05-26 11:12:27',
'model_id': 'test_model',
'prediction_confidence': 0.1,
'comment': 'Test comment'
}
expected_output = DataFrame({
'prediction': [0],
'response_time': [0],
'timestamp': [input_data['timestamp']],
'model_id': [input_data['model_id']],
'prediction_confidence': [input_data['prediction_confidence']],
'prediction_status': ['Bad'],
'comment': [input_data['comment']]
})
# Act
result = await gates_activity.format_default_prediction(input_data)
result = await gates.format_default_prediction(input_data)
assert result == expected_output.to_dict()
# 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(
gates
):
async def test_get_last_timestamp_with_data(gates_activity):
# Arrange
input_data = {
'data': {
'variable': ['variable1', 'variable2'],
'value': [1, 2],
'timestamp': ['2021-01-01', '2021-01-02']
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']
}
}
result = await gates.get_last_timestamp(input_data)
assert result == '2021-01-02'
# 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

229
tests/laborious/activities/test_postgres.py
View File

@@ -1,132 +1,159 @@
from unittest.mock import ANY, MagicMock, patch
from pandas import DataFrame
from pytest import fixture
from pytest import mark
from sientia_do.notifications.models import NotificationLevel
from unittest.mock import MagicMock, patch
from pytest import fixture, mark
import pandas as pd
from laborious.activities.postgres import Postgres
@fixture
@patch("laborious.activities.postgres.ThreadedConnectionPool")
def postgres_client(mock_pool):
@patch("laborious.activities.postgres.create_engine")
def postgres_activity(_mock_create_engine):
return Postgres(
host="localhost",
port=5432,
user="postgres",
password="postgres",
dbname="postgres",
user="test_user",
password="test_password",
dbname="test_db",
min_connections=1,
max_connections=10,
max_connections=5,
logger=MagicMock(),
notification_handler=MagicMock(),
notification_handler=MagicMock()
)
@mark.asyncio
@patch("laborious.activities.postgres.read_sql_query",
return_value=DataFrame([{"a": 1, "b": 2}]))
async def test_load_custom_query_success(mock_read_sql_query, postgres_client):
query = "SELECT * FROM test"
result = await postgres_client.load_custom_query(query)
assert result is not None
assert len(result) > 0
assert result == {'a': {0: 1}, 'b': {0: 2}}
postgres_client.notification_handler.build_and_send_notification.assert_not_called()
@patch("laborious.activities.postgres.read_sql_query")
async def test_load_custom_query_none_data(mock_read_sql_query, postgres_activity):
query = "SELECT * FROM test_table LIMIT 1"
mock_read_sql_query.return_value = None
result = await postgres_activity.load_custom_query(query)
assert isinstance(result, dict)
assert len(result) == 0
@mark.asyncio
@patch("laborious.activities.postgres.read_sql_query",
side_effect=Exception("Error fetching data from query"))
async def test_load_custom_query_error(mock_read_sql_query, postgres_client):
query = "SELECT * FROM test"
result = await postgres_client.load_custom_query(query)
assert result == {}
postgres_client.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="ERROR_LOADING_CUSTOM_QUERY",
message="Error fetching data from query: Error fetching data from query",
block="load_custom_query",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
@patch("laborious.activities.postgres.read_sql_query")
async def test_load_custom_query_date_converted(mock_read_sql_query, postgres_activity):
query = "SELECT * FROM test_table LIMIT 1"
mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]})
mock_data['date'] = pd.to_datetime('2022-01-01')
mock_read_sql_query.return_value = mock_data
result = await postgres_activity.load_custom_query(query)
assert isinstance(result, dict)
assert len(result) == 3
assert "column1" in result
assert "column2" in result
assert "date" in result
assert result['date'] == {0: '2022-01-01 00:00:00'}
@mark.asyncio
async def test_repeat_last_prediction_success(postgres_client):
query_items = {"schema": "test", "table_name": "test", "model": 1}
await postgres_client.repeat_last_prediction(query_items)
postgres_client.notification_handler.build_and_send_notification.assert_not_called()
postgres_client.pool.getconn.assert_called_once()
postgres_client.pool.putconn.assert_called_once()
@patch("laborious.activities.postgres.read_sql_query")
async def test_load_custom_query_success(mock_read_sql_query, postgres_activity):
query = "SELECT * FROM test_table LIMIT 1"
mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]})
postgres_client.pool.getconn.return_value.cursor.assert_called_once()
postgres_client.pool.getconn.return_value.cursor.return_value.execute.assert_called_once_with(
f"""
INSERT INTO \"{query_items['schema']}\".{query_items['table_name']} (model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, created_at)
SELECT model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, NOW()
FROM \"{query_items['schema']}\".{query_items['table_name']}
WHERE model_id = {query_items['model']}
ORDER BY timestamp DESC
LIMIT 1;
"""
)
postgres_client.pool.getconn.return_value.commit.assert_called_once()
postgres_client.pool.getconn.return_value.cursor.return_value.close.assert_called_once()
mock_read_sql_query.return_value = mock_data
result = await postgres_activity.load_custom_query(query)
assert isinstance(result, dict)
assert len(result) == 2
assert "column1" in result
assert "column2" in result
postgres_activity.logger.info.assert_called()
@mark.asyncio
async def test_repeat_last_prediction_error(postgres_client):
postgres_client.pool.getconn.return_value.cursor.return_value.execute.side_effect = Exception(
"Error repeating last prediction")
query_items = {"schema": "test", "table_name": "test", "model": 1}
await postgres_client.repeat_last_prediction(query_items)
postgres_client.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="ERROR_REPEATING_LAST_PREDICTION",
message="Error repeating last prediction: Error repeating last prediction",
block="repeat_last_prediction",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
postgres_client.pool.getconn.assert_called_once()
postgres_client.pool.putconn.assert_called_once()
async def test_load_custom_query_error(postgres_activity):
query = "SELECT * FROM non_existent_table"
error_msg = "Table not found"
with patch("laborious.activities.postgres.read_sql_query", side_effect=ValueError(error_msg)):
result = await postgres_activity.load_custom_query(query)
assert isinstance(result, dict)
assert len(result) == 0
postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
postgres_activity.logger.error.assert_called()
@mark.asyncio
@patch("laborious.activities.postgres.DataFrame")
async def test_export_data_to_postgres_success(mock_dataframe, postgres_client):
data = {"schema": "test", "table_name": "test",
"data": {"a": [1, 2, 3], "b": [4, 5, 6]}}
await postgres_client.export_data_to_postgres(data)
postgres_client.notification_handler.build_and_send_notification.assert_not_called()
postgres_client.pool.getconn.assert_called_once()
postgres_client.pool.putconn.assert_called_once()
async def test_repeat_last_prediction_success(postgres_activity):
query_items = {
"schema": "public",
"table_name": "predictions",
"model": 1
}
mock_dataframe.assert_called_once_with(data["data"])
mock_dataframe.return_value.to_sql.assert_called_once_with(
data["table_name"],
postgres_client.pool.getconn.return_value,
schema=data["schema"],
if_exists="append",
index=False
)
postgres_client.pool.getconn.return_value.commit.assert_called_once()
with patch("sqlalchemy.orm.session.Session.execute") as mock_execute:
await postgres_activity.repeat_last_prediction(query_items)
mock_execute.assert_called_once()
postgres_activity.logger.info.assert_called()
@mark.asyncio
@patch("laborious.activities.postgres.DataFrame", return_value=MagicMock(
to_sql=MagicMock(side_effect=Exception("Error exporting data to postgres"))
))
async def test_export_data_to_postgres_error(mock_dataframe, postgres_client):
data = {"schema": "test", "table_name": "test",
"data": {"a": [1, 2, 3], "b": [4, 5, 6]}}
await postgres_client.export_data_to_postgres(data)
postgres_client.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id="ERROR_EXPORTING_DATA_TO_POSTGRES",
message="Error exporting data to postgres: Error exporting data to postgres",
block="export_data_to_postgres",
level=NotificationLevel.ERROR,
attachment_content=ANY
)
async def test_repeat_last_prediction_error(postgres_activity):
query_items = {
"schema": "public",
"table_name": "predictions",
"model": 1
}
error_msg = "Database error"
postgres_client.pool.getconn.assert_called_once()
postgres_client.pool.putconn.assert_called_once()
with patch("sqlalchemy.orm.session.Session.execute", side_effect=ValueError(error_msg)):
await postgres_activity.repeat_last_prediction(query_items)
postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
postgres_activity.logger.error.assert_called()
@mark.asyncio
async def test_export_data_to_postgres_success(postgres_activity):
input_data = {
"schema": "public",
"table_name": "test_table",
"data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]})
}
with patch("laborious.activities.postgres.DataFrame.to_sql") as mock_to_sql:
await postgres_activity.export_data_to_postgres(input_data)
mock_to_sql.assert_called_once()
postgres_activity.logger.debug.assert_called()
@mark.asyncio
async def test_export_data_to_postgres_error(postgres_activity):
input_data = {
"schema": "public",
"table_name": "test_table",
"data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]})
}
error_msg = "Export failed"
with patch("laborious.activities.postgres.DataFrame.to_sql", side_effect=ValueError(error_msg)):
await postgres_activity.export_data_to_postgres(input_data)
postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
postgres_activity.logger.error.assert_called()
@mark.asyncio
async def test_close(postgres_activity):
postgres_activity.close()
postgres_activity.engine.dispose.assert_called_once()
@mark.asyncio
async def test_del(postgres_activity):
postgres_activity.close = MagicMock()
postgres_activity.__del__()
postgres_activity.close.assert_called_once()