aba4a3f1a5
This commit renames the 'laborious' package to 'model_manager' across the entire project. This includes renaming directories, modules, references in code, configuration files, and documentation to reflect the new package name. This change improves clarity and consistency within the project.
598 lines
17 KiB
Python
598 lines
17 KiB
Python
from unittest.mock import MagicMock, ANY, patch
|
|
from pytest import fixture, mark
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from model_manager.activities.gates import Gates
|
|
|
|
|
|
@fixture
|
|
def gates_activity():
|
|
gates = Gates(
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
)
|
|
gates.error = MagicMock()
|
|
gates.debug = MagicMock()
|
|
gates.info = MagicMock()
|
|
gates.warning = MagicMock()
|
|
gates.critical = MagicMock()
|
|
gates.send_notification = MagicMock()
|
|
return gates
|
|
|
|
|
|
metadata = {
|
|
"metadata": {
|
|
"model_id": "test_model",
|
|
"model_name": "test_model",
|
|
"workflow_name": "test_workflow",
|
|
"schema_name": "test_schedule",
|
|
},
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_input_gate_invalid_filter(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {
|
|
'INVALID_FILTER': {'POLICY': 'STOP'}
|
|
},
|
|
'data': {'value': [1, 2, 3]},
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.input_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
gates_activity.error.assert_called_once_with(
|
|
"Filter INVALID_FILTER not found", metadata['metadata']
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('model_manager.activities.gates.input_filter_functions')
|
|
async def test_input_gate_filter_exception(mock_input_filter_functions, gates_activity):
|
|
# Arrange
|
|
mock_input_filter_functions.__contains__.return_value = True
|
|
mock_input_filter_functions.__getitem__.return_value = MagicMock(
|
|
side_effect=Exception("Test error"))
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {
|
|
'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
|
|
},
|
|
'data': {'value': []},
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.input_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
gates_activity.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id="INTPUT_GATE_ERROR__EMPTY_DATA",
|
|
message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \n Test error",
|
|
block="input_gate",
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_input_gate_no_filters(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {},
|
|
'data': {'value': [1, 2, 3]},
|
|
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.input_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
gates_activity.debug.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_input_gate_with_filter(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {
|
|
'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
|
|
},
|
|
'data': {'value': []},
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.input_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == ('STOP', -1, "Input data with bad quality")
|
|
gates_activity.debug.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_response_gate_invalid_filter(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {
|
|
'INVALID_FILTER': {'POLICY': 'STOP'}
|
|
},
|
|
'data': {'content': {'message': 'success'}},
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_response_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('model_manager.activities.gates.mlflow_response_filter_functions')
|
|
async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter_functions,
|
|
gates_activity):
|
|
# Arrange
|
|
mock_mlflow_response_filter_functions.__contains__.return_value = True
|
|
mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock(
|
|
side_effect=Exception("Test error"))
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {
|
|
'INVALID_FILTER': {'POLICY': 'STOP'}
|
|
},
|
|
'data': {'content': {'message': 'success'}},
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_response_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
gates_activity.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id="MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER",
|
|
message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error",
|
|
block="mlflow_gate",
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_response_gate_no_filters(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {},
|
|
'data': {'content': {'message': 'success'}},
|
|
'type': 'test',
|
|
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_response_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
gates_activity.debug.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_response_gate_with_filter(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {
|
|
'API_ERROR': {'policy': 'STOP'}
|
|
},
|
|
'data': {
|
|
'success': False,
|
|
'content': {
|
|
'message': 'API error occurred',
|
|
'traceback': 'error trace'
|
|
}
|
|
},
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_response_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == ('STOP', -1, "API error occurred")
|
|
gates_activity.debug.assert_called()
|
|
gates_activity.send_notification.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_content_gate_invalid_filter(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {
|
|
'INVALID_FILTER': {'POLICY': 'STOP'}
|
|
},
|
|
'data': {'value': [1, 2, 3]},
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_content_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('model_manager.activities.gates.mlflow_content_filter_functions')
|
|
async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_functions,
|
|
gates_activity):
|
|
# Arrange
|
|
mock_mlflow_content_filter_functions.__contains__.return_value = True
|
|
mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock(
|
|
side_effect=Exception("Test error"))
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {
|
|
'API_ERROR': {'POLICY': 'STOP'}
|
|
},
|
|
'data': {
|
|
'success': False,
|
|
'content': {
|
|
'message': 'API error occurred',
|
|
'traceback': 'error trace'
|
|
}
|
|
},
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_content_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
gates_activity.debug.assert_called()
|
|
gates_activity.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id="MLFLOW_GATE_CONTENT_FILTER__API_ERROR",
|
|
message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error",
|
|
block="mlflow_gate",
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_content_gate_no_filters(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {},
|
|
'data': {'value': [1, 2, 3]},
|
|
'type': 'test',
|
|
'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_content_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, "")
|
|
gates_activity.debug.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_content_gate_with_filter(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {
|
|
'NAN_VALUES': {'policy': 'STOP', 'config': {}}
|
|
},
|
|
'data': {'value': [None, None, None]},
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_content_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (
|
|
'STOP', -1, "Transformed data not passed the content filter")
|
|
gates_activity.debug.assert_called()
|
|
gates_activity.send_notification.assert_called()
|
|
|
|
|
|
def test_get_prediction_store_policy_invalid_policy(gates_activity):
|
|
# Arrange
|
|
prediction_store_policy = 'INVALID_POLICY'
|
|
|
|
# Act
|
|
policy_type, policy_value = gates_activity.get_prediction_store_policy(
|
|
prediction_store_policy, metadata)
|
|
|
|
# Assert
|
|
assert policy_type == 'lts'
|
|
assert policy_value == 1
|
|
|
|
|
|
def test_get_prediction_store_policy_invalid_policy_value(gates_activity):
|
|
# Arrange
|
|
prediction_store_policy = 'abc:INVALID_VALUE'
|
|
|
|
# Act
|
|
policy_type, policy_value = gates_activity.get_prediction_store_policy(
|
|
prediction_store_policy, metadata)
|
|
|
|
# Assert
|
|
assert policy_type == 'lts'
|
|
assert policy_value == 1
|
|
|
|
|
|
def test_get_prediction_store_policy_valid_policy_type(gates_activity):
|
|
# Arrange
|
|
prediction_store_policy = 'abc:1'
|
|
|
|
# Act
|
|
policy_type, policy_value = gates_activity.get_prediction_store_policy(
|
|
prediction_store_policy, metadata)
|
|
|
|
# Assert
|
|
assert policy_type == 'lts'
|
|
assert policy_value == 1
|
|
|
|
|
|
def test_get_prediction_store_policy_valid_policy(gates_activity):
|
|
# Arrange
|
|
prediction_store_policy = 'erl:1'
|
|
|
|
# Act
|
|
policy_type, policy_value = gates_activity.get_prediction_store_policy(
|
|
prediction_store_policy, metadata)
|
|
|
|
# Assert
|
|
assert policy_type == 'erl'
|
|
assert policy_value == 1
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_format_prediction_no_timestamp(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'data': {
|
|
'prediction': {
|
|
'2023-05-26 11:12:27': 1
|
|
},
|
|
'response_time': {
|
|
'2023-05-26 11:12:27': 0.1
|
|
}
|
|
},
|
|
'model_id': 'test_model',
|
|
'prediction_confidence': 0.9,
|
|
'prediction_store_policy': 'lts:1'
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.format_prediction(input_data)
|
|
|
|
# Assert
|
|
assert result['prediction'] == {0: 1}
|
|
assert result['response_time'] == {0: ANY}
|
|
assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
|
|
assert result['model_id'] == {0: 'test_model'}
|
|
assert result['prediction_confidence'] == {0: 0.9}
|
|
assert result['prediction_status'] == {0: 'Good'}
|
|
assert result['comments'] == {0: ""}
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_format_prediction_with_timestamp_erl(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'data': {
|
|
'prediction': {
|
|
'2023-05-26 11:12:27': 1,
|
|
'2023-05-26 11:12:28': 2,
|
|
'2023-05-26 11:12:29': 3,
|
|
},
|
|
'response_time': {
|
|
'2023-05-26 11:12:27': 0.1,
|
|
'2023-05-26 11:12:28': 0.2,
|
|
'2023-05-26 11:12:29': 0.3,
|
|
}
|
|
},
|
|
'model_id': 'test_model',
|
|
'prediction_confidence': 0.9,
|
|
'prediction_store_policy': 'erl:2'
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.format_prediction(input_data)
|
|
|
|
# Assert
|
|
assert result['prediction'] == {0: 2, 1: 1}
|
|
assert result['response_time'] == {0: 0.2, 1: 0.1}
|
|
assert result['timestamp'] == {
|
|
0: '2023-05-26 11:12:28', 1: '2023-05-26 11:12:27'}
|
|
assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
|
|
assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
|
|
assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
|
|
assert result['comments'] == {0: "", 1: ""}
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_format_prediction_with_timestamp_lts(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'data': {
|
|
'prediction': {
|
|
'2023-05-26 11:12:27': 1,
|
|
'2023-05-26 11:12:28': 2,
|
|
'2023-05-26 11:12:29': 3,
|
|
},
|
|
'response_time': {
|
|
'2023-05-26 11:12:27': 0.1,
|
|
'2023-05-26 11:12:28': 0.2,
|
|
'2023-05-26 11:12:29': 0.3,
|
|
}
|
|
},
|
|
'model_id': 'test_model',
|
|
'prediction_confidence': 0.9,
|
|
'prediction_store_policy': 'lts:2'
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.format_prediction(input_data)
|
|
|
|
# Assert
|
|
assert result['prediction'] == {0: 3, 1: 2}
|
|
assert result['response_time'] == {0: 0.3, 1: 0.2}
|
|
assert result['timestamp'] == {
|
|
0: '2023-05-26 11:12:29', 1: '2023-05-26 11:12:28'}
|
|
assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
|
|
assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
|
|
assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
|
|
assert result['comments'] == {0: "", 1: ""}
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_format_prediction_with_timestamp_invalid_policy(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'data': {'prediction': [1, 2, 3],
|
|
'response_time': [0.1, 0.2, 0.3],
|
|
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']},
|
|
'model_id': 'test_model',
|
|
'prediction_confidence': 0.9,
|
|
'prediction_store_policy': 'lts:2'
|
|
}
|
|
gates_activity.get_prediction_store_policy = MagicMock(
|
|
return_value=('invalid', 1))
|
|
|
|
try:
|
|
result = await gates_activity.format_prediction(input_data)
|
|
except ValueError as e:
|
|
assert str(e) == "Invalid policy type: invalid"
|
|
else:
|
|
assert False, "Expected ValueError"
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_format_default_prediction(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'timestamp': '2023-05-26 11:12:27',
|
|
'model_id': 'test_model',
|
|
'prediction_confidence': 0.1,
|
|
'comment': 'Test comment'
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.format_default_prediction(input_data)
|
|
|
|
# Assert
|
|
assert result['prediction'] == {0: 0}
|
|
assert result['response_time'] == {0: 0}
|
|
assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
|
|
assert result['model_id'] == {0: 'test_model'}
|
|
assert result['prediction_confidence'] == {0: 0.1}
|
|
assert result['prediction_status'] == {0: 'Bad'}
|
|
assert result['comments'] == {0: 'Test comment'}
|
|
gates_activity.debug.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_get_last_timestamp_with_data(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'data': {
|
|
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']
|
|
}
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.get_last_timestamp(input_data)
|
|
|
|
# Assert
|
|
assert result == '2023-05-26 11:12:28'
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_get_last_timestamp_no_data(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
'data': {},
|
|
**metadata
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.get_last_timestamp(input_data)
|
|
|
|
# Assert
|
|
assert isinstance(result, str) # Should be a timestamp string
|
|
assert len(result) > 0
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('model_manager.activities.gates.metrics')
|
|
async def test_write_metrics(mock_metrics, gates_activity):
|
|
"""Test write_metrics method."""
|
|
input_data = {
|
|
**metadata,
|
|
'prediction': {
|
|
'prediction': [1, 2, 3],
|
|
'prediction_confidence': [0.9, 0.8, 0.7],
|
|
'response_time': [0.1, 0.2, 0.3]
|
|
}
|
|
}
|
|
await gates_activity.write_metrics(input_data)
|
|
mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.assert_called_once_with(
|
|
pod_id=gates_activity.pod_id,
|
|
model_name=metadata['metadata']['model_name'],
|
|
pipeline_name=metadata['metadata']['workflow_name']
|
|
)
|
|
mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.return_value.inc.assert_called_once_with()
|
|
|
|
mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.assert_called_once_with(
|
|
pod_id=gates_activity.pod_id,
|
|
model_name=metadata['metadata']['model_name'],
|
|
pipeline_name=metadata['metadata']['workflow_name']
|
|
)
|
|
mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.return_value.set.assert_called_once_with(
|
|
0.9
|
|
)
|
|
|
|
mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.assert_called_once_with(
|
|
pod_id=gates_activity.pod_id,
|
|
model_name=metadata['metadata']['model_name'],
|
|
pipeline_name=metadata['metadata']['workflow_name']
|
|
)
|
|
mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
|
|
0.1
|
|
)
|