1177 lines
44 KiB
Python
1177 lines
44 KiB
Python
from datetime import UTC, datetime
|
|
from unittest.mock import ANY, MagicMock, call, patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from pandas import DataFrame, Timestamp
|
|
|
|
from model_manager.utils.repository.model_repository import MLFlowRepository
|
|
|
|
|
|
@pytest.fixture
|
|
def mlflow_repository():
|
|
with patch(
|
|
'model_manager.utils.repository.model_repository.ModelServing', autospec=True
|
|
) as mock_model_serving:
|
|
mock_instance = mock_model_serving.return_value
|
|
mock_instance.get_transformed_data = MagicMock()
|
|
|
|
repo = MLFlowRepository(
|
|
host='http://localhost:5000', username='admin', password='admin', logger=MagicMock()
|
|
)
|
|
return repo
|
|
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'model_id': 'test_model',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
'schema_name': 'test_schedule',
|
|
},
|
|
}
|
|
|
|
|
|
class Any:
|
|
pass
|
|
|
|
|
|
invalid_cases = [
|
|
({'value': {'2024-01-01 12:00:00': 1, 2024: 2}}),
|
|
({'value': {'2024-01-01': 1, '2024-01-02': 2}}),
|
|
({'value': {Any(): 1, Any(): 2}}),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize('data', invalid_cases)
|
|
def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data):
|
|
input_data = DataFrame(data)
|
|
|
|
with pytest.raises(ValueError) as e:
|
|
mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
|
|
|
|
assert (
|
|
str(e)
|
|
== 'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S'
|
|
)
|
|
|
|
|
|
valid_cases = [
|
|
(
|
|
{'value': {'2024-01-01 12:00:00+0000': 1, '2024-01-02 12:00:00+0000': 2}},
|
|
['2024-01-01 12:00:00+0000', '2024-01-02 12:00:00+0000'],
|
|
),
|
|
(
|
|
{
|
|
'value': {
|
|
datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC): 1,
|
|
datetime(2025, 1, 2, 12, 0, 0, tzinfo=UTC): 2,
|
|
}
|
|
},
|
|
['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000'],
|
|
),
|
|
(
|
|
{
|
|
'value': {
|
|
Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=UTC): 1,
|
|
Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=UTC): 2,
|
|
}
|
|
},
|
|
['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000'],
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize('data,expected', valid_cases)
|
|
def test_detect_and_parse_datetime_index_valid_format(mlflow_repository, data, expected):
|
|
input_data = DataFrame(data)
|
|
|
|
response = mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
|
|
|
|
assert response.index.tolist() == expected
|
|
|
|
|
|
def test_transform_success(mlflow_repository):
|
|
data = MagicMock()
|
|
model_name = 'model'
|
|
|
|
mlflow_repository.detect_and_parse_datetime_index = MagicMock()
|
|
|
|
output = mlflow_repository.transform(model_name, data, {}, metadata['metadata'])
|
|
|
|
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
|
|
model_name, data, 0, 'sklearn', False, 'model', 'predict'
|
|
)
|
|
|
|
mlflow_repository.detect_and_parse_datetime_index.assert_called_once_with(
|
|
mlflow_repository.model_serving.get_cached_transform.return_value, metadata['metadata']
|
|
)
|
|
|
|
assert output == {
|
|
'success': True,
|
|
'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value,
|
|
}
|
|
|
|
|
|
def test_transform_error(mlflow_repository):
|
|
data = MagicMock()
|
|
model_name = 'model'
|
|
|
|
mlflow_repository.model_serving.get_cached_transform.side_effect = Exception('error')
|
|
|
|
output = mlflow_repository.transform(model_name, data, {}, metadata['metadata'])
|
|
|
|
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
|
|
model_name, data, 0, 'sklearn', False, 'model', 'predict'
|
|
)
|
|
|
|
assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}}
|
|
|
|
|
|
def test_predict_success(mlflow_repository):
|
|
data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
|
|
model_name = 'model'
|
|
mlflow_repository.model_serving.get_cached_predict.return_value = np.array([2, 3])
|
|
|
|
output = mlflow_repository.predict(model_name, data, {}, metadata['metadata'])
|
|
|
|
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
|
|
model_name, data, 0, 'pyfunc', False, 'model'
|
|
)
|
|
|
|
assert output['success'] is True
|
|
assert output['content'] == {
|
|
'prediction': {'index_1': 2, 'index_2': 3},
|
|
'response_time': {'index_1': ANY, 'index_2': ANY},
|
|
}
|
|
|
|
|
|
def test_predict_error(mlflow_repository):
|
|
data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
|
|
model_name = 'model'
|
|
|
|
mlflow_repository.model_serving.get_cached_predict = MagicMock(side_effect=Exception('error'))
|
|
|
|
output = mlflow_repository.predict(model_name, data, {}, metadata['metadata'])
|
|
|
|
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
|
|
model_name, data, 0, 'pyfunc', False, 'model'
|
|
)
|
|
|
|
assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}}
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.mlflow')
|
|
def test_get_experiment_by_run_id(mlflow, mlflow_repository):
|
|
mlflow.get_run.return_value = MagicMock(
|
|
info=MagicMock(
|
|
experiment_id='0',
|
|
)
|
|
)
|
|
mlflow.get_experiment.return_value = MagicMock()
|
|
mlflow.get_experiment.return_value.name = 'test'
|
|
|
|
output = mlflow_repository.get_experiment_by_run_id('0')
|
|
assert output == 'test'
|
|
mlflow.get_run.assert_called_once_with('0')
|
|
mlflow.get_experiment.assert_called_once_with('0')
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.mlflow')
|
|
def test_get_next_run_name(mlflow, mlflow_repository):
|
|
mlflow.search_runs.return_value = [1, 2, 3]
|
|
output = mlflow_repository.get_next_run_name('run')
|
|
assert output == 'run-4'
|
|
mlflow.search_runs.assert_called_once_with(
|
|
experiment_names=['run'],
|
|
order_by=['start_time desc'],
|
|
)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.mlflow')
|
|
def test_get_experiment_success(mlflow, mlflow_repository):
|
|
mlflow.get_experiment_by_name.return_value = MagicMock(experiment_id='0')
|
|
|
|
output = mlflow_repository.get_experiment('test')
|
|
|
|
assert output == '0'
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.mlflow')
|
|
def test_get_experiment_error(mlflow, mlflow_repository):
|
|
mlflow.get_experiment_by_name.return_value = None
|
|
|
|
try:
|
|
mlflow_repository.get_experiment('test')
|
|
except ValueError as e:
|
|
assert str(e) == 'Experiment test not found'
|
|
else:
|
|
raise AssertionError('Expected exception')
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.mlflow')
|
|
def test_get_experiment_last_run(mlflow, mlflow_repository):
|
|
mlflow.search_runs.return_value = DataFrame(
|
|
{
|
|
'params.retrain': ['True', 'False', 'True', 'False'],
|
|
'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'],
|
|
'run_id': ['0', '1', '2', '3'],
|
|
}
|
|
)
|
|
|
|
output = mlflow_repository.get_experiment_last_run(0)
|
|
|
|
mlflow.search_runs.assert_called_once_with(
|
|
experiment_ids=[0],
|
|
filter_string='',
|
|
output_format='pandas',
|
|
)
|
|
|
|
assert output == '2'
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.mlflow')
|
|
def test_get_experiment_last_run_error(mlflow, mlflow_repository):
|
|
mlflow.search_runs.return_value = []
|
|
|
|
try:
|
|
mlflow_repository.get_experiment_last_run(0)
|
|
except ValueError as e:
|
|
assert str(e) == 'Runs is not a pandas DataFrame'
|
|
else:
|
|
raise AssertionError('Expected exception')
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.mlflow.sklearn')
|
|
@patch('model_manager.utils.repository.model_repository.mlflow.set_experiment')
|
|
def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
|
|
mlflow_repository.model_serving.get_model_info = MagicMock(return_value='0')
|
|
mlflow_repository.model_serving.get_model_uri = MagicMock(return_value='test')
|
|
mlflow_repository.get_experiment_by_run_id = MagicMock()
|
|
|
|
data_model_mock = MagicMock()
|
|
prediction_model_mock = MagicMock()
|
|
|
|
sklearn.load_model.side_effect = [data_model_mock, prediction_model_mock]
|
|
|
|
data_model_mock.fit.return_value = data_model_mock
|
|
data_model_mock.predict.return_value = DataFrame(
|
|
{
|
|
'x': [10, 20, 30],
|
|
}
|
|
)
|
|
data_model_mock.target_variable = 'y'
|
|
|
|
prediction_model_mock.fit.return_value = prediction_model_mock
|
|
|
|
data = DataFrame({'x': [1, 2, 3], 'y': [4, 5, 6]})
|
|
|
|
output = mlflow_repository.create_model_experiment('test', data)
|
|
|
|
mlflow_repository.model_serving.get_model_info.assert_called_once_with('test')
|
|
mlflow_repository.model_serving.get_model_uri.assert_called_once_with('0', prediction=False)
|
|
|
|
sklearn.load_model.assert_has_calls(
|
|
[
|
|
call(mlflow_repository.model_serving.get_model_uri.return_value),
|
|
call('models:/test/production'),
|
|
]
|
|
)
|
|
assert sklearn.load_model.call_count == 2
|
|
|
|
data_model_mock.fit.assert_called_once_with(data)
|
|
data_model_mock.predict.assert_called_once_with(data)
|
|
|
|
fit_args = prediction_model_mock.fit.call_args[0][0]
|
|
assert fit_args.equals(
|
|
DataFrame(
|
|
{
|
|
'x': [10, 20, 30],
|
|
'y': [4, 5, 6],
|
|
}
|
|
)
|
|
)
|
|
|
|
mlflow_repository.get_experiment_by_run_id.assert_called_once_with('0')
|
|
|
|
set_experiment.assert_called_once_with(mlflow_repository.get_experiment_by_run_id.return_value)
|
|
|
|
assert output == (
|
|
prediction_model_mock,
|
|
data_model_mock,
|
|
mlflow_repository.get_experiment_by_run_id.return_value,
|
|
)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.path.exists')
|
|
@patch('model_manager.utils.repository.model_repository.remove')
|
|
@patch('model_manager.utils.repository.model_repository.mlflow.start_run')
|
|
@patch('model_manager.utils.repository.model_repository.mlflow.log_param')
|
|
@patch('model_manager.utils.repository.model_repository.mlflow.sklearn.log_model')
|
|
@patch('model_manager.utils.repository.model_repository.mlflow.log_artifact')
|
|
def test_perform_model_retrain(
|
|
log_artifact, log_model, log_param, start_run, mock_remove, mock_path_exists, mlflow_repository
|
|
):
|
|
# Create mock models with attributes to test the for loops (lines 268-274)
|
|
prediction_model_mock = MagicMock()
|
|
prediction_model_mock.__dict__ = {'model': 'pred_model', 'param1': 'value1', 'param2': 'value2'}
|
|
|
|
data_model_mock = MagicMock()
|
|
data_model_mock.__dict__ = {'model': 'data_model', 'param3': 'value3', 'param4': 'value4'}
|
|
|
|
experiment = 'test'
|
|
model_name = 'test'
|
|
data = MagicMock()
|
|
|
|
mlflow_repository.get_next_run_name = MagicMock(return_value='test-1')
|
|
run = MagicMock()
|
|
start_run.__enter__.return_value = run
|
|
mock_path_exists.return_value = True
|
|
|
|
output = mlflow_repository.perform_model_retrain(
|
|
prediction_model_mock, data_model_mock, experiment, model_name, data
|
|
)
|
|
|
|
mlflow_repository.get_next_run_name.assert_called_once_with(experiment)
|
|
start_run.assert_called_once_with(
|
|
run_name='test-1', description='Retrain model test with new data'
|
|
)
|
|
|
|
log_model.assert_has_calls(
|
|
[
|
|
call(data_model_mock, 'data_model'),
|
|
call(prediction_model_mock, 'prediction_model'),
|
|
]
|
|
)
|
|
|
|
data.to_csv.assert_called_once_with('temp/raw_data_test.csv', index=True)
|
|
|
|
log_artifact.assert_called_once_with('temp/raw_data_test.csv')
|
|
|
|
# Verify that model attributes were logged (excluding 'model' key)
|
|
log_param.assert_has_calls(
|
|
[
|
|
call('param1', 'value1'), # from prediction_model
|
|
call('param2', 'value2'), # from prediction_model
|
|
call('param3', 'value3'), # from data_model
|
|
call('param4', 'value4'), # from data_model
|
|
call('retrain', True),
|
|
],
|
|
any_order=True,
|
|
)
|
|
|
|
# Verify temp file cleanup
|
|
mock_path_exists.assert_called_once_with('temp/raw_data_test.csv')
|
|
mock_remove.assert_called_once_with('temp/raw_data_test.csv')
|
|
|
|
assert output == ('Model retrained successfully', experiment)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.path.exists')
|
|
@patch('model_manager.utils.repository.model_repository.remove')
|
|
@patch('model_manager.utils.repository.model_repository.mlflow.start_run')
|
|
@patch('model_manager.utils.repository.model_repository.mlflow.log_param')
|
|
@patch('model_manager.utils.repository.model_repository.mlflow.sklearn.log_model')
|
|
@patch('model_manager.utils.repository.model_repository.mlflow.log_artifact')
|
|
def test_perform_model_retrain_file_not_exists(
|
|
log_artifact, log_model, log_param, start_run, mock_remove, mock_path_exists, mlflow_repository
|
|
):
|
|
"""Test perform_model_retrain when temp file doesn't exist (line 291->294 branch)."""
|
|
prediction_model_mock = MagicMock()
|
|
prediction_model_mock.__dict__ = {'model': 'pred_model'}
|
|
|
|
data_model_mock = MagicMock()
|
|
data_model_mock.__dict__ = {'model': 'data_model'}
|
|
|
|
experiment = 'test'
|
|
model_name = 'test'
|
|
data = MagicMock()
|
|
|
|
mlflow_repository.get_next_run_name = MagicMock(return_value='test-1')
|
|
run = MagicMock()
|
|
start_run.__enter__.return_value = run
|
|
mock_path_exists.return_value = False # File doesn't exist
|
|
|
|
output = mlflow_repository.perform_model_retrain(
|
|
prediction_model_mock, data_model_mock, experiment, model_name, data
|
|
)
|
|
|
|
# Verify temp file cleanup was checked but not executed
|
|
mock_path_exists.assert_called_once_with('temp/raw_data_test.csv')
|
|
mock_remove.assert_not_called() # Should not be called when file doesn't exist
|
|
|
|
assert output == ('Model retrained successfully', experiment)
|
|
|
|
|
|
def test_retrain_model(mlflow_repository):
|
|
data = MagicMock()
|
|
model_name = 'test'
|
|
|
|
mlflow_repository.create_model_experiment = MagicMock(
|
|
return_value=('data_model', 'prediction_model', '0')
|
|
)
|
|
|
|
mlflow_repository.perform_model_retrain = MagicMock(return_value='Model retrained successfully')
|
|
|
|
output = mlflow_repository.retrain_model(data, model_name)
|
|
|
|
mlflow_repository.create_model_experiment.assert_called_once_with(model_name, data)
|
|
|
|
mlflow_repository.perform_model_retrain.assert_called_once_with(
|
|
'data_model', 'prediction_model', '0', model_name, data
|
|
)
|
|
|
|
assert output == 'Model retrained successfully'
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.mlflow')
|
|
def test_update_production_model_by_run_id(mlflow, mlflow_repository):
|
|
client_mock = MagicMock()
|
|
mlflow.tracking.MlflowClient.return_value = client_mock
|
|
|
|
client_mock.get_registered_model.return_value = MagicMock(
|
|
latest_versions=[
|
|
MagicMock(version='1'),
|
|
MagicMock(version='2'),
|
|
MagicMock(version='3'),
|
|
]
|
|
)
|
|
output = mlflow_repository.update_production_model_by_run_id('0', 'test')
|
|
|
|
mlflow.register_model.assert_called_once_with(
|
|
'runs:/0/prediction_model',
|
|
'test',
|
|
)
|
|
|
|
mlflow.tracking.MlflowClient.assert_called_once()
|
|
client_mock.get_registered_model.assert_called_once_with('test')
|
|
client_mock.transition_model_version_stage.assert_called_once_with(
|
|
name='test',
|
|
version='3',
|
|
stage='Production',
|
|
archive_existing_versions=True,
|
|
)
|
|
|
|
assert output == {
|
|
'model_name': 'test',
|
|
'version': '3',
|
|
'mlflow_run_id': '0',
|
|
}
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.mlflow')
|
|
def test_update_production_model_by_run_id_error(mlflow, mlflow_repository):
|
|
mlflow.tracking.MlflowClient.return_value = MagicMock(
|
|
get_registered_model=MagicMock(return_value=MagicMock(latest_versions={}))
|
|
)
|
|
|
|
try:
|
|
mlflow_repository.update_production_model_by_run_id('0', 'test')
|
|
except Exception as e: # noqa: BLE001
|
|
assert str(e) == 'Model versions is not a list'
|
|
else:
|
|
raise AssertionError('Expected exception')
|
|
|
|
|
|
def test_update_production_model(mlflow_repository):
|
|
connector = mlflow_repository
|
|
|
|
with patch.object(connector, 'get_experiment', return_value='0') as get_experiment:
|
|
with patch.object(
|
|
connector, 'get_experiment_last_run', return_value='2'
|
|
) as get_experiment_last_run:
|
|
with patch.object(
|
|
connector,
|
|
'update_production_model_by_run_id',
|
|
return_value={'model_name': 'test', 'version': '3', 'mlflow_run_id': '0'},
|
|
) as update_production_model_by_run_id:
|
|
output = connector.update_production_model('0', 'test')
|
|
|
|
get_experiment.assert_called_once_with('0')
|
|
get_experiment_last_run.assert_called_once_with('0')
|
|
update_production_model_by_run_id.assert_called_once_with('2', 'test')
|
|
|
|
assert output == {
|
|
'model_name': 'test',
|
|
'version': '3',
|
|
'mlflow_run_id': '0',
|
|
'mlflow_experiment_id': '0',
|
|
}
|
|
|
|
|
|
# ========== Tests for Model Artifact Generation Methods ==========
|
|
|
|
|
|
def test_get_next_run_name_new(mlflow_repository):
|
|
"""Test get_next_run_name generates correct run name based on existing runs."""
|
|
mlflow_repository.model_serving.search_runs_by_name.return_value = [
|
|
MagicMock(),
|
|
MagicMock(),
|
|
MagicMock(),
|
|
]
|
|
|
|
result = mlflow_repository.get_next_run_name_new('test_experiment')
|
|
|
|
mlflow_repository.model_serving.search_runs_by_name.assert_called_once_with(
|
|
experiment_names=['test_experiment'], order_by=['start_time desc']
|
|
)
|
|
assert result == 'test_experiment-4'
|
|
|
|
|
|
def test_get_next_run_name_new_first_run(mlflow_repository):
|
|
"""Test get_next_run_name for first run (no existing runs)."""
|
|
mlflow_repository.model_serving.search_runs_by_name.return_value = []
|
|
|
|
result = mlflow_repository.get_next_run_name_new('new_experiment')
|
|
|
|
assert result == 'new_experiment-1'
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_generate_artifacts_success(mock_path, mlflow_repository):
|
|
"""Test generate_artifacts successfully creates all artifacts."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
# Mock data
|
|
params = MagicMock(spec=TrainModelParams)
|
|
params.target_variable = 'target'
|
|
params.variable_columns = ['feat1', 'feat2']
|
|
params.experiment_name = 'test_exp'
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.run_name = 'test_run-1'
|
|
data.params = params
|
|
data.x_train = DataFrame({'feat1': [1, 2], 'feat2': [3, 4]})
|
|
data.y_train = DataFrame({'target': [5, 6]})
|
|
data.x_test = DataFrame({'feat1': [7, 8], 'feat2': [9, 10]})
|
|
data.y_test = DataFrame({'target': [11, 12]})
|
|
data.regr = MagicMock()
|
|
data.regr.predict = MagicMock(return_value=np.array([5.1, 6.1]))
|
|
data.y_pred = np.array([11.1, 12.1])
|
|
|
|
# Mock path operations
|
|
mock_path.exists.return_value = True
|
|
mock_path.join.side_effect = lambda *args: '/'.join(args)
|
|
|
|
# Mock private methods
|
|
mlflow_repository._get_reports_directory = MagicMock(return_value='/reports')
|
|
mlflow_repository._create_run_directory = MagicMock(return_value='/reports/test_run-1_20231010')
|
|
mlflow_repository._setup_run_directory = MagicMock()
|
|
mlflow_repository._generate_report = MagicMock(return_value=data)
|
|
|
|
result = mlflow_repository.generate_artifacts(data)
|
|
|
|
# Assertions
|
|
mlflow_repository._get_reports_directory.assert_called_once()
|
|
mlflow_repository._create_run_directory.assert_called_once_with('/reports', 'test_run-1')
|
|
mlflow_repository._setup_run_directory.assert_called_once()
|
|
mlflow_repository._generate_report.assert_called_once()
|
|
assert result == data
|
|
|
|
|
|
def test_generate_artifacts_missing_run_name(mlflow_repository):
|
|
"""Test generate_artifacts raises ValueError when run_name is not set."""
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.run_name = None
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
mlflow_repository.generate_artifacts(data)
|
|
|
|
assert 'run_name must be set before generating artifacts' in str(exc_info.value)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_generate_artifacts_reports_directory_not_exists(mock_path, mlflow_repository):
|
|
"""Test generate_artifacts raises FileNotFoundError when reports directory doesn't exist."""
|
|
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.target_variable = 'target'
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.run_name = 'test_run-1'
|
|
data.params = params
|
|
data.x_train = DataFrame({'feat1': [1]})
|
|
data.y_train = DataFrame({'target': [2]})
|
|
data.x_test = DataFrame({'feat1': [3]})
|
|
data.y_test = DataFrame({'target': [4]})
|
|
data.regr = MagicMock()
|
|
data.y_pred = np.array([4.1])
|
|
|
|
mlflow_repository._get_reports_directory = MagicMock(return_value='/reports')
|
|
mock_path.exists.return_value = False
|
|
|
|
with pytest.raises(FileNotFoundError) as exc_info:
|
|
mlflow_repository.generate_artifacts(data)
|
|
|
|
assert 'Reports directory does not exist' in str(exc_info.value)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_generate_artifacts_header_file_not_exists(mock_path, mlflow_repository):
|
|
"""Test generate_artifacts raises FileNotFoundError when header.html doesn't exist."""
|
|
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.target_variable = 'target'
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.run_name = 'test_run-1'
|
|
data.params = params
|
|
data.x_train = DataFrame({'feat1': [1]})
|
|
data.y_train = DataFrame({'target': [2]})
|
|
data.x_test = DataFrame({'feat1': [3]})
|
|
data.y_test = DataFrame({'target': [4]})
|
|
data.regr = MagicMock()
|
|
data.regr.predict = MagicMock(return_value=np.array([2.1]))
|
|
data.y_pred = np.array([4.1])
|
|
|
|
mlflow_repository._get_reports_directory = MagicMock(return_value='/reports')
|
|
mlflow_repository._create_run_directory = MagicMock(return_value='/reports/test_run-1_20231010')
|
|
|
|
# First call returns True (reports dir exists), second returns False (header.html doesn't exist)
|
|
mock_path.exists.side_effect = [True, False]
|
|
mock_path.join.side_effect = lambda *args: '/'.join(args)
|
|
|
|
with pytest.raises(FileNotFoundError) as exc_info:
|
|
mlflow_repository.generate_artifacts(data)
|
|
|
|
assert 'Header file does not exist' in str(exc_info.value)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_save_run_success(mock_path, mlflow_repository):
|
|
"""Test save_run successfully logs all parameters, metrics, models, and 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.train_size = 80
|
|
params.removed_intervals = [(1, 10), (20, 30)]
|
|
params.experiment_name = 'test_exp'
|
|
params.target_variable = 'target'
|
|
params.variable_columns = ['feat1', 'feat2']
|
|
params.lag_train = 5
|
|
params.lag_val = 3
|
|
params.window = 10
|
|
params.low_lim = 0.0
|
|
params.upp_lim = 1.0
|
|
params.include_ar = True
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.run_name = 'test_run-1'
|
|
data.params = params
|
|
data.report_path = '/reports/report.html'
|
|
data.train_data_path = '/reports/train.csv'
|
|
data.test_data_path = '/reports/test.csv'
|
|
data.mse_val = 0.123
|
|
data.r2_val = 0.987
|
|
data.mae_val = 0.456
|
|
data.scaler_dict = {'scaler': 'minmax'}
|
|
data.process_data = MagicMock()
|
|
data.regr = MagicMock()
|
|
|
|
mock_path.exists.return_value = True
|
|
|
|
mlflow_repository.save_run(data)
|
|
|
|
# Verify experiment was set
|
|
mlflow_repository.model_serving.set_experiment.assert_called_once_with('test_exp')
|
|
|
|
# Verify parameters were logged
|
|
assert mlflow_repository.model_serving.log_param.call_count == 13
|
|
|
|
# Verify metrics were logged
|
|
mlflow_repository.model_serving.log_metric.assert_any_call('MSE', 0.123)
|
|
mlflow_repository.model_serving.log_metric.assert_any_call('R2', 0.987)
|
|
mlflow_repository.model_serving.log_metric.assert_any_call('MAE', 0.456)
|
|
|
|
# Verify models were logged
|
|
mlflow_repository.model_serving.log_model.assert_any_call(data.process_data, 'data_model')
|
|
mlflow_repository.model_serving.log_model.assert_any_call(data.regr, 'prediction_model')
|
|
|
|
# Verify artifacts were logged
|
|
mlflow_repository.model_serving.log_artifact.assert_any_call('/reports/report.html')
|
|
mlflow_repository.model_serving.log_artifact.assert_any_call('/reports/train.csv')
|
|
mlflow_repository.model_serving.log_artifact.assert_any_call('/reports/test.csv')
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_save_run_missing_report_path(mock_path, mlflow_repository):
|
|
"""Test save_run raises ValueError when report_path is missing."""
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.report_path = None
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
mlflow_repository.save_run(data)
|
|
|
|
assert 'Report file does not exist' in str(exc_info.value)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_save_run_missing_metrics(mock_path, mlflow_repository):
|
|
"""Test save_run raises ValueError when metrics are None."""
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.report_path = '/reports/report.html'
|
|
data.train_data_path = '/reports/train.csv'
|
|
data.test_data_path = '/reports/test.csv'
|
|
data.mse_val = None
|
|
data.r2_val = 0.987
|
|
data.mae_val = 0.456
|
|
|
|
mock_path.exists.return_value = True
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
mlflow_repository.save_run(data)
|
|
|
|
assert 'One or more metrics (MSE, R2, MAE) are None' in str(exc_info.value)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_save_run_mlflow_error(mock_path, mlflow_repository):
|
|
"""Test save_run handles MLflow errors 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.train_size = 80
|
|
params.removed_intervals = []
|
|
params.experiment_name = 'test_exp'
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.run_name = 'test_run-1'
|
|
data.params = params
|
|
data.report_path = '/reports/report.html'
|
|
data.train_data_path = '/reports/train.csv'
|
|
data.test_data_path = '/reports/test.csv'
|
|
data.mse_val = 0.123
|
|
data.r2_val = 0.987
|
|
data.mae_val = 0.456
|
|
|
|
mock_path.exists.return_value = True
|
|
mlflow_repository.model_serving.set_experiment.side_effect = Exception(
|
|
'MLflow connection error'
|
|
)
|
|
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
mlflow_repository.save_run(data)
|
|
|
|
assert 'Failed to save run' in str(exc_info.value)
|
|
assert 'MLflow connection error' in str(exc_info.value)
|
|
|
|
|
|
# ========== Additional Tests for 100% Coverage ==========
|
|
|
|
|
|
def test_init_artifacts_data_empty_x_train(mlflow_repository):
|
|
"""Test _init_artifacts_data raises ValueError when x_train is empty."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.params = params
|
|
data.x_train = DataFrame() # Empty DataFrame
|
|
data.y_train = DataFrame({'target': [1]})
|
|
data.x_test = DataFrame({'feat1': [1]})
|
|
data.y_test = DataFrame({'target': [1]})
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
mlflow_repository._init_artifacts_data(data)
|
|
|
|
assert 'Training features (x_train) are empty' in str(exc_info.value)
|
|
|
|
|
|
def test_init_artifacts_data_empty_y_train(mlflow_repository):
|
|
"""Test _init_artifacts_data raises ValueError when y_train is empty."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.params = params
|
|
data.x_train = DataFrame({'feat1': [1]})
|
|
data.y_train = DataFrame() # Empty DataFrame
|
|
data.x_test = DataFrame({'feat1': [1]})
|
|
data.y_test = DataFrame({'target': [1]})
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
mlflow_repository._init_artifacts_data(data)
|
|
|
|
assert 'Training target (y_train) is empty' in str(exc_info.value)
|
|
|
|
|
|
def test_init_artifacts_data_empty_x_test(mlflow_repository):
|
|
"""Test _init_artifacts_data raises ValueError when x_test is empty."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.params = params
|
|
data.x_train = DataFrame({'feat1': [1]})
|
|
data.y_train = DataFrame({'target': [1]})
|
|
data.x_test = DataFrame() # Empty DataFrame
|
|
data.y_test = DataFrame({'target': [1]})
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
mlflow_repository._init_artifacts_data(data)
|
|
|
|
assert 'Test features (x_test) are empty' in str(exc_info.value)
|
|
|
|
|
|
def test_init_artifacts_data_empty_y_test(mlflow_repository):
|
|
"""Test _init_artifacts_data raises ValueError when y_test is empty."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.params = params
|
|
data.x_train = DataFrame({'feat1': [1]})
|
|
data.y_train = DataFrame({'target': [1]})
|
|
data.x_test = DataFrame({'feat1': [1]})
|
|
data.y_test = DataFrame() # Empty DataFrame
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
mlflow_repository._init_artifacts_data(data)
|
|
|
|
assert 'Test target (y_test) is empty' in str(exc_info.value)
|
|
|
|
|
|
def test_init_artifacts_data_none_y_pred(mlflow_repository):
|
|
"""Test _init_artifacts_data raises ValueError when y_pred is None."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.params = params
|
|
data.x_train = DataFrame({'feat1': [1]})
|
|
data.y_train = DataFrame({'target': [1]})
|
|
data.x_test = DataFrame({'feat1': [1]})
|
|
data.y_test = DataFrame({'target': [1]})
|
|
data.y_pred = None
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
mlflow_repository._init_artifacts_data(data)
|
|
|
|
assert 'Test predictions (y_pred) are None' in str(exc_info.value)
|
|
|
|
|
|
def test_init_artifacts_data_success(mlflow_repository):
|
|
"""Test _init_artifacts_data successfully prepares data."""
|
|
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.target_variable = 'target'
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.params = params
|
|
data.x_train = DataFrame({'feat1': [1, 2]})
|
|
data.y_train = DataFrame({'target': [3, 4]})
|
|
data.x_test = DataFrame({'feat1': [5, 6]})
|
|
data.y_test = DataFrame({'target': [7, 8]})
|
|
data.regr = MagicMock()
|
|
data.regr.predict = MagicMock(return_value=np.array([3.1, 4.1]))
|
|
data.y_pred = np.array([7.1, 8.1])
|
|
|
|
reference_data, current_data = mlflow_repository._init_artifacts_data(data)
|
|
|
|
assert 'target' in reference_data.columns
|
|
assert 'prediction' in reference_data.columns
|
|
assert 'target' in current_data.columns
|
|
assert 'prediction' in current_data.columns
|
|
assert len(reference_data) == 2
|
|
assert len(current_data) == 2
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.makedirs')
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_create_run_directory_success(mock_path, mock_makedirs, mlflow_repository):
|
|
"""Test _create_run_directory successfully creates directory."""
|
|
mock_path.join.return_value = '/reports/test_run_20231010_123456_123456'
|
|
|
|
result = mlflow_repository._create_run_directory('/reports', 'test_run')
|
|
|
|
mock_makedirs.assert_called_once_with('/reports/test_run_20231010_123456_123456', exist_ok=True)
|
|
assert result == '/reports/test_run_20231010_123456_123456'
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.makedirs')
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_create_run_directory_permission_error(mock_path, mock_makedirs, mlflow_repository):
|
|
"""Test _create_run_directory handles PermissionError."""
|
|
mock_path.join.return_value = '/reports/test_run_20231010'
|
|
mock_makedirs.side_effect = PermissionError('Permission denied')
|
|
|
|
with pytest.raises(PermissionError) as exc_info:
|
|
mlflow_repository._create_run_directory('/reports', 'test_run')
|
|
|
|
assert 'Permission denied when creating directory' in str(exc_info.value)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.makedirs')
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_create_run_directory_os_error(mock_path, mock_makedirs, mlflow_repository):
|
|
"""Test _create_run_directory handles OSError."""
|
|
mock_path.join.return_value = '/reports/test_run_20231010'
|
|
mock_makedirs.side_effect = OSError('Disk full')
|
|
|
|
with pytest.raises(OSError) as exc_info:
|
|
mlflow_repository._create_run_directory('/reports', 'test_run')
|
|
|
|
assert 'Failed to create directory' in str(exc_info.value)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.shutil')
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_setup_run_directory_success(mock_path, mock_shutil, mlflow_repository):
|
|
"""Test _setup_run_directory successfully sets up directory."""
|
|
mock_path.join.side_effect = lambda *args: '/'.join(args)
|
|
mock_open = MagicMock()
|
|
|
|
with patch('builtins.open', mock_open):
|
|
mlflow_repository._setup_run_directory('/run_dir', '/reports/header.html')
|
|
|
|
assert mock_open.call_count == 3 # 3 empty files
|
|
mock_shutil.copy.assert_called_once_with('/reports/header.html', '/run_dir/header.html')
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.shutil')
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_setup_run_directory_file_not_found(mock_path, mock_shutil, mlflow_repository):
|
|
"""Test _setup_run_directory handles FileNotFoundError."""
|
|
mock_path.join.side_effect = lambda *args: '/'.join(args)
|
|
mock_shutil.copy.side_effect = FileNotFoundError('Header not found')
|
|
|
|
mock_open = MagicMock()
|
|
with patch('builtins.open', mock_open):
|
|
with pytest.raises(FileNotFoundError) as exc_info:
|
|
mlflow_repository._setup_run_directory('/run_dir', '/reports/header.html')
|
|
|
|
assert 'Header file not found' in str(exc_info.value)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.shutil')
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_setup_run_directory_permission_error(mock_path, mock_shutil, mlflow_repository):
|
|
"""Test _setup_run_directory handles PermissionError."""
|
|
mock_path.join.side_effect = lambda *args: '/'.join(args)
|
|
|
|
mock_open = MagicMock()
|
|
mock_open.side_effect = PermissionError('Permission denied')
|
|
|
|
with patch('builtins.open', mock_open):
|
|
with pytest.raises(PermissionError) as exc_info:
|
|
mlflow_repository._setup_run_directory('/run_dir', '/reports/header.html')
|
|
|
|
assert 'Permission denied when setting up directory' in str(exc_info.value)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.shutil')
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_setup_run_directory_os_error(mock_path, mock_shutil, mlflow_repository):
|
|
"""Test _setup_run_directory handles OSError."""
|
|
mock_path.join.side_effect = lambda *args: '/'.join(args)
|
|
|
|
mock_open = MagicMock()
|
|
mock_open.side_effect = OSError('Disk error')
|
|
|
|
with patch('builtins.open', mock_open):
|
|
with pytest.raises(OSError) as exc_info:
|
|
mlflow_repository._setup_run_directory('/run_dir', '/reports/header.html')
|
|
|
|
assert 'Failed to setup run directory' in str(exc_info.value)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.Reports')
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_generate_report_success(mock_path, mock_reports, mlflow_repository):
|
|
"""Test _generate_report successfully generates all reports."""
|
|
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.variable_columns = ['feat1', 'feat2']
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.params = params
|
|
data.run_dir = '/run_dir'
|
|
|
|
reference_data = DataFrame(
|
|
{'feat1': [1.0], 'feat2': [2.0], 'target': [3.0], 'prediction': [3.1]}
|
|
)
|
|
current_data = DataFrame({'feat1': [4.0], 'feat2': [5.0], 'target': [6.0], 'prediction': [6.1]})
|
|
|
|
mock_path.join.side_effect = lambda *args: '/'.join(args)
|
|
mock_report_instance = MagicMock()
|
|
mock_reports.return_value = mock_report_instance
|
|
|
|
# Mock DataFrame.to_csv to avoid actual file writing
|
|
with patch.object(DataFrame, 'to_csv'):
|
|
result = mlflow_repository._generate_report(reference_data, current_data, data)
|
|
|
|
mock_reports.assert_called_once()
|
|
mock_report_instance.add_data_quality_section.assert_called_once()
|
|
mock_report_instance.add_data_drift_section.assert_called_once()
|
|
mock_report_instance.add_regression_section.assert_called_once()
|
|
mock_report_instance.save_all_sections_html.assert_called_once()
|
|
assert result == data
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_generate_report_value_error(mock_path, mlflow_repository):
|
|
"""Test _generate_report handles ValueError from data conversion."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.params = params
|
|
data.run_dir = '/run_dir'
|
|
|
|
# DataFrame with non-numeric data
|
|
reference_data = DataFrame({'feat1': ['a', 'b']})
|
|
current_data = DataFrame({'feat1': ['c', 'd']})
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
mlflow_repository._generate_report(reference_data, current_data, data)
|
|
|
|
assert 'Failed to convert data to float64' in str(exc_info.value)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.Reports')
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_generate_report_permission_error(mock_path, mock_reports, mlflow_repository):
|
|
"""Test _generate_report handles PermissionError."""
|
|
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.variable_columns = ['feat1']
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.params = params
|
|
data.run_dir = '/run_dir'
|
|
|
|
reference_data = DataFrame({'feat1': [1.0]})
|
|
current_data = DataFrame({'feat1': [2.0]})
|
|
|
|
mock_path.join.side_effect = lambda *args: '/'.join(args)
|
|
mock_report_instance = MagicMock()
|
|
mock_reports.return_value = mock_report_instance
|
|
mock_report_instance.save_all_sections_html.side_effect = PermissionError('Permission denied')
|
|
|
|
with pytest.raises(PermissionError) as exc_info:
|
|
mlflow_repository._generate_report(reference_data, current_data, data)
|
|
|
|
assert 'Permission denied when writing report files' in str(exc_info.value)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.Reports')
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_generate_report_os_error(mock_path, mock_reports, mlflow_repository):
|
|
"""Test _generate_report handles OSError."""
|
|
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.variable_columns = ['feat1']
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.params = params
|
|
data.run_dir = '/run_dir'
|
|
|
|
reference_data = DataFrame({'feat1': [1.0]})
|
|
current_data = DataFrame({'feat1': [2.0]})
|
|
|
|
mock_path.join.side_effect = lambda *args: '/'.join(args)
|
|
mock_report_instance = MagicMock()
|
|
mock_reports.return_value = mock_report_instance
|
|
mock_report_instance.save_all_sections_html.side_effect = OSError('Disk error')
|
|
|
|
with pytest.raises(OSError) as exc_info:
|
|
mlflow_repository._generate_report(reference_data, current_data, data)
|
|
|
|
assert 'Failed to generate report' in str(exc_info.value)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.Reports')
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_generate_report_run_dir_none(mock_path, mock_reports, mlflow_repository):
|
|
"""Test _generate_report raises ValueError when run_dir is None."""
|
|
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.variable_columns = ['feat1']
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.params = params
|
|
data.run_dir = None # Not set
|
|
|
|
reference_data = DataFrame({'feat1': [1.0]})
|
|
current_data = DataFrame({'feat1': [2.0]})
|
|
|
|
mock_report_instance = MagicMock()
|
|
mock_reports.return_value = mock_report_instance
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
mlflow_repository._generate_report(reference_data, current_data, data)
|
|
|
|
assert 'run_dir is not set after directory creation' in str(exc_info.value)
|
|
|
|
|
|
def test_get_reports_directory(mlflow_repository):
|
|
"""Test _get_reports_directory returns correct path."""
|
|
result = mlflow_repository._get_reports_directory()
|
|
|
|
assert result.endswith('model_manager/reports')
|
|
assert 'model_manager' in result
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_save_run_missing_train_data_path(mock_path, mlflow_repository):
|
|
"""Test save_run raises ValueError when train_data_path is missing."""
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.report_path = '/reports/report.html'
|
|
data.train_data_path = None
|
|
|
|
mock_path.exists.return_value = True
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
mlflow_repository.save_run(data)
|
|
|
|
assert 'Training data file does not exist' in str(exc_info.value)
|
|
|
|
|
|
@patch('model_manager.utils.repository.model_repository.path')
|
|
def test_save_run_missing_test_data_path(mock_path, mlflow_repository):
|
|
"""Test save_run raises ValueError when test_data_path is missing."""
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
data = MagicMock(spec=TrainModelResult)
|
|
data.report_path = '/reports/report.html'
|
|
data.train_data_path = '/reports/train.csv'
|
|
data.test_data_path = None
|
|
|
|
mock_path.exists.return_value = True
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
mlflow_repository.save_run(data)
|
|
|
|
assert 'Test data file does not exist' in str(exc_info.value)
|