vitor-aignosi
192834bde3
QTZPOC-20: Refactor model import workflow, introduce atomic model ID generation, and clarify MinIO expiry. Fuses bundle download, decryption, MLflow experiment creation, and artifact upload into a single activity to prevent cross-worker issues. Implements an atomic counter for MongoDB model IDs, replacing the max(id) + 1 scan to avoid collisions during concurrent imports. Adds comprehensive documentation for MinIO object expiry mechanisms and refines lifecycle rule application. Improves error handling with clearer responsibility classification and centralizes MLflow repository injection.
658 lines
20 KiB
Python
658 lines
20 KiB
Python
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
|
|
|
import numpy as np
|
|
from pytest import fixture, mark, raises
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
|
|
|
|
from laborious.activities.mlflow import MLFlow
|
|
|
|
|
|
@fixture(autouse=True)
|
|
def _passthrough_from_dict():
|
|
with patch(
|
|
'laborious.activities.mlflow.MinioDataFramePayload.from_dict', side_effect=lambda x: x
|
|
):
|
|
yield
|
|
|
|
|
|
@patch('laborious.activities.mlflow.MinioRepository')
|
|
def test___init__(mock_minio_repository):
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
metrics_controller = AsyncMock()
|
|
mock_mlflow_repository = MagicMock()
|
|
minio_repo = mock_minio_repository(
|
|
endpoint='localhost:9000',
|
|
access_key='minio',
|
|
secret_key='minio123',
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
bucket='test',
|
|
)
|
|
mlflow = MLFlow(
|
|
mlflow_repository=mock_mlflow_repository,
|
|
minio_repository=minio_repo,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
assert mlflow.model_monitoring_repository is mock_mlflow_repository
|
|
|
|
mock_minio_repository.assert_called_once_with(
|
|
endpoint='localhost:9000',
|
|
access_key='minio',
|
|
secret_key='minio123',
|
|
logger=ANY,
|
|
notification_handler=ANY,
|
|
metrics_controller=ANY,
|
|
bucket='test',
|
|
)
|
|
|
|
|
|
def test___init___requires_mlflow_repository():
|
|
"""The failure moves to the signature, at construction time, instead of a `None` that only
|
|
surfaces silently at first use."""
|
|
minio_repo = MagicMock()
|
|
logger = MagicMock()
|
|
with raises(TypeError):
|
|
MLFlow(minio_repository=minio_repo, logger=logger)
|
|
|
|
|
|
@fixture
|
|
@patch('laborious.activities.mlflow.MinioRepository')
|
|
def mlflow(mock_minio_repository):
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
metrics_controller = AsyncMock()
|
|
minio_repo = mock_minio_repository(
|
|
endpoint='localhost:9000',
|
|
access_key='minio',
|
|
secret_key='minio123',
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
bucket='test',
|
|
)
|
|
mlflow = MLFlow(
|
|
mlflow_repository=MagicMock(),
|
|
minio_repository=minio_repo,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
mlflow.model_monitoring_repository = AsyncMock()
|
|
mlflow.minio_repository = AsyncMock()
|
|
|
|
mlflow.send_notification = MagicMock()
|
|
mlflow.emit_metric = AsyncMock()
|
|
mlflow.send_notification_async = AsyncMock()
|
|
mlflow.error = MagicMock()
|
|
mlflow.debug = MagicMock()
|
|
mlflow.info = MagicMock()
|
|
mlflow.warning = MagicMock()
|
|
mlflow.critical = MagicMock()
|
|
|
|
return mlflow
|
|
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'model_id': 'test_model',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
'schema_name': 'test_schedule',
|
|
},
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
@patch(
|
|
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
)
|
|
async def test_request_transform_success(mock_from_dataframe, mlflow):
|
|
data_mock = MagicMock()
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=data_mock)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
|
|
transform_response = {'success': True, 'content': MagicMock()}
|
|
mlflow.model_monitoring_repository.transform.return_value = transform_response
|
|
|
|
data_mock.sort_values.return_value = data_mock
|
|
data_mock.drop_duplicates.return_value = data_mock
|
|
data_mock.pivot.return_value = data_mock
|
|
|
|
response_data = await mlflow.request_transform(input_data)
|
|
|
|
mlflow.model_monitoring_repository.transform.assert_called_once_with(
|
|
'test_model', data_mock, {}, metadata['metadata']
|
|
)
|
|
mock_from_dataframe.assert_called_once()
|
|
assert response_data == mock_from_dataframe.return_value
|
|
|
|
|
|
@mark.asyncio
|
|
@patch(
|
|
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
)
|
|
async def test_request_transform_failure(mock_from_dataframe, mlflow):
|
|
data_mock = MagicMock()
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=data_mock)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
|
|
transform_response = {'success': False, 'message': 'Transform failed'}
|
|
mlflow.model_monitoring_repository.transform.return_value = transform_response
|
|
|
|
data_mock.sort_values.return_value = data_mock
|
|
data_mock.drop_duplicates.return_value = data_mock
|
|
data_mock.pivot.return_value = data_mock
|
|
|
|
response_data = await mlflow.request_transform(input_data)
|
|
|
|
mock_from_dataframe.assert_called_once_with(
|
|
dataframe=None,
|
|
minio_repo=mlflow.minio_repository,
|
|
model_name='test_model',
|
|
operation='transform',
|
|
status=transform_response,
|
|
workflow_metadata=metadata['metadata'],
|
|
last_timestamp=payload.last_timestamp,
|
|
logger=mlflow.logger,
|
|
)
|
|
assert response_data == mock_from_dataframe.return_value
|
|
|
|
|
|
@mark.asyncio
|
|
@patch(
|
|
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
)
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow):
|
|
data_mock = MagicMock()
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=data_mock)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
|
|
predict_response = {'success': True, 'content': MagicMock()}
|
|
mlflow.model_monitoring_repository.predict.return_value = predict_response
|
|
|
|
response_data = await mlflow.request_predict(input_data)
|
|
|
|
data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
|
|
mock_to_datetime.assert_called_once_with(
|
|
data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
|
|
)
|
|
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
|
|
|
|
mlflow.model_monitoring_repository.predict.assert_called_once_with(
|
|
'test_model', data_mock, {}, metadata['metadata']
|
|
)
|
|
mock_from_dataframe.assert_called_once()
|
|
assert response_data == mock_from_dataframe.return_value
|
|
|
|
|
|
@mark.asyncio
|
|
@patch(
|
|
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
)
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_request_predict_skips_timestamp_reinsertion_when_configured(
|
|
mock_to_datetime, mock_from_dataframe, mlflow
|
|
):
|
|
"""skip_predict_timestamp=True leaves the caller's index alone — no reinsertion round-trip."""
|
|
data_mock = MagicMock()
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=data_mock)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {'skip_predict_timestamp': True},
|
|
}
|
|
|
|
predict_response = {'success': True, 'content': MagicMock()}
|
|
mlflow.model_monitoring_repository.predict.return_value = predict_response
|
|
|
|
await mlflow.request_predict(input_data)
|
|
|
|
data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
|
|
mock_to_datetime.assert_not_called()
|
|
mlflow.model_monitoring_repository.predict.assert_called_once_with(
|
|
'test_model', data_mock, {'skip_predict_timestamp': True}, metadata['metadata']
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch(
|
|
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
)
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, mlflow):
|
|
data_mock = MagicMock()
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=data_mock)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
|
|
predict_response = {'success': False, 'message': 'Predict failed'}
|
|
mlflow.model_monitoring_repository.predict.return_value = predict_response
|
|
|
|
response_data = await mlflow.request_predict(input_data)
|
|
|
|
mock_from_dataframe.assert_called_once_with(
|
|
dataframe=None,
|
|
minio_repo=mlflow.minio_repository,
|
|
model_name='test_model',
|
|
operation='predict',
|
|
status=predict_response,
|
|
workflow_metadata=metadata['metadata'],
|
|
last_timestamp=payload.last_timestamp,
|
|
logger=mlflow.logger,
|
|
)
|
|
assert response_data == mock_from_dataframe.return_value
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlflow):
|
|
mlflow.model_monitoring_repository.retrain_model.return_value = {
|
|
'success': True,
|
|
'experiment': 'test_experiment',
|
|
'message': 'Model retrained successfully.',
|
|
}
|
|
|
|
raw_data = MagicMock(columns=['variable', 'timestamp', 'value'])
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=raw_data)
|
|
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
}
|
|
)
|
|
|
|
timestamp = raw_data.__getitem__.return_value.max.return_value
|
|
|
|
raw_data.sort_values.assert_not_called()
|
|
raw_data.drop_duplicates.assert_called_once_with(subset=['variable', 'timestamp'], keep='first')
|
|
raw_data = raw_data.drop_duplicates.return_value
|
|
|
|
raw_data.drop.assert_has_calls(
|
|
[
|
|
call(columns=['model_id'], inplace=True, errors='ignore'),
|
|
call(columns=['created_at'], inplace=True, errors='ignore'),
|
|
]
|
|
)
|
|
raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
|
|
raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
|
|
|
|
raw_data = raw_data.pivot.return_value
|
|
|
|
raw_data.__setitem__.assert_has_calls(
|
|
[
|
|
call('timestamp', raw_data.index),
|
|
call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
|
|
call('timestamp', mock_to_datetime.return_value),
|
|
]
|
|
)
|
|
mock_to_datetime.assert_has_calls(
|
|
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
|
|
)
|
|
mock_to_datetime.assert_has_calls(
|
|
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
|
|
)
|
|
|
|
mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
|
|
data=raw_data,
|
|
model_name='test_model',
|
|
model_config={
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
assert response == {
|
|
'success': True,
|
|
'experiment': 'test_experiment',
|
|
'message': 'Model retrained successfully.',
|
|
'timestamp': timestamp,
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_retrain_model_success_with_payload_data(mock_to_datetime, mlflow):
|
|
mlflow.model_monitoring_repository.retrain_model.return_value = {
|
|
'success': True,
|
|
'experiment': 'test_experiment',
|
|
'message': 'Model retrained successfully.',
|
|
}
|
|
|
|
raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=raw_data)
|
|
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
}
|
|
)
|
|
|
|
assert response['success'] is True
|
|
mlflow.minio_repository.download_file.assert_not_called()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow):
|
|
mlflow.model_monitoring_repository.retrain_model.return_value = {
|
|
'success': False,
|
|
'traceback': 'test_traceback',
|
|
'message': 'Model retrained failed.',
|
|
}
|
|
|
|
raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=raw_data)
|
|
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
}
|
|
)
|
|
|
|
timestamp = raw_data.__getitem__.return_value.max.return_value
|
|
|
|
raw_data.sort_values.assert_called_once_with('created_at', ascending=False)
|
|
raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with(
|
|
subset=['variable', 'timestamp'], keep='first'
|
|
)
|
|
raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value
|
|
|
|
raw_data.drop.assert_has_calls(
|
|
[
|
|
call(columns=['model_id'], inplace=True, errors='ignore'),
|
|
call(columns=['created_at'], inplace=True, errors='ignore'),
|
|
]
|
|
)
|
|
raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
|
|
raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
|
|
|
|
raw_data = raw_data.pivot.return_value
|
|
|
|
raw_data.__setitem__.assert_has_calls(
|
|
[
|
|
call('timestamp', raw_data.index),
|
|
call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
|
|
call('timestamp', mock_to_datetime.return_value),
|
|
]
|
|
)
|
|
mock_to_datetime.assert_has_calls(
|
|
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
|
|
)
|
|
mock_to_datetime.assert_has_calls(
|
|
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
|
|
)
|
|
|
|
mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
|
|
data=raw_data,
|
|
model_name='test_model',
|
|
model_config={
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
mlflow.send_notification_async.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='RETRAIN_MODEL_ERROR',
|
|
message='Error retraining model test_model: Model retrained failed.',
|
|
block='retrain_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
assert response == {
|
|
'success': False,
|
|
'traceback': 'test_traceback',
|
|
'message': 'Model retrained failed.',
|
|
'timestamp': timestamp,
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_retrain_model_data_error(mlflow):
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
}
|
|
)
|
|
|
|
assert response == {
|
|
'success': False,
|
|
'message': "Error loading retrain data: 'data'",
|
|
'traceback': ANY,
|
|
'timestamp': ANY,
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_retrain_model_data_error_no_minio_repository(mlflow):
|
|
mlflow.minio_repository = None
|
|
|
|
with raises(ValueError) as e:
|
|
await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'object_key': 'test_object_key',
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
}
|
|
)
|
|
|
|
assert str(e.value) == 'Minio repository not initialized'
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_update_production_model(mlflow):
|
|
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', metadata=metadata['metadata']
|
|
)
|
|
|
|
assert response == mlflow.model_monitoring_repository.update_production_model.return_value
|
|
|
|
|
|
@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:
|
|
assert str(e) == 'Error updating production model'
|
|
mlflow.send_notification_async.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')
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_get_reference_data_success(mock_to_datetime, mlflow):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'model_name': 'test_model',
|
|
}
|
|
|
|
# Mock reference data DataFrame
|
|
mock_reference_data = MagicMock()
|
|
mock_reference_data.__getitem__.return_value = MagicMock()
|
|
mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
|
|
mock_reference_data.to_dict.return_value = [
|
|
{'timestamp': '2023-05-26 11:12:27', 'value': 1.0},
|
|
{'timestamp': '2023-05-26 11:12:28', 'value': 2.0},
|
|
]
|
|
|
|
mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = mock_reference_data
|
|
|
|
# Act
|
|
result = await mlflow.get_reference_data(input_data)
|
|
|
|
# Assert
|
|
mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
|
|
model_name='test_model',
|
|
artifact_path='evaluation_data.csv',
|
|
metadata=metadata['metadata'],
|
|
)
|
|
mock_to_datetime.assert_called_once_with(mock_reference_data.__getitem__.return_value)
|
|
|
|
mock_reference_data.to_dict.assert_called_once_with(orient='records')
|
|
assert result == mock_reference_data.to_dict.return_value
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_get_reference_data_not_found(mlflow):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'model_name': 'test_model',
|
|
}
|
|
|
|
mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = None
|
|
|
|
# Act
|
|
result = await mlflow.get_reference_data(input_data)
|
|
|
|
# Assert
|
|
mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
|
|
model_name='test_model',
|
|
artifact_path='evaluation_data.csv',
|
|
metadata=metadata['metadata'],
|
|
)
|
|
mlflow.warning.assert_called_once_with(
|
|
'Reference data not found for model test_model', metadata['metadata']
|
|
)
|
|
assert result is None
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_get_reference_data_exception(mlflow):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'model_name': 'test_model',
|
|
}
|
|
|
|
mlflow.model_monitoring_repository.load_artifact_dataframe.side_effect = Exception(
|
|
'Error loading artifact'
|
|
)
|
|
|
|
# Act & Assert
|
|
with raises(Exception) as e:
|
|
await mlflow.get_reference_data(input_data)
|
|
|
|
assert str(e.value) == 'Error loading artifact'
|
|
mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
|
|
model_name='test_model',
|
|
artifact_path='evaluation_data.csv',
|
|
metadata=metadata['metadata'],
|
|
)
|