vitor-aignosi
e4fc12e529
QTZPOC-20: Enhance system robustness and test coverage across bundle processing, activity error handling, and worker operations. * Strengthen bundle reader validation with new tests for archive integrity, security vulnerabilities (e.g., path traversal, duplicate entries, special file types), and I/O error resilience. * Improve error handling and status propagation in PI Web API, MLflow, and OPC activities, including partial write scenarios and OPC session management. * Add comprehensive tests for worker startup, configuration, graceful shutdown, and Prometheus metrics server initialization. * Introduce a new utility for intelligent DataFrame debugging, summarizing large dataframes for concise logging, and refine exception chain summarization. * Update code coverage configuration to exclude @overload decorator lines.
663 lines
21 KiB
Python
663 lines
21 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.MLFlowRepository')
|
|
@patch('laborious.activities.mlflow.MinioRepository')
|
|
def test___init__(mock_minio_repository, mock_mlflow_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_host='http://localhost',
|
|
mlflow_port=5000,
|
|
mlflow_username='admin',
|
|
mlflow_password='admin',
|
|
minio_repository=minio_repo,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
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, ANY, ANY
|
|
)
|
|
|
|
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',
|
|
)
|
|
|
|
|
|
@fixture
|
|
@patch('laborious.activities.mlflow.MLFlowRepository')
|
|
@patch('laborious.activities.mlflow.MinioRepository')
|
|
def mlflow(mock_minio_repository, mock_mlflow_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_host='http://localhost:5000',
|
|
mlflow_port=5000,
|
|
mlflow_username='admin',
|
|
mlflow_password='admin',
|
|
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'],
|
|
)
|