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

SIENTIAPDE-1250: Refactor: Remove 'laborious' directory from test structure and update README.md accordingly.

Browse Source
This commit is contained in:
Bruno Domingues
2025-10-07 10:02:38 -03:00
parent 851f34c80e
commit 2628dc92d2
23 changed files with 11 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
tests/activities/__init__.py Normal file
View File

139
tests/activities/test_activities.py Normal file
View File

@@ -0,0 +1,139 @@
from unittest.mock import ANY, MagicMock, patch
from pytest import mark
from sientia_do.temporal.activities.postgres import Postgres
from model_manager.activities.activities import Activities
from model_manager.activities.gates import Gates
from model_manager.activities.mlflow import MLFlow
@patch('model_manager.activities.activities.Postgres.__init__')
@patch('model_manager.activities.activities.MLFlow.__init__')
@patch('model_manager.activities.activities.MinIO.__init__')
@patch('model_manager.activities.activities.Gates.__init__')
def test___init__(mock_gates_init, mock_minio_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'}
minio_config = {
'endpoint_url': 'http://localhost:9000',
'access_key': 'minioadmin',
'secret_key': 'minioadmin',
'region': 'us-east-1',
'use_ssl': False,
'max_retry_attempts': 3,
'retry_mode': 'adaptive',
'connect_timeout': 10,
'read_timeout': 60,
}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
minio_config=minio_config,
logger=logger,
notification_handler=notification_handler,
)
assert isinstance(activities, Activities)
assert isinstance(activities, Postgres)
assert isinstance(activities, MLFlow)
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_minio_init.assert_called_once_with(
ANY,
endpoint_url=minio_config['endpoint_url'],
access_key=minio_config['access_key'],
secret_key=minio_config['secret_key'],
region=minio_config['region'],
use_ssl=minio_config['use_ssl'],
max_retry_attempts=minio_config['max_retry_attempts'],
retry_mode=minio_config['retry_mode'],
connect_timeout=minio_config['connect_timeout'],
read_timeout=minio_config['read_timeout'],
logger=logger,
notification_handler=notification_handler,
)
mock_gates_init.assert_called_once_with(
ANY, logger=logger, notification_handler=notification_handler
)
@mark.asyncio
@patch('model_manager.activities.activities.Postgres', return_value=MagicMock())
@patch('model_manager.activities.activities.MLFlow', return_value=MagicMock())
async def test_shutdown(_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'}
minio_config = {
'endpoint_url': 'http://localhost:9000',
'access_key': 'minioadmin',
'secret_key': 'minioadmin',
'region': 'us-east-1',
'use_ssl': False,
'max_retry_attempts': 3,
'retry_mode': 'adaptive',
'connect_timeout': 10,
'read_timeout': 60,
}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
minio_config=minio_config,
logger=logger,
notification_handler=notification_handler,
)
await activities.shutdown()
mock_postgres_init.close.assert_called_once()

630
tests/activities/test_gates.py Normal file
View File

