vitor-aignosi
381856f5ca
SIENTIAPDE-1712 Enhance logging across various classes by adding logger parameters and improving debug statements. This update includes adjustments in Gates, MLFlow, ModelMetrics, and MLFlowRepository classes for better traceability and observability during operations.
1009 lines
30 KiB
Python
1009 lines
30 KiB
Python
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
|
|
|
from pandas import DataFrame
|
|
from pytest import fixture, mark
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
|
|
from laborious.activities.gates import Gates
|
|
|
|
|
|
@fixture(autouse=True)
|
|
def _passthrough_from_dict():
|
|
with patch(
|
|
'laborious.activities.gates.MinioDataFramePayload.from_dict', side_effect=lambda x: x
|
|
):
|
|
yield
|
|
|
|
|
|
def _minio_payload(retrieve_return, status=None):
|
|
"""
|
|
Build a MinioDataFramePayload-like test double with async retrieve.
|
|
|
|
Args:
|
|
retrieve_return: Value returned from await retrieve(minio_repo, metadata).
|
|
status: Optional status dict for MLflow response gate (payload.status).
|
|
|
|
Return:
|
|
MagicMock: Object with async retrieve and optional status.
|
|
"""
|
|
p = MagicMock()
|
|
p.retrieve = AsyncMock(return_value=retrieve_return)
|
|
p.status = status
|
|
return p
|
|
|
|
|
|
@fixture
|
|
def gates_activity():
|
|
gates = Gates(
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
metrics_controller=AsyncMock(),
|
|
)
|
|
gates.error = MagicMock()
|
|
gates.debug = MagicMock()
|
|
gates.info = MagicMock()
|
|
gates.warning = MagicMock()
|
|
gates.critical = MagicMock()
|
|
gates.send_notification = MagicMock()
|
|
gates.send_notification_async = AsyncMock()
|
|
gates.emit_metric = AsyncMock()
|
|
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': _minio_payload(DataFrame({'value': [1, 2, 3]})),
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.input_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, '')
|
|
gates_activity.error.assert_called_once_with(
|
|
'Filter INVALID_FILTER not found', metadata['metadata']
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.gates.input_filter_functions')
|
|
async def test_input_gate_filter_exception(mock_input_filter_functions, gates_activity):
|
|
# Arrange
|
|
mock_input_filter_functions.__contains__.return_value = True
|
|
mock_input_filter_functions.__getitem__.return_value = MagicMock(
|
|
side_effect=Exception('Test error')
|
|
)
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}},
|
|
'data': _minio_payload(DataFrame({'value': []})),
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.input_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, '')
|
|
gates_activity.send_notification_async.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': _minio_payload(DataFrame({'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': _minio_payload(DataFrame({'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_input_gate_with_filter_lowercase_keys(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
|
|
'data': _minio_payload(DataFrame({'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')
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_input_gate_with_filter_capitalized_keys(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {'EMPTY_DATA': {'Policy': 'STOP', 'Config': {}}},
|
|
'data': _minio_payload(DataFrame({'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')
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_input_gate_with_filter_not_caught(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}},
|
|
'data': _minio_payload(DataFrame({'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.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': _minio_payload(
|
|
{'content': {'message': 'success'}},
|
|
status={'success': True},
|
|
),
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_response_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, '')
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.gates.mlflow_response_filter_functions')
|
|
async def test_mlflow_response_gate_filter_exception(
|
|
mock_mlflow_response_filter_functions, gates_activity
|
|
):
|
|
# Arrange
|
|
mock_mlflow_response_filter_functions.__contains__.return_value = True
|
|
mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock(
|
|
side_effect=Exception('Test error')
|
|
)
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
|
|
'data': _minio_payload(
|
|
{'content': {'message': 'success'}},
|
|
status={'success': True},
|
|
),
|
|
'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_async.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': _minio_payload(
|
|
{'content': {'message': 'success'}},
|
|
status={'success': True},
|
|
),
|
|
'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': _minio_payload(
|
|
{'content': {'message': 'API error occurred', 'traceback': 'error trace'}},
|
|
status={'success': False, 'message': 'API error occurred'},
|
|
),
|
|
'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_async.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_response_gate_with_filter_capitalized_keys(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {'API_ERROR': {'Policy': 'STOP'}},
|
|
'data': _minio_payload(
|
|
{'content': {'message': 'API error occurred', 'traceback': 'error trace'}},
|
|
status={'success': False, 'message': 'API error occurred'},
|
|
),
|
|
'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')
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_response_gate_with_filter_not_caught(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {'API_ERROR': {'POLICY': 'STOP'}},
|
|
'data': _minio_payload(
|
|
{'content': {'message': 'success'}},
|
|
status={'success': True},
|
|
),
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', '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_content_gate_invalid_filter(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
|
|
'data': _minio_payload(DataFrame({'value': [1, 2, 3]})),
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.mlflow_content_gate(input_data)
|
|
|
|
# Assert
|
|
assert result == (None, 0, '')
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.gates.mlflow_content_filter_functions')
|
|
async def test_mlflow_content_gate_filter_exception(
|
|
mock_mlflow_content_filter_functions, gates_activity
|
|
):
|
|
# Arrange
|
|
mock_mlflow_content_filter_functions.__contains__.return_value = True
|
|
mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock(
|
|
side_effect=Exception('Test error')
|
|
)
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {'API_ERROR': {'POLICY': 'STOP'}},
|
|
'data': _minio_payload(DataFrame({'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, '')
|
|
gates_activity.debug.assert_called()
|
|
gates_activity.send_notification_async.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': _minio_payload(DataFrame({'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': _minio_payload(DataFrame({'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_async.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_content_gate_with_filter_not_caught(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {'API_ERROR': {'POLICY': 'STOP'}},
|
|
'data': _minio_payload(DataFrame({'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, '')
|
|
gates_activity.debug.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_mlflow_content_gate_filter_returns_false(gates_activity):
|
|
input_data = {
|
|
**metadata,
|
|
'filters': {'NAN_VALUES': {'POLICY': 'STOP', 'CONFIG': {}}},
|
|
'data': _minio_payload(DataFrame({'value': [1, 2, 3]})),
|
|
'type': 'test',
|
|
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
|
}
|
|
|
|
result = await gates_activity.mlflow_content_gate(input_data)
|
|
|
|
assert result == (None, 0, '')
|
|
gates_activity.debug.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': _minio_payload(
|
|
DataFrame(
|
|
{
|
|
'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',
|
|
'timestamp': '2023-05-26 11:12:27',
|
|
}
|
|
|
|
# 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': _minio_payload(
|
|
DataFrame(
|
|
{
|
|
'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',
|
|
'timestamp': '2023-05-26 11:12:27',
|
|
}
|
|
|
|
# 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': _minio_payload(
|
|
DataFrame(
|
|
{
|
|
'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',
|
|
'timestamp': '2023-05-26 11:12:27',
|
|
}
|
|
|
|
# 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': _minio_payload(
|
|
DataFrame(
|
|
{
|
|
'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',
|
|
'timestamp': '2023-05-26 11:12:27',
|
|
}
|
|
gates_activity.get_prediction_store_policy = MagicMock(return_value=('invalid', 1))
|
|
|
|
try:
|
|
await gates_activity.format_prediction(input_data)
|
|
except ValueError as e:
|
|
assert str(e) == 'Invalid policy type: invalid'
|
|
else:
|
|
raise AssertionError('Expected ValueError')
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.gates.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock)
|
|
async def test_format_transformed_data_single_row(mock_from_dataframe, gates_activity):
|
|
# Arrange
|
|
payload_result = MagicMock()
|
|
mock_from_dataframe.return_value = payload_result
|
|
input_data = {
|
|
**metadata,
|
|
'data': _minio_payload(
|
|
DataFrame(
|
|
{
|
|
'var1': {'2023-05-26 11:12:27': 1.0},
|
|
'var2': {'2023-05-26 11:12:27': 2.0},
|
|
}
|
|
)
|
|
),
|
|
'model_id': 'test_model',
|
|
'model_name': 'test_model',
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.format_transformed_data(input_data)
|
|
|
|
# Assert
|
|
assert result is payload_result
|
|
mock_from_dataframe.assert_called_once()
|
|
kwargs = mock_from_dataframe.call_args.kwargs
|
|
assert kwargs['model_name'] == 'test_model'
|
|
assert kwargs['operation'] == 'transform'
|
|
assert kwargs['workflow_metadata'] == metadata['metadata']
|
|
assert kwargs['minio_repo'] is gates_activity.minio_repository
|
|
assert 'dataframe' in kwargs
|
|
gates_activity.info.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.gates.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock)
|
|
async def test_format_transformed_data_multiple_rows(mock_from_dataframe, gates_activity):
|
|
# Arrange
|
|
payload_result = MagicMock()
|
|
mock_from_dataframe.return_value = payload_result
|
|
input_data = {
|
|
**metadata,
|
|
'data': _minio_payload(
|
|
DataFrame(
|
|
{
|
|
'var1': {
|
|
'2023-05-26 11:12:27': 1.0,
|
|
'2023-05-26 11:12:28': 2.0,
|
|
},
|
|
'var2': {
|
|
'2023-05-26 11:12:27': 3.0,
|
|
'2023-05-26 11:12:28': 4.0,
|
|
},
|
|
}
|
|
)
|
|
),
|
|
'model_id': 'test_model',
|
|
'model_name': 'test_model',
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.format_transformed_data(input_data)
|
|
|
|
# Assert
|
|
assert result is payload_result
|
|
mock_from_dataframe.assert_called_once()
|
|
kwargs = mock_from_dataframe.call_args.kwargs
|
|
assert kwargs['model_name'] == 'test_model'
|
|
assert kwargs['operation'] == 'transform'
|
|
assert kwargs['workflow_metadata'] == metadata['metadata']
|
|
assert kwargs['minio_repo'] is gates_activity.minio_repository
|
|
assert 'dataframe' in kwargs
|
|
gates_activity.info.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.gates.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock)
|
|
async def test_format_transformed_data_empty_data(mock_from_dataframe, gates_activity):
|
|
# Arrange
|
|
payload_result = MagicMock()
|
|
mock_from_dataframe.return_value = payload_result
|
|
input_data = {
|
|
**metadata,
|
|
'data': _minio_payload(DataFrame()),
|
|
'model_id': 'test_model',
|
|
'model_name': 'test_model',
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.format_transformed_data(input_data)
|
|
|
|
# Assert
|
|
assert result is payload_result
|
|
mock_from_dataframe.assert_called_once()
|
|
kwargs = mock_from_dataframe.call_args.kwargs
|
|
assert kwargs['model_name'] == 'test_model'
|
|
assert kwargs['operation'] == 'transform'
|
|
assert kwargs['workflow_metadata'] == metadata['metadata']
|
|
assert kwargs['minio_repo'] is gates_activity.minio_repository
|
|
assert 'dataframe' in kwargs
|
|
gates_activity.info.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,
|
|
'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_format_retrain_report(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'experiment_response': {
|
|
'success': True,
|
|
'timestamp': '2023-05-26 11:12:27',
|
|
'message': 'success',
|
|
},
|
|
'update_report': {
|
|
'version': '1.0.0',
|
|
'mlflow_run_id': 'test_mlflow_run_id',
|
|
'mlflow_experiment_id': 'test_mlflow_experiment_id',
|
|
},
|
|
'model_id': 'test_model',
|
|
'model_name': 'test_model',
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.format_retrain_report(input_data)
|
|
|
|
# Assert
|
|
assert result['model_id'] == {0: 'test_model'}
|
|
assert result['model_name'] == {0: 'test_model'}
|
|
assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
|
|
assert result['status'] == {0: 'success'}
|
|
assert result['version'] == {0: '1.0.0'}
|
|
assert result['mlflow_run_id'] == {0: 'test_mlflow_run_id'}
|
|
assert result['mlflow_experiment_id'] == {0: 'test_mlflow_experiment_id'}
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_format_retrain_report_failure(gates_activity):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'experiment_response': {
|
|
'success': False,
|
|
'timestamp': '2023-05-26 11:12:27',
|
|
'message': 'failure',
|
|
},
|
|
'update_report': {
|
|
'version': '1.0.0',
|
|
'mlflow_run_id': 'test_mlflow_run_id',
|
|
'mlflow_experiment_id': 'test_mlflow_experiment_id',
|
|
},
|
|
'model_id': 'test_model',
|
|
'model_name': 'test_model',
|
|
}
|
|
|
|
# Act
|
|
result = await gates_activity.format_retrain_report(input_data)
|
|
|
|
# Assert
|
|
assert result['model_id'] == {0: 'test_model'}
|
|
assert result['model_name'] == {0: 'test_model'}
|
|
assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
|
|
assert result['status'] == {0: 'failure'}
|
|
assert 'version' not in result
|
|
assert 'mlflow_run_id' not in result
|
|
assert 'mlflow_experiment_id' not in result
|
|
gates_activity.info.assert_called()
|
|
gates_activity.debug.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.gates.metrics')
|
|
async def test_write_metrics(mock_metrics, gates_activity):
|
|
"""Test write_metrics method."""
|
|
input_data = {
|
|
**metadata,
|
|
'prediction': {
|
|
'prediction': [1, 2, 3],
|
|
'prediction_confidence': [0.9, 0.8, 0.7],
|
|
'response_time': [0.1, 0.2, 0.3],
|
|
},
|
|
'opc_metrics': {'server1': {'tag1': 0.1, 'tag2': 0.2}},
|
|
}
|
|
await gates_activity.write_metrics(input_data)
|
|
core_tags = {
|
|
'pod_id': gates_activity.pod_id,
|
|
'runtime': gates_activity.runtime,
|
|
'operation_type': 'predict',
|
|
'model_name': metadata['metadata']['model_name'],
|
|
'workflow_name': metadata['metadata']['workflow_name'],
|
|
}
|
|
gates_activity.emit_metric.assert_has_calls(
|
|
[
|
|
call(
|
|
metric_object=mock_metrics.PREDICTIONS_WRITTEN_COUNT,
|
|
tags=core_tags,
|
|
),
|
|
]
|
|
)
|
|
gates_activity.emit_metric.assert_has_calls(
|
|
[
|
|
call(
|
|
metric_object=mock_metrics.PREDICTION_CONFIDENCE_MONITOR,
|
|
method='set',
|
|
tags=core_tags,
|
|
value=0.9,
|
|
),
|
|
]
|
|
)
|
|
gates_activity.emit_metric.assert_has_calls(
|
|
[
|
|
call(
|
|
metric_object=mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR,
|
|
method='observe',
|
|
tags=core_tags,
|
|
value=0.1,
|
|
),
|
|
]
|
|
)
|
|
gates_activity.emit_metric.assert_has_calls(
|
|
[
|
|
call(
|
|
metric_object=mock_metrics.PREDICTION_OPC_WRITING_COUNT,
|
|
tags={
|
|
**core_tags,
|
|
'opc_server_id': 'server1',
|
|
'tag': 'tag1',
|
|
},
|
|
),
|
|
]
|
|
)
|
|
gates_activity.emit_metric.assert_has_calls(
|
|
[
|
|
call(
|
|
metric_object=mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
|
|
method='observe',
|
|
tags={
|
|
**core_tags,
|
|
'opc_server_id': 'server1',
|
|
'tag': 'tag1',
|
|
},
|
|
value=0.1,
|
|
),
|
|
]
|
|
)
|
|
gates_activity.emit_metric.assert_has_calls(
|
|
[
|
|
call(
|
|
metric_object=mock_metrics.PREDICTION_OPC_WRITING_COUNT,
|
|
tags={
|
|
**core_tags,
|
|
'opc_server_id': 'server1',
|
|
'tag': 'tag2',
|
|
},
|
|
),
|
|
]
|
|
)
|
|
gates_activity.emit_metric.assert_has_calls(
|
|
[
|
|
call(
|
|
metric_object=mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
|
|
method='observe',
|
|
tags={
|
|
**core_tags,
|
|
'opc_server_id': 'server1',
|
|
'tag': 'tag2',
|
|
},
|
|
value=0.2,
|
|
),
|
|
]
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.gates.metrics')
|
|
async def test_write_metrics_with_none_opc_response_time(mock_metrics, gates_activity):
|
|
"""Test write_metrics method with None response_time in opc_metrics."""
|
|
input_data = {
|
|
**metadata,
|
|
'prediction': {
|
|
'prediction': [1],
|
|
'prediction_confidence': [0.9],
|
|
'response_time': [0.1],
|
|
},
|
|
'opc_metrics': {'server1': {'tag1': 0.1, 'tag2': None}},
|
|
}
|
|
await gates_activity.write_metrics(input_data)
|
|
|
|
# Verify that metrics for tag1 are emitted
|
|
gates_activity.emit_metric.assert_any_call(
|
|
metric_object=mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
|
|
method='observe',
|
|
tags={
|
|
'pod_id': gates_activity.pod_id,
|
|
'runtime': gates_activity.runtime,
|
|
'operation_type': 'predict',
|
|
'model_name': metadata['metadata']['model_name'],
|
|
'workflow_name': metadata['metadata']['workflow_name'],
|
|
'opc_server_id': 'server1',
|
|
'tag': 'tag1',
|
|
},
|
|
value=0.1,
|
|
)
|
|
|
|
# Verify that metrics for tag2 (with None response_time) are NOT emitted
|
|
calls = [
|
|
c
|
|
for c in gates_activity.emit_metric.call_args_list
|
|
if len(c[1].get('tags', {})) > 0 and c[1]['tags'].get('tag') == 'tag2'
|
|
]
|
|
assert len(calls) == 0, 'Metrics should not be emitted for None response_time'
|