697 lines
23 KiB
Python
697 lines
23 KiB
Python
from unittest.mock import ANY, MagicMock, patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
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')
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_save_model_success(mlflow):
|
|
"""Test save_model successfully saves model and artifacts to MLflow."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
# Mock train result
|
|
params = MagicMock(spec=TrainModelParams)
|
|
params.experiment_name = 'test_experiment'
|
|
|
|
train_result = MagicMock(spec=TrainModelResult)
|
|
train_result.params = params
|
|
train_result.run_name = None # Will be set by get_next_run_name
|
|
|
|
# Mock repository methods
|
|
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
|
|
mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
|
|
mlflow.model_monitoring_repository.save_run.return_value = None
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'train_result': train_result,
|
|
}
|
|
|
|
# Call the method
|
|
response = await mlflow.save_model(input_data)
|
|
|
|
# Verify repository methods were called
|
|
mlflow.model_monitoring_repository.get_next_run_name.assert_called_once_with('test_experiment')
|
|
mlflow.model_monitoring_repository.generate_artifacts.assert_called_once_with(train_result)
|
|
mlflow.model_monitoring_repository.save_run.assert_called_once_with(train_result)
|
|
|
|
# Verify response - now returns TrainModelResult directly
|
|
assert response == train_result
|
|
assert response.run_name == 'test_experiment-1'
|
|
assert train_result.run_name == 'test_experiment-1'
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_save_model_get_next_run_name_error(mlflow):
|
|
"""Test save_model handles error during get_next_run_name."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
params.experiment_name = 'test_experiment'
|
|
|
|
train_result = MagicMock(spec=TrainModelResult)
|
|
train_result.params = params
|
|
|
|
# Mock error in get_next_run_name
|
|
mlflow.model_monitoring_repository.get_next_run_name.side_effect = Exception(
|
|
'MLflow connection error'
|
|
)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'train_result': train_result,
|
|
}
|
|
|
|
# Call the method - should raise exception
|
|
with pytest.raises(Exception, match='MLflow connection error'):
|
|
await mlflow.save_model(input_data)
|
|
|
|
# Verify notification was sent
|
|
mlflow.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='SAVE_MODEL_ERROR',
|
|
message=ANY,
|
|
block='save_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_save_model_generate_artifacts_error(mlflow):
|
|
"""Test save_model handles error during generate_artifacts."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
params.experiment_name = 'test_experiment'
|
|
|
|
train_result = MagicMock(spec=TrainModelResult)
|
|
train_result.params = params
|
|
|
|
# Mock successful get_next_run_name but error in generate_artifacts
|
|
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
|
|
mlflow.model_monitoring_repository.generate_artifacts.side_effect = FileNotFoundError(
|
|
'Reports directory does not exist'
|
|
)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'train_result': train_result,
|
|
}
|
|
|
|
# Call the method - should raise exception
|
|
with pytest.raises(FileNotFoundError, match='Reports directory does not exist'):
|
|
await mlflow.save_model(input_data)
|
|
|
|
# Verify notification was sent
|
|
mlflow.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='SAVE_MODEL_ERROR',
|
|
message=ANY,
|
|
block='save_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_save_model_save_run_error(mlflow):
|
|
"""Test save_model handles error during save_run."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
params.experiment_name = 'test_experiment'
|
|
|
|
train_result = MagicMock(spec=TrainModelResult)
|
|
train_result.params = params
|
|
|
|
# Mock successful get_next_run_name and generate_artifacts but error in save_run
|
|
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
|
|
mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
|
|
mlflow.model_monitoring_repository.save_run.side_effect = ValueError(
|
|
'One or more metrics (MSE, R2, MAE) are None'
|
|
)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'train_result': train_result,
|
|
}
|
|
|
|
# Call the method - should raise exception
|
|
with pytest.raises(ValueError, match=r'One or more metrics \(MSE, R2, MAE\) are None'):
|
|
await mlflow.save_model(input_data)
|
|
|
|
# Verify notification was sent
|
|
mlflow.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='SAVE_MODEL_ERROR',
|
|
message=ANY,
|
|
block='save_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_save_model_missing_metadata(mlflow):
|
|
"""Test save_model handles missing metadata gracefully."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
params.experiment_name = 'test_experiment'
|
|
|
|
train_result = MagicMock(spec=TrainModelResult)
|
|
train_result.params = params
|
|
|
|
# Mock repository methods
|
|
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
|
|
mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
|
|
mlflow.model_monitoring_repository.save_run.return_value = None
|
|
|
|
# Input data without metadata
|
|
input_data = {
|
|
'train_result': train_result,
|
|
}
|
|
|
|
# Call the method
|
|
response = await mlflow.save_model(input_data)
|
|
|
|
# Verify it still works (metadata defaults to {}) - returns TrainModelResult directly
|
|
assert response == train_result
|
|
assert response.run_name == 'test_experiment-1'
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_save_model_complete_flow(mlflow):
|
|
"""Test save_model complete flow with all steps."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
params.experiment_name = 'production_model'
|
|
|
|
train_result = MagicMock(spec=TrainModelResult)
|
|
train_result.params = params
|
|
train_result.run_name = None
|
|
train_result.run_dir = None
|
|
train_result.report_path = None
|
|
|
|
# Mock complete flow
|
|
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'production_model-5'
|
|
|
|
# After generate_artifacts, paths should be set
|
|
updated_result = MagicMock(spec=TrainModelResult)
|
|
updated_result.params = params
|
|
updated_result.run_name = 'production_model-5'
|
|
updated_result.run_dir = '/reports/production_model-5_20231010'
|
|
updated_result.report_path = '/reports/production_model-5_20231010/report.html'
|
|
updated_result.train_data_path = '/reports/production_model-5_20231010/train_data.csv'
|
|
updated_result.test_data_path = '/reports/production_model-5_20231010/test_data.csv'
|
|
|
|
mlflow.model_monitoring_repository.generate_artifacts.return_value = updated_result
|
|
mlflow.model_monitoring_repository.save_run.return_value = None
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'train_result': train_result,
|
|
}
|
|
|
|
# Call the method
|
|
response = await mlflow.save_model(input_data)
|
|
|
|
# Verify complete flow
|
|
mlflow.model_monitoring_repository.get_next_run_name.assert_called_once_with('production_model')
|
|
mlflow.model_monitoring_repository.generate_artifacts.assert_called_once()
|
|
mlflow.model_monitoring_repository.save_run.assert_called_once_with(updated_result)
|
|
|
|
# Verify response - returns TrainModelResult directly
|
|
assert response == updated_result
|
|
assert response.run_name == 'production_model-5'
|
|
assert response.run_dir is not None
|
|
assert response.report_path is not None
|
|
|
|
|
|
# ============================================================================
|
|
# Tests for cleanup_run_directory
|
|
# ============================================================================
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_success(mlflow):
|
|
"""Test successful cleanup of run directory."""
|
|
input_data = {
|
|
**metadata,
|
|
'run_dir': 'test_run_dir',
|
|
}
|
|
|
|
# Patch os and shutil inside the activity method
|
|
with (
|
|
patch('os.path.exists', return_value=True) as mock_exists,
|
|
patch('shutil.rmtree') as mock_rmtree,
|
|
):
|
|
# Call the method
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Verify directory existence was checked
|
|
mock_exists.assert_called_once_with('test_run_dir')
|
|
|
|
# Verify shutil.rmtree was called
|
|
mock_rmtree.assert_called_once_with('test_run_dir')
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_already_deleted(mlflow):
|
|
"""Test cleanup when directory is already deleted (idempotent)."""
|
|
input_data = {
|
|
**metadata,
|
|
'run_dir': 'already_deleted_dir',
|
|
}
|
|
|
|
# Patch os and shutil inside the activity method
|
|
with (
|
|
patch('os.path.exists', return_value=False) as mock_exists,
|
|
patch('shutil.rmtree') as mock_rmtree,
|
|
):
|
|
# Call the method - should not raise error
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Verify directory existence was checked
|
|
mock_exists.assert_called_once_with('already_deleted_dir')
|
|
|
|
# Verify shutil.rmtree was NOT called
|
|
mock_rmtree.assert_not_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_no_run_dir(mlflow):
|
|
"""Test cleanup when no run_dir is provided."""
|
|
input_data = {
|
|
**metadata,
|
|
# No run_dir key
|
|
}
|
|
|
|
# Call the method - should not raise error
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Should complete without errors
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_empty_run_dir(mlflow):
|
|
"""Test cleanup when run_dir is empty string."""
|
|
input_data = {
|
|
**metadata,
|
|
'run_dir': '',
|
|
}
|
|
|
|
# Call the method - should not raise error
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Should complete without errors
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_none_run_dir(mlflow):
|
|
"""Test cleanup when run_dir is None."""
|
|
input_data = {
|
|
**metadata,
|
|
'run_dir': None,
|
|
}
|
|
|
|
# Call the method - should not raise error
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Should complete without errors
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_error(mlflow):
|
|
"""Test cleanup handles errors correctly."""
|
|
input_data = {
|
|
**metadata,
|
|
'run_dir': 'error_dir',
|
|
}
|
|
|
|
# Patch os and shutil inside the activity method
|
|
with (
|
|
patch('os.path.exists', return_value=True),
|
|
patch('shutil.rmtree', side_effect=PermissionError('Permission denied')),
|
|
):
|
|
# Call the method - should raise exception
|
|
with pytest.raises(PermissionError, match='Permission denied'):
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Verify notification was sent
|
|
mlflow.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='CLEANUP_RUN_DIRECTORY_ERROR',
|
|
message=ANY,
|
|
block='cleanup_run_directory',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_missing_metadata(mlflow):
|
|
"""Test cleanup handles missing metadata gracefully."""
|
|
input_data = {
|
|
'run_dir': 'no_metadata_dir',
|
|
}
|
|
|
|
# Patch os and shutil inside the activity method
|
|
with patch('os.path.exists', return_value=True), patch('shutil.rmtree') as mock_rmtree:
|
|
# Call the method
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Verify directory was deleted
|
|
mock_rmtree.assert_called_once_with('no_metadata_dir')
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_oserror(mlflow):
|
|
"""Test cleanup handles OSError correctly."""
|
|
input_data = {
|
|
**metadata,
|
|
'run_dir': 'os_error_dir',
|
|
}
|
|
|
|
# Patch os and shutil inside the activity method
|
|
with (
|
|
patch('os.path.exists', return_value=True),
|
|
patch('shutil.rmtree', side_effect=OSError('Directory not empty')),
|
|
):
|
|
# Call the method - should raise exception
|
|
with pytest.raises(OSError, match='Directory not empty'):
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Verify notification was sent
|
|
mlflow.send_notification.assert_called_once()
|
|
call_args = mlflow.send_notification.call_args[1]
|
|
assert call_args['notification_id'] == 'CLEANUP_RUN_DIRECTORY_ERROR'
|
|
assert call_args['level'] == NotificationLevel.ERROR
|
|
assert 'Directory not empty' in call_args['message']
|