@@ -0,0 +1,630 @@
from unittest.mock import ANY, MagicMock, 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_input_gate_filter_returns_false(gates_activity):
"""Test to cover line 129 branch when filter returns False (filter passes)."""
# Arrange - Use data that will NOT trigger EMPTY_DATA filter (has data)
input_data = {
**metadata,
'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
'data': {'value': [1, 2, 3, 4, 5]}, # Has data, filter returns False
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.input_gate(input_data)
# Assert - Filter returns False, so no policy is added to filter_output
assert result == (None, 0, '') # No filter triggered
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_response_gate_filter_returns_false(gates_activity):
"""Test to cover line 208 branch when filter returns False (no API error)."""
# Arrange - Use data that will NOT trigger API_ERROR filter (success=True)
input_data = {
**metadata,
'filters': {'API_ERROR': {'policy': 'STOP'}},
'data': {
'success': True, # Success=True, filter returns False
'content': {'message': 'Operation successful', 'result': 'data'},
},
'type': 'transform',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_response_gate(input_data)
# Assert - Filter returns False, so no policy is added to filter_output
assert result == (None, 0, '') # No filter triggered
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': {'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()
@mark.asyncio
async def test_mlflow_content_gate_filter_returns_false(gates_activity):
"""Test to cover line 293 branch when filter returns False (no NaN values)."""
# Arrange - Use data that will NOT trigger NAN_VALUES filter (no NaN)
input_data = {
**metadata,
'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}},
'data': {'value': [1.0, 2.0, 3.0, 4.0, 5.0]}, # All valid numbers, no NaN
'type': 'predict',
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
}
# Act
result = await gates_activity.mlflow_content_gate(input_data)
# Assert - Filter returns False, so no policy is added to filter_output
assert result == (None, 0, '') # No filter triggered
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': {
'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:
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
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
)

328
tests/activities/test_minio.py Normal file
View File

@@ -0,0 +1,328 @@
from io import BytesIO
from unittest.mock import MagicMock, patch
from pytest import fixture, mark, raises
from model_manager.activities.minio import MinIO
@patch('model_manager.activities.minio.boto3.client')
def test___init__(mock_boto3_client):
"""Test MinIO initialization with correct configuration."""
mock_client = MagicMock()
mock_boto3_client.return_value = mock_client
logger = MagicMock()
notification_handler = MagicMock()
minio = MinIO(
endpoint_url='http://localhost:9000',
access_key='minioadmin',
secret_key='minioadmin',
region='us-east-1',
use_ssl=False,
max_retry_attempts=3,
retry_mode='adaptive',
connect_timeout=10,
read_timeout=60,
logger=logger,
notification_handler=notification_handler,
)
assert minio.endpoint_url == 'http://localhost:9000'
assert minio.access_key == 'minioadmin'
assert minio.secret_key == 'minioadmin'
assert minio.region == 'us-east-1'
assert minio.use_ssl is False
assert minio.max_retry_attempts == 3
assert minio.retry_mode == 'adaptive'
assert minio.connect_timeout == 10
assert minio.read_timeout == 60
# Verify boto3 client was created with correct parameters
mock_boto3_client.assert_called_once()
call_kwargs = mock_boto3_client.call_args[1]
assert call_kwargs['endpoint_url'] == 'http://localhost:9000'
assert call_kwargs['aws_access_key_id'] == 'minioadmin'
assert call_kwargs['aws_secret_access_key'] == 'minioadmin'
assert call_kwargs['use_ssl'] is False
@patch('model_manager.activities.minio.boto3.client')
def test___init___failure(mock_boto3_client):
"""Test MinIO initialization failure handling."""
mock_boto3_client.side_effect = Exception('Connection failed')
logger = MagicMock()
notification_handler = MagicMock()
with raises(ConnectionError, match='Failed to initialize MinIO client'):
MinIO(
endpoint_url='http://localhost:9000',
access_key='minioadmin',
secret_key='minioadmin',
region='us-east-1',
use_ssl=False,
max_retry_attempts=3,
retry_mode='adaptive',
connect_timeout=10,
read_timeout=60,
logger=logger,
notification_handler=notification_handler,
)
@fixture
@patch('model_manager.activities.minio.boto3.client')
def minio(mock_boto3_client):
"""Fixture to create a MinIO instance for testing."""
mock_client = MagicMock()
mock_boto3_client.return_value = mock_client
logger = MagicMock()
notification_handler = MagicMock()
minio_instance = MinIO(
endpoint_url='http://localhost:9000',
access_key='minioadmin',
secret_key='minioadmin',
region='us-east-1',
use_ssl=False,
max_retry_attempts=3,
retry_mode='adaptive',
connect_timeout=10,
read_timeout=60,
logger=logger,
notification_handler=notification_handler,
)
minio_instance.send_notification = MagicMock()
minio_instance.minio_client = mock_client
return minio_instance
metadata = {
'metadata': {
'workflow_name': 'test_workflow',
'model_name': 'test_model',
'model_id': 'test_model_id',
}
}
@mark.asyncio
async def test_fetch_file_from_minio_success(minio):
"""Test successful file fetch from MinIO."""
# Arrange
test_content = b'test file content'
mock_response = {'Body': MagicMock()}
mock_response['Body'].__enter__ = MagicMock(
return_value=MagicMock(read=MagicMock(return_value=test_content))
)
mock_response['Body'].__exit__ = MagicMock(return_value=None)
minio.minio_client.get_object.return_value = mock_response
input_data = {
'metadata': metadata['metadata'],
'bucket_name': 'test-bucket',
'file_name': 'test-file.txt',
}
# Act
result = await minio.fetch_file_from_minio(input_data)
# Assert
assert isinstance(result, BytesIO)
result.seek(0)
assert result.read() == test_content
minio.minio_client.get_object.assert_called_once_with(Bucket='test-bucket', Key='test-file.txt')
@mark.asyncio
async def test_fetch_file_from_minio_file_not_found(minio):
"""Test file fetch when file doesn't exist."""
# Arrange
minio.minio_client.get_object.side_effect = Exception(
'NoSuchKey: The specified key does not exist'
)
input_data = {
'metadata': metadata['metadata'],
'bucket_name': 'test-bucket',
'file_name': 'nonexistent.txt',
}
# Act & Assert
with raises(OSError, match='Error fetching file from MinIO'):
await minio.fetch_file_from_minio(input_data)
# Verify notification was sent
minio.send_notification.assert_called_once()
call_kwargs = minio.send_notification.call_args[1]
assert call_kwargs['notification_id'] == 'FETCH_FILE_FROM_MINIO_ERROR'
assert call_kwargs['block'] == 'fetch_file_from_minio'
@mark.asyncio
async def test_fetch_file_from_minio_network_error(minio):
"""Test file fetch with network error."""
# Arrange
minio.minio_client.get_object.side_effect = Exception('Network timeout')
input_data = {
'metadata': metadata['metadata'],
'bucket_name': 'test-bucket',
'file_name': 'test-file.txt',
}
# Act & Assert
with raises(OSError, match='Error fetching file from MinIO'):
await minio.fetch_file_from_minio(input_data)
minio.send_notification.assert_called_once()
@mark.asyncio
async def test_delete_file_from_minio_success(minio):
"""Test successful file deletion from MinIO."""
# Arrange
minio.minio_client.delete_object.return_value = None
input_data = {
'metadata': metadata['metadata'],
'bucket_name': 'test-bucket',
'file_name': 'test-file.txt',
}
# Act
result = await minio.delete_file_from_minio(input_data)
# Assert
assert result is None
minio.minio_client.delete_object.assert_called_once_with(
Bucket='test-bucket', Key='test-file.txt'
)
@mark.asyncio
async def test_delete_file_from_minio_idempotent(minio):
"""Test that delete is idempotent (no error if file doesn't exist)."""
# Arrange
# MinIO delete_object is idempotent - no error if file doesn't exist
minio.minio_client.delete_object.return_value = None
input_data = {
'metadata': metadata['metadata'],
'bucket_name': 'test-bucket',
'file_name': 'nonexistent.txt',
}
# Act
result = await minio.delete_file_from_minio(input_data)
# Assert
assert result is None
minio.minio_client.delete_object.assert_called_once()
@mark.asyncio
async def test_delete_file_from_minio_access_denied(minio):
"""Test file deletion with access denied error."""
# Arrange
minio.minio_client.delete_object.side_effect = Exception('AccessDenied: Access Denied')
input_data = {
'metadata': metadata['metadata'],
'bucket_name': 'test-bucket',
'file_name': 'test-file.txt',
}
# Act & Assert
with raises(OSError, match='Error deleting file from MinIO'):
await minio.delete_file_from_minio(input_data)
# Verify notification was sent
minio.send_notification.assert_called_once()
call_kwargs = minio.send_notification.call_args[1]
assert call_kwargs['notification_id'] == 'DELETE_FILE_FROM_MINIO_ERROR'
assert call_kwargs['block'] == 'delete_file_from_minio'
@mark.asyncio
async def test_delete_file_from_minio_network_error(minio):
"""Test file deletion with network error."""
# Arrange
minio.minio_client.delete_object.side_effect = Exception('Connection timeout')
input_data = {
'metadata': metadata['metadata'],
'bucket_name': 'test-bucket',
'file_name': 'test-file.txt',
}
# Act & Assert
with raises(OSError, match='Error deleting file from MinIO'):
await minio.delete_file_from_minio(input_data)
minio.send_notification.assert_called_once()
@mark.asyncio
async def test_fetch_file_from_minio_large_file(minio):
"""Test fetching a large file from MinIO."""
# Arrange
# Simulate a 10MB file
large_content = b'x' * (10 * 1024 * 1024)
mock_response = {'Body': MagicMock()}
mock_response['Body'].__enter__ = MagicMock(
return_value=MagicMock(read=MagicMock(return_value=large_content))
)
mock_response['Body'].__exit__ = MagicMock(return_value=None)
minio.minio_client.get_object.return_value = mock_response
input_data = {
'metadata': metadata['metadata'],
'bucket_name': 'test-bucket',
'file_name': 'large-file.bin',
}
# Act
result = await minio.fetch_file_from_minio(input_data)
# Assert
assert isinstance(result, BytesIO)
result.seek(0)
assert len(result.read()) == 10 * 1024 * 1024
@mark.asyncio
async def test_fetch_file_from_minio_empty_file(minio):
"""Test fetching an empty file from MinIO."""
# Arrange
empty_content = b''
mock_response = {'Body': MagicMock()}
mock_response['Body'].__enter__ = MagicMock(
return_value=MagicMock(read=MagicMock(return_value=empty_content))
)
mock_response['Body'].__exit__ = MagicMock(return_value=None)
minio.minio_client.get_object.return_value = mock_response
input_data = {
'metadata': metadata['metadata'],
'bucket_name': 'test-bucket',
'file_name': 'empty-file.txt',
}
# Act
result = await minio.fetch_file_from_minio(input_data)
# Assert
assert isinstance(result, BytesIO)
result.seek(0)
assert result.read() == b''

301
tests/activities/test_mlflow.py Normal file
View File

@@ -0,0 +1,301 @@
from unittest.mock import ANY, MagicMock, patch
import numpy as np
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from model_manager.activities.mlflow import MLFlow
@patch('model_manager.activities.mlflow.MLFlowRepository')
def test___init__(mock_mlflow_repository):
mlflow = MLFlow(
mlflow_host='http://localhost',
mlflow_port=5000,
mlflow_username='admin',
mlflow_password='admin',
logger=MagicMock(),
notification_handler=MagicMock(),
)
assert mlflow.mlflow_host == 'http://localhost'
assert mlflow.mlflow_port == 5000
assert mlflow.mlflow_username == 'admin'
assert mlflow.mlflow_password == 'admin'
mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY)
@fixture
@patch('model_manager.activities.mlflow.MLFlowRepository')
def mlflow(mock_mlflow_repository):
mlflow = MLFlow(
mlflow_host='http://localhost:5000',
mlflow_port=5000,
mlflow_username='admin',
mlflow_password='admin',
logger=MagicMock(),
notification_handler=MagicMock(),
)
mlflow.send_notification = MagicMock()
return mlflow
metadata = {
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@mark.asyncio
@patch('model_manager.activities.mlflow.DataFrame')
@patch('model_manager.activities.mlflow.max')
async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02'
# Mock input data
input_data = {
**metadata,
'data': [
{
'timestamp': '2024-01-01',
'variable': 'var1',
'value': 1.0,
'created_at': '2024-01-01 12:00:00',
},
{
'timestamp': '2024-01-01',
'variable': 'var2',
'value': 2.0,
'created_at': '2024-01-01 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var1',
'value': 3.0,
'created_at': '2024-01-02 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var2',
'value': 4.0,
'created_at': '2024-01-02 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var1',
'value': 1.0,
'created_at': '2024-01-01 12:00:00',
},
{
'timestamp': '2024-01-02',
'variable': 'var2',
'value': 1.0,
'created_at': '2024-01-01 12:00:00',
},
],
'model_name': 'test_model',
'model_config': {},
}
# Mock the transform response
expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']}
mlflow.model_monitoring_repository.transform.return_value = expected_response
mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value
mock_dataframe.return_value.drop_duplicates.return_value = mock_dataframe.return_value
# Call the method
response_data = await mlflow.request_transform(input_data)
# Verify the data was correctly transformed
mock_dataframe.assert_called_once_with(input_data['data'])
mock_dataframe.return_value.pivot.assert_called_once_with(
index='timestamp', columns='variable', values='value'
)
mock_dataframe = mock_dataframe.return_value.pivot.return_value
mock_dataframe.fillna.assert_called_once_with(np.nan, inplace=True)
# mock_dataframe.reset_index.assert_called_once()
mock_dataframe.columns.name = None
# Verify the response
assert response_data == expected_response
# Verify the repository was called with correct arguments
mlflow.model_monitoring_repository.transform.assert_called_once_with(
'test_model', mock_dataframe, {}, metadata['metadata']
)
@mark.asyncio
@patch('model_manager.activities.mlflow.DataFrame')
@patch('model_manager.activities.mlflow.to_datetime')
@patch('model_manager.activities.mlflow.max')
async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02'
# Mock input data
input_data = {
**metadata,
'data': {
'variable': {
'2024-01-01': 'var1',
'2024-01-02': 'var2',
'2024-01-03': 'var1',
'2024-01-04': 'var2',
},
'value': {'2024-01-01': 1.0, '2024-01-02': 2.0, '2024-01-03': 3.0, '2024-01-04': 4.0},
},
'model_name': 'test_model',
'model_config': {},
}
# Mock the predict response
expected_response = {'prediction': [0.5, 0.6]}
mlflow.model_monitoring_repository.predict.return_value = expected_response
# Call the method
response_data = await mlflow.request_predict(input_data)
mock_dataframe.assert_called_once_with(input_data['data'])
mock_dataframe.return_value.replace.assert_called_once_with(np.nan, None, inplace=True)
mock_dataframe.return_value.__setitem__.assert_any_call(
'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
)
mock_to_datetime.assert_called_once_with(
mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
)
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
# Verify the response
assert response_data == expected_response
# Verify the repository was called with correct arguments
mlflow.model_monitoring_repository.predict.assert_called_once_with(
'test_model', mock_dataframe.return_value, {}, metadata['metadata']
)
@mark.asyncio
async def test_retrain_model(mlflow):
data = {
'model_id': [4, 5, 6, 7],
'created_at': [1, 2, 3, 4],
'timestamp': [1, 1, 2, 2],
'variable': ['var1', 'var2', 'var1', 'var2'],
'value': [1, 2, 3, 4],
}
mlflow.model_monitoring_repository.retrain_model.return_value = (
'Model retrained successfully',
'test',
)
response = await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'})
mlflow.model_monitoring_repository.retrain_model.assert_called_once()
assert response == {
'status': 'Model retrained successfully',
'timestamp': 2,
'experiment': 'test',
}
@mark.asyncio
async def test_retrain_model_error(mlflow):
mlflow.model_monitoring_repository.retrain_model.side_effect = Exception(
'Error retraining model'
)
data = {
'model_id': [4, 5, 6, 7],
'created_at': [1, 2, 3, 4],
'timestamp': [1, 1, 2, 2],
'variable': ['var1', 'var2', 'var1', 'var2'],
'value': [1, 2, 3, 4],
}
try:
await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'})
except Exception as e: # noqa: BLE001
assert str(e) == 'Error retraining model'
mlflow.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='RETRAIN_MODEL_ERROR',
message='Error retraining model test_model: Error retraining model',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
else:
raise AssertionError('No exception raised')
@mark.asyncio
async def test_update_production_model(mlflow):
mlflow.model_monitoring_repository.update_production_model.return_value = {
'data1': 1,
'data2': 2,
}
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 1,
'experiment': 'test',
'timestamp': 2,
'status': 'success',
}
response = await mlflow.update_production_model(input_data)
mlflow.model_monitoring_repository.update_production_model.assert_called_once_with(
experiment='test', model_name='test_model'
)
assert response == {
'data1': {0: 1},
'data2': {0: 2},
'model_id': {0: 1},
'model_name': {0: 'test_model'},
'timestamp': {0: 2},
'status': {0: 'success'},
}
@mark.asyncio
async def test_update_production_model_error(mlflow):
mlflow.model_monitoring_repository.update_production_model.side_effect = Exception(
'Error updating production model'
)
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 1,
'experiment': 'test',
'timestamp': 2,
'status': 'success',
}
try:
await mlflow.update_production_model(input_data)
except Exception as e: # noqa: BLE001
assert str(e) == 'Error updating production model'
mlflow.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
message='Error updating production model test_model: Error updating production model',
block='update_production_model',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
else:
raise AssertionError('No exception raised